chore: import upstream snapshot with attribution

tensorflow/lite/python/optimize/BUILD

@@ -0,0 +1,109 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "calibration_wrapper_lib",
+    srcs = ["calibration_wrapper.cc"],
+    hdrs = ["calibration_wrapper.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/lite:offset_buffer",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:shared_library",
+        "//tensorflow/lite:string_util",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/python/interpreter_wrapper:numpy",
+        "//tensorflow/lite/python/interpreter_wrapper:python_error_reporter",
+        "//tensorflow/lite/python/interpreter_wrapper:python_utils",
+        "//tensorflow/lite/tools/optimize:quantization_wrapper_utils",
+        "//tensorflow/lite/tools/optimize:quantize_model",
+        "//tensorflow/lite/tools/optimize/calibration:calibration_reader",
+        "//tensorflow/lite/tools/optimize/calibration:calibrator_lib",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/types:optional",
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+pybind_extension(
+    name = "_pywrap_tensorflow_lite_calibration_wrapper",
+    srcs = [
+        "calibration_wrapper_pybind11.cc",
+    ],
+    hdrs = ["calibration_wrapper.h"],
+    additional_stubgen_deps = [
+        "//third_party/py/numpy:numpy",
+    ],
+    common_lib_packages = [
+        "litert/python",
+        "tensorflow/lite/python",
+    ],
+    enable_stub_generation = True,
+    link_in_framework = True,
+    pytype_srcs = [
+        "_pywrap_tensorflow_lite_calibration_wrapper.pyi",
+    ],
+    wrap_py_init = True,
+    deps = [
+        ":calibration_wrapper_lib",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework_stable",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",
+    ],
+)
+
+py_library(
+    name = "calibrator",
+    srcs = [
+        "calibrator.py",
+    ],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":_pywrap_tensorflow_lite_calibration_wrapper",  # buildcleaner: keep
+        "//tensorflow/lite/python:convert_phase",
+        "//tensorflow/lite/python:interpreter",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/util:lazy_loader",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_test(
+    name = "calibrator_test",
+    srcs = ["calibrator_test.py"],
+    data = [
+        ":test_data",
+        "//tensorflow/lite:testdata/multi_add.bin",
+    ],
+    strict_deps = True,
+    tags = ["no_oss"],
+    deps = [
+        ":calibrator",
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/lite/python:schema_py",
+        "//tensorflow/lite/tools:flatbuffer_utils",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:resource_loader",
+    ],
+)

tensorflow/lite/python/optimize/_pywrap_tensorflow_lite_calibration_wrapper.pyi

@@ -0,0 +1,38 @@
+# 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.
+# ==============================================================================
+
+from typing import Callable, overload
+
+class CalibrationWrapper:
+    def __init__(self, arg0: object, arg1: list[str], arg2: list[Callable[[int], None]]) -> None: ...
+    def Calibrate(self) -> object: ...
+    @overload
+    def FeedTensor(self, arg0: object, arg1: str) -> object: ...
+    @overload
+    def FeedTensor(self, arg0: object) -> object: ...
+    @overload
+    def Prepare(self, arg0: object, arg1: str) -> object: ...
+    @overload
+    def Prepare(self, arg0: object) -> object: ...
+    @overload
+    def Prepare(self, arg0: str) -> object: ...
+    @overload
+    def Prepare(self) -> object: ...
+    @overload
+    def QuantizeModel(self, arg0: int, arg1: int, arg2: bool, arg3: int, arg4: int, arg5: bool, arg6: bool) -> object: ...
+    @overload
+    def QuantizeModel(self, arg0: int, arg1: int, arg2: bool, arg3: str) -> object: ...
+
+def AddIntermediateTensors(arg0: object) -> object: ...

tensorflow/lite/python/optimize/calibration_wrapper.cc

@@ -0,0 +1,844 @@
+/* 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.
+==============================================================================*/
+// clang-format off
+// This #include needs to precede the inclusion of any other TF Lite header
+// file that might depend on the non-mutable schema_generated.h, directly,
+// e.g. core/api/op_resolver.h, or indirectly, e.g. core/subgraph.h.
+// That's because "tensorflow/lite/schema/mutable/schema_generated.h"
+// and "tensorflow/lite/schema/schema_generated.h" both use the same
+// header guard macro (FLATBUFFERS_GENERATED_SCHEMA_TFLITE_H_), but have
+// different contents (the former is a superset of the latter). In particular
+// the one in mutable/ is built with the "--gen-mutable" and "--gen-object-api"
+// flags to the flatbuffer schema compiler which cause some additional
+// (non-virtual) accessor methods and API functions to be declared.
+// The code here uses those methods, so we need to make sure that we get
+// the mutable variant of this header.
+#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
+
+#include "tensorflow/lite/python/optimize/calibration_wrapper.h"
+// clang-format on
+
+// NOLINTBEGIN
+// Nolint disables warnings about header file ordering caused by
+// `mutable/schema_generated.h`.
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <functional>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+// NOLINTEND
+
+#include "absl/algorithm/container.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/lite/offset_buffer.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/core/model_builder.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
+#include "tensorflow/lite/python/interpreter_wrapper/python_error_reporter.h"
+#include "tensorflow/lite/python/interpreter_wrapper/python_utils.h"
+#include "tensorflow/lite/shared_library.h"
+#include "tensorflow/lite/string_util.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_reader.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibrator.h"
+#include "tensorflow/lite/tools/optimize/quantization_wrapper_utils.h"
+#include "tensorflow/lite/tools/optimize/quantize_model.h"
+
+#define TFLITE_PY_CHECK(x)               \
+  if ((x) != kTfLiteOk) {                \
+    return error_reporter_->exception(); \
+  }
+
+#define TFLITE_PY_ENSURE_VALID_INTERPRETER()                               \
+  if (!interpreter_) {                                                     \
+    PyErr_SetString(PyExc_ValueError, "Interpreter was not initialized."); \
+    return nullptr;                                                        \
+  }
+
+namespace tflite {
+namespace calibration_wrapper {
+namespace {
+
+using ::tflite::interpreter_wrapper::PythonErrorReporter;
+using ::tflite::ops::builtin::BuiltinOpResolver;
+using ::tflite::optimize::AddIntermediateTensorsToFusedOp;
+using ::tflite::optimize::QuantizeModelAllOperators;
+using ::tflite::optimize::calibration::BuildLoggingInterpreter;
+using ::tflite::optimize::calibration::CalibrationReader;
+using ::tflite::python::ImportNumpy;
+using ::tflite::python_utils::ConvertFromPyString;
+using ::tflite::python_utils::ConvertToPyString;
+using ::tflite::python_utils::PyDecrefDeleter;
+using ::tflite::python_utils::TfLiteTypeFromPyArray;
+using ::tflite::python_utils::TfLiteTypeFromPyType;
+
+std::unique_ptr<ModelT> CreateMutableModel(const Model& model) {
+  auto copied_model = std::make_unique<ModelT>();
+  model.UnPackTo(copied_model.get(), nullptr);
+  return copied_model;
+}
+
+bool NoOpModel(const FlatBufferModel& model) {
+  return model->subgraphs()->size() == 1 &&
+         (!model->subgraphs()->begin()->operators() ||
+          model->subgraphs()->begin()->operators()->size() == 0);
+}
+
+inline TensorType TfLiteTypeToSchemaType(TfLiteType type) {
+  switch (type) {
+    case kTfLiteNoType:
+      return TensorType_FLOAT32;  // TODO(b/129336260): No schema type for none.
+    case kTfLiteFloat32:
+      return TensorType_FLOAT32;
+    case kTfLiteFloat16:
+      return TensorType_FLOAT16;
+    case kTfLiteBFloat16:
+      return TensorType_BFLOAT16;
+    case kTfLiteFloat64:
+      return TensorType_FLOAT64;
+    case kTfLiteInt32:
+      return TensorType_INT32;
+    case kTfLiteUInt32:
+      return TensorType_UINT32;
+    case kTfLiteInt2:
+      return TensorType_INT2;
+    case kTfLiteUInt4:
+      return TensorType_UINT4;
+    case kTfLiteFloat8E4M3FN:
+      return TensorType_FLOAT8_E4M3FN;
+    case kTfLiteFloat8E5M2:
+      return TensorType_FLOAT8_E5M2;
+    case kTfLiteInt4:
+      return TensorType_INT4;
+    case kTfLiteUInt8:
+      return TensorType_UINT8;
+    case kTfLiteInt8:
+      return TensorType_INT8;
+    case kTfLiteInt64:
+      return TensorType_INT64;
+    case kTfLiteUInt64:
+      return TensorType_UINT64;
+    case kTfLiteString:
+      return TensorType_STRING;
+    case kTfLiteBool:
+      return TensorType_BOOL;
+    case kTfLiteInt16:
+      return TensorType_INT16;
+    case kTfLiteUInt16:
+      return TensorType_UINT16;
+    case kTfLiteComplex64:
+      return TensorType_COMPLEX64;
+    case kTfLiteComplex128:
+      return TensorType_COMPLEX128;
+    case kTfLiteResource:
+      return TensorType_RESOURCE;
+    case kTfLiteVariant:
+      return TensorType_VARIANT;
+  }
+  // No default to get compiler error when new type is introduced.
+}
+
+bool RegisterCustomOpByName(const char* registerer_name,
+                            MutableOpResolver* resolver) {
+  // Registerer functions take a pointer to a BuiltinOpResolver as an input
+  // parameter and return void.
+  // TODO(b/137576229): We should implement this functionality in a more
+  // principled way.
+  typedef void (*RegistererFunctionType)(MutableOpResolver*);
+
+  // Look for the Registerer function by name.
+  RegistererFunctionType registerer = reinterpret_cast<RegistererFunctionType>(
+      SharedLibrary::GetSymbol(registerer_name));
+
+  // Fail in an informative way if the function was not found.
+  if (registerer == nullptr) {
+    PyErr_Format(PyExc_ValueError,
+                 "Looking up symbol '%s' failed with error '%s'.",
+                 registerer_name, SharedLibrary::GetError());
+    return false;
+  }
+
+  // Call the registerer with the resolver.
+  registerer(resolver);
+  return true;
+}
+
+// Returns the dimension from the stored list in the PyObject. If the given
+// PyObject is not a list, it will return absl::optional and set the Python
+// error message to notify users.
+std::optional<std::vector<int>> ConvertInputShapeToVector(
+    PyObject* input_shapes, size_t index) {
+  PyObject* shape = PyList_GetItem(input_shapes, index);
+  if (!shape || !PyList_Check(shape)) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid %ld input shape: expected to be a list.", index);
+    return std::nullopt;
+  }
+  size_t size = PyList_Size(shape);
+  std::vector<int> dims(size);
+  for (size_t dim_index = 0; dim_index < size; ++dim_index) {
+    PyObject* dim = PyList_GetItem(shape, dim_index);
+    dims[dim_index] = PyLong_AsLong(dim);
+  }
+  return dims;
+}
+
+// Finds the starting position of the offset buffer within `model_buffer` if the
+// `model_buffer` can be split into base buffer and offset buffer. Returns
+// `std::nullopt` iff offset buffer is not used or there were no buffers with
+// valid offset. Assumes `model_buffer` is valid.
+std::optional<int64_t> GetOffsetBufferStartPosition(
+    const absl::string_view model_buffer) {
+  const Model& model = *GetModel(model_buffer.data());
+
+  if (!FlatBufferModel::CheckBufferOutsideModel(&model)) {
+    // Means the offset buffer is not used, e.g.
+    // `_experimental_use_buffer_offset` is not set.
+    return std::nullopt;
+  }
+
+  const int64_t int64_max = std::numeric_limits<int64_t>::max();
+  const int64_t min_offset = absl::c_accumulate(
+      *model.buffers(), /*init=*/int64_max,
+      /*binary_op=*/[](const int64_t acc, const Buffer* buffer) -> int64_t {
+        const int64_t buffer_offset = buffer->offset();
+        return IsValidBufferOffset(buffer_offset) ? std::min(acc, buffer_offset)
+                                                  : acc;
+      });
+  if (min_offset == int64_max) {
+    // Means there were no buffers with valid offset.
+    return std::nullopt;
+  }
+  return min_offset;
+}
+
+// Splits the model buffer into base buffer and offset buffer. Offset buffer may
+// exist when `_experimental_use_buffer_offset` is set.
+std::pair<absl::string_view, absl::string_view> SplitOffsetBuffer(
+    const absl::string_view model_buffer) {
+  const std::optional<int64_t> offset_buffer_pos =
+      GetOffsetBufferStartPosition(model_buffer);
+  if (offset_buffer_pos == std::nullopt) {
+    return {model_buffer, absl::string_view(model_buffer.data(), /*len=*/0)};
+  }
+
+  const absl::string_view base_buffer(model_buffer.data(), *offset_buffer_pos);
+
+  const int64_t offset_buffer_length = model_buffer.size() - *offset_buffer_pos;
+  const absl::string_view offset_buffer(
+      model_buffer.data() + *offset_buffer_pos, offset_buffer_length);
+
+  return {base_buffer, offset_buffer};
+}
+
+// Merges `base_buffer` with the `offset_buffer` that contains the actual tensor
+// buffer data.
+std::string MergeOffsetBuffer(const absl::string_view base_buffer,
+                              const absl::string_view offset_buffer) {
+  return absl::StrCat(base_buffer, offset_buffer);
+}
+
+// Updates buffer offsets in `base_buffer` by `offset_diff`.
+std::string UpdateBufferOffsets(const absl::string_view base_buffer,
+                                const int64_t offset_diff) {
+  std::string result_buffer(base_buffer);
+
+  Model* mutable_model = GetMutableModel(result_buffer.data());
+  for (Buffer* buffer : *mutable_model->mutable_buffers()) {
+    if (const int64_t offset = buffer->offset(); IsValidBufferOffset(offset)) {
+      buffer->mutate_offset(offset + offset_diff);
+    }
+  }
+
+  return result_buffer;
+}
+
+}  // namespace
+
+PyObject* AddIntermediateTensors(PyObject* data) {
+  char* buf = nullptr;
+  Py_ssize_t length;
+  std::unique_ptr<PythonErrorReporter> error_reporter(new PythonErrorReporter);
+  ImportNumpy();
+
+  if (ConvertFromPyString(data, &buf, &length) == -1) {
+    return nullptr;
+  }
+
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromBuffer(buf, length, error_reporter.get());
+  if (!model) {
+    PyErr_Format(PyExc_ValueError, "Invalid model");
+    return nullptr;
+  }
+
+  const auto [base_buffer, offset_buffer] =
+      SplitOffsetBuffer(/*model_buffer=*/absl::string_view(buf, length));
+
+  flatbuffers::FlatBufferBuilder builder;
+  auto tflite_model = CreateMutableModel(*model->GetModel());
+  if (AddIntermediateTensorsToFusedOp(&builder, tflite_model.get()) !=
+      kTfLiteOk) {
+    error_reporter->exception();
+    return nullptr;
+  }
+
+  const int64_t result_base_buffer_size = builder.GetSize();
+  if (result_base_buffer_size == 0) {
+    // When AddIntermediateTensorsToFusedOp early returns, return the model as
+    // it is.
+    return ConvertToPyString(buf, length);
+  }
+
+  const int64_t offset_diff =
+      result_base_buffer_size - static_cast<int64_t>(base_buffer.size());
+  const std::string updated_result_base_buffer = UpdateBufferOffsets(
+      /*base_buffer=*/absl::string_view(
+          reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
+          builder.GetSize()),
+      offset_diff);
+
+  const std::string result_buffer =
+      MergeOffsetBuffer(updated_result_base_buffer, offset_buffer);
+
+  return ConvertToPyString(result_buffer.data(), result_buffer.size());
+}
+
+CalibrationWrapper::CalibrationWrapper(
+    std::unique_ptr<Interpreter> interpreter,
+    std::unique_ptr<BuiltinOpResolver> resolver,
+    std::unique_ptr<PythonErrorReporter> error_reporter,
+    std::unique_ptr<FlatBufferModel> model,
+    std::unique_ptr<CalibrationReader> reader,
+    std::unique_ptr<std::string> model_str)
+    : interpreter_(std::move(interpreter)),
+      error_reporter_(std::move(error_reporter)),
+      resolver_(std::move(resolver)),
+      model_(std::move(model)),
+      reader_(std::move(reader)),
+      model_str_(std::move(model_str)) {}
+
+CalibrationWrapper::~CalibrationWrapper() = default;
+
+PyObject* CalibrationWrapper::Prepare() {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+  TFLITE_PY_CHECK(interpreter_->AllocateTensors());
+  TFLITE_PY_CHECK(interpreter_->ResetVariableTensors());
+  Py_RETURN_NONE;
+}
+
+PyObject* CalibrationWrapper::Prepare(std::string signature_key) {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+  impl::SignatureRunner* runner =
+      interpreter_->GetSignatureRunner(signature_key.c_str());
+  if (runner == nullptr) {
+    PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
+                 signature_key.c_str());
+    return nullptr;
+  }
+  TFLITE_PY_CHECK(runner->AllocateTensors());
+  TFLITE_PY_CHECK(interpreter_->ResetVariableTensors());
+  Py_RETURN_NONE;
+}
+
+PyObject* CalibrationWrapper::Prepare(PyObject* input_shapes,
+                                      std::string signature_key) {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+  if (!PyList_Check(input_shapes)) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input shapes: expected shapes to be a list.");
+    return nullptr;
+  }
+  const int subgraph_index =
+      interpreter_->GetSubgraphIndexFromSignature(signature_key.c_str());
+  if (subgraph_index == -1) {
+    PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
+                 signature_key.c_str());
+    return nullptr;
+  }
+  auto* subgraph = interpreter_->subgraph(subgraph_index);
+
+  const size_t inputs_size = PyList_Size(input_shapes);
+  if (inputs_size != subgraph->inputs().size()) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input shapes: expected %ld items got %ld items.",
+                 subgraph->inputs().size(), inputs_size);
+    return nullptr;
+  }
+
+  for (size_t i = 0; i < inputs_size; ++i) {
+    std::optional<std::vector<int>> dims =
+        ConvertInputShapeToVector(input_shapes, i);
+    if (!dims.has_value()) {
+      return nullptr;
+    }
+    int input_tensor_idx = subgraph->inputs()[i];
+    if (subgraph->ResizeInputTensor(input_tensor_idx, *dims) != kTfLiteOk) {
+      PyErr_Format(PyExc_ValueError, "Failed to resize %ld input tensor.", i);
+      return nullptr;
+    }
+  }
+
+  return Prepare(signature_key);
+}
+
+PyObject* CalibrationWrapper::Prepare(PyObject* input_shapes) {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+  if (!PyList_Check(input_shapes)) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input shapes: expected shapes to be a list.");
+    return nullptr;
+  }
+
+  const size_t inputs_size = PyList_Size(input_shapes);
+  if (inputs_size != interpreter_->inputs().size()) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input shapes: expected %ld items got %ld items.",
+                 interpreter_->inputs().size(), inputs_size);
+    return nullptr;
+  }
+
+  for (size_t i = 0; i < inputs_size; ++i) {
+    std::optional<std::vector<int>> dims =
+        ConvertInputShapeToVector(input_shapes, i);
+    if (!dims.has_value()) {
+      return nullptr;
+    }
+    int input_tensor_idx = interpreter_->inputs()[i];
+    if (interpreter_->ResizeInputTensor(input_tensor_idx, *dims) != kTfLiteOk) {
+      PyErr_Format(PyExc_ValueError, "Failed to resize %ld input tensor.", i);
+      return nullptr;
+    }
+  }
+
+  return Prepare();
+}
+
+PyObject* CalibrationWrapper::FeedTensor(PyObject* input_value,
+                                         std::string signature_key) {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+  if (!PyList_Check(input_value)) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input type: expected input to be a list.");
+    return nullptr;
+  }
+  const int subgraph_index =
+      interpreter_->GetSubgraphIndexFromSignature(signature_key.c_str());
+  if (subgraph_index == -1) {
+    PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
+                 signature_key.c_str());
+    return nullptr;
+  }
+  const size_t inputs_size = PyList_Size(input_value);
+
+  auto* subgraph = interpreter_->subgraph(subgraph_index);
+  if (inputs_size != subgraph->inputs().size()) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input size: expected %ld items got %ld items.",
+                 subgraph->inputs().size(), inputs_size);
+    return nullptr;
+  }
+
+  for (size_t i = 0; i < inputs_size; ++i) {
+    PyObject* input = PyList_GetItem(input_value, i);
+    if (!input) {
+      return nullptr;
+    }
+    int input_tensor_idx = subgraph->inputs()[i];
+    if (!SetTensor(input_tensor_idx, input, signature_key)) {
+      return nullptr;
+    }
+  }
+
+  TFLITE_PY_CHECK(subgraph->Invoke());
+  Py_RETURN_NONE;
+}
+
+PyObject* CalibrationWrapper::FeedTensor(PyObject* input_value) {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+  if (!PyList_Check(input_value)) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input type: expected input to be a list.");
+    return nullptr;
+  }
+
+  const size_t inputs_size = PyList_Size(input_value);
+
+  if (inputs_size != interpreter_->inputs().size()) {
+    PyErr_Format(PyExc_ValueError,
+                 "Invalid input size: expected %ld items got %ld items.",
+                 interpreter_->inputs().size(), inputs_size);
+    return nullptr;
+  }
+
+  for (size_t i = 0; i < inputs_size; ++i) {
+    PyObject* input = PyList_GetItem(input_value, i);
+    if (!input) {
+      return nullptr;
+    }
+    int input_tensor_idx = interpreter_->inputs()[i];
+    if (!SetTensor(input_tensor_idx, input)) {
+      return nullptr;
+    }
+  }
+
+  TFLITE_PY_CHECK(interpreter_->Invoke());
+  Py_RETURN_NONE;
+}
+
+PyObject* CalibrationWrapper::SetTensor(int index, PyObject* value,
+                                        std::string signature_key) {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+  std::unique_ptr<PyObject, PyDecrefDeleter> array_safe(
+      PyArray_FromAny(value, nullptr, 0, 0, NPY_ARRAY_CARRAY, nullptr));
+  if (!array_safe) {
+    PyErr_SetString(PyExc_ValueError,
+                    "Failed to convert value into readable tensor.");
+    return nullptr;
+  }
+
+  PyArrayObject* array = reinterpret_cast<PyArrayObject*>(array_safe.get());
+
+  const int subgraph_index =
+      interpreter_->GetSubgraphIndexFromSignature(signature_key.c_str());
+  if (subgraph_index == -1) {
+    PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
+                 signature_key.c_str());
+    return nullptr;
+  }
+  auto* subgraph = interpreter_->subgraph(subgraph_index);
+  const TfLiteTensor* tensor = subgraph->tensor(index);
+
+  if (TfLiteTypeFromPyArray(array) != tensor->type) {
+    PyErr_Format(PyExc_ValueError,
+                 "Cannot set tensor: "
+                 "Got value of type %s "
+                 "but expected type %s for input %d, name: %s ",
+                 TfLiteTypeGetName(TfLiteTypeFromPyArray(array)),
+                 TfLiteTypeGetName(tensor->type), index, tensor->name);
+    return nullptr;
+  }
+
+  if (PyArray_NDIM(array) != tensor->dims->size) {
+    PyErr_Format(PyExc_ValueError,
+                 "Cannot set tensor: Dimension count mismatch, expected %d "
+                 "but found %d",
+                 tensor->dims->size, PyArray_NDIM(array));
+    return nullptr;
+  }
+
+  std::vector<int> dims(PyArray_NDIM(array));
+  bool has_unknown_dims = false;
+  for (int j = 0; j < PyArray_NDIM(array); ++j) {
+    // Ensure the calibration data input shape is the same as the model input
+    // shape unless the dimension is unknown.
+    if (tensor->dims_signature != nullptr &&
+        tensor->dims_signature->size == tensor->dims->size &&
+        tensor->dims_signature->data[j] == -1) {
+      has_unknown_dims = true;
+    } else if (tensor->dims->data[j] != PyArray_SHAPE(array)[j]) {
+      PyErr_Format(PyExc_ValueError,
+                   "Cannot set tensor: Size mismatch, expected %d for dim "
+                   "%d but found %ld",
+                   tensor->dims->data[j], j, PyArray_SHAPE(array)[j]);
+      return nullptr;
+    }
+    dims[j] = PyArray_SHAPE(array)[j];
+  }
+
+  // Resize the input tensor if there are unknown dimensions.
+  if (has_unknown_dims) {
+    // Does strict checking on the `ResizeInputTensor` call.
+    TFLITE_PY_CHECK(subgraph->ResizeInputTensorStrict(index, dims));
+    TFLITE_PY_CHECK(subgraph->AllocateTensors());
+  }
+
+  // Re-read the updated tensor after the allocation is done.
+  tensor = subgraph->tensor(index);
+
+  size_t size = PyArray_NBYTES(array);
+
+  if (tensor->type == kTfLiteString) {
+    DynamicBuffer buffer;
+    buffer.AddString(reinterpret_cast<const char*>(PyArray_BYTES(array)), size);
+    buffer.WriteToTensor(subgraph->tensor(index), /*new_shape=*/nullptr);
+    Py_RETURN_NONE;
+  }
+
+  if (size != tensor->bytes) {
+    PyErr_Format(PyExc_ValueError,
+                 "numpy array had %zu bytes but expected %zu bytes.", size,
+                 tensor->bytes);
+    return nullptr;
+  }
+  memcpy(tensor->data.raw, PyArray_DATA(array), size);
+  Py_RETURN_NONE;
+}
+
+PyObject* CalibrationWrapper::SetTensor(int index, PyObject* value) {
+  TFLITE_PY_ENSURE_VALID_INTERPRETER();
+
+  std::unique_ptr<PyObject, PyDecrefDeleter> array_safe(
+      PyArray_FromAny(value, nullptr, 0, 0, NPY_ARRAY_CARRAY, nullptr));
+  if (!array_safe) {
+    PyErr_SetString(PyExc_ValueError,
+                    "Failed to convert value into readable tensor.");
+    return nullptr;
+  }
+
+  PyArrayObject* array = reinterpret_cast<PyArrayObject*>(array_safe.get());
+  const TfLiteTensor* tensor = interpreter_->tensor(index);
+
+  if (TfLiteTypeFromPyArray(array) != tensor->type) {
+    PyErr_Format(PyExc_ValueError,
+                 "Cannot set tensor: "
+                 "Got value of type %s "
+                 "but expected type %s for input %d, name: %s ",
+                 TfLiteTypeGetName(TfLiteTypeFromPyArray(array)),
+                 TfLiteTypeGetName(tensor->type), index, tensor->name);
+    return nullptr;
+  }
+
+  if (PyArray_NDIM(array) != tensor->dims->size) {
+    PyErr_Format(
+        PyExc_ValueError,
+        "Cannot set tensor: Dimension count mismatch, expected %d but found %d",
+        tensor->dims->size, PyArray_NDIM(array));
+    return nullptr;
+  }
+
+  std::vector<int> dims(PyArray_NDIM(array));
+  bool has_unknown_dims = false;
+  for (int j = 0; j < PyArray_NDIM(array); ++j) {
+    // Ensure the calibration data input shape is the same as the model input
+    // shape unless the dimension is unknown.
+    if (tensor->dims_signature != nullptr &&
+        tensor->dims_signature->size == tensor->dims->size &&
+        tensor->dims_signature->data[j] == -1) {
+      has_unknown_dims = true;
+    } else if (tensor->dims->data[j] != PyArray_SHAPE(array)[j]) {
+      PyErr_Format(PyExc_ValueError,
+                   "Cannot set tensor: Size mismatch, expected %d for dim "
+                   "%d but found %ld",
+                   tensor->dims->data[j], j, PyArray_SHAPE(array)[j]);
+      return nullptr;
+    }
+    dims[j] = PyArray_SHAPE(array)[j];
+  }
+
+  // Resize the input tensor if there are unknown dimensions.
+  if (has_unknown_dims) {
+    // Does strict checking on the `ResizeInputTensor` call.
+    TFLITE_PY_CHECK(interpreter_->ResizeInputTensorStrict(index, dims));
+    TFLITE_PY_CHECK(interpreter_->AllocateTensors());
+  }
+
+  // Re-read the updated tensor after the allocation is done.
+  tensor = interpreter_->tensor(index);
+
+  size_t size = PyArray_NBYTES(array);
+
+  if (tensor->type == kTfLiteString) {
+    DynamicBuffer buffer;
+    buffer.AddString(reinterpret_cast<const char*>(PyArray_BYTES(array)), size);
+    buffer.WriteToTensor(interpreter_->tensor(index), /*new_shape=*/nullptr);
+    Py_RETURN_NONE;
+  }
+
+  if (size != tensor->bytes) {
+    PyErr_Format(PyExc_ValueError,
+                 "numpy array had %zu bytes but expected %zu bytes.", size,
+                 tensor->bytes);
+    return nullptr;
+  }
+  memcpy(tensor->data.raw, PyArray_DATA(array), size);
+  Py_RETURN_NONE;
+}
+
+PyObject* CalibrationWrapper::Calibrate() {
+  const auto [base_buffer, offset_buffer] =
+      SplitOffsetBuffer(/*model_buffer=*/absl::string_view(
+          reinterpret_cast<const char*>(model_->allocation()->base()),
+          model_->allocation()->bytes()));
+
+  auto tflite_model = CreateMutableModel(*model_->GetModel());
+  reader_->AddCalibrationToModel(tflite_model.get(), /*update=*/false);
+  flatbuffers::FlatBufferBuilder builder;
+  auto loc = Model::Pack(builder, tflite_model.get());
+  FinishModelBuffer(builder, loc);
+
+  const int64_t result_base_buffer_size = builder.GetSize();
+  const int64_t offset_diff =
+      result_base_buffer_size - static_cast<int64_t>(base_buffer.size());
+  const std::string updated_result_base_buffer = UpdateBufferOffsets(
+      /*base_buffer=*/absl::string_view(
+          reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
+          result_base_buffer_size),
+      offset_diff);
+  const std::string result_buffer =
+      MergeOffsetBuffer(updated_result_base_buffer, offset_buffer);
+
+  return ConvertToPyString(result_buffer.data(), result_buffer.size());
+}
+
+PyObject* CalibrationWrapper::QuantizeModel(int input_py_type,
+                                            int output_py_type,
+                                            bool allow_float,
+                                            int activations_py_type,
+                                            int bias_py_type) {
+  return QuantizeModel(
+      input_py_type, output_py_type, allow_float, activations_py_type,
+      bias_py_type,
+      /*disable_per_channel=*/false,
+      /*disable_per_channel_quantization_for_dense_layers=*/false);
+}
+
+PyObject* CalibrationWrapper::QuantizeModel(
+    int input_py_type, int output_py_type, bool allow_float,
+    int activations_py_type, int bias_py_type, bool disable_per_channel,
+    bool disable_per_channel_quantization_for_dense_layers) {
+  if (NoOpModel(*model_)) {
+    return ConvertToPyString(model_str_->data(), model_str_->size());
+  }
+
+  TfLiteType input_type = TfLiteTypeFromPyType(input_py_type);
+  TfLiteType output_type = TfLiteTypeFromPyType(output_py_type);
+  TfLiteType activations_type = TfLiteTypeFromPyType(activations_py_type);
+  TfLiteType bias_type = TfLiteTypeFromPyType(bias_py_type);
+
+  if (input_type == kTfLiteNoType || output_type == kTfLiteNoType) {
+    PyErr_SetString(PyExc_ValueError,
+                    "Input/output type cannot be kTfLiteNoType");
+    return nullptr;
+  }
+  auto tflite_model = CreateMutableModel(*model_->GetModel());
+  reader_->AddCalibrationToModel(tflite_model.get(), /*update=*/false);
+  flatbuffers::FlatBufferBuilder builder;
+  auto status = kTfLiteOk;
+
+  status = QuantizeModelAllOperators(
+      &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type),
+      TfLiteTypeToSchemaType(output_type), allow_float,
+      TfLiteTypeToSchemaType(activations_type),
+      TfLiteTypeToSchemaType(bias_type), disable_per_channel,
+      disable_per_channel_quantization_for_dense_layers, error_reporter_.get());
+
+  if (status != kTfLiteOk) {
+    error_reporter_->exception();
+    return nullptr;
+  }
+
+  return ConvertToPyString(
+      reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
+      builder.GetSize());
+}
+
+PyObject* CalibrationWrapper::QuantizeModel(int input_py_type,
+                                            int output_py_type,
+                                            bool allow_float,
+                                            const char* operator_output_name) {
+  string op_name = std::string(operator_output_name);
+
+  TfLiteType input_type = TfLiteTypeFromPyType(input_py_type);
+  TfLiteType output_type = TfLiteTypeFromPyType(output_py_type);
+  if (input_type == kTfLiteNoType || output_type == kTfLiteNoType) {
+    PyErr_SetString(PyExc_ValueError,
+                    "Input/output type cannot be kTfLiteNoType");
+    return nullptr;
+  }
+  auto tflite_model = CreateMutableModel(*model_->GetModel());
+  reader_->AddCalibrationToModel(tflite_model.get(), /*update=*/false);
+  flatbuffers::FlatBufferBuilder builder;
+  auto status = optimize::QuantizeModel(
+      &builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type),
+      TfLiteTypeToSchemaType(output_type), allow_float, {op_name},
+      /*activations_type=*/TensorType_INT8, /*bias_type=*/TensorType_INT32,
+      error_reporter_.get());
+  if (status != kTfLiteOk) {
+    error_reporter_->exception();
+    return nullptr;
+  }
+
+  return ConvertToPyString(
+      reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
+      builder.GetSize());
+}
+
+/*static*/ CalibrationWrapper* CalibrationWrapper::CreateWrapperCPPFromBuffer(
+    PyObject* data, const std::vector<std::string>& registerers_by_name,
+    const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
+    std::string* error_msg) {
+  char* buf = nullptr;
+  Py_ssize_t length;
+  std::unique_ptr<PythonErrorReporter> error_reporter(new PythonErrorReporter);
+  ImportNumpy();
+
+  if (ConvertFromPyString(data, &buf, &length) == -1) {
+    *error_msg = "Failed to convert from python string";
+    return nullptr;
+  }
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::BuildFromBuffer(buf, length, error_reporter.get());
+  if (!model) {
+    *error_msg = "Invalid model";
+    return nullptr;
+  }
+
+  auto resolver = std::make_unique<BuiltinOpResolver>();
+  for (const auto& registerer : registerers_by_name) {
+    if (!RegisterCustomOpByName(registerer.c_str(), resolver.get())) {
+      *error_msg =
+          absl::StrFormat("Looking up symbol '%s' failed with error '%s'.",
+                          registerer.c_str(), SharedLibrary::GetError());
+      return nullptr;
+    }
+  }
+  for (const auto& registerer : registerers_by_func) {
+    registerer(reinterpret_cast<uintptr_t>(resolver.get()));
+  }
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status =
+      BuildLoggingInterpreter(*model, *resolver, &interpreter, &reader);
+  if (status != kTfLiteOk) {
+    *error_msg = error_reporter->message();
+    return nullptr;
+  }
+
+  auto model_str = std::make_unique<std::string>(buf, length);
+  // If we are not going to use this string during quantization, reset the
+  // pointer and release the memory.
+  if (!NoOpModel(*model)) {
+    model_str.reset();
+  }
+
+  auto wrapper = new CalibrationWrapper(
+      std::move(interpreter), std::move(resolver), std::move(error_reporter),
+      std::move(model), std::move(reader), std::move(model_str));
+  return wrapper;
+}
+
+}  // namespace calibration_wrapper
+}  // namespace tflite

