parameters: (
      inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, axis=0,
  )
docstring: r"""Gather tensors across the specified group.
      Only root process will receive the final result.
  
      Args:
          inp: Input tensor.
          group: The process group to work on.
              The default group is WORLD which means all processes available.
              You can use a list of process ranks to create new group to work on it, e.g. [1, 3, 5].
          device: The specific device to execute this operator.
              None default device means the device of inp will be used.
              Specify "gpu0:1" to execute this operator on diffrent cuda stream,
              1 is stream id, and default stream id is 0.
          axis: The concat axis for collective_comm result
  
      Examples:
  
          .. code-block::
  
             input = Tensor([rank])
             # Rank 0 # input: Tensor([0])
             # Rank 1 # input: Tensor([1])
             output = gather(input)
             # Rank 0 # output: Tensor([0 1])
             # Rank 1 # output: None
  
             input = Tensor([rank])
             group = Group([1, 0]) # first rank is root
             output = gather(input, group)
             # Rank 0 # output: None
             # Rank 1 # output: Tensor([1 0])
docstring_summary: r"""Gather tensors across the specified group.
      Only root process will receive the final result.
docstring_tokens: r Gather tensors across the specified group . Only root process will receive the final result .
function: def gather(
      inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, axis=0,
  ) -> Tensor:
      r"""Gather tensors across the specified group.
      Only root process will receive the final result.
  
      Args:
          inp: Input tensor.
          group: The process group to work on.
              The default group is WORLD which means all processes available.
              You can use a list of process ranks to create new group to work on it, e.g. [1, 3, 5].
          device: The specific device to execute this operator.
              None default device means the device of inp will be used.
              Specify "gpu0:1" to execute this operator on diffrent cuda stream,
              1 is stream id, and default stream id is 0.
          axis: The concat axis for collective_comm result
  
      Examples:
  
          .. code-block::
  
             input = Tensor([rank])
             # Rank 0 # input: Tensor([0])
             # Rank 1 # input: Tensor([1])
             output = gather(input)
             # Rank 0 # output: Tensor([0 1])
             # Rank 1 # output: None
  
             input = Tensor([rank])
             group = Group([1, 0]) # first rank is root
             output = gather(input, group)
             # Rank 0 # output: None
             # Rank 1 # output: Tensor([1 0])
      """
      assert (
          axis < inp.ndim
      ), "your concat_axis exceeds the dim of the tensor, the tensor shape is {}".format(
          inp.shape
      )
  
      out = _Gather(group, device)(inp)
  
      if group.rank == 0:
          if axis == 0:
              return out
          else:
              group_size = group.size
              transformed_shape = list(inp._tuple_shape)
              transformed_shape[axis] *= group_size
              n, *shp = out._tuple_shape
              index = (
                  [_ for _ in range(1, axis)]
                  + [axis, 0]
                  + [_ for _ in range(axis + 1, out.ndim + 1)]
              )
              return (
                  out.reshape(group_size, n // group_size, *shp)
                  .transpose(index)
                  .reshape(transformed_shape)
              )
      else:
          _save_output_for_autodiff(inp, out)
function_tokens: def gather ( inp : Tensor , group : Optional [ Group ] = WORLD , device : Optional [ str ] = None , axis = 0 , ) -> Tensor : assert ( axis < inp . ndim ) , "your concat_axis exceeds the dim of the tensor, the tensor shape is {}" . format ( inp . shape ) out = _Gather ( group , device ) ( inp ) if group . rank == 0 : if axis == 0 : return out else : group_size = group . size transformed_shape = list ( inp . _tuple_shape ) transformed_shape [ axis ] *= group_size n , * shp = out . _tuple_shape index = ( [ _ for _ in range ( 1 , axis ) ] + [ axis , 0 ] + [ _ for _ in range ( axis + 1 , out . ndim + 1 ) ] ) return ( out . reshape ( group_size , n // group_size , * shp ) . transpose ( index ) . reshape ( transformed_shape ) ) else : _save_output_for_autodiff ( inp , out )