code: def check_mix_gather_S1S0(device_mesh, rank):
      tensor_to_check = torch.arange(64).reshape((8, 8)).cuda()
      (f, b) = (0, 1)
      f_target_pair = (f, [1])
      b_target_pair = (b, [0])
      gather_dim, logical_process_axes = mix_gather_simulator(
          f_target_pair, b_target_pair
      )
      tensor_slice = [2, 4]
      rank_slice = 4
      f_start = (rank % rank_slice) * tensor_slice[0]
      b_start = (rank // rank_slice) * tensor_slice[1]
      tensor_to_comm = (
          tensor_to_check[
              f_start : f_start + tensor_slice[0],
              b_start : b_start + tensor_slice[1],
          ]
          .contiguous()
          .cuda()
      )
  
      dim_partition_dict = {0: [1], 1: [0]}
  
      # DistSpec:
      #     shard_sequence: S1,S0
      #     device_mesh_shape: (2, 4)
      source_spec = ShardingSpec(
          device_mesh,
          tensor_to_check.shape,
          dim_partition_dict=dim_partition_dict,
      )
  
      comm_spec = CommSpec(
          CollectiveCommPattern.MIXGATHER_FWD_SPLIT_BWD,
          sharding_spec=source_spec,
          gather_dim=gather_dim,
          logical_process_axis=logical_process_axes,
          forward_only=True,
          mix_gather=True,
      )
      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_mix_gather.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.tensor.utils import mix_gather_simulator
  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_mix_gather_S1S0(device_mesh, rank) without return types