# 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
        # in a AUTO_REUSE scope.
        # Note: Repeatedly 'with tf.device():' slows down the graph
        # construction. Therefore we directly use the cached device_stack here.
        tf.get_default_graph()._device_function_stack = (
            mesh_impl.graph_device_function_stacks[physical_pnum])

        if tf.get_variable_scope().reuse == tf.AUTO_REUSE:
          slice_var = tf.get_variable(
              initializer=zero_tensor,
              trainable=self._variable.trainable,
              collections=["TPU_VAR"],
              dtype=variable.slice_dtype,
              name=slice_var_name)
        else:
          slice_var = tf.Variable(
              initial_value=zero_tensor,
              trainable=self._variable.trainable,
              collections=["TPU_VAR"],
              dtype=variable.slice_dtype,
              name=slice_var_name,
              expected_shape=slice_shape)

        slices.append(slice_var)

      # Restore the initial stack
      tf.get_default_graph()._device_function_stack = init_device_stack
      # pylint: enable=protected-access

      self._laid_out_tensor = mesh_impl.LaidOutTensor(
          [tpu_variables.ReplicatedVariable(base_name, slices)])

# 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

        # in a AUTO_REUSE scope.
        # Note: Repeatedly 'with tf.device():' slows down the graph
        # construction. Therefore we directly use the cached device_stack here.
        tf.get_default_graph()._device_function_stack = (
            mesh_impl.graph_device_function_stacks[physical_pnum])

        if tf.get_variable_scope().reuse == tf.AUTO_REUSE:
          slice_var = tf.get_variable(
              initializer=zero_tensor,
              trainable=self._variable.trainable,
              collections=["TPU_VAR"],
              dtype=variable.slice_dtype,
              name=slice_var_name)
        else:
          slice_var = tf.Variable(
              initial_value=zero_tensor,
              trainable=self._variable.trainable,
              collections=["TPU_VAR"],
              dtype=variable.slice_dtype,
              name=slice_var_name,
              expected_shape=slice_shape)

        slices.append(slice_var)

      # Restore the initial stack
      tf.get_default_graph()._device_function_stack = init_device_stack
      # pylint: enable=protected-access

      self._laid_out_tensor = mesh_impl.LaidOutTensor(
          [tpu_variables.ReplicatedVariable(base_name, slices)])

      with tf.device(variable.master_device), utils.outside_all_rewrites():