Grouped allconcat (like MPI allgather followed by concat).

    Args:
      x: a LaidOutTensor
      mesh_axis: an integer - the mesh axis along which to group
      concat_axis: an integer (the Tensor axis along which to concatenate)
    Returns:
      a LaidOutTensor

def allconcat(self, x, mesh_axis, concat_axis):
    """Grouped allconcat (like MPI allgather followed by concat).

    Args:
      x: a LaidOutTensor
      mesh_axis: an integer - the mesh axis along which to group
      concat_axis: an integer (the Tensor axis along which to concatenate)
    Returns:
      a LaidOutTensor
    """
    return self._collective_with_groups(
        x, [mesh_axis],
        functools.partial(allconcat_ring, concat_axis=concat_axis))