chore: import upstream snapshot with attribution

tensorflow/lite/python/interpreter_wrapper/BUILD

@@ -0,0 +1,112 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "numpy",
+    srcs = ["numpy.cc"],
+    hdrs = ["numpy.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/core/c:common",
+        "//third_party/py/numpy:headers",
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+cc_library(
+    name = "interpreter_wrapper_lib",
+    srcs = ["interpreter_wrapper.cc"],
+    hdrs = [
+        "interpreter_wrapper.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":numpy",
+        ":python_error_reporter",
+        ":python_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:shared_library",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
+        "//tensorflow/lite/kernels:reference_ops",
+        "//tensorflow/lite/kernels/internal:compatibility",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+cc_library(
+    name = "python_error_reporter",
+    srcs = ["python_error_reporter.cc"],
+    hdrs = ["python_error_reporter.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/lite:stateful_error_reporter",
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+cc_library(
+    name = "python_utils",
+    srcs = ["python_utils.cc"],
+    hdrs = ["python_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+config_setting(
+    name = "tflite_pip_with_flex",
+    define_values = {
+        "tflite_pip_with_flex": "true",
+    },
+)
+
+pybind_extension(
+    name = "_pywrap_tensorflow_interpreter_wrapper",
+    srcs = [
+        "interpreter_wrapper_pybind11.cc",
+    ],
+    hdrs = ["interpreter_wrapper.h"],
+    additional_stubgen_deps = [
+        "//third_party/py/numpy:numpy",
+    ],
+    common_lib_packages = [
+        "litert/python",
+        "tensorflow/lite/python",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    enable_stub_generation = True,
+    link_in_framework = True,
+    pytype_srcs = [
+        "_pywrap_tensorflow_interpreter_wrapper.pyi",
+    ],
+    wrap_py_init = True,
+    deps = [
+        ":interpreter_wrapper_lib",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework_stable",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",
+    ] + select({
+        ":tflite_pip_with_flex": ["//tensorflow/lite/delegates/flex:delegate"],
+        "//conditions:default": [],
+    }),
+)

tensorflow/lite/python/interpreter_wrapper/_pywrap_tensorflow_interpreter_wrapper.pyi

@@ -0,0 +1,47 @@
+# 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 InterpreterWrapper:
+    def __init__(self, *args, **kwargs) -> None: ...
+    def AllocateTensors(self, subgraph_index: int = ...) -> object: ...
+    def GetSignatureDefs(self) -> object: ...
+    def GetSubgraphIndexFromSignature(self, arg0: str) -> object: ...
+    def GetTensor(self, tensor_index: int, subgraph_index: int = ...) -> object: ...
+    def InputIndices(self) -> object: ...
+    def Invoke(self, subgraph_index: int = ...) -> object: ...
+    def ModifyGraphWithDelegate(self, arg0: int) -> object: ...
+    def NodeInputs(self, arg0: int) -> object: ...
+    def NodeName(self, arg0: int) -> str: ...
+    def NodeOutputs(self, arg0: int) -> object: ...
+    def NumNodes(self) -> int: ...
+    def NumSubgraphs(self) -> int: ...
+    def NumTensors(self, arg0: int) -> int: ...
+    def OutputIndices(self) -> object: ...
+    def ResetVariableTensors(self) -> object: ...
+    def ResizeInputTensor(self, i: int, value: object, strict: bool, subgraph_index: int = ...) -> object: ...
+    def SetNumThreads(self, arg0: int) -> object: ...
+    def SetTensor(self, i: int, value: object, subgraph_index: int = ...) -> object: ...
+    def TensorName(self, arg0: int, arg1: int) -> str: ...
+    def TensorQuantization(self, arg0: int, arg1: int) -> object: ...
+    def TensorQuantizationParameters(self, arg0: int, arg1: int) -> object: ...
+    def TensorSize(self, arg0: int, arg1: int) -> object: ...
+    def TensorSizeSignature(self, arg0: int, arg1: int) -> object: ...
+    def TensorSparsityParameters(self, arg0: int, arg1: int) -> object: ...
+    def TensorType(self, arg0: int, arg1: int) -> object: ...
+    def interpreter(self) -> int: ...
+    def tensor(self, base_object: object, tensor_index: int, subgraph_index: int = ...) -> object: ...
+
+def CreateWrapperFromBuffer(*args, **kwargs): ...
+def CreateWrapperFromFile(*args, **kwargs): ...

tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h

@@ -0,0 +1,169 @@
+/* Copyright 2018 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_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_
+#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <vector>
+
+// 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>
+
+#include "tensorflow/lite/core/interpreter.h"
+
+struct TfLiteDelegate;
+
+// We forward declare TFLite classes here to avoid exposing them to SWIG.
+namespace tflite {
+class MutableOpResolver;
+
+namespace impl {
+class FlatBufferModel;
+}
+
+namespace interpreter_wrapper {
+
+class PythonErrorReporter;
+
+class InterpreterWrapper {
+ public:
+  using Model = impl::FlatBufferModel;
+
+  // SWIG caller takes ownership of pointer.
+  static InterpreterWrapper* CreateWrapperCPPFromFile(
+      const char* model_path, int op_resolver_id,
+      const std::vector<std::string>& registerers, std::string* error_msg,
+      bool preserve_all_tensors, bool disable_delegate_clustering);
+  static InterpreterWrapper* CreateWrapperCPPFromFile(
+      const char* model_path, int op_resolver_id,
+      const std::vector<std::string>& registerers_by_name,
+      const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
+      std::string* error_msg, bool preserve_all_tensors,
+      bool disable_delegate_clustering, int num_threads,
+      bool default_delegate_latest_features,
+      bool compress_quantization_zero_points,
+      bool disable_delegate_node_fusion = false,
+      bool force_delegate_node_profiling = false);
+
+  // SWIG caller takes ownership of pointer.
+  static InterpreterWrapper* CreateWrapperCPPFromBuffer(
+      PyObject* data, int op_resolver_id,
+      const std::vector<std::string>& registerers, std::string* error_msg,
+      bool preserve_all_tensors, bool disable_delegate_clustering);
+  static InterpreterWrapper* CreateWrapperCPPFromBuffer(
+      PyObject* data, int op_resolver_id,
+      const std::vector<std::string>& registerers_by_name,
+      const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
+      std::string* error_msg, bool preserve_all_tensors,
+      bool disable_delegate_clustering, int num_threads,
+      bool default_delegate_latest_features,
+      bool compress_quantization_zero_points,
+      bool disable_delegate_node_fusion = false,
+      bool force_delegate_node_profiling = false);
+
+  ~InterpreterWrapper();
+  PyObject* AllocateTensors(int subgraph_index);
+  PyObject* Invoke(int subgraph_index);
+
+  PyObject* InputIndices() const;
+  PyObject* OutputIndices() const;
+  PyObject* ResizeInputTensor(int i, PyObject* value, bool strict,
+                              int subgraph_index);
+
+  int NumTensors(int subgraph_index) const;
+  int NumSubgraphs() const;
+  std::string TensorName(int tensor_index, int subgraph_index) const;
+  PyObject* TensorType(int tensor_index, int subgraph_index) const;
+  PyObject* TensorSize(int tensor_index, int subgraph_index) const;
+  PyObject* TensorSizeSignature(int tensor_index, int subgraph_index) const;
+  PyObject* TensorSparsityParameters(int tensor_index,
+                                     int subgraph_index) const;
+  // Deprecated in favor of TensorQuantizationScales, below.
+  PyObject* TensorQuantization(int tensor_index, int subgraph_index) const;
+  PyObject* TensorQuantizationParameters(int tensor_index,
+                                         int subgraph_index) const;
+  PyObject* SetTensor(int tensor_index, PyObject* value, int subgraph_index);
+  PyObject* GetTensor(int tensor_index, int subgraph_index) const;
+  PyObject* GetSubgraphIndexFromSignature(const char* signature_key);
+  PyObject* GetSignatureDefs() const;
+  PyObject* ResetVariableTensors();
+
+  int NumNodes() const;
+  std::string NodeName(int i) const;
+  PyObject* NodeInputs(int i) const;
+  PyObject* NodeOutputs(int i) const;
+
+  // Returns a reference to tensor index as a numpy array from subgraph. The
+  // base_object should be the interpreter object providing the memory.
+  PyObject* tensor(PyObject* base_object, int tensor_index, int subgraph_index);
+
+  PyObject* SetNumThreads(int num_threads);
+
+  // Adds a delegate to the interpreter.
+  PyObject* ModifyGraphWithDelegate(TfLiteDelegate* delegate);
+
+  // Experimental and subject to change.
+  //
+  // Returns a pointer to the underlying interpreter.
+  Interpreter* interpreter() { return interpreter_.get(); }
+
+ private:
+  // Helper function to construct an `InterpreterWrapper` object.
+  // It only returns InterpreterWrapper if it can construct an `Interpreter`.
+  // Otherwise it returns `nullptr`.
+  static InterpreterWrapper* CreateInterpreterWrapper(
+      std::unique_ptr<Model> model, int op_resolver_id,
+      std::unique_ptr<PythonErrorReporter> error_reporter,
+      const std::vector<std::string>& registerers_by_name,
+      const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
+      std::string* error_msg, bool preserve_all_tensors,
+      bool disable_delegate_clustering, int num_threads,
+      bool default_delegate_latest_features,
+      bool compress_quantization_zero_points,
+      bool disable_delegate_node_fusion = false,
+      bool force_delegate_node_profiling = false);
+
+  InterpreterWrapper(std::unique_ptr<Model> model,
+                     std::unique_ptr<PythonErrorReporter> error_reporter,
+                     std::unique_ptr<tflite::MutableOpResolver> resolver,
+                     std::unique_ptr<Interpreter> interpreter);
+
+  // InterpreterWrapper is not copyable or assignable. We avoid the use of
+  // InterpreterWrapper() = delete here for SWIG compatibility.
+  InterpreterWrapper();
+  InterpreterWrapper(const InterpreterWrapper& rhs);
+
+  // Helper function to resize an input tensor.
+  PyObject* ResizeInputTensorImpl(int i, PyObject* value);
+
+  // The public functions which creates `InterpreterWrapper` should ensure all
+  // these member variables are initialized successfully. Otherwise it should
+  // report the error and return `nullptr`.
+  const std::unique_ptr<Model> model_;
+  const std::unique_ptr<PythonErrorReporter> error_reporter_;
+  const std::unique_ptr<tflite::MutableOpResolver> resolver_;
+  const std::unique_ptr<Interpreter> interpreter_;
+};
+
+}  // namespace interpreter_wrapper
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_

tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc

@@ -0,0 +1,259 @@
+/* Copyright 2020 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 <cstdint>
+#include <functional>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include "pybind11/functional.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+
+namespace py = pybind11;
+using tflite::interpreter_wrapper::InterpreterWrapper;
+
+PYBIND11_MODULE(_pywrap_tensorflow_interpreter_wrapper, m) {
+  m.doc() = R"pbdoc(
+    _pywrap_tensorflow_interpreter_wrapper
+    -----
+  )pbdoc";
+
+  // pybind11 suggests to convert factory functions into constructors, but
+  // when bytes are provided the wrapper will be confused which
+  // constructor to call.
+  m.def("CreateWrapperFromFile",
+        [](const std::string& model_path, int op_resolver_id,
+           const std::vector<std::string>& registerers,
+           bool preserve_all_tensors, bool disable_delegate_clustering) {
+          std::string error;
+          auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromFile(
+              model_path.c_str(), op_resolver_id, registerers, &error,
+              preserve_all_tensors, disable_delegate_clustering);
+          if (!wrapper) {
+            throw std::invalid_argument(error);
+          }
+          return wrapper;
+        });
+  m.def(
+      "CreateWrapperFromFile",
+      [](const std::string& model_path, int op_resolver_id,
+         const std::vector<std::string>& registerers_by_name,
+         const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
+         bool preserve_all_tensors, bool disable_delegate_clustering,
+         int num_threads, bool default_delegate_latest_features,
+         bool compress_quantization_zero_points,
+         bool disable_delegate_node_fusion = false,
+         bool force_delegate_node_profiling = false) {
+        std::string error;
+        auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromFile(
+            model_path.c_str(), op_resolver_id, registerers_by_name,
+            registerers_by_func, &error, preserve_all_tensors,
+            disable_delegate_clustering, num_threads,
+            default_delegate_latest_features, compress_quantization_zero_points,
+            disable_delegate_node_fusion, force_delegate_node_profiling);
+        if (!wrapper) {
+          throw std::invalid_argument(error);
+        }
+        return wrapper;
+      });
+  m.def("CreateWrapperFromBuffer",
+        [](const py::bytes& data, int op_resolver_id,
+           const std::vector<std::string>& registerers,
+           bool preserve_all_tensors, bool disable_delegate_clustering) {
+          std::string error;
+          auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromBuffer(
+              data.ptr(), op_resolver_id, registerers, &error,
+              preserve_all_tensors, disable_delegate_clustering);
+          if (!wrapper) {
+            throw std::invalid_argument(error);
+          }
+          return wrapper;
+        });
+  m.def(
+      "CreateWrapperFromBuffer",
+      [](const py::bytes& data, int op_resolver_id,
+         const std::vector<std::string>& registerers_by_name,
+         const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
+         bool preserve_all_tensors, bool disable_delegate_clustering,
+         int num_threads, bool default_delegate_latest_features,
+         bool compress_quantization_zero_points,
+         bool disable_delegate_node_fusion = false,
+         bool force_delegate_node_profiling = false) {
+        std::string error;
+        auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromBuffer(
+            data.ptr(), op_resolver_id, registerers_by_name,
+            registerers_by_func, &error, preserve_all_tensors,
+            disable_delegate_clustering, num_threads,
+            default_delegate_latest_features, compress_quantization_zero_points,
+            disable_delegate_node_fusion, force_delegate_node_profiling);
+        if (!wrapper) {
+          throw std::invalid_argument(error);
+        }
+        return wrapper;
+      });
+  py::class_<InterpreterWrapper>(m, "InterpreterWrapper", py::module_local())
+      .def(
+          "AllocateTensors",
+          [](InterpreterWrapper& self, int subgraph_index) {
+            return tensorflow::PyoOrThrow(self.AllocateTensors(subgraph_index));
+          },
+          // LINT.IfChange
+          py::arg("subgraph_index") = -1)
+      // LINT.ThenChange(//tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc)
+      .def(
+          "Invoke",
+          [](InterpreterWrapper& self, int subgraph_index) {
+            return tensorflow::PyoOrThrow(self.Invoke(subgraph_index));
+          },
+          py::arg("subgraph_index") = 0)
+      .def("InputIndices",
+           [](const InterpreterWrapper& self) {
+             return tensorflow::PyoOrThrow(self.InputIndices());
+           })
+      .def("OutputIndices",
+           [](InterpreterWrapper& self) {
+             return tensorflow::PyoOrThrow(self.OutputIndices());
+           })
+      .def(
+          "ResizeInputTensor",
+          [](InterpreterWrapper& self, int i, py::handle& value, bool strict,
+             int subgraph_index) {
+            return tensorflow::PyoOrThrow(
+                self.ResizeInputTensor(i, value.ptr(), strict, subgraph_index));
+          },
+          py::arg("i"), py::arg("value"), py::arg("strict"),
+          py::arg("subgraph_index") = 0)
+      .def("NumTensors", &InterpreterWrapper::NumTensors)
+      .def("NumSubgraphs", &InterpreterWrapper::NumSubgraphs)
+      .def("TensorName", &InterpreterWrapper::TensorName)
+      .def("TensorType",
+           [](const InterpreterWrapper& self, int tensor_index,
+              int subgraph_index) {
+             return tensorflow::PyoOrThrow(
+                 self.TensorType(tensor_index, subgraph_index));
+           })
+      .def("TensorSize",
+           [](const InterpreterWrapper& self, int tensor_index,
+              int subgraph_index) {
+             return tensorflow::PyoOrThrow(
+                 self.TensorSize(tensor_index, subgraph_index));
+           })
+      .def("TensorSizeSignature",
+           [](const InterpreterWrapper& self, int tensor_index,
+              int subgraph_index) {
+             return tensorflow::PyoOrThrow(
+                 self.TensorSizeSignature(tensor_index, subgraph_index));
+           })
+      .def("TensorSparsityParameters",
+           [](const InterpreterWrapper& self, int tensor_index,
+              int subgraph_index) {
+             return tensorflow::PyoOrThrow(
+                 self.TensorSparsityParameters(tensor_index, subgraph_index));
+           })
+      .def(
+          "TensorQuantization",
+          [](const InterpreterWrapper& self, int tensor_index,
+             int subgraph_index) {
+            return tensorflow::PyoOrThrow(
+                self.TensorQuantization(tensor_index, subgraph_index));
+          },
+          R"pbdoc(
+            Deprecated in favor of TensorQuantizationParameters.
+          )pbdoc")
+      .def("TensorQuantizationParameters",
+           [](InterpreterWrapper& self, int tensor_index, int subgraph_index) {
+             return tensorflow::PyoOrThrow(self.TensorQuantizationParameters(
+                 tensor_index, subgraph_index));
+           })
+      .def(
+          "SetTensor",
+          [](InterpreterWrapper& self, int tensor_index, py::handle& value,
+             int subgraph_index) {
+            return tensorflow::PyoOrThrow(
+                self.SetTensor(tensor_index, value.ptr(), subgraph_index));
+          },
+          py::arg("i"), py::arg("value"), py::arg("subgraph_index") = 0)
+      .def(
+          "GetTensor",
+          [](const InterpreterWrapper& self, int tensor_index,
+             int subgraph_index) {
+            return tensorflow::PyoOrThrow(
+                self.GetTensor(tensor_index, subgraph_index));
+          },
+          py::arg("tensor_index"), py::arg("subgraph_index") = 0)
+      .def("GetSubgraphIndexFromSignature",
+           [](InterpreterWrapper& self, const char* signature_key) {
+             return tensorflow::PyoOrThrow(
+                 self.GetSubgraphIndexFromSignature(signature_key));
+           })
+      .def("GetSignatureDefs",
+           [](InterpreterWrapper& self) {
+             return tensorflow::PyoOrThrow(self.GetSignatureDefs());
+           })
+      .def("ResetVariableTensors",
+           [](InterpreterWrapper& self) {
+             return tensorflow::PyoOrThrow(self.ResetVariableTensors());
+           })
+      .def("NumNodes", &InterpreterWrapper::NumNodes)
+      .def("NodeName", &InterpreterWrapper::NodeName)
+      .def("NodeInputs",
+           [](const InterpreterWrapper& self, int i) {
+             return tensorflow::PyoOrThrow(self.NodeInputs(i));
+           })
+      .def("NodeOutputs",
+           [](const InterpreterWrapper& self, int i) {
+             return tensorflow::PyoOrThrow(self.NodeOutputs(i));
+           })
+      .def(
+          "tensor",
+          [](InterpreterWrapper& self, py::handle& base_object,
+             int tensor_index, int subgraph_index) {
+            return tensorflow::PyoOrThrow(
+                self.tensor(base_object.ptr(), tensor_index, subgraph_index));
+          },
+          R"pbdoc(
+            Returns a reference to tensor index as a numpy array from subgraph.
+            The base_object should be the interpreter object providing the
+            memory.
+          )pbdoc",
+          py::arg("base_object"), py::arg("tensor_index"),
+          py::arg("subgraph_index") = 0)
+      .def(
+          "ModifyGraphWithDelegate",
+          // Address of the delegate is passed as an argument.
+          [](InterpreterWrapper& self, uintptr_t delegate_ptr) {
+            return tensorflow::PyoOrThrow(self.ModifyGraphWithDelegate(
+                reinterpret_cast<TfLiteDelegate*>(delegate_ptr)));
+          },
+          R"pbdoc(
+            Adds a delegate to the interpreter.
+          )pbdoc")
+      .def(
+          "SetNumThreads",
+          [](InterpreterWrapper& self, int num_threads) {
+            return tensorflow::PyoOrThrow(self.SetNumThreads(num_threads));
+          },
+          R"pbdoc(
+             ask the interpreter to set the number of threads to use.
+          )pbdoc")
+      .def("interpreter", [](InterpreterWrapper& self) {
+        return reinterpret_cast<intptr_t>(self.interpreter());
+      });
+}

tensorflow/lite/python/interpreter_wrapper/numpy.cc

@@ -0,0 +1,327 @@
+/* Copyright 2019 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 <cstddef>
+#include <cstdint>
+#include <memory>
+
+#define TFLITE_IMPORT_NUMPY  // See numpy.h for explanation.
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
+
+namespace tflite {
+namespace python {
+
+void ImportNumpy() { import_array1(); }
+
+}  // namespace python
+
+namespace python_utils {
+
+struct PyObjectDereferencer {
+  void operator()(PyObject* py_object) const { Py_DECREF(py_object); }
+};
+using UniquePyObjectRef = std::unique_ptr<PyObject, PyObjectDereferencer>;
+
+int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) {
+  switch (tf_lite_type) {
+    case kTfLiteFloat32:
+      return NPY_FLOAT32;
+    case kTfLiteFloat16:
+      return NPY_FLOAT16;
+    case kTfLiteBFloat16:
+      // TODO(b/329491949): Supports other ml_dtypes user-defined types.
+      return NPY_USERDEF;
+    case kTfLiteFloat64:
+      return NPY_FLOAT64;
+    case kTfLiteInt32:
+      return NPY_INT32;
+    case kTfLiteUInt32:
+      return NPY_UINT32;
+    case kTfLiteUInt16:
+      return NPY_UINT16;
+    case kTfLiteInt16:
+      return NPY_INT16;
+    case kTfLiteInt4:
+      // TODO(b/246806634): NPY_INT4 currently doesn't exist
+      return NPY_BYTE;
+    case kTfLiteUInt4:
+      return NPY_UINT8;
+    case kTfLiteFloat8E4M3FN:
+    case kTfLiteFloat8E5M2:
+      return NPY_BYTE;
+    case kTfLiteInt2:
+      // TODO(b/246806634): NPY_INT2 currently doesn't exist
+      return NPY_BYTE;
+    case kTfLiteUInt8:
+      return NPY_UINT8;
+    case kTfLiteInt8:
+      return NPY_INT8;
+    case kTfLiteInt64:
+      return NPY_INT64;
+    case kTfLiteUInt64:
+      return NPY_UINT64;
+    case kTfLiteString:
+      return NPY_STRING;
+    case kTfLiteBool:
+      return NPY_BOOL;
+    case kTfLiteComplex64:
+      return NPY_COMPLEX64;
+    case kTfLiteComplex128:
+      return NPY_COMPLEX128;
+    case kTfLiteResource:
+    case kTfLiteVariant:
+      return NPY_OBJECT;
+    case kTfLiteNoType:
+      return NPY_NOTYPE;
+      // Avoid default so compiler errors created when new types are made.
+  }
+  return NPY_NOTYPE;
+}  // NOLINT(direct import ndarraytypes.h cannot be used here)
+
+TfLiteType TfLiteTypeFromPyType(int py_type) {
+  switch (py_type) {
+    case NPY_FLOAT32:
+      return kTfLiteFloat32;
+    case NPY_FLOAT16:
+      return kTfLiteFloat16;
+    case NPY_FLOAT64:
+      return kTfLiteFloat64;
+    case NPY_INT32:
+      return kTfLiteInt32;
+    case NPY_UINT32:
+      return kTfLiteUInt32;
+    case NPY_INT16:
+      return kTfLiteInt16;
+    case NPY_UINT16:
+      return kTfLiteUInt16;
+    case NPY_UINT8:
+      return kTfLiteUInt8;
+    case NPY_INT8:
+      return kTfLiteInt8;
+    case NPY_INT64:
+      return kTfLiteInt64;
+    case NPY_UINT64:
+      return kTfLiteUInt64;
+    case NPY_BOOL:
+      return kTfLiteBool;
+    case NPY_OBJECT:
+    case NPY_STRING:
+    case NPY_UNICODE:
+      return kTfLiteString;
+    case NPY_COMPLEX64:
+      return kTfLiteComplex64;
+    case NPY_COMPLEX128:
+      return kTfLiteComplex128;
+    case NPY_USERDEF:
+      // User-defined types are defined in ml_dtypes. (bfloat16, float8, etc.)
+      // For now, we only support bfloat16.
+      return kTfLiteBFloat16;
+  }
+  return kTfLiteNoType;
+}  // NOLINT(direct import ndarraytypes.h cannot be used here)
+
+TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) {
+  int pyarray_type = PyArray_TYPE(array);
+  return TfLiteTypeFromPyType(pyarray_type);
+}
+
+#if PY_VERSION_HEX >= 0x03030000
+bool FillStringBufferFromPyUnicode(PyObject* value,
+                                   DynamicBuffer* dynamic_buffer) {
+  Py_ssize_t len = -1;
+  const char* buf = PyUnicode_AsUTF8AndSize(value, &len);
+  if (buf == nullptr) {
+    PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8AndSize() failed.");
+    return false;
+  }
+  dynamic_buffer->AddString(buf, len);
+  return true;
+}
+#else
+bool FillStringBufferFromPyUnicode(PyObject* value,
+                                   DynamicBuffer* dynamic_buffer) {
+  UniquePyObjectRef utemp(PyUnicode_AsUTF8String(value));
+  if (!utemp) {
+    PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8String() failed.");
+    return false;
+  }
+  char* buf = nullptr;
+  Py_ssize_t len = -1;
+  if (PyBytes_AsStringAndSize(utemp.get(), &buf, &len) == -1) {
+    PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
+    return false;
+  }
+  dynamic_buffer->AddString(buf, len);
+  return true;
+}
+#endif
+
+bool FillStringBufferFromPyString(PyObject* value,
+                                  DynamicBuffer* dynamic_buffer) {
+  if (PyUnicode_Check(value)) {
+    return FillStringBufferFromPyUnicode(value, dynamic_buffer);
+  }
+
+  char* buf = nullptr;
+  Py_ssize_t len = -1;
+  if (PyBytes_AsStringAndSize(value, &buf, &len) == -1) {
+    PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
+    return false;
+  }
+  dynamic_buffer->AddString(buf, len);
+  return true;
+}
+
+bool FillStringBufferWithPyArray(PyObject* value,
+                                 DynamicBuffer* dynamic_buffer) {
+  if (!PyArray_Check(value)) {
+    PyErr_Format(PyExc_ValueError,
+                 "Passed in value type is not a numpy array, got type %s.",
+                 value->ob_type->tp_name);
+    return false;
+  }
+
+  PyArrayObject* array = reinterpret_cast<PyArrayObject*>(value);
+  switch (PyArray_TYPE(array)) {
+    case NPY_OBJECT:
+    case NPY_STRING:
+    case NPY_UNICODE: {
+      if (PyArray_NDIM(array) == 0) {
+        dynamic_buffer->AddString(static_cast<char*>(PyArray_DATA(array)),
+                                  PyArray_NBYTES(array));
+        return true;
+      }
+      UniquePyObjectRef iter(PyArray_IterNew(value));
+      while (PyArray_ITER_NOTDONE(iter.get())) {
+        UniquePyObjectRef item(PyArray_GETITEM(
+            array, reinterpret_cast<char*>(PyArray_ITER_DATA(iter.get()))));
+
+        if (!FillStringBufferFromPyString(item.get(), dynamic_buffer)) {
+          return false;
+        }
+
+        PyArray_ITER_NEXT(iter.get());
+      }
+      return true;
+    }
+    default:
+      break;
+  }
+
+  PyErr_Format(PyExc_ValueError,
+               "Cannot use numpy array of type %d for string tensor.",
+               PyArray_TYPE(array));
+  return false;
+}
+
+// Helper function to pack int8/uint8 numpy array data into an INT4/UINT4
+// tensor.
+PyObject* Set4BitTensor(TfLiteTensor* tensor, PyArrayObject* array,
+                        int tensor_index) {
+  TfLiteType incoming_type = TfLiteTypeFromPyArray(array);
+  if (tensor->type == kTfLiteInt4) {
+    if (incoming_type != kTfLiteInt8) {
+      PyErr_Format(PyExc_ValueError,
+                   "Cannot set tensor:"
+                   " Expected a numpy array of int8 for INT4 input "
+                   "%d, name: %s, but got %s",
+                   tensor_index, tensor->name,
+                   TfLiteTypeGetName(incoming_type));
+      return nullptr;
+    }
+  } else if (tensor->type == kTfLiteUInt4) {
+    if (incoming_type != kTfLiteUInt8) {
+      PyErr_Format(PyExc_ValueError,
+                   "Cannot set tensor:"
+                   " Expected a numpy array of uint8 for UINT4 input "
+                   "%d, name: %s, but got %s",
+                   tensor_index, tensor->name,
+                   TfLiteTypeGetName(incoming_type));
+      return nullptr;
+    }
+  }
+
+  size_t num_elements = 1;
+  for (int i = 0; i < tensor->dims->size; ++i) {
+    num_elements *= tensor->dims->data[i];
+  }
+  size_t expected_packed_bytes = (num_elements + 1) / 2;
+  size_t actual_numpy_bytes = PyArray_NBYTES(array);
+
+  const char* tensor_type_name = TfLiteTypeGetName(tensor->type);
+
+  if (actual_numpy_bytes != num_elements) {
+    PyErr_Format(PyExc_ValueError,
+                 "Cannot set tensor:"
+                 " Numpy array for %s input %d, name: %s, has %zu bytes, "
+                 "but expected %zu bytes for %zu elements",
+                 tensor_type_name, tensor_index, tensor->name,
+                 actual_numpy_bytes, num_elements, num_elements);
+    return nullptr;
+  }
+
+  if (tensor->data.raw == nullptr && tensor->bytes) {
+    PyErr_Format(PyExc_ValueError,
+                 "Cannot set tensor:"
+                 " Tensor is unallocated. Try calling allocate_tensors()"
+                 " first for input %d, name: %s",
+                 tensor_index, tensor->name);
+    return nullptr;
+  }
+
+  // Pack the int8/uint8 array into int4/uint4
+  uint8_t* packed_data = reinterpret_cast<uint8_t*>(tensor->data.raw);
+
+  if (tensor->type == kTfLiteInt4) {
+    int8_t* numpy_data = reinterpret_cast<int8_t*>(PyArray_DATA(array));
+    for (size_t i = 0; i < expected_packed_bytes; ++i) {
+      int8_t first_nibble = numpy_data[2 * i];
+      int8_t second_nibble =
+          (2 * i + 1 < num_elements) ? numpy_data[2 * i + 1] : 0;
+      if ((first_nibble < -8 || first_nibble > 7) ||
+          (second_nibble < -8 || second_nibble > 7)) {
+        PyErr_Format(PyExc_ValueError,
+                     "Cannot set tensor:"
+                     " Values for INT4 input must be between -8 and 7.");
+        return nullptr;
+      }
+      // Pack the two int8 values into a single byte. The first nibble
+      // occupies the lower 4 bits and the second nibble occupies the upper 4
+      // bits. We mask the first nibble with 0x0F to ensure only the lower 4
+      // bits are used, handling potential sign extension in the int8 value.
+      packed_data[i] = (first_nibble & 0x0F) | (second_nibble << 4);
+    }
+  } else {  // kTfLiteUInt4
+    uint8_t* numpy_data = reinterpret_cast<uint8_t*>(PyArray_DATA(array));
+    for (size_t i = 0; i < expected_packed_bytes; ++i) {
+      uint8_t first_nibble = numpy_data[2 * i];
+      uint8_t second_nibble =
+          (2 * i + 1 < num_elements) ? numpy_data[2 * i + 1] : 0;
+      if (first_nibble > 15 || second_nibble > 15) {
+        PyErr_Format(PyExc_ValueError,
+                     "Cannot set tensor:"
+                     " Values for UINT4 input must be between 0 and 15.");
+        return nullptr;
+      }
+      packed_data[i] = (first_nibble & 0x0F) | (second_nibble << 4);
+    }
+  }
+  Py_RETURN_NONE;
+}
+
+}  // namespace python_utils
+}  // namespace tflite

tensorflow/lite/python/interpreter_wrapper/numpy.h

@@ -0,0 +1,83 @@
+/* Copyright 2019 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_INTERPRETER_WRAPPER_NUMPY_H_
+#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_NUMPY_H_
+
+#ifdef PyArray_Type
+#error "Numpy cannot be included before numpy.h."
+#endif
+
+// Disallow Numpy 1.7 deprecated symbols.
+#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+
+// To handle PyArray_* calles, numpy defines a static lookup table called
+// PyArray_API, or PY_ARRAY_UNIQUE_SYMBOL, if defined. This causes the
+// PyArray_* pointers to be different for different translation units, unless
+// we take care of selectivel defined NO_IMPORT_ARRAY.
+//
+// Virtually every usage will define NO_IMPORT_ARRAY, and will have access to
+// the lookup table via:
+//   extern void **PyArray_API;
+// In numpy.cc we will define TFLITE_IMPORT_NUMPY, effectively disabling that
+// and instead using:
+//   void **PyArray_API;
+// which is initialized when ImportNumpy() is called.
+//
+// If we don't define PY_ARRAY_UNIQUE_SYMBOL then PyArray_API is a static
+// variable, which causes strange crashes when the pointers are used across
+// translation unit boundaries.
+//
+// For more info see https://sourceforge.net/p/numpy/mailman/message/5700519
+// See also tensorflow/compiler/xla/tsl/python/lib/core/numpy.h for a similar
+// approach.
+#define PY_ARRAY_UNIQUE_SYMBOL _tflite_numpy_api
+#ifndef TFLITE_IMPORT_NUMPY
+#define NO_IMPORT_ARRAY
+#endif
+
+#include <Python.h>
+
+#include "numpy/arrayobject.h"  // IWYU pragma: export
+#include "numpy/ufuncobject.h"  // IWYU pragma: export
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/string_util.h"
+
+namespace tflite {
+namespace python {
+
+void ImportNumpy();
+
+}  // namespace python
+
+namespace python_utils {
+
+int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type);
+
+TfLiteType TfLiteTypeFromPyType(int py_type);
+
+TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array);
+
+bool FillStringBufferWithPyArray(PyObject* value,
+                                 DynamicBuffer* dynamic_buffer);
+
+// Helper function to pack int8/uint8 numpy array data into an INT4/UINT4
+// tensor.
+PyObject* Set4BitTensor(TfLiteTensor* tensor, PyArrayObject* array,
+                        int tensor_index);
+
+}  // namespace python_utils
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_NUMPY_H_

tensorflow/lite/python/interpreter_wrapper/python_error_reporter.cc

@@ -0,0 +1,49 @@
+/* Copyright 2019 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 "tensorflow/lite/python/interpreter_wrapper/python_error_reporter.h"
+
+#include <Python.h>
+
+#include <cstdarg>
+#include <cstdio>
+#include <string>
+
+namespace tflite {
+namespace interpreter_wrapper {
+
+// Report an error message
+int PythonErrorReporter::Report(const char* format, va_list args) {
+  char buf[1024];
+  int formatted = vsnprintf(buf, sizeof(buf), format, args);
+  buffer_ << buf;
+  return formatted;
+}
+
+// Set's a Python runtime exception with the last error.
+PyObject* PythonErrorReporter::exception() {
+  std::string last_message = message();
+  PyErr_SetString(PyExc_RuntimeError, last_message.c_str());
+  return nullptr;
+}
+
+// Gets the last error message and clears the buffer.
+std::string PythonErrorReporter::message() {
+  std::string value = buffer_.str();
+  buffer_.clear();
+  return value;
+}
+}  // namespace interpreter_wrapper
+}  // namespace tflite

tensorflow/lite/python/interpreter_wrapper/python_error_reporter.h

@@ -0,0 +1,49 @@
+/* Copyright 2019 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_INTERPRETER_WRAPPER_PYTHON_ERROR_REPORTER_H_
+#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_ERROR_REPORTER_H_
+
+#include <Python.h>
+
+#include <sstream>
+#include <string>
+
+#include "tensorflow/lite/stateful_error_reporter.h"
+
+namespace tflite {
+namespace interpreter_wrapper {
+
+class PythonErrorReporter : public tflite::StatefulErrorReporter {
+ public:
+  PythonErrorReporter() {}
+
+  // Report an error message
+  int Report(const char* format, va_list args) override;
+
+  // Sets a Python runtime exception with the last error and
+  // clears the error message buffer.
+  PyObject* exception();
+
+  // Gets the last error message and clears the buffer.
+  std::string message() override;
+
+ private:
+  std::stringstream buffer_;
+};
+
+}  // namespace interpreter_wrapper
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_ERROR_REPORTER_H_

tensorflow/lite/python/interpreter_wrapper/python_utils.cc

@@ -0,0 +1,45 @@
+/* Copyright 2019 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 "tensorflow/lite/python/interpreter_wrapper/python_utils.h"
+
+#include <cstddef>
+
+namespace tflite {
+namespace python_utils {
+
+int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length) {
+#if PY_MAJOR_VERSION >= 3
+  if (PyUnicode_Check(obj)) {
+    // const_cast<> is for CPython 3.7 finally adding const to the API.
+    *data = const_cast<char*>(PyUnicode_AsUTF8AndSize(obj, length));
+    return *data == nullptr ? -1 : 0;
+  }
+  return PyBytes_AsStringAndSize(obj, data, length);
+#else
+  return PyString_AsStringAndSize(obj, data, length);
+#endif
+}
+
+PyObject* ConvertToPyString(const char* data, size_t length) {
+#if PY_MAJOR_VERSION >= 3
+  return PyBytes_FromStringAndSize(data, length);
+#else
+  return PyString_FromStringAndSize(data, length);
+#endif
+}
+
+}  // namespace python_utils
+}  // namespace tflite

tensorflow/lite/python/interpreter_wrapper/python_utils.h

@@ -0,0 +1,35 @@
+/* Copyright 2019 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_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
+#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
+
+#include <Python.h>
+
+#include <cstddef>
+
+namespace tflite {
+namespace python_utils {
+
+struct PyDecrefDeleter {
+  void operator()(PyObject* p) const { Py_DECREF(p); }
+};
+
+int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length);
+PyObject* ConvertToPyString(const char* data, size_t length);
+
+}  // namespace python_utils
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_