<commit_before># Copyright 2017 Google Inc. 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.
"""Tools for analyzing a TensorFlow graph.

This module exports the function determine_ready_tensors_and_table_initializers
which analyzes a TensorFlow graph to determine which tensors and table
initializers are "ready".  The concept of readiness arises as tf.Transform
works by building a single TF graph containing placeholders for the outputs
of analyzers.  These placeholders are progressively replaced by constants in
a number of phases, where in each phase we run some analyzers and replace their
outputs with constants.  We analyze the structure of the graph to determine
which analyzers to run in each phase.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import itertools
import uuid

# GOOGLE-INITIALIZATION

import six
import tensorflow as tf
from tensorflow_transform import analyzer_nodes
from tensorflow_transform import tf_utils
from tensorflow.python.framework import function_def_to_graph  # pylint: disable=g-direct-tensorflow-import

_INITIALIZABLE_TABLE_OP_TYPES = [
    'CuckooTable',
    'CuckooTableV2',
    'HashTable',
    'HashTableV2',
    'IndexTable',
    'IndexTableV2',
]

_TABLE_INIT_OP_TYPES = [
    'InitializeTable',
    'InitializeTableV2',
    'InitializeTableFromTextFile',
    'InitializeTableFromTextFileV2',
    'LookupTableImport',
    'LookupTableImportV2',
]


def _decompose_tensor(tensor):
  """Returns the raw components of a Tensor / SparseTensor / RaggedTensor."""
  if isinstance(tensor, tf.SparseTensor):
    yield tensor.indices
    yield tensor.values
    yield tensor.dense_shape
  elif isinstance(tensor, tf.RaggedTensor):
    yield tensor.row_splits
    v = tensor.values  # values is a potentially ragged tensor.
    while isinstance(v, tf.RaggedTensor):
      yield v.row_splits
      v = v.values
    yield v
  else:
    yield tensor


def _get_func_graph_for_name(graph, func_name):
  """Returns the FuncGraph associated to the given func_name if possible."""
  func = graph._get_function(str(func_name))  # pylint: disable=protected-access
  if not func:
    raise ValueError(
        'Function {} does not exist in the graph.'.format(func_name))
  if not hasattr(func, 'graph'):
    # This is a _DefinedFunction.
    return function_def_to_graph.function_def_to_graph(func.definition)
  else:
    return func.graph


class _UnexpectedPlaceholderError(Exception):

  def __init__(self, op, func_graph_name):
    tensor = op.outputs[0]
    msg = 'An unexpected placeholder was encountered ({})'.format(tensor)
    super(_UnexpectedPlaceholderError, self).__init__(msg)
    self.tensor = tensor
    self.func_graph_name = func_graph_name

<commit_msg>Modifies the graph_tools logic for finding a FuncGraph for a function name, which may require iterating over outer graphs due to a change in TF 2.1.<commit_after># Copyright 2017 Google Inc. 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.
"""Tools for analyzing a TensorFlow graph.

This module exports the function determine_ready_tensors_and_table_initializers
which analyzes a TensorFlow graph to determine which tensors and table
initializers are "ready".  The concept of readiness arises as tf.Transform
works by building a single TF graph containing placeholders for the outputs
of analyzers.  These placeholders are progressively replaced by constants in
a number of phases, where in each phase we run some analyzers and replace their
outputs with constants.  We analyze the structure of the graph to determine
which analyzers to run in each phase.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import itertools
import uuid

# GOOGLE-INITIALIZATION

import six
import tensorflow as tf
from tensorflow_transform import analyzer_nodes
from tensorflow_transform import tf_utils
from tensorflow.python.framework import function_def_to_graph  # pylint: disable=g-direct-tensorflow-import

_INITIALIZABLE_TABLE_OP_TYPES = [
    'CuckooTable',
    'CuckooTableV2',
    'HashTable',
    'HashTableV2',
    'IndexTable',
    'IndexTableV2',
]

_TABLE_INIT_OP_TYPES = [
    'InitializeTable',
    'InitializeTableV2',
    'InitializeTableFromTextFile',
    'InitializeTableFromTextFileV2',
    'LookupTableImport',
    'LookupTableImportV2',
]


def _decompose_tensor(tensor):
  """Returns the raw components of a Tensor / SparseTensor / RaggedTensor."""
  if isinstance(tensor, tf.SparseTensor):
    yield tensor.indices
    yield tensor.values
    yield tensor.dense_shape
  elif isinstance(tensor, tf.RaggedTensor):
    yield tensor.row_splits
    v = tensor.values  # values is a potentially ragged tensor.
    while isinstance(v, tf.RaggedTensor):
      yield v.row_splits
      v = v.values
    yield v
  else:
    yield tensor


def _get_func_graph_for_name(graph, func_name):
  """Returns the FuncGraph associated to the given func_name if possible."""
  while graph is not None:
    func = graph._get_function(str(func_name))  # pylint: disable=protected-access
    if func is not None:
      if hasattr(func, 'graph'):
        return func.graph
      func_graph = function_def_to_graph.function_def_to_graph(func.definition)
      if func_graph is not None:
        return func_graph
    if hasattr(graph, 'outer_graph'):
      graph = graph.outer_graph
    else:
      raise ValueError(
          'Function {} does not exist in the graph.'.format(func_name))


class _UnexpectedPlaceholderError(Exception):

  def __init__(self, op, func_graph_name):
    tensor = op.outputs[0]
    msg = 'An unexpected placeholder was encountered ({})'.format(tensor)
    super(_UnexpectedPlaceholderError, self).__init__(msg)
    self.tensor = tensor
    self.func_graph_name = func_graph_name


{'fraction_non_alphanumeric': 0.06721715196291468, 'fraction_numerical': 0.004635665652614805, 'mean_word_length': 4.331274131274132, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5537133', 'n_tokens_mistral': 1943, 'n_tokens_neox': 1910, 'n_words': 711}