code: def test_device_mesh():
      physical_mesh_id = torch.arange(0, 16)
      mesh_shape = (4, 4)
      # [[0, 1, 2, 3],
      #  [4, 5, 6, 7],
      #  [8, 9, 10,11],
      #  [12,13,14,15]]
      device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
      assert device_mesh.global_rank_to_local_rank(5) == [1, 1]
      assert device_mesh.global_rank_to_local_rank(11) == [2, 3]
      assert device_mesh.get_ranks_in_process_group(
          axis=1, global_rank=2
      ) == [0, 1, 2, 3]
  
file_path: ColossalAI/tests/test_device/test_device_mesh.py
imports: import pytest
  import torch
  import torch.distributed
  import colossalai
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.testing import rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 1 for calling a global function with example usage: test_device_mesh() without return types