chore: import upstream snapshot with attribution

tensorflow/lite/python/authoring/BUILD

@@ -0,0 +1,34 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "authoring",
+    srcs = [
+        "authoring.py",
+    ],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
+        "//tensorflow/lite/python:convert",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "authoring_test",
+    srcs = ["authoring_test.py"],
+    strict_deps = True,
+    deps = [
+        ":authoring",
+        #internal proto upb dep
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/lite/python:lite",
+    ],
+)

tensorflow/lite/python/authoring/authoring.py

@@ -0,0 +1,301 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""TensorFlow Authoring tool package for TFLite compatibility.
+
+WARNING: The package is experimental and subject to change.
+
+This package provides a way to check TFLite compatibility at model authoring
+time.
+
+Example:
+    @tf.lite.experimental.authoring.compatible
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[None], dtype=tf.float32)
+    ])
+    def f(x):
+      return tf.cosh(x)
+
+    result = f(tf.constant([0.0]))
+
+    > COMPATIBILITY WARNING: op 'tf.Cosh' require(s) "Select TF Ops" for model
+    > conversion for TensorFlow Lite.
+    > Op: tf.Cosh
+    >   - tensorflow/python/framework/op_def_library.py:xxx
+    >   - tensorflow/python/ops/gen_math_ops.py:xxx
+    >   - simple_authoring.py:xxx
+"""
+import functools
+from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
+# pylint: disable=g-import-not-at-top
+from tensorflow.lite.python import convert
+from tensorflow.lite.python import lite
+from tensorflow.python.util.tf_export import tf_export as _tf_export
+
+
+_CUSTOM_OPS_HDR = "Custom ops: "
+_TF_OPS_HDR = "TF Select ops: "
+_AUTHORING_ERROR_HDR = "COMPATIBILITY ERROR"
+_AUTHORING_WARNING_HDR = "COMPATIBILITY WARNING"
+_FUNC_GRAPH_SRC_PATH = "tensorflow/python/framework/func_graph.py"
+
+
+class CompatibilityError(Exception):
+  """Raised when an error occurs with TFLite compatibility."""
+  pass
+
+
+class _Compatible:
+  """A decorator class to check TFLite compatibility created by `lite.experimental.authoring.compatible`."""
+
+  def __init__(self,
+               target,
+               converter_target_spec=None,
+               converter_allow_custom_ops=None,
+               raise_exception=False):
+    """Initialize the decorator object.
+
+    Here is the description of the object variables.
+    - _func     : decorated function.
+    - _obj_func : for class object, we need to use this object to provide `self`
+                  instance as 1 first argument.
+    - _verified : whether the compatibility is checked or not.
+
+    Args:
+      target: decorated function.
+      converter_target_spec : target_spec of TFLite converter parameter.
+      converter_allow_custom_ops : allow_custom_ops of TFLite converter
+          parameter.
+      raise_exception : to raise an exception on compatibility issues.
+          User need to use get_compatibility_log() to check details.
+    """
+    functools.update_wrapper(self, target)
+    self._func = target
+    self._obj_func = None
+    self._verified = False
+    self._log_messages = []
+    self._raise_exception = raise_exception
+    self._converter_target_spec = converter_target_spec
+    self._converter_allow_custom_ops = converter_allow_custom_ops
+
+  def __get__(self, instance, cls):
+    """A Python descriptor interface."""
+    self._obj_func = self._func.__get__(instance, cls)
+    return self
+
+  def _get_func(self):
+    """Returns decorated function object.
+
+    For a class method, use self._obj_func to provide `self` instance.
+    """
+    if self._obj_func is not None:
+      return self._obj_func
+    else:
+      return self._func
+
+  def __call__(self, *args, **kwargs):  # pylint: disable=g-doc-args
+    """Calls decorated function object.
+
+    Also verifies if the function is compatible with TFLite.
+
+    Returns:
+      A execution result of the decorated function.
+    """
+
+    if not self._verified:
+      model = self._get_func()
+      concrete_func = model.get_concrete_function(*args, **kwargs)
+      converter = lite.TFLiteConverterV2.from_concrete_functions(
+          [concrete_func], model)
+      # Set provided converter parameters
+      if self._converter_target_spec is not None:
+        converter.target_spec = self._converter_target_spec
+      if self._converter_allow_custom_ops is not None:
+        converter.allow_custom_ops = self._converter_allow_custom_ops
+      try:
+        converter.convert()
+      except convert.ConverterError as err:
+        self._decode_error(err)
+      finally:
+        self._verified = True
+
+    return self._get_func()(*args, **kwargs)
+
+  def get_concrete_function(self, *args, **kwargs):
+    """Returns a concrete function of the decorated function."""
+    return self._get_func().get_concrete_function(*args, **kwargs)
+
+  def _get_location_string(self, location):
+    """Dump location of ConveterError.errors.location."""
+    callstack = []
+    for single_call in reversed(location.call):
+      if (location.type ==
+          converter_error_data_pb2.ConverterErrorData.CALLSITELOC):
+        callstack.append(
+            f"  - {single_call.source.filename}:{single_call.source.line}")
+      else:
+        callstack.append(str(single_call))
+    callstack_dump = "\n".join(callstack)
+    return callstack_dump
+
+  def _dump_error_details(self, ops, locations):
+    """Dump the list of ops and locations."""
+    for i in range(0, len(ops)):
+      callstack_dump = self._get_location_string(locations[i])
+      err_string = f"Op: {ops[i]}\n{callstack_dump}\n"
+      self._log(err_string)
+
+  def _decode_error_legacy(self, err):
+    """Parses the given legacy ConverterError for OSS."""
+    for line in str(err).splitlines():
+      # Check custom op usage error.
+      if line.startswith(_CUSTOM_OPS_HDR):
+        custom_ops = line[len(_CUSTOM_OPS_HDR):]
+        err_string = (
+            f"{_AUTHORING_ERROR_HDR}: op '{custom_ops}' is(are) not natively "
+            "supported by TensorFlow Lite. You need to provide a custom "
+            "operator. https://www.tensorflow.org/lite/guide/ops_custom")
+        self._log(err_string)
+      # Check TensorFlow op usage error.
+      elif line.startswith(_TF_OPS_HDR):
+        tf_ops = line[len(_TF_OPS_HDR):]
+        err_string = (
+            f"{_AUTHORING_WARNING_HDR}: op '{tf_ops}' require(s) \"Select TF "
+            "Ops\" for model conversion for TensorFlow Lite. "
+            "https://www.tensorflow.org/lite/guide/ops_select")
+        self._log(err_string)
+
+  def _decode_converter_error(self, err):
+    """Parses the given ConverterError which has detailed error information."""
+    custom_ops = []
+    custom_ops_location = []
+    tf_ops = []
+    tf_ops_location = []
+    gpu_not_compatible_ops = []
+    for err in err.errors:
+      # Check custom op usage error.
+      if err.error_code == converter_error_data_pb2.ConverterErrorData.ERROR_NEEDS_CUSTOM_OPS:
+        custom_ops.append(err.operator.name)
+        custom_ops_location.append(err.location)
+      # Check TensorFlow op usage error.
+      elif err.error_code == converter_error_data_pb2.ConverterErrorData.ERROR_NEEDS_FLEX_OPS:
+        tf_ops.append(err.operator.name)
+        tf_ops_location.append(err.location)
+      # Check GPU delegate compatibility error.
+      elif err.error_code == converter_error_data_pb2.ConverterErrorData.ERROR_GPU_NOT_COMPATIBLE:
+        gpu_not_compatible_ops.append(err.operator.name)
+        # Log the first line of ConveterError.errors.error_message only
+        # since the seond line is "Error code: xxxx"
+        self._log(err.error_message.splitlines()[0])
+        self._log(self._get_location_string(err.location) + "\n")
+      else:
+        # Log other errors.
+        self._log(f"{_AUTHORING_ERROR_HDR}: {err.error_message}")
+        self._log(self._get_location_string(err.location) + "\n")
+
+    if custom_ops:
+      custom_ops_str = ", ".join(sorted(custom_ops))
+      err_string = (
+          f"{_AUTHORING_ERROR_HDR}: op '{custom_ops_str}' is(are) not natively "
+          "supported by TensorFlow Lite. You need to provide a custom "
+          "operator. https://www.tensorflow.org/lite/guide/ops_custom")
+      self._log(err_string)
+      self._dump_error_details(custom_ops, custom_ops_location)
+
+    if tf_ops:
+      tf_ops_str = ", ".join(sorted(tf_ops))
+      err_string = (
+          f"{_AUTHORING_WARNING_HDR}: op '{tf_ops_str}' require(s) \"Select TF"
+          " Ops\" for model conversion for TensorFlow Lite. "
+          "https://www.tensorflow.org/lite/guide/ops_select")
+      self._log(err_string)
+      self._dump_error_details(tf_ops, tf_ops_location)
+
+    if gpu_not_compatible_ops:
+      not_compatible_ops_str = ", ".join(sorted(gpu_not_compatible_ops))
+      err_string = (
+          f"{_AUTHORING_WARNING_HDR}: op '{not_compatible_ops_str}' aren't "
+          "compatible with TensorFlow Lite GPU delegate. "
+          "https://www.tensorflow.org/lite/performance/gpu")
+      self._log(err_string)
+
+  def _decode_error(self, err):
+    """Parses the given ConverterError and generates compatibility warnings."""
+    if hasattr(err, "errors"):
+      self._decode_converter_error(err)
+    else:
+      self._decode_error_legacy(err)
+
+    if self._raise_exception and self._log_messages:
+      raise CompatibilityError(f"CompatibilityException at {repr(self._func)}")
+
+  def _log(self, message):
+    """Log and print authoring warning / error message."""
+    self._log_messages.append(message)
+    print(message)
+
+  def get_compatibility_log(self):
+    """Returns list of compatibility log messages.
+
+    WARNING: This method should only be used for unit tests.
+
+    Returns:
+      The list of log messages by the recent compatibility check.
+    Raises:
+      RuntimeError: when the compatibility was NOT checked.
+    """
+    if not self._verified:
+      raise RuntimeError("target compatibility isn't verified yet")
+    return self._log_messages
+
+
+@_tf_export("lite.experimental.authoring.compatible")
+def compatible(target=None, converter_target_spec=None, **kwargs):
+  """Wraps `tf.function` into a callable function with TFLite compatibility checking.
+
+  Example:
+
+  ```python
+  @tf.lite.experimental.authoring.compatible
+  @tf.function(input_signature=[
+      tf.TensorSpec(shape=[None], dtype=tf.float32)
+  ])
+  def f(x):
+      return tf.cosh(x)
+
+  result = f(tf.constant([0.0]))
+  # COMPATIBILITY WARNING: op 'tf.Cosh' require(s) "Select TF Ops" for model
+  # conversion for TensorFlow Lite.
+  # Op: tf.Cosh
+  #   - tensorflow/python/framework/op_def_library.py:748
+  #   - tensorflow/python/ops/gen_math_ops.py:2458
+  #   - <stdin>:6
+  ```
+
+  WARNING: Experimental interface, subject to change.
+
+  Args:
+    target: A `tf.function` to decorate.
+    converter_target_spec : target_spec of TFLite converter parameter.
+    **kwargs: The keyword arguments of the decorator class _Compatible.
+
+  Returns:
+     A callable object of `tf.lite.experimental.authoring._Compatible`.
+  """
+  if target is None:
+    def wrapper(target):
+      return _Compatible(target, converter_target_spec, **kwargs)
+    return wrapper
+  else:
+    return _Compatible(target, converter_target_spec, **kwargs)

tensorflow/lite/python/authoring/authoring_test.py

@@ -0,0 +1,298 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Unit tests for authoring package."""
+# pylint: disable=g-direct-tensorflow-import
+
+import tensorflow as tf
+
+from tensorflow.lite.python.authoring import authoring
+
+
+class TFLiteAuthoringTest(tf.test.TestCase):
+
+  def test_simple_cosh(self):
+    @authoring.compatible
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[None], dtype=tf.float32)
+    ])
+    def f(x):
+      return tf.cosh(x)
+
+    result = f(tf.constant([0.0]))
+    log_messages = f.get_compatibility_log()
+    self.assertEqual(result, tf.constant([1.0]))
+    self.assertIn(
+        "COMPATIBILITY WARNING: op 'tf.Cosh' require(s) \"Select TF Ops\" for "
+        "model conversion for TensorFlow Lite. "
+        "https://www.tensorflow.org/lite/guide/ops_select", log_messages)
+
+    # Check the op location ends with filename of the this test.
+    self.assertIn("authoring_test.py", log_messages[-1])
+
+  def test_simple_cosh_raises_CompatibilityError(self):
+    @authoring.compatible(raise_exception=True)
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[None], dtype=tf.float32)
+    ])
+    def f(x):
+      return tf.cosh(x)
+
+    # Check if the CompatibilityError exception is raised.
+    with self.assertRaises(authoring.CompatibilityError):
+      result = f(tf.constant([0.0]))
+      del result
+    log_messages = f.get_compatibility_log()
+    self.assertIn(
+        "COMPATIBILITY WARNING: op 'tf.Cosh' require(s) \"Select TF Ops\" for "
+        "model conversion for TensorFlow Lite. "
+        "https://www.tensorflow.org/lite/guide/ops_select", log_messages)
+
+  def test_flex_compatibility(self):
+    @authoring.compatible
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[3, 3, 3, 3, 3], dtype=tf.float32)
+    ])
+    def f(inp):
+      tanh = tf.math.tanh(inp)
+      conv3d = tf.nn.conv3d(
+          tanh,
+          tf.ones([3, 3, 3, 3, 3]),
+          strides=[1, 1, 1, 1, 1],
+          padding="SAME")
+      erf = tf.math.erf(conv3d)
+      output = tf.math.tanh(erf)
+      return output
+
+    f(tf.ones(shape=(3, 3, 3, 3, 3), dtype=tf.float32))
+    log_messages = f.get_compatibility_log()
+    self.assertIn(
+        "COMPATIBILITY WARNING: op 'tf.Erf' require(s) \"Select TF Ops\" for "
+        "model conversion for TensorFlow Lite. "
+        "https://www.tensorflow.org/lite/guide/ops_select", log_messages)
+
+  def test_compatibility_error_generic(self):
+    @authoring.compatible
+    @tf.function
+    def f():
+      dataset = tf.data.Dataset.range(3)
+      dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
+      return dataset
+
+    f()
+    log_messages = f.get_compatibility_log()
+    self.assertIn(
+        "COMPATIBILITY ERROR: op 'tf.DummySeedGenerator, tf.RangeDataset,"
+        " tf.ShuffleDatasetV3' is(are) not natively supported by TensorFlow"
+        " Lite. You need to provide a custom operator."
+        " https://www.tensorflow.org/lite/guide/ops_custom",
+        log_messages,
+    )
+
+  def test_compatibility_error_custom(self):
+    target_spec = tf.lite.TargetSpec()
+    target_spec.supported_ops = [
+        tf.lite.OpsSet.TFLITE_BUILTINS,
+        tf.lite.OpsSet.SELECT_TF_OPS,
+    ]
+    @authoring.compatible(converter_target_spec=target_spec)
+    @tf.function
+    def f():
+      dataset = tf.data.Dataset.range(3)
+      dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
+      return dataset
+
+    f()
+    log_messages = f.get_compatibility_log()
+    self.assertIn(
+        "COMPATIBILITY ERROR: op 'tf.DummySeedGenerator, tf.RangeDataset, "
+        "tf.ShuffleDatasetV3' is(are) not natively supported by "
+        "TensorFlow Lite. You need to provide a custom operator. "
+        "https://www.tensorflow.org/lite/guide/ops_custom", log_messages)
+
+  def test_simple_variable(self):
+    external_var = tf.Variable(1.0)
+    @authoring.compatible
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[None], dtype=tf.float32)
+    ])
+    def f(x):
+      return x * external_var
+
+    result = f(tf.constant(2.0, shape=(1)))
+    log_messages = f.get_compatibility_log()
+
+    self.assertEqual(result, tf.constant([2.0]))
+    self.assertEmpty(log_messages)
+
+  def test_class_method(self):
+    class Model(tf.Module):
+
+      @authoring.compatible
+      @tf.function(input_signature=[
+          tf.TensorSpec(shape=[None], dtype=tf.float32)
+      ])
+      def eval(self, x):
+        return tf.cosh(x)
+
+    m = Model()
+    result = m.eval(tf.constant([0.0]))
+    log_messages = m.eval.get_compatibility_log()
+
+    self.assertEqual(result, tf.constant([1.0]))
+    self.assertIn(
+        "COMPATIBILITY WARNING: op 'tf.Cosh' require(s) \"Select TF Ops\" for "
+        "model conversion for TensorFlow Lite. "
+        "https://www.tensorflow.org/lite/guide/ops_select", log_messages)
+
+  def test_decorated_function_type(self):
+    @authoring.compatible
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[None], dtype=tf.float32)
+    ])
+    def func(x):
+      return tf.cos(x)
+
+    result = func(tf.constant([0.0]))
+    self.assertEqual(result, tf.constant([1.0]))
+
+    # Check if the decorator keeps __name__ attribute.
+    self.assertEqual(func.__name__, "func")
+
+    # Check if the decorator works with get_concrete_function method.
+    converter = tf.lite.TFLiteConverter.from_concrete_functions(
+        [func.get_concrete_function()], func)
+    converter.convert()
+
+  def test_decorated_class_method_type(self):
+    class Model(tf.Module):
+
+      @authoring.compatible
+      @tf.function(input_signature=[
+          tf.TensorSpec(shape=[None], dtype=tf.float32)
+      ])
+      def eval(self, x):
+        return tf.cos(x)
+
+    m = Model()
+    result = m.eval(tf.constant([0.0]))
+    self.assertEqual(result, tf.constant([1.0]))
+
+    # Check if the decorator keeps __name__ attribute.
+    self.assertEqual(m.eval.__name__, "eval")
+
+    # Check if the decorator works with get_concrete_function method.
+    converter = tf.lite.TFLiteConverter.from_concrete_functions(
+        [m.eval.get_concrete_function()], m)
+    converter.convert()
+
+  def test_simple_cosh_multiple(self):
+    @authoring.compatible
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[None], dtype=tf.float32)
+    ])
+    def f(x):
+      return tf.cosh(x)
+
+    f(tf.constant([1.0]))
+    f(tf.constant([2.0]))
+    f(tf.constant([3.0]))
+    warning_messages = f.get_compatibility_log()
+
+    # Test if compatiblility checks happens only once.
+    # The number of warning_messages will be 2 by op location detail.
+    self.assertEqual(2, len(warning_messages))
+
+  def test_user_tf_ops_all_filtered(self):
+    target_spec = tf.lite.TargetSpec()
+    target_spec.supported_ops = [
+        tf.lite.OpsSet.TFLITE_BUILTINS,
+        tf.lite.OpsSet.SELECT_TF_OPS,
+    ]
+    target_spec.experimental_select_user_tf_ops = [
+        "RangeDataset", "DummySeedGenerator", "ShuffleDatasetV3"
+    ]
+    @authoring.compatible(converter_target_spec=target_spec)
+    @tf.function
+    def f():
+      dataset = tf.data.Dataset.range(3)
+      dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
+      return dataset
+
+    f()
+    log_messages = f.get_compatibility_log()
+    self.assertEmpty(log_messages)
+
+  def test_user_tf_ops_partial_filtered(self):
+    target_spec = tf.lite.TargetSpec()
+    target_spec.supported_ops = [
+        tf.lite.OpsSet.TFLITE_BUILTINS,
+        tf.lite.OpsSet.SELECT_TF_OPS,
+    ]
+    target_spec.experimental_select_user_tf_ops = [
+        "DummySeedGenerator"
+    ]
+    @authoring.compatible(converter_target_spec=target_spec)
+    @tf.function
+    def f():
+      dataset = tf.data.Dataset.range(3)
+      dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
+      return dataset
+
+    f()
+    log_messages = f.get_compatibility_log()
+    self.assertIn(
+        "COMPATIBILITY ERROR: op 'tf.RangeDataset, tf.ShuffleDatasetV3' is(are) "
+        "not natively supported by TensorFlow Lite. You need to provide a "
+        "custom operator. https://www.tensorflow.org/lite/guide/ops_custom",
+        log_messages)
+
+  def test_allow_custom_ops(self):
+    target_spec = tf.lite.TargetSpec()
+    target_spec.supported_ops = [
+        tf.lite.OpsSet.TFLITE_BUILTINS,
+        tf.lite.OpsSet.SELECT_TF_OPS,
+    ]
+    @authoring.compatible(
+        converter_allow_custom_ops=True, converter_target_spec=target_spec)
+    @tf.function
+    def f():
+      dataset = tf.data.Dataset.range(3)
+      dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
+      return dataset
+
+    f()
+    log_messages = f.get_compatibility_log()
+    self.assertEmpty(log_messages)
+
+  def test_gpu_compatible(self):
+    target_spec = tf.lite.TargetSpec()
+    target_spec.supported_ops = [
+        tf.lite.OpsSet.TFLITE_BUILTINS,
+    ]
+    target_spec.experimental_supported_backends = ["GPU"]
+
+    @authoring.compatible(converter_target_spec=target_spec)
+    @tf.function(
+        input_signature=[tf.TensorSpec(shape=[4, 4], dtype=tf.float32)])
+    def func(x):
+      return tf.cos(x)
+
+    func(tf.ones(shape=(4, 4), dtype=tf.float32))
+    log_messages = func.get_compatibility_log()
+    self.assertEmpty(log_messages)
+
+
+if __name__ == "__main__":
+  tf.test.main()