code: def check_linear_col_plus_row(
      lazy_init: bool, seq_parallel: bool, overlap: bool
  ):
      ctx = LazyInitContext() if lazy_init else nullcontext()
  
      linear_1 = nn.Linear(32, 128).cuda()
      linear_2 = nn.Linear(128, 32).cuda()
  
      with ctx:
          linear_1_copy = nn.Linear(32, 128).cuda()
          linear_2_copy = nn.Linear(128, 32).cuda()
      linear_col = Linear1D_Col.from_native_module(
          linear_1_copy,
          process_group=None,
          gather_output=False,
          seq_parallel=seq_parallel,
          overlap=overlap,
      )
      linear_row = Linear1D_Row.from_native_module(
          linear_2_copy,
          process_group=None,
          parallel_input=True,
          seq_parallel=seq_parallel,
      )
  
      linear_1.load_state_dict(linear_col.state_dict())
      linear_col.load_state_dict(linear_1.state_dict())
      linear_2.load_state_dict(linear_row.state_dict())
      linear_row.load_state_dict(linear_2.state_dict())
  
      # check computation correctness
      # [batch_size, seq_len, hidden_size]
      x = torch.rand(2, 4, 32).cuda()
      x_for_unshard = x.expand_as(x.clone())
      x_for_unshard.requires_grad_(True)
      x_for_shard = (
          x.expand_as(x.clone())
          if seq_parallel is False
          else torch.chunk(x.clone(), 2, dim=1)[dist.get_rank()]
      )
      x_for_shard.requires_grad_(True)
  
      # run forward
      unshard_out = linear_2(linear_1(x_for_unshard))
      shard_out = linear_row(linear_col(x_for_shard))
      target_out = (
          unshard_out
          if seq_parallel is False
          else torch.chunk(unshard_out.clone(), 2, dim=1)[
              dist.get_rank()
          ]
      )
      assert_close(target_out, shard_out)
  
      # check backward correctness
      unshard_out.sum().backward()
      shard_out.sum().backward()
  
      rank = dist.get_rank()
      target_1_grad = torch.chunk(linear_1.weight.grad, 2, dim=0)[rank]
      assert_close(target_1_grad, linear_col.weight.grad)
  
      # check the input gradients
      assert x_for_shard.grad is not None
      assert x_for_unshard.grad is not None
      target_unshard_gard = (
          x_for_unshard.grad
          if seq_parallel is False
          else torch.chunk(x_for_unshard.grad.clone(), 2, dim=1)[
              dist.get_rank()
          ]
      )
      assert_close(target_unshard_gard, x_for_shard.grad)
  
file_path: ColossalAI/tests/test_shardformer/test_layer/test_linear_1d.py
imports: from contextlib import nullcontext
  import torch
  import torch.distributed
  import torch.nn
  from torch.testing import assert_close
  import colossalai
  from colossalai.lazy import LazyInitContext
  from colossalai.shardformer.layer import Linear1D_Col, Linear1D_Row
  from colossalai.tensor.d_tensor import is_distributed_tensor
  from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 3 for calling a global function with example usage: check_linear_col_plus_row(lazy_init, seq_parallel, overlap) without return types