class_docstr: Mesh implementation for TPU using SIMD and MPI operations.
class_docstr_tok: ["Mesh","implementation","for","TPU","using","SIMD","and","MPI","operations","."]
docstr: 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.
code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import os
  import gin
  from mesh_tensorflow import ops_with_redefined_builtins
  from mesh_tensorflow import tpu_variables
  from mesh_tensorflow import utils
  from six.moves import xrange
  import tensorflow.compat.v1
  from tensorflow.python.tpu.ops import tpu_ops
label_desc: Use image node_id 7 for calling the SimdMeshImpl obj's underlying member method code with example usage: obj.allreduce(x, mesh_axes, reduction_fn_string) and returns: x, self, x