# Lint as: python3
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility for exporting an InferenceGraph proto from model params."""

import collections
import contextlib
import re
import lingvo.compat as tf
from lingvo.core import base_model
from lingvo.core import bfloat16_variables
from lingvo.core import inference_graph_pb2
from lingvo.core import py_utils
import six

from google.protobuf import text_format

FLAGS = tf.flags.FLAGS

# InferenceDeviceOptions contains options to configure inference on the device.
# device: Device to infer on.
# retain_device_placement: If true, the specified device in the generated
#   inference graph nodes will be retained. Otherwise, the specified device
#   will be cleared, so that the runtime can choose automatically.
# var_options: Options on handling variables. For TPUs, variables can be
#   either placed on device through 'ON_DEVICE' option, or treated as
#   constants with AS_CONSTANTS.
# gen_init_op: Whether to serialize initialization ops for the device. For TPUs,
#   servers can be initialized globally once, in which case this should be
#   turned off to avoid tripping initialization checks.
# dtype_override: Whether to override the dtype to use for activations and
# weights in the model. Options supported are None or tf.bfloat16.
InferenceDeviceOptions = collections.namedtuple('InferenceDeviceOptions', [
    'device', 'retain_device_placement', 'var_options', 'gen_init_op',
    'dtype_override', 'fprop_dtype_override'
])

_CONST_GUARANTEE = None


@contextlib.contextmanager
def NoConstGuaranteeScope():
  """Disallow const gauranteeing variable with-in scope."""
  global _CONST_GUARANTEE
  var_scope = tf.get_variable_scope()
  old_caching_device = var_scope.caching_device
  old_val = _CONST_GUARANTEE
  var_scope.set_caching_device(None)
  _CONST_GUARANTEE = False
  yield
  _CONST_GUARANTEE = old_val
  var_scope.set_caching_device(old_caching_device)


# Marks variable as constants for compilation
def MaybeGuaranteeConstGetter(getter, name, *args, **kwargs):
  global _CONST_GUARANTEE
  if _CONST_GUARANTEE:
    with tf.control_dependencies(None):
      return tf.guarantee_const(
          getter(name, *args, **kwargs), name=name + '/GuaranteeConst')
  else:
    return getter(name, *args, **kwargs)


@contextlib.contextmanager
def ConstGuaranteeScope():
  """Treats all variables under this scope as constants."""
  global _CONST_GUARANTEE
  var_scope = tf.get_variable_scope()
  old_custom_getter = var_scope.custom_getter
  old_caching_device = var_scope.caching_device
  old_val = _CONST_GUARANTEE
  var_scope.set_custom_getter(MaybeGuaranteeConstGetter)
  var_scope.set_caching_device(lambda op: op.device)
  _CONST_GUARANTEE = True
  yield
  _CONST_GUARANTEE = old_val
  var_scope.set_custom_getter(old_custom_getter)
  var_scope.set_caching_device(old_caching_device)


@contextlib.contextmanager
def _DummyScope():
  yield None


def _GetVarName(v):
  return v.name[:-len(':0')]


def _MakeVariableDictionary(variables):
  """Returns a dictionary with name -> tf.Variable() mapping."""
  vars_dict = {}
  for v in variables:
    vars_dict[_GetVarName(v)] = v
  return vars_dict


def IsTpu(device_options):
  return device_options.device == 'tpu'


def ShouldForceBfloat16ForWeightsAndActivations(device_options):
  return device_options.dtype_override == tf.bfloat16


def ShouldForceBfloat16ForActivations(device_options):
  return device_options.fprop_dtype_override == tf.bfloat16


def ConvertSubgraphDictToProto(subgraphs_dict):
  """Converts dict of subgraphs/feeds/fetches to InferenceGraph.

  Args:
    subgraphs_dict: Dict of (fetches, feeds) where each fetches/feeds is a
      NestedMap.

  Returns:
    Equivalent InferenceGraph.
  """
  # Build the output inference graph.
  inference_graph_proto = inference_graph_pb2.InferenceGraph()
  for subgraph_name, tensors in subgraphs_dict.items():
    fetches = tensors[0]
    feeds = tensors[1]

    # Rewrite fetches and feeds to map to their tensor name instead of
    # Tensor instance.
    named_fetches = {k: v.name for k, v in fetches.items() if v is not None}
    named_feeds = {k: v.name for k, v in feeds.items() if v is not None}

    # Export as subgraph.
    inference_graph_proto.subgraphs[subgraph_name].fetches.update(named_fetches)
    inference_graph_proto.subgraphs[subgraph_name].feeds.update(named_feeds)
  return inference_graph_proto


def GetOutputOpNames(graph,
                     inference_graph_proto,
                     subgraphs=None,
                     preserve_colocation_nodes=True,
                     preserve_saver_restore_nodes=False,
                     preserve_extra_ops=None):
  """Gets output op names from an inference graph.

  Args:
    graph: The tf graph.
    inference_graph_proto: an InferenceGraph proto.
    subgraphs: an optional list of subgraph names. If provided, only output ops
      from these subgraphs are preserved. Otherwise, all subgraphs are included.
    preserve_colocation_nodes: a Python bool, default to True. Preserves nodes
      colocating with the closure of output ops in the returned array.
    preserve_saver_restore_nodes: a Python bool, default to False. Preserves
      nodes for restoring according to inference_graph_proto.saver_def.
    preserve_extra_ops: an optional list of extra op names to preserve as long
      as they present in the graph.

  Returns:
    Array of tf op names that should be preserved in the graph.
  """
  output_op_names = set()

  def _GetOpName(tensor_or_op_name):
    """Returns the op name of the given node name."""
    # Tensor names have format <op_name>:<output_index>. Some inference
    # graphs put tensors and others put ops in the feeds/fetches (depends
    # on how it is used). We differentiate here. We still do the lookup in
    # the graph to sanity check (versus relying on the text manipulation).
    # If this logic ever breaks, TensorFlow will raise a ValueError with
    # a description of the syntax of each.
    if re.search(r':[0-9]+$', tensor_or_op_name):
      # Tensor-name.
      t = graph.get_tensor_by_name(tensor_or_op_name)
      return t.op.name
    else:
      op = graph.get_operation_by_name(tensor_or_op_name)
      return op.name

  for subgraph_name, subgraph in inference_graph_proto.subgraphs.items():
    if subgraphs and subgraph_name not in subgraphs:
      tf.logging.info('Skip subgraph %s.', subgraph_name)
      continue
    # Sometimes feeds aren't connected to any outputs but keep them in the graph
    # anyways to avoid errors.
    for tensor_or_op_name in (list(subgraph.feeds.values()) +
                              list(subgraph.fetches.values())):
      output_op_names.add(_GetOpName(tensor_or_op_name))

  if preserve_saver_restore_nodes:
    # Only nodes for restoring is preserved. saver_def.save_tensor_name is
    # skipped because it's only used for saving.
    saver_def = inference_graph_proto.saver_def
    for op_name in [saver_def.filename_tensor_name, saver_def.restore_op_name]:
      try:
        output_op_names.add(_GetOpName(op_name))
      except KeyError:
        tf.logging.info('Op/tensor %s not in the graph. Ignoring.' % op_name)

  if not preserve_colocation_nodes and not preserve_extra_ops:
    return sorted(list(output_op_names))

  # We also need to preserve any nodes that are used for colocation.
  # E.g., a node may have this attr:
  #   attr {
  #     key: "_class"
  #     value {
  #       list {
  #         s: "loc:@inference/embedding_lookup/Read/ReadVariableOp"
  #       }
  #     }
  #   }
  #
  # In this case, we need to make sure the node
  # inference/embedding_lookup/Read/ReadVariableOp is not pruned.
  #
  # TODO(zhifengc): It's possible that it's better to fix in
  # tf.graph_util.extract_sub_graph.
  graph_def = tf.graph_util.extract_sub_graph(graph.as_graph_def(),
                                              list(output_op_names))
  reachable_vars = [node.name for node in graph_def.node]

  for node in graph.get_operations():
    if preserve_extra_ops and node.name in preserve_extra_ops:
      output_op_names.add(node.name)
    elif preserve_colocation_nodes and '_class' in node.node_def.attr:
      for loc in node.node_def.attr['_class'].list.s:
        loc = six.ensure_text(loc, 'utf-8')
        if loc.startswith('loc:@'):
          loc_name = loc[5:]
          if loc_name not in reachable_vars:
            # Skip nodes that cannot be reached from the pruned graph.
            continue
          output_op_names.add(node.name)

  return sorted(list(output_op_names))


def _ParamExists(param_obj, param_name):
  """Tests whether param_name is contained in param_obj."""
  if not param_obj:
    return
  for k, _ in param_obj.IterParams():
    if k == param_name:
      return True
  return False


def _FreezeGraphFromCheckpoint(graph, saver, checkpoint, output_op_names):
  """Freezes a graph from a checkpoint.

  Args:
    graph: tf.Graph.
    saver: The tf.Saver to use for restoration.
    checkpoint: The checkpoint to restore.
    output_op_names: Names of output ops.

  Returns:
    Resulting tf.GraphDef.
  """
  sess = tf.Session(graph=graph, config=py_utils.SessionConfig())
  saver.restore(sess, checkpoint)
  return tf.graph_util.convert_variables_to_constants(
      sess, graph.as_graph_def(), output_op_names)


def _FreezeDefaults(graph, output_op_names):
  """Default initializes a graph and freezes it.

  Args:
    graph: tf.Graph.
    output_op_names: Names of output ops.

  Returns:
    Resulting tf.GraphDef.
  """
  with tf.Session(graph=graph, config=py_utils.SessionConfig()) as sess:
    sess.run(graph.get_operation_by_name('init_all_variables'))
    return tf.graph_util.convert_variables_to_constants(sess,
                                                        graph.as_graph_def(),
                                                        output_op_names)


class InferenceGraphExporter:
  """Class for exporting inference graphs."""

  @classmethod
  def Export(cls,
             model_cfg,
             model_task_name=None,
             device_options=InferenceDeviceOptions(
                 device='',
                 retain_device_placement=False,
                 var_options=None,
                 gen_init_op=True,
                 dtype_override=None,
                 fprop_dtype_override=None),
             freeze_checkpoint=None,
             freeze_defaults=False,
             export_path=None,
             subgraph_filter=None,
             random_seed=None,
             disable_packed_input=True):
    """Exports a InferenceGraph proto with piecewise subgraphs.

    Sets FLAGS.enable_asserts to False unless user explicitly sets it to True.

    Note: Enable FLAGS.pin_vars_to_cpu (default false) to make weight-sharing
    and multi-core inference on TPUs work properly.

    Args:
      model_cfg: a Params instance as returned by
        model_registry.GetParams(modelname, 'Test') or model_params.Model().
      model_task_name: The task to generate an inference graph for. Should be
        None for single-task models.
      device_options: Device options for the accelerator used for serving.
      freeze_checkpoint: The checkpoint to load. Loads and freezes the model if
        given.
      freeze_defaults: Default initializes the graph and freeze. Useful for
        early testing of downstream tools without having a checkpoint.
      export_path: If not None, write the inference graph in ASCII to this path.
      subgraph_filter: A string or a list of subgraph names. If not None or
        empty, export only this list of inference subgraphs.
      random_seed: Fixes the random seed in the exported inference graph.
      disable_packed_input: Disable packed input for inference writing purposes.

    Returns:
      InferenceGraph proto.

    Raises:
      ValueError: if the model does not support the listed subgraphs.
    """
    assert issubclass(model_cfg.cls, base_model.BaseModel)
    if device_options.dtype_override and device_options.fprop_dtype_override:
      raise ValueError(
          'device_options{dtype_override,fprop_dtype_override) can not both be'
          'set.')
    if subgraph_filter and not isinstance(subgraph_filter, (tuple, list)):
      subgraph_filter = [subgraph_filter]

    # Disable assertions unless user explicitly enables it.
    if FLAGS['enable_asserts'].using_default_value:
      FLAGS.enable_asserts = False

    # TODO(laurenzo): Work out how much we need to specify here in terms of
    # cluster configuration.
    cls._SetClusterParams(model_cfg.cluster, device_options)

    # Configure the model.
    model_cfg.random_seed = random_seed
    model_cfg.is_inference = True

    if disable_packed_input:

      def _DisablePackedInput(task):
        if (_ParamExists(task, 'encoder') and
            _ParamExists(task.encoder, 'packed_input')):
          task.encoder.packed_input = False
        if (_ParamExists(task, 'decoder') and
            _ParamExists(task.decoder, 'packed_input')):
          task.decoder.packed_input = False

      if issubclass(model_cfg.cls, base_model.MultiTaskModel):
        for _, task_param in model_cfg.task_params.IterParams():
          _DisablePackedInput(task_param)
      else:
        _DisablePackedInput(model_cfg.task)

    tf.logging.debug('Model %s params:', model_cfg.name)
    for line in model_cfg.ToText().split('\n'):
      tf.logging.debug('%s', line)

    # Instantiate the graph.
    graph = tf.Graph()
    with graph.as_default():
      tf.random.set_seed(random_seed)
      cluster = model_cfg.cluster.Instantiate()
      device = cluster.GetPlacer()
      tpu_const_scope = _DummyScope()
      if (IsTpu(device_options) and
          device_options.var_options == 'AS_CONSTANTS'):
        # Do not specify devices for variables if we are marking them as
        # constants.
        device = ''
        tpu_const_scope = ConstGuaranteeScope()

      with cluster, tf.device(device), tpu_const_scope:

        bfloat16_override = ShouldForceBfloat16ForWeightsAndActivations(
            device_options)

        if bfloat16_override:
          py_utils.UpdateDtype(model_cfg, tf.bfloat16)
          py_utils.UpdateFpropDtype(model_cfg, tf.bfloat16)

        act_bfloat16_override = ShouldForceBfloat16ForActivations(
            device_options)
        if act_bfloat16_override:
          py_utils.UpdateFpropDtype(model_cfg, tf.bfloat16)

        # Hard-code TPU-related flags prior to instantiating model.
        old_enable_asserts = FLAGS.enable_asserts
        old_xla_device = FLAGS.xla_device
        if IsTpu(device_options):
          FLAGS.enable_asserts = False
          FLAGS.xla_device = 'tpu'

        try:
          mdl = model_cfg.Instantiate()
          task = mdl.GetTask(model_task_name)

          variables_to_restore = (
              _MakeVariableDictionary(tf.global_variables()) if not mdl.ema else
              mdl.ema.variables_to_restore(mdl.variables_for_ema))

          if bfloat16_override:
            saver_var_spec = (
                bfloat16_variables
                .get_saver_spec_for_variables_with_bf16_overrides(
                    variables_to_restore))
          else:
            saver_var_spec = variables_to_restore

          saver = tf.train.Saver(saver_var_spec)
          tf.variables_initializer(
              tf.global_variables(), name='init_all_variables')
          if IsTpu(device_options) and device_options.gen_init_op:
            tf.group(tf.tpu.initialize_system(), name='tpu_init_op')

          if freeze_checkpoint or freeze_defaults:
            # Replace variables with tensors using tf.identity in theta before
            # freezing to avoid the graph referencing types of DT_RESOURCE.
            def AddIdentityToTheta(layer):
              layer._private_theta = layer._private_theta.Transform(tf.identity)  # pylint: disable=protected-access
              layer.children.Transform(AddIdentityToTheta)

            AddIdentityToTheta(task)

          inference_graph_proto = inference_graph_pb2.InferenceGraph()
          subgraphs_proto = task.Inference()
          if isinstance(subgraphs_proto, dict):
            subgraphs_proto = ConvertSubgraphDictToProto(subgraphs_proto)
          for name, subgraph in subgraphs_proto.subgraphs.items():
            if not subgraph_filter or name in subgraph_filter:
              inference_graph_proto.subgraphs[name].CopyFrom(subgraph)

          # Yes, graph collections are bad, however this seems to be the
          # easiest way to get this assets registered from
          # TextFileInitializer.
          assets_collection = tf.compat.v1.get_collection(
              tf.compat.v1.GraphKeys.ASSET_FILEPATHS)
          for asset in assets_collection:
            if asset.op.type == 'Const' and asset.op.get_attr(
                'dtype') == tf.dtypes.string:
              constant_value = asset.op.get_attr('value')
              if constant_value.string_val:
                tf.logging.info('Found asset file_path: %s',
                                constant_value.string_val[0])
                asset_file_def = inference_graph_proto.asset_file_def.add()
                asset_file_def.tensor_info.name = asset.name
                asset_file_def.filename = constant_value.string_val[0]

          # Add a table init op and global variable init op to the graph.
          # Tables can be declared anywhere in the graph, so this op has to be
          # added last.
          tf.tables_initializer(name='init_all_tables')
        finally:
          # Reset TPU-related flags after model instantiation.
          FLAGS.enable_asserts = old_enable_asserts
          FLAGS.xla_device = old_xla_device

    tf.logging.info('Graph contains ops: %r',
                    [op.name for op in graph.get_operations()])

    # Collection defs
    if not tf.executing_eagerly():
      meta_graph = tf.train.export_meta_graph(graph=graph)
      for key in meta_graph.collection_def:
        tf.logging.info('copying collection %s', key)
        inference_graph_proto.collection_def[key].CopyFrom(
            meta_graph.collection_def[key])
    else:
      tf.logging.warning('Not exporting collection defs '
                         'since operating in eager mode.')

    # Freezing.
    if freeze_defaults or freeze_checkpoint:
      output_op_names = GetOutputOpNames(
          graph,
          inference_graph_proto,
          preserve_colocation_nodes=False,
          preserve_saver_restore_nodes=False)
      if cls._DeviceSupportsFreezing(device_options):
        raise ValueError('freeze_checkpoint cannot be used with device ' +
                         device_options.device)
      if freeze_checkpoint:
        tf.logging.info('Freezing graph from checkpoint: %s', freeze_checkpoint)
        graph_def = _FreezeGraphFromCheckpoint(graph, saver, freeze_checkpoint,
                                               output_op_names)
      elif freeze_defaults:
        tf.logging.info('Default initializing graph and freezing.')
        graph_def = _FreezeDefaults(graph, output_op_names)
    else:
      inference_graph_proto.saver_def.CopyFrom(saver.as_saver_def())
      output_op_names = GetOutputOpNames(graph, inference_graph_proto)

      # Prune the graph to just the parts we need.
      # To support restoring, we have to not prune out the restore node.
      output_op_names.append('init_all_tables')
      output_op_names.append('init_all_variables')
      output_op_names.append('save/control_dependency')
      output_op_names.append('save/restore_all')
      if IsTpu(device_options) and device_options.gen_init_op:
        output_op_names.append('tpu_init_op')
      graph_def = graph.as_graph_def()
      tf.logging.info('Pruning graph to output ops: %r', output_op_names)
      graph_def = tf.graph_util.extract_sub_graph(graph_def, output_op_names)

    if not device_options.retain_device_placement:
      # Clear the device so that the runtime can choose.
      tf.logging.info('Clearing device placement for: %s',
                      device_options.device)
      for node in graph_def.node:
        node.ClearField('device')
      for function in graph_def.library.function:
        for node_def in function.node_def:
          node_def.ClearField('device')

    inference_graph_proto.graph_def.CopyFrom(graph_def)

    if export_path:
      with tf.io.gfile.GFile(export_path, 'w') as f:
        f.write(text_format.MessageToString(inference_graph_proto))
    return inference_graph_proto

  @classmethod
  def _SetClusterParams(cls, cluster_params, device_options):
    """Sets cluster params.

    Args:
      cluster_params: Model().cluster config.
      device_options: InferenceDeviceOptions.
    """

    def Update(p):
      """Update cluster params `p`."""
      p.name = '/job:localhost'
      p.replicas = 1
      p.tpus_per_replica = 1 if IsTpu(device_options) else 0
      p.gpus_per_replica = 0
      p.devices_per_split = 1

    cluster_params.mode = 'sync'
    cluster_params.job = 'decoder'
    cluster_params.add_summary = False
    cluster_params.do_eval = True
    Update(cluster_params.controller)
    Update(cluster_params.worker)
    Update(cluster_params.ps)
    Update(cluster_params.evaler)
    Update(cluster_params.decoder)
    Update(cluster_params.input)

  @classmethod
  def _DeviceSupportsFreezing(cls, device_options):
    return IsTpu(device_options)

[((45, 25, 48, 2), 'collections.namedtuple', 'collections.namedtuple', ({(45, 48, 45, 72): '"""InferenceDeviceOptions"""', (45, 74, 48, 1): "['device', 'retain_device_placement', 'var_options', 'gen_init_op',\n    'dtype_override', 'fprop_dtype_override']"}, {}), "('InferenceDeviceOptions', ['device',\n    'retain_device_placement', 'var_options', 'gen_init_op',\n    'dtype_override', 'fprop_dtype_override'])", False, 'import collections\n'), ((57, 14, 57, 37), 'lingvo.compat.get_variable_scope', 'tf.get_variable_scope', ({}, {}), '()', True, 'import lingvo.compat as tf\n'), ((82, 14, 82, 37), 'lingvo.compat.get_variable_scope', 'tf.get_variable_scope', ({}, {}), '()', True, 'import lingvo.compat as tf\n'), ((135, 26, 135, 62), 'lingvo.core.inference_graph_pb2.InferenceGraph', 'inference_graph_pb2.InferenceGraph', ({}, {}), '()', False, 'from lingvo.core import inference_graph_pb2\n'), ((184, 7, 184, 48), 're.search', 're.search', ({(184, 17, 184, 28): '""":[0-9]+$"""', (184, 30, 184, 47): 'tensor_or_op_name'}, {}), "(':[0-9]+$', tensor_or_op_name)", False, 'import re\n'), ((381, 4, 381, 56), 'lingvo.compat.logging.debug', 'tf.logging.debug', ({(381, 21, 381, 39): '"""Model %s params:"""', (381, 41, 381, 55): 'model_cfg.name'}, {}), "('Model %s params:', model_cfg.name)", True, 'import lingvo.compat as tf\n'), ((386, 12, 386, 22), 'lingvo.compat.Graph', 'tf.Graph', ({}, {}), '()', True, 'import lingvo.compat as tf\n'), ((71, 9, 71, 38), 'lingvo.compat.control_dependencies', 'tf.control_dependencies', ({(71, 33, 71, 37): 'None'}, {}), '(None)', True, 'import lingvo.compat as tf\n'), ((194, 6, 194, 57), 'lingvo.compat.logging.info', 'tf.logging.info', ({(194, 22, 194, 41): '"""Skip subgraph %s."""', (194, 43, 194, 56): 'subgraph_name'}, {}), "('Skip subgraph %s.', subgraph_name)", True, 'import lingvo.compat as tf\n'), ((273, 40, 273, 64), 'lingvo.core.py_utils.SessionConfig', 'py_utils.SessionConfig', ({}, {}), '()', False, 'from lingvo.core import py_utils\n'), ((383, 6, 383, 34), 'lingvo.compat.logging.debug', 'tf.logging.debug', ({(383, 23, 383, 27): '"""%s"""', (383, 29, 383, 33): 'line'}, {}), "('%s', line)", True, 'import lingvo.compat as tf\n'), ((388, 6, 388, 37), 'lingvo.compat.random.set_seed', 'tf.random.set_seed', ({(388, 25, 388, 36): 'random_seed'}, {}), '(random_seed)', True, 'import lingvo.compat as tf\n'), ((488, 11, 488, 33), 'lingvo.compat.executing_eagerly', 'tf.executing_eagerly', ({}, {}), '()', True, 'import lingvo.compat as tf\n'), ((489, 19, 489, 58), 'lingvo.compat.train.export_meta_graph', 'tf.train.export_meta_graph', (), '', True, 'import lingvo.compat as tf\n'), ((495, 6, 496, 58), 'lingvo.compat.logging.warning', 'tf.logging.warning', ({(495, 25, 496, 57): '"""Not exporting collection defs since operating in eager mode."""'}, {}), "(\n    'Not exporting collection defs since operating in eager mode.')", True, 'import lingvo.compat as tf\n'), ((528, 6, 528, 73), 'lingvo.compat.logging.info', 'tf.logging.info', ({(528, 22, 528, 55): '"""Pruning graph to output ops: %r"""', (528, 57, 528, 72): 'output_op_names'}, {}), "('Pruning graph to output ops: %r', output_op_names)", True, 'import lingvo.compat as tf\n'), ((529, 18, 529, 77), 'lingvo.compat.graph_util.extract_sub_graph', 'tf.graph_util.extract_sub_graph', ({(529, 50, 529, 59): 'graph_def', (529, 61, 529, 76): 'output_op_names'}, {}), '(graph_def, output_op_names)', True, 'import lingvo.compat as tf\n'), ((533, 6, 534, 44), 'lingvo.compat.logging.info', 'tf.logging.info', ({(533, 22, 533, 57): '"""Clearing device placement for: %s"""', (534, 22, 534, 43): 'device_options.device'}, {}), "('Clearing device placement for: %s', device_options.device)", True, 'import lingvo.compat as tf\n'), ((289, 38, 289, 62), 'lingvo.core.py_utils.SessionConfig', 'py_utils.SessionConfig', ({}, {}), '()', False, 'from lingvo.core import py_utils\n'), ((399, 20, 399, 37), 'lingvo.compat.device', 'tf.device', ({(399, 30, 399, 36): 'device'}, {}), '(device)', True, 'import lingvo.compat as tf\n'), ((491, 8, 491, 53), 'lingvo.compat.logging.info', 'tf.logging.info', ({(491, 24, 491, 47): '"""copying collection %s"""', (491, 49, 491, 52): 'key'}, {}), "('copying collection %s', key)", True, 'import lingvo.compat as tf\n'), ((509, 8, 509, 80), 'lingvo.compat.logging.info', 'tf.logging.info', ({(509, 24, 509, 60): '"""Freezing graph from checkpoint: %s"""', (509, 62, 509, 79): 'freeze_checkpoint'}, {}), "('Freezing graph from checkpoint: %s', freeze_checkpoint)", True, 'import lingvo.compat as tf\n'), ((544, 11, 544, 46), 'lingvo.compat.io.gfile.GFile', 'tf.io.gfile.GFile', ({(544, 29, 544, 40): 'export_path', (544, 42, 544, 45): '"""w"""'}, {}), "(export_path, 'w')", True, 'import lingvo.compat as tf\n'), ((210, 8, 210, 77), 'lingvo.compat.logging.info', 'tf.logging.info', ({(210, 24, 210, 76): "('Op/tensor %s not in the graph. Ignoring.' % op_name)"}, {}), "('Op/tensor %s not in the graph. Ignoring.' % op_name)", True, 'import lingvo.compat as tf\n'), ((240, 14, 240, 43), 'six.ensure_text', 'six.ensure_text', ({(240, 30, 240, 33): 'loc', (240, 35, 240, 42): '"""utf-8"""'}, {}), "(loc, 'utf-8')", False, 'import six\n'), ((405, 10, 405, 54), 'lingvo.core.py_utils.UpdateDtype', 'py_utils.UpdateDtype', ({(405, 31, 405, 40): 'model_cfg', (405, 42, 405, 53): 'tf.bfloat16'}, {}), '(model_cfg, tf.bfloat16)', False, 'from lingvo.core import py_utils\n'), ((406, 10, 406, 59), 'lingvo.core.py_utils.UpdateFpropDtype', 'py_utils.UpdateFpropDtype', ({(406, 36, 406, 45): 'model_cfg', (406, 47, 406, 58): 'tf.bfloat16'}, {}), '(model_cfg, tf.bfloat16)', False, 'from lingvo.core import py_utils\n'), ((411, 10, 411, 59), 'lingvo.core.py_utils.UpdateFpropDtype', 'py_utils.UpdateFpropDtype', ({(411, 36, 411, 45): 'model_cfg', (411, 47, 411, 58): 'tf.bfloat16'}, {}), '(model_cfg, tf.bfloat16)', False, 'from lingvo.core import py_utils\n'), ((436, 18, 436, 48), 'lingvo.compat.train.Saver', 'tf.train.Saver', ({(436, 33, 436, 47): 'saver_var_spec'}, {}), '(saver_var_spec)', True, 'import lingvo.compat as tf\n'), ((451, 34, 451, 70), 'lingvo.core.inference_graph_pb2.InferenceGraph', 'inference_graph_pb2.InferenceGraph', ({}, {}), '()', False, 'from lingvo.core import inference_graph_pb2\n'), ((462, 30, 463, 53), 'lingvo.compat.compat.v1.get_collection', 'tf.compat.v1.get_collection', ({(463, 14, 463, 52): 'tf.compat.v1.GraphKeys.ASSET_FILEPATHS'}, {}), '(tf.compat.v1.GraphKeys.ASSET_FILEPATHS)', True, 'import lingvo.compat as tf\n'), ((478, 10, 478, 55), 'lingvo.compat.tables_initializer', 'tf.tables_initializer', (), '', True, 'import lingvo.compat as tf\n'), ((513, 8, 513, 67), 'lingvo.compat.logging.info', 'tf.logging.info', ({(513, 24, 513, 66): '"""Default initializing graph and freezing."""'}, {}), "('Default initializing graph and freezing.')", True, 'import lingvo.compat as tf\n'), ((545, 16, 545, 66), 'google.protobuf.text_format.MessageToString', 'text_format.MessageToString', ({(545, 44, 545, 65): 'inference_graph_proto'}, {}), '(inference_graph_proto)', False, 'from google.protobuf import text_format\n'), ((430, 16, 432, 41), 'lingvo.core.bfloat16_variables.get_saver_spec_for_variables_with_bf16_overrides', 'bfloat16_variables.get_saver_spec_for_variables_with_bf16_overrides', ({(432, 20, 432, 40): 'variables_to_restore'}, {}), '(\n    variables_to_restore)', False, 'from lingvo.core import bfloat16_variables\n'), ((438, 14, 438, 35), 'lingvo.compat.global_variables', 'tf.global_variables', ({}, {}), '()', True, 'import lingvo.compat as tf\n'), ((425, 38, 425, 59), 'lingvo.compat.global_variables', 'tf.global_variables', ({}, {}), '()', True, 'import lingvo.compat as tf\n'), ((440, 21, 440, 47), 'lingvo.compat.tpu.initialize_system', 'tf.tpu.initialize_system', ({}, {}), '()', True, 'import lingvo.compat as tf\n'), ((469, 16, 470, 61), 'lingvo.compat.logging.info', 'tf.logging.info', ({(469, 32, 469, 59): '"""Found asset file_path: %s"""', (470, 32, 470, 60): 'constant_value.string_val[0]'}, {}), "('Found asset file_path: %s', constant_value.string_val[0])", True, 'import lingvo.compat as tf\n')]