code: def test_reduce_scatter_parity(self) -> None:
      if self.DISTRIBUTED_BACKEND == "ucc":
          self.skipTest(
              "torch_ucc does NOT support `torch.distributed._reduce_scatter_base` as of 2022/06/15"
          )
      from torch.distributed.distributed_c10d import (
          reduce_scatter,
          _reduce_scatter_base,
      )  # NOQA
  
      for tensor_model_parallel_world_size in range(
          2, 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,
          )
          tensor_model_parallel_rank = (
              parallel_state.get_tensor_model_parallel_rank()
          )
          cur_tensor_model_device = torch.device(
              f"cuda:{tensor_model_parallel_rank}"
          )
          with torch.no_grad():
              input = torch.cat(
                  [
                      i
                      * torch.ones(
                          self.tensor_shape,
                          dtype=torch.float32,
                          device=cur_tensor_model_device,
                      )
                      for i in range(tensor_model_parallel_world_size)
                  ]
              )
              input_list = [
                  t.clone()
                  for t in input.chunk(tensor_model_parallel_world_size)
              ]
          output = torch.empty(
              self.tensor_shape,
              device=cur_tensor_model_device,
              dtype=torch.float32,
              requires_grad=False,
          )
          reduce_scatter(
              output,
              input_list,
              group=parallel_state.get_tensor_model_parallel_group(),
          )
  
          output_for_base = torch.empty(
              self.tensor_shape,
              device=cur_tensor_model_device,
              dtype=torch.float32,
              requires_grad=False,
          )
          _reduce_scatter_base(
              output_for_base,
              input,
              group=parallel_state.get_tensor_model_parallel_group(),
          )
  
          msg = f"tensor_model_parallel_world_size: {tensor_model_parallel_world_size}"
          self.assertEqual(output, output_for_base, msg=msg)
          self.assertEqual(input, torch.cat(input_list), msg=msg)
          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 3 for calling the TensorParallelLayerTestBase obj's underlying member method code with example usage: obj.test_reduce_scatter_parity() without return types