tensorflow/lite/python/optimize/calibration_wrapper.h

@@ -0,0 +1,138 @@
+/* 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_OPTIMIZE_CALIBRATION_WRAPPER_H_
+#define TENSORFLOW_LITE_PYTHON_OPTIMIZE_CALIBRATION_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"
+
+// We forward declare TFLite classes here to avoid exposing them to SWIG.
+namespace tflite {
+namespace ops {
+namespace builtin {
+class BuiltinOpResolver;
+}  // namespace builtin
+}  // namespace ops
+
+namespace impl {
+class FlatBufferModel;
+}
+
+namespace interpreter_wrapper {
+class PythonErrorReporter;
+}  // namespace interpreter_wrapper
+
+namespace optimize {
+namespace calibration {
+class CalibrationReader;
+}  // namespace calibration
+}  // namespace optimize
+
+namespace calibration_wrapper {
+
+PyObject* AddIntermediateTensors(PyObject* data);
+
+class CalibrationWrapper {
+ public:
+  // SWIG caller takes ownership of pointer.
+  static CalibrationWrapper* CreateWrapperCPPFromBuffer(
+      PyObject* data, const std::vector<std::string>& registerers_by_name,
+      const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
+      std::string* error_msg);
+  ~CalibrationWrapper();
+
+  // Allocates the primary subgraph's tensors.
+  PyObject* Prepare();
+
+  // Allocates the tensors of the given signature, defined by the signature
+  // key.
+  PyObject* Prepare(std::string signature_key);
+
+  // Allocates the primary subgraph's tensors with the given input shapes.
+  PyObject* Prepare(PyObject* input_shapes);
+
+  // Allocates the tensors of the given signature with the given input
+  // shapes, defined by the signature key.
+  PyObject* Prepare(PyObject* input_shapes, std::string signature_key);
+
+  // Sets the given input tensors to the primary subgraph.
+  PyObject* FeedTensor(PyObject* input_value);
+
+  // Sets the given input tensor to the given signature, defined by the
+  // signature key.
+  PyObject* FeedTensor(PyObject* input_value, std::string signature_key);
+
+  // Allows quantizing only the operator that produces the tensor.
+  PyObject* QuantizeModel(int input_py_type, int output_py_type,
+                          bool allow_float, int activations_py_type,
+                          int bias_py_type);
+
+  // Allows quantizing only the operator that produces the tensor with name
+  // operator_output_name. (This can be used to help debug.).
+  // TODO(suharshs): Allow providing multiple names.
+  PyObject* QuantizeModel(int input_py_type, int output_py_type,
+                          bool allow_float, const char* operator_output_name);
+
+  // Disables per-channel quantization, can be used to produce smaller
+  // models but may cause accuracy issues.
+  PyObject* QuantizeModel(
+      int input_py_type, int output_py_type, bool allow_float,
+      int activations_py_type, int bias_py_type, bool disable_per_channel,
+      bool disable_per_channel_quantization_for_dense_layers);
+
+  // Writes the in-memory calibration results to the model flatbuffer. The
+  // produced model is as same as the original input model, but the min/max
+  // in the quantization field.
+  PyObject* Calibrate();
+
+ private:
+  // CalibrationWrapper is not copyable or assignable. We avoid the use of
+  // CalibrationWrapper() = delete here for SWIG compatibility.
+  CalibrationWrapper(
+      std::unique_ptr<Interpreter> interpreter,
+      std::unique_ptr<ops::builtin::BuiltinOpResolver> resolver,
+      std::unique_ptr<interpreter_wrapper::PythonErrorReporter> error_reporter,
+      std::unique_ptr<impl::FlatBufferModel> model,
+      std::unique_ptr<optimize::calibration::CalibrationReader> reader,
+      std::unique_ptr<std::string> model_str_);
+
+  CalibrationWrapper(const CalibrationWrapper& rhs);
+
+  PyObject* SetTensor(int index, PyObject* value);
+  PyObject* SetTensor(int index, PyObject* value, std::string signature_key);
+
+  std::unique_ptr<Interpreter> interpreter_;
+  std::unique_ptr<interpreter_wrapper::PythonErrorReporter> error_reporter_;
+  std::unique_ptr<ops::builtin::BuiltinOpResolver> resolver_;
+  std::unique_ptr<impl::FlatBufferModel> model_;
+  std::unique_ptr<optimize::calibration::CalibrationReader> reader_;
+  std::unique_ptr<std::string> model_str_;
+};
+
+}  // namespace calibration_wrapper
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_PYTHON_OPTIMIZE_CALIBRATION_WRAPPER_H_

tensorflow/lite/python/optimize/calibration_wrapper_pybind11.cc

@@ -0,0 +1,102 @@
+/* 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/optimize/calibration_wrapper.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+
+namespace py = pybind11;
+using tflite::calibration_wrapper::AddIntermediateTensors;
+using tflite::calibration_wrapper::CalibrationWrapper;
+
+PYBIND11_MODULE(_pywrap_tensorflow_lite_calibration_wrapper, m) {
+  m.doc() = R"pbdoc(
+    _pywrap_tensorflow_lite_calibration_wrapper
+    -----
+  )pbdoc";
+  m.def("AddIntermediateTensors", [](py::handle& data) {
+    return tensorflow::PyoOrThrow(AddIntermediateTensors(data.ptr()));
+  });
+  py::class_<CalibrationWrapper>(m, "CalibrationWrapper")
+      .def(py::init([](py::handle& data,
+                       const std::vector<std::string>& registerers_by_name,
+                       const std::vector<std::function<void(uintptr_t)>>&
+                           registerers_by_func) {
+        std::string error;
+        auto* wrapper = ::CalibrationWrapper::CreateWrapperCPPFromBuffer(
+            data.ptr(), registerers_by_name, registerers_by_func, &error);
+        if (!wrapper) {
+          throw std::invalid_argument(error);  // throws ValueError in Python
+        }
+        return wrapper;
+      }))
+      .def("Prepare",
+           [](CalibrationWrapper& self, py::handle& input_shapes,
+              std::string signature_key) {
+             return tensorflow::PyoOrThrow(
+                 self.Prepare(input_shapes.ptr(), signature_key));
+           })
+      .def("Prepare",
+           [](CalibrationWrapper& self, py::handle& input_shapes) {
+             return tensorflow::PyoOrThrow(self.Prepare(input_shapes.ptr()));
+           })
+      .def("Prepare",
+           [](CalibrationWrapper& self, std::string signature_key) {
+             return tensorflow::PyoOrThrow(self.Prepare(signature_key));
+           })
+      .def("Prepare",
+           [](CalibrationWrapper& self) {
+             return tensorflow::PyoOrThrow(self.Prepare());
+           })
+      .def("FeedTensor",
+           [](CalibrationWrapper& self, py::handle& input_value,
+              std::string signature_key) {
+             return tensorflow::PyoOrThrow(
+                 self.FeedTensor(input_value.ptr(), signature_key));
+           })
+      .def("FeedTensor",
+           [](CalibrationWrapper& self, py::handle& input_value) {
+             return tensorflow::PyoOrThrow(self.FeedTensor(input_value.ptr()));
+           })
+      .def("QuantizeModel",
+           [](CalibrationWrapper& self, int input_py_type, int output_py_type,
+              bool allow_float, int activations_py_type, int bias_py_type,
+              bool disable_per_channel,
+              bool disable_per_channel_quantization_for_dense_layers) {
+             return tensorflow::PyoOrThrow(self.QuantizeModel(
+                 input_py_type, output_py_type, allow_float,
+                 activations_py_type, bias_py_type, disable_per_channel,
+                 disable_per_channel_quantization_for_dense_layers));
+           })
+      .def("QuantizeModel",
+           [](CalibrationWrapper& self, int input_py_type, int output_py_type,
+              bool allow_float, const char* operator_output_name) {
+             return tensorflow::PyoOrThrow(
+                 self.QuantizeModel(input_py_type, output_py_type, allow_float,
+                                    operator_output_name));
+           })
+      .def("Calibrate", [](CalibrationWrapper& self) {
+        return tensorflow::PyoOrThrow(self.Calibrate());
+      });
+}

tensorflow/lite/python/optimize/calibrator.py

@@ -0,0 +1,259 @@
+# 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.
+# ==============================================================================
+"""Python wrapper for post training quantization with calibration."""
+import numpy as np
+
+from tensorflow.lite.python.convert_phase import Component
+from tensorflow.lite.python.convert_phase import convert_phase
+from tensorflow.lite.python.convert_phase import SubComponent
+from tensorflow.lite.python.interpreter import Interpreter
+from tensorflow.python.framework import dtypes
+from tensorflow.python.util.lazy_loader import LazyLoader
+
+# Lazy load since some of the performance benchmark skylark rules
+# break dependencies. Must use double quotes to match code internal rewrite
+# rule.
+_calibration_wrapper = LazyLoader(
+    "_calibration_wrapper",
+    globals(),
+    (
+        "tensorflow.lite.python.optimize."
+        "_pywrap_tensorflow_lite_calibration_wrapper"
+    ),
+)
+
+
+def add_intermediate_tensors(model_content):
+  """Adds intermediate tensors to fused op if needed."""
+  return _calibration_wrapper.AddIntermediateTensors(model_content)
+
+
+class Calibrator:
+  """Calibrates a floating point model and then quantizes it.
+
+  This is an internal class, not a public interface.
+  """
+
+  def __init__(
+      self,
+      model_content,
+      custom_op_registerers_by_name=None,
+      custom_op_registerers_by_func=None,
+  ):
+    """Constructor.
+
+    Args:
+      model_content: Content of a TF-Lite Flatbuffer file.
+      custom_op_registerers_by_name: List of str (symbol names) that take a
+        pointer to a MutableOpResolver and register custom ops.
+      custom_op_registerers_by_func: List of functions that take a pointer to a
+        MutableOpResolver and register custom ops.
+
+    Raises:
+      ValueError: If the calibrator was unable to open the model.
+    """
+    if not model_content:
+      raise ValueError("`model_content` must be specified.")
+    if custom_op_registerers_by_name is None:
+      custom_op_registerers_by_name = []
+    if custom_op_registerers_by_func is None:
+      custom_op_registerers_by_func = []
+    try:
+      self._calibrator = _calibration_wrapper.CalibrationWrapper(
+          model_content,
+          custom_op_registerers_by_name,
+          custom_op_registerers_by_func,
+      )
+      self._model_content = model_content
+    except Exception as e:
+      raise ValueError("Failed to parse the model: %s." % e)
+    if not self._calibrator:
+      raise ValueError("Failed to parse the model.")
+    self._interpreter = None
+
+  def _create_input_array_from_dict(self, signature_key, inputs):
+    input_array = []
+    signature_runner = self._interpreter.get_signature_runner(signature_key)
+    input_details = sorted(
+        signature_runner.get_input_details().items(),
+        key=lambda item: item[1]["index"],
+    )
+    for input_name, _ in input_details:
+      input_array.append(inputs[input_name])
+    return input_array
+
+  def _feed_tensors(self, dataset_gen, resize_input):
+    """Feed tensors to the calibrator."""
+    initialized = {}
+
+    for sample in dataset_gen():
+      if isinstance(sample, tuple):
+        if not isinstance(sample[1], dict):
+          raise ValueError(
+              "You need to provide either a dictionary with input "
+              "names and values in the second argument in the "
+              "tuple"
+          )
+        # Convert signature based inputs to the tensor index based data.
+        if self._interpreter is None:
+          self._interpreter = Interpreter(model_content=self._model_content)
+        signature_key = sample[0]
+        input_array = self._create_input_array_from_dict(
+            signature_key, sample[1]
+        )
+      elif isinstance(sample, dict):
+        # Convert signature based inputs to the tensor index based data.
+        if self._interpreter is None:
+          self._interpreter = Interpreter(model_content=self._model_content)
+        signature_key = None
+        input_array = self._create_input_array_from_dict(None, sample)
+      elif isinstance(sample, list):
+        signature_key = None
+        input_array = sample
+      else:
+        raise ValueError(
+            "You need to provide either a dictionary with input "
+            "names and values, a tuple with signature key and a "
+            "dictionary with input names and values, or an array "
+            "with input values in the order of input tensors of "
+            "the graph in the representative_dataset function. "
+            "Unsupported value from dataset: {}.".format(sample)
+        )
+
+      if signature_key not in initialized:
+        initialized[signature_key] = True
+        if resize_input:
+          if signature_key is not None:
+            self._calibrator.Prepare(
+                [list(s.shape) for s in input_array], signature_key
+            )
+          else:
+            self._calibrator.Prepare([list(s.shape) for s in input_array])
+        else:
+          if signature_key is not None:
+            self._calibrator.Prepare(signature_key)
+          else:
+            self._calibrator.Prepare()
+      if signature_key is not None:
+        self._calibrator.FeedTensor(input_array, signature_key)
+      else:
+        self._calibrator.FeedTensor(input_array)
+
+  @convert_phase(
+      Component.OPTIMIZE_TFLITE_MODEL,
+      SubComponent.QUANTIZE_USING_DEPRECATED_QUANTIZER,
+  )
+  def calibrate_and_quantize(
+      self,
+      dataset_gen,
+      input_type,
+      output_type,
+      allow_float,
+      activations_type=dtypes.int8,
+      bias_type=dtypes.int32,
+      resize_input=True,
+      disable_per_channel=False,
+      disable_per_channel_quantization_for_dense_layers=False,
+  ):
+    """Calibrates the model with specified generator and then quantizes it.
+
+    The input shapes of the calibrator are resized with the calibration data if
+    `resize_input` is set.
+
+    Returns:
+      A quantized model.
+
+    Args:
+      dataset_gen: A generator that generates calibration samples.
+      input_type: A tf.dtype representing the desired real-value input type.
+      output_type: A tf.dtype representing the desired real-value output type.
+      allow_float: A boolean. False if the resulting model cannot perform float
+        computation, useful when targeting an integer-only backend. If False, an
+        error will be thrown if an operation cannot be quantized, otherwise the
+        model will fallback to float ops.
+      activations_type: A tf.dtype representing the desired type for
+        activations.
+      bias_type: A tf.dtype representing the desired type for bias.
+      resize_input: A boolean. True if the shape of the sample data is different
+        from the input.
+      disable_per_channel: A boolean. True if disabling per-channel
+        quantization.
+      disable_per_channel_quantization_for_dense_layers: A boolean. True if
+        disabling per-channel quantization only in Dense layers.
+    """
+    self._feed_tensors(dataset_gen, resize_input)
+    return self._calibrator.QuantizeModel(
+        np.dtype(input_type.as_numpy_dtype()).num,
+        np.dtype(output_type.as_numpy_dtype()).num,
+        allow_float,
+        np.dtype(activations_type.as_numpy_dtype()).num,
+        np.dtype(bias_type.as_numpy_dtype()).num,
+        disable_per_channel,
+        disable_per_channel_quantization_for_dense_layers,
+    )
+
+  @convert_phase(
+      Component.OPTIMIZE_TFLITE_MODEL,
+      SubComponent.QUANTIZE_USING_DEPRECATED_QUANTIZER,
+  )
+  def calibrate_and_quantize_single(
+      self,
+      dataset_gen,
+      input_type,
+      output_type,
+      allow_float,
+      op_output_name,
+      resize_input=True,
+  ):
+    """Calibrates the model with specified generator and then quantizes it.
+
+    Only the single op with output op_output_name will be quantized.
+    The input shapes of the calibrator are resized with the calibration data.
+
+    Returns:
+      A quantized model.
+
+    Args:
+      dataset_gen: A generator that generates calibration samples.
+      input_type: A tf.dtype representing the desired real-value input type.
+      output_type: A tf.dtype representing the desired real-value output type.
+      allow_float: A boolean. False if the resulting model cannot perform float
+        computation, useful when targeting an integer-only backend. If False, an
+        error will be thrown if an operation cannot be quantized, otherwise the
+        model will fallback to float ops.
+      op_output_name: A string, only this op will be quantized.
+      resize_input: A boolean. True if the shape of the sample data is different
+        from the input.
+    """
+    self._feed_tensors(dataset_gen, resize_input)
+    return self._calibrator.QuantizeModel(
+        np.dtype(input_type.as_numpy_dtype()).num,
+        np.dtype(output_type.as_numpy_dtype()).num,
+        allow_float,
+        op_output_name,
+    )
+
+  @convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.CALIBRATE)
+  def calibrate(self, dataset_gen):
+    """Calibrates the model with specified generator.
+
+    Returns:
+      A model with min and max calibration stats.
+
+    Args:
+      dataset_gen: A generator that generates calibration samples.
+    """
+    self._feed_tensors(dataset_gen, resize_input=True)
+    return self._calibrator.Calibrate()

tensorflow/lite/python/optimize/calibrator_test.py

@@ -0,0 +1,311 @@
+# 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.
+# ==============================================================================
+"""Tests for Calibrator."""
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.lite.python import lite
+from tensorflow.lite.python import schema_py_generated as schema_fb
+from tensorflow.lite.python.optimize import calibrator as _calibrator
+from tensorflow.lite.tools import flatbuffer_utils
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import resource_loader
+from tensorflow.python.platform import test
+
+
+def _uses_buffer_offset(model: schema_fb.ModelT) -> bool:
+  """Determines whether the model is using an offset buffer.
+
+  Args:
+    model: A TFLite model.
+
+  Returns:
+    True iff the model is using offset buffers. Offset buffers are enabled by
+    the flag `_experimental_use_buffer_offset`.
+  """
+  if not model.metadata:
+    return False
+
+  return any(
+      map(
+          lambda metadata: metadata.name.decode('utf-8') == 'buffer_location',
+          model.metadata,
+      )
+  )
+
+
+class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      # Activation type Int8
+      ('UseActivationTypeInt8', dtypes.int8),
+      # Activation type Int16
+      ('UseActivationTypeInt16', dtypes.int16),
+  )
+  def test_calibration_with_quantization(self, activations_type):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    # Input generator for the model.
+    def input_gen():
+      for _ in range(10):
+        yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
+
+    quantized_model = quantizer.calibrate_and_quantize(
+        input_gen, dtypes.float32, dtypes.float32, False, activations_type
+    )
+    self.assertIsNotNone(quantized_model)
+
+  @parameterized.named_parameters(
+      # Activation type Int8
+      ('UseActivationTypeInt8', dtypes.int8),
+      # Activation type Int16
+      ('UseActivationTypeInt16', dtypes.int16),
+  )
+  def test_calibration_with_quantization_allow_float(self, activations_type):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    # Input generator for the model.
+    def input_gen():
+      for _ in range(10):
+        yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
+
+    quantized_model = quantizer.calibrate_and_quantize(
+        input_gen, dtypes.float32, dtypes.float32, True, activations_type
+    )
+    self.assertIsNotNone(quantized_model)
+
+  def test_calibration_with_quantization_single_op(self):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    # Input generator for the model.
+    def input_gen():
+      for _ in range(10):
+        yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
+
+    quantized_model = quantizer.calibrate_and_quantize_single(
+        input_gen, dtypes.float32, dtypes.float32, True, 'conv2d_8/BiasAdd'
+    )
+    self.assertIsNotNone(quantized_model)
+
+  def test_calibration_with_string_input(self):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/string_input_flex_model.bin'
+    )
+    with open(model_path, 'rb') as fp:
+      model_with_string_input = fp.read()
+    quantizer = _calibrator.Calibrator(model_with_string_input)
+
+    # Input generator for the model.
+    def input_gen():
+      for i in range(10):
+        yield [np.array('Test' + str(i))]
+
+    quantized_model = quantizer.calibrate_and_quantize_single(
+        input_gen, dtypes.float32, dtypes.float32, True, 'Identity'
+    )
+    self.assertIsNotNone(quantized_model)
+
+  @parameterized.named_parameters(
+      # Activation type Int8
+      ('UseActivationTypeInt8 - EnableMlirQuantizer', dtypes.int8),
+      # Activation type Int16
+      ('UseActivationTypeInt16 - DisableEnableMlirQuantizer', dtypes.int16),
+  )
+  def test_calibration_with_quantization_multiple_inputs(
+      self, activations_type
+  ):
+    # Load multi add model from test data.
+    # This model has 4 inputs of size (1, 8, 8, 3).
+    model_path = resource_loader.get_path_to_datafile(
+        '../../testdata/multi_add.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    # Input generator for the model.
+    def input_gen():
+      for _ in range(10):
+        yield [np.ones(shape=(1, 8, 8, 3), dtype=np.float32) for _ in range(4)]
+
+    quantized_model = quantizer.calibrate_and_quantize(
+        input_gen, dtypes.float32, dtypes.float32, False, activations_type
+    )
+    self.assertIsNotNone(quantized_model)
+
+  def test_invalid_model_buffer(self):
+    float_model = b'\0' * 100
+    with self.assertRaisesRegex(ValueError, 'Failed to parse the model'):
+      _calibrator.Calibrator(float_model)
+
+  # TODO(fengliuai): enable mlir quantizer
+  def test_empty_calibrator_gen(self):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    def empty_input_gen():
+      for i in ():
+        yield i
+
+    with self.assertRaises(RuntimeError):
+      quantizer.calibrate_and_quantize(
+          empty_input_gen, dtypes.float32, dtypes.float32, False
+      )
+
+  def test_invalid_shape_calibrator_gen(self):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    # Input generator with incorrect shape.
+    def input_gen():
+      for _ in range(10):
+        yield [np.ones(shape=(1, 2, 2, 3), dtype=np.float32)]
+
+    with self.assertRaisesRegex(ValueError, 'Size mismatch'):
+      quantizer.calibrate_and_quantize(
+          input_gen,
+          dtypes.float32,
+          dtypes.float32,
+          False,
+          activations_type=dtypes.int8,
+          bias_type=dtypes.int32,
+          resize_input=False,
+      )
+
+  def test_invalid_type_calibrator_gen(self):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    # Input generator with incorrect type.
+    def input_gen():
+      for _ in range(10):
+        yield [np.ones(shape=(1, 5, 5, 3), dtype=np.int32)]
+
+    with self.assertRaises(ValueError):
+      quantizer.calibrate_and_quantize(
+          input_gen, dtypes.float32, dtypes.float32, False, dtypes.int8
+      )
+
+  def test_calibration(self):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    float_model = open(model_path, 'rb').read()
+    quantizer = _calibrator.Calibrator(float_model)
+
+    # Input generator for the model.
+    def input_gen():
+      for _ in range(10):
+        yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
+
+    quantized_model = quantizer.calibrate(input_gen)
+    self.assertIsNotNone(quantized_model)
+
+  def test_add_intermediate_tensors(self):
+    model_path = resource_loader.get_path_to_datafile(
+        'test_data/mobilenet_like_model.bin'
+    )
+    model = open(model_path, 'rb').read()
+    added_model = _calibrator.add_intermediate_tensors(model)
+    self.assertIsNotNone(added_model)
+
+  def test_calibrate_model_with_offset_buffer(self):
+    # Define a simple model to run calibration with.
+    class MatMulModel(tf.Module):
+
+      def __init__(self):
+        # Use ones for predictable calibration results.
+        self.filter = np.ones((4, 3)).astype(np.float32)
+
+      @tf.function(
+          input_signature=[tf.TensorSpec(shape=(1, 4), dtype=dtypes.float32)]
+      )
+      def __call__(self, input_tensor: tf.Tensor) -> tf.Tensor:
+        output_tensor = tf.linalg.matmul(input_tensor, self.filter)
+        return {'output': output_tensor}
+
+    model = MatMulModel()
+    saved_model_path = self.create_tempdir().full_path
+    tf.saved_model.save(model, saved_model_path)
+
+    converter = lite.TFLiteConverter.from_saved_model(saved_model_path)
+    # Enable the use of buffer offsets.
+    # pylint: disable=protected-access
+    converter._experimental_use_buffer_offset = True
+    # pylint: enable=protected-access
+    converter.exclude_conversion_metadata = True
+
+    model_serialized = converter.convert()
+
+    model = flatbuffer_utils.convert_bytearray_to_object(model_serialized)
+    self.assertTrue(_uses_buffer_offset(model))
+
+    quantizer = _calibrator.Calibrator(model_serialized)
+
+    # Input generator for the model.
+    def input_gen():
+      for _ in range(2):
+        yield [np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float32)]
+
+    calibrated_model_serialized = quantizer.calibrate(input_gen)
+    self.assertIsNotNone(calibrated_model_serialized)
+
+    calibrated_model = flatbuffer_utils.convert_bytearray_to_object(
+        calibrated_model_serialized
+    )
+    self.assertTrue(_uses_buffer_offset(calibrated_model))
+
+    # Confirm that the tensors are correctly calibrated.
+    subgraph = calibrated_model.subgraphs[0]
+
+    matmul_input_tensor = subgraph.tensors[0]
+    self.assertAllClose(matmul_input_tensor.quantization.min, [1.0])
+    self.assertAllClose(matmul_input_tensor.quantization.max, [1.0])
+
+    matmul_filter_tensor = subgraph.tensors[1]
+    self.assertAllClose(matmul_filter_tensor.quantization.min, [1.0])
+    self.assertAllClose(matmul_filter_tensor.quantization.max, [1.0])
+
+    # The matmul is performed with all ones so the output is expected to be 4s.
+    matmul_output_tensor = subgraph.tensors[2]
+    self.assertAllClose(matmul_output_tensor.quantization.min, [4.0])
+    self.assertAllClose(matmul_output_tensor.quantization.max, [4.0])
+
+
+if __name__ == '__main__':
+  test.main()