# Path: mesh_tensorflow/ops_with_redefined_builtins.py

# 
# Path: mesh_tensorflow/tpu_variables.py
#   TF_23 = True
#   TF_23 = False
# def _handle_graph(handle):
# def _enclosing_tpu_context():
#   def __init__(self, name, variables):
#   def handle(self):
#   def _assign_dependencies(self):
#   def initializer(self):
#   def graph(self):
#   def _shared_name(self):
#   def _unique_id(self):
#   def name(self):
#   def dtype(self):
#   def shape(self):
#   def get_shape(self):
#   def to_proto(self, export_scope=None):
#   def constraint(self):
#   def op(self):
#   def _read_variable_op(self):
#   def read_value(self):
#   def assign(self, value, use_locking=None, name=None, read_value=False):
#   def assign_add(self, delta, use_locking=None, name=None, read_value=True):
#   def assign_sub(self, delta, use_locking=None, name=None, read_value=True):
#   def get(self):
#   def _in_graph_mode(self):
#   def _should_act_as_resource_variable(self):
#   def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
# def _tensor_conversion(var, dtype=None, name=None, as_ref=False):
# class ReplicatedVariable(VariableBase):
# 
# Path: mesh_tensorflow/utils.py
# def outside_all_rewrites():
#   def __init__(self, devices, init_usage=None):
#   def device_function(self, var):
# def create_host_call(model_dir):
#   def maybe_cast(tensor):
#   def host_call_fn(global_step, *args):
# def topology_rank(topology):
# def remove_summaries():
# class BalancedVariablePlacer(object):
# SCALAR_SUMMARIES_COLLECTION_KEY = 'mtf_scalar_summaries'


import math
import os
import gin
import tensorflow.compat.v1 as tf
from mesh_tensorflow import ops_with_redefined_builtins as mtf
from mesh_tensorflow import tpu_variables
from mesh_tensorflow import utils
from six.moves import xrange  # pylint: disable=redefined-builtin
from tensorflow.python.tpu.ops import tpu_ops  # pylint: disable=g-direct-tensorflow-import
         for name, s in zip(self._names, projected_shapes)])

  @property
  def logical_to_physical(self):
    """List of physical processor numbers."""
    return list(self._logical_to_physical)

  @property
  def mesh_shape(self):
    return self._mesh_shape

  @classmethod
  def spec_to_mesh_shape(cls, spec, num_processors):
    """Compute mesh shape even without knowing the physical shape.

    This is useful in cases where the mesh shape must be computed before
    you know the physical_shape.

    Args:
      spec: a list of (string, list-of-integers) pairs
      num_processors: an integer
    Returns:
      a mtf.Shape
    """
    logical_ndims = len(spec)
    names = [p[0] for p in spec]
    sizes = [p[1] for p in spec]
    sizes = [None if s is None else mtf.list_product(s) for s in sizes]
    if logical_ndims > 0 and sizes[0] is None:
      sizes[0] = num_processors // mtf.list_product(sizes[1:])

# Path: mesh_tensorflow/ops_with_redefined_builtins.py

# 
# Path: mesh_tensorflow/tpu_variables.py
#   TF_23 = True
#   TF_23 = False
# def _handle_graph(handle):
# def _enclosing_tpu_context():
#   def __init__(self, name, variables):
#   def handle(self):
#   def _assign_dependencies(self):
#   def initializer(self):
#   def graph(self):
#   def _shared_name(self):
#   def _unique_id(self):
#   def name(self):
#   def dtype(self):
#   def shape(self):
#   def get_shape(self):
#   def to_proto(self, export_scope=None):
#   def constraint(self):
#   def op(self):
#   def _read_variable_op(self):
#   def read_value(self):
#   def assign(self, value, use_locking=None, name=None, read_value=False):
#   def assign_add(self, delta, use_locking=None, name=None, read_value=True):
#   def assign_sub(self, delta, use_locking=None, name=None, read_value=True):
#   def get(self):
#   def _in_graph_mode(self):
#   def _should_act_as_resource_variable(self):
#   def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
# def _tensor_conversion(var, dtype=None, name=None, as_ref=False):
# class ReplicatedVariable(VariableBase):
# 
# Path: mesh_tensorflow/utils.py
# def outside_all_rewrites():
#   def __init__(self, devices, init_usage=None):
#   def device_function(self, var):
# def create_host_call(model_dir):
#   def maybe_cast(tensor):
#   def host_call_fn(global_step, *args):
# def topology_rank(topology):
# def remove_summaries():
# class BalancedVariablePlacer(object):
# SCALAR_SUMMARIES_COLLECTION_KEY = 'mtf_scalar_summaries'


# Path: mesh_tensorflow/simd_mesh_impl.py
import math
import os
import gin
import tensorflow.compat.v1 as tf
from mesh_tensorflow import ops_with_redefined_builtins as mtf
from mesh_tensorflow import tpu_variables
from mesh_tensorflow import utils
from six.moves import xrange  # pylint: disable=redefined-builtin
from tensorflow.python.tpu.ops import tpu_ops  # pylint: disable=g-direct-tensorflow-import

         for name, s in zip(self._names, projected_shapes)])

  @property
  def logical_to_physical(self):
    """List of physical processor numbers."""
    return list(self._logical_to_physical)

  @property
  def mesh_shape(self):
    return self._mesh_shape

  @classmethod
  def spec_to_mesh_shape(cls, spec, num_processors):
    """Compute mesh shape even without knowing the physical shape.

    This is useful in cases where the mesh shape must be computed before
    you know the physical_shape.

    Args:
      spec: a list of (string, list-of-integers) pairs
      num_processors: an integer
    Returns:
      a mtf.Shape
    """
    logical_ndims = len(spec)
    names = [p[0] for p in spec]
    sizes = [p[1] for p in spec]
    sizes = [None if s is None else mtf.list_product(s) for s in sizes]
    if logical_ndims > 0 and sizes[0] is None:
      sizes[0] = num_processors // mtf.list_product(sizes[1:])

    if mtf.list_product(sizes) != num_processors: