chore: import upstream snapshot with attribution

tensorflow/lite/python/metrics/BUILD

@@ -0,0 +1,158 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "if_portable", "pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "metrics_wrapper_lib",
+    srcs = if_portable(
+        if_false = ["wrapper/metrics_wrapper_nonportable.cc"],
+        if_true = ["wrapper/metrics_wrapper_portable.cc"],
+    ),
+    hdrs = ["wrapper/metrics_wrapper.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:private"],
+    deps = [
+        "@xla//third_party/python_runtime:headers",
+    ] + if_portable(
+        if_false = [
+            "//learning/brain/google/monitoring:metrics_exporter",
+        ],
+        if_true = [],
+    ),
+)
+
+pybind_extension(
+    name = "_pywrap_tensorflow_lite_metrics_wrapper",
+    srcs = ["wrapper/metrics_wrapper_pybind11.cc"],
+    hdrs = ["wrapper/metrics_wrapper.h"],
+    common_lib_packages = [
+        "litert/python",
+        "tensorflow/lite/python",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_tensorflow_lite_metrics_wrapper.pyi",
+    ],
+    visibility = ["//visibility:public"],
+    wrap_py_init = True,
+    deps = [
+        ":metrics_wrapper_lib",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@com_google_protobuf//:protobuf",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",
+    ],
+)
+
+pytype_strict_library(
+    name = "metrics_wrapper",
+    srcs = ["wrapper/metrics_wrapper.py"],
+    deps = [
+        ":_pywrap_tensorflow_lite_metrics_wrapper",
+        "//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
+        "//tensorflow/compiler/mlir/lite/python:wrap_converter",
+    ],
+)
+
+py_test(
+    name = "metrics_wrapper_test",
+    srcs = ["wrapper/metrics_wrapper_test.py"],
+    strict_deps = True,
+    deps = [
+        ":metrics_wrapper",
+        #internal proto upb dep
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/lite/python:convert",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+pytype_strict_library(
+    name = "metrics_interface",
+    srcs = ["metrics_interface.py"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//tensorflow/lite/tools/pip_package:__subpackages__"],
+)
+
+genrule(
+    name = "metrics_py_gen",
+    srcs = if_portable(
+        if_false = ["metrics_nonportable.py"],
+        if_true = ["metrics_portable.py"],
+    ),
+    outs = ["metrics.py"],
+    cmd = (
+        "cat $(SRCS) > $(OUTS)"
+    ),
+    compatible_with = get_compatible_with_portable(),
+)
+
+pytype_strict_library(
+    name = "metrics",
+    srcs = ["metrics.py"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//tensorflow/lite:__subpackages__"],
+    deps = if_portable(
+        if_false = [
+            "//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
+            ":metrics_wrapper",
+            "//tensorflow/python/eager:monitoring",
+        ],
+        if_true = [],
+    ) + [":metrics_interface"],
+)
+
+py_test(
+    name = "metrics_test",
+    srcs = if_portable(
+        if_false = ["metrics_nonportable_test.py"],
+        if_true = ["metrics_portable_test.py"],
+    ),
+    data = [
+        "//tensorflow/lite/python/testdata/control_flow_v1_saved_model:saved_model.pb",
+    ],
+    main = if_portable(
+        if_false = "metrics_nonportable_test.py",
+        if_true = "metrics_portable_test.py",
+    ),
+    strict_deps = True,
+    tags = ["notap"],  # TODO(b/373657707): Remove once we debug the failure.
+    deps = [
+        ":metrics",
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/lite/python:convert",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:monitoring",
+        "//tensorflow/python/framework:convert_to_constants",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:importer",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:string_ops",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:resource_loader",
+        "//tensorflow/python/saved_model",
+        "//tensorflow/python/trackable:autotrackable",
+    ],
+)

tensorflow/lite/python/metrics/_pywrap_tensorflow_lite_metrics_wrapper.pyi

@@ -0,0 +1,18 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+class MetricsWrapper:
+    def __init__(self, arg0: str) -> None: ...
+    def ExportMetrics(self) -> object: ...

tensorflow/lite/python/metrics/metrics_interface.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+"""Python TFLite metrics helper interface."""
+import abc
+
+
+class TFLiteMetricsInterface(metaclass=abc.ABCMeta):
+  """Abstract class for TFLiteMetrics."""
+
+  @abc.abstractmethod
+  def increase_counter_debugger_creation(self):
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def increase_counter_interpreter_creation(self):
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def increase_counter_converter_attempt(self):
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def increase_counter_converter_success(self):
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def set_converter_param(self, name, value):
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def set_converter_error(self, error_data):
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def set_converter_latency(self, value):
+    raise NotImplementedError

tensorflow/lite/python/metrics/metrics_nonportable.py

@@ -0,0 +1,130 @@
+# 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.
+# ==============================================================================
+"""Python TFLite metrics helper."""
+from typing import Optional, Text
+import uuid
+
+from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
+from tensorflow.lite.python.metrics import metrics_interface
+from tensorflow.lite.python.metrics.wrapper import metrics_wrapper
+from tensorflow.python.eager import monitoring
+
+_counter_debugger_creation = monitoring.Counter(
+    '/tensorflow/lite/quantization_debugger/created',
+    'Counter for the number of debugger created.')
+
+_counter_interpreter_creation = monitoring.Counter(
+    '/tensorflow/lite/interpreter/created',
+    'Counter for number of interpreter created in Python.', 'language')
+
+# The following are conversion metrics. Attempt and success are kept separated
+# instead of using a single metric with a label because the converter may
+# raise exceptions if conversion failed. That may lead to cases when we are
+# unable to capture the conversion attempt. Increasing attempt count at the
+# beginning of conversion process and the success count at the end is more
+# suitable in these cases.
+_counter_conversion_attempt = monitoring.Counter(
+    '/tensorflow/lite/convert/attempt',
+    'Counter for number of conversion attempts.')
+
+_counter_conversion_success = monitoring.Counter(
+    '/tensorflow/lite/convert/success',
+    'Counter for number of successful conversions.')
+
+_gauge_conversion_params = monitoring.StringGauge(
+    '/tensorflow/lite/convert/params',
+    'Gauge for keeping conversion parameters.', 'name')
+
+_gauge_conversion_errors = monitoring.StringGauge(
+    '/tensorflow/lite/convert/errors',
+    'Gauge for collecting conversion errors. The value represents the error '
+    'message.', 'component', 'subcomponent', 'op_name', 'error_code')
+
+_gauge_conversion_latency = monitoring.IntGauge(
+    '/tensorflow/lite/convert/latency', 'Conversion latency in ms.')
+
+
+class TFLiteMetrics(metrics_interface.TFLiteMetricsInterface):
+  """TFLite metrics helper for prod (borg) environment.
+
+  Attributes:
+    model_hash: A string containing the hash of the model binary.
+    model_path: A string containing the path of the model for debugging
+      purposes.
+  """
+
+  def __init__(self,
+               model_hash: Optional[Text] = None,
+               model_path: Optional[Text] = None) -> None:
+    del self  # Temporarily removing self until parameter logic is implemented.
+    if model_hash and not model_path or not model_hash and model_path:
+      raise ValueError('Both model metadata(model_hash, model_path) should be '
+                       'given at the same time.')
+    if model_hash:
+      # TODO(b/180400857): Create stub once the service is implemented.
+      pass
+
+  def increase_counter_debugger_creation(self):
+    _counter_debugger_creation.get_cell().increase_by(1)
+
+  def increase_counter_interpreter_creation(self):
+    _counter_interpreter_creation.get_cell('python').increase_by(1)
+
+  def increase_counter_converter_attempt(self):
+    _counter_conversion_attempt.get_cell().increase_by(1)
+
+  def increase_counter_converter_success(self):
+    _counter_conversion_success.get_cell().increase_by(1)
+
+  def set_converter_param(self, name, value):
+    _gauge_conversion_params.get_cell(name).set(value)
+
+  def set_converter_error(
+      self, error_data: converter_error_data_pb2.ConverterErrorData):
+    error_code_str = converter_error_data_pb2.ConverterErrorData.ErrorCode.Name(
+        error_data.error_code)
+    _gauge_conversion_errors.get_cell(
+        error_data.component,
+        error_data.subcomponent,
+        error_data.operator.name,
+        error_code_str,
+    ).set(error_data.error_message)
+
+  def set_converter_latency(self, value):
+    _gauge_conversion_latency.get_cell().set(value)
+
+
+class TFLiteConverterMetrics(TFLiteMetrics):
+  """Similar to TFLiteMetrics but specialized for converter.
+
+  A unique session id will be created for each new TFLiteConverterMetrics.
+  """
+
+  def __init__(self) -> None:
+    super(TFLiteConverterMetrics, self).__init__()
+    session_id = uuid.uuid4().hex
+    self._metrics_exporter = metrics_wrapper.MetricsWrapper(session_id)
+    self._exported = False
+
+  def __del__(self):
+    if not self._exported:
+      self.export_metrics()
+
+  def set_export_required(self):
+    self._exported = False
+
+  def export_metrics(self):
+    self._metrics_exporter.ExportMetrics()
+    self._exported = True

tensorflow/lite/python/metrics/metrics_nonportable_test.py

@@ -0,0 +1,570 @@
+# 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 Lite Python metrics helper TFLiteMetrics check."""
+import gc
+import os
+import tempfile
+import time
+from unittest import mock
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
+from tensorflow.core.framework import graph_pb2
+from tensorflow.lite.python import lite
+from tensorflow.lite.python.convert import ConverterError
+from tensorflow.lite.python.convert import register_custom_opdefs
+from tensorflow.lite.python.metrics import metrics
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+from tensorflow.python.eager import monitoring
+from tensorflow.python.framework import convert_to_constants
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework.importer import import_graph_def
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import resource_loader
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import saved_model
+from tensorflow.python.trackable import autotrackable
+
+
+class MetricsNonportableTest(test_util.TensorFlowTestCase):
+
+  def test_TFLiteMetrics_creation_no_arg_success(self):
+    metrics.TFLiteMetrics()
+
+  def test_TFLiteMetrics_creation_arg_success(self):
+    metrics.TFLiteMetrics('hash', '/path/to/model')
+
+  def test_TFLiteMetrics_creation_fails_with_only_hash(self):
+    with self.assertRaises(ValueError):
+      metrics.TFLiteMetrics(model_hash='hash')
+
+  def test_TFLiteMetrics_creation_fail2_with_only_model_path(self):
+    with self.assertRaises(ValueError):
+      metrics.TFLiteMetrics(model_path='/path/to/model')
+
+  def test_debugger_creation_counter_increase_multiple_same_topic_success(self):
+    try:
+      stub = metrics.TFLiteMetrics()
+      stub.increase_counter_debugger_creation()
+      self.assertEqual(metrics._counter_debugger_creation.get_cell().value(), 1)
+      stub2 = metrics.TFLiteMetrics()
+      stub2.increase_counter_debugger_creation()
+      self.assertEqual(metrics._counter_debugger_creation.get_cell().value(), 2)
+      del stub
+      gc.collect()
+      stub2.increase_counter_debugger_creation()
+      self.assertEqual(metrics._counter_debugger_creation.get_cell().value(), 3)
+    except:
+      raise Exception('No exception should be raised.')
+
+  def test_interpreter_creation_counter_increase_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.increase_counter_interpreter_creation()
+    self.assertEqual(
+        metrics._counter_interpreter_creation.get_cell('python').value(), 1)
+
+  def test_converter_attempt_counter_increase_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.increase_counter_converter_attempt()
+    self.assertEqual(metrics._counter_conversion_attempt.get_cell().value(), 1)
+
+  def test_converter_success_counter_increase_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.increase_counter_converter_success()
+    self.assertEqual(metrics._counter_conversion_success.get_cell().value(), 1)
+
+  def test_converter_params_set_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.set_converter_param('name', 'value')
+    self.assertEqual(
+        metrics._gauge_conversion_params.get_cell('name').value(), 'value')
+
+  def test_converter_params_multiple_set_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.set_converter_param('name', 'value')
+    stub.set_converter_param('name', 'value1')
+    self.assertEqual(
+        metrics._gauge_conversion_params.get_cell('name').value(), 'value1')
+
+  def test_converter_params_multiple_label_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.set_converter_param('name1', 'value1')
+    stub.set_converter_param('name2', 'value2')
+    self.assertEqual(
+        metrics._gauge_conversion_params.get_cell('name1').value(), 'value1')
+    self.assertEqual(
+        metrics._gauge_conversion_params.get_cell('name2').value(), 'value2')
+
+  def test_converter_params_set_latency(self):
+    stub = metrics.TFLiteMetrics()
+    stub.set_converter_latency(34566)
+    self.assertEqual(metrics._gauge_conversion_latency.get_cell().value(),
+                     34566)
+
+
+class ConverterMetricsTest(test_util.TensorFlowTestCase):
+  """Testing conversion metrics."""
+
+  def _constructGraphDef(self):
+    with ops.Graph().as_default():
+      in_tensor = array_ops.placeholder(
+          shape=[None, 16, 16, 3], dtype=dtypes.float32, name='in_tensor')
+      math_ops.add(in_tensor, in_tensor, name='add')
+      sess = session.Session()
+
+    return (
+        convert_to_constants.convert_variables_to_constants_from_session_graph(
+            sess, sess.graph_def, ['add']))
+
+  def test_conversion_from_constructor_success(self):
+    frozen_graph_def = self._constructGraphDef()
+
+    # Check metrics when conversion successed.
+    converter = lite.TFLiteConverter(frozen_graph_def, None, None,
+                                     [('in_tensor', [2, 16, 16, 3])], ['add'])
+    mock_metrics = mock.create_autospec(
+        metrics.TFLiteConverterMetrics, instance=True)
+    converter._tflite_metrics = mock_metrics
+    tflite_model = converter.convert()
+    self.assertIsNotNone(tflite_model)
+    mock_metrics.assert_has_calls([
+        mock.call.increase_counter_converter_attempt(),
+        mock.call.increase_counter_converter_success(),
+        mock.call.export_metrics(),
+        mock.call.set_converter_param('input_format', '1'),
+        mock.call.set_converter_param('allow_custom_ops', 'False'),
+        mock.call.set_converter_param('api_version', '1'),
+    ], any_order=True)  # pyformat: disable
+
+  def test_conversion_from_constructor_fail(self):
+    frozen_graph_def = self._constructGraphDef()
+
+    # Check metrics when conversion failed.
+    converter = lite.TFLiteConverter(frozen_graph_def, None, None,
+                                     [('wrong_tensor', [2, 16, 16, 3])],
+                                     ['add'])
+    mock_metrics = mock.create_autospec(
+        metrics.TFLiteConverterMetrics, instance=True)
+    converter._tflite_metrics = mock_metrics
+    with self.assertRaises(ConverterError):
+      converter.convert()
+    mock_metrics.assert_has_calls([
+        mock.call.increase_counter_converter_attempt(),
+        mock.call.set_converter_param('output_format', '2'),
+        mock.call.set_converter_param('select_user_tf_ops', 'None'),
+        mock.call.set_converter_param('post_training_quantize', 'False'),
+    ], any_order=True)  # pyformat: disable
+    mock_metrics.increase_counter_converter_success.assert_not_called()
+
+  def _getIntegerQuantizeModel(self):
+    np.random.seed(0)
+
+    root = autotrackable.AutoTrackable()
+
+    @tf.function(
+        input_signature=[tf.TensorSpec(shape=[1, 5, 5, 3], dtype=tf.float32)])
+    def func(inp):
+      conv = tf.nn.conv2d(
+          inp, tf.ones([3, 3, 3, 16]), strides=[1, 1, 1, 1], padding='SAME')
+      output = tf.nn.relu(conv, name='output')
+      return output
+
+    def calibration_gen():
+      for _ in range(5):
+        yield [np.random.uniform(-1, 1, size=(1, 5, 5, 3)).astype(np.float32)]
+
+    root.f = func
+    to_save = root.f.get_concrete_function()
+    return (root, to_save, calibration_gen)
+
+  def test_conversion_from_frozen_graph_v2(self):
+    model, func, calibration_gen = self._getIntegerQuantizeModel()
+
+    quantized_converter = lite.TFLiteConverterV2.from_concrete_functions([func],
+                                                                         model)
+    mock_metrics = mock.create_autospec(
+        metrics.TFLiteConverterMetrics, instance=True)
+    quantized_converter._tflite_metrics = mock_metrics
+    quantized_converter.optimizations = [lite.Optimize.DEFAULT]
+    quantized_converter.representative_dataset = calibration_gen
+    quantized_tflite_model = quantized_converter.convert()
+    self.assertIsNotNone(quantized_tflite_model)
+    mock_metrics.assert_has_calls([
+        mock.call.increase_counter_converter_attempt(),
+        mock.call.increase_counter_converter_success(),
+        mock.call.set_converter_param(
+            'optimization_post_training_integer_quantize', 'True'),
+        mock.call.set_converter_param('inference_type', 'tf.int8'),
+        mock.call.set_converter_param('select_user_tf_ops', 'None'),
+        mock.call.set_converter_param('activations_type', 'tf.int8'),
+    ], any_order=True)  # pyformat: disable
+
+  def test_conversion_from_keras_v2(self):
+    x = [-1, 0, 1, 2, 3, 4]
+    y = [-3, -1, 1, 3, 5, 7]
+    model = tf.keras.models.Sequential(
+        [tf.keras.layers.Dense(units=1, input_shape=[1])])
+    model.compile(optimizer='sgd', loss='mean_squared_error')
+    model.fit(x, y, epochs=1)
+    converter = lite.TFLiteConverterV2.from_keras_model(model)
+    mock_metrics = mock.create_autospec(
+        metrics.TFLiteConverterMetrics, instance=True)
+    converter._tflite_metrics = mock_metrics
+    converter.convert()
+    mock_metrics.assert_has_calls([
+        mock.call.increase_counter_converter_attempt(),
+        mock.call.increase_counter_converter_success(),
+        mock.call.export_metrics(),
+        mock.call.set_converter_param('inference_type', 'tf.float32'),
+        mock.call.set_converter_param('target_ops', 'TFLITE_BUILTINS'),
+        mock.call.set_converter_param('optimization_default', 'False'),
+    ], any_order=True)  # pyformat: disable
+
+  def _createV1SavedModel(self, shape):
+    """Create a simple SavedModel."""
+    saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
+    with tf.Graph().as_default():
+      with tf.compat.v1.Session() as sess:
+        in_tensor_1 = tf.compat.v1.placeholder(
+            shape=shape, dtype=tf.float32, name='inputB')
+        in_tensor_2 = tf.compat.v1.placeholder(
+            shape=shape, dtype=tf.float32, name='inputA')
+        variable_node = tf.Variable(1.0, name='variable_node')
+        out_tensor = in_tensor_1 + in_tensor_2 * variable_node
+        inputs = {'x': in_tensor_1, 'y': in_tensor_2}
+        outputs = {'z': out_tensor}
+        sess.run(tf.compat.v1.variables_initializer([variable_node]))
+        saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
+    return saved_model_dir
+
+  def test_conversion_from_saved_model(self):
+    saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
+    converter = lite.TFLiteSavedModelConverter(saved_model_dir, set(['serve']),
+                                               ['serving_default'])
+    converter.experimental_new_converter = True
+    mock_metrics = mock.create_autospec(
+        metrics.TFLiteConverterMetrics, instance=True)
+    converter._tflite_metrics = mock_metrics
+    time.process_time = mock.Mock(side_effect=np.arange(1, 1000, 2).tolist())
+    converter.convert()
+    mock_metrics.assert_has_calls([
+        mock.call.increase_counter_converter_attempt(),
+        mock.call.increase_counter_converter_success(),
+        mock.call.set_converter_latency(2000),
+        mock.call.export_metrics(),
+    ], any_order=True)  # pyformat: disable
+
+  def disable_converter_counter_metrics(self, tflite_metrics):
+
+    def empty_func():
+      pass
+
+    tflite_metrics.increase_counter_converter_attempt = empty_func
+    tflite_metrics.increase_counter_converter_success = empty_func
+
+  def test_export_at_conversion_done(self):
+    saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
+
+    converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
+    tflite_metrics = converter._tflite_metrics
+    mock_exporter = mock.MagicMock()
+    tflite_metrics._metrics_exporter = mock_exporter
+    self.disable_converter_counter_metrics(tflite_metrics)
+    mock_exporter.ExportMetrics.assert_not_called()
+    converter.convert()
+    mock_exporter.ExportMetrics.assert_called_once()
+    tflite_metrics.__del__()
+    mock_exporter.ExportMetrics.assert_called_once()
+
+  def test_export_at_exit(self):
+    saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
+    converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
+    tflite_metrics = converter._tflite_metrics
+    mock_exporter = mock.MagicMock()
+    tflite_metrics._metrics_exporter = mock_exporter
+    self.disable_converter_counter_metrics(tflite_metrics)
+    mock_exporter.ExportMetrics.assert_not_called()
+    tflite_metrics.__del__()
+    mock_exporter.ExportMetrics.assert_called_once()
+
+
+def mock_ngrams(data, width, axis=-1, string_separator=' ', name=None):
+  """This mock Ngrams lack the width attr, causing conversion to fail."""
+
+  experimental_implements = [
+      'name: "tftext:Ngrams"',
+      'attr { key: "axis" value { i: %d } }' % axis,
+      'attr { key: "reduction_type" value { s: "STRING_JOIN" } }',
+      'attr { key: "string_separator" value { s: "%s" } }' % string_separator,
+  ]
+  experimental_implements = ' '.join(experimental_implements)
+
+  @tf.function(experimental_implements=experimental_implements)
+  def func(data):
+    with ops.name_scope(name, 'NGrams', [data, width]):
+      data = ragged_tensor.convert_to_tensor_or_ragged_tensor(data, name='data')
+      slices = []
+      for start in range(width):
+        stop = None if start - width + 1 == 0 else start - width + 1
+        if axis >= 0:
+          idx = [slice(None)] * axis + [slice(start, stop)]
+        else:
+          idx = [Ellipsis, slice(start, stop)] + [slice(None)] * (-axis - 1)
+        slices.append(data[idx])
+
+      # Stack the slices.
+      stack_axis = axis + 1 if axis >= 0 else axis
+      windowed_data = array_ops_stack.stack(slices, stack_axis)
+
+      return string_ops.reduce_join(
+          windowed_data, axis=axis, separator=string_separator)
+
+  return func(data)
+
+
+class ConverterErrorMetricTest(test_util.TensorFlowTestCase,
+                               parameterized.TestCase):
+  """Testing conversion error metric."""
+
+  def setUp(self):
+    super(ConverterErrorMetricTest, self).setUp()
+
+    # Mock metrics instance except errors so other test cases are not affected.
+    mock_attempt = mock.create_autospec(monitoring.Counter, instance=True)
+    self._counter_conversion_attempt = metrics._counter_conversion_attempt
+    metrics._counter_conversion_attempt = mock_attempt
+
+    mock_success = mock.create_autospec(monitoring.Counter, instance=True)
+    self._counter_conversion_success = metrics._counter_conversion_success
+    metrics._counter_conversion_success = mock_success
+
+    mock_params = mock.create_autospec(monitoring.StringGauge, instance=True)
+    self._gauge_conversion_params = metrics._gauge_conversion_params
+    metrics._gauge_conversion_params = mock_params
+
+  def tearDown(self):
+    super(ConverterErrorMetricTest, self).tearDown()
+    # # Restore metrics instances.
+    metrics._counter_conversion_attempt = self._counter_conversion_attempt
+    metrics._counter_conversion_success = self._counter_conversion_success
+    metrics._gauge_conversion_params = self._gauge_conversion_params
+
+  def convert_and_check_location_info(self,
+                                      converter,
+                                      expected_type,
+                                      expected_sources=None):
+    # The custom attribute of ConverterError can't be accessed with
+    # assertRaises so use try-catch block instead.
+    try:
+      tflite_model = converter.convert()
+      self.assertIsNone(tflite_model)
+    except ConverterError as converter_error:
+      # pylint: disable=g-assert-in-except
+      self.assertLen(converter_error.errors, 1)
+      location = converter_error.errors[0].location
+      self.assertEqual(location.type, expected_type)
+
+      if expected_sources:
+        debug_string = str(location)
+        for source in expected_sources:
+          self.assertIn(source, debug_string)
+      # pylint: enable=g-assert-in-except
+
+  def test_failure_at_PrepareCompositeFunctionsPass(self):
+    if context.is_tfrt_enabled():
+      self.skipTest('This test crashed with TFRT.')
+
+    class NgramsLayer(tf.keras.layers.Layer):
+
+      def call(self, input_tensor, **kwargs):
+        return mock_ngrams(input_tensor, width=2, axis=-1, string_separator=' ')
+
+    # Registers a fake WhitespaceTokenizeWithOffsets so the TFText fusing logic
+    # is enable in MLIR side.
+    custom_opdefs_str = (
+        'name: \'WhitespaceTokenizeWithOffsets\' input_arg: {name: \'Input1\' '
+        'type: DT_FLOAT} input_arg: {name: \'Input2\' type: DT_FLOAT} '
+        'output_arg: {name: \'Output\' type: DT_FLOAT}')
+    register_custom_opdefs([custom_opdefs_str])
+
+    model = tf.keras.models.Sequential([NgramsLayer()])
+    model.predict(tf.constant(['test']))
+    converter = lite.TFLiteConverterV2.from_keras_model(model)
+    converter.allow_custom_ops = True
+    self.convert_and_check_location_info(
+        converter, converter_error_data_pb2.ConverterErrorData.UNKNOWNLOC)
+    exported_error = metrics._gauge_conversion_errors.get_cell(
+        'CONVERT_TF_TO_TFLITE_MODEL',
+        'PrepareCompositeFunctionsPass',
+        'tf.Const',
+        'UNKNOWN',
+    ).value()
+    self.assertEqual(exported_error,
+                     "\'width\' attribute is not set or not an integer")
+
+  def test_need_flex_ops(self):
+
+    def create_graph_with_custom_add(opname='CustomAdd'):
+      custom_opdefs_str = (
+          'name: \'' + opname +
+          '\' input_arg: {name: \'Input1\' type: DT_FLOAT} '
+          'input_arg: {name: \'Input2\' type: DT_FLOAT} output_arg: {name: '
+          '\'Output\' type: DT_FLOAT}')
+
+      # Create a graph that has one add op.
+      new_graph = graph_pb2.GraphDef()
+      with ops.Graph().as_default():
+        with session.Session() as sess:
+          in_tensor = array_ops.placeholder(
+              shape=[1, 16, 16, 3], dtype=dtypes.float32, name='input')
+          out_tensor = in_tensor + in_tensor
+          inputs = {'x': in_tensor}
+          outputs = {'z': out_tensor}
+
+          new_graph.CopyFrom(sess.graph_def)
+
+      # Rename Add op name to opname.
+      for node in new_graph.node:
+        if node.op.startswith('Add'):
+          node.op = opname
+          del node.attr['T']
+
+      # Register custom op defs to import modified graph def.
+      register_custom_opdefs([custom_opdefs_str])
+
+      return (new_graph, inputs, outputs)
+
+    new_graph, inputs, outputs = create_graph_with_custom_add()
+
+    # Import to load the custom opdef.
+    saved_model_dir = os.path.join(self.get_temp_dir(), 'model')
+    with ops.Graph().as_default():
+      with session.Session() as sess:
+        import_graph_def(new_graph, name='')
+        saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
+
+    converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
+    self.convert_and_check_location_info(
+        converter,
+        converter_error_data_pb2.ConverterErrorData.NAMELOC,
+        expected_sources='add')
+    exported_error = metrics._gauge_conversion_errors.get_cell(
+        'CONVERT_TF_TO_TFLITE_MODEL', 'CONVERT_SAVED_MODEL', 'tf.CustomAdd',
+        'ERROR_NEEDS_CUSTOM_OPS').value()
+    self.assertIn(
+        "'tf.CustomAdd' op is neither a custom op nor a flex op\n",
+        exported_error,
+    )
+    self.assertIn('Error code: ERROR_NEEDS_CUSTOM_OPS', exported_error)
+
+  def test_unsupported_control_flow_v1(self):
+    filename = resource_loader.get_path_to_datafile(
+        '../testdata/control_flow_v1_saved_model')
+    converter = lite.TFLiteConverterV2.from_saved_model(filename)
+    self.convert_and_check_location_info(
+        converter, converter_error_data_pb2.ConverterErrorData.UNKNOWNLOC)
+    exported_error = metrics._gauge_conversion_errors.get_cell(
+        'CONVERT_TF_TO_TFLITE_MODEL', 'CONVERT_SAVED_MODEL', '',
+        'ERROR_UNSUPPORTED_CONTROL_FLOW_V1').value()
+    self.assertEqual(
+        exported_error,
+        'Merge only has 4 inputs, while only merge nodes with two inputs '
+        'supported.\n\tFailed to functionalize Control Flow V1 ops. Consider '
+        'using Control Flow V2 ops instead. See https://www.tensorflow.org/'
+        'api_docs/python/tf/compat/v1/enable_control_flow_v2.')
+
+  def test_location_from_concrete_functions(self):
+
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[None, None, 2, 3, 3], dtype=tf.complex64),
+        tf.TensorSpec(shape=[None, None, 1, 3, 3], dtype=tf.complex64),
+    ])
+    def model(a, b):
+      return tf.add(a, b, name='add')
+
+    converter = lite.TFLiteConverterV2.from_concrete_functions(
+        [model.get_concrete_function()], model)
+    self.convert_and_check_location_info(
+        converter,
+        converter_error_data_pb2.ConverterErrorData.CALLSITELOC,
+        expected_sources=[
+            'tensorflow/lite/python/metrics/metrics_nonportable_test.py',
+        ])
+
+  def test_location_from_saved_model(self):
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+
+      class Adder(tf.Module):
+
+        @tf.function(input_signature=[
+            tf.TensorSpec(shape=[None, None, 2, 3, 3], dtype=tf.complex64),
+            tf.TensorSpec(shape=[None, None, 1, 3, 3], dtype=tf.complex64),
+        ])
+        def serving_default(self, a, b):
+          return tf.add(a, b, name='add')
+
+      tf.saved_model.save(
+          Adder(),
+          tmp_dir,
+          options=tf.saved_model.SaveOptions(save_debug_info=True))
+
+      converter = lite.TFLiteConverterV2.from_saved_model(tmp_dir)
+      self.convert_and_check_location_info(
+          converter,
+          converter_error_data_pb2.ConverterErrorData.CALLSITELOC,
+          expected_sources=[
+              'tensorflow/lite/python/metrics/metrics_nonportable_test.py',
+          ])
+
+  @parameterized.named_parameters(
+      ('_WithoutLoweringToSavedModel', False, None),
+      ('_WithLoweringToSavedModel', True,
+       'tensorflow/lite/python/metrics/metrics_nonportable_test.py'))
+  def test_location_from_keras_model(self, lower_to_saved_model,
+                                     expected_source):
+    input_tensor1 = tf.keras.layers.Input(
+        shape=[None, None, 2, 3, 3], dtype=tf.complex64)
+    input_tensor2 = tf.keras.layers.Input(
+        shape=[None, None, 2, 3, 3], dtype=tf.complex64)
+    output = tf.keras.layers.Add()([input_tensor1, input_tensor2])
+    model = tf.keras.Model(
+        inputs=[input_tensor1, input_tensor2], outputs=output)
+    model.compile(
+        optimizer='adam',
+        loss='sparse_categorical_crossentropy',
+        metrics=['accuracy'])
+
+    converter = lite.TFLiteConverterV2.from_keras_model(model)
+    converter.experimental_lower_to_saved_model = lower_to_saved_model
+    # The location does not contain callsite to the current file.
+    self.convert_and_check_location_info(
+        converter,
+        converter_error_data_pb2.ConverterErrorData.CALLSITELOC,
+        expected_sources=[expected_source] if expected_source else None)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/lite/python/metrics/metrics_portable.py

