code: def _affine_weight_init_test_impl(
      self, init_device: str, is_column_parallel: bool
  ) -> None:
      dim = int(not is_column_parallel)
      for tensor_model_parallel_world_size in range(
          1, self.world_size + 1
      ):
          if self.world_size % tensor_model_parallel_world_size:
              continue
          parallel_state.initialize_model_parallel(
              tensor_model_parallel_size_=tensor_model_parallel_world_size
          )
          input_size: int = (
              self.INPUT_SIZE_COEFF * tensor_model_parallel_world_size
          )
          output_size: int = (
              self.OUTPUT_SIZE_COEFF * tensor_model_parallel_world_size
          )
  
          weight_shape = (
              (self.OUTPUT_SIZE_COEFF, input_size)
              if is_column_parallel
              else (output_size, self.INPUT_SIZE_COEFF)
          )
          weight = torch.empty(weight_shape)
          set_random_seed(self.SEED)
  
          sharding_dim_size = (
              self.OUTPUT_SIZE_COEFF
              if is_column_parallel
              else self.INPUT_SIZE_COEFF
          )
  
          if init_device == "cpu":
              layers._initialize_affine_weight_cpu(
                  weight,
                  output_size,
                  input_size,
                  sharding_dim_size,
                  dim,
                  nn.init.normal_,
                  params_dtype=torch.float32,
              )
          else:
              layers._initialize_affine_weight_gpu(
                  weight, torch.nn.init.normal_, dim
              )
          # Target
          set_random_seed(self.SEED)
          if init_device == "cpu":
              main_weight = torch.empty(output_size, input_size)
              nn.init.normal_(main_weight)
              curr_weight = torch.split(
                  main_weight, sharding_dim_size, dim=dim
              )[parallel_state.get_tensor_model_parallel_rank()]
          else:
              curr_weight = torch.empty(*weight_shape)
              nn.init.normal_(curr_weight)
  
          self.assertEqual(
              curr_weight,
              weight,
              msg=f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}",
          )
          parallel_state.destroy_model_parallel()
  
imports: import logging
  import unittest
  import typing
  import torch
  import torch.nn
  from torch.testing._internal import common_utils
  from apex.transformer import parallel_state
  from apex.transformer.tensor_parallel import layers
  from apex.transformer.testing.commons import set_random_seed
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
  from apex.transformer.testing.distributed_test_base import UccDistributedTestBase
label_desc: Use image node_id 5 for calling the TensorParallelLayerTestBase obj's underlying member method code with example usage: obj._affine_weight_init_test_impl(init_device, is_column_parallel) without return types