Home
last modified time | relevance | path

Searched refs:mesh_tensor (Results 1 – 6 of 6) sorted by relevance

/aosp_15_r20/external/pytorch/test/distributed/
H A Dtest_device_mesh.py77 mesh_tensor = torch.arange(4).reshape(2, 2)
80 DeviceMesh(device_type, mesh_tensor)
145 mesh_tensor = torch.arange(4).reshape(2, 2)
147 mesh = DeviceMesh(self.device_type, mesh_tensor)
237 mesh_tensor = torch.arange(4).reshape(2, 2)
238 mesh = DeviceMesh(device_type, mesh_tensor)
250 mesh_tensor = torch.arange(8).reshape(2, 2, 2)
251 mesh = DeviceMesh(self.device_type, mesh_tensor)
257 self.assertTrue(dim < mesh_tensor.ndim)
258 dim_ranks = mesh_tensor.swapdims(-1, dim).reshape(-1, 2)
[all …]
/aosp_15_r20/external/pytorch/test/distributed/_tensor/
H A Dtest_utils.py33 mesh_tensor = torch.arange(self.world_size).reshape(4, 2)
34 mesh = DeviceMesh(self.device_type, mesh_tensor)
70 mesh_tensor = torch.arange(self.world_size)
71 device_mesh = DeviceMesh(self.device_type, mesh_tensor)
102 mesh_tensor = torch.arange(self.world_size).reshape(2, 4)
103 device_mesh = DeviceMesh(self.device_type, mesh_tensor)
H A Dtest_dtensor.py616 mesh_tensor = torch.arange(self.world_size).reshape(2, 4)
618 mesh = DeviceMesh(self.device_type, mesh_tensor)
640 mesh_tensor = torch.arange(self.world_size).reshape(2, 2, 2)
641 mesh = DeviceMesh(self.device_type, mesh_tensor)
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Ddtensor_sharded_prefix_kernel.cc40 const Tensor* mesh_tensor; in Compute() local
46 OP_REQUIRES_OK(ctx, ctx->input("mesh", &mesh_tensor)); in Compute()
59 const std::string& mesh_str = mesh_tensor->scalar<tstring>()(); in Compute()
/aosp_15_r20/external/pytorch/torch/distributed/
H A Ddevice_mesh.py102 mesh_tensor = device_mesh.mesh
114 mesh_tensor = mesh_tensor.flatten(
135 mesh_dims_remained_idx = list(range(mesh_tensor.ndim))
142 pg_ranks_by_dim = mesh_tensor.permute(
/aosp_15_r20/external/pytorch/test/distributed/_composable/fsdp/
H A Dtest_fully_shard_init.py872 mesh_tensor = torch.tensor(
881 mesh=mesh_tensor,