@@ -0,0 +1,70 @@
+# 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.
+# ==============================================================================
+"""Python TFLite metrics helper."""
+import os
+from typing import Optional, Text
+
+# pylint: disable=g-import-not-at-top
+if not os.path.splitext(__file__)[0].endswith(
+    os.path.join('tflite_runtime', 'metrics_portable')):
+  # This file is part of tensorflow package.
+  from tensorflow.lite.python.metrics import metrics_interface  # type: ignore
+else:
+  # This file is part of tflite_runtime package.
+  from tflite_runtime import metrics_interface  # type: ignore
+# pylint: enable=g-import-not-at-top
+
+
+class TFLiteMetrics(metrics_interface.TFLiteMetricsInterface):
+  """TFLite metrics helper."""
+
+  def __init__(self,
+               model_hash: Optional[Text] = None,
+               model_path: Optional[Text] = None) -> None:
+    pass
+
+  def increase_counter_debugger_creation(self):
+    pass
+
+  def increase_counter_interpreter_creation(self):
+    pass
+
+  def increase_counter_converter_attempt(self):
+    pass
+
+  def increase_counter_converter_success(self):
+    pass
+
+  def set_converter_param(self, name, value):
+    pass
+
+  def set_converter_error(self, error_data):
+    pass
+
+  def set_converter_latency(self, value):
+    pass
+
+
+class TFLiteConverterMetrics(TFLiteMetrics):
+  """Similar to TFLiteMetrics but specialized for converter."""
+
+  def __del__(self):
+    pass
+
+  def set_export_required(self):
+    pass
+
+  def export_metrics(self):
+    pass

