code: def _reduce_scatter(input_, dim=1, process_group=None):
      """Do reduce-scatter operation.
  
      Args:
          input_ (`torch.Tensor`): The input tensor from sequence parallel region.
          dim (int): The dimension to perform reduce-scatter.
          process_group (`torch.distributed.ProcessGroup`): The process group used for collective communication.
      """
      world_size = dist.get_world_size(process_group)
      if world_size == 1:
          return input_
  
      # reduce-scatter
      new_shape = list(input_.shape)
      assert (
          new_shape[dim] % dist.get_world_size(process_group) == 0
      ), f"The dimension to split ({new_shape[dim]}) is not a multiple of tensor parallel size ({dist.get_world_size(process_group)}). "
      new_shape[dim] = new_shape[dim] // world_size
      output = torch.empty(
          new_shape, dtype=input_.dtype, device=input_.device
      )
      dist.reduce_scatter(output, input_, group=process_group)
  
      return output
  
imports: import torch
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 5 for calling a global function with example usage: _reduce_scatter(input_, dim, process_group) and returns: output, input_