code: def test_all_to_all_1d(self):
      # transpose on a 2D tensor distributed over N nodes:
      mesh = DeviceMesh(self.device_type, torch.arange(self.world_size))
      tensor_shape = [3, 3]
      input_tensor_list = [
          torch.ones(*tensor_shape, device=self.device_type)
          * (rank + self.rank * self.world_size)
          for rank in range(self.world_size)
      ]
      expected_tensor_list = [
          torch.ones(tensor_shape, device=self.device_type)
          * (self.rank + rank * self.world_size)  # i.e. transpose
          for rank in range(self.world_size)
      ]
      for scatter_dim in range(len(tensor_shape)):
          output_tensor_list = [
              torch.empty_like(input_tensor_list[idx])
              for idx in range(len(input_tensor_list))
          ]
          # scatter on dim > 0 would generate non-contiguous tensor, verify that works
          mesh_all_to_all(
              output_tensor_list, input_tensor_list, mesh, mesh_dim=0
          )
          output_tensor = torch.cat(output_tensor_list, dim=scatter_dim)
          expected_tensor = torch.cat(
              expected_tensor_list, dim=scatter_dim
          )
  
          self.assertEqual(output_tensor, expected_tensor)
  
imports: import os
  import torch
  import torch.distributed._functional_collectives
  from torch.distributed._tensor import DTensor
  from torch.distributed._tensor._collective_utils import mesh_all_to_all, mesh_broadcast, mesh_scatter
  from torch.distributed._tensor.placement_types import _Partial, Shard
  from torch.distributed.device_mesh import _mesh_resources, DeviceMesh, init_device_mesh
  from torch.distributed.distributed_c10d import get_global_rank, get_world_size, init_process_group, is_initialized, is_nccl_available, ProcessGroup
  from torch.testing._internal.common_utils import run_tests
  from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase, with_comms
  from torch.testing._internal.distributed.fake_pg import FakeStore
label_desc: Use image node_id 10 for calling the DeviceMeshCollectiveTest obj's underlying member method code with example usage: obj.test_all_to_all_1d() without return types