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 allconcat (like MPI allgather followed by concat).
  
  TODO(noam): inefficient - replace with a XLA allconcat when available
  
  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)
    stack: a boolean - whether to stack instead of concat
  Returns:
    a LaidOutTensor
code: def allconcat(self, x, mesh_axis, concat_axis, stack=False):
      """Grouped allconcat (like MPI allgather followed by concat).
  
      TODO(noam): inefficient - replace with a XLA allconcat when available
  
      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)
        stack: a boolean - whether to stack instead of concat
      Returns:
        a LaidOutTensor
      """
      x = x.to_laid_out_tensor()
      coord = self.laid_out_pcoord(mesh_axis)
      t = x.one_slice
      old_shape = t.shape.as_list()
      num_parts = self.shape[mesh_axis].size
      t = tf.expand_dims(t, concat_axis)
      t *= tf.reshape(
          tf.one_hot(coord.one_slice, num_parts, dtype=t.dtype),
          [
              num_parts if i == concat_axis else 1
              for i in xrange(len(old_shape) + 1)
          ],
      )
      if not stack:
          new_shape = old_shape[:]
          new_shape[concat_axis] *= num_parts
          t = tf.reshape(t, new_shape)
      return self.allreduce(self.LaidOutTensor([t]), [mesh_axis], "SUM")
  
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 8 for calling the SimdMeshImpl obj's underlying member method code with example usage: obj.allconcat(x, mesh_axis, concat_axis, stack) and returns: self