code: def _reduce_scatter_along_first_dim(
      input_: torch.Tensor,
  ) -> torch.Tensor:
      """Reduce-scatter the input tensor across model parallel group."""
      world_size = get_tensor_model_parallel_world_size()
      # Bypass the function if we are using only 1 GPU.
      if world_size == 1:
          return input_
  
      shape = list(input_.shape)
      assert shape[0] % world_size == 0
      shape[0] //= world_size
      output = torch.empty(
          shape, dtype=input_.dtype, device=torch.cuda.current_device()
      )
      torch.distributed.reduce_scatter_tensor(
          output,
          input_.contiguous(),
          group=get_tensor_model_parallel_group(),
      )
      return output
  
imports: import torch
  from apex.transformer.parallel_state import get_tensor_model_parallel_group
  from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
  from apex.transformer.parallel_state import get_tensor_model_parallel_rank
  from apex.transformer.tensor_parallel.utils import split_tensor_along_last_dim
label_desc: Use image node_id 6 for calling a global function with example usage: _reduce_scatter_along_first_dim(input_) and returns: output, input_