code: def allreduce(self, x, mesh_axes, reduction_fn_string):
      
      return self._collective_with_groups(
          x, mesh_axes, functools.partial(
              allreduce_ring, reduction_fn_string=reduction_fn_string))
docs: 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