code: def check_comm(rank, world_size, port):
      disable_existing_loggers()
      launch(
          config={},
          rank=rank,
          world_size=world_size,
          host="localhost",
          port=port,
          backend="nccl",
      )
  
      physical_mesh_id = torch.arange(0, 8)
      assert rank == dist.get_rank()
  
      mesh_shape = (2, 4)
      # [[0, 1, 2, 3],
      #  [4, 5, 6, 7]]
      device_mesh = DeviceMesh(
          physical_mesh_id,
          mesh_shape,
          init_process_group=True,
          need_flatten=True,
      )
  
      check_mix_gather_S0S1(device_mesh, rank)
  
      check_two_all_gather_S0S1(device_mesh, rank)
  
      check_mix_gather_S1S0(device_mesh, rank)
  
      check_two_all_gather_S1S0(device_mesh, rank)
  
      check_mix_gather_S01R(device_mesh, rank)
  
      check_two_all_gather_S01R(device_mesh, rank)
  
      check_mix_gather_RS01(device_mesh, rank)
  
      check_two_all_gather_RS01(device_mesh, rank)
  
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 9 for calling a global function with example usage: check_comm(rank, world_size, port) without return types