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 alltoall (like MPI alltoall with splitting and concatenation).
  
  Args:
    x: a LaidOutTensor
    mesh_axis: an integer the mesh axis along which to group
    split_axis: an integer (the Tensor axis along which to split)
    concat_axis: an integer (the Tensor axis along which to concatenate)
  Returns:
    a LaidOutTensor
code: def alltoall(self, x, mesh_axis, split_axis, concat_axis):
      """Grouped alltoall (like MPI alltoall with splitting and concatenation).
  
      Args:
        x: a LaidOutTensor
        mesh_axis: an integer the mesh axis along which to group
        split_axis: an integer (the Tensor axis along which to split)
        concat_axis: an integer (the Tensor axis along which to concatenate)
      Returns:
        a LaidOutTensor
      """
      x = x.to_laid_out_tensor()
      t = x.one_slice
      group_assignment = self._create_group_assignment([mesh_axis])
      dtype = t.dtype
      if dtype == tf.float32:
          # There seems to be a bug with float32 alltoall.
          # Do it in bfloat16 until the bug is fixed.
          # TODO(noam): file a bug
          t = tf.to_bfloat16(t)
      t = tpu_ops.all_to_all(
          t,
          concat_dimension=concat_axis,
          split_dimension=split_axis,
          split_count=len(group_assignment[0]),
          group_assignment=group_assignment,
      )
      t = tf.cast(t, dtype)
      x = self.LaidOutTensor([t])
      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 9 for calling the SimdMeshImpl obj's underlying member method code with example usage: obj.alltoall(x, mesh_axis, split_axis, concat_axis) and returns: x