class_docstr: Mesh implementation for TPU using SIMD and MPI operations.
class_docstr_tok: ["Mesh","implementation","for","TPU","using","SIMD","and","MPI","operations","."]
docstr: Create a SimdMeshImpl.
  
  Args:
    shape: an input to mtf.convert_to_shape()
    layout: an input to mtf.convert_to_layout_rules()
    devices: deprecated
    device_assignment: a tf.tpu.experimental.DeviceAssignment -
      devices must be asssigned in lexicographic order
    logical_to_physical: an optional permutation representing the mapping
      from logical cores to "physical" cores, where the physical cores are
      listed in lexicographic order in the physical mesh, and the logical
      cores are listed in lexicographic order in the logical mesh.
      Default is lexicographic order.
    allreduce_in_bfloat16_max_group_size: an integer.  Allreduces of bfloat16
      tensors are done in float32 if the group size exceeds this value.
code: def __init__(
      self,
      shape,
      layout,
      devices=None,
      device_assignment=None,
      logical_to_physical=None,
      allreduce_in_bfloat16_max_group_size=8,
  ):
      """Create a SimdMeshImpl.
  
      Args:
        shape: an input to mtf.convert_to_shape()
        layout: an input to mtf.convert_to_layout_rules()
        devices: deprecated
        device_assignment: a tf.tpu.experimental.DeviceAssignment -
          devices must be asssigned in lexicographic order
        logical_to_physical: an optional permutation representing the mapping
          from logical cores to "physical" cores, where the physical cores are
          listed in lexicographic order in the physical mesh, and the logical
          cores are listed in lexicographic order in the logical mesh.
          Default is lexicographic order.
        allreduce_in_bfloat16_max_group_size: an integer.  Allreduces of bfloat16
          tensors are done in float32 if the group size exceeds this value.
      """
      super(SimdMeshImpl, self).__init__(shape, layout)
      if devices is not None:
          tf.logging.warning(
              "SimdMeshImpl ignoring devices %s" % devices
          )
      self._device_assignment = device_assignment
      tf.logging.info(
          "SimdMeshImpl init: {0} {1}".format(shape, layout)
      )
      tf.logging.info(
          "Device Assignment: {0}".format(device_assignment)
      )
      if logical_to_physical is None:
          # TODO(noam): maybe use auto_logical_to_physical_tpu() here
          logical_to_physical = list(range(self.size))
      if sorted(logical_to_physical) != list(range(self.size)):
          raise ValueError(
              "logical_to_physical must be a permutation on range(shape.size)"
              " shape=%s logical_to_physical=%s"
              % (shape, logical_to_physical)
          )
      self._logical_to_physical = logical_to_physical
      self._physical_to_logical = [None] * self.size
      for logical, physical in enumerate(self._logical_to_physical):
          self._physical_to_logical[physical] = logical
      self._pnum_tensor = None
      self.graph_device_function_stacks = []
      self.copy_master_to_slice_ops = []
      self._allreduce_in_bfloat16_max_group_size = (
          allreduce_in_bfloat16_max_group_size
      )
  
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 1 to create a new SimdMeshImpl object from inherited base classes: mtf with example: obj = SimdMeshImpl(shape, layout, devices, device_assignment, logical_to_physical, allreduce_in_bfloat16_max_group_size)