class_docstr: Creates mtf graph, mesh, and mesh implementation.
class_docstr_tok: ["Creates","mtf","graph",",","mesh",",","and","mesh","implementation","."]
code: def __init__(self, sess, use_tpu, mesh_shape, layout_rules):
      super(MeshContext, self).__init__()
      self._use_tpu = use_tpu
      self._mesh_shape = mtf.convert_to_shape(mesh_shape)
      self._layout_rules = layout_rules
  
      self._d_assignment = None
      self._num_hosts = None
      self._num_cores = None
  
      self._cpu_devices, self._gpu_devices = self._list_cpu_gpu_devices(
          sess
      )
  
      if self._use_tpu:
          topology = sess.run(tpu.initialize_system())
          topo_object = tpu.Topology(serialized=topology)
          self._num_cores = int(np.prod(topo_object.mesh_shape))
          self._num_hosts = int(topo_object.num_tasks)
          num_cores_per_host = int(self._num_cores // self._num_hosts)
          assert num_cores_per_host == int(
              topo_object.num_tpus_per_task
          )
  
          # Get a device_assignment object for mtf.
          self._d_assignment = device_assignment.device_assignment(
              topology,
              computation_shape=[
                  1,
              ]
              * mtf.utils.topology_rank(topology),
              num_replicas=self._num_cores,
          )
  
          self._mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
              self._mesh_shape,
              self._layout_rules,
              None,
              self._d_assignment,
          )
      else:
          self._mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
              self._mesh_shape, self._layout_rules, self._gpu_devices
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import re
  import mesh_tensorflow
  import numpy
  import six
  from six.moves import range
  from six.moves import zip
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from mesh_tensorflow.experimental import input_reader
  from mesh_tensorflow.experimental import unet
  from tensorflow.contrib import summary
  from tensorflow.contrib import tpu
  from tensorflow.contrib.tpu.python.tpu import device_assignment
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.platform import flags
  from tensorflow.python.tpu.ops import tpu_ops
label_desc: Use image node_id 1 to create a new MeshContext object from inherited base classes: object with example: obj = MeshContext(sess, use_tpu, mesh_shape, layout_rules)