code: def check_dtensor(rank, world_size, port):
      disable_existing_loggers()
      launch(
          config={},
          rank=rank,
          world_size=world_size,
          host="localhost",
          port=port,
          backend="nccl",
      )
      test_model = TestModel(8, 8).to("cuda")
      original_tensor = torch.rand(4, 8).to("cuda")
      compare_output = test_model(original_tensor)
  
      device_mesh = DeviceMesh(
          torch.Tensor([0, 1, 2, 3]), (2, 2), init_process_group=True
      )
      target_sharding_spec = ShardingSpec(
          dim_size=original_tensor.dim(), dim_partition_dict={0: [0]}
      )
      d_tensor = distribute_tensor(
          original_tensor, device_mesh, target_sharding_spec
      )
  
      assert get_global_shape(d_tensor) == original_tensor.shape
      assert d_tensor.dtype == original_tensor.dtype
  
      if rank in (0, 1):
          assert d_tensor.equal(original_tensor.narrow(0, 0, 2))
      elif rank in (2, 3):
          assert d_tensor.equal(original_tensor.narrow(0, 2, 2))
      else:
          raise ValueError(f"rank {rank} is not in the device mesh")
      assert to_global(d_tensor).equal(original_tensor)
      output = test_model(d_tensor)
  
      if rank in (0, 1):
          assert output.equal(compare_output.narrow(0, 0, 2))
      elif rank in (2, 3):
          assert output.equal(compare_output.narrow(0, 2, 2))
      else:
          raise ValueError(f"rank {rank} is not in the device mesh")
  
      new_sharding_spec = ShardingSpec(
          dim_size=original_tensor.dim(), dim_partition_dict={0: [0, 1]}
      )
      d_tensor = redistribute(d_tensor, device_mesh, new_sharding_spec)
  
      if rank == 0:
          assert d_tensor.equal(original_tensor.narrow(0, 0, 1))
      elif rank == 1:
          assert d_tensor.equal(original_tensor.narrow(0, 1, 1))
      elif rank == 2:
          assert d_tensor.equal(original_tensor.narrow(0, 2, 1))
      elif rank == 3:
          assert d_tensor.equal(original_tensor.narrow(0, 3, 1))
      else:
          raise ValueError(f"rank {rank} is not in the device mesh")
  
      dtensor_from_local = distribute_tensor(
          original_tensor, device_mesh, new_sharding_spec
      )
  
      if rank == 0:
          assert dtensor_from_local.equal(
              original_tensor.narrow(0, 0, 1)
          )
      elif rank == 1:
          assert dtensor_from_local.equal(
              original_tensor.narrow(0, 1, 1)
          )
      elif rank == 2:
          assert dtensor_from_local.equal(
              original_tensor.narrow(0, 2, 1)
          )
      elif rank == 3:
          assert dtensor_from_local.equal(
              original_tensor.narrow(0, 3, 1)
          )
      else:
          raise ValueError(f"rank {rank} is not in the device mesh")
  
file_path: colossalai/tests/test_tensor/test_dtensor/test_dtensor.py
imports: import torch
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.initialize import launch
  from colossalai.logging import disable_existing_loggers
  from colossalai.tensor.d_tensor import ShardingSpec, distribute_tensor, get_global_shape, redistribute, to_global
  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: check_dtensor(rank, world_size, port) without return types