tensorflow/lite/python/metrics/metrics_portable_test.py

@@ -0,0 +1,48 @@
+# 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 Lite Python metrics helpr TFLiteMetrics check."""
+from tensorflow.lite.python.metrics import metrics
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class MetricsPortableTest(test_util.TensorFlowTestCase):
+
+  def test_TFLiteMetrics_creation_success(self):
+    metrics.TFLiteMetrics()
+
+  def test_debugger_creation_counter_increase_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.increase_counter_debugger_creation()
+
+  def test_interpreter_creation_counter_increase_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.increase_counter_interpreter_creation()
+
+  def test_converter_attempt_counter_increase_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.increase_counter_converter_attempt()
+
+  def test_converter_success_counter_increase_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.increase_counter_converter_success()
+
+  def test_converter_params_set_success(self):
+    stub = metrics.TFLiteMetrics()
+    stub.set_converter_param('name', 'value')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h

@@ -0,0 +1,58 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_PYTHON_METRICS_WRAPPER_METRICS_WRAPPER_H_
+#define TENSORFLOW_LITE_PYTHON_METRICS_WRAPPER_METRICS_WRAPPER_H_
+
+#include <memory>
+#include <string>
+
+// Place `<locale>` before <Python.h> to avoid build failures in macOS.
+#include <locale>
+
+// The empty line above is on purpose as otherwise clang-format will
+// automatically move <Python.h> before <locale>.
+#include <Python.h>
+
+// We forward declare TFLite classes here to avoid exposing them to SWIG.
+namespace tensorflow {
+namespace monitoring {
+class MetricsExporter;
+}  // namespace monitoring
+}  // namespace tensorflow
+
+namespace tflite {
+namespace metrics_wrapper {
+
+class MetricsWrapper {
+ public:
+  using MetricsExporter = tensorflow::monitoring::MetricsExporter;
+  // SWIG caller takes ownership of pointer.
+  static MetricsWrapper* CreateMetricsWrapper(const std::string& session_id);
+
+  ~MetricsWrapper();
+
+  // Export metrics with Streamz.
+  PyObject* ExportMetrics();
+
+ private:
+  explicit MetricsWrapper(std::unique_ptr<MetricsExporter> exporter);
+
+  const std::unique_ptr<MetricsExporter> exporter_;
+};
+
+}  // namespace metrics_wrapper
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_PYTHON_METRICS_WRAPPER_METRICS_WRAPPER_H_

tensorflow/lite/python/metrics/wrapper/metrics_wrapper.py

@@ -0,0 +1,34 @@
+# 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.
+# ==============================================================================
+"""Stub to make pywrap metrics wrapper accessible."""
+
+from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
+from tensorflow.compiler.mlir.lite.python import wrap_converter
+from tensorflow.lite.python.metrics._pywrap_tensorflow_lite_metrics_wrapper import MetricsWrapper  # pylint: disable=unused-import
+
+
+def retrieve_collected_errors():
+  """Returns and clears the list of collected errors in ErrorCollector.
+
+  The RetrieveCollectedErrors function in C++ returns a list of serialized proto
+  messages. This function will convert them to ConverterErrorData instances.
+
+  Returns:
+    A list of ConverterErrorData.
+  """
+  serialized_message_list = wrap_converter.wrapped_retrieve_collected_errors()
+  return list(
+      map(converter_error_data_pb2.ConverterErrorData.FromString,
+          serialized_message_list))

tensorflow/lite/python/metrics/wrapper/metrics_wrapper_nonportable.cc

@@ -0,0 +1,63 @@
+/* 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.
+==============================================================================*/
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "learning/brain/google/monitoring/metrics_exporter.h"
+#include "tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h"
+
+namespace tflite {
+namespace metrics_wrapper {
+
+MetricsWrapper::MetricsWrapper(std::unique_ptr<MetricsExporter> exporter)
+    : exporter_(std::move(exporter)) {}
+
+MetricsWrapper::~MetricsWrapper() {}
+
+MetricsWrapper* MetricsWrapper::CreateMetricsWrapper(
+    const std::string& session_id) {
+  MetricsExporter::Options options;
+  options.export_at_exit = false;
+  {
+    // Add session_id to root label list.
+    std::vector<streamz::Entity::Label> root_labels;
+    streamz::Entity::Label session_id_label;
+    session_id_label.set_key("session_id");
+    session_id_label.set_string_value(session_id);
+    root_labels.push_back(session_id_label);
+    options.entity_labels = root_labels;
+  }
+
+  std::unique_ptr<MetricsExporter> exporter(new MetricsExporter(options));
+
+  MetricsWrapper* wrapper = new MetricsWrapper(std::move(exporter));
+  return wrapper;
+}
+
+PyObject* MetricsWrapper::ExportMetrics() {
+  if (!exporter_) {
+    PyErr_SetString(PyExc_ValueError, "MetricsExporter was not initialized.");
+    return nullptr;
+  }
+
+  exporter_->ExportMetrics();
+
+  Py_RETURN_NONE;
+}
+
+}  // namespace metrics_wrapper
+}  // namespace tflite

tensorflow/lite/python/metrics/wrapper/metrics_wrapper_portable.cc

@@ -0,0 +1,60 @@
+/* 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.
+==============================================================================*/
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h"
+
+namespace tensorflow {
+namespace monitoring {
+class MetricsExporter {
+ public:
+  MetricsExporter() {}
+  void ExportMetrics() {}
+};
+}  // namespace monitoring
+}  // namespace tensorflow
+
+namespace tflite {
+namespace metrics_wrapper {
+
+MetricsWrapper::MetricsWrapper(std::unique_ptr<MetricsExporter> exporter)
+    : exporter_(std::move(exporter)) {}
+
+MetricsWrapper::~MetricsWrapper() {}
+
+MetricsWrapper* MetricsWrapper::CreateMetricsWrapper(
+    const std::string& session_id) {
+  std::unique_ptr<MetricsExporter> exporter(new MetricsExporter());
+
+  MetricsWrapper* wrapper = new MetricsWrapper(std::move(exporter));
+  return wrapper;
+}
+
+PyObject* MetricsWrapper::ExportMetrics() {
+  if (!exporter_) {
+    PyErr_SetString(PyExc_ValueError, "MetricsExporter was not initialized.");
+    return nullptr;
+  }
+
+  exporter_->ExportMetrics();
+
+  Py_RETURN_NONE;
+}
+
+
+}  // namespace metrics_wrapper
+}  // namespace tflite

tensorflow/lite/python/metrics/wrapper/metrics_wrapper_pybind11.cc

@@ -0,0 +1,44 @@
+/* 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.
+==============================================================================*/
+
+#include <string>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+
+namespace py = pybind11;
+using tflite::metrics_wrapper::MetricsWrapper;
+
+PYBIND11_MODULE(_pywrap_tensorflow_lite_metrics_wrapper, m) {
+  m.doc() = R"pbdoc(
+    _pywrap_tensorflow_lite_metrics_wrapper
+    -----
+  )pbdoc";
+
+  py::class_<MetricsWrapper>(m, "MetricsWrapper")
+      .def(py::init([](const std::string& session_id) {
+        auto* wrapper = MetricsWrapper::CreateMetricsWrapper(session_id);
+        if (!wrapper) {
+          throw std::invalid_argument("Failed to created MetricsWrapper");
+        }
+        return wrapper;
+      }))
+      .def("ExportMetrics", [](MetricsWrapper& self) {
+        return tensorflow::PyoOrThrow(self.ExportMetrics());
+      });
+}

tensorflow/lite/python/metrics/wrapper/metrics_wrapper_test.py

@@ -0,0 +1,50 @@
+# 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.
+# ==============================================================================
+"""TFLite metrics_wrapper module test cases."""
+
+import tensorflow as tf
+
+from tensorflow.lite.python import lite
+from tensorflow.lite.python.convert import ConverterError
+from tensorflow.lite.python.metrics.wrapper import metrics_wrapper
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class MetricsWrapperTest(test_util.TensorFlowTestCase):
+
+  def test_basic_retrieve_collected_errors_empty(self):
+    errors = metrics_wrapper.retrieve_collected_errors()
+    self.assertEmpty(errors)
+
+  def test_basic_retrieve_collected_errors_not_empty(self):
+
+    @tf.function(
+        input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
+    def func(x):
+      return tf.cosh(x)
+
+    converter = lite.TFLiteConverterV2.from_concrete_functions(
+        [func.get_concrete_function()], func)
+    try:
+      converter.convert()
+    except ConverterError as err:
+      # retrieve_collected_errors is already captured in err.errors
+      captured_errors = err.errors
+    self.assertNotEmpty(captured_errors)
+
+
+if __name__ == "__main__":
+  test.main()