code: def check_process_group_mesh_with_gpc():
      from colossalai.legacy.context import ParallelMode
      from colossalai.legacy.core import global_context as gpc
  
      DP_DIM, PP_DIM, TP_DIM = 0, 1, 2
      pg_mesh = ProcessGroupMesh(1, 2, 2)
  
      # check world size
      assert gpc.get_world_size(ParallelMode.TENSOR) == pg_mesh.size(
          TP_DIM
      ), f"{gpc.get_world_size(ParallelMode.TENSOR)} != {pg_mesh.size(TP_DIM)}"
      assert gpc.get_world_size(ParallelMode.PIPELINE) == pg_mesh.size(
          PP_DIM
      )
      assert gpc.get_world_size(ParallelMode.DATA) == pg_mesh.size(
          DP_DIM
      )
  
      # check locak rank (coordinate)
      assert gpc.get_local_rank(
          ParallelMode.TENSOR
      ) == pg_mesh.coordinate(
          TP_DIM
      ), f"{gpc.get_local_rank(ParallelMode.TENSOR)} != {pg_mesh.coordinate(TP_DIM)}"
      assert gpc.get_local_rank(
          ParallelMode.PIPELINE
      ) == pg_mesh.coordinate(PP_DIM)
      assert gpc.get_local_rank(
          ParallelMode.DATA
      ) == pg_mesh.coordinate(DP_DIM)
  
      # check ranks in group
      tp_group = pg_mesh.get_group_along_axis(TP_DIM)
      assert gpc.get_ranks_in_group(
          ParallelMode.TENSOR
      ) == pg_mesh.get_ranks_in_group(tp_group)
      pp_group = pg_mesh.get_group_along_axis(PP_DIM)
      assert gpc.get_ranks_in_group(
          ParallelMode.PIPELINE
      ) == pg_mesh.get_ranks_in_group(pp_group)
      dp_group = pg_mesh.get_group_along_axis(DP_DIM)
      assert gpc.get_ranks_in_group(
          ParallelMode.DATA
      ) == pg_mesh.get_ranks_in_group(dp_group)
  
      # check prev rank
      coord = pg_mesh.coordinate()
      if not gpc.is_first_rank(ParallelMode.TENSOR):
          assert coord[TP_DIM] != 0
          prev_coord = (
              coord[:TP_DIM]
              + (coord[TP_DIM] - 1,)
              + coord[TP_DIM + 1 :]
          )
          assert gpc.get_prev_global_rank(
              ParallelMode.TENSOR
          ) == pg_mesh.ravel(prev_coord, pg_mesh.shape)
      if not gpc.is_first_rank(ParallelMode.PIPELINE):
          assert coord[PP_DIM] != 0
          prev_coord = (
              coord[:PP_DIM]
              + (coord[PP_DIM] - 1,)
              + coord[PP_DIM + 1 :]
          )
          assert gpc.get_prev_global_rank(
              ParallelMode.PIPELINE
          ) == pg_mesh.ravel(prev_coord, pg_mesh.shape)
  
      # check next rank
      if not gpc.is_last_rank(ParallelMode.TENSOR):
          assert coord[TP_DIM] != pg_mesh.size(TP_DIM) - 1
          next_coord = (
              coord[:TP_DIM]
              + (coord[TP_DIM] + 1,)
              + coord[TP_DIM + 1 :]
          )
          assert gpc.get_next_global_rank(
              ParallelMode.TENSOR
          ) == pg_mesh.ravel(next_coord, pg_mesh.shape)
      if not gpc.is_last_rank(ParallelMode.PIPELINE):
          assert coord[PP_DIM] != pg_mesh.size(PP_DIM) - 1
          next_coord = (
              coord[:PP_DIM]
              + (coord[PP_DIM] + 1,)
              + coord[PP_DIM + 1 :]
          )
          assert gpc.get_next_global_rank(
              ParallelMode.PIPELINE
          ) == pg_mesh.ravel(next_coord, pg_mesh.shape)
  
file_path: ColossalAI/tests/test_cluster/test_process_group_mesh.py
imports: import pytest
  import torch.distributed
  import colossalai
  from colossalai.cluster import ProcessGroupMesh
  from colossalai.testing import spawn
label_desc: Use image node_id 1 for calling a global function with example usage: check_process_group_mesh_with_gpc() without return types