# Copyright 2017 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.

# ==============================================================================

"""Define tflite op hints (intrinsic operations).

This essentially allows defining a TensorFlow API for tflite operations in

Python with hints on how they are represented in TensorFlow Lite. This basically

is a form of tflite intrinsic. It wraps a subpart of a TensorFlow execution

graph and is useful for LSTMs and other complicated TensorFlow constructions

that are difficult to pattern match in TOCO, but are represented by a single

accelerated tflite op.

Example:

def tflite_cool_activation(input):

# A cool activation function.

custom = tf.lite.OpHint("cool_activation")

input, = custom.add_inputs(input)

output = tf.sigmoid(input) * input

output, = custom.add_outputs(output)

return output

image = tf.compat.v1.placeholder(tf.float32, (1, 16, 16, 1))

output = tf.identity(tflite_cool_activation(image))

session = tf.compat.v1.Session()

graphdef_to_convert = tf.lite.convert_op_hints_to_stubs(session)

tflite_graph = tf.compat.v1.lite.toco_convert(

graphdef_to_convert, [image], [output])

with open("/tmp/graph.fb", "wb") as fp:

fp.write(tflite_graph)

How does it work?:

OpHint is a helper that you use when defining a vanilla python function.

It allows you to wrap arguments with tf.identities with some custom attributes.

These attributes allow you to find the original block of ops that was created.

For example, if you use cool_activation above you essentially get:

a_input = tf.identity()

result = tf.multiply(tf.sigmoid(a_input), a_input)

output = tf.identity()

a_input, output are identities that have parameters representing

what argument they are, what the name of the function they should turn into

in tf lite as well as a guid that uniquely identifies a particular invocation.

Once you have built your whole tensorflow graph, you can run it and train it

as usual, but after you have done that, you need to convert the graph into

a form that replaces these subgraphs wrapped in identities to stub ops. These

ops don't actually exist in the normal TensorFlow runtime, but will be

understood by toco later.

"""

# TODO(aselle): Make this use generic graph transformations.

# TODO(aselle): _tensor_name_base should be called _tensor_name_to_op_name.

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import collections as _collections

import copy as _copy

import json as _json

import uuid as _uuid

import six as _six

from tensorflow.core.framework import attr_value_pb2 as _attr_value_pb2

from tensorflow.core.framework import graph_pb2 as _graph_pb2

from tensorflow.core.framework import node_def_pb2 as _node_def_pb2

from tensorflow.python.framework import ops as _ops

# TODO(aselle): publicize these apis if we continue to use these.

from tensorflow.python.framework.graph_util_impl import _bfs_for_reachable_nodes

from tensorflow.python.framework.graph_util_impl import _extract_graph_summary

from tensorflow.python.ops import array_ops as _array_ops

from tensorflow.python.util import compat as _compat

from tensorflow.python.util.all_util import remove_undocumented

from tensorflow.python.util.tf_export import tf_export as _tf_export

@_tf_export(v1=["lite.OpHint"])

class OpHint(object):

return [self._outputs.add(arg) for arg in args]

class _LiteOperand(object):

raise RuntimeError("Unimplemented abstract method.")

class _LiteSingleOperand(_LiteOperand):

return str(self.name)

class _LiteAggregateOperand(_LiteOperand):

return s

class _LiteFuncCall(object):

(self.function_name, self.uuid, self.level, inputs_str, outputs_str))

def _find_all_hints_in_nodes(nodes):

return func_calls

def _extract_topology_sequence_mapping(nodes):

return dict(

(_tensor_name_base(node.name), idx) for idx, node in enumerate(nodes))

def _find_children_hints_in_while_loop(function_def, nodes_mapping):

return ordered_children_hints, new_nodes

def _find_children_hints(call, graph_def):

return children_hints, out, function_def_nodes

def _tensor_name_base(full_tensor_name):

return full_tensor_name.split(":")[0]

def _tensorflow_output_name(tensor_name, output_index):

return tensor_name if output_index == 0 else "%s:%d" % (tensor_name,

output_index)

# TODO(aselle): This should be converted to grappler in the future.

def _check_subgraph_closed(n, reachable_by_input, input_nodes_set,

]

# TODO(aselle): This should be converted to grappler in the future.

def _convert_single_op_hint_to_stub(call,