chore: import upstream snapshot with attribution

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()