text: # Copyright 2015 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.
  # ==============================================================================
  """Simple graph matching functions."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from six import string_types
  
  from tensorflow.contrib.graph_editor import select
  from tensorflow.python.framework import ops as tf_ops
  
  __all__ = [
      "op_type",
      "OpMatcher",
  ]
  
  
  def _make_graph_match(graph_match):
    """Convert to a OpMatcher instance."""
    if graph_match is None:
      return None
    if not isinstance(graph_match, OpMatcher):
      graph_match = OpMatcher(graph_match)
    return graph_match
  
  
  def op_type(op_types, op=None):
    """Check if an op is of the given type.
  
    Args:
      op_types: tuple of strings containing the types to check against.
        For instance: ("Add", "Const")
      op: the operation to check (or None).
    Returns:
      if op is not None, return True if the op is of the correct type.
      if op is None, return a lambda function which does the type checking.
    """
    if isinstance(op_types, string_types):
      op_types = (op_types)
    if op is None:
      return lambda op: op.node_def.op in op_types
    else:
      return op.node_def.op in op_types
  
  
  class OpMatcher(object):
    """Graph match class."""
  
    def __init__(self, positive_filter):
      """Graph match constructor."""
      self.positive_filters = []
      self.input_op_matches = None
      self.control_input_op_matches = None
      self.output_op_matches = None
      positive_filter = self._finalize_positive_filter(positive_filter)
      self.positive_filters.append(positive_filter)
  
    def _finalize_positive_filter(self, elem):
      """Convert to a filter function."""
      if select.can_be_regex(elem):
        regex_ = select.make_regex(elem)
        return lambda op, regex=regex_: regex.search(op.name) is not None
      elif isinstance(elem, tf_ops.Operation):
        return lambda op, match_op=elem: op is match_op
      elif callable(elem):
        return elem
      elif elem is True:
        return lambda op: True
      else:
        raise ValueError("Cannot finalize the positive filter: {}".format(elem))
  
    def __call__(self, op):
      """Evaluate if the op matches or not."""
      if not isinstance(op, tf_ops.Operation):
        raise TypeError("Expect tf.Operation, got: {}".format(type(op)))
      for positive_filter in self.positive_filters:
        if not positive_filter(op):
          return False
      if self.input_op_matches is not None:
        if len(op.inputs) != len(self.input_op_matches):
          return False
        for input_t, input_op_match in zip(op.inputs, self.input_op_matches):
          if input_op_match is None:
            continue
          if not input_op_match(input_t.op):
            return False
      if self.control_input_op_matches is not None:
        if len(op.control_inputs) != len(self.control_input_op_matches):
          return False
        for cinput_op, cinput_op_match in zip(op.control_inputs,
                                              self.control_input_op_matches):
          if cinput_op_match is None:
            continue
          if not cinput_op_match(cinput_op):
            return False
      if self.output_op_matches is not None:
        if len(op.outputs) != len(self.output_op_matches):
          return False
        for output_t, output_op_matches in zip(op.outputs,
                                               self.output_op_matches):
          if output_op_matches is None:
            continue
          if len(output_t.consumers()) != len(output_op_matches):
            return False
          for consumer_op, consumer_op_match in zip(output_t.consumers(),
                                                    output_op_matches):
            if consumer_op_match is None:
              continue
            if not consumer_op_match(consumer_op):
              return False
      return True
  
    def input_ops(self, *args):
      """Add input matches."""
      if self.input_op_matches is not None:
        raise ValueError("input_op_matches is already set.")
      self.input_op_matches = []
      for input_match in args:
        self.input_op_matches.append(_make_graph_match(input_match))
      return self
  
    def control_input_ops(self, *args):
      """Add input matches."""
      if self.control_input_op_matches is not None:
        raise ValueError("control_input_op_matches is already set.")
      self.control_input_op_matches = []
      for input_match in args:
        self.control_input_op_matches.append(_make_graph_match(input_match))
      return self
  
    def output_ops(self, *args):
      """Add output matches."""
      if self.output_op_matches is not None:
        raise ValueError("output_op_matches is already set.")
      self.output_op_matches = []
      for consumer_op_matches in args:
        if consumer_op_matches is None:
          self.output_op_matches.append(None)
        if not isinstance(consumer_op_matches, list):
          consumer_op_matches = [consumer_op_matches]
        consumer_op_matches = [_make_graph_match(consumer_op_match)
                               for consumer_op_match in consumer_op_matches]
        self.output_op_matches.append(consumer_op_matches)
      return self
  