Supported Select TensorFlow operators

TensorFlow core operators

The following is an exhaustive list of TensorFlow core operations that are supported by TensorFlow Lite runtime with the Select TensorFlow Ops feature.

TensorFlow Text and SentencePiece operators

The following TensorFlow Text and SentencePiece operators are supported if you use the Python API for conversion and import those libraries.

TF.Text operators:

  • CaseFoldUTF8
  • ConstrainedSequence
  • MaxSpanningTree
  • NormalizeUTF8
  • NormalizeUTF8WithOffsetsMap
  • RegexSplitWithOffsets
  • RougeL
  • SentenceFragments
  • SentencepieceOp
  • SentencepieceTokenizeOp
  • SentencepieceTokenizeWithOffsetsOp
  • SentencepieceDetokenizeOp
  • SentencepieceVocabSizeOp
  • SplitMergeTokenizeWithOffsets
  • UnicodeScriptTokenizeWithOffsets
  • WhitespaceTokenizeWithOffsets
  • WordpieceTokenizeWithOffsets

SentencePiece operators:

  • SentencepieceGetPieceSize
  • SentencepiecePieceToId
  • SentencepieceIdToPiece
  • SentencepieceEncodeDense
  • SentencepieceEncodeSparse
  • SentencepieceDecode

The following snippet shows how to convert models with the above operators:

import tensorflow as tf
# These imports are required to load operators' definition.
import tensorflow_text as tf_text
import sentencepiece as spm

converter = tf.lite.TFLiteConverter.from_keras_model(your_model)
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
model_data = converter.convert()

On the runtime side, it is also required to link the TensorFlow Text or SentencePiece library into the final app or binary.

User's defined Operators

If you created your own TensorFlow operators, you can also convert models containing them to TensorFlow Lite by listing required operators in the experimental_select_user_tf_ops as following:

import tensorflow as tf

ops_module = tf.load_op_library('./your_ops_library.so')

converter = tf.lite.TFLiteConverter.from_saved_model(your_model)
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
converter.target_spec.experimental_select_user_tf_ops = [
    'your_op_name1',
    'your_op_name2'
]
model_data = converter.convert()

On the runtime side, it is also required to link your operators library into the final app or binary.

Add TensorFlow core operators to the allowed list.

If you hit the case where the TensorFlow core operators are not in the above allowed list, you can report the feature request at here with the names of the TensorFlow core operators, not listed in the allowed list.

You can also create own your pull request from the source code. For example, if you want to add the raw_ops.StringToNumber op in the allowed list, there are three places to update like this commit.

(1) Add the operator kernel source code to the portable_extended_ops_group2 BUILD rule.

filegroup(
    name = "portable_extended_ops_group2",
    srcs = [
        ...
+       "string_to_number_op.cc",

        ...
    ],
)

In order to find the relevant operator kernel source file under the tensorflow/core/kernels directory, you can search the source code location, which contains the following kernel declaration with the operator name:

REGISTER_KERNEL_BUILDER(Name("StringToNumber")                 \
                            .Device(DEVICE_CPU)                \
                            .TypeConstraint<type>("out_type"), \
                        StringToNumberOp<type>)

If there are any header files under the tensorflow/core/kernels directory, required in the operator kernel source code, you need to add the header file into the portable_extended_ops_headers BUILD rule as the follows:

filegroup(
    name = "portable_extended_ops_headers",
    srcs = [
        ...
+       "string_util.h",

        ...
    ],
)

(2) Add the operator name to the allowed list.

The allowed list is defined in the tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops.cc. The TensorFlow core operator name is need to be listed in order to be allowed through the Select TF option.

static const std::set<std::string>* allowlisted_flex_ops =
    new std::set<std::string>({
        ...
+       "StringToNumber",

        ...
    });

Since the above list is sorted in alphabetical order, it makes sure to place the name in the right place.

(3) Add the operator name to this guide page.

To show the operator inclusion to the other developers, this guide page should be updated as well. This page is located at the tensorflow/lite/g3doc/guide/op_select_allowlist.md.

## TensorFlow core operators

The following is an exhaustive list of TensorFlow core operations that are
supported by TensorFlow Lite runtime with the Select TensorFlow Ops feature.

...
+*   `raw_ops.StringToNumber`
...

Since the above list is sorted in alphabetical order, it makes sure to place the name in the right place.