code: def test_initialize_affine_weight(tensor_model_parallel_size):
      mpu.initialize_model_parallel(tensor_model_parallel_size)
      if torch.distributed.get_rank() == 0:
          print(
              "> testing initialize_affine_weight with model parallel "
              "size: {}".format(tensor_model_parallel_size)
          )
      tensor_model_parallel_size = (
          mpu.get_tensor_model_parallel_world_size()
      )
  
      seed = 12345
      input_size_coeff = 13
      input_size = input_size_coeff * tensor_model_parallel_size
      output_size_coeff = 17
      output_size = output_size_coeff * tensor_model_parallel_size
  
      # ---------------
      # Column parallel
      # ---------------
      weight = torch.empty(output_size_coeff, input_size)
      set_random_seed(seed)
      layers._initialize_affine_weight(
          weight,
          output_size,
          input_size,
          output_size_coeff,
          0,
          torch.nn.init.normal_,
      )
      # Target.
      set_random_seed(seed)
      master_weight = torch.empty(output_size, input_size)
      torch.nn.init.normal_(master_weight)
      rank = mpu.get_tensor_model_parallel_rank()
      my_weight = (
          torch.split(master_weight, output_size_coeff, dim=0)[rank]
          .contiguous()
          .clone()
      )
  
      # Compare.
      error = weight.sub(my_weight).abs().max()
      torch.distributed.barrier()
      print(
          "   column parallel max error (should be zero) on global rank "
          "{}: {}".format(torch.distributed.get_rank(), error)
      )
      assert error < 1.0e-6
  
      # ------------
      # Row parallel
      # ------------
      weight = torch.empty(output_size, input_size_coeff)
      set_random_seed(seed)
      mpu.layers._initialize_affine_weight(
          weight,
          output_size,
          input_size,
          input_size_coeff,
          1,
          torch.nn.init.normal_,
      )
      # Target.
      set_random_seed(seed)
      master_weight = torch.empty(output_size, input_size)
      torch.nn.init.normal_(master_weight)
      rank = mpu.get_tensor_model_parallel_rank()
      my_weight = (
          torch.split(master_weight, input_size_coeff, dim=1)[rank]
          .contiguous()
          .clone()
      )
  
      # Compare.
      error = weight.sub(my_weight).abs().max()
      torch.distributed.barrier()
      print(
          "   row parallel max error (should be zero) on global rank "
          "{}: {}".format(torch.distributed.get_rank(), error)
      )
      assert error < 1.0e-6
  
      # Reset groups
      mpu.destroy_model_parallel()
  
      torch.distributed.barrier()
      if torch.distributed.get_rank() == 0:
          print(" >> passed the test :-)")
  
imports: from mpu import layers
  from commons import set_random_seed
  from commons import print_separator
  from commons import initialize_distributed
  import mpu
  from torch.nn.parameter import Parameter
  import torch.nn.init
  import torch
  import random
  import sys
label_desc: Use image node_id 2 for calling a global function with example usage: test_initialize_affine_weight(tensor_model_parallel_size) without return types