code: def check_all_to_all(device_mesh, rank):
      # tensor to comm
      if rank in (0, 1):
          sharded_tensor_0 = torch.zeros(2, 1)
          sharded_tensor_1 = torch.ones(2, 1)
          # tensor([[0., 1.],
          #         [0., 1.]])
          tensor_to_comm = torch.cat(
              (sharded_tensor_0, sharded_tensor_1), 1
          ).cuda()
      if rank in (2, 3):
          sharded_tensor_0 = torch.ones(2, 1) * 2
          sharded_tensor_1 = torch.ones(2, 1) * 3
          # tensor([[2., 3.],
          #         [2., 3.]])
          tensor_to_comm = torch.cat(
              (sharded_tensor_0, sharded_tensor_1), 1
          ).cuda()
  
      if rank in (0, 1):
          # tensor([[0.],
          #         [0.],
          #         [2.],
          #         [2.]])
          tensor_to_check = torch.tensor(
              [[0], [0], [2], [2]], dtype=tensor_to_comm.dtype
          ).cuda()
      if rank in (2, 3):
          # tensor([[1.],
          #         [1.],
          #         [3.],
          #         [3.]])
          tensor_to_check = torch.tensor(
              [[1], [1], [3], [3]], dtype=tensor_to_comm.dtype
          ).cuda()
  
      # test shard
      dim_partition_dict = {0: [0]}
  
      # DistSpec:
      #     shard_sequence: S0,R
      #     device_mesh_shape: (2, 2)
      sharding_spec = ShardingSpec(
          device_mesh,
          torch.Size((4, 2)),
          dim_partition_dict=dim_partition_dict,
      )
  
      # CommSpec:(comm_pattern:shard, shard_dim:1, logical_process_axis:1)
      comm_spec = CommSpec(
          CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD,
          sharding_spec,
          gather_dim=0,
          shard_dim=1,
          logical_process_axis=0,
      )
      tensor_to_comm = comm_spec.covert_spec_to_action(tensor_to_comm)
  
      assert tensor_to_comm.equal(tensor_to_check)
  
file_path: ColossalAI/tests/test_tensor/test_comm_spec_apply.py
imports: import pytest
  import torch
  import torch.distributed
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.initialize import launch
  from colossalai.logging import disable_existing_loggers
  from colossalai.tensor.shape_consistency import CollectiveCommPattern, CommSpec
  from colossalai.tensor.sharding_spec import ShardingSpec
  from colossalai.testing import rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 3 for calling a global function with example usage: check_all_to_all(device_mesh, rank) without return types