code: #!/usr/bin/env python
  
  import torch
  import torch.nn as nn
  from colossalai.nn import CheckpointModule
  from .utils.dummy_data_generator import DummyDataGenerator
  from .registry import non_distributed_component_funcs
  
  
  class NetWithRepeatedlyComputedLayers(CheckpointModule):
      """
      This model is to test with layers which go through forward pass multiple times.
      In this model, the fc1 and fc2 call forward twice
      """
  
      def __init__(self, checkpoint=False) -> None:
          super().__init__(checkpoint=checkpoint)
          self.fc1 = nn.Linear(5, 5)
          self.fc2 = nn.Linear(5, 5)
          self.fc3 = nn.Linear(5, 2)
          self.layers = [self.fc1, self.fc2, self.fc1, self.fc2, self.fc3]
  
      def forward(self, x):
          for layer in self.layers:
              x = layer(x)
          return x
  
  
  class DummyDataLoader(DummyDataGenerator):
  
      def generate(self):
          data = torch.rand(16, 5)
          label = torch.randint(low=0, high=2, size=(16,))
          return data, label
  
  
  @non_distributed_component_funcs.register(name='repeated_computed_layers')
  def get_training_components():
  
      def model_builder(checkpoint=True):
          return NetWithRepeatedlyComputedLayers(checkpoint)
  
      trainloader = DummyDataLoader()
      testloader = DummyDataLoader()
  
      criterion = torch.nn.CrossEntropyLoss()
      return model_builder, trainloader, testloader, torch.optim.Adam, criterion
   import torch
  from colossalai.tensor import ColoTensor
  from numpy import allclose
  
  
  def test_tensor_indexing():
      torch_t = torch.randn(2, 3)
      colo_t = ColoTensor.init_from_torch_tensor(torch_t)
      assert allclose(torch_t[:, 1], colo_t[:, 1].torch_tensor())
  
  
  def test_lazy_init_tensor():
      lazy_t = ColoTensor(2, 3, dtype=torch.float32, requires_grad=True)
      assert lazy_t._torch_tensor.numel() == 0
      assert lazy_t.numel() == 6 == lazy_t.torch_tensor().numel()
  
  
  def test_wrapped_tensor_func():
      t_ref = torch.randn(4, 5)
      t = ColoTensor.init_from_torch_tensor(t_ref.clone())
  
      # non-func attr
      assert t.is_cuda == t_ref.is_cuda
  
      # TODO I don't find out a tensor function which returns None.
  
      # return 1 torch.Tensor
      t_abs = t.abs()
      assert isinstance(t_abs, ColoTensor) and torch.equal(t_abs.torch_tensor(), t_ref.abs())
  
      # return 1 non-torch.Tensor
      assert t.dim() == t_ref.dim()
  
      # return >1 torch.Tensor
      t_split1, t_split2 = t.split(2)
      assert isinstance(t_split1, ColoTensor) and isinstance(t_split2, ColoTensor)
  
  
  def test_operand():
      t_ref = torch.randn(4, 5)
      t = ColoTensor.init_from_torch_tensor(t_ref.clone())
  
      t_ref_res = t_ref + t_ref
      t_res = t + t
      assert torch.allclose(t_ref_res, t_res)
  
apis: torch.rand torch.nn.Linear torch.randint torch.nn.CrossEntropyLoss
apis: torch.randn torch.allclose