Class OpHint

tensorflow/lite/python/op_hint.py:93–462

A class that helps build tflite function invocations. It allows you to take a bunch of TensorFlow ops and annotate the construction such that toco knows how to convert it to tflite. This embeds a pseudo function in a TensorFlow graph. This allows embedding high-level API usage information i

@_tf_export(v1=["lite.OpHint"])
class OpHint(object):
  """A class that helps build tflite function invocations.

  It allows you to take a bunch of TensorFlow ops and annotate the construction
  such that toco knows how to convert it to tflite. This embeds a pseudo
  function in a TensorFlow graph. This allows embedding high-level API usage
  information in a lower level TensorFlow implementation so that an alternative
  implementation can be substituted later.

  Essentially, any "input" into this pseudo op is fed into an identity, and
  attributes are added to that input before being used by the constituent ops
  that make up the pseudo op. A similar process is done to any output that
  is to be exported from the current op.

  """
  # TODO(aselle): When TensorFlow functions functionality works for arbitrary
  # constructs, this mechanism can be retired and changed to use python defun's.

  # Attr constants that are used for representation in the GraphDef. These
  # will be used on every Identity op that is involved in a total OpHint.

  # Name of the OpHint function (cosmetic).
  FUNCTION_NAME_ATTR = "_tflite_function_name"
  # UUID of the function (each OpHint gets a new uuid).
  FUNCTION_UUID_ATTR = "_tflite_function_uuid"
  # The input index of the input (or nothing if it is an output).
  FUNCTION_INPUT_INDEX_ATTR = "_tflite_function_input_index"
  # The output index of the output (or nothing if it is an input).
  FUNCTION_OUTPUT_INDEX_ATTR = "_tflite_function_output_index"
  # An index that orders aggregate arguments. Aggregate arguments are ones
  # that are separate but will be fused horizontally. For example a static LSTM
  # has a lstm cell for each time step. Each one has a separate opHint, but a
  # fused SequentialLSTM will treat this as a single tensor.
  FUNCTION_SORT_INDEX_ATTR = "_tflite_function_sort_index"
  # The way in which multiple parts of the aggregate argument will be joined
  # into a fused operand. Valid options are OpHint.AGGREGATE_FIRST,
  # OpHint.AGGREGATE_LAST, OpHint.AGGREGATE_STACK.
  FUNCTION_AGGREGATE_ATTR = "_tflite_function_aggregate"
  # On fused OpHint stub, the order of inputs that the final LSTM call will
  # have. What this means is that the TensorFlow order might be
  # "foo", "bar", "stuff" and you might want the TF lite op order to be
  # "stuff", "foo", "bar", -1 (where -1 is unused). So you would set this
  # attribute to [2, 0, 1, -1].
  TFLITE_INPUT_INDICES = "_tflite_input_indices"
  # OpHint level.
  FUNCTION_LEVEL_ATTR = "_tflite_ophint_level"
  # Ophint internal mapping, this is for high level Ophint only.
  # This basically contains three kinds of mapping:
  #   1) How parental ophinted inputs map to the first child ophinted inputs;
  #   2) How internal children nodes are connected;
  #   3) How parental ophinted outputs map to the last child ophinted outputs.
  CHILDREN_INPUTS_MAPPINGS = "_tflite_children_ophint_inputs_mapping"

  # Types of aggregations
  #  stack: stacks all ophints with matching tags. i.e. for a static rnn.
  #   specifically, this is good for an input or output to a static rnn cell.
  AGGREGATE_STACK = "stack"
  # first: only takes the first output (one with lowest sort index)