Grouped allreduce, (across the given dimensions).

    Args:
      x: a LaidOutTensor
      mesh_axes: a list of integers - the mesh dimensions to be reduced
      reduction_fn_string: "SUM" or "MAX"
    Returns:
      a LaidOutTensor

def allreduce(self, x, mesh_axes, reduction_fn_string):
    """Grouped allreduce, (across the given dimensions).

    Args:
      x: a LaidOutTensor
      mesh_axes: a list of integers - the mesh dimensions to be reduced
      reduction_fn_string: "SUM" or "MAX"
    Returns:
      a LaidOutTensor
    """
    return self._collective_with_groups(
        x, mesh_axes, functools.partial(
            allreduce_ring, reduction_fn_string=reduction_fn_string))