code: def alltoall(self, x, mesh_axis, split_axis, concat_axis):
      
      return self._collective_with_groups(
          x, [mesh_axis],
          functools.partial(
              alltoall_ring, split_axis=split_axis, concat_axis=concat_axis))
docs: Grouped alltoall.
  
      Args:
        x: a LaidOutTensor
        mesh_axis: an integer the mesh axis along which to group
        split_axis: an integer (the Tensor axis along which to split)
        concat_axis: an integer (the Tensor axis along which to concatenate)
      Returns:
        a LaidOutTensor