code: def run():
      group = ProcessGroup(tp_degree=dist.get_world_size())
      rank = dist.get_rank()
      size = dist.get_world_size()
      depth = int(math.sqrt(size))
      assert depth == math.sqrt(size)
      x = torch.rand(8, 8).cuda()
      old_dist_spec = ReplicaSpec()
      row_spec = ShardSpec([0], [size])
      col_spec = ShardSpec([-1], [size])
      mat_spec = ShardSpec([0, 1], [depth, depth])
      row_shard = DistSpecManager._shard_as(
          x, old_dist_spec, row_spec, group
      )
      assert torch.equal(x.chunk(size, 0)[rank], row_shard)
      assert torch.equal(
          x, DistSpecManager._gather(row_shard, row_spec, group)
      )
      col_shard = DistSpecManager._all_to_all(
          row_shard, row_spec, col_spec, group
      )
      assert torch.equal(x.chunk(size, -1)[rank], col_shard)
      assert torch.equal(
          x, DistSpecManager._gather(col_shard, col_spec, group)
      )
      mat_shard = DistSpecManager._shard_as(
          x, old_dist_spec, mat_spec, group
      )
      assert torch.equal(
          x.chunk(depth, 0)[rank // depth].chunk(depth, 1)[
              rank % depth
          ],
          mat_shard,
      )
      assert torch.equal(
          x, DistSpecManager._gather(mat_shard, mat_spec, group)
      )
  
file_path: ColossalAI/tests/test_legacy/test_tensor/core/test_dist_spec_mgr.py
imports: import math
  import pytest
  import torch
  import torch.distributed
  import colossalai
  from colossalai.legacy.tensor import DistSpecManager, ProcessGroup, ReplicaSpec, ShardSpec
  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: run() without return types