parameters: (
      inp: Tensor,
      group: Optional[Group] = WORLD,
      device: Optional[str] = None,
      split_axis: int = 0,
      concat_axis: int = 0,
  )
return_statement: return (
          out.reshape(group_size, n // group_size, *shp)
          .transpose(index)
          .reshape(transformed_shape)
      )
docstring: r"""Each process scatter input tensor to all processes and return gathered tensor.
  
      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.
          split_axis: The axis that collectivecomm will split data
              the default axis is 0
  
      Returns:
          Result tensor.
  
      Examples:
  
          .. code-block::
  
             input = Tensor([0 1]) + rank*2
             # Rank 0 # input: Tensor([0 1])
             # Rank 1 # input: Tensor([2 3])
             output = all_to_all(input)
             # Rank 0 # output: Tensor([0 2])
             # Rank 1 # output: Tensor([1 3])
  
             input = Tensor([0 1]) + rank*2
             group = Group([1, 0])
             output = all_to_all(input, group)
             # Rank 0 # output: Tensor([0 3])
             # Rank 1 # output: Tensor([2 1])
docstring_summary: r"""Each process scatter input tensor to all processes and return gathered tensor.
docstring_tokens: r Each process scatter input tensor to all processes and return gathered tensor .
function: def all_to_all(
      inp: Tensor,
      group: Optional[Group] = WORLD,
      device: Optional[str] = None,
      split_axis: int = 0,
      concat_axis: int = 0,
  ) -> Tensor:
      r"""Each process scatter input tensor to all processes and return gathered tensor.
  
      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.
          split_axis: The axis that collectivecomm will split data
              the default axis is 0
  
      Returns:
          Result tensor.
  
      Examples:
  
          .. code-block::
  
             input = Tensor([0 1]) + rank*2
             # Rank 0 # input: Tensor([0 1])
             # Rank 1 # input: Tensor([2 3])
             output = all_to_all(input)
             # Rank 0 # output: Tensor([0 2])
             # Rank 1 # output: Tensor([1 3])
  
             input = Tensor([0 1]) + rank*2
             group = Group([1, 0])
             output = all_to_all(input, group)
             # Rank 0 # output: Tensor([0 3])
             # Rank 1 # output: Tensor([2 1])
      """
      group_size = group.size if group is not None else 1
      assert (
          list(inp._tuple_shape)[split_axis] % group_size == 0
      ), "current axis: {} can't devided by group size".format(split_axis)
      origin_shape = inp._tuple_shape
      if split_axis != 0:
          k_new_shape = list(inp._tuple_shape)
          k_new_shape[split_axis] //= group_size
          k_new_shape[0] *= group_size
          new_shape = list(inp._tuple_shape)
          new_shape[split_axis] //= group_size
          new_shape.insert(split_axis, group_size)
          index = (
              [split_axis]
              + [_ for _ in range(0, split_axis)]
              + [_ for _ in range(split_axis + 1, inp.ndim + 1)]
          )
          inp = inp.reshape(new_shape).transpose(index).reshape(k_new_shape)
  
      mode = CollectiveComm.Mode.ALL_TO_ALL
      out = collective_comm(inp, mode, group, device)
  
      if concat_axis == 0:
          return out
  
      transformed_shape = list(origin_shape)
      transformed_shape[concat_axis] *= group_size
      transformed_shape[split_axis] //= group_size
  
      n, *shp = out._tuple_shape
      index = (
          [_ for _ in range(1, concat_axis)]
          + [concat_axis, 0]
          + [_ for _ in range(concat_axis + 1, out.ndim + 1)]
      )
      return (
          out.reshape(group_size, n // group_size, *shp)
          .transpose(index)
          .reshape(transformed_shape)
      )
function_tokens: def all_to_all ( inp : Tensor , group : Optional [ Group ] = WORLD , device : Optional [ str ] = None , split_axis : int = 0 , concat_axis : int = 0 , ) -> Tensor : group_size = group . size if group is not None else 1 assert ( list ( inp . _tuple_shape ) [ split_axis ] % group_size == 0 ) , "current axis: {} can't devided by group size" . format ( split_axis ) origin_shape = inp . _tuple_shape if split_axis != 0 : k_new_shape = list ( inp . _tuple_shape ) k_new_shape [ split_axis ] //= group_size k_new_shape [ 0 ] *= group_size new_shape = list ( inp . _tuple_shape ) new_shape [ split_axis ] //= group_size new_shape . insert ( split_axis , group_size ) index = ( [ split_axis ] + [ _ for _ in range ( 0 , split_axis ) ] + [ _ for _ in range ( split_axis + 1 , inp . ndim + 1 ) ] ) inp = inp . reshape ( new_shape ) . transpose ( index ) . reshape ( k_new_shape ) mode = CollectiveComm . Mode . ALL_TO_ALL out = collective_comm ( inp , mode , group , device ) if concat_axis == 0 : return out transformed_shape = list ( origin_shape ) transformed_shape [ concat_axis ] *= group_size transformed_shape [ split_axis ] //= group_size n , * shp = out . _tuple_shape index = ( [ _ for _ in range ( 1 , concat_axis ) ] + [ concat_axis , 0 ] + [ _ for _ in range ( concat_axis + 1 , out . ndim + 1 ) ] ) return ( out . reshape ( group_size , n // group_size , * shp ) . transpose ( index ) . reshape ( transformed_shape ) )