Grouped allreduce, (summed across the given dimensions).

    Args:
      x: a LaidOutTensor
      mesh_axes: a list of integers
      reduction_fn_string: "SUM"
    Returns:
      a LaidOutTensor
    Raises:
      ValueError: if the reduction is not yet implemented.

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

    Args:
      x: a LaidOutTensor
      mesh_axes: a list of integers
      reduction_fn_string: "SUM"
    Returns:
      a LaidOutTensor
    Raises:
      ValueError: if the reduction is not yet implemented.
    """
    if not mesh_axes:
      return x
    x = x.to_laid_out_tensor()
    if reduction_fn_string == "SUM":
      group_assignment = self._create_group_assignment(mesh_axes)
      group_size = len(group_assignment[0])
      tf_in = x.one_slice
      dtype = tf_in.dtype
      if dtype == tf.float32:
        cast_to_float32 = False
      elif dtype == tf.bfloat16:
        cast_to_float32 = (
            group_size > self._allreduce_in_bfloat16_max_group_size)
      else:
        tf.logging.info("Casting %s to float32 for allreduce" % tf_in.dtype)
        cast_to_float32 = True
      if cast_to_float32:
        tf_in = tf.cast(tf_in, tf.float32)
      tf_out = tpu_ops.cross_replica_sum(tf_in, group_assignment)
      if cast_to_float32:
        tf_out = tf.cast(tf_out, dtype)
      return self.LaidOutTensor([tf_out])
    else:
      for axis in mesh_axes:
        x = self.allconcat(x, axis, 0, stack=True)
        x = self.LaidOutTensor(
            [mtf.reduction_fn(reduction_fn_string)(x.one_slice, 0)])
      return x