chore: import upstream snapshot with attribution

tensorflow/lite/tools/optimize/BUILD

@@ -0,0 +1,341 @@
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/lite:special_rules.bzl", "tflite_portable_test_suite")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "reduced_precision_support",
+    srcs = [],
+    hdrs = [
+        "reduced_precision_support.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/lite/tools/optimize:reduced_precision_metadata",
+    ],
+)
+
+tf_cc_test(
+    name = "reduced_precision_support_test",
+    srcs = ["reduced_precision_support_test.cc"],
+    tags = [
+        "tflite_not_portable_android",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":reduced_precision_support",
+        "//tensorflow/core/platform:platform_port",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "//tensorflow/lite/testing:util",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "modify_model_interface",
+    srcs = ["modify_model_interface.cc"],
+    hdrs = ["modify_model_interface.h"],
+    deps = [
+        ":model_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal:compatibility",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/memory",
+        "@flatbuffers",
+    ],
+)
+
+tf_cc_test(
+    name = "modify_model_interface_test",
+    srcs = ["modify_model_interface_test.cc"],
+    tags = [
+        "tflite_not_portable_android",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":modify_model_interface",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "@com_google_absl//absl/memory",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+cc_binary(
+    name = "modify_model_interface_main",
+    srcs = ["modify_model_interface_main.cc"],
+    deps = [
+        ":modify_model_interface",
+        ":quantize_model",
+    ],
+)
+
+cc_library(
+    name = "quantization_wrapper_utils",
+    srcs = ["quantization_wrapper_utils.cc"],
+    hdrs = ["quantization_wrapper_utils.h"],
+    deps = [
+        ":operator_property",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@flatbuffers",
+    ],
+)
+
+tf_cc_test(
+    name = "quantization_wrapper_utils_test",
+    srcs = ["quantization_wrapper_utils_test.cc"],
+    tags = [
+        "tflite_not_portable_android",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":quantization_wrapper_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest",
+        "@flatbuffers",
+    ],
+)
+
+tf_cc_test(
+    name = "quantization_wrapper_utils_custom_test",
+    srcs = [
+        "quantization_wrapper_utils.cc",
+        "quantization_wrapper_utils.h",
+        "quantization_wrapper_utils_custom_test.cc",
+    ],
+    defines = [
+        "TFLITE_CUSTOM_LSTM",
+    ],
+    tags = [
+        "tflite_not_portable_android",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":operator_property",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "quantization_wrapper",
+    srcs = ["quantization_wrapper.cc"],
+    hdrs = ["quantization_wrapper.h"],
+    deps = [
+        ":quantization_wrapper_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/tools/optimize:quantize_model",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "quantization_utils",
+    srcs = ["quantization_utils.cc"],
+    hdrs = ["quantization_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":model_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal:cppmath",
+        "//tensorflow/lite/kernels/internal:quantization_util",
+        "//tensorflow/lite/kernels/internal:tensor_utils",
+        "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+cc_library(
+    name = "model_utils",
+    srcs = ["model_utils.cc"],
+    hdrs = ["model_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":operator_property",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/kernels/internal:tensor_utils",
+        "//tensorflow/lite/kernels/internal:types",
+        "//tensorflow/lite/schema:schema_conversion_utils",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+tf_cc_test(
+    name = "model_utils_test",
+    srcs = ["model_utils_test.cc"],
+    tags = [
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":model_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest_main",
+        "@flatbuffers",
+    ],
+)
+
+alias(
+    name = "operator_property",
+    actual = "//tensorflow/compiler/mlir/lite/tools/optimize:operator_property",
+)
+
+tf_cc_test(
+    name = "quantization_utils_test",
+    srcs = ["quantization_utils_test.cc"],
+    args = [
+        "--test_model_file=$(location //tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin)",
+    ],
+    data = [
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin",
+    ],
+    tags = [
+        "tflite_not_portable_android",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":quantization_utils",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:test_util",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "//tensorflow/lite/testing:util",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest",
+        "@eigen_archive//:eigen3",
+        "@flatbuffers",
+    ],
+)
+
+cc_library(
+    name = "quantize_model",
+    srcs = ["quantize_model.cc"],
+    hdrs = ["quantize_model.h"],
+    deps = [
+        ":model_utils",
+        ":operator_property",
+        ":quantization_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/kernels/internal:cppmath",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers",
+    ],
+)
+
+tf_cc_test(
+    name = "quantize_model_test",
+    srcs = ["quantize_model_test.cc"],
+    args = [
+        "--test_model_file=$(location //tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin)",
+    ],
+    data = [
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/add_with_const_input.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/argmax.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/broadcast_to.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/concat.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/fc.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/fc_qat.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/gather_nd.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_calibrated.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_calibrated2.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_quantized.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/lstm_quantized2.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/maximum.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/minimum.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/mixed.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/mixed16x8.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/multi_input_add_reshape.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/pack.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/resource_vars_calibrated.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_avg_pool_min_minus_5_max_plus_5.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_no_bias.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_0_max_plus_10.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_conv_weights_min_minus_127_max_plus_127.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/single_softmax_min_minus_5_max_plus_5.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/split.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/svdf_calibrated.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/svdf_quantized.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/transpose.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/unidirectional_sequence_lstm_calibrated.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/unidirectional_sequence_lstm_quantized.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/unpack.bin",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:testdata/where.bin",
+    ],
+    tags = [
+        "tflite_not_portable_android",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":quantize_model",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:test_util",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/schema:schema_utils",
+        "//tensorflow/lite/testing:util",
+        "@com_google_googletest//:gtest",
+        "@flatbuffers",
+    ],
+)
+
+tflite_portable_test_suite()

tensorflow/lite/tools/optimize/calibration/BUILD

@@ -0,0 +1,197 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts")
+load("//tensorflow/lite:special_rules.bzl", "tflite_portable_test_suite")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "builtin_logging_op",
+    srcs = ["builtin_logging_ops/lstm.cc"],
+    hdrs = ["builtin_logging_ops/lstm.h"],
+    copts = tflite_copts(),
+    deps = [
+        ":calibration_logger",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework_stable",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels:lstm_shared",
+        "//tensorflow/lite/kernels:op_macros",
+        "//tensorflow/lite/kernels/internal:reference",
+        "//tensorflow/lite/kernels/internal:tensor_utils",
+        "@ruy//ruy/profiler:instrumentation",
+    ],
+)
+
+cc_library(
+    name = "custom_logging_op",
+    srcs = ["custom_logging_ops/lstm.cc"],
+    hdrs = ["custom_logging_ops/lstm.h"],
+    copts = tflite_copts(),
+    deps = [
+        ":calibration_logger",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework_stable",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:kernel_util",
+        "//tensorflow/lite/kernels:lstm_shared",
+        "//tensorflow/lite/kernels:op_macros",
+        "//tensorflow/lite/kernels/internal:reference",
+        "//tensorflow/lite/kernels/internal:tensor_utils",
+    ],
+)
+
+cc_library(
+    name = "calibrator_lib",
+    srcs = ["calibrator.cc"],
+    hdrs = ["calibrator.h"],
+    copts = tflite_copts(),
+    deps = [
+        ":builtin_logging_op",
+        ":calibration_common",
+        ":calibration_logger",
+        ":calibration_reader",
+        ":custom_logging_op",
+        ":logging_op",
+        ":logging_op_resolver",
+        "//tensorflow/compiler/mlir/lite:allocation",
+        "//tensorflow/compiler/mlir/lite/schema:schema_utils",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@flatbuffers",
+    ],
+)
+
+tf_cc_test(
+    name = "calibrator_test",
+    srcs = ["calibrator_test.cc"],
+    args = [
+        "--test_model_file=$(location //tensorflow/lite:testdata/multi_add.bin)",
+    ],
+    data = [
+        "//tensorflow/lite:testdata/call_once_mul.bin",
+        "//tensorflow/lite:testdata/custom_lstm.bin",
+        "//tensorflow/lite:testdata/lstm.bin",
+        "//tensorflow/lite:testdata/multi_add.bin",
+        "//tensorflow/lite:testdata/multi_subgraphs_while.bin",
+        "//tensorflow/lite:testdata/unidirectional_sequence_lstm.bin",
+    ],
+    tags = [
+        "tflite_not_portable_android",
+        "tflite_not_portable_ios",
+    ],
+    deps = [
+        ":calibration_reader",
+        ":calibrator_lib",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/kernels:builtin_ops",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
+cc_library(
+    name = "logging_op_resolver",
+    srcs = ["logging_op_resolver.cc"],
+    hdrs = ["logging_op_resolver.h"],
+    copts = tflite_copts(),
+    deps = [
+        ":calibration_common",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_test(
+    name = "logging_op_resolver_test",
+    srcs = ["logging_op_resolver_test.cc"],
+    deps = [
+        ":calibration_common",
+        ":logging_op_resolver",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "calibration_reader",
+    srcs = ["calibration_reader.cc"],
+    hdrs = ["calibration_reader.h"],
+    copts = tflite_copts(),
+    deps = [
+        ":calibration_logger",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:common",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:c_api_types",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
+cc_library(
+    name = "calibration_logger",
+    srcs = ["calibration_logger.cc"],
+    hdrs = ["calibration_logger.h"],
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/c:c_api_types",
+        "//tensorflow/lite/core/api",
+        "//tensorflow/lite/core/c:common",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
+cc_library(
+    name = "calibration_common",
+    hdrs = ["calibration_common.h"],
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/lite:framework",
+    ],
+)
+
+cc_library(
+    name = "logging_op",
+    hdrs = ["logging_op.h"],
+    copts = tflite_copts(),
+    deps = [
+        ":calibration_logger",
+        "//tensorflow/lite/core/c:common",
+    ],
+)
+
+tflite_portable_test_suite()

tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.cc

@@ -0,0 +1,657 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.h"
+
+#include <algorithm>
+#include <cstdio>
+#include <vector>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/lstm_shared.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_logger.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+namespace builtin {
+
+namespace {
+
+inline void CalculateLstmGateFloat(
+    const float* input, const float* input_to_gate_weights,
+    const float* aux_input, const float* aux_input_to_gate_weights,
+    const float* output_state, const float* recurrent_to_gate_weights,
+    const float* cell_state, const float* cell_to_gate_weights,
+    const float* layer_norm_coefficients, const float* gate_bias,
+    const int n_batch, const int n_input, const int n_aux_input,
+    const int n_output, const int n_cell,
+    const TfLiteFusedActivation activation, float* gate,
+    const bool is_input_all_zeros, const bool is_aux_input_all_zeros,
+    Logger* logger, int intermediate_tensor_index, const int subgraph_index,
+    ErrorReporter* error_reporter) {
+  const bool use_peephole = (cell_to_gate_weights != nullptr);
+  const bool use_layer_norm = (layer_norm_coefficients != nullptr);
+
+  // Initialize scratch buffers with bias for regular lstm or initialize with
+  // zero for layer norm lstm.
+  if (use_layer_norm) {
+    std::fill_n(gate, n_cell * n_batch, 0.0f);
+  } else {
+    tensor_utils::VectorBatchVectorAssign(gate_bias, n_cell, n_batch, gate);
+  }
+  // For each batch and cell: compute input_weight * input.
+  // Skip if input is all zeros.
+  if (!is_input_all_zeros) {
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        input_to_gate_weights, n_cell, n_input, input, n_batch, gate);
+  }
+  // For each batch and cell: compute aux_input_weight * aux_input.
+  // Skip if auxiliary input is not available or all zeros.
+  if (!is_aux_input_all_zeros) {
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(aux_input_to_gate_weights,
+                                                      n_cell, n_aux_input,
+                                                      aux_input, n_batch, gate);
+  }
+  // For each batch and cell: compute recurrent_weight * output_state.
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+      recurrent_to_gate_weights, n_cell, n_output, output_state, n_batch, gate);
+  // For each batch and cell: compute cell_weight .* cell_state (peephole LSTM)
+  if (use_peephole) {
+    tensor_utils::VectorBatchVectorCwiseProductAccumulate(
+        cell_to_gate_weights, n_cell, cell_state, n_batch, gate);
+  }
+  // Do layer normalization (if layer norm LSTM)
+  if (use_layer_norm) {
+    logger->LogTensorValue(subgraph_index, intermediate_tensor_index, gate,
+                           n_cell * n_batch, error_reporter);
+
+    tensor_utils::MeanStddevNormalization(gate, gate, n_cell, n_batch);
+    tensor_utils::VectorBatchVectorCwiseProduct(layer_norm_coefficients, n_cell,
+                                                gate, n_batch, gate);
+    tensor_utils::VectorBatchVectorAdd(gate_bias, n_cell, n_batch, gate);
+  }
+  // Apply activation
+  tensor_utils::ApplyActivationToVector(gate, n_batch * n_cell, activation,
+                                        gate);
+}
+
+// TODO(b/159066113): This is the exact same function as UpdateLstmCellFloat in
+// kernels/lstm_eval.cc, make that public and remove this.
+void UpdateLstmCellFloat(int n_batch, int n_cell, float* cell_state,
+                         const float* input_gate, float* forget_gate,
+                         const float* cell_gate, bool use_cifg, float clip) {
+  tensor_utils::VectorVectorCwiseProduct(forget_gate, cell_state,
+                                         n_batch * n_cell, cell_state);
+
+  if (use_cifg) {
+    // With CIFG, input_gate = 1-forget_gate. Use the forget_gate array as
+    // scratch, as input_gate array is not allocated in this case. (Be careful
+    // not to write to the scratch before reading the forget gate data.)
+    float* scratch = forget_gate;
+    tensor_utils::Sub1Vector(forget_gate, n_batch * n_cell, scratch);
+    tensor_utils::VectorVectorCwiseProductAccumulate(
+        cell_gate, scratch, n_batch * n_cell, cell_state);
+  } else {
+    tensor_utils::VectorVectorCwiseProductAccumulate(
+        cell_gate, input_gate, n_batch * n_cell, cell_state);
+  }
+  if (clip > 0.0f) {
+    tensor_utils::CwiseClipping(cell_state, n_batch * n_cell, clip);
+  }
+}
+
+void CalculateLstmOutputCalibration(
+    int n_batch, int n_cell, int n_output, const float* cell_state,
+    const float* output_gate, TfLiteFusedActivation activation,
+    const float* projection_weights, const float* projection_bias,
+    const float proj_clip, float* output_state, float* scratch, Logger* logger,
+    int intermediate_tensor_index, const int subgraph_index,
+    ErrorReporter* error_reporter) {
+  tensor_utils::ApplyActivationToVector(cell_state, n_batch * n_cell,
+                                        activation, scratch);
+  tensor_utils::VectorVectorCwiseProduct(output_gate, scratch, n_batch * n_cell,
+                                         scratch);
+
+  logger->LogTensorValue(subgraph_index, intermediate_tensor_index, scratch,
+                         n_cell * n_batch, error_reporter);
+
+  const bool use_projection = (projection_weights != nullptr);
+  const bool use_projection_bias = (projection_bias != nullptr);
+
+  if (use_projection) {
+    if (use_projection_bias) {
+      tensor_utils::VectorBatchVectorAssign(projection_bias, n_output, n_batch,
+                                            output_state);
+    } else {
+      std::fill_n(output_state, n_batch * n_output, 0.0f);
+    }
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        projection_weights, n_output, n_cell, scratch, n_batch, output_state);
+    if (proj_clip > 0.0f) {
+      tensor_utils::CwiseClipping(output_state, n_batch * n_output, proj_clip);
+    }
+  } else {
+    std::copy_n(scratch, n_batch * n_output, output_state);
+  }
+}
+
+inline void LstmStepCalibration(
+    const float* input_ptr, const float* input_to_input_weights_ptr,
+    const float* input_to_forget_weights_ptr,
+    const float* input_to_cell_weights_ptr,
+    const float* input_to_output_weights_ptr, const float* aux_input_ptr,
+    const float* aux_input_to_input_weights_ptr,
+    const float* aux_input_to_forget_weights_ptr,
+    const float* aux_input_to_cell_weights_ptr,
+    const float* aux_input_to_output_weights_ptr,
+    const float* recurrent_to_input_weights_ptr,
+    const float* recurrent_to_forget_weights_ptr,
+    const float* recurrent_to_cell_weights_ptr,
+    const float* recurrent_to_output_weights_ptr,
+    const float* cell_to_input_weights_ptr,
+    const float* cell_to_forget_weights_ptr,
+    const float* cell_to_output_weights_ptr,
+    const float* input_layer_norm_coefficients_ptr,
+    const float* forget_layer_norm_coefficients_ptr,
+    const float* cell_layer_norm_coefficients_ptr,
+    const float* output_layer_norm_coefficients_ptr,
+    const float* input_gate_bias_ptr, const float* forget_gate_bias_ptr,
+    const float* cell_gate_bias_ptr, const float* output_gate_bias_ptr,
+    const float* projection_weights_ptr, const float* projection_bias_ptr,
+    const TfLiteLSTMParams* params, int n_batch, int n_cell, int n_input,
+    int n_aux_input, int n_output, int output_batch_leading_dim,
+    float* output_state_ptr, float* cell_state_ptr, float* scratch0,
+    float* scratch1, float* scratch2, float* scratch3, float* output_ptr,
+    Logger* logger, const std::vector<int>& intermediate_tensor_indexes,
+    const int subgraph_index, ErrorReporter* error_reporter) {
+  ruy::profiler::ScopeLabel label("LstmStepCalibration");
+  // Since we have already checked that weights are all there or none, we can
+  // check the existence of only one to the get the condition.
+  const bool use_cifg = (input_to_input_weights_ptr == nullptr);
+
+  // Make named scratch buffers.
+  float* input_gate_scratch = scratch0;
+  float* forget_gate_scratch = scratch1;
+  float* cell_gate_scratch = scratch2;
+  float* output_gate_scratch = scratch3;
+
+  // Check if inputs are all zeros so we can skip some computations.
+  const bool is_input_all_zeros =
+      tensor_utils::IsZeroVector(input_ptr, n_batch * n_input);
+  const bool is_aux_input_all_zeros =
+      (aux_input_ptr == nullptr ||
+       tensor_utils::IsZeroVector(aux_input_ptr, n_batch * n_aux_input));
+  if (!use_cifg) {
+    // Calculate the input gate. (If not CIFG.)
+    CalculateLstmGateFloat(
+        input_ptr, input_to_input_weights_ptr, aux_input_ptr,
+        aux_input_to_input_weights_ptr, output_state_ptr,
+        recurrent_to_input_weights_ptr, cell_state_ptr,
+        cell_to_input_weights_ptr, input_layer_norm_coefficients_ptr,
+        input_gate_bias_ptr, n_batch, n_input, n_aux_input, n_output, n_cell,
+        /*activation=*/kTfLiteActSigmoid, input_gate_scratch,
+        is_input_all_zeros, is_aux_input_all_zeros, logger,
+        intermediate_tensor_indexes[0], subgraph_index, error_reporter);
+  }
+  // Calculate the forget gate.
+  CalculateLstmGateFloat(
+      input_ptr, input_to_forget_weights_ptr, aux_input_ptr,
+      aux_input_to_forget_weights_ptr, output_state_ptr,
+      recurrent_to_forget_weights_ptr, cell_state_ptr,
+      cell_to_forget_weights_ptr, forget_layer_norm_coefficients_ptr,
+      forget_gate_bias_ptr, n_batch, n_input, n_aux_input, n_output, n_cell,
+      /*activation=*/kTfLiteActSigmoid, forget_gate_scratch, is_input_all_zeros,
+      is_aux_input_all_zeros, logger, intermediate_tensor_indexes[1],
+      subgraph_index, error_reporter);
+  // Calculate the cell update gate.
+  CalculateLstmGateFloat(
+      input_ptr, input_to_cell_weights_ptr, aux_input_ptr,
+      aux_input_to_cell_weights_ptr, output_state_ptr,
+      recurrent_to_cell_weights_ptr, /*cell_state=*/nullptr,
+      /*cell_to_gate_weights=*/nullptr, cell_layer_norm_coefficients_ptr,
+      cell_gate_bias_ptr, n_batch, n_input, n_aux_input, n_output, n_cell,
+      params->activation, cell_gate_scratch, is_input_all_zeros,
+      is_aux_input_all_zeros, logger, intermediate_tensor_indexes[2],
+      subgraph_index, error_reporter);
+  // Update the cell state.
+  UpdateLstmCellFloat(n_batch, n_cell, cell_state_ptr, input_gate_scratch,
+                      forget_gate_scratch, cell_gate_scratch, use_cifg,
+                      params->cell_clip);
+  // Calculate output gate.
+  CalculateLstmGateFloat(
+      input_ptr, input_to_output_weights_ptr, aux_input_ptr,
+      aux_input_to_output_weights_ptr, output_state_ptr,
+      recurrent_to_output_weights_ptr, cell_state_ptr,
+      cell_to_output_weights_ptr, output_layer_norm_coefficients_ptr,
+      output_gate_bias_ptr, n_batch, n_input, n_aux_input, n_output, n_cell,
+      /*activation=*/kTfLiteActSigmoid, output_gate_scratch, is_input_all_zeros,
+      is_aux_input_all_zeros, logger, intermediate_tensor_indexes[3],
+      subgraph_index, error_reporter);
+  // Update the output state.
+  CalculateLstmOutputCalibration(
+      n_batch, n_cell, n_output, cell_state_ptr, output_gate_scratch,
+      params->activation, projection_weights_ptr, projection_bias_ptr,
+      params->proj_clip, output_state_ptr, scratch2, logger,
+      intermediate_tensor_indexes[4], subgraph_index, error_reporter);
+  // Copy output state to the output. Note that the output's rows may not be
+  // contiguous (output_batch_leading_dim != n_output).
+  for (int b = 0; b < n_batch; b++) {
+    std::copy_n(output_state_ptr + b * n_output, n_output,
+                output_ptr + b * output_batch_leading_dim);
+  }
+}
+
+TfLiteStatus EvalCalibration(
+    const TfLiteTensor* input, const TfLiteTensor* input_to_input_weights,
+    const TfLiteTensor* input_to_forget_weights,
+    const TfLiteTensor* input_to_cell_weights,
+    const TfLiteTensor* input_to_output_weights,
+    const TfLiteTensor* recurrent_to_input_weights,
+    const TfLiteTensor* recurrent_to_forget_weights,
+    const TfLiteTensor* recurrent_to_cell_weights,
+    const TfLiteTensor* recurrent_to_output_weights,
+    const TfLiteTensor* cell_to_input_weights,
+    const TfLiteTensor* cell_to_forget_weights,
+    const TfLiteTensor* cell_to_output_weights,
+    const TfLiteTensor* input_layer_norm_coefficients,
+    const TfLiteTensor* forget_layer_norm_coefficients,
+    const TfLiteTensor* cell_layer_norm_coefficients,
+    const TfLiteTensor* output_layer_norm_coefficients,
+    const TfLiteTensor* aux_input,
+    const TfLiteTensor* aux_input_to_input_weights,
+    const TfLiteTensor* aux_input_to_forget_weights,
+    const TfLiteTensor* aux_input_to_cell_weights,
+    const TfLiteTensor* aux_input_to_output_weights,
+    const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
+    const TfLiteTensor* cell_gate_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
+    const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
+    int output_offset, TfLiteTensor* scratch_buffer, TfLiteTensor* output_state,
+    TfLiteTensor* cell_state, TfLiteTensor* output, Logger* logger,
+    const std::vector<int>& intermediate_tensor_indexes,
+    const int subgraph_index, ErrorReporter* error_reporter) {
+  TF_LITE_ASSERT(input->dims->size >= 2 && input->dims->size <= 3);
+  int max_time, n_batch;
+  if (input->dims->size == 3) {
+    max_time = (time_major) ? input->dims->data[0] : input->dims->data[1];
+    n_batch = (time_major) ? input->dims->data[1] : input->dims->data[0];
+  } else {
+    max_time = 1;
+    n_batch = input->dims->data[0];
+  }
+  const int n_input = input->dims->data[input->dims->size - 1];
+  const int aux_input_size =
+      (aux_input) ? aux_input->dims->data[aux_input->dims->size - 1] : 0;
+
+  // n_cell and n_output will be the same size when there is no projection.
+  const int n_cell = input_to_output_weights->dims->data[0];
+  const int n_output = recurrent_to_output_weights->dims->data[1];
+
+  // Since we have already checked that weights are all there or none, we can
+  // check the existence of only one to the get the condition.
+  const bool use_cifg = (input_to_input_weights == nullptr);
+
+  // Index the scratch buffers pointers to the global scratch buffer.
+  float* scratch_buffer_ptr = GetTensorData<float>(scratch_buffer);
+  float* input_gate_scratch = nullptr;
+  float* cell_gate_scratch = nullptr;
+  float* forget_gate_scratch = nullptr;
+  float* output_gate_scratch = nullptr;
+  if (use_cifg) {
+    cell_gate_scratch = scratch_buffer_ptr;
+    forget_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
+    output_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
+  } else {
+    input_gate_scratch = scratch_buffer_ptr;
+    cell_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
+    forget_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
+    output_gate_scratch = scratch_buffer_ptr + 3 * n_cell * n_batch;
+  }
+
+  const int output_batch_leading_dim =
+      output->dims->data[output->dims->size - 1];
+  if (time_major) {
+    // Loop through the sequence.
+    const int input_step = n_batch * n_input;
+    const int output_step = n_batch * output_batch_leading_dim;
+    for (int t = 0; t < max_time; t++) {
+      // If this is the forward_sequence, step forward, otherwise step
+      // backwards.
+      const int t_rel = forward_sequence ? t : max_time - t - 1;
+      const float* input_ptr = GetTensorData<float>(input) + t_rel * input_step;
+      const float* aux_input_ptr = nullptr;
+      if (aux_input) {
+        aux_input_ptr = GetTensorData<float>(aux_input) + t_rel * input_step;
+      }
+      float* output_ptr_time =
+          GetTensorData<float>(output) + t_rel * output_step + output_offset;
+
+      LstmStepCalibration(
+          input_ptr, GetTensorData<float>(input_to_input_weights),
+          GetTensorData<float>(input_to_forget_weights),
+          GetTensorData<float>(input_to_cell_weights),
+          GetTensorData<float>(input_to_output_weights), aux_input_ptr,
+          GetTensorData<float>(aux_input_to_input_weights),
+          GetTensorData<float>(aux_input_to_forget_weights),
+          GetTensorData<float>(aux_input_to_cell_weights),
+          GetTensorData<float>(aux_input_to_output_weights),
+          GetTensorData<float>(recurrent_to_input_weights),
+          GetTensorData<float>(recurrent_to_forget_weights),
+          GetTensorData<float>(recurrent_to_cell_weights),
+          GetTensorData<float>(recurrent_to_output_weights),
+          GetTensorData<float>(cell_to_input_weights),
+          GetTensorData<float>(cell_to_forget_weights),
+          GetTensorData<float>(cell_to_output_weights),
+          GetTensorData<float>(input_layer_norm_coefficients),
+          GetTensorData<float>(forget_layer_norm_coefficients),
+          GetTensorData<float>(cell_layer_norm_coefficients),
+          GetTensorData<float>(output_layer_norm_coefficients),
+          GetTensorData<float>(input_gate_bias),
+          GetTensorData<float>(forget_gate_bias),
+          GetTensorData<float>(cell_gate_bias),
+          GetTensorData<float>(output_gate_bias),
+          GetTensorData<float>(projection_weights),
+          GetTensorData<float>(projection_bias), params, n_batch, n_cell,
+          n_input, aux_input_size, n_output, output_batch_leading_dim,
+          GetTensorData<float>(output_state), GetTensorData<float>(cell_state),
+          input_gate_scratch, forget_gate_scratch, cell_gate_scratch,
+          output_gate_scratch, output_ptr_time, logger,
+          intermediate_tensor_indexes, subgraph_index, error_reporter);
+    }
+  } else {
+    for (int b = 0; b < n_batch; b++) {
+      const int input_step = n_input;
+      const int output_step = output_batch_leading_dim;
+      for (int t = 0; t < max_time; t++) {
+        // If this is the forward_sequence, step forward, otherwise step
+        // backwards.
+        const int t_rel = forward_sequence ? t : max_time - t - 1;
+        const int time_offset = b * max_time + t_rel;
+        const float* input_ptr =
+            GetTensorData<float>(input) + time_offset * input_step;
+        const float* aux_input_ptr = nullptr;
+        if (aux_input) {
+          aux_input_ptr =
+              GetTensorData<float>(aux_input) + time_offset * input_step;
+        }
+        float* output_ptr = GetTensorData<float>(output) +
+                            time_offset * output_step + output_offset;
+
+        // Offset the {output,cell}_state pointers to the right batch.
+        float* output_state_ptr =
+            GetTensorData<float>(output_state) + b * output_batch_leading_dim;
+        float* cell_state_ptr = GetTensorData<float>(cell_state) + b * n_cell;
+        // Offset the scratch pointers to the right batch.
+        float* input_gate_scratch_ptr =
+            input_gate_scratch ? input_gate_scratch + b * n_cell : nullptr;
+        float* forget_gate_scratch_ptr = forget_gate_scratch + b * n_cell;
+        float* cell_gate_scratch_ptr = cell_gate_scratch + b * n_cell;
+        float* output_gate_scratch_ptr = output_gate_scratch + b * n_cell;
+
+        LstmStepCalibration(
+            input_ptr, GetTensorData<float>(input_to_input_weights),
+            GetTensorData<float>(input_to_forget_weights),
+            GetTensorData<float>(input_to_cell_weights),
+            GetTensorData<float>(input_to_output_weights), aux_input_ptr,
+            GetTensorData<float>(aux_input_to_input_weights),
+            GetTensorData<float>(aux_input_to_forget_weights),
+            GetTensorData<float>(aux_input_to_cell_weights),
+            GetTensorData<float>(aux_input_to_output_weights),
+            GetTensorData<float>(recurrent_to_input_weights),
+            GetTensorData<float>(recurrent_to_forget_weights),
+            GetTensorData<float>(recurrent_to_cell_weights),
+            GetTensorData<float>(recurrent_to_output_weights),
+            GetTensorData<float>(cell_to_input_weights),
+            GetTensorData<float>(cell_to_forget_weights),
+            GetTensorData<float>(cell_to_output_weights),
+            GetTensorData<float>(input_layer_norm_coefficients),
+            GetTensorData<float>(forget_layer_norm_coefficients),
+            GetTensorData<float>(cell_layer_norm_coefficients),
+            GetTensorData<float>(output_layer_norm_coefficients),
+            GetTensorData<float>(input_gate_bias),
+            GetTensorData<float>(forget_gate_bias),
+            GetTensorData<float>(cell_gate_bias),
+            GetTensorData<float>(output_gate_bias),
+            GetTensorData<float>(projection_weights),
+            GetTensorData<float>(projection_bias), params, /*n_batch=*/1,
+            n_cell, n_input, aux_input_size, n_output, output_batch_leading_dim,
+            output_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
+            forget_gate_scratch_ptr, cell_gate_scratch_ptr,
+            output_gate_scratch_ptr, output_ptr, logger,
+            intermediate_tensor_indexes, subgraph_index, error_reporter);
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
+struct OpData {
+  // Which kernel type to use. Full kernel (24 inputs) or basic kernel (5
+  // inputs).
+  // Please note the 20-input full kernel is deprecated and only kept
+  // here for backward compatibility.
+  TfLiteLSTMKernelType kernel_type;
+
+  // If the lstm is layer norm.
+  bool use_layer_norm;
+
+  // These fields are only used by full kernel.
+  int scratch_tensor_index;
+};
+
+// Resize the output, state tensors based on the sizes of the input tensors.
+// Allocate a temporary scratch tensor. Also check that the sizes of the input
+// tensors match each other.
+TfLiteStatus lstm_eval(TfLiteContext* context, int subgraph_index,
+                       TfLiteNode* node, LSTMType lstm_type, Logger* logger,
+                       ErrorReporter* error_reporter) {
+  const TfLiteTensor* input;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kInputTensor, &input));
+
+  const TfLiteTensor* input_to_input_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kInputToInputWeightsTensor);
+  const TfLiteTensor* input_to_forget_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kInputToForgetWeightsTensor,
+                   &input_to_forget_weights));
+  const TfLiteTensor* input_to_cell_weights;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kInputToCellWeightsTensor,
+                            &input_to_cell_weights));
+  const TfLiteTensor* input_to_output_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kInputToOutputWeightsTensor,
+                   &input_to_output_weights));
+
+  const TfLiteTensor* recurrent_to_input_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kRecurrentToInputWeightsTensor);
+  const TfLiteTensor* recurrent_to_forget_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kRecurrentToForgetWeightsTensor,
+                   &recurrent_to_forget_weights));
+  const TfLiteTensor* recurrent_to_cell_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kRecurrentToCellWeightsTensor,
+                   &recurrent_to_cell_weights));
+  const TfLiteTensor* recurrent_to_output_weights;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node,
+                   ops::builtin::lstm::full::kRecurrentToOutputWeightsTensor,
+                   &recurrent_to_output_weights));
+
+  const TfLiteTensor* cell_to_input_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kCellToInputWeightsTensor);
+  const TfLiteTensor* cell_to_forget_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kCellToForgetWeightsTensor);
+  const TfLiteTensor* cell_to_output_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kCellToOutputWeightsTensor);
+
+  const TfLiteTensor* input_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kInputLayerNormCoefficientsTensor);
+  const TfLiteTensor* forget_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kForgetLayerNormCoefficientsTensor);
+  const TfLiteTensor* cell_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kCellLayerNormCoefficientsTensor);
+  const TfLiteTensor* output_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kOutputLayerNormCoefficientsTensor);
+
+  const TfLiteTensor* input_gate_bias = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kInputGateBiasTensor);
+  const TfLiteTensor* forget_gate_bias;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kForgetGateBiasTensor,
+                            &forget_gate_bias));
+  const TfLiteTensor* cell_gate_bias;
+  TF_LITE_ENSURE_OK(
+      context,
+      GetInputSafe(context, node, ops::builtin::lstm::full::kCellGateBiasTensor,
+                   &cell_gate_bias));
+  const TfLiteTensor* output_gate_bias;
+  TF_LITE_ENSURE_OK(
+      context, GetInputSafe(context, node,
+                            ops::builtin::lstm::full::kOutputGateBiasTensor,
+                            &output_gate_bias));
+
+  const TfLiteTensor* projection_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kProjectionWeightsTensor);
+  const TfLiteTensor* projection_bias = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kProjectionBiasTensor);
+
+  // Index the scratch buffers pointers to the global scratch buffer.
+  TfLiteTensor* scratch_buffer;
+  TF_LITE_ENSURE_OK(
+      context, GetTemporarySafe(context, node, /*index=*/0, &scratch_buffer));
+
+  TfLiteTensor* output_state = GetVariableInput(
+      context, node, ops::builtin::lstm::full::kOutputStateTensor);
+  TF_LITE_ENSURE(context, output_state != nullptr);
+  TfLiteTensor* cell_state = GetVariableInput(
+      context, node, ops::builtin::lstm::full::kCellStateTensor);
+  TF_LITE_ENSURE(context, cell_state != nullptr);
+
+  TfLiteTensor* output;
+  TF_LITE_ENSURE_OK(
+      context, GetOutputSafe(context, node,
+                             ops::builtin::lstm::full::kOutputTensor, &output));
+
+  std::vector<int> intermediate_tensor_indexes(node->intermediates->size);
+  // LSTM expect 5 intermediate tensors.
+  TF_LITE_ENSURE_EQ(context, node->intermediates->size, 5);
+  for (int i = 0; i < node->intermediates->size; ++i) {
+    intermediate_tensor_indexes[i] = node->intermediates->data[i];
+  }
+
+  TfLiteLSTMParams lstm_params;
+  bool time_major = true;
+  switch (lstm_type) {
+    case LSTMType::kLSTM: {
+      lstm_params = *(static_cast<TfLiteLSTMParams*>(node->builtin_data));
+      time_major = true;
+      break;
+    }
+    case LSTMType::kUnidirectionalSequenceLSTM: {
+      const auto* params = static_cast<TfLiteUnidirectionalSequenceLSTMParams*>(
+          node->builtin_data);
+      // Copy out the LSTM specific params so they can be passed in the
+      // function.
+      lstm_params.activation = params->activation;
+      lstm_params.cell_clip = params->cell_clip;
+      lstm_params.proj_clip = params->proj_clip;
+      lstm_params.asymmetric_quantize_inputs =
+          params->asymmetric_quantize_inputs;
+      time_major = params->time_major;
+      break;
+    }
+    default:
+      return kTfLiteError;
+  }
+
+  switch (input_to_output_weights->type) {
+    case kTfLiteFloat32: {
+      return EvalCalibration(
+          input, input_to_input_weights, input_to_forget_weights,
+          input_to_cell_weights, input_to_output_weights,
+          recurrent_to_input_weights, recurrent_to_forget_weights,
+          recurrent_to_cell_weights, recurrent_to_output_weights,
+          cell_to_input_weights, cell_to_forget_weights, cell_to_output_weights,
+          input_layer_norm_coefficients, forget_layer_norm_coefficients,
+          cell_layer_norm_coefficients, output_layer_norm_coefficients,
+          /*aux_input=*/nullptr,
+          /*aux_input_to_input_weights=*/nullptr,
+          /*aux_input_to_forget_weights=*/nullptr,
+          /*aux_input_to_cell_weights=*/nullptr,
+          /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
+          forget_gate_bias, cell_gate_bias, output_gate_bias,
+          projection_weights, projection_bias, &lstm_params,
+          /*forward_sequence=*/true,
+          /*time_major=*/time_major,
+          /*output_offset=*/0, scratch_buffer, output_state, cell_state, output,
+          logger, intermediate_tensor_indexes, subgraph_index, error_reporter);
+    }
+    case kTfLiteUInt8:
+    case kTfLiteInt8:
+    default:
+      printf("Error. Only float model can be calibrated\n");
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+}  // namespace
+
+TfLiteStatus lstm_logging_kernel(TfLiteContext* context,
+                                 const int subgraph_index, TfLiteNode* node,
+                                 Logger* logger,
+                                 ErrorReporter* error_reporter) {
+  return lstm_eval(context, subgraph_index, node, LSTMType::kLSTM, logger,
+                   error_reporter);
+}
+
+TfLiteStatus unidirectional_sequence_lstm_logging_kernel(
+    TfLiteContext* context, const int subgraph_index, TfLiteNode* node,
+    Logger* logger, ErrorReporter* error_reporter) {
+  return lstm_eval(context, subgraph_index, node,
+                   LSTMType::kUnidirectionalSequenceLSTM, logger,
+                   error_reporter);
+}
+
+}  // namespace builtin
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.h

@@ -0,0 +1,45 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_BUILTIN_LOGGING_OPS_LSTM_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_BUILTIN_LOGGING_OPS_LSTM_H_
+
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_logger.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+namespace builtin {
+
+enum class LSTMType {
+  kLSTM,
+  kUnidirectionalSequenceLSTM,
+};
+
+TfLiteStatus lstm_logging_kernel(TfLiteContext* context, int subgraph_index,
+                                 TfLiteNode* node, Logger* logger,
+                                 ErrorReporter* error_reporter);
+
+TfLiteStatus unidirectional_sequence_lstm_logging_kernel(
+    TfLiteContext* context, int subgraph_index, TfLiteNode* node,
+    Logger* logger, ErrorReporter* error_reporter);
+
+}  // namespace builtin
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_BUILTIN_LOGGING_OPS_LSTM_H_

tensorflow/lite/tools/optimize/calibration/calibration_common.h

@@ -0,0 +1,76 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_COMMON_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_COMMON_H_
+
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "tensorflow/lite/mutable_op_resolver.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+using BuiltinOperatorKey = std::pair<BuiltinOperator, int>;
+
+using CustomOperatorKey = std::pair<std::string, int>;
+
+using BuiltinOpsSet = std::unordered_set<
+    BuiltinOperatorKey,
+    op_resolver_hasher::OperatorKeyHasher<BuiltinOperatorKey>>;
+
+using CustomOpsSet = std::unordered_set<
+    CustomOperatorKey,
+    op_resolver_hasher::OperatorKeyHasher<CustomOperatorKey>>;
+
+template <typename T>
+class BuiltinOpsMap
+    : public std::unordered_map<
+          BuiltinOperatorKey, T,
+          op_resolver_hasher::OperatorKeyHasher<BuiltinOperatorKey>> {};
+
+template <typename T>
+class CustomOpsMap
+    : public std::unordered_map<
+          CustomOperatorKey, T,
+          op_resolver_hasher::OperatorKeyHasher<CustomOperatorKey>> {};
+
+// An alias for |TfLiteRegistration.invoke|.
+using KernelEvalFuncPtr = TfLiteStatus (*)(TfLiteContext*, TfLiteNode*);
+
+enum class OperatorTensorType { kNone, kInput, kOutput, kIntermediate };
+
+// Information about an operator in the TfLite graph.
+struct OperatorInfo {
+  int subgraph_index;
+  int node_index;
+  std::string name;
+  BuiltinOperator builtin_op_code;
+  bool is_custom_op;
+  std::vector<int> inputs;
+  std::vector<int> outputs;
+  // Inputs that need to be logged.
+  std::vector<int> loggable_inputs;
+  // Outputs that need to be logged.
+  std::vector<int> loggable_outputs;
+  const TfLiteRegistration* registration;
+  int version;
+};
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_COMMON_H_

tensorflow/lite/tools/optimize/calibration/calibration_logger.cc

@@ -0,0 +1,57 @@
+/* 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/tools/optimize/calibration/calibration_logger.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+TfLiteStatus MinMax::Update(const float* values, size_t tensor_size,
+                            ErrorReporter* error_reporter) {
+  if (tensor_size <= 0) return kTfLiteOk;
+
+  // TODO(shashishekhar): Make it possible to use weighted/moving average.
+  bool has_nan_value = false;
+  for (size_t i = 0; i < tensor_size; ++i) {
+    const float value = values[i];
+    if (std::isnan(value)) {
+      has_nan_value = true;
+      continue;
+    }
+    has_values_ = true;
+    min_ = std::min<float>(min_, value);
+    max_ = std::max<float>(max_, value);
+  }
+  if (has_nan_value) {
+    TFLITE_LOG(TFLITE_LOG_WARNING,
+               "Model resulted in Nan value during calibration. Please "
+               "consider making sure that model results in all real-values "
+               "during inference with provided dataset.");
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/calibration/calibration_logger.h

@@ -0,0 +1,74 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_LOGGER_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_LOGGER_H_
+
+#include <limits>
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/common.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+class MinMax {
+ public:
+  TfLiteStatus Update(const float* values, size_t tensor_size,
+                      ErrorReporter* error_reporter);
+
+  bool HasValues() const { return has_values_; }
+
+  TfLiteStatus Get(float* min_val, float* max_val) const {
+    if (!has_values_) return kTfLiteError;
+    *min_val = min_;
+    *max_val = max_;
+    return kTfLiteOk;
+  }
+
+ private:
+  bool has_values_ = false;
+  float min_ = std::numeric_limits<float>::max();
+  float max_ = std::numeric_limits<float>::min();
+};
+
+// Captures min max values for tensors.
+class Logger {
+ public:
+  // Log the value for tensor at |tensor_index| which has |tensor_values|
+  TfLiteStatus LogTensorValue(int subgraph_index, int tensor_index,
+                              const float* tensor_values, size_t tensor_size,
+                              ErrorReporter* error_reporter) {
+    std::tuple<int, int> key{subgraph_index, tensor_index};
+    return tensor_id_to_stats_map_[key].Update(tensor_values, tensor_size,
+                                               error_reporter);
+  }
+
+  // Returns a map from tensor_index -> observed min max values.
+  const absl::flat_hash_map<std::tuple<int, int>, MinMax>&
+  GetCalibrationValues() const {
+    return tensor_id_to_stats_map_;
+  }
+
+ private:
+  absl::flat_hash_map<std::tuple<int, int>, MinMax> tensor_id_to_stats_map_;
+};
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_LOGGER_H_

tensorflow/lite/tools/optimize/calibration/calibration_reader.cc

@@ -0,0 +1,80 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/calibration/calibration_reader.h"
+
+#include <memory>
+#include <tuple>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+TfLiteStatus CalibrationReader::GetTensorStatsAsMap(
+    absl::flat_hash_map<std::tuple<int, int>, CalibrationStats>*
+        tensor_id_to_stats_map) const {
+  tensor_id_to_stats_map->clear();
+  for (const auto& tensorid_stat : logger_->GetCalibrationValues()) {
+    auto minmax = tensorid_stat.second;
+    CalibrationReader::CalibrationStats stats;
+    TF_LITE_ENSURE_STATUS(minmax.Get(&stats.min, &stats.max));
+    tensor_id_to_stats_map->insert({tensorid_stat.first, stats});
+  }
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus CalibrationReader::AddCalibrationToModel(ModelT* model,
+                                                      const bool update) const {
+  if (!model || model->subgraphs.empty()) {
+    return kTfLiteError;
+  }
+  for (const auto& tensorid_stat : logger_->GetCalibrationValues()) {
+    int subgraph_index, tensor_index;
+    std::tie(subgraph_index, tensor_index) = tensorid_stat.first;
+    const auto& subgraph = model->subgraphs[subgraph_index];
+    auto minmax = tensorid_stat.second;
+    float min, max;
+    TfLiteStatus status = minmax.Get(&min, &max);
+    if (status != kTfLiteOk) continue;
+    if (update) {
+      auto tensor = subgraph->tensors[tensor_index].get();
+      if (tensor->quantization) {
+        if (!tensor->quantization->min.empty()) {
+          const float existing_min = tensor->quantization->min[0];
+          min = min < existing_min ? min : existing_min;
+        }
+        if (!tensor->quantization->max.empty()) {
+          const float existing_max = tensor->quantization->max[0];
+          max = max > existing_max ? max : existing_max;
+        }
+      }
+    }
+    auto quant_params = std::make_unique<tflite::QuantizationParametersT>();
+    quant_params->min.push_back(min);
+    quant_params->max.push_back(max);
+    subgraph->tensors[tensor_index]->quantization = std::move(quant_params);
+  }
+
+  return kTfLiteOk;
+}
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/calibration/calibration_reader.h

@@ -0,0 +1,62 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_READER_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_READER_H_
+
+#include <tuple>
+
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/lite/core/c/c_api_types.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_logger.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+// Warning: This is not a public API and subject to change.
+//
+// Reads calibrator data collected by running the interpreter through
+// a calibration set.
+class CalibrationReader {
+ public:
+  struct CalibrationStats {
+    float min;
+    float max;
+  };
+  explicit CalibrationReader(const Logger* logger) : logger_(logger) {}
+
+  // Gets a map from tensor index to recorded calibration values.
+  virtual TfLiteStatus GetTensorStatsAsMap(
+      absl::flat_hash_map<std::tuple<int, int>, CalibrationStats>*
+          tensor_id_to_stats_map) const;
+
+  // Annotates the tensors in the given model with statistics captured during
+  // calibration.
+  // "update" is a flag: when set to true, the min/max are updated, instead of
+  // being overwritten.
+  virtual TfLiteStatus AddCalibrationToModel(ModelT* model, bool update) const;
+
+  virtual ~CalibrationReader() {}
+
+ private:
+  const Logger* logger_;
+};
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATION_READER_H_

tensorflow/lite/tools/optimize/calibration/calibrator.cc

@@ -0,0 +1,497 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/calibration/calibrator.h"
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/vector.h"  // from @flatbuffers
+#include "tensorflow/compiler/mlir/lite/allocation.h"
+#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/api/op_resolver.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/op_resolver.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/stderr_reporter.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/tools/optimize/calibration/builtin_logging_ops/lstm.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_common.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_logger.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_reader.h"
+#include "tensorflow/lite/tools/optimize/calibration/custom_logging_ops/lstm.h"
+#include "tensorflow/lite/tools/optimize/calibration/logging_op.h"
+#include "tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+namespace {
+
+// Calibrator is used to hold information that can be accessed during kernel
+// invocations.
+// TfLite kernel invocations are C functions and cannot look at the global
+// structure of the graph. Calibrator allows the kernel invoke functions to
+// access the global structure of graph and know which node is currently being
+// executed. This also allows us to write a simple kernel invoke wrapper
+// (see LoggingEval) that can work for most builtin ops.
+class Calibrator {
+ public:
+  Calibrator(const std::unordered_map<const TfLiteNode*, OperatorInfo>&
+                 node_ptr_opinfo_map,
+             std::unique_ptr<LoggingOpResolver> logging_op_resolver,
+             ErrorReporter* error_reporter)
+      : node_ptr_opinfo_map_(node_ptr_opinfo_map),
+        logging_op_resolver_(std::move(logging_op_resolver)),
+        error_reporter_(error_reporter) {
+    logger_ = std::make_unique<Logger>();
+  }
+
+  // Returns the wrapped kernel invoke function |TfLiteRegistration.invoke|.
+  KernelEvalFuncPtr GetKernelInvoke(const TfLiteNode* node) const;
+
+  // Gets the instance of logger associated with the current context.
+  Logger* GetLogger() const { return logger_.get(); }
+
+  // Gets the error reporter.
+  ErrorReporter* GetErrorReporter() const { return error_reporter_; }
+
+  // Gets the operator information about the given TfLiteNode.
+  const OperatorInfo& GetOpInfo(const TfLiteNode* node) const {
+    return node_ptr_opinfo_map_.at(node);
+  }
+
+  std::vector<const TfLiteNode*> GetNodesUnderCalibration() {
+    std::vector<const TfLiteNode*> nodes;
+    nodes.reserve(node_ptr_opinfo_map_.size());
+    for (const auto& entry : node_ptr_opinfo_map_) {
+      nodes.push_back(entry.first);
+    }
+    return nodes;
+  }
+
+ private:
+  std::unordered_map<const TfLiteNode*, OperatorInfo> node_ptr_opinfo_map_;
+  std::unique_ptr<LoggingOpResolver> logging_op_resolver_;
+  const std::unordered_map<int, OperatorInfo> index_opinfo_;
+  std::unique_ptr<Logger> logger_;
+  ErrorReporter* error_reporter_;
+};
+
+KernelEvalFuncPtr Calibrator::GetKernelInvoke(const TfLiteNode* node) const {
+  auto op_info = node_ptr_opinfo_map_.at(node);
+  if (op_info.is_custom_op) {
+    return logging_op_resolver_->GetWrappedKernelInvoke(op_info.name.c_str(),
+                                                        op_info.version);
+  }
+  return logging_op_resolver_->GetWrappedKernelInvoke(op_info.builtin_op_code,
+                                                      op_info.version);
+}
+
+// A registry of |Calibrator| objects per |TfLiteContext|.
+// This global registry is needed to access |Calibrator| objects in the kernel
+// invoke functions i.e. |TfLiteRegistration.invoke|.
+// Kernel invoke functions are C functions that have limited access to
+// |TfLiteContext|. Kernel invoke functions don't have access to global state of
+// graph. That means during a kernel invocation, the function cannot know which
+// node it was invoked for. E.g. in case of a model with |Conv| op at two
+// locations, there is no easy way for the Conv.invoke function to disambiguate
+// the calls.
+//
+// For calibration we solve this problem by creating a map of calibrators
+// per |TfLiteContext|. This map is |GlobalCalibrationRegistry|.
+//
+// This registry is then accessed using a global getter function:
+// |GetCalibratorRegistry|.
+// E.g.
+// TfLiteStatus SomeKernelInvokeFn(TfLiteContext* context, TfLiteNode* node) {
+//   .... code ....
+//   auto registry = GetCalibratorRegistry();
+//   auto calibrator = registry->GetCalibrator(context);
+//   ..... code ....
+//  }
+//
+// This way the kernel invoke functions can get the access to the Calibrator
+// object associated with the |TfLiteContext|.
+class GlobalCalibratorRegistry {
+ public:
+  // Get the |Calibrator| associated with given context, returns null if no
+  // calibrator is associated with the given context.
+  Calibrator* GetCalibrator(const TfLiteNode* node) const {
+    if (node_to_calibrator_.find(node) == node_to_calibrator_.cend()) {
+      return nullptr;
+    }
+    return node_to_calibrator_.at(node);
+  }
+
+  // Removes the association between calibrator and context.
+  // Note: This deletes the calibrator as well.
+  void RemoveCalibrator(const TfLiteContext* context) {
+    Calibrator* calibrator = calibrator_registry_.at(context).get();
+    auto nodes = calibrator->GetNodesUnderCalibration();
+    for (auto node : nodes) {
+      node_to_calibrator_.erase(node);
+    }
+    calibrator_registry_.erase(context);
+  }
+
+  // Creates an instance of |Calibrator|.
+  // Registry owns the |Calibrator| object which can be deleted by calling
+  // |RemoveCalibrator|.
+  TfLiteStatus CreateCalibrator(
+      const TfLiteContext* context,
+      const std::unordered_map<const TfLiteNode*, OperatorInfo>& node_to_opinfo,
+      std::unique_ptr<LoggingOpResolver> logging_op_resolver,
+      Calibrator** calibrator_ptr, ErrorReporter* reporter) {
+    if (calibrator_registry_.find(context) != calibrator_registry_.cend()) {
+      reporter->Report(
+          "Failed to create calibrator, context already registered.");
+      return kTfLiteError;
+    }
+    auto calibrator = std::make_unique<Calibrator>(
+        node_to_opinfo, std::move(logging_op_resolver), reporter);
+    calibrator_registry_[context] = std::move(calibrator);
+    *calibrator_ptr = calibrator_registry_.at(context).get();
+    for (const auto& entry : node_to_opinfo) {
+      node_to_calibrator_[entry.first] = *calibrator_ptr;
+    }
+    return kTfLiteOk;
+  }
+
+ private:
+  absl::flat_hash_map<const TfLiteContext*, std::unique_ptr<Calibrator>>
+      calibrator_registry_;
+  absl::flat_hash_map<const TfLiteNode*, Calibrator*> node_to_calibrator_;
+};
+
+GlobalCalibratorRegistry* GetCalibratorRegistry() {
+  static GlobalCalibratorRegistry* registry = new GlobalCalibratorRegistry();
+  return registry;
+}
+
+// Get the logging kernel if there are any.
+// TODO(jianlijianli): extend this to support multiple recipe for the same
+// model.
+logging_kernel_func_ptr GetLoggingEvalFunc(TfLiteContext* context,
+                                           TfLiteNode* node,
+                                           int builtin_op_code) {
+  switch (builtin_op_code) {
+    case BuiltinOperator_LSTM: {
+      if (node->intermediates->size == 12) {
+        return tflite::optimize::calibration::custom::lstm_logging_kernel;
+      }
+      return tflite::optimize::calibration::builtin::lstm_logging_kernel;
+    }
+    case BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM:
+      return tflite::optimize::calibration::builtin::
+          unidirectional_sequence_lstm_logging_kernel;
+    default:
+      return nullptr;
+  }
+}
+
+// A wrapper implementation for |TfLiteRegistration.invoke| that logs inputs,
+// invokes the wrapped implementation and then logs the outputs.
+TfLiteStatus LoggingEval(TfLiteContext* context, TfLiteNode* node) {
+  Calibrator* calibrator = GetCalibratorRegistry()->GetCalibrator(node);
+
+  if (!calibrator) {
+    TF_LITE_KERNEL_LOG(context, "No calibrator found for context.");
+    return kTfLiteError;
+  }
+
+  auto kernel_invoke = calibrator->GetKernelInvoke(node);
+  auto logger = calibrator->GetLogger();
+  auto op_info = calibrator->GetOpInfo(node);
+  auto error_reporter = calibrator->GetErrorReporter();
+
+  for (int i : op_info.loggable_inputs) {
+    auto tensor = context->tensors[i];
+    TF_LITE_ENSURE_STATUS(
+        logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
+                               tensor.bytes / sizeof(float), error_reporter));
+  }
+  auto builtin_op_code = calibrator->GetOpInfo(node).builtin_op_code;
+  auto kernel_invoke_intermediate =
+      GetLoggingEvalFunc(context, node, builtin_op_code);
+  if (kernel_invoke_intermediate == nullptr) {
+    TF_LITE_ENSURE_STATUS(kernel_invoke(context, node));
+  } else {
+    TF_LITE_ENSURE_STATUS(
+        kernel_invoke_intermediate(context, op_info.subgraph_index, node,
+                                   calibrator->GetLogger(), error_reporter));
+  }
+
+  // TODO(shashishekhar): An intermediate tensor in graph will get logged twice
+  // once as an input and second time as output. This doesn't change the min max
+  // values but is inefficient.
+  // Using moving average will also break this.
+
+  // Log input again to make sure the state tensors are captured after lstm
+  // cell.
+  for (int i : op_info.loggable_inputs) {
+    auto tensor = context->tensors[i];
+    TF_LITE_ENSURE_STATUS(
+        logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
+                               tensor.bytes / sizeof(float), error_reporter));
+  }
+
+  for (int i : op_info.loggable_outputs) {
+    auto tensor = context->tensors[i];
+    TF_LITE_ENSURE_STATUS(
+        logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
+                               tensor.bytes / sizeof(float), error_reporter));
+  }
+
+  return kTfLiteOk;
+}
+
+// Returns the loggable tensors. Not all inputs and outputs need to be logged.
+// For example, const weight tensors which have buffers associated with them
+// don't need to be logged.
+std::vector<int> GetLoggableTensorIndices(
+    const std::vector<int>& tensor_indices,
+    const flatbuffers::Vector<flatbuffers::Offset<Tensor>>* tensors,
+    const flatbuffers::Vector<flatbuffers::Offset<Buffer>>* tensor_buffers) {
+  std::vector<int> loggable;
+  for (auto tensor_index : tensor_indices) {
+    if (tensor_index == kTfLiteOptionalTensor) {
+      continue;
+    }
+    auto tensor = tensors->Get(tensor_index);
+    auto buffer_index = tensor->buffer();
+    const bool has_no_buffer =
+        (tensor_buffers->Get(buffer_index) == nullptr) ||
+        (tensor_buffers->Get(buffer_index)->data() == nullptr) ||
+        (tensor_buffers->Get(buffer_index)->data()->size() == 0);
+    if (has_no_buffer && tensor->type() == tflite::TensorType_FLOAT32) {
+      loggable.push_back(tensor_index);
+    }
+  }
+  return loggable;
+}
+
+// Creates a mapping between the static model graph and the runtime TfLiteNode*
+// nodes in the graph for the given context.
+// This is done by querying the TfLiteContext for node and registrations using
+// the |NodeInfoDelegateObserver|.
+TfLiteStatus GetNodeOpInfoMapAndContext(
+    const absl::flat_hash_map<std::tuple<int, int>, OperatorInfo>&
+        node_to_opinfo,
+    tflite::Interpreter* const interpreter,
+    std::unordered_map<const TfLiteNode*, OperatorInfo>* node_ptr_opinfo_map,
+    TfLiteContext** context) {
+  *context = interpreter->primary_subgraph().context();
+
+  // Since we only consider the primary subgraph while populating
+  // node_to_opinfo, do the same here.
+  // Because Flex delegate can merge multiple op nodes into one Delegate node if
+  // they are located in a row, the size of the execution plan can be lesser
+  // than the size of the graph's op nodes.
+  TF_LITE_ENSURE(*context,
+                 interpreter->execution_plan().size() <= node_to_opinfo.size());
+  for (const auto& entry : node_to_opinfo) {
+    auto op_info = entry.second;
+    int subgraph_index, op_index;
+    std::tie(subgraph_index, op_index) = entry.first;
+    const auto* node_and_reg =
+        interpreter->node_and_registration(subgraph_index, op_index);
+    op_info.registration = &node_and_reg->second;
+    node_ptr_opinfo_map->insert({&node_and_reg->first, op_info});
+  }
+  return kTfLiteOk;
+}
+
+string GetOpName(const tflite::OperatorCode& opcode) {
+  if (opcode.custom_code() != nullptr) {
+    return opcode.custom_code()->str();
+  }
+  return tflite::EnumNamesBuiltinOperator()[GetBuiltinCode(&opcode)];
+}
+
+// A |CalibrationReader| that owns the Calibrator.
+class Reader : public CalibrationReader {
+ public:
+  Reader(const TfLiteContext* context, const Logger* logger)
+      : CalibrationReader(logger), context_(context) {}
+
+  ~Reader() override { GetCalibratorRegistry()->RemoveCalibrator(context_); }
+
+ private:
+  const TfLiteContext* context_;
+};
+
+bool HasInputs(BuiltinOperator code) {
+  switch (code) {
+    case BuiltinOperator_CALL_ONCE:
+    case BuiltinOperator_VAR_HANDLE:
+    // Custom ops, including Flex ops, might not have inputs.
+    case BuiltinOperator_CUSTOM:
+      return false;
+    default:
+      return true;
+  }
+}
+
+bool HasOutputs(BuiltinOperator code) {
+  switch (code) {
+    case BuiltinOperator_ASSIGN_VARIABLE:
+    case BuiltinOperator_CALL_ONCE:
+    // Custom ops, including Flex ops, might not have outputs.
+    case BuiltinOperator_CUSTOM:
+      return false;
+    default:
+      return true;
+  }
+}
+
+}  // namespace
+
+TfLiteStatus BuildLoggingInterpreter(
+    const FlatBufferModel& model, const OpResolver& op_resolver,
+    std::unique_ptr<Interpreter>* interpreter,
+    std::unique_ptr<CalibrationReader>* calibration_reader) {
+  return BuildLoggingInterpreter(model.GetModel(), model.error_reporter(),
+                                 op_resolver, interpreter, calibration_reader,
+                                 model.allocation());
+}
+
+TfLiteStatus BuildLoggingInterpreter(
+    const tflite::Model* tflite_model, ErrorReporter* error_reporter,
+    const OpResolver& op_resolver, std::unique_ptr<Interpreter>* interpreter,
+    std::unique_ptr<CalibrationReader>* calibration_reader,
+    const Allocation* allocation) {
+  if (error_reporter == nullptr) {
+    // Make sure error_reporter is valid.
+    error_reporter = DefaultErrorReporter();
+  }
+  auto subgraphs = tflite_model->subgraphs();
+  auto tensor_buffers = tflite_model->buffers();
+
+  // Populate the node index to operator info map.
+  // We want to collect this information so we can use it during runtime to
+  // log details of which inputs and outputs.
+  // At runtime TFLite kernel invoke functions can only look into their
+  // own node in the graph (TFLiteNode*) and some limited context information.
+  absl::flat_hash_map<std::tuple<int, int>, OperatorInfo> node_to_opinfo;
+  BuiltinOpsSet builtin_op_and_versions;
+  CustomOpsSet custom_op_and_versions;
+
+  for (size_t subgraph_index = 0; subgraph_index < subgraphs->size();
+       subgraph_index++) {
+    auto subgraph = subgraphs->Get(subgraph_index);
+    auto operator_codes = tflite_model->operator_codes();
+    auto operators = subgraph->operators();
+    auto tensors = subgraph->tensors();
+    if (!operators) {
+      continue;
+    }
+
+    for (size_t i = 0; i < operators->size(); i++) {
+      OperatorInfo op_info;
+      op_info.subgraph_index = subgraph_index;
+      op_info.node_index = i;
+      auto op = operators->Get(i);
+      auto operator_code = operator_codes->Get(op->opcode_index());
+      op_info.builtin_op_code = GetBuiltinCode(operator_code);
+      op_info.name = GetOpName(*operator_code);
+      op_info.is_custom_op = operator_code->custom_code() != nullptr;
+      op_info.version = operator_code->version();
+
+      auto op_inputs = op->inputs();
+      auto op_outputs = op->outputs();
+      if (op_inputs) {
+        op_info.inputs = std::vector<int>(op_inputs->begin(), op_inputs->end());
+      } else if (HasInputs(op_info.builtin_op_code)) {
+        TFLITE_LOG(TFLITE_LOG_WARNING, "Op %s missing inputs",
+                   op_info.name.c_str());
+      }
+      if (op_outputs) {
+        op_info.outputs =
+            std::vector<int>(op_outputs->begin(), op_outputs->end());
+      } else if (HasOutputs(op_info.builtin_op_code)) {
+        TFLITE_LOG(TFLITE_LOG_WARNING, "Op %s missing outputs",
+                   op_info.name.c_str());
+      }
+      op_info.loggable_inputs =
+          GetLoggableTensorIndices(op_info.inputs, tensors, tensor_buffers);
+      op_info.loggable_outputs =
+          GetLoggableTensorIndices(op_info.outputs, tensors, tensor_buffers);
+      if (op_info.is_custom_op) {
+        op_info.registration =
+            op_resolver.FindOp(op_info.name.c_str(), operator_code->version());
+        custom_op_and_versions.insert(
+            {op_info.name.c_str(), operator_code->version()});
+      } else {
+        op_info.registration = op_resolver.FindOp(GetBuiltinCode(operator_code),
+                                                  operator_code->version());
+        builtin_op_and_versions.insert(
+            {op_info.builtin_op_code, operator_code->version()});
+      }
+      std::tuple<int, int> key{subgraph_index, i};
+      node_to_opinfo[key] = op_info;
+    }
+  }
+
+  // Prepare the logging op resolver to use |LoggingEval| for kernel
+  // invocations.
+  auto logging_op_resolver = std::make_unique<LoggingOpResolver>(
+      builtin_op_and_versions, custom_op_and_versions, op_resolver, LoggingEval,
+      error_reporter);
+  tflite::InterpreterBuilder(tflite_model, *logging_op_resolver, error_reporter,
+                             /*options_experimental=*/nullptr,
+                             allocation)(interpreter);
+
+  if (!(*interpreter)) {
+    error_reporter->Report("Failed to construct interpreter");
+    return kTfLiteError;
+  }
+
+  // Compute the mapping between runtime and static graph structure, i.e.
+  // (TfLiteContext, TfLiteNode) -> OperatorInfo
+  std::unordered_map<const TfLiteNode*, OperatorInfo> node_ptr_opinfo_map;
+  TfLiteContext* context = nullptr;
+  TF_LITE_ENSURE_STATUS(GetNodeOpInfoMapAndContext(
+      node_to_opinfo, interpreter->get(), &node_ptr_opinfo_map, &context));
+
+  Calibrator* calibrator = nullptr;
+  // Register a calibrator object for the context. This can be accessed
+  // during invocations by the logging kernels.
+  TF_LITE_ENSURE_STATUS(GetCalibratorRegistry()->CreateCalibrator(
+      context, node_ptr_opinfo_map, std::move(logging_op_resolver), &calibrator,
+      error_reporter));
+  *calibration_reader = std::unique_ptr<CalibrationReader>(
+      new Reader(context, calibrator->GetLogger()));
+
+  return kTfLiteOk;
+}
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/calibration/calibrator.h

@@ -0,0 +1,77 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATOR_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATOR_H_
+
+#include <memory>
+
+#include "tensorflow/compiler/mlir/lite/allocation.h"
+#include "tensorflow/lite/allocation.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/api/op_resolver.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_reader.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+// Warning: This is not a public API and subject to change.
+
+// Builds a interpreter that logs the calibration data in memory.
+// The calibration data can be recovered using |calibration_reader|.
+//
+// Sample usage:
+// std::unique_ptr<Interpreter> interpreter;
+// std::unique_ptr<CalibrationReader> calibration_reader;
+// BuiltinOpResolver resolver = ...
+// FlatBufferModel model = ..
+//
+// BuildLoggingInterpreter(model, resolver, &interpreter,
+//  &calibration_reader);
+//
+//
+// * Allocate tensors...
+// * Call interpreter->invoke on calibration dataset.
+//
+// Calibration data can be read either directly by calling
+// std::unordered_map<int,  CalibrationStats>> tensor_index_to_stats;
+// calibration_reader->GetTensorStatsAsMap(&tensor_index_to_stats);
+//
+// or adding calibration data to model itself.
+// ModelT * original_floating_point_model = ...
+// calibration_reader->AddCalibrationToModel(original_floating_point_model,
+// false);
+//
+TfLiteStatus BuildLoggingInterpreter(
+    const FlatBufferModel& model, const OpResolver& op_resolver,
+    std::unique_ptr<Interpreter>* interpreter,
+    std::unique_ptr<CalibrationReader>* calibration_reader);
+
+// Same as above, except gets separate tflite::Model and ErrorReporter pointers.
+TfLiteStatus BuildLoggingInterpreter(
+    const tflite::Model* model, ErrorReporter* error_reporter,
+    const OpResolver& op_resolver, std::unique_ptr<Interpreter>* interpreter,
+    std::unique_ptr<CalibrationReader>* calibration_reader,
+    const Allocation* allocation = nullptr);
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CALIBRATOR_H_

tensorflow/lite/tools/optimize/calibration/calibrator_test.cc

@@ -0,0 +1,733 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/calibration/calibrator.h"
+
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <iostream>
+#include <memory>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/command_line_flags.h"
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/kernels/register.h"
+#include "tensorflow/lite/interpreter.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_reader.h"
+
+namespace {
+std::string* g_test_model_dir = nullptr;
+}  // namespace
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+namespace {
+
+std::unique_ptr<FlatBufferModel> ReadModel(const string& model_name) {
+  auto model_path = tensorflow::io::JoinPath(*g_test_model_dir, model_name);
+  return FlatBufferModel::BuildFromFile(model_path.c_str());
+}
+
+TEST(CalibratorTest, CalibrationStatsAreCollected) {
+  auto model = ReadModel("multi_add.bin");
+  ASSERT_TRUE(model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(
+      *model, ops::builtin::BuiltinOpResolver{}, &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_TRUE(stats.empty());
+
+  status = interpreter->AllocateTensors();
+  ASSERT_EQ(kTfLiteOk, status);
+  // Model does the following:
+  // 0        1       2        3
+  // |        |__ ____|        |
+  // |           |             |
+  // |          Add(tensor:4)  |
+  // |____ ______|______ ______|
+  //      |             |
+  //      Add          Add
+  //      |             |
+  //    Output:5      Output:6
+
+  const size_t tensor_size = 1 * 8 * 8 * 3;
+
+  std::vector<float> ones(tensor_size, 1.0f);
+  // Fill input tensor i with i+1, i.e. input[0] = 1.0f, input[1] = 2.0f,
+  // input[2] = 3.0f
+
+  for (size_t i = 0; i < interpreter->inputs().size(); i++) {
+    int input_tensor_idx = interpreter->inputs()[i];
+    TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_EQ(tensor->bytes, tensor_size * sizeof(float));
+    for (size_t j = 0; j < tensor_size; j++) {
+      tensor->data.f[j] = i + 1;
+    }
+  }
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+  const float eps = 1e-6f;
+  // Verify that tensor 5: is 6
+  // Verify that tensor 6: is 9
+  TfLiteTensor* tensor = interpreter->tensor(interpreter->outputs()[0]);
+  for (size_t i = 0; i < tensor_size; i++) {
+    EXPECT_NEAR(tensor->data.f[i], 6.0f, eps);
+  }
+  tensor = interpreter->tensor(interpreter->outputs()[1]);
+  for (size_t i = 0; i < tensor_size; i++) {
+    EXPECT_NEAR(tensor->data.f[i], 9.0f, eps);
+  }
+
+  // Verify that min max of tensors.
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_EQ(7, stats.size());
+  // Check inputs
+  for (int tensor_idx = 0; tensor_idx < 4; tensor_idx++) {
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.min, tensor_idx + 1, eps);
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.max, tensor_idx + 1, eps);
+  }
+  // Check tensor 4 max.
+  EXPECT_NEAR(stats.find({0, 4})->second.min, 5, eps);
+  EXPECT_NEAR(stats.find({0, 4})->second.max, 5, eps);
+
+  // Check outputs
+  EXPECT_NEAR(stats.find({0, 5})->second.min, 6, eps);
+  EXPECT_NEAR(stats.find({0, 5})->second.max, 6, eps);
+
+  EXPECT_NEAR(stats.find({0, 6})->second.min, 9, eps);
+  EXPECT_NEAR(stats.find({0, 6})->second.max, 9, eps);
+}
+
+TEST(CalibratorTest, MultipleInvokes) {
+  auto model = ReadModel("multi_add.bin");
+  ASSERT_TRUE(model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(
+      *model, ops::builtin::BuiltinOpResolver{}, &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  status = interpreter->AllocateTensors();
+
+  EXPECT_EQ(kTfLiteOk, status);
+  const size_t tensor_size = 1 * 8 * 8 * 3;
+  // Fill input tensor i with i+1, i.e. input[0] = 1.0f, input[1] = 2.0f,
+  // input[2] = 3.0f
+
+  for (size_t i = 0; i < interpreter->inputs().size(); i++) {
+    int input_tensor_idx = interpreter->inputs()[i];
+    TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_EQ(tensor->bytes, tensor_size * sizeof(float));
+    for (size_t j = 0; j < tensor_size; j++) {
+      tensor->data.f[j] = i + 1;
+    }
+  }
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+  const float eps = 1e-6f;
+  // Verify that min max of tensors.
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_EQ(7, stats.size());
+  const float expected_values[7] = {
+      1.0f,  // input 0
+      2.0f,  // input 1
+      3.0f,  // input 2
+      4.0f,  // input 3
+      5.0f,  // Add(1, 2)
+      6.0f,  // Output 5: Add(0, Add(1,2))
+      9.0f,  // Output 6: Add(Add(1,2), 3)
+  };
+  for (int tensor_idx = 0; tensor_idx < 7; tensor_idx++) {
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.min,
+                expected_values[tensor_idx], eps);
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.max,
+                expected_values[tensor_idx], eps);
+  }
+  // Set input[0][0] = 1.5 and input[0][1] = 0.5 this should change the values
+  // only for input[0] and tensor 4 and outputs 5, 6.
+  TfLiteTensor* input0 = interpreter->tensor(0);
+  input0->data.f[0] = 1.5f;
+  input0->data.f[1] = 0.5f;
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_EQ(7, stats.size());
+  EXPECT_NEAR(stats.find({0, 0})->second.min, 0.5f, eps);
+  EXPECT_NEAR(stats.find({0, 0})->second.max, 1.5f, eps);
+
+  for (int tensor_idx = 1; tensor_idx < 5; tensor_idx++) {
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.min,
+                expected_values[tensor_idx], eps);
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.max,
+                expected_values[tensor_idx], eps);
+  }
+
+  EXPECT_NEAR(stats.find({0, 5})->second.min, 5.5f, eps);
+  EXPECT_NEAR(stats.find({0, 5})->second.max, 6.5f, eps);
+
+  EXPECT_NEAR(stats.find({0, 6})->second.min, 9.0f, eps);
+  EXPECT_NEAR(stats.find({0, 6})->second.max, 9.0f, eps);
+}
+
+TEST(CalibratorTest, UpdateMinMax) {
+  auto flatbuffer_model = ReadModel("multi_add.bin");
+  ASSERT_TRUE(flatbuffer_model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(*flatbuffer_model,
+                                        ops::builtin::BuiltinOpResolver{},
+                                        &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+  auto readonly_model = flatbuffer_model->GetModel();
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  status = interpreter->AllocateTensors();
+
+  EXPECT_EQ(kTfLiteOk, status);
+  const size_t tensor_size = 1 * 8 * 8 * 3;
+  for (size_t i = 0; i < interpreter->inputs().size(); i++) {
+    int input_tensor_idx = interpreter->inputs()[i];
+    TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_EQ(tensor->bytes, tensor_size * sizeof(float));
+    for (size_t j = 0; j < tensor_size; j++) {
+      tensor->data.f[j] = i + 1;
+    }
+  }
+  auto input_0_quant_params =
+      std::make_unique<tflite::QuantizationParametersT>();
+  input_0_quant_params->min.push_back(0.5);
+  input_0_quant_params->max.push_back(1.5);
+  model.subgraphs[0]->tensors[0]->quantization =
+      std::move(input_0_quant_params);
+
+  // Invoke with update == true.
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+  const float eps = 1e-6f;
+  // Verify that min max of tensors.
+  const float expected_min[7] = {
+      0.5f,  // input 0
+      2.0f,  // input 1
+      3.0f,  // input 2
+      4.0f,  // input 3
+      5.0f,  // Add(1, 2)
+      6.0f,  // Output 5: Add(0, Add(1,2))
+      9.0f,  // Output 6: Add(Add(1,2), 3)
+  };
+  const float expected_max[7] = {
+      1.5f,  // input 0
+      2.0f,  // input 1
+      3.0f,  // input 2
+      4.0f,  // input 3
+      5.0f,  // Add(1, 2)
+      6.0f,  // Output 5: Add(0, Add(1,2))
+      9.0f,  // Output 6: Add(Add(1,2), 3)
+  };
+  status = reader->AddCalibrationToModel(&model, /*update=*/true);
+  for (int tensor_idx = 0; tensor_idx < 7; tensor_idx++) {
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->min[0],
+                expected_min[tensor_idx], eps);
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->max[0],
+                expected_max[tensor_idx], eps);
+  }
+
+  // Invoke with update == false;
+  // Verify that min max of tensors.
+  const float expected_value[7] = {
+      1.0f,  // input 0
+      2.0f,  // input 1
+      3.0f,  // input 2
+      4.0f,  // input 3
+      5.0f,  // Add(1, 2)
+      6.0f,  // Output 5: Add(0, Add(1,2))
+      9.0f,  // Output 6: Add(Add(1,2), 3)
+  };
+  status = reader->AddCalibrationToModel(&model, /*update=*/false);
+  for (int tensor_idx = 0; tensor_idx < 7; tensor_idx++) {
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->min[0],
+                expected_value[tensor_idx], eps);
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->max[0],
+                expected_value[tensor_idx], eps);
+  }
+}
+
+TEST(CalibratorTest, HandleNanValues) {
+  auto flatbuffer_model = ReadModel("multi_add.bin");
+  ASSERT_TRUE(flatbuffer_model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(*flatbuffer_model,
+                                        ops::builtin::BuiltinOpResolver{},
+                                        &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+  auto readonly_model = flatbuffer_model->GetModel();
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  status = interpreter->AllocateTensors();
+
+  EXPECT_EQ(kTfLiteOk, status);
+  const size_t tensor_size = 1 * 8 * 8 * 3;
+  for (size_t i = 0; i < interpreter->inputs().size(); i++) {
+    int input_tensor_idx = interpreter->inputs()[i];
+    TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_EQ(tensor->bytes, tensor_size * sizeof(float));
+    for (size_t j = 0; j < tensor_size; j++) {
+      if (j % 2 == 0) {
+        tensor->data.f[j] = NAN;
+      } else {
+        tensor->data.f[j] = i + 1;
+      }
+    }
+  }
+  auto input_0_quant_params =
+      std::make_unique<tflite::QuantizationParametersT>();
+  input_0_quant_params->min.push_back(0.5);
+  input_0_quant_params->max.push_back(1.5);
+  model.subgraphs[0]->tensors[0]->quantization =
+      std::move(input_0_quant_params);
+
+  // Invoke with update == true.
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+  const float eps = 1e-6f;
+  // Verify that min max of tensors.
+  const float expected_min[7] = {
+      0.5f,  // input 0
+      2.0f,  // input 1
+      3.0f,  // input 2
+      4.0f,  // input 3
+      5.0f,  // Add(1, 2)
+      6.0f,  // Output 5: Add(0, Add(1,2))
+      9.0f,  // Output 6: Add(Add(1,2), 3)
+  };
+  const float expected_max[7] = {
+      1.5f,  // input 0
+      2.0f,  // input 1
+      3.0f,  // input 2
+      4.0f,  // input 3
+      5.0f,  // Add(1, 2)
+      6.0f,  // Output 5: Add(0, Add(1,2))
+      9.0f,  // Output 6: Add(Add(1,2), 3)
+  };
+  status = reader->AddCalibrationToModel(&model, /*update=*/true);
+  for (int tensor_idx = 0; tensor_idx < 7; tensor_idx++) {
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->min[0],
+                expected_min[tensor_idx], eps);
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->max[0],
+                expected_max[tensor_idx], eps);
+  }
+
+  // Invoke with update == false;
+  // Verify that min max of tensors.
+  const float expected_value[7] = {
+      1.0f,  // input 0
+      2.0f,  // input 1
+      3.0f,  // input 2
+      4.0f,  // input 3
+      5.0f,  // Add(1, 2)
+      6.0f,  // Output 5: Add(0, Add(1,2))
+      9.0f,  // Output 6: Add(Add(1,2), 3)
+  };
+  status = reader->AddCalibrationToModel(&model, /*update=*/false);
+  for (int tensor_idx = 0; tensor_idx < 7; tensor_idx++) {
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->min[0],
+                expected_value[tensor_idx], eps);
+    EXPECT_NEAR(model.subgraphs[0]->tensors[tensor_idx]->quantization->max[0],
+                expected_value[tensor_idx], eps);
+  }
+}
+
+TEST(CalibratorTest, LSTM) {
+  auto flatbuffer_model = ReadModel("lstm.bin");
+  ASSERT_TRUE(flatbuffer_model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(*flatbuffer_model,
+                                        ops::builtin::BuiltinOpResolver{},
+                                        &interpreter, &reader);
+  EXPECT_EQ(status, kTfLiteOk);
+
+  auto readonly_model = flatbuffer_model->GetModel();
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  status = interpreter->AllocateTensors();
+
+  EXPECT_EQ(kTfLiteOk, status);
+  const std::vector<float> lstm_input = {0.3, 0.2};
+  int input_tensor_idx = interpreter->inputs()[0];
+  TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+  for (size_t j = 0; j < lstm_input.size(); j++) {
+    tensor->data.f[j] = lstm_input[j];
+  }
+
+  ASSERT_EQ(interpreter->Invoke(), kTfLiteOk);
+
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  EXPECT_EQ(reader->GetTensorStatsAsMap(&stats), kTfLiteOk);
+
+  // Check the results.
+  const float eps = 1e-6f;
+  const absl::flat_hash_map<std::tuple<int, int>,
+                            CalibrationReader::CalibrationStats>
+      expected_calibration_result = {
+          // Input.
+          {{0, 0}, {0.200000, 0.300000}},
+          // State.
+          {{0, 18}, {0.000000, 0.468415}},
+          // State.
+          {{0, 19}, {0.000000, 0.424350}},
+          // Output.
+          {{0, 24}, {0.265968, 0.468415}},
+          // Intemediate_0.
+          {{0, 25}, {0.080045, 0.170588}},
+          // Intemediate_1.
+          {{0, 26}, {0.080045, 0.170588}},
+          // Intemediate_2.
+          {{0, 27}, {0.080045, 0.170588}},
+          // Intemediate_3.
+          {{0, 28}, {0.080045, 0.170588}},
+          // Intemediate_4.
+          {{0, 29}, {0.000000, 0.270944}},
+      };
+  EXPECT_EQ(expected_calibration_result.size(), stats.size());
+  for (const auto& e : stats) {
+    auto expected_result = expected_calibration_result.find(e.first)->second;
+    EXPECT_NEAR(e.second.min, expected_result.min, eps);
+    EXPECT_NEAR(e.second.max, expected_result.max, eps);
+  }
+}
+
+TEST(CalibratorTest, UnidirectionalSequenceLSTM) {
+  auto flatbuffer_model = ReadModel("unidirectional_sequence_lstm.bin");
+  ASSERT_TRUE(flatbuffer_model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(*flatbuffer_model,
+                                        ops::builtin::BuiltinOpResolver{},
+                                        &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+
+  auto readonly_model = flatbuffer_model->GetModel();
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  EXPECT_EQ(interpreter->AllocateTensors(), kTfLiteOk);
+  const std::vector<float> lstm_input = {0.3, 0.2, 0.9, 0.8};
+  int input_tensor_idx = interpreter->inputs()[0];
+  TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+  for (size_t j = 0; j < lstm_input.size(); j++) {
+    tensor->data.f[j] = lstm_input[j];
+  }
+
+  ASSERT_EQ(interpreter->Invoke(), kTfLiteOk);
+
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  EXPECT_EQ(reader->GetTensorStatsAsMap(&stats), kTfLiteOk);
+
+  // Check the results.
+  const float eps = 1e-6f;
+  const absl::flat_hash_map<std::tuple<int, int>,
+                            CalibrationReader::CalibrationStats>
+      expected_calibration_result = {
+          // Input.
+          {{0, 0}, {0.200000, 0.900000}},
+          // State.
+          {{0, 18}, {0.000000, 0.520999}},
+          // State.
+          {{0, 19}, {0.000000, 0.711364}},
+          // Output.
+          {{0, 24}, {0.247992, 0.520999}},
+          // Intemediate_0.
+          {{0, 25}, {0.080045, 0.824241}},
+          // Intemediate_1.
+          {{0, 26}, {0.080045, 0.824241}},
+          // Intemediate_2.
+          {{0, 27}, {0.080045, 0.824241}},
+          // Intemediate_3.
+          {{0, 28}, {0.080045, 0.824241}},
+          // Intemediate_4.
+          {{0, 29}, {0.000000, 0.413618}},
+      };
+  EXPECT_EQ(expected_calibration_result.size(), stats.size());
+  for (const auto& e : stats) {
+    auto expected_result = expected_calibration_result.find(e.first)->second;
+    EXPECT_NEAR(e.second.min, expected_result.min, eps);
+    EXPECT_NEAR(e.second.max, expected_result.max, eps);
+  }
+}
+
+TEST(CalibratorTest, CustomLSTM) {
+  auto flatbuffer_model = ReadModel("custom_lstm.bin");
+  ASSERT_TRUE(flatbuffer_model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(*flatbuffer_model,
+                                        ops::builtin::BuiltinOpResolver{},
+                                        &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+
+  auto readonly_model = flatbuffer_model->GetModel();
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  EXPECT_EQ(interpreter->AllocateTensors(), kTfLiteOk);
+  const std::vector<float> lstm_input = {0.3, 0.2, 0.9, 0.8};
+  int input_tensor_idx = interpreter->inputs()[0];
+  TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+  for (size_t j = 0; j < lstm_input.size(); j++) {
+    tensor->data.f[j] = lstm_input[j];
+  }
+
+  ASSERT_EQ(interpreter->Invoke(), kTfLiteOk);
+
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  EXPECT_EQ(reader->GetTensorStatsAsMap(&stats), kTfLiteOk);
+
+  // Check the results.
+  const float eps = 1e-6f;
+  const absl::flat_hash_map<std::tuple<int, int>,
+                            CalibrationReader::CalibrationStats>
+      expected_calibration_result = {
+          // input.
+          {{0, 0}, {0.200000, 0.300000}},
+          // state.
+          {{0, 18}, {0.000000, 0.468415}},
+          // state.
+          {{0, 19}, {0.000000, 0.424349}},
+          // output.
+          {{0, 24}, {0.265968, 0.468415}},
+          // intermediate 0.
+          {{0, 25}, {0.080045, 0.170588}},
+          // intermediate 1.
+          {{0, 26}, {0.080045, 0.170588}},
+          // intermediate 2.
+          {{0, 27}, {0.000000, 0.000000}},
+          // intermediate 3.
+          {{0, 28}, {0.080045, 0.170588}},
+          // intermediate 4.
+          {{0, 29}, {0.080045, 0.170588}},
+          // intermediate 5.
+          {{0, 30}, {0.000000, 0.000000}},
+          // intermediate 6.
+          {{0, 31}, {0.080045, 0.170588}},
+          // intermediate 7.
+          {{0, 32}, {0.080045, 0.170588}},
+          // intermediate 8.
+          {{0, 33}, {0.000000, 0.000000}},
+          // intermediate 9.
+          {{0, 34}, {0.080045, 0.170588}},
+          // intermediate 10.
+          {{0, 35}, {0.080045, 0.170588}},
+          // intermediate 11.
+          {{0, 36}, {0.000000, 0.000000}},
+      };
+  EXPECT_EQ(expected_calibration_result.size(), stats.size());
+  for (const auto& e : stats) {
+    auto expected_result = expected_calibration_result.find(e.first)->second;
+    EXPECT_NEAR(e.second.min, expected_result.min, eps);
+    EXPECT_NEAR(e.second.max, expected_result.max, eps);
+  }
+}
+
+TEST(CalibratorTest, CalibrationWithMultipleSubgraphs) {
+  auto model = ReadModel("multi_subgraphs_while.bin");
+  ASSERT_TRUE(model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(
+      *model, ops::builtin::BuiltinOpResolver{}, &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_TRUE(stats.empty());
+
+  status = interpreter->AllocateTensors();
+  ASSERT_EQ(kTfLiteOk, status);
+  // The model accepts a float input, and multiply this by 2 for two times with
+  // while loop.
+  // Model does the following:
+  //          serving_default_inp:0
+  //                    |
+  // While(cond: subgraph/1, body: subgraph/2)
+  //            /                  \
+  //    PartitionedCall:0   PartitionedCall:1
+  //      loop variable      input multiplier
+
+  const size_t tensor_size = 1;
+
+  for (size_t i = 0; i < interpreter->inputs().size(); i++) {
+    int input_tensor_idx = interpreter->inputs()[i];
+    TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_EQ(tensor->bytes, tensor_size * sizeof(int));
+    for (size_t j = 0; j < tensor_size; j++) {
+      tensor->data.f[j] = i + 1;
+    }
+  }
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+
+  // Verify that min max of tensors.
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_EQ(4, stats.size());
+
+  // Check the results.
+  const float eps = 1e-6f;
+  const absl::flat_hash_map<std::tuple<int, int>,
+                            CalibrationReader::CalibrationStats>
+      expected_calibration_result = {
+          // input.
+          {{0, 0}, {1.0, 1.0}},
+          // output.
+          {{0, 4}, {4.0, 4.0}},
+          // loop_multiply_input.
+          {{2, 2}, {1.0, 2.0}},
+          // loop_multiply_output.
+          {{2, 6}, {2.0, 4.0}},
+      };
+  EXPECT_EQ(expected_calibration_result.size(), stats.size());
+  for (const auto& e : stats) {
+    auto expected_result = expected_calibration_result.find(e.first)->second;
+    EXPECT_NEAR(e.second.min, expected_result.min, eps);
+    EXPECT_NEAR(e.second.max, expected_result.max, eps);
+  }
+}
+
+TEST(CalibratorTest, CalibrationWithCallOnce) {
+  auto model = ReadModel("call_once_mul.bin");
+  ASSERT_TRUE(model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(
+      *model, ops::builtin::BuiltinOpResolver{}, &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_TRUE(stats.empty());
+
+  status = interpreter->AllocateTensors();
+  ASSERT_EQ(kTfLiteOk, status);
+  // Model does the following:
+  // serving_default_inp:0 -> CallOnce() -> ReadVariableOp() -> Mul()
+  //                              |                 |
+  //                       AssignVariableOp()  AssignVariableOp
+  const size_t tensor_size = 1;
+  for (size_t i = 0; i < interpreter->inputs().size(); i++) {
+    int input_tensor_idx = interpreter->inputs()[i];
+    TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_EQ(tensor->bytes, tensor_size * sizeof(int));
+    for (size_t j = 0; j < tensor_size; j++) {
+      tensor->data.f[j] = i + 1;
+    }
+  }
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+
+  // Verify that min max of tensors.
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  EXPECT_EQ(3, stats.size());
+
+  // Check the results.
+  const float eps = 1e-6f;
+  const absl::flat_hash_map<std::tuple<int, int>,
+                            CalibrationReader::CalibrationStats>
+      expected_calibration_result = {// input.
+                                     {{0, 0}, {1.0, 1.0}},
+                                     // readvariableop.
+                                     {{0, 2}, {2.0, 2.0}},
+                                     // mul output.
+                                     {{0, 3}, {2.0, 2.0}}};
+  EXPECT_EQ(expected_calibration_result.size(), stats.size());
+  for (const auto& e : stats) {
+    auto expected_result = expected_calibration_result.find(e.first)->second;
+    EXPECT_NEAR(e.second.min, expected_result.min, eps);
+    EXPECT_NEAR(e.second.max, expected_result.max, eps);
+  }
+}
+}  // namespace
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+
+int main(int argc, char** argv) {
+  std::string model_file;
+  const std::vector<tensorflow::Flag> flag_list = {
+      tensorflow::Flag("test_model_file", &model_file,
+                       "Path to test tflite model file."),
+  };
+
+  const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
+  if (!parse_result) {
+    std::cerr << "Required test_model_file\n";
+    std::abort();
+  }
+  g_test_model_dir = new std::string(tensorflow::io::Dirname(model_file));
+  ::tensorflow::port::InitMain(argv[0], &argc, &argv);
+  return RUN_ALL_TESTS();
+}

tensorflow/lite/tools/optimize/calibration/custom_logging_ops/lstm.cc

@@ -0,0 +1,662 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/calibration/custom_logging_ops/lstm.h"
+
+#include <algorithm>
+#include <cstdio>
+#include <vector>
+
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+#include "tensorflow/lite/kernels/lstm_shared.h"
+#include "tensorflow/lite/kernels/op_macros.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_logger.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+namespace custom {
+
+namespace {
+
+inline void LstmStepWithAuxInput(
+    const float* input_ptr, const float* input_to_input_weights_ptr,
+    const float* input_to_forget_weights_ptr,
+    const float* input_to_cell_weights_ptr,
+    const float* input_to_output_weights_ptr, const float* aux_input_ptr,
+    const float* aux_input_to_input_weights_ptr,
+    const float* aux_input_to_forget_weights_ptr,
+    const float* aux_input_to_cell_weights_ptr,
+    const float* aux_input_to_output_weights_ptr,
+    const float* recurrent_to_input_weights_ptr,
+    const float* recurrent_to_forget_weights_ptr,
+    const float* recurrent_to_cell_weights_ptr,
+    const float* recurrent_to_output_weights_ptr,
+    const float* cell_to_input_weights_ptr,
+    const float* cell_to_forget_weights_ptr,
+    const float* cell_to_output_weights_ptr,
+    const float* input_layer_norm_coefficients_ptr,
+    const float* forget_layer_norm_coefficients_ptr,
+    const float* cell_layer_norm_coefficients_ptr,
+    const float* output_layer_norm_coefficients_ptr,
+    const float* input_gate_bias_ptr, const float* forget_gate_bias_ptr,
+    const float* cell_bias_ptr, const float* output_gate_bias_ptr,
+    const float* projection_weights_ptr, const float* projection_bias_ptr,
+    const TfLiteLSTMParams* params, int n_batch, int n_cell, int n_input,
+    int n_aux_input, int n_output, int output_batch_leading_dim,
+    float* output_state_ptr, float* cell_state_ptr, float* input_gate_scratch,
+    float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch,
+    float* output_ptr, Logger* logger,
+    const std::vector<int>& intemediate_tensor_indexes,
+    const int subgraph_index, ErrorReporter* error_reporter) {
+  // Since we have already checked that weights are all there or none, we can
+  // check the existence of only one to the get the condition.
+  const bool use_cifg = (input_to_input_weights_ptr == nullptr);
+  const bool use_peephole = (cell_to_output_weights_ptr != nullptr);
+  const bool use_layer_norm = (forget_layer_norm_coefficients_ptr != nullptr);
+
+  // Initialize scratch buffers with bias for regular lstm or initialize with
+  // zero for layer norm lstm.
+  if (use_layer_norm) {
+    if (!use_cifg) {
+      std::fill_n(input_gate_scratch, n_cell * n_batch, 0.0f);
+    }
+    std::fill_n(forget_gate_scratch, n_cell * n_batch, 0.0f);
+    std::fill_n(cell_scratch, n_cell * n_batch, 0.0f);
+    std::fill_n(output_gate_scratch, n_cell * n_batch, 0.0f);
+  } else {
+    if (!use_cifg) {
+      tensor_utils::VectorBatchVectorAssign(input_gate_bias_ptr, n_cell,
+                                            n_batch, input_gate_scratch);
+    }
+    tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch,
+                                          forget_gate_scratch);
+    tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch,
+                                          cell_scratch);
+    tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, n_cell, n_batch,
+                                          output_gate_scratch);
+  }
+
+  // For each batch and cell: compute input_weight * input.
+  if (!use_cifg) {
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        input_to_input_weights_ptr, n_cell, n_input, input_ptr, n_batch,
+        input_gate_scratch);
+  }
+
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+      input_to_forget_weights_ptr, n_cell, n_input, input_ptr, n_batch,
+      forget_gate_scratch);
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(input_to_cell_weights_ptr,
+                                                    n_cell, n_input, input_ptr,
+                                                    n_batch, cell_scratch);
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+      input_to_output_weights_ptr, n_cell, n_input, input_ptr, n_batch,
+      output_gate_scratch);
+
+  {
+    // calibration.
+    if (!use_cifg) {
+      logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[1],
+                             input_gate_scratch, n_cell * n_batch,
+                             error_reporter);
+    }
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[4],
+                           forget_gate_scratch, n_cell * n_batch,
+                           error_reporter);
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[7],
+                           cell_scratch, n_cell * n_batch, error_reporter);
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[10],
+                           output_gate_scratch, n_cell * n_batch,
+                           error_reporter);
+  }
+
+  // If auxiliary input is available then compute aux_input_weight * aux_input
+  if (aux_input_ptr != nullptr) {
+    if (!use_cifg) {
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          aux_input_to_input_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
+          n_batch, input_gate_scratch);
+    }
+
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        aux_input_to_forget_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
+        n_batch, forget_gate_scratch);
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        aux_input_to_cell_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
+        n_batch, cell_scratch);
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        aux_input_to_output_weights_ptr, n_cell, n_aux_input, aux_input_ptr,
+        n_batch, output_gate_scratch);
+  }
+
+  // For each batch and cell: compute recurrent_weight * output_state.
+  if (!use_cifg) {
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        recurrent_to_input_weights_ptr, n_cell, n_output, output_state_ptr,
+        n_batch, input_gate_scratch);
+  }
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+      recurrent_to_forget_weights_ptr, n_cell, n_output, output_state_ptr,
+      n_batch, forget_gate_scratch);
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+      recurrent_to_cell_weights_ptr, n_cell, n_output, output_state_ptr,
+      n_batch, cell_scratch);
+  tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+      recurrent_to_output_weights_ptr, n_cell, n_output, output_state_ptr,
+      n_batch, output_gate_scratch);
+  {
+    // calibrition.
+    if (!use_cifg) {
+      std::vector<float> temp_input(n_batch * n_cell);
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          recurrent_to_input_weights_ptr, n_cell, n_output, output_state_ptr,
+          n_batch, temp_input.data());
+      logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[2],
+                             temp_input.data(), n_cell * n_batch,
+                             error_reporter);
+    }
+    std::vector<float> temp_forget(n_batch * n_cell);
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        recurrent_to_forget_weights_ptr, n_cell, n_output, output_state_ptr,
+        n_batch, temp_forget.data());
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[5],
+                           temp_forget.data(), n_cell * n_batch,
+                           error_reporter);
+
+    std::vector<float> temp_cell(n_batch * n_cell);
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        recurrent_to_cell_weights_ptr, n_cell, n_output, output_state_ptr,
+        n_batch, temp_cell.data());
+
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[8],
+                           temp_cell.data(), n_cell * n_batch, error_reporter);
+
+    std::vector<float> temp_output(n_batch * n_cell);
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        recurrent_to_output_weights_ptr, n_cell, n_output, output_state_ptr,
+        n_batch, temp_output.data());
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[11],
+                           temp_output.data(), n_cell * n_batch,
+                           error_reporter);
+  }
+
+  // For each batch and cell: update input gate.
+  if (!use_cifg) {
+    if (use_peephole) {
+      tensor_utils::VectorBatchVectorCwiseProductAccumulate(
+          cell_to_input_weights_ptr, n_cell, cell_state_ptr, n_batch,
+          input_gate_scratch);
+    }
+    if (use_layer_norm) {
+      logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[0],
+                             input_gate_scratch, n_cell * n_batch,
+                             error_reporter);
+      tensor_utils::MeanStddevNormalization(
+          input_gate_scratch, input_gate_scratch, n_cell, n_batch);
+      tensor_utils::VectorBatchVectorCwiseProduct(
+          input_layer_norm_coefficients_ptr, n_cell, input_gate_scratch,
+          n_batch, input_gate_scratch);
+      tensor_utils::VectorBatchVectorAdd(input_gate_bias_ptr, n_cell, n_batch,
+                                         input_gate_scratch);
+    }
+    tensor_utils::ApplySigmoidToVector(input_gate_scratch, n_cell * n_batch,
+                                       input_gate_scratch);
+  }
+
+  // For each batch and cell: update forget gate.
+  if (use_peephole) {
+    tensor_utils::VectorBatchVectorCwiseProductAccumulate(
+        cell_to_forget_weights_ptr, n_cell, cell_state_ptr, n_batch,
+        forget_gate_scratch);
+  }
+  if (use_layer_norm) {
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[3],
+                           forget_gate_scratch, n_cell * n_batch,
+                           error_reporter);
+    tensor_utils::MeanStddevNormalization(forget_gate_scratch,
+                                          forget_gate_scratch, n_cell, n_batch);
+    tensor_utils::VectorBatchVectorCwiseProduct(
+        forget_layer_norm_coefficients_ptr, n_cell, forget_gate_scratch,
+        n_batch, forget_gate_scratch);
+    tensor_utils::VectorBatchVectorAdd(forget_gate_bias_ptr, n_cell, n_batch,
+                                       forget_gate_scratch);
+  }
+  tensor_utils::ApplySigmoidToVector(forget_gate_scratch, n_cell * n_batch,
+                                     forget_gate_scratch);
+
+  // For each batch and cell: update the cell.
+  tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_ptr,
+                                         n_batch * n_cell, cell_state_ptr);
+  if (use_layer_norm) {
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[6],
+                           cell_scratch, n_cell * n_batch, error_reporter);
+    tensor_utils::MeanStddevNormalization(cell_scratch, cell_scratch, n_cell,
+                                          n_batch);
+    tensor_utils::VectorBatchVectorCwiseProduct(
+        cell_layer_norm_coefficients_ptr, n_cell, cell_scratch, n_batch,
+        cell_scratch);
+    tensor_utils::VectorBatchVectorAdd(cell_bias_ptr, n_cell, n_batch,
+                                       cell_scratch);
+  }
+  tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell,
+                                        params->activation, cell_scratch);
+  if (use_cifg) {
+    tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell,
+                             forget_gate_scratch);
+    tensor_utils::VectorVectorCwiseProductAccumulate(
+        cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_ptr);
+  } else {
+    tensor_utils::VectorVectorCwiseProductAccumulate(
+        cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr);
+  }
+  if (params->cell_clip > 0.0) {
+    tensor_utils::CwiseClipping(cell_state_ptr, n_batch * n_cell,
+                                params->cell_clip);
+  }
+
+  // For each batch and cell: update the output gate.
+  if (use_peephole) {
+    tensor_utils::VectorBatchVectorCwiseProductAccumulate(
+        cell_to_output_weights_ptr, n_cell, cell_state_ptr, n_batch,
+        output_gate_scratch);
+  }
+  if (use_layer_norm) {
+    logger->LogTensorValue(subgraph_index, intemediate_tensor_indexes[9],
+                           output_gate_scratch, n_cell * n_batch,
+                           error_reporter);
+    tensor_utils::MeanStddevNormalization(output_gate_scratch,
+                                          output_gate_scratch, n_cell, n_batch);
+    tensor_utils::VectorBatchVectorCwiseProduct(
+        output_layer_norm_coefficients_ptr, n_cell, output_gate_scratch,
+        n_batch, output_gate_scratch);
+    tensor_utils::VectorBatchVectorAdd(output_gate_bias_ptr, n_cell, n_batch,
+                                       output_gate_scratch);
+  }
+  tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell,
+                                     output_gate_scratch);
+  tensor_utils::ApplyActivationToVector(cell_state_ptr, n_batch * n_cell,
+                                        params->activation, cell_scratch);
+  tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch,
+                                         n_batch * n_cell, output_gate_scratch);
+
+  const bool use_projection_weight = (projection_weights_ptr != nullptr);
+  const bool use_projection_bias = (projection_bias_ptr != nullptr);
+
+  // For each batch: update the projection and output_state. Note that since
+  // the output batch rows may not be contiguous (output_batch_leading_dim !=
+  // n_output), we unroll batched operations.
+  if (use_projection_weight) {
+    if (use_projection_bias) {
+      for (int k = 0; k < n_batch; k++) {
+        std::copy_n(projection_bias_ptr, n_output,
+                    output_ptr + k * output_batch_leading_dim);
+      }
+    } else {
+      for (int k = 0; k < n_batch; k++) {
+        std::fill_n(output_ptr + k * output_batch_leading_dim, n_output, 0.0f);
+      }
+    }
+    for (int k = 0; k < n_batch; k++) {
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          projection_weights_ptr, n_output, n_cell,
+          output_gate_scratch + k * n_cell,
+          /*n_batch=*/1, output_ptr + k * output_batch_leading_dim);
+      if (params->proj_clip > 0.0) {
+        tensor_utils::CwiseClipping(output_ptr + k * output_batch_leading_dim,
+                                    n_output, params->proj_clip);
+      }
+    }
+  } else {
+    for (int k = 0; k < n_batch; k++) {
+      std::copy_n(output_gate_scratch + k * n_output, n_output,
+                  output_ptr + k * output_batch_leading_dim);
+    }
+  }
+  for (int k = 0; k < n_batch; k++) {
+    std::copy_n(output_ptr + k * output_batch_leading_dim, n_output,
+                output_state_ptr + k * n_output);
+  }
+}
+
+TfLiteStatus EvalFloat(
+    const TfLiteTensor* input, const TfLiteTensor* input_to_input_weights,
+    const TfLiteTensor* input_to_forget_weights,
+    const TfLiteTensor* input_to_cell_weights,
+    const TfLiteTensor* input_to_output_weights,
+    const TfLiteTensor* recurrent_to_input_weights,
+    const TfLiteTensor* recurrent_to_forget_weights,
+    const TfLiteTensor* recurrent_to_cell_weights,
+    const TfLiteTensor* recurrent_to_output_weights,
+    const TfLiteTensor* cell_to_input_weights,
+    const TfLiteTensor* cell_to_forget_weights,
+    const TfLiteTensor* cell_to_output_weights,
+    const TfLiteTensor* input_layer_norm_coefficients,
+    const TfLiteTensor* forget_layer_norm_coefficients,
+    const TfLiteTensor* cell_layer_norm_coefficients,
+    const TfLiteTensor* output_layer_norm_coefficients,
+    const TfLiteTensor* aux_input,
+    const TfLiteTensor* aux_input_to_input_weights,
+    const TfLiteTensor* aux_input_to_forget_weights,
+    const TfLiteTensor* aux_input_to_cell_weights,
+    const TfLiteTensor* aux_input_to_output_weights,
+    const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias,
+    const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias,
+    const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias,
+    const TfLiteLSTMParams* params, bool forward_sequence, bool time_major,
+    int output_offset, TfLiteTensor* scratch_buffer,
+    TfLiteTensor* activation_state, TfLiteTensor* cell_state,
+    TfLiteTensor* output, Logger* logger,
+    const std::vector<int>& intemediate_tensor_indexes,
+    const int subgraph_index, ErrorReporter* error_reporter) {
+  TF_LITE_ASSERT(input->dims->size >= 2 && input->dims->size <= 3);
+  int max_time, n_batch;
+  if (input->dims->size == 3) {
+    max_time = (time_major) ? input->dims->data[0] : input->dims->data[1];
+    n_batch = (time_major) ? input->dims->data[1] : input->dims->data[0];
+  } else {
+    max_time = 1;
+    n_batch = input->dims->data[0];
+  }
+  const int n_input = input->dims->data[input->dims->size - 1];
+  const int aux_input_size =
+      (aux_input) ? aux_input->dims->data[aux_input->dims->size - 1] : 0;
+
+  // n_cell and n_output will be the same size when there is no projection.
+  const int n_cell = input_to_output_weights->dims->data[0];
+  const int n_output = recurrent_to_output_weights->dims->data[1];
+
+  // Since we have already checked that weights are all there or none, we can
+  // check the existence of only one to the get the condition.
+  const bool use_cifg = (input_to_input_weights == nullptr);
+
+  // Index the scratch buffers pointers to the global scratch buffer.
+  float* scratch_buffer_ptr = GetTensorData<float>(scratch_buffer);
+  float* input_gate_scratch = nullptr;
+  float* cell_scratch = nullptr;
+  float* forget_gate_scratch = nullptr;
+  float* output_gate_scratch = nullptr;
+  if (use_cifg) {
+    cell_scratch = scratch_buffer_ptr;
+    forget_gate_scratch = scratch_buffer_ptr + n_cell * n_batch;
+    output_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
+  } else {
+    input_gate_scratch = scratch_buffer_ptr;
+    cell_scratch = scratch_buffer_ptr + n_cell * n_batch;
+    forget_gate_scratch = scratch_buffer_ptr + 2 * n_cell * n_batch;
+    output_gate_scratch = scratch_buffer_ptr + 3 * n_cell * n_batch;
+  }
+
+  const int output_batch_leading_dim =
+      output->dims->data[output->dims->size - 1];
+  if (time_major) {
+    // Loop through the sequence.
+    const int input_step = n_batch * n_input;
+    const int output_step = n_batch * output_batch_leading_dim;
+    for (int t = 0; t < max_time; t++) {
+      // If this is the forward_sequence, step forward, otherwise step
+      // backwards.
+      const int t_rel = forward_sequence ? t : max_time - t - 1;
+      const float* input_ptr = GetTensorData<float>(input) + t_rel * input_step;
+      const float* aux_input_ptr = nullptr;
+      if (aux_input) {
+        aux_input_ptr = GetTensorData<float>(aux_input) + t_rel * input_step;
+      }
+      float* output_ptr_time =
+          GetTensorData<float>(output) + t_rel * output_step + output_offset;
+
+      LstmStepWithAuxInput(
+          input_ptr, GetTensorData<float>(input_to_input_weights),
+          GetTensorData<float>(input_to_forget_weights),
+          GetTensorData<float>(input_to_cell_weights),
+          GetTensorData<float>(input_to_output_weights), aux_input_ptr,
+          GetTensorData<float>(aux_input_to_input_weights),
+          GetTensorData<float>(aux_input_to_forget_weights),
+          GetTensorData<float>(aux_input_to_cell_weights),
+          GetTensorData<float>(aux_input_to_output_weights),
+          GetTensorData<float>(recurrent_to_input_weights),
+          GetTensorData<float>(recurrent_to_forget_weights),
+          GetTensorData<float>(recurrent_to_cell_weights),
+          GetTensorData<float>(recurrent_to_output_weights),
+          GetTensorData<float>(cell_to_input_weights),
+          GetTensorData<float>(cell_to_forget_weights),
+          GetTensorData<float>(cell_to_output_weights),
+          GetTensorData<float>(input_layer_norm_coefficients),
+          GetTensorData<float>(forget_layer_norm_coefficients),
+          GetTensorData<float>(cell_layer_norm_coefficients),
+          GetTensorData<float>(output_layer_norm_coefficients),
+          GetTensorData<float>(input_gate_bias),
+          GetTensorData<float>(forget_gate_bias),
+          GetTensorData<float>(cell_bias),
+          GetTensorData<float>(output_gate_bias),
+          GetTensorData<float>(projection_weights),
+          GetTensorData<float>(projection_bias), params, n_batch, n_cell,
+          n_input, aux_input_size, n_output, output_batch_leading_dim,
+          GetTensorData<float>(activation_state),
+          GetTensorData<float>(cell_state), input_gate_scratch,
+          forget_gate_scratch, cell_scratch, output_gate_scratch,
+          output_ptr_time, logger, intemediate_tensor_indexes, subgraph_index,
+          error_reporter);
+    }
+  } else {
+    for (int b = 0; b < n_batch; b++) {
+      const int input_step = n_input;
+      const int output_step = output_batch_leading_dim;
+      for (int t = 0; t < max_time; t++) {
+        // If this is the forward_sequence, step forward, otherwise step
+        // backwards.
+        const int t_rel = forward_sequence ? t : max_time - t - 1;
+        const int time_offset = b * max_time + t_rel;
+        const float* input_ptr =
+            GetTensorData<float>(input) + time_offset * input_step;
+        const float* aux_input_ptr = nullptr;
+        if (aux_input) {
+          aux_input_ptr =
+              GetTensorData<float>(aux_input) + time_offset * input_step;
+        }
+        float* output_ptr = GetTensorData<float>(output) +
+                            time_offset * output_step + output_offset;
+
+        // Offset the {activation,cell}_state pointers to the right batch.
+        float* activation_state_ptr = GetTensorData<float>(activation_state) +
+                                      b * output_batch_leading_dim;
+        float* cell_state_ptr = GetTensorData<float>(cell_state) + b * n_cell;
+        // Offset the scratch pointers to the right batch.
+        float* input_gate_scratch_ptr =
+            input_gate_scratch ? input_gate_scratch + b * n_cell : nullptr;
+        float* forget_gate_scratch_ptr = forget_gate_scratch + b * n_cell;
+        float* cell_scratch_ptr = cell_scratch + b * n_cell;
+        float* output_gate_scratch_ptr = output_gate_scratch + b * n_cell;
+
+        LstmStepWithAuxInput(
+            input_ptr, GetTensorData<float>(input_to_input_weights),
+            GetTensorData<float>(input_to_forget_weights),
+            GetTensorData<float>(input_to_cell_weights),
+            GetTensorData<float>(input_to_output_weights), aux_input_ptr,
+            GetTensorData<float>(aux_input_to_input_weights),
+            GetTensorData<float>(aux_input_to_forget_weights),
+            GetTensorData<float>(aux_input_to_cell_weights),
+            GetTensorData<float>(aux_input_to_output_weights),
+            GetTensorData<float>(recurrent_to_input_weights),
+            GetTensorData<float>(recurrent_to_forget_weights),
+            GetTensorData<float>(recurrent_to_cell_weights),
+            GetTensorData<float>(recurrent_to_output_weights),
+            GetTensorData<float>(cell_to_input_weights),
+            GetTensorData<float>(cell_to_forget_weights),
+            GetTensorData<float>(cell_to_output_weights),
+            GetTensorData<float>(input_layer_norm_coefficients),
+            GetTensorData<float>(forget_layer_norm_coefficients),
+            GetTensorData<float>(cell_layer_norm_coefficients),
+            GetTensorData<float>(output_layer_norm_coefficients),
+            GetTensorData<float>(input_gate_bias),
+            GetTensorData<float>(forget_gate_bias),
+            GetTensorData<float>(cell_bias),
+            GetTensorData<float>(output_gate_bias),
+            GetTensorData<float>(projection_weights),
+            GetTensorData<float>(projection_bias), params, /*n_batch=*/1,
+            n_cell, n_input, aux_input_size, n_output, output_batch_leading_dim,
+            activation_state_ptr, cell_state_ptr, input_gate_scratch_ptr,
+            forget_gate_scratch_ptr, cell_scratch_ptr, output_gate_scratch_ptr,
+            output_ptr, logger, intemediate_tensor_indexes, subgraph_index,
+            error_reporter);
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+
+struct OpData {
+  // Which kernel type to use. Full kernel (24 inputs) or basic kernel (5
+  // inputs).
+  // Please note the 20-input full kernel is deprecated and only kept
+  // here for backward compatibility.
+  TfLiteLSTMKernelType kernel_type;
+
+  // If the lstm is layer norm.
+  bool use_layer_norm;
+
+  // These fields are only used by full kernel.
+  int scratch_tensor_index;
+};
+
+// Resize the output, state tensors based on the sizes of the input tensors.
+// Allocate a temporary scratch tensor. Also check that the sizes of the input
+// tensors match each other.
+TfLiteStatus lstm_eval(TfLiteContext* context, const int subgraph_index,
+                       TfLiteNode* node, Logger* logger,
+                       ErrorReporter* error_reporter) {
+  const auto* params = static_cast<TfLiteLSTMParams*>(node->builtin_data);
+
+  const TfLiteTensor* input =
+      GetInput(context, node, ops::builtin::lstm::full::kInputTensor);
+
+  const TfLiteTensor* input_to_input_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kInputToInputWeightsTensor);
+  const TfLiteTensor* input_to_forget_weights = GetInput(
+      context, node, ops::builtin::lstm::full::kInputToForgetWeightsTensor);
+  const TfLiteTensor* input_to_cell_weights = GetInput(
+      context, node, ops::builtin::lstm::full::kInputToCellWeightsTensor);
+  const TfLiteTensor* input_to_output_weights = GetInput(
+      context, node, ops::builtin::lstm::full::kInputToOutputWeightsTensor);
+
+  const TfLiteTensor* recurrent_to_input_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kRecurrentToInputWeightsTensor);
+  const TfLiteTensor* recurrent_to_forget_weights = GetInput(
+      context, node, ops::builtin::lstm::full::kRecurrentToForgetWeightsTensor);
+  const TfLiteTensor* recurrent_to_cell_weights = GetInput(
+      context, node, ops::builtin::lstm::full::kRecurrentToCellWeightsTensor);
+  const TfLiteTensor* recurrent_to_output_weights = GetInput(
+      context, node, ops::builtin::lstm::full::kRecurrentToOutputWeightsTensor);
+
+  const TfLiteTensor* cell_to_input_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kCellToInputWeightsTensor);
+  const TfLiteTensor* cell_to_forget_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kCellToForgetWeightsTensor);
+  const TfLiteTensor* cell_to_output_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kCellToOutputWeightsTensor);
+
+  const TfLiteTensor* input_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kInputLayerNormCoefficientsTensor);
+  const TfLiteTensor* forget_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kForgetLayerNormCoefficientsTensor);
+  const TfLiteTensor* cell_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kCellLayerNormCoefficientsTensor);
+  const TfLiteTensor* output_layer_norm_coefficients = GetOptionalInputTensor(
+      context, node,
+      ops::builtin::lstm::full::kOutputLayerNormCoefficientsTensor);
+
+  const TfLiteTensor* input_gate_bias = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kInputGateBiasTensor);
+  const TfLiteTensor* forget_gate_bias =
+      GetInput(context, node, ops::builtin::lstm::full::kForgetGateBiasTensor);
+  const TfLiteTensor* cell_bias =
+      GetInput(context, node, ops::builtin::lstm::full::kCellGateBiasTensor);
+  const TfLiteTensor* output_gate_bias =
+      GetInput(context, node, ops::builtin::lstm::full::kOutputGateBiasTensor);
+
+  const TfLiteTensor* projection_weights = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kProjectionWeightsTensor);
+  const TfLiteTensor* projection_bias = GetOptionalInputTensor(
+      context, node, ops::builtin::lstm::full::kProjectionBiasTensor);
+
+  // Index the scratch buffers pointers to the global scratch buffer.
+  TfLiteTensor* scratch_buffer = GetTemporary(context, node, /*index=*/0);
+
+  TfLiteTensor* activation_state = GetVariableInput(
+      context, node, ops::builtin::lstm::full::kOutputStateTensor);
+  TF_LITE_ENSURE(context, activation_state != nullptr);
+  TfLiteTensor* cell_state = GetVariableInput(
+      context, node, ops::builtin::lstm::full::kCellStateTensor);
+  TF_LITE_ENSURE(context, cell_state != nullptr);
+
+  TfLiteTensor* output =
+      GetOutput(context, node, ops::builtin::lstm::full::kOutputTensor);
+
+  std::vector<int> intemediate_tensor_indexes(node->intermediates->size);
+  for (int i = 0; i < node->intermediates->size; ++i) {
+    intemediate_tensor_indexes[i] = node->intermediates->data[i];
+  }
+
+  switch (input_to_output_weights->type) {
+    case kTfLiteFloat32: {
+      return EvalFloat(
+          input, input_to_input_weights, input_to_forget_weights,
+          input_to_cell_weights, input_to_output_weights,
+          recurrent_to_input_weights, recurrent_to_forget_weights,
+          recurrent_to_cell_weights, recurrent_to_output_weights,
+          cell_to_input_weights, cell_to_forget_weights, cell_to_output_weights,
+          input_layer_norm_coefficients, forget_layer_norm_coefficients,
+          cell_layer_norm_coefficients, output_layer_norm_coefficients,
+          /*aux_input=*/nullptr,
+          /*aux_input_to_input_weights=*/nullptr,
+          /*aux_input_to_forget_weights=*/nullptr,
+          /*aux_input_to_cell_weights=*/nullptr,
+          /*aux_input_to_output_weights=*/nullptr, input_gate_bias,
+          forget_gate_bias, cell_bias, output_gate_bias, projection_weights,
+          projection_bias, params, /*forward_sequence=*/true,
+          /*time_major=*/true,
+          /*output_offset=*/0, scratch_buffer, activation_state, cell_state,
+          output, logger, intemediate_tensor_indexes, subgraph_index,
+          error_reporter);
+    }
+    case kTfLiteUInt8:
+    case kTfLiteInt8:
+    default:
+      printf("Error. Only float model can be calibrated\n");
+      return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+}  // namespace
+
+TfLiteStatus lstm_logging_kernel(TfLiteContext* context,
+                                 const int subgraph_index, TfLiteNode* node,
+                                 Logger* logger,
+                                 ErrorReporter* error_reporter) {
+  return lstm_eval(context, subgraph_index, node, logger, error_reporter);
+}
+
+}  // namespace custom
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/calibration/custom_logging_ops/lstm.h

@@ -0,0 +1,36 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_CUSTOM_LOGGING_OPS_LSTM_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CUSTOM_LOGGING_OPS_LSTM_H_
+
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_logger.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+namespace custom {
+
+TfLiteStatus lstm_logging_kernel(TfLiteContext* context,
+                                 const int subgraph_index, TfLiteNode* node,
+                                 Logger* logger, ErrorReporter* error_reporter);
+
+}  // namespace custom
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_CUSTOM_LOGGING_OPS_LSTM_H_

tensorflow/lite/tools/optimize/calibration/logging_op.h

@@ -0,0 +1,35 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_LOGGING_OP_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_LOGGING_OP_H_
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_logger.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+typedef TfLiteStatus (*logging_kernel_func_ptr)(TfLiteContext* context,
+                                                const int subgraph_index,
+                                                TfLiteNode* node,
+                                                Logger* logger,
+                                                ErrorReporter* error_reporter);
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_LOGGING_OP_H_

tensorflow/lite/tools/optimize/calibration/logging_op_resolver.cc

@@ -0,0 +1,122 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/api/op_resolver.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_common.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+LoggingOpResolver::LoggingOpResolver(
+    const BuiltinOpsSet& builtin_ops_to_replace,
+    const CustomOpsSet& custom_ops_to_replace, const OpResolver& base_resolver,
+    KernelEvalFuncPtr logging_eval_fn, ErrorReporter* error_reporter) {
+  std::vector<std::string> unresolved_builtin_ops;
+  std::vector<std::string> unresolved_custom_ops;
+
+  for (const auto& op_and_version : builtin_ops_to_replace) {
+    const TfLiteRegistration* base_registration =
+        base_resolver.FindOp(op_and_version.first, op_and_version.second);
+    if (!base_registration) {
+      unresolved_builtin_ops.push_back(
+          EnumNameBuiltinOperator(op_and_version.first));
+      continue;
+    }
+    BuiltinOperatorKey key = op_and_version;
+    builtin_op_evalfn_map_[key] = base_registration->invoke;
+    auto logging_registration =
+        std::make_unique<TfLiteRegistration>(*base_registration);
+    logging_registration->invoke = logging_eval_fn;
+    builtin_op_registration_map_[key] = std::move(logging_registration);
+  }
+  for (const auto& op_and_version : custom_ops_to_replace) {
+    const TfLiteRegistration* base_registration = base_resolver.FindOp(
+        op_and_version.first.c_str(), op_and_version.second);
+    if (!base_registration) {
+      if (!IsFlexOp(op_and_version.first.c_str()))
+        unresolved_custom_ops.push_back(op_and_version.first.c_str());
+      continue;
+    }
+    CustomOperatorKey key = op_and_version;
+    custom_op_evalfn_map_[key] = base_registration->invoke;
+    auto logging_registration =
+        std::make_unique<TfLiteRegistration>(*base_registration);
+    logging_registration->invoke = logging_eval_fn;
+    custom_op_registration_map_[key] = std::move(logging_registration);
+  }
+
+  if (!unresolved_builtin_ops.empty() || !unresolved_custom_ops.empty()) {
+    if (!error_reporter) return;
+    std::string error_message =
+        "Failed to initialize op resolver for calibration:";
+    if (!unresolved_builtin_ops.empty())
+      absl::StrAppend(&error_message, "\nThere are unresolved builtin ops: [",
+                      absl::StrJoin(unresolved_builtin_ops, ", "), "]");
+    if (!unresolved_custom_ops.empty()) {
+      absl::StrAppend(&error_message, "\nThere are unresolved custom ops: [",
+                      absl::StrJoin(unresolved_custom_ops, ", "), "]");
+    }
+    TF_LITE_REPORT_ERROR(error_reporter, error_message.c_str());
+  }
+}
+
+const TfLiteRegistration* LoggingOpResolver::FindOp(BuiltinOperator op,
+                                                    int version) const {
+  BuiltinOperatorKey key = {op, version};
+  if (builtin_op_registration_map_.find(key) !=
+      builtin_op_registration_map_.end()) {
+    return builtin_op_registration_map_.at(key).get();
+  }
+
+  return nullptr;
+}
+
+KernelEvalFuncPtr LoggingOpResolver::GetWrappedKernelInvoke(BuiltinOperator op,
+                                                            int version) const {
+  return builtin_op_evalfn_map_.at({op, version});
+}
+
+const TfLiteRegistration* LoggingOpResolver::FindOp(const char* op,
+                                                    int version) const {
+  CustomOperatorKey key = {op, version};
+  if (custom_op_registration_map_.find(key) !=
+      custom_op_registration_map_.end()) {
+    return custom_op_registration_map_.at(key).get();
+  }
+
+  return nullptr;
+}
+
+KernelEvalFuncPtr LoggingOpResolver::GetWrappedKernelInvoke(const char* op,
+                                                            int version) const {
+  return custom_op_evalfn_map_.at({op, version});
+}
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h

@@ -0,0 +1,68 @@
+/* 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_TOOLS_OPTIMIZE_CALIBRATION_LOGGING_OP_RESOLVER_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_LOGGING_OP_RESOLVER_H_
+
+#include <set>
+#include <unordered_map>
+
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/api/op_resolver.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+#include "tensorflow/lite/op_resolver.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_common.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+
+// A resolver that replaces the kernel invocations with a wrapper
+// eval function.
+class LoggingOpResolver : public OpResolver {
+ public:
+  // Creates an instance of |LoggingOpResolver|.
+  // All |TfLiteRegistration.invoke| functions are replaced by
+  // |logging_eval_fn|.
+  // TODO(shashishekhar): This interface needs to change for
+  // BuiltinOps that need special logging implementations.
+  LoggingOpResolver(const BuiltinOpsSet& builtin_ops_to_replace,
+                    const CustomOpsSet& custom_ops_to_replace,
+                    const OpResolver& base_resolver,
+                    KernelEvalFuncPtr logging_eval_fn,
+                    ErrorReporter* error_reporter);
+
+  const TfLiteRegistration* FindOp(BuiltinOperator op,
+                                   int version) const override;
+  KernelEvalFuncPtr GetWrappedKernelInvoke(BuiltinOperator op,
+                                           int version) const;
+
+  const TfLiteRegistration* FindOp(const char* op, int version) const override;
+  KernelEvalFuncPtr GetWrappedKernelInvoke(const char* op, int version) const;
+
+ private:
+  BuiltinOpsMap<std::unique_ptr<TfLiteRegistration>>
+      builtin_op_registration_map_;
+  BuiltinOpsMap<KernelEvalFuncPtr> builtin_op_evalfn_map_;
+  CustomOpsMap<std::unique_ptr<TfLiteRegistration>> custom_op_registration_map_;
+  CustomOpsMap<KernelEvalFuncPtr> custom_op_evalfn_map_;
+};
+
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_CALIBRATION_LOGGING_OP_RESOLVER_H_

tensorflow/lite/tools/optimize/calibration/logging_op_resolver_test.cc

@@ -0,0 +1,223 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/calibration/logging_op_resolver.h"
+
+#include <string>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/mutable_op_resolver.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/tools/optimize/calibration/calibration_common.h"
+
+namespace tflite {
+namespace optimize {
+namespace calibration {
+namespace {
+
+TfLiteStatus ConvPrepare(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus ConvEval(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus AddPrepare(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus AddEval(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus CustomPrepare(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus CustomEval(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TfLiteStatus WrappingInvoke(TfLiteContext* context, TfLiteNode* node) {
+  return kTfLiteOk;
+}
+
+TEST(LoggingOpResolverTest, KernelInvokesAreReplaced) {
+  MutableOpResolver base_resolver;
+  TfLiteRegistration conv_registration = {};
+  conv_registration.prepare = ConvPrepare;
+  conv_registration.invoke = ConvEval;
+
+  base_resolver.AddBuiltin(BuiltinOperator_CONV_2D, &conv_registration);
+
+  TfLiteRegistration add_registration = {};
+  add_registration.prepare = AddPrepare;
+  add_registration.invoke = AddEval;
+
+  base_resolver.AddBuiltin(BuiltinOperator_ADD, &add_registration);
+  BuiltinOpsSet ops_to_replace = {
+      {BuiltinOperator_CONV_2D, /*version*/ 1},
+      {BuiltinOperator_ADD, /*version*/ 1},
+  };
+
+  LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+
+  auto reg = resolver.FindOp(BuiltinOperator_CONV_2D, 1);
+
+  EXPECT_EQ(reg->builtin_code, BuiltinOperator_CONV_2D);
+  EXPECT_TRUE(reg->prepare == ConvPrepare);
+  EXPECT_TRUE(reg->invoke == WrappingInvoke);
+
+  reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+
+  EXPECT_EQ(reg->builtin_code, BuiltinOperator_ADD);
+  EXPECT_TRUE(reg->prepare == AddPrepare);
+  EXPECT_TRUE(reg->invoke == WrappingInvoke);
+}
+
+TEST(LoggingOpResolverTest, OriginalKernelInvokesAreRetained) {
+  MutableOpResolver base_resolver;
+  TfLiteRegistration conv_registration = {};
+  conv_registration.prepare = ConvPrepare;
+  conv_registration.invoke = ConvEval;
+
+  base_resolver.AddBuiltin(BuiltinOperator_CONV_2D, &conv_registration);
+
+  TfLiteRegistration add_registration = {};
+  add_registration.prepare = AddPrepare;
+  add_registration.invoke = AddEval;
+
+  base_resolver.AddBuiltin(BuiltinOperator_ADD, &add_registration);
+  BuiltinOpsSet ops_to_replace = {
+      {BuiltinOperator_CONV_2D, /*version*/ 1},
+      {BuiltinOperator_ADD, /*version*/ 1},
+  };
+
+  LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+  auto kernel_invoke =
+      resolver.GetWrappedKernelInvoke(BuiltinOperator_CONV_2D, 1);
+  EXPECT_TRUE(kernel_invoke == ConvEval);
+  kernel_invoke = resolver.GetWrappedKernelInvoke(BuiltinOperator_ADD, 1);
+  EXPECT_TRUE(kernel_invoke == AddEval);
+}
+
+TEST(LoggingOpResolverTest, OnlyOpsInReplacementSetAreReplaces) {
+  MutableOpResolver base_resolver;
+  TfLiteRegistration conv_registration = {};
+  conv_registration.prepare = ConvPrepare;
+  conv_registration.invoke = ConvEval;
+
+  base_resolver.AddBuiltin(BuiltinOperator_CONV_2D, &conv_registration);
+
+  TfLiteRegistration add_registration = {};
+  add_registration.prepare = AddPrepare;
+  add_registration.invoke = AddEval;
+
+  base_resolver.AddBuiltin(BuiltinOperator_ADD, &add_registration);
+  // Only replace conv2d
+  BuiltinOpsSet ops_to_replace = {
+      {BuiltinOperator_CONV_2D, /*version*/ 1},
+  };
+
+  LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+  auto reg = resolver.FindOp(BuiltinOperator_CONV_2D, 1);
+  EXPECT_EQ(reg->builtin_code, BuiltinOperator_CONV_2D);
+  EXPECT_TRUE(reg->prepare == ConvPrepare);
+  EXPECT_TRUE(reg->invoke == WrappingInvoke);
+
+  reg = resolver.FindOp(BuiltinOperator_ADD, 1);
+  EXPECT_EQ(nullptr, reg);
+}
+
+TEST(LoggingOpResolverTest, CustomOps) {
+  MutableOpResolver base_resolver;
+  TfLiteRegistration custom_registration = {};
+  custom_registration.prepare = CustomPrepare;
+  custom_registration.invoke = CustomEval;
+
+  std::string custom_op_name = "custom";
+  base_resolver.AddCustom(custom_op_name.c_str(), &custom_registration);
+
+  CustomOpsSet ops_to_replace = {
+      {custom_op_name, /*version*/ 1},
+  };
+
+  LoggingOpResolver resolver(BuiltinOpsSet(), ops_to_replace, base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+
+  auto reg = resolver.FindOp(custom_op_name.c_str(), 1);
+
+  EXPECT_EQ(reg->builtin_code, BuiltinOperator_CUSTOM);
+  EXPECT_EQ(reg->custom_name, custom_op_name.c_str());
+  EXPECT_TRUE(reg->prepare == CustomPrepare);
+  EXPECT_TRUE(reg->invoke == WrappingInvoke);
+}
+
+TEST(LoggingOpResolverTest, UnresolvedCustomOps) {
+  // No custom op registration.
+  MutableOpResolver base_resolver;
+
+  std::string custom_op_name = "unresolved_custom_op";
+
+  CustomOpsSet ops_to_replace = {
+      {custom_op_name, /*version*/ 1},
+  };
+
+  // Expect no death.
+  LoggingOpResolver(BuiltinOpsSet(), ops_to_replace, base_resolver,
+                    WrappingInvoke, /*error_reporter=*/nullptr);
+}
+
+TEST(LoggingOpResolverTest, UnresolvedBuiltinOps) {
+  // No builtin op registration.
+  MutableOpResolver base_resolver;
+
+  BuiltinOpsSet ops_to_replace = {
+      {BuiltinOperator_CONV_2D, /*version*/ 1},
+      {BuiltinOperator_ADD, /*version*/ 1},
+  };
+
+  // Expect no death.
+  LoggingOpResolver resolver(ops_to_replace, CustomOpsSet(), base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+}
+
+TEST(LoggingOpResolverTest, FlexOps) {
+  // No flex op registration.
+  MutableOpResolver base_resolver;
+
+  std::string custom_op_name = "FlexAdd";
+
+  CustomOpsSet ops_to_replace = {
+      {custom_op_name, /*version*/ 1},
+  };
+
+  LoggingOpResolver resolver(BuiltinOpsSet(), ops_to_replace, base_resolver,
+                             WrappingInvoke, /*error_reporter=*/nullptr);
+
+  auto reg = resolver.FindOp(custom_op_name.c_str(), 1);
+
+  EXPECT_TRUE(!reg);
+}
+
+}  // namespace
+}  // namespace calibration
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/debugging/README.md

@@ -0,0 +1,198 @@
+# TensorFlow Lite Quantization Debugger
+
+[TOC]
+
+## Overview
+
+When a quantized model is produced, it requires tedious and manual custom code
+to debug the model in order to:
+
+1. Verify if the quantized model is working as expected (spot errors, check
+   accuracy, etc).
+2. Compare the quantized model and the original float model.
+
+This is now feasible using the TensorFlow Lite Quantization Debugger, as shown
+below.
+
+Note: Currently, this workflow is only supported for full integer (int8)
+quantization. The debug model produced using this workflow should only be used
+for debugging purposes only (and not for inference).
+
+## Analysis with quantized model only
+
+### Produce a debug model
+
+Modify the
+[TFLite full integer (int8) quantization steps](https://www.tensorflow.org/lite/performance/post_training_quantization#full_integer_quantization)
+as shown below to produce a debug model (used for debugging purposes only, and
+not inference)
+
+#### How does this work?
+
+With the help of the MLIR quantizer's debug mode feature, the debug model
+produced has both the original float operators (or ops) and the quantized ops.
+Additionally, `NumericVerify` ops are added to compare the outputs of the
+original float and quantized ops and to also collect statistics. It has the name
+in the format of `NumericVerify/{original tensor name}:{original tensor id}`
+
+```python
+# for mlir_quantize
+from tensorflow.lite.python import convert
+
+# set full-integer quantization parameters as usual.
+converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+converter.representative_dataset = calibration_gen
+
+# Create a TFLite model with new quantizer and numeric verify ops. Rather than
+# calling convert() only, calibrate model first and call `mlir_quantize` to run
+# the actual quantization, with `enable_numeric_verify` set to `True`.
+converter.optimizations = [tf.lite.Optimize.DEFAULT]
+converter._experimental_calibrate_only = True
+calibrated = converter.convert()
+return convert.mlir_quantize(calibrated, enable_numeric_verify=True)
+```
+
+### Run debugger with debug model
+
+Initialize debugger with the debug model. This can be done in two ways.
+
+```python
+from tensorflow.lite.tools.optimize.debugging.python import debugger
+
+# `debug_dataset` accpets the same type as `converter.representative_dataset`.
+quant_debugger = debugger.QuantizationDebugger(
+    quant_debug_model_content=quant_debug_model,
+    debug_dataset=data_gen)
+
+# OR
+
+quant_debugger = debugger.QuantizationDebugger(
+    quant_debug_model_path='/path/to/debug_model.tflite',
+    debug_dataset=data_gen)
+
+quant_debugger.run()
+```
+
+### Inspect statistics
+
+When you call `quant_debugger.run()`, `quant_debugger.layer_statistics` is
+filled with aggregated statistics for each `NumericVerify` ops. Some metrics
+(i.e. stddev, mean square error) are calculated by default.
+
+#### Example output
+```python
+# `quant_debugger.layer_statistics.metrics` is defaultdict, convert it to dict
+# for readable output.
+import pprint
+for layer_name, metrics in quant_debugger.layer_statistics.items():
+  print(layer_name)
+  pprint.pprint(dict(metrics))
+```
+
+```python
+# ...
+
+NumericVerify/sequential/dense/MatMul;sequential/dense/BiasAdd3:77
+{'max_abs_error': 0.05089309,
+ 'mean_error': -0.00017149668,
+ 'mean_squared_error': 0.00040816222,
+ 'num_elements': 256.0,
+ 'stddev': 0.02009948}
+NumericVerify/sequential/dense_1/MatMul;sequential/dense_1/BiasAdd3:81
+{'max_abs_error': 0.09744112,
+ 'mean_error': 0.0048679365,
+ 'mean_squared_error': 0.0036721828,
+ 'num_elements': 10.0,
+ 'stddev': 0.055745363}
+NumericVerify/Identity2:85
+{'max_abs_error': 0.0036417267,
+ 'mean_error': -0.00068773015,
+ 'mean_squared_error': 3.439951e-06,
+ 'num_elements': 10.0,
+ 'stddev': 0.0016223773}
+
+# ...
+```
+
+## Adding custom metrics
+
+More metrics can be added by passing `QuantizationDebugOptions` to the
+initializer. For example, if you want to add mean absolute error, use following
+snippet.
+
+```python
+debug_options = debugger.QuantizationDebugOptions(
+    layer_debug_metrics={
+        'mean_abs_error': lambda diffs: np.mean(np.abs(diffs))
+    })
+
+quant_debugger = debugger.QuantizationDebugger(
+    quant_debug_model_content=quant_debug_model,
+    debug_dataset=data_gen,
+    debug_options=debug_options
+)
+quant_debugger.run()
+```
+
+Now `quant_debugger.layer_statistics` includes mean absoulte error for each
+layer.
+
+## Analysis with float and quantized models
+
+In addition to single model analysis, the output of original float model and
+quantized model can be compared when both models are given. This can be done
+by providing a float model, and metrics to compare outputs. This can be `argmax`
+for classification models, bit for more complex models like detection more
+complicated logic should be given.
+
+```python
+# functions for model_debug_metrics gets all output tensors from float and
+# quantized models, and returns a single metric value.
+debug_options = debugger.QuantizationDebugOptions(
+    model_debug_metrics={
+        'argmax_accuracy': lambda f, q: np.argmax(f[0]) == np.argmax(q[0])
+    })
+
+float_model = converter.convert()  # converted without any optimizations.
+
+quant_debugger = debugger.QuantizationDebugger(
+    quant_debug_model_content=quant_debug_model,
+    float_model_content=float_model,  # can pass `float_model_path` instead.
+    debug_dataset=data_gen,
+    debug_options=debug_options
+)
+quant_debugger.run()
+```
+
+The result is a single number per metric, so it's easier to inspect.
+
+```python
+>>> quant_debugger.model_statistics
+{'argmax_accuracy': 0.89}
+```
+
+## Advanced usage: Export stats to csv, and import to pandas
+
+`quant_debugger.layer_statistics_dump` function accepts file-like object, and
+exports layer statistics to csv. This can be imported to other tools like
+`pandas` for further processing. The exported data also has name of the op,
+originating tensor ID, and quantization parameters (scales and zero points) for
+quantized layer.
+
+Note: scales and zero points are lists, and imported to `pandas` as text by
+default. Additional processing to parse them is required before processing.
+
+```python
+import pandas as pd
+import yaml  # used to parse lists
+
+with open('/path/to/stats.csv', 'w') as f:
+  quant_debugger.layer_statistics_dump(f)
+
+data = pd.read_csv(
+    '/path/to/stats.csv',
+    converters={
+        'scales': yaml.safe_load,
+        'zero_points': yaml.safe_load
+    })
+```

tensorflow/lite/tools/optimize/debugging/python/BUILD

@@ -0,0 +1,43 @@
+# QuantizationDebugger for TFLite accuracy tooling.
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "debugger",
+    srcs = ["debugger.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/lite/python:convert",
+        "//tensorflow/lite/python:interpreter",
+        "//tensorflow/lite/python/metrics",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_test(
+    name = "debugger_test",
+    srcs = [
+        "debugger_test.py",
+    ],
+    exec_properties = {"cpp_link.mem": "16g"},
+    strict_deps = True,
+    deps = [
+        ":debugger",
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/lite/python:convert",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/lite/python/metrics",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/trackable:autotrackable",
+    ],
+)

tensorflow/lite/tools/optimize/debugging/python/debugger.py

@@ -0,0 +1,549 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Python TF-Lite QuantizationDebugger."""
+import collections
+import csv
+import re
+from typing import (Any, Callable, Dict, IO, Iterable, List, Mapping, Optional,
+                    Sequence, Tuple)
+
+import numpy as np
+
+from tensorflow.lite.python import convert
+from tensorflow.lite.python import interpreter as _interpreter
+from tensorflow.lite.python.metrics import metrics as metrics_stub  # type: ignore
+from tensorflow.python.util import tf_export
+
+
+# TODO(b/198099651): move converter implementation out of lite.py
+TFLiteConverter = Any  # importing tf.lite creates circular dependency
+
+# Returns metrics based on difference of values for quantized/float ops.
+_DEFAULT_LAYER_DEBUG_METRICS = {
+    'num_elements': lambda diffs: diffs.size,
+    'stddev': np.std,
+    'mean_error': np.average,
+    'max_abs_error': lambda diffs: np.max(np.abs(diffs)),
+    'mean_squared_error': lambda diffs: np.average(diffs**2),
+}
+
+_NUMERIC_VERIFY_OP_NAME = 'NumericVerify'
+
+
+def _get_quant_params(
+    tensor_detail: Mapping[str, Any]) -> Optional[Tuple[float, int]]:
+  """Returns first scale and zero point from tensor detail, if present."""
+  quant_params = tensor_detail['quantization_parameters']
+  if not quant_params:
+    return None
+  if quant_params['scales'] and quant_params['zero_points']:
+    return (quant_params['scales'][0], quant_params['zero_points'][0])
+  return None
+
+
+@tf_export.tf_export('lite.experimental.QuantizationDebugOptions')
+class QuantizationDebugOptions:
+  """Debug options to set up a given QuantizationDebugger."""
+
+  def __init__(self,
+               layer_debug_metrics: Optional[Mapping[str,
+                                                     Callable[[np.ndarray],
+                                                              float]]] = None,
+               model_debug_metrics: Optional[Mapping[
+                   str, Callable[[Sequence[np.ndarray], Sequence[np.ndarray]],
+                                 float]]] = None,
+               layer_direct_compare_metrics: Optional[Mapping[str, Callable[
+                   [Sequence[np.ndarray], Sequence[np.ndarray], float, int],
+                   float]]] = None,
+               denylisted_ops: Optional[List[str]] = None,
+               denylisted_nodes: Optional[List[str]] = None,
+               fully_quantize: bool = False) -> None:
+    """Initializes debugger options.
+
+    Args:
+      layer_debug_metrics: a dict to specify layer debug functions
+        {function_name_str: function} where the function accepts result of
+          NumericVerify Op, which is value difference between float and
+          dequantized op results. The function returns single scalar value.
+      model_debug_metrics: a dict to specify model debug functions
+        {function_name_str: function} where the function accepts outputs from
+          two models, and returns single scalar value for a metric. (e.g.
+          accuracy, IoU)
+      layer_direct_compare_metrics: a dict to specify layer debug functions
+        {function_name_str: function}. The signature is different from that of
+          `layer_debug_metrics`, and this one gets passed (original float value,
+          original quantized value, scale, zero point). The function's
+          implementation is responsible for correctly dequantize the quantized
+          value to compare. Use this one when comparing diff is not enough.
+          (Note) quantized value is passed as int8, so cast to int32 is needed.
+      denylisted_ops: a list of op names which is expected to be removed from
+        quantization.
+      denylisted_nodes: a list of op's output tensor names to be removed from
+        quantization.
+      fully_quantize: Bool indicating whether to fully quantize the model.
+        Besides model body, the input/output will be quantized as well.
+        Corresponding to mlir_quantize's fully_quantize parameter.
+
+    Raises:
+      ValueError: when there are duplicate keys
+    """
+    self.layer_debug_metrics = layer_debug_metrics
+    self.model_debug_metrics = model_debug_metrics
+    self.layer_direct_compare_metrics = layer_direct_compare_metrics
+
+    keys = []
+    for metrics in [
+        layer_debug_metrics, model_debug_metrics, layer_direct_compare_metrics
+    ]:
+      if metrics is not None:
+        keys.extend(metrics.keys())
+    if len(keys) != len(set(keys)):
+      raise ValueError('Provided metrics have duplicate keys.')
+
+    self.denylisted_ops = denylisted_ops
+    self.denylisted_nodes = denylisted_nodes
+    self.fully_quantize = fully_quantize
+
+
+@tf_export.tf_export('lite.experimental.QuantizationDebugger')
+class QuantizationDebugger:
+  """Debugger for Quantized TensorFlow Lite debug mode models.
+
+  This can run the TensorFlow Lite converted models equipped with debug ops and
+  collect debug information. This debugger calculates statistics from
+  user-defined post-processing functions as well as default ones.
+  """
+
+  def __init__(self,
+               quant_debug_model_path: Optional[str] = None,
+               quant_debug_model_content: Optional[bytes] = None,
+               float_model_path: Optional[str] = None,
+               float_model_content: Optional[bytes] = None,
+               debug_dataset: Optional[Callable[
+                   [], Iterable[Sequence[np.ndarray]]]] = None,
+               debug_options: Optional[QuantizationDebugOptions] = None,
+               converter: Optional[TFLiteConverter] = None) -> None:
+    """Runs the TFLite debugging model with given debug options.
+
+    Args:
+      quant_debug_model_path: Path to the quantized debug TFLite model file.
+      quant_debug_model_content: Content of the quantized debug TFLite model.
+      float_model_path: Path to float TFLite model file.
+      float_model_content: Content of the float TFLite model.
+      debug_dataset: a factory function that returns dataset generator which is
+        used to generate input samples (list of np.ndarray) for the model. The
+        generated elements must have same types and shape as inputs to the
+        model.
+      debug_options: Debug options to debug the given model.
+      converter: Optional, use converter instead of quantized model.
+
+    Raises:
+      ValueError: If the debugger was unable to be created.
+
+    Attributes:
+      layer_statistics: results of error metrics for each NumericVerify op
+        results. in {layer_name: {metric_name: metric}} format.
+      model_statistics: results of error metrics for difference between float
+        and quantized models. in {metric_name: metric} format.
+    """
+    self._data_gen = debug_dataset
+    self._debug_options = debug_options or QuantizationDebugOptions()
+    self.converter = None
+    self.calibrated_model = None
+    self.float_model = None
+    self._float_interpreter = None
+    if converter is not None:
+      if self._debug_options.model_debug_metrics:
+        old_optimizations = converter.optimizations
+        self.converter = self._set_converter_options_for_float(converter)
+        self.float_model = self.converter.convert()
+        converter.optimizations = old_optimizations
+
+      self.converter = self._set_converter_options_for_calibration(converter)
+      self.calibrated_model = self.converter.convert()
+      # Converter should be already set up with all options
+      self._init_from_converter(
+          self._debug_options,
+          self.converter,
+          self.calibrated_model,
+          float_model=self.float_model)
+    else:
+      self._quant_interpreter = _interpreter.Interpreter(
+          quant_debug_model_path,
+          quant_debug_model_content,
+          experimental_preserve_all_tensors=(
+              self._debug_options.layer_direct_compare_metrics is not None))
+      if self._debug_options.model_debug_metrics:
+        self._float_interpreter = _interpreter.Interpreter(
+            float_model_path, float_model_content)
+    self._initialize_stats()
+
+  @property
+  def options(self) -> QuantizationDebugOptions:
+    return self._debug_options
+
+  @options.setter
+  def options(self, options: QuantizationDebugOptions) -> None:
+    self._debug_options = options
+    if not self.converter or not self.calibrated_model:
+      return
+    self._init_from_converter(
+        self._debug_options,
+        self.converter,
+        self.calibrated_model,
+        float_model=self.float_model)
+    self._initialize_stats()
+
+  def _initialize_stats(self):
+    """Helper function initializes stats."""
+    # TODO(b/177749613) : Fix the dependency on tf.lite._get_ops_details()
+    # Following code is needed to get op's name from the output tensor index,
+    # since NumericVerify op only provides its quantized input tensor index.
+    self._defining_op = dict()
+    for op_info in self._quant_interpreter._get_ops_details():  # pylint: disable=protected-access
+      self._defining_op.update(
+          {tensor_idx: op_info['index'] for tensor_idx in op_info['outputs']})
+
+    self._numeric_verify_tensor_details = None
+    self._numeric_verify_op_details = None
+    if not self._get_numeric_verify_tensor_details():
+      raise ValueError('Please check if the quantized model is in debug mode')
+
+    self._layer_debug_metrics = _DEFAULT_LAYER_DEBUG_METRICS.copy()
+    if self._debug_options.layer_debug_metrics:
+      self._layer_debug_metrics.update(self._debug_options.layer_debug_metrics)
+
+    self.layer_statistics = None
+    self.model_statistics = None
+
+    self._metrics = metrics_stub.TFLiteMetrics()
+    self._metrics.increase_counter_debugger_creation()
+
+  def _get_quantized_model(self, is_debug: bool) -> bytes:
+    if not self.converter:
+      raise ValueError('No converter found, use this function with the '
+                       'converter option in the constructor.')
+
+    return convert.mlir_quantize(
+        self.calibrated_model,
+        disable_per_channel=self.converter._experimental_disable_per_channel,  # pylint: disable=protected-access
+        fully_quantize=self._debug_options.fully_quantize,
+        enable_numeric_verify=is_debug,
+        denylisted_ops=self._debug_options.denylisted_ops,
+        denylisted_nodes=self._debug_options.denylisted_nodes)
+
+  def get_nondebug_quantized_model(self) -> bytes:
+    """Returns a non-instrumented quantized model.
+
+    Convert the quantized model with the initialized converter and
+    return bytes for nondebug model. The model will not be instrumented with
+    numeric verification operations.
+
+    Returns:
+      Model bytes corresponding to the model.
+    Raises:
+      ValueError: if converter is not passed to the debugger.
+    """
+    return self._get_quantized_model(is_debug=False)
+
+  def get_debug_quantized_model(self) -> bytes:
+    """Returns an instrumented quantized model.
+
+    Convert the quantized model with the initialized converter and
+    return bytes for model. The model will be instrumented with numeric
+    verification operations and should only be used for debugging.
+
+    Returns:
+      Model bytes corresponding to the model.
+    Raises:
+      ValueError: if converter is not passed to the debugger.
+    """
+    return self._get_quantized_model(is_debug=True)
+
+  def _init_from_converter(self,
+                           options: QuantizationDebugOptions,
+                           converter: TFLiteConverter,
+                           calibrated_model: Optional[bytes] = None,
+                           float_model: Optional[bytes] = None) -> None:
+    """Convert the model and apply options.
+
+    Converts the quantized model and initializes a quantized model interpreter
+    with the quantized model. Returns a float model interpreter if float model
+    is provided.
+
+    Args:
+      options: a QuantizationDebugOptions object.
+      converter: an initialized tf.lite.TFLiteConverter.
+      calibrated_model: Calibrated model bytes.
+      float_model: Float model bytes.
+    """
+    self.quant_model = convert.mlir_quantize(
+        calibrated_model,
+        disable_per_channel=converter._experimental_disable_per_channel,  # pylint: disable=protected-access
+        fully_quantize=options.fully_quantize,
+        enable_numeric_verify=True,
+        denylisted_ops=options.denylisted_ops,
+        denylisted_nodes=options.denylisted_nodes)
+    self._quant_interpreter = _interpreter.Interpreter(
+        model_content=self.quant_model)
+    self._float_interpreter = None
+    if float_model is not None:
+      self._float_interpreter = _interpreter.Interpreter(
+          model_content=float_model)
+
+  def _set_converter_options_for_float(
+      self, converter: TFLiteConverter) -> TFLiteConverter:
+    """Verify converter options and set required experimental options."""
+    if converter.optimizations:
+      converter.optimizations = []
+    return converter
+
+  def _set_converter_options_for_calibration(
+      self, converter: TFLiteConverter) -> TFLiteConverter:
+    """Verify converter options and set required experimental options."""
+    if not converter.optimizations:
+      raise ValueError(
+          'converter object must set optimizations to lite.Optimize.DEFAULT')
+    if not converter.representative_dataset:
+      raise ValueError('converter object must set representative_dataset')
+
+    converter.experimental_mlir_quantizer = True
+    converter._experimental_calibrate_only = True  # pylint: disable=protected-access
+    return converter
+
+  def run(self) -> None:
+    """Runs models and gets metrics."""
+    self.layer_statistics = self._collect_layer_statistics()
+    if self._debug_options.model_debug_metrics:
+      self.model_statistics = self._collect_model_statistics()
+
+  def _collect_layer_statistics(self) -> Dict[str, Dict[str, float]]:
+    """Collects layer statistics by applying layer debug metrics.
+
+    For all data from the given RepresentativeDataset, collect statistics per
+    example by getting the NumericVerify op results in _quant_interpreter
+    and calculating layer debug metrics on the results.
+
+    Returns:
+      aggregated per-layer statistics of NumericVerify results.
+      {layer_name: {metric_name: metric}}
+    """
+    layer_statistics = collections.defaultdict(
+        lambda: collections.defaultdict(list))
+
+    initialize = True
+    for tensor_data in self._data_gen():
+      self._set_input_tensors(self._quant_interpreter, tensor_data, initialize)
+      initialize = False
+
+      # Run the model.
+      self._quant_interpreter.invoke()
+
+      # Collect the statistics of this invoke result.
+      for tensor_detail in self._get_numeric_verify_tensor_details():
+        tensor_name = tensor_detail['name']  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+        diffs = self._quant_interpreter.get_tensor(tensor_detail['index'])  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+        for metric_name, metric_fn in self._layer_debug_metrics.items():
+          layer_statistics[tensor_name][metric_name].append(metric_fn(diffs))
+
+      if self._debug_options.layer_direct_compare_metrics is not None:
+        for tensor_detail in self._get_numeric_verify_tensor_details():
+          tensor_name = tensor_detail['name']  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+          op_idx = self._defining_op[tensor_detail['index']]  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+          op_detail = self._quant_interpreter._get_op_details(op_idx)  # pylint: disable=protected-access
+          q_idx, f_idx = op_detail['inputs']
+          quant_input_detail = self._quant_interpreter._get_tensor_details(  # pylint: disable=protected-access
+              q_idx, subgraph_index=0)
+          for (metric_name, metric_fn
+              ) in self._debug_options.layer_direct_compare_metrics.items():
+            layer_statistics[tensor_name][metric_name].append(
+                metric_fn(
+                    self._quant_interpreter.get_tensor(f_idx),
+                    self._quant_interpreter.get_tensor(q_idx),
+                    quant_input_detail['quantization_parameters']['scales'][0],
+                    quant_input_detail['quantization_parameters']['zero_points']
+                    [0]))
+
+    # Calculate final aggregated metrics for each layer.
+    for metrics in layer_statistics.values():
+      for metric_name in metrics:
+        metrics[metric_name] = np.nanmean(metrics[metric_name])
+
+    return layer_statistics
+
+  def _collect_model_statistics(self) -> Dict[str, float]:
+    """Collects model output metrics.
+
+    For all data from the given RepresentativeDataset, collect all model output
+    results from float model & quantized debug model, and calculate metrics
+    by using model output functions. As a result, self.model_results is filled,
+
+    where self.model_results[model_output_function_name] = `aggregated model
+    output function value` (a scalar).
+
+    Returns:
+      aggregated per-model output discrepancy metrics.
+      {metric_name: aggregated_metric}
+    """
+
+    model_statistics = collections.defaultdict(list)
+
+    initialize = True
+    for tensor_data in self._data_gen():
+      # Run quantized debug model and collect output results.
+      self._set_input_tensors(self._quant_interpreter, tensor_data, initialize)
+      self._quant_interpreter.invoke()
+      quant_tensor_data = self._get_output_tensors(self._quant_interpreter)
+
+      # Run float model if it's initialized.
+      float_tensor_data = []
+      if self._float_interpreter:
+        self._set_input_tensors(
+            self._float_interpreter, tensor_data, initialize)
+        self._float_interpreter.invoke()
+        float_tensor_data = self._get_output_tensors(self._float_interpreter)
+
+      initialize = False
+
+      # Calculate the metrics.
+      for (metric_name,
+           metric_fn) in self._debug_options.model_debug_metrics.items():
+        model_statistics[metric_name].append(
+            metric_fn(float_tensor_data, quant_tensor_data))
+
+    # Calculate final aggregated metrics for each outputs.
+    return {
+        metric_name: np.mean(metric)
+        for metric_name, metric in model_statistics.items()
+    }
+
+  def _set_input_tensors(self, interpreter: _interpreter.Interpreter,
+                         tensor_data: Sequence[np.ndarray],
+                         initialize: bool) -> None:
+    """Sets input tensors into TFLite model Interpreter.
+
+    Args:
+      interpreter: a tf.lite.Interpreter object with allocated tensors.
+      tensor_data: a list of Numpy array data.
+      initialize: set to true when input is first set for the interpreter, to
+        set input shapes and allocate tensors.
+
+    Raises:
+      ValueError: when inputs can't be set, or size of provided inputs does not
+      match size of model inputs.
+    """
+    input_details = interpreter.get_input_details()
+    if len(input_details) != len(tensor_data):
+      raise ValueError(
+          'Number of inputs provided ({}) does not match number of inputs to '
+          'the model ({})'.format(len(tensor_data), len(input_details)))
+
+    if initialize:
+      for input_detail, tensor in zip(input_details, tensor_data):
+        interpreter.resize_tensor_input(input_detail['index'], tensor.shape)
+      interpreter.allocate_tensors()
+
+    for input_detail, tensor in zip(input_details, tensor_data):
+      if tensor.dtype == np.float32 and input_detail['dtype'] == np.int8:
+        quant_params = _get_quant_params(input_detail)
+        if quant_params:
+          scale, zero_point = quant_params
+          tensor = np.round((tensor / scale) + zero_point).astype(np.int8)
+      interpreter.set_tensor(input_detail['index'], tensor)
+
+  def _get_output_tensors(
+      self, interpreter: _interpreter.Interpreter) -> List[np.ndarray]:
+    """Returns output tensors of given TFLite model Interpreter.
+
+    Args:
+      interpreter: a tf.lite.Interpreter object with allocated tensors.
+
+    Returns:
+      a list of numpy arrays representing output tensor results.
+    """
+
+    outputs = []
+    for output_detail in interpreter.get_output_details():
+      tensor = interpreter.get_tensor(output_detail['index'])
+      if output_detail['dtype'] == np.int8:
+        quant_params = _get_quant_params(output_detail)
+        if quant_params:
+          scale, zero_point = quant_params
+          tensor = ((tensor.astype(np.float32) - zero_point) * scale).astype(
+              np.float32)
+      outputs.append(tensor)
+
+    return outputs
+
+  def _get_numeric_verify_tensor_details(self) -> List[str]:
+    """Returns all names of all tensors from NumericVerify op."""
+    # pylint: disable=protected-access
+    if not self._numeric_verify_tensor_details:
+      self._numeric_verify_tensor_details = []
+      self._numeric_verify_op_details = {}
+      for op_info in self._quant_interpreter._get_ops_details():
+        if op_info['op_name'] == _NUMERIC_VERIFY_OP_NAME:
+          self._numeric_verify_tensor_details.append(
+              self._quant_interpreter._get_tensor_details(
+                  op_info['outputs'][0], subgraph_index=0))
+          tensor_name = self._numeric_verify_tensor_details[-1]['name']
+          self._numeric_verify_op_details[tensor_name] = op_info
+    # pylint: enable=protected-access
+    return self._numeric_verify_tensor_details
+
+  def _get_operand_name_and_index(self,
+                                  numeric_verify_name: str) -> Tuple[str, int]:
+    """Gets the index and name of NumericVerify Op's quantized input tensor.
+
+    Args:
+      numeric_verify_name: name of the NumericVerify op's output tensor. It has
+        format of `NumericVerify/{quantized_tensor_name}:{quantized_tensor_idx}`
+
+    Returns:
+      Tuple of (tensor_name, tensor_idx) for quantized op's output tensor.
+    """
+    tensor_name, tensor_idx = numeric_verify_name.rsplit(':', 1)
+    float_tensor_name = tensor_name[len(_NUMERIC_VERIFY_OP_NAME) + 1:]
+    if re.match(r'\d', float_tensor_name[-1]):
+      float_tensor_name = float_tensor_name[:-1]
+
+    return (float_tensor_name, int(tensor_idx))
+
+  def layer_statistics_dump(self, file: IO[str]) -> None:
+    """Dumps layer statistics into file, in csv format.
+
+    Args:
+      file: file, or file-like object to write.
+    """
+    # order of `fields` is the order of fields in csv.
+    fields = ['op_name', 'tensor_idx'] + list(self._layer_debug_metrics.keys())
+    if self._debug_options.layer_direct_compare_metrics is not None:
+      fields += list(self._debug_options.layer_direct_compare_metrics.keys())
+    fields += ['scale', 'zero_point', 'tensor_name']
+    writer = csv.DictWriter(file, fields)
+    writer.writeheader()
+    if self.layer_statistics:
+      for name, metrics in self.layer_statistics.items():
+        data = metrics.copy()
+        (data['tensor_name'], _) = self._get_operand_name_and_index(name)
+        data['tensor_idx'] = self._numeric_verify_op_details[name]['inputs'][0]
+        data['op_name'] = self._quant_interpreter._get_op_details(  # pylint: disable=protected-access
+            self._defining_op[data['tensor_idx']])['op_name']
+        details = self._quant_interpreter._get_tensor_details(  # pylint: disable=protected-access
+            data['tensor_idx'], subgraph_index=0)
+        data['scale'], data['zero_point'] = (
+            details['quantization_parameters']['scales'][0],
+            details['quantization_parameters']['zero_points'][0])
+        writer.writerow(data)

tensorflow/lite/tools/optimize/debugging/python/debugger_test.py

@@ -0,0 +1,437 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for QuantizationDebugger."""
+
+import csv
+import io
+import re
+from unittest import mock
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.lite.python import convert
+from tensorflow.lite.python import lite
+from tensorflow.lite.python.metrics import metrics
+from tensorflow.lite.tools.optimize.debugging.python import debugger
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.trackable import autotrackable
+
+
+def _get_model():
+  """Returns somple model with Conv2D and representative dataset gen."""
+  root = autotrackable.AutoTrackable()
+  kernel_in = np.array([-2, -1, 1, 2], dtype=np.float32).reshape((2, 2, 1, 1))
+
+  @tf.function(
+      input_signature=[tf.TensorSpec(shape=[1, 3, 3, 1], dtype=tf.float32)])
+  def func(inp):
+    kernel = tf.constant(kernel_in, dtype=tf.float32)
+    conv = tf.nn.conv2d(inp, kernel, strides=1, padding='SAME')
+    output = tf.nn.relu(conv, name='output')
+    return output
+
+  root.f = func
+  to_save = root.f.get_concrete_function()
+  return (root, to_save)
+
+
+def _calibration_gen():
+  for i in range(5):
+    yield [np.arange(9).reshape((1, 3, 3, 1)).astype(np.float32) * i]
+
+
+def _convert_model(model, func):
+  """Converts TF model to TFLite float model."""
+  converter = lite.TFLiteConverterV2.from_concrete_functions([func], model)
+  # TODO(b/191205988): Explicitly disable saved model lowering in conversion.
+  converter.experimental_lower_to_saved_model = False
+  return converter.convert()
+
+
+def _quantize_converter(model, func, calibration_gen, debug=True):
+  """Returns a converter appropriate for the function and debug configs."""
+  converter = lite.TFLiteConverterV2.from_concrete_functions([func], model)
+  converter.target_spec.supported_ops = [lite.OpsSet.TFLITE_BUILTINS_INT8]
+  converter.representative_dataset = calibration_gen
+
+  # TODO(b/191205988): Explicitly disable saved model lowering in conversion.
+  converter.experimental_lower_to_saved_model = False
+
+  # Create a TFLite model with new quantizer and numeric verify ops.
+  converter.optimizations = [lite.Optimize.DEFAULT]
+  converter.experimental_new_quantizer = True
+  if debug:
+    converter._experimental_calibrate_only = True
+  return converter
+
+
+def _quantize_model(model,
+                    func,
+                    calibration_gen,
+                    quantized_io=False,
+                    debug=True):
+  """Quantizes model, in debug or normal mode."""
+  converter = _quantize_converter(model, func, calibration_gen, debug)
+  if debug:
+    calibrated = converter.convert()
+    return convert.mlir_quantize(
+        calibrated, enable_numeric_verify=True, fully_quantize=quantized_io)
+  else:
+    return converter.convert()
+
+
+def _dummy_fn(*unused_args):
+  return 0.0
+
+
+class QuantizationDebugOptionsTest(test_util.TensorFlowTestCase,
+                                   parameterized.TestCase):
+
+  @test_util.run_v2_only
+  def test_init_duplicate_keys_raises_ValueError(self):
+    with self.assertRaises(ValueError):
+      debugger.QuantizationDebugOptions(
+          layer_debug_metrics={
+              'a': _dummy_fn,
+              'b': _dummy_fn
+          },
+          model_debug_metrics={
+              'c': _dummy_fn,
+              'd': _dummy_fn
+          },
+          layer_direct_compare_metrics={
+              'a': _dummy_fn,
+              'e': _dummy_fn
+          })
+
+    with self.assertRaises(ValueError):
+      debugger.QuantizationDebugOptions(
+          layer_debug_metrics={
+              'a': _dummy_fn,
+              'b': _dummy_fn
+          },
+          layer_direct_compare_metrics={
+              'a': _dummy_fn,
+              'e': _dummy_fn
+          })
+
+
+class QuantizationDebuggerTest(test_util.TensorFlowTestCase,
+                               parameterized.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.tf_model_root, cls.tf_model = _get_model()
+    cls.float_model = _convert_model(cls.tf_model_root, cls.tf_model)
+    cls.debug_model_float = _quantize_model(
+        cls.tf_model_root, cls.tf_model, _calibration_gen, quantized_io=False)
+    cls.debug_model_int8 = _quantize_model(
+        cls.tf_model_root, cls.tf_model, _calibration_gen, quantized_io=True)
+
+  @parameterized.named_parameters(
+      ('float_io', False, False),
+      ('quantized_io', True, False),
+      ('float_io_from_converter', False, True),
+      ('quantized_io_from_converter', True, True),
+  )
+  @test_util.run_v2_only
+  def test_layer_metrics(self, quantized_io, from_converter):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))})
+    if not from_converter:
+      if quantized_io:
+        debug_model = QuantizationDebuggerTest.debug_model_int8
+      else:
+        debug_model = QuantizationDebuggerTest.debug_model_float
+      quant_debugger = debugger.QuantizationDebugger(
+          quant_debug_model_content=debug_model,
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+    else:
+      options.fully_quantize = quantized_io
+      quant_debugger = debugger.QuantizationDebugger(
+          converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                        _calibration_gen),
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+
+    quant_debugger.run()
+
+    expected_quant_io_metrics = {
+        'num_elements': 9,
+        'stddev': 0.03850026,
+        'mean_error': 0.01673192,
+        'max_abs_error': 0.10039272,
+        'mean_squared_error': 0.0027558778,
+        'l1_norm': 0.023704167,
+    }
+    expected_float_io_metrics = {
+        'num_elements': 9,
+        'stddev': 0.050998904,
+        'mean_error': 0.007843441,
+        'max_abs_error': 0.105881885,
+        'mean_squared_error': 0.004357292,
+        'l1_norm': 0.035729896,
+    }
+    expected_metrics = (
+        expected_quant_io_metrics
+        if quantized_io else expected_float_io_metrics)
+    self.assertLen(quant_debugger.layer_statistics, 1)
+    actual_metrics = next(iter(quant_debugger.layer_statistics.values()))
+
+    self.assertCountEqual(expected_metrics.keys(), actual_metrics.keys())
+    for key, value in expected_metrics.items():
+      self.assertAlmostEqual(value, actual_metrics[key], places=5)
+
+    buffer = io.StringIO()
+    quant_debugger.layer_statistics_dump(buffer)
+    reader = csv.DictReader(buffer.getvalue().split())
+    actual_values = next(iter(reader))
+
+    expected_values = expected_metrics.copy()
+    expected_values.update({
+        'op_name': 'CONV_2D',
+        'tensor_idx': 7,
+        'scale': 0.15686275,
+        'zero_point': -128,
+        'tensor_name': r'Identity[1-9]?$'
+    })
+    for key, value in expected_values.items():
+      if isinstance(value, str):
+        self.assertIsNotNone(
+            re.match(value, actual_values[key]),
+            'String is different from expected string. Please fix test code if'
+            " it's being affected by graph manipulation changes.")
+      elif isinstance(value, list):
+        self.assertAlmostEqual(
+            value[0], float(actual_values[key][1:-1]), places=5)
+      else:
+        self.assertAlmostEqual(value, float(actual_values[key]), places=5)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_model_metrics(self, quantized_io):
+    if quantized_io:
+      debug_model = QuantizationDebuggerTest.debug_model_int8
+    else:
+      debug_model = QuantizationDebuggerTest.debug_model_float
+    options = debugger.QuantizationDebugOptions(
+        model_debug_metrics={'stdev': lambda x, y: np.std(x[0] - y[0])})
+    quant_debugger = debugger.QuantizationDebugger(
+        quant_debug_model_content=debug_model,
+        float_model_content=QuantizationDebuggerTest.float_model,
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+    quant_debugger.run()
+
+    expected_metrics = {'stdev': 0.050998904}
+    actual_metrics = quant_debugger.model_statistics
+
+    self.assertCountEqual(expected_metrics.keys(), actual_metrics.keys())
+    for key, value in expected_metrics.items():
+      self.assertAlmostEqual(value, actual_metrics[key], places=5)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_layer_direct_compare_metrics(self, quantized_io):
+
+    def _corr(float_values, quant_values, scale, zero_point):
+      dequant_values = (quant_values.astype(np.int32) - zero_point) * scale
+      return np.corrcoef(float_values.flatten(), dequant_values.flatten())[0, 1]
+
+    if quantized_io:
+      debug_model = QuantizationDebuggerTest.debug_model_int8
+    else:
+      debug_model = QuantizationDebuggerTest.debug_model_float
+
+    options = debugger.QuantizationDebugOptions(
+        layer_direct_compare_metrics={'corr': _corr})
+    quant_debugger = debugger.QuantizationDebugger(
+        quant_debug_model_content=debug_model,
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+    quant_debugger.run()
+
+    expected_metrics = {
+        'corr': 0.99999,
+    }
+    self.assertLen(quant_debugger.layer_statistics, 1)
+    actual_metrics = next(iter(quant_debugger.layer_statistics.values()))
+
+    for key, value in expected_metrics.items():
+      self.assertAlmostEqual(value, actual_metrics[key], places=4)
+
+  @test_util.run_v2_only
+  def test_wrong_input_raises_ValueError(self):
+
+    def wrong_calibration_gen():
+      for _ in range(5):
+        yield [
+            np.ones((1, 3, 3, 1), dtype=np.float32),
+            np.ones((1, 3, 3, 1), dtype=np.float32)
+        ]
+
+    quant_debugger = debugger.QuantizationDebugger(
+        quant_debug_model_content=QuantizationDebuggerTest.debug_model_float,
+        debug_dataset=wrong_calibration_gen)
+    with self.assertRaisesRegex(
+        ValueError, r'inputs provided \(2\).+inputs to the model \(1\)'):
+      quant_debugger.run()
+
+  @test_util.run_v2_only
+  def test_non_debug_model_raises_ValueError(self):
+    normal_quant_model = _quantize_model(
+        QuantizationDebuggerTest.tf_model_root,
+        QuantizationDebuggerTest.tf_model,
+        _calibration_gen,
+        debug=False)
+
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      debugger.QuantizationDebugger(
+          quant_debug_model_content=normal_quant_model,
+          debug_dataset=_calibration_gen)
+
+  @parameterized.named_parameters(
+      ('empty quantization parameter', {
+          'quantization_parameters': {}
+      }, None),
+      ('empty scales/zero points', {
+          'quantization_parameters': {
+              'scales': [],
+              'zero_points': []
+          }
+      }, None),
+      ('invalid scales/zero points', {
+          'quantization_parameters': {
+              'scales': [1.0],
+              'zero_points': []
+          }
+      }, None),
+      ('correct case', {
+          'quantization_parameters': {
+              'scales': [0.5, 1.0],
+              'zero_points': [42, 7]
+          }
+      }, (0.5, 42)),
+  )
+  def test_get_quant_params(self, tensor_detail, expected_value):
+    self.assertEqual(debugger._get_quant_params(tensor_detail), expected_value)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_ops_from_option_setter(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io)
+    quant_debugger = debugger.QuantizationDebugger(
+        converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                      _calibration_gen),
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+
+    options.denylisted_ops = ['CONV_2D']
+    # TODO(b/195084873): The exception is expected to check whether selective
+    # quantization was done properly, since after the selective quantization
+    # the model will have no quantized layers thus have no NumericVerify ops,
+    # resulted in this exception. Marked with a bug to fix this in more
+    # straightforward way.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      quant_debugger.options = options
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_ops_from_option_constructor(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io,
+        denylisted_ops=['CONV_2D'])
+    # TODO(b/195084873): Count the number of NumericVerify op.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      _ = debugger.QuantizationDebugger(
+          converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                        _calibration_gen),
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_nodes_from_option_setter(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io)
+    quant_debugger = debugger.QuantizationDebugger(
+        converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                      _calibration_gen),
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+
+    options.denylisted_nodes = ['Identity']
+    # TODO(b/195084873): Count the number of NumericVerify op.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      quant_debugger.options = options
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_nodes_from_option_constructor(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io,
+        denylisted_nodes=['Identity'])
+    # TODO(b/195084873): Count the number of NumericVerify op.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      _ = debugger.QuantizationDebugger(
+          converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                        _calibration_gen),
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+
+  @mock.patch.object(metrics.TFLiteMetrics,
+                     'increase_counter_debugger_creation')
+  def test_creation_counter(self, increase_call):
+    debug_model = QuantizationDebuggerTest.debug_model_float
+    debugger.QuantizationDebugger(
+        quant_debug_model_content=debug_model, debug_dataset=_calibration_gen)
+    increase_call.assert_called_once()
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/lite/tools/optimize/g3doc/quantize_weights.md

@@ -0,0 +1,70 @@
+# TFLite Quantize Weights Tool
+
+## Recommended usage
+
+The Quantize Weights transformation is integrated with
+[tflite_convert](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/convert/cmdline_reference.md#transformation-flags).
+
+The recommended way of invoking this tool is by simply adding the
+`--post_training_quantize` flag to your original tflite_convert invocation. For
+example,
+
+```
+tflite_convert \
+  --output_file=/tmp/foo.tflite \
+  --saved_model_dir=/tmp/saved_model \
+  --post_training_quantize
+```
+
+## Overview
+
+The Quantize Weights tool provides a simple way to quantize the weights for a
+float TFLite model.
+
+TODO(raghuramank): Add link to weight quantization tutorial.
+
+### Size reduction
+
+float32 weights will be converted to 8 bit integers. This results in a model
+that is around 1/4th the size of the original model.
+
+### Latency reduction
+
+TFLite also has "hybrid" kernels implemented for many operations. These "hybrid"
+kernels take 8 bit integer weights and float inputs, dynamically quantize the
+inputs tensor (based on the input tensor's min and max elements), and does
+computations using the 8 bit integer values. This results in a 2-4x reduction in
+latency for "hybrid" kernels. In this mode the inference type is still FLOAT
+since the inputs and output to each operation is still float.
+
+For operations that do not yet have "hybrid" kernels implemented, we introduce a
+Dequantize operation after 8 bit integer weights. These convert weights back to
+float32 during inference to allow original float32 kernels to run. Since we
+cache dequantized results, the result of each of this dequantized path will be
+on-par with the original float model.
+
+TODO(yunluli): Fill in latency results from latency experiments.
+
+### Accuracy
+
+Since this technique quantizes weights after the model has already been trained,
+there can be accuracy drops depending on the model. For common CNN networks, the
+observed accuracy drops are small and can be seen below.
+
+TODO(yunluli): Fill in accuracy results from accuracy experiments.
+
+## Direct usage
+
+One can also invoke the Quantize Weights directly via C++ if they have a float
+`::tflite::Model` that they want to convert. They must provide a
+`flatbuffers::FlatBufferBuilder` which owns the underlying buffer of the created
+model. Here is an example invocation:
+
+```
+::tflite::Model* input_model = ...;
+flatbuffers::FlatBufferBuilder builder;
+TfLiteStatus status = ::tflite::optimize::QuantizeWeights(&builder, input_model);
+CHECK(status, kTfLiteStatusOk);
+const uint8_t* buffer = builder->GetBufferPointer();
+tflite::Model* output_model = ::tflite::GetModel(buffer);
+```

tensorflow/lite/tools/optimize/model_utils.cc

@@ -0,0 +1,194 @@
+/* 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/tools/optimize/model_utils.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <fstream>
+#include <ios>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/compiler/mlir/lite/tools/optimize/operator_property.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/schema/schema_conversion_utils.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
+
+namespace tflite {
+namespace optimize {
+namespace utils {
+
+namespace {
+
+// LINT.IfChange(GetOrInsertOpCodeIndex)
+// Returns the index of the OpCode.
+// If a OpCode doesn't exist, adds it and returns its index.
+int32_t GetOrInsertOpCodeIndex(ModelT* model, const BuiltinOperator& op_code,
+                               int32_t version) {
+  for (size_t i = 0; i < model->operator_codes.size(); ++i) {
+    if (GetBuiltinCode(model->operator_codes[i].get()) == op_code) {
+      return i;
+    }
+  }
+  model->operator_codes.push_back(std::make_unique<OperatorCodeT>());
+  int op_code_idx = model->operator_codes.size() - 1;
+  model->operator_codes[op_code_idx]->builtin_code = op_code;
+  model->operator_codes[op_code_idx]->deprecated_builtin_code =
+      ConvertBuiltinCodeToDeprecatedBuiltinCode(op_code);
+  // Version 2 and onwards supports INT8 inputs.
+  model->operator_codes[op_code_idx]->version = version;
+
+  // Return the index of the newly placed OperatorCodeT.
+  return op_code_idx;
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc:GetOrInsertOpCodeIndex)
+
+}  // namespace
+
+// LINT.IfChange(MakeDequantizeOperator)
+// Creates a Dequantize OperatorT object.
+void MakeDequantizeOperator(ModelT* model, std::unique_ptr<OperatorT>* op,
+                            int32_t input, int32_t output) {
+  OperatorT* op_raw = new OperatorT;
+  // Version 2 and onwards supports INT8 inputs.
+  op_raw->opcode_index =
+      GetOrInsertOpCodeIndex(model, BuiltinOperator_DEQUANTIZE, 2);
+  op_raw->inputs = {input};
+  op_raw->outputs = {output};
+
+  op->reset(op_raw);
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc:MakeDequantizeOperator)
+
+// Creates a Quantize OperatorT object.
+void MakeQuantizeOperator(ModelT* model, std::unique_ptr<OperatorT>* op,
+                          int32_t input, int32_t output) {
+  OperatorT* op_raw = new OperatorT;
+  op_raw->opcode_index =
+      GetOrInsertOpCodeIndex(model, BuiltinOperator_QUANTIZE, 1);
+  op_raw->inputs = {input};
+  op_raw->outputs = {output};
+
+  op->reset(op_raw);
+}
+
+// LINT.IfChange(MakeTensor)
+// Create a new TensorT object without quantization parameters.
+void MakeTensor(const string& name, const std::vector<int32_t>& shape,
+                const std::vector<int32_t>& shape_signature,
+                const TensorType& type, std::unique_ptr<TensorT>* tensor) {
+  TensorT* tensor_raw = new TensorT;
+  tensor_raw->name = name;
+  tensor_raw->shape = shape;
+  if (!shape_signature.empty()) {
+    tensor_raw->shape_signature = shape_signature;
+  }
+  tensor_raw->type = type;
+
+  tensor->reset(tensor_raw);
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc:MakeTensor)
+
+// Create a new TensorT object with quantization parameters.
+void MakeTensorWithQuantParam(const string& name,
+                              const std::vector<int32_t>& shape,
+                              const std::vector<int32_t>& shape_signature,
+                              const TensorType& type, float scale,
+                              int64_t zero_point,
+                              std::unique_ptr<TensorT>* tensor) {
+  MakeTensor(name, shape, shape_signature, type, tensor);
+  (*tensor)->quantization = std::make_unique<QuantizationParametersT>();
+  (*tensor)->quantization->scale.push_back(scale);
+  (*tensor)->quantization->zero_point.push_back(zero_point);
+}
+
+bool QuantizationParametersExist(const TensorT* tensor) {
+  return tensor->quantization != nullptr &&
+         !tensor->quantization->scale.empty() &&
+         !tensor->quantization->zero_point.empty();
+}
+
+bool HasBuffer(const ModelT* model, const SubGraphT* subgraph,
+               int tensor_index) {
+  const int buffer_index = subgraph->tensors[tensor_index]->buffer;
+  BufferT* buffer = model->buffers[buffer_index].get();
+  if (buffer == nullptr || buffer->data.empty()) {
+    return false;
+  }
+  return true;
+}
+
+// LINT.IfChange(HasMinMax)
+bool HasMinMax(const TensorT* tensor) {
+  return tensor->quantization && !tensor->quantization->min.empty() &&
+         !tensor->quantization->max.empty();
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc:HasMinMax)
+
+void SetOperatorCodeVersion(ModelT* model) {
+  for (int subgraph_idx = 0, end = model->subgraphs.size(); subgraph_idx < end;
+       subgraph_idx++) {
+    SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get();
+    // Iterate backward to avoid messing with index.
+    for (int op_idx = subgraph->operators.size() - 1; op_idx >= 0; op_idx--) {
+      OperatorT* op = subgraph->operators[op_idx].get();
+      OperatorCodeT* op_code = model->operator_codes[op->opcode_index].get();
+      operator_property::OperatorProperty property =
+          operator_property::GetOperatorProperty(model, subgraph_idx, op_idx);
+      if (property.quantizable && op_code->version < property.version) {
+        // Only update the versions of quantizable operations if the original
+        // version is lesser than minimum quantized one mentioned by
+        // OperatorProperty.
+        op_code->version = property.version;
+      }
+    }
+  }
+}
+
+void WriteFile(const std::string& out_file, const uint8_t* bytes,
+               size_t num_bytes) {
+  std::fstream stream(out_file, std::ios::binary | std::ios::out);
+  for (size_t i = 0; i < num_bytes; i++) {
+    stream << bytes[i];
+  }
+  TFLITE_DCHECK(!stream.bad() && !stream.fail());
+}
+
+std::unique_ptr<flatbuffers::FlatBufferBuilder> FinishModel(
+    const tflite::ModelT* model) {
+  std::unique_ptr<flatbuffers::FlatBufferBuilder> builder(
+      new flatbuffers::FlatBufferBuilder());
+  auto packed_model = tflite::Model::Pack(*builder, model);
+  tflite::FinishModelBuffer(*builder, packed_model);
+  return builder;
+}
+
+std::unique_ptr<tflite::ModelT> CreateMutableModelFromFile(
+    const string& model_filepath) {
+  auto fb_model =
+      tflite::FlatBufferModel::BuildFromFile(model_filepath.c_str());
+  auto tflite_model = fb_model->GetModel();
+  auto copied_model = std::make_unique<tflite::ModelT>();
+  tflite_model->UnPackTo(copied_model.get(), nullptr);
+  return copied_model;
+}
+
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/model_utils.h

@@ -0,0 +1,88 @@
+/* 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_TOOLS_OPTIMIZE_MODEL_UTILS_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_MODEL_UTILS_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/memory/memory.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace optimize {
+namespace utils {
+
+// LINT.IfChange(MakeDequantizeOperator)
+// Creates a Dequantize OperatorT object.
+void MakeDequantizeOperator(ModelT* model, std::unique_ptr<OperatorT>* op,
+                            int32_t input, int32_t output);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc:MakeDequantizeOperator)
+
+// Creates a Quantize OperatorT object.
+void MakeQuantizeOperator(ModelT* model, std::unique_ptr<OperatorT>* op,
+                          int32_t input, int32_t output);
+
+// LINT.IfChange(MakeTensor)
+// Create a new TensorT object without quantization parameters.
+void MakeTensor(const string& name, const std::vector<int32_t>& shape,
+                const std::vector<int32_t>& shape_signature,
+                const TensorType& type, std::unique_ptr<TensorT>* tensor);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc:MakeTensor)
+
+// Create a new TensorT object with quantization parameters.
+void MakeTensorWithQuantParam(const string& name,
+                              const std::vector<int32_t>& shape,
+                              const std::vector<int32_t>& shape_signature,
+                              const TensorType& type, float scale,
+                              int64_t zero_point,
+                              std::unique_ptr<TensorT>* tensor);
+
+// Checks if the tensor has scale and zero point populated.
+bool QuantizationParametersExist(const TensorT* tensor);
+
+bool HasBuffer(const ModelT* model, const SubGraphT* subgraph,
+               int tensor_index);
+
+// LINT.IfChange(HasMinMax)
+bool HasMinMax(const TensorT* tensor);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils.cc:HasMinMax)
+
+// Set version of OperatorCode. The version will only be applied for operations
+// that have been quantized.
+void SetOperatorCodeVersion(ModelT* model);
+
+// Writes model buffer to file.
+void WriteFile(const std::string& out_file, const uint8_t* bytes,
+               size_t num_bytes);
+
+// Finishes model buffer and returns its builder.
+std::unique_ptr<flatbuffers::FlatBufferBuilder> FinishModel(
+    const tflite::ModelT* model);
+
+// Reads TensorFlow Lite model from the given path.
+std::unique_ptr<tflite::ModelT> CreateMutableModelFromFile(
+    const string& model_filepath);
+
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_MODEL_UTILS_H_

tensorflow/lite/tools/optimize/model_utils_test.cc

@@ -0,0 +1,68 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/model_utils.h"
+
+#include <memory>
+#include <utility>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace optimize {
+namespace utils {
+namespace {
+
+TEST(ModelUtilsTest, QuantizationParametersExist) {
+  TensorT tensor;
+  tensor.quantization = std::make_unique<QuantizationParametersT>();
+  tensor.quantization->scale.push_back(0.5);
+  tensor.quantization->scale.push_back(1.5);
+  EXPECT_FALSE(QuantizationParametersExist(&tensor));
+  tensor.quantization->zero_point.push_back(1);
+  tensor.quantization->zero_point.push_back(-1);
+  EXPECT_TRUE(QuantizationParametersExist(&tensor));
+}
+
+TEST(ModelUtilsTest, HasBuffer) {
+  tflite::ModelT model;
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto tensor = std::make_unique<tflite::TensorT>();
+  tensor->buffer = 0;
+  subgraph->tensors.push_back(std::move(tensor));
+  model.subgraphs.push_back(std::move(subgraph));
+  auto buffer = std::make_unique<tflite::BufferT>();
+  model.buffers.push_back(std::move(buffer));
+  EXPECT_FALSE(HasBuffer(&model, model.subgraphs[0].get(), 0));
+  model.buffers[0]->data = {0, 1, 2, 3};
+  EXPECT_TRUE(HasBuffer(&model, model.subgraphs[0].get(), 0));
+}
+
+// LINT.IfChange(HasMinMaxTest)
+TEST(ModelUtilsTest, HasMinMax) {
+  TensorT tensor;
+  tensor.quantization = std::make_unique<QuantizationParametersT>();
+  tensor.quantization->min.push_back(0.5);
+  EXPECT_FALSE(HasMinMax(&tensor));
+  tensor.quantization->max.push_back(1.5);
+  EXPECT_TRUE(HasMinMax(&tensor));
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/model_utils_test.cc:HasMinMaxTest)
+
+}  // namespace
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/modify_model_interface.cc

@@ -0,0 +1,499 @@
+/* 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 "tensorflow/lite/tools/optimize/modify_model_interface.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "flatbuffers/flexbuffers.h"
+#include "absl/container/flat_hash_map.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/error_reporter.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
+#include "tensorflow/lite/tools/optimize/model_utils.h"
+
+namespace tflite {
+namespace optimize {
+
+namespace {
+
+// Structure to hold input tensor, op and output tensor.
+// op must be either quantize or dequantize.
+struct TensorOpTensor {
+  size_t subgraph_index;  // index of the subgraph.
+  int32_t input_index;    // index of the input tensor.
+  int32_t op_index;       // index of the op.
+  int32_t output_index;   // index of the output tensor.
+  int32_t model_index;    // index of the added tensor in the model.
+};
+
+// Finds float tensors that are model inputs and is consumed by a quantize Op.
+// The returned TensorOpTensor should have reverse order.
+std::vector<TensorOpTensor> GetInputTensors(const TensorType& input_type,
+                                            ModelT* model,
+                                            ErrorReporter* error_reporter) {
+  std::vector<TensorOpTensor> result;
+  // Get all input tensors.
+  for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size();
+       subgraph_idx++) {
+    SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get();
+    absl::flat_hash_map<TensorT*, int> input_tensors;
+    for (size_t input_idx = 0; input_idx < subgraph->inputs.size();
+         input_idx++) {
+      TensorT* tensor = subgraph->tensors[subgraph->inputs[input_idx]].get();
+      if (tensor->type == TensorType_FLOAT32) {
+        input_tensors.insert({tensor, input_idx});
+      }
+    }
+
+    for (int32_t op_idx = subgraph->operators.size() - 1; op_idx >= 0;
+         op_idx--) {
+      OperatorT* op = subgraph->operators[op_idx].get();
+      const BuiltinOperator op_code =
+          GetBuiltinCode(model->operator_codes[op->opcode_index].get());
+      TensorT* input_tensor = subgraph->tensors[op->inputs[0]].get();
+      if (input_tensors.find(input_tensor) == input_tensors.end()) {
+        continue;
+      }
+      if (op_code != BuiltinOperator_QUANTIZE) {
+        // Currently only supports int8 and int16 quantized models.
+        TF_LITE_REPORT_ERROR(
+            error_reporter,
+            "modify_model_interface called on a model without quant/dequant.");
+        return {};
+      }
+      if (op->inputs.size() != 1) {
+        continue;
+      }
+      if (op->outputs.size() != 1) {
+        continue;
+      }
+      const int model_input_index = input_tensors[input_tensor];
+      TensorT* quant_output = subgraph->tensors[op->outputs[0]].get();
+      if (quant_output->type != TensorType_INT8 &&
+          quant_output->type != TensorType_INT16) {
+        TF_LITE_REPORT_ERROR(error_reporter,
+                             "modify_model_interface currently only supports "
+                             "int8 and int16 quantized models.");
+      }
+
+      // The input type must be the same as the model quantization type
+      if (input_type != quant_output->type) {
+        // An exception, allow for UINT8 input type for INT8 quantized model.
+        if (!(input_type == TensorType_UINT8 &&
+              quant_output->type == TensorType_INT8)) {
+          TF_LITE_REPORT_ERROR(
+              error_reporter,
+              "The %s input type is incompatible with %s quantized models. "
+              "To resolve this error, change the input_type to a compatible "
+              "one. "
+              "See: modify_model_interface.cc",
+              EnumNameTensorType(input_type),
+              EnumNameTensorType(quant_output->type));
+        }
+      }
+      if (quant_output->quantization == nullptr) {
+        continue;
+      }
+      result.push_back({subgraph_idx, op->inputs[0], op_idx, op->outputs[0],
+                        model_input_index});
+    }
+  }
+  return result;
+}
+
+// Finds float tensors that are model output and is consumed by a dequantize Op.
+// The returned TensorOpTensor should have reverse order.
+std::vector<TensorOpTensor> GetOutputTensors(const TensorType& output_type,
+                                             ModelT* model,
+                                             ErrorReporter* error_reporter) {
+  std::vector<TensorOpTensor> result;
+  // Get all output tensors.
+  for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size();
+       subgraph_idx++) {
+    SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get();
+    absl::flat_hash_map<TensorT*, int> output_tensors;
+    for (size_t output_idx = 0; output_idx < subgraph->outputs.size();
+         output_idx++) {
+      TensorT* tensor = subgraph->tensors[subgraph->outputs[output_idx]].get();
+      if (tensor->type == TensorType_FLOAT32) {
+        output_tensors.insert({tensor, output_idx});
+      }
+    }
+
+    for (int32_t op_idx = subgraph->operators.size() - 1; op_idx >= 0;
+         op_idx--) {
+      OperatorT* op = subgraph->operators[op_idx].get();
+      const BuiltinOperator op_code =
+          GetBuiltinCode(model->operator_codes[op->opcode_index].get());
+      TensorT* output_tensor = subgraph->tensors[op->outputs[0]].get();
+      if (output_tensors.find(output_tensor) == output_tensors.end()) {
+        continue;
+      }
+      if (op_code != BuiltinOperator_DEQUANTIZE) {
+        // Currently only supports int8 and int16 quantized models.
+        TF_LITE_REPORT_ERROR(
+            error_reporter,
+            "modify_model_interface called on a model without quant/dequant.");
+        return {};
+      }
+      if (op->inputs.size() != 1) {
+        continue;
+      }
+      if (op->outputs.size() != 1) {
+        continue;
+      }
+      const int model_output_index = output_tensors[output_tensor];
+      TensorT* dequant_input = subgraph->tensors[op->inputs[0]].get();
+      if (dequant_input->type != TensorType_INT8 &&
+          dequant_input->type != TensorType_INT16) {
+        // Currently only supports int8 and int16 quantized models.
+        TF_LITE_REPORT_ERROR(error_reporter,
+                             "modify_model_interface currently only supports "
+                             "int8 and int16 quantized models.");
+        return {};
+      }
+      if (output_type != dequant_input->type) {
+        // An exception, allow for UINT8 input type for INT8 quantized model.
+        if (!(output_type == TensorType_UINT8 &&
+              dequant_input->type == TensorType_INT8)) {
+          TF_LITE_REPORT_ERROR(
+              error_reporter,
+              "The %s output type is incompatible with %s quantized models. "
+              "To resolve this error, change the output_type to a compatible "
+              "one. "
+              "See: modify_model_interface.cc",
+              EnumNameTensorType(output_type),
+              EnumNameTensorType(dequant_input->type));
+        }
+      }
+      if (dequant_input->quantization == nullptr) {
+        continue;
+      }
+      result.push_back({subgraph_idx, op->inputs[0], op_idx, op->outputs[0],
+                        model_output_index});
+    }
+  }
+  return result;
+}
+
+TfLiteStatus SetInputTypeToUINT8(ModelT* model,
+                                 const std::vector<TensorOpTensor>& inputs) {
+  // If the input type is uint8, change float to uint8.
+  for (auto tot : inputs) {
+    SubGraphT* subgraph = model->subgraphs.at(tot.subgraph_index).get();
+    TensorT* quant_tensor = subgraph->tensors[tot.output_index].get();
+    const float quant_tensor_scale = quant_tensor->quantization->scale[0];
+    const int quant_tensor_zp = quant_tensor->quantization->zero_point[0];
+    TensorT* float_tensor = subgraph->tensors[tot.input_index].get();
+    float_tensor->type = TensorType_UINT8;
+    if (float_tensor->quantization == nullptr) {
+      float_tensor->quantization = std::make_unique<QuantizationParametersT>();
+    }
+    float_tensor->quantization->scale.push_back(quant_tensor_scale);
+    float_tensor->quantization->zero_point.push_back(quant_tensor_zp + 128);
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus SetOutputTypeToUINT8(ModelT* model,
+                                  const std::vector<TensorOpTensor>& outputs) {
+  // Find Quant op code index.
+  size_t quant_op_index = 0;
+  for (size_t i = 0; i < model->operator_codes.size(); ++i) {
+    if (GetBuiltinCode(model->operator_codes[i].get()) ==
+        BuiltinOperator_QUANTIZE) {
+      quant_op_index = i;
+    }
+  }
+  // If the output type is uint8, change float to uint8.
+  for (auto tot : outputs) {
+    SubGraphT* subgraph = model->subgraphs.at(tot.subgraph_index).get();
+    TensorT* quant_tensor = subgraph->tensors[tot.input_index].get();
+    const float quant_tensor_scale = quant_tensor->quantization->scale[0];
+    const int quant_tensor_zp = quant_tensor->quantization->zero_point[0];
+    TensorT* float_tensor = subgraph->tensors[tot.output_index].get();
+    float_tensor->type = TensorType_UINT8;
+    if (float_tensor->quantization == nullptr) {
+      float_tensor->quantization = std::make_unique<QuantizationParametersT>();
+    }
+    float_tensor->quantization->scale.push_back(quant_tensor_scale);
+    float_tensor->quantization->zero_point.push_back(quant_tensor_zp + 128);
+
+    // Change op from dequant (int8 to float) to quant (int8 to uint8)
+    OperatorT* op = subgraph->operators[tot.op_index].get();
+    op->opcode_index = quant_op_index;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus RemoveInputTensor(ModelT* model,
+                               const std::vector<TensorOpTensor>& inputs,
+                               int32_t original_number_tensors) {
+  // Consistency check to make sure that erase start from the end.
+  int last_op_index = std::numeric_limits<int32_t>::max();
+  int last_tensor_index = std::numeric_limits<int32_t>::max();
+  for (auto tot : inputs) {
+    TFLITE_DCHECK(tot.input_index < last_tensor_index);
+    TFLITE_DCHECK(tot.op_index < last_op_index);
+    last_tensor_index = tot.input_index;
+    last_op_index = tot.op_index;
+  }
+  // Removes the input tensor and the related operator.
+  for (auto tot : inputs) {
+    SubGraphT* subgraph = model->subgraphs.at(tot.subgraph_index).get();
+    TFLITE_DCHECK(tot.input_index < subgraph->tensors.size());
+    TFLITE_DCHECK(tot.op_index < subgraph->operators.size());
+    if (tot.input_index >= original_number_tensors) {
+      subgraph->tensors.erase(subgraph->tensors.begin() + tot.input_index);
+    }
+    subgraph->operators.erase(subgraph->operators.begin() + tot.op_index);
+    subgraph->inputs[tot.model_index] = tot.output_index;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus RemoveOutputTensor(ModelT* model,
+                                const std::vector<TensorOpTensor>& outputs,
+                                int32_t original_number_tensors) {
+  // Consistency check to make sure that erase start from the end.
+  int last_op_index = std::numeric_limits<int32_t>::max();
+  int last_tensor_index = std::numeric_limits<int32_t>::max();
+  for (auto tot : outputs) {
+    TFLITE_DCHECK(tot.output_index < last_tensor_index);
+    TFLITE_DCHECK(tot.op_index < last_op_index);
+    last_tensor_index = tot.output_index;
+    last_op_index = tot.op_index;
+  }
+  // Removes the output tensor and the related operator.
+  for (auto tot : outputs) {
+    SubGraphT* subgraph = model->subgraphs.at(tot.subgraph_index).get();
+    TFLITE_DCHECK(tot.output_index < subgraph->tensors.size());
+    TFLITE_DCHECK(tot.op_index < subgraph->operators.size());
+    if (tot.output_index >= original_number_tensors) {
+      subgraph->tensors.erase(subgraph->tensors.begin() + tot.output_index);
+    }
+    subgraph->operators.erase(subgraph->operators.begin() + tot.op_index);
+    subgraph->outputs[tot.model_index] = tot.input_index;
+  }
+  return kTfLiteOk;
+}
+
+int GetOriginalNumberOfTensors(const TensorType& input_type,
+                               const TensorType& output_type, ModelT* model,
+                               ErrorReporter* error_reporter) {
+  std::vector<TensorOpTensor> outputs =
+      GetOutputTensors(output_type, model, error_reporter);
+  std::vector<TensorOpTensor> inputs =
+      GetInputTensors(input_type, model, error_reporter);
+  return model->subgraphs[0]->tensors.size() - outputs.size() - inputs.size();
+}
+
+}  // namespace
+
+TfLiteStatus ModifyModelInterface(flatbuffers::FlatBufferBuilder* builder,
+                                  ModelT* model, const TensorType& input_type,
+                                  const TensorType& output_type) {
+  tflite::StderrReporter error_reporter;
+  const int original_number_tensors = GetOriginalNumberOfTensors(
+      input_type, output_type, model, &error_reporter);
+  // Finds float tensors that are model output and are consumed by a float to
+  // int8/int16 quantize Op. Do output first since the tensors are added into
+  // input first.,
+  std::vector<TensorOpTensor> outputs =
+      GetOutputTensors(output_type, model, &error_reporter);
+  switch (output_type) {
+    case TensorType_UINT8:
+      SetOutputTypeToUINT8(model, outputs);
+      break;
+    case TensorType_INT8:
+    case TensorType_INT16:
+      RemoveOutputTensor(model, outputs, original_number_tensors);
+      break;
+    default:
+      return kTfLiteError;
+  }
+
+  // Find float tensors that are model input and is consumed by a float to
+  // int8/int16 quantize Op.
+  std::vector<TensorOpTensor> inputs =
+      GetInputTensors(input_type, model, &error_reporter);
+  switch (input_type) {
+    case TensorType_UINT8:
+      SetInputTypeToUINT8(model, inputs);
+      break;
+    case TensorType_INT8:
+    case TensorType_INT16:
+      RemoveInputTensor(model, inputs, original_number_tensors);
+      break;
+    default:
+      return kTfLiteError;
+  }
+
+  // Write to builder.
+  flatbuffers::Offset<Model> output_model_location =
+      Model::Pack(*builder, model);
+  FinishModelBuffer(*builder, output_model_location);
+
+  return kTfLiteOk;
+}
+
+TfLiteStatus ModifyModelInterface(const string& input_file,
+                                  const string& output_file,
+                                  const TensorType& input_type,
+                                  const TensorType& output_type) {
+  // Consistency Check
+  if (input_type != tflite::TensorType_INT8 &&
+      input_type != tflite::TensorType_UINT8 &&
+      input_type != tflite::TensorType_INT16) {
+    return kTfLiteError;
+  }
+  if (output_type != tflite::TensorType_INT8 &&
+      output_type != tflite::TensorType_UINT8 &&
+      output_type != tflite::TensorType_INT16) {
+    return kTfLiteError;
+  }
+
+  // Create model.
+  auto tflite_model = utils::CreateMutableModelFromFile(input_file);
+
+  auto model_builder = utils::FinishModel(tflite_model.get());
+
+  auto fixed_point_model_builder =
+      std::make_unique<flatbuffers::FlatBufferBuilder>();
+  flatbuffers::FlatBufferBuilder builder;
+
+  auto status = ModifyModelInterface(&builder, tflite_model.get(), input_type,
+                                     output_type);
+  TFLITE_DCHECK_EQ(status, kTfLiteOk);
+
+  utils::WriteFile(output_file, builder.GetBufferPointer(), builder.GetSize());
+
+  return kTfLiteOk;
+}
+
+namespace {
+void AddUint8Dequant(
+    const std::unordered_map<string, std::pair<float, int32_t>>& quant_params,
+    ModelT* model) {
+  for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size();
+       subgraph_idx++) {
+    SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get();
+    // Add dequant to input tensors.
+    for (size_t input_idx = 0; input_idx < subgraph->inputs.size();
+         input_idx++) {
+      const int32_t tensor_idx = subgraph->inputs[input_idx];
+      TensorT* tensor = subgraph->tensors[tensor_idx].get();
+      if (tensor->type != TensorType_FLOAT32) {
+        continue;
+      }
+      if (quant_params.find(tensor->name) != quant_params.end()) {
+        // Add uint8 tensor
+        const string added_tensor_name = tensor->name + "_uint8";
+        std::unique_ptr<TensorT> leading_op_input;
+        const std::pair<float, int32_t>& provided_quant_params =
+            quant_params.at(string(tensor->name));
+        utils::MakeTensorWithQuantParam(
+            added_tensor_name, tensor->shape, tensor->shape_signature,
+            TensorType_UINT8, provided_quant_params.first,
+            provided_quant_params.second, &leading_op_input);
+        const int32_t leading_op_input_idx = subgraph->tensors.size();
+        subgraph->tensors.push_back(std::move(leading_op_input));
+
+        // Create the leading op, which is deqantize Op.
+        std::unique_ptr<OperatorT> leading_op;
+        utils::MakeDequantizeOperator(model, &leading_op, leading_op_input_idx,
+                                      tensor_idx);
+
+        // Insert the new op at the start of the model.
+        subgraph->operators.insert(subgraph->operators.begin(),
+                                   std::move(leading_op));
+      }
+    }
+  }
+}
+
+void AddUint8Quant(
+    const std::unordered_map<string, std::pair<float, int32_t>>& quant_params,
+    ModelT* model) {
+  for (size_t subgraph_idx = 0; subgraph_idx < model->subgraphs.size();
+       subgraph_idx++) {
+    SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get();
+    // Add quant to output tensors.
+    for (size_t output_idx = 0; output_idx < subgraph->outputs.size();
+         output_idx++) {
+      const int32_t tensor_idx = subgraph->outputs[output_idx];
+      TensorT* tensor = subgraph->tensors[tensor_idx].get();
+      if (tensor->type != TensorType_FLOAT32) {
+        continue;
+      }
+      if (quant_params.find(tensor->name) != quant_params.end()) {
+        // Add uint8 tensor
+        const string added_tensor_name = tensor->name + "_uint8";
+        std::unique_ptr<TensorT> tailing_op_output;
+        const std::pair<float, int32_t>& provided_quant_params =
+            quant_params.at(string(tensor->name));
+        utils::MakeTensorWithQuantParam(
+            added_tensor_name, tensor->shape, tensor->shape_signature,
+            TensorType_UINT8, provided_quant_params.first,
+            provided_quant_params.second, &tailing_op_output);
+        const int32_t tailing_op_output_idx = subgraph->tensors.size();
+        subgraph->tensors.push_back(std::move(tailing_op_output));
+
+        // Create the tailing op, which is Qantize Op.
+        std::unique_ptr<OperatorT> tailing_op;
+        utils::MakeQuantizeOperator(model, &tailing_op, tensor_idx,
+                                    tailing_op_output_idx);
+
+        // Insert the new op at the end of the model.
+        subgraph->operators.push_back(std::move(tailing_op));
+      }
+    }
+  }
+}
+}  // namespace
+
+TfLiteStatus Uint8QuantizeModelInputsOutputs(
+    flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
+    const std::unordered_map<string, std::pair<float, int32_t>>&
+        input_quant_params,
+    const std::unordered_map<string, std::pair<float, int32_t>>&
+        output_quant_params) {
+  std::unique_ptr<ModelT> model;
+  model.reset(input_model->UnPack());
+  // Add Dequant for inputs.
+  AddUint8Dequant(input_quant_params, model.get());
+
+  // Add Quant for outputs.
+  AddUint8Quant(output_quant_params, model.get());
+
+  // Output model.
+  flatbuffers::Offset<Model> output_model_location =
+      Model::Pack(*builder, model.get());
+  FinishModelBuffer(*builder, output_model_location);
+
+  return kTfLiteOk;
+}
+
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/modify_model_interface.h

@@ -0,0 +1,68 @@
+/* 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_TOOLS_OPTIMIZE_MODIFY_MODEL_INTERFACE_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_MODIFY_MODEL_INTERFACE_H_
+
+#include <cstdint>
+#include <unordered_map>
+#include <utility>
+
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace optimize {
+
+// Changes the interface of a quantized model. This method allows the users to
+// replace float interface with other types. Currently only int8, int16 and
+// uint8 are supported.
+//
+// This method populates the builder with the new model.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus ModifyModelInterface(flatbuffers::FlatBufferBuilder* builder,
+                                  ModelT* model, const TensorType& input_type,
+                                  const TensorType& output_type);
+
+// Same as above but allows input file path and output file path.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus ModifyModelInterface(const string& input_file,
+                                  const string& output_file,
+                                  const TensorType& input_type,
+                                  const TensorType& output_type);
+
+// Adds uint8 quantize ops for specified inputs and uint8 dequantize ops for
+// specified outputs for a float model. The scale and zero point of uint8
+// tensors are provided through quant_params.
+//   - input_quant_params has a map between tensor name and the
+//     <scale and zero_point> pair for inputs.
+//   - output_quant_params has a map between tensor name and the
+//     <scale and zero_point> pair for inputs.
+// For the inputs/output tensors for the model, if its quantization parameters
+// are not provided, that tensor is not affected.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus Uint8QuantizeModelInputsOutputs(
+    flatbuffers::FlatBufferBuilder* builder, const Model* input_model,
+    const std::unordered_map<string, std::pair<float, int32_t>>&
+        input_quant_params,
+    const std::unordered_map<string, std::pair<float, int32_t>>&
+        output_quant_params);
+
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_MODIFY_MODEL_INTERFACE_H_

tensorflow/lite/tools/optimize/modify_model_interface_main.cc

@@ -0,0 +1,51 @@
+/* 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 <wrapper_internal_exception_macros.h>
+
+#include <cstdio>
+#include <string>
+#include <unordered_map>
+
+#include "tensorflow/lite/tools/optimize/modify_model_interface.h"
+//
+// Note: This is a private API, subject to change.
+int main(int argc, char** argv) {
+  if (argc != 5) {
+    printf(
+        "Wrong number of arguments. Example: modify_model_interface_main "
+        "${input} ${output} ${input_interface} ${output_interface}");
+    return 1;
+  }
+
+  const std::unordered_map<std::string, tflite::TensorType> supported_types{
+      {"uint8", tflite::TensorType_UINT8},
+      {"int8", tflite::TensorType_INT8},
+      {"int16", tflite::TensorType_INT16}};
+
+  tflite::TensorType input = tflite::TensorType_INT8;
+  tflite::TensorType output = tflite::TensorType_INT8;
+
+  try {
+    input = supported_types.at(argv[3]);
+    output = supported_types.at(argv[4]);
+  } catch (const std::out_of_range&) {
+    printf("Only supports uint8, int8 and int16 for input and output types");
+    return 1;
+  }
+
+  tflite::optimize::ModifyModelInterface(argv[1], argv[2], input, output);
+
+  return 0;
+}

tensorflow/lite/tools/optimize/modify_model_interface_test.cc

@@ -0,0 +1,634 @@
+/* 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 "tensorflow/lite/tools/optimize/modify_model_interface.h"
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
+
+namespace tflite {
+namespace optimize {
+namespace {
+
+// Create a model with 1 quant, 1 FC, 1 dequant
+std::unique_ptr<ModelT> CreateQuantizedModelSingleInputOutput(
+    const TensorType& quantization_type) {
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  auto quant_op_code = std::make_unique<OperatorCodeT>();
+  auto quant_op = std::make_unique<OperatorT>();
+  auto fc_op_code = std::make_unique<OperatorCodeT>();
+  auto fc_op = std::make_unique<OperatorT>();
+  auto dequant_op_code = std::make_unique<OperatorCodeT>();
+  auto dequant_op = std::make_unique<OperatorT>();
+
+  model->subgraphs.push_back(std::move(subgraph));
+
+  // Op code
+  quant_op_code->builtin_code = BuiltinOperator_QUANTIZE;
+  quant_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_QUANTIZE);
+  quant_op_code->version = 2;
+
+  fc_op_code->builtin_code = BuiltinOperator_FULLY_CONNECTED;
+  fc_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_FULLY_CONNECTED);
+  fc_op_code->version = 2;
+
+  dequant_op_code->builtin_code = BuiltinOperator_DEQUANTIZE;
+  dequant_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_DEQUANTIZE);
+  dequant_op_code->version = 2;
+
+  // Op.
+  quant_op->opcode_index = 0;
+  quant_op->inputs = {0};
+  quant_op->outputs = {1};
+
+  fc_op->opcode_index = 1;
+  fc_op->inputs = {1};
+  fc_op->outputs = {2};
+
+  dequant_op->opcode_index = 2;
+  dequant_op->inputs = {2};
+  dequant_op->outputs = {3};
+
+  model->subgraphs[0]->operators.push_back(std::move(quant_op));
+  model->subgraphs[0]->operators.push_back(std::move(fc_op));
+  model->subgraphs[0]->operators.push_back(std::move(dequant_op));
+
+  model->operator_codes.push_back(std::move(quant_op_code));
+  model->operator_codes.push_back(std::move(fc_op_code));
+  model->operator_codes.push_back(std::move(dequant_op_code));
+
+  // Model input/output.
+  model->subgraphs[0]->inputs = {0};
+  model->subgraphs[0]->outputs = {3};
+
+  // Tensors
+  auto tensor_0 = std::make_unique<TensorT>();
+  tensor_0->name = "tensor_0";
+  tensor_0->shape = {};
+  tensor_0->type = TensorType_FLOAT32;
+
+  auto tensor_1 = std::make_unique<TensorT>();
+  tensor_1->quantization = std::make_unique<QuantizationParametersT>();
+  tensor_1->quantization->scale.push_back(0.35);
+  tensor_1->quantization->zero_point.push_back(28);
+  tensor_1->name = "tensor_1";
+  tensor_1->shape = {};
+  tensor_1->type = quantization_type;
+
+  auto tensor_2 = std::make_unique<TensorT>();
+  tensor_2->quantization = std::make_unique<QuantizationParametersT>();
+  tensor_2->quantization->scale.push_back(0.12);
+  tensor_2->quantization->zero_point.push_back(50);
+  tensor_2->name = "tensor_2";
+  tensor_2->shape = {};
+  tensor_2->type = quantization_type;
+
+  auto tensor_3 = std::make_unique<TensorT>();
+  tensor_3->name = "tensor_3";
+  tensor_3->shape = {};
+  tensor_3->type = TensorType_FLOAT32;
+
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_0));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_1));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_2));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_3));
+
+  // Buffer
+  model->buffers.push_back(std::move(buffer));
+
+  return model;
+}
+
+// Create a model with 2 quant, 1 FC, 2 dequant
+// The model mimics the behavior of the quantize_model.cc.
+std::unique_ptr<ModelT> CreateQuantizedModelMultipleInputOutput(
+    const TensorType& quantization_type) {
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  auto quant_op_code = std::make_unique<OperatorCodeT>();
+  auto quant_op_1 = std::make_unique<OperatorT>();
+  auto quant_op_2 = std::make_unique<OperatorT>();
+  auto fc_op_code = std::make_unique<OperatorCodeT>();
+  auto fc_op = std::make_unique<OperatorT>();
+  auto dequant_op_code = std::make_unique<OperatorCodeT>();
+  auto dequant_op_1 = std::make_unique<OperatorT>();
+  auto dequant_op_2 = std::make_unique<OperatorT>();
+
+  model->subgraphs.push_back(std::move(subgraph));
+
+  // Op code
+  quant_op_code->builtin_code = BuiltinOperator_QUANTIZE;
+  quant_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_QUANTIZE);
+  quant_op_code->version = 2;
+
+  fc_op_code->builtin_code = BuiltinOperator_FULLY_CONNECTED;
+  fc_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_FULLY_CONNECTED);
+  fc_op_code->version = 2;
+
+  dequant_op_code->builtin_code = BuiltinOperator_DEQUANTIZE;
+  dequant_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_DEQUANTIZE);
+  dequant_op_code->version = 2;
+
+  // Op.
+  quant_op_1->opcode_index = 0;
+  quant_op_1->inputs = {0};
+  quant_op_1->outputs = {2};
+  quant_op_2->opcode_index = 0;
+  quant_op_2->inputs = {1};
+  quant_op_2->outputs = {3};
+
+  fc_op->opcode_index = 1;
+  fc_op->inputs = {2, 3};
+  fc_op->outputs = {4, 5};
+
+  dequant_op_1->opcode_index = 2;
+  dequant_op_1->inputs = {4};
+  dequant_op_1->outputs = {6};
+  dequant_op_2->opcode_index = 2;
+  dequant_op_2->inputs = {5};
+  dequant_op_2->outputs = {7};
+
+  model->subgraphs[0]->operators.push_back(std::move(quant_op_1));
+  model->subgraphs[0]->operators.push_back(std::move(quant_op_2));
+  model->subgraphs[0]->operators.push_back(std::move(fc_op));
+  model->subgraphs[0]->operators.push_back(std::move(dequant_op_1));
+  model->subgraphs[0]->operators.push_back(std::move(dequant_op_2));
+
+  model->operator_codes.push_back(std::move(quant_op_code));
+  model->operator_codes.push_back(std::move(fc_op_code));
+  model->operator_codes.push_back(std::move(dequant_op_code));
+
+  // Model input/output.
+  model->subgraphs[0]->inputs = {0, 1};
+  model->subgraphs[0]->outputs = {6, 7};
+
+  // Tensors
+  auto tensor_0 = std::make_unique<TensorT>();
+  tensor_0->name = "tensor_0";
+  tensor_0->shape = {};
+  tensor_0->type = TensorType_FLOAT32;
+
+  auto tensor_1 = std::make_unique<TensorT>();
+  tensor_1->name = "tensor_1";
+  tensor_1->shape = {};
+  tensor_1->type = TensorType_FLOAT32;
+
+  auto tensor_2 = std::make_unique<TensorT>();
+  tensor_2->quantization = std::make_unique<QuantizationParametersT>();
+  tensor_2->quantization->scale.push_back(0.35);
+  tensor_2->quantization->zero_point.push_back(28);
+  tensor_2->name = "tensor_2";
+  tensor_2->shape = {};
+  tensor_2->type = quantization_type;
+
+  auto tensor_3 = std::make_unique<TensorT>();
+  tensor_3->quantization = std::make_unique<QuantizationParametersT>();
+  tensor_3->quantization->scale.push_back(0.12);
+  tensor_3->quantization->zero_point.push_back(50);
+  tensor_3->name = "tensor_3";
+  tensor_3->shape = {};
+  tensor_3->type = quantization_type;
+
+  auto tensor_4 = std::make_unique<TensorT>();
+  tensor_4->quantization = std::make_unique<QuantizationParametersT>();
+  tensor_4->quantization->scale.push_back(0.45);
+  tensor_4->quantization->zero_point.push_back(28);
+  tensor_4->name = "tensor_4";
+  tensor_4->shape = {};
+  tensor_4->type = quantization_type;
+
+  auto tensor_5 = std::make_unique<TensorT>();
+  tensor_5->quantization = std::make_unique<QuantizationParametersT>();
+  tensor_5->quantization->scale.push_back(0.22);
+  tensor_5->quantization->zero_point.push_back(50);
+  tensor_5->name = "tensor_5";
+  tensor_5->shape = {};
+  tensor_5->type = quantization_type;
+
+  auto tensor_6 = std::make_unique<TensorT>();
+  tensor_6->name = "tensor_6";
+  tensor_6->shape = {};
+  tensor_6->type = TensorType_FLOAT32;
+
+  auto tensor_7 = std::make_unique<TensorT>();
+  tensor_7->name = "tensor_7";
+  tensor_7->shape = {};
+  tensor_7->type = TensorType_FLOAT32;
+
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_0));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_1));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_2));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_3));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_4));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_5));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_6));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_7));
+
+  // Buffer
+  model->buffers.push_back(std::move(buffer));
+
+  return model;
+}
+
+// Create a model with 1 FC.
+std::unique_ptr<ModelT> CreateFloatModel() {
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  auto fc_op_code = std::make_unique<OperatorCodeT>();
+  auto fc_op = std::make_unique<OperatorT>();
+
+  model->subgraphs.push_back(std::move(subgraph));
+
+  // Op code
+  fc_op_code->builtin_code = BuiltinOperator_FULLY_CONNECTED;
+  fc_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_FULLY_CONNECTED);
+  fc_op_code->version = 2;
+
+  // Op.
+  fc_op->opcode_index = 0;
+  fc_op->inputs = {0};
+  fc_op->outputs = {1};
+
+  model->subgraphs[0]->operators.push_back(std::move(fc_op));
+  model->operator_codes.push_back(std::move(fc_op_code));
+
+  // Model input/output.
+  model->subgraphs[0]->inputs = {0};
+  model->subgraphs[0]->outputs = {1};
+
+  // Tensors
+  auto tensor_0 = std::make_unique<TensorT>();
+  tensor_0->name = "tensor_0";
+  tensor_0->shape = {};
+  tensor_0->type = TensorType_FLOAT32;
+
+  auto tensor_1 = std::make_unique<TensorT>();
+  tensor_1->name = "tensor_1";
+  tensor_1->shape = {};
+  tensor_1->type = TensorType_FLOAT32;
+
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_0));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor_1));
+
+  // Buffer
+  model->buffers.push_back(std::move(buffer));
+
+  return model;
+}
+
+struct ModelInterface : ::testing::TestWithParam<tflite::TensorType> {};
+
+TEST_P(ModelInterface, SingleInputOutput) {
+  TensorType quantization_type = GetParam();
+
+  auto model = CreateQuantizedModelSingleInputOutput(quantization_type);
+
+  // Change model type.
+  flatbuffers::FlatBufferBuilder builder;
+  EXPECT_EQ(ModifyModelInterface(&builder, model.get(), quantization_type,
+                                 quantization_type),
+            kTfLiteOk);
+
+  // Verify results.
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->inputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->inputs[0], 1);
+  EXPECT_EQ(model->subgraphs[0]->outputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->outputs[0], 2);
+  EXPECT_EQ(model->operator_codes.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->operators[0]->opcode_index, 1);
+
+  auto fc_op = model->subgraphs[0]->operators[0].get();
+
+  auto input = model->subgraphs[0]->tensors[fc_op->inputs[0]].get();
+  EXPECT_EQ(input->name, "tensor_1");
+  EXPECT_EQ(input->type, quantization_type);
+  EXPECT_FLOAT_EQ(input->quantization->scale[0], 0.35);
+  EXPECT_EQ(input->quantization->zero_point[0], 28);
+
+  auto output = model->subgraphs[0]->tensors[fc_op->outputs[0]].get();
+  EXPECT_EQ(output->name, "tensor_2");
+  EXPECT_EQ(output->type, quantization_type);
+  EXPECT_FLOAT_EQ(output->quantization->scale[0], 0.12);
+  EXPECT_EQ(output->quantization->zero_point[0], 50);
+}
+
+TEST_P(ModelInterface, MutipleInputOutput) {
+  TensorType quantization_type = GetParam();
+
+  auto model = CreateQuantizedModelMultipleInputOutput(quantization_type);
+
+  // Change model type.
+  flatbuffers::FlatBufferBuilder builder;
+  EXPECT_EQ(ModifyModelInterface(&builder, model.get(), quantization_type,
+                                 quantization_type),
+            kTfLiteOk);
+
+  // Verify results.
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 6);
+  EXPECT_EQ(model->subgraphs[0]->inputs.size(), 2);
+  EXPECT_EQ(model->subgraphs[0]->inputs[0], 2);
+  EXPECT_EQ(model->subgraphs[0]->inputs[1], 3);
+  EXPECT_EQ(model->subgraphs[0]->outputs.size(), 2);
+  EXPECT_EQ(model->subgraphs[0]->outputs[0], 4);
+  EXPECT_EQ(model->subgraphs[0]->outputs[1], 5);
+  EXPECT_EQ(model->operator_codes.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->operators[0]->opcode_index, 1);
+
+  auto fc_op = model->subgraphs[0]->operators[0].get();
+
+  auto input_1 = model->subgraphs[0]->tensors[fc_op->inputs[0]].get();
+  EXPECT_EQ(input_1->name, "tensor_2");
+  EXPECT_EQ(input_1->type, quantization_type);
+  EXPECT_FLOAT_EQ(input_1->quantization->scale[0], 0.35);
+  EXPECT_EQ(input_1->quantization->zero_point[0], 28);
+
+  auto input_2 = model->subgraphs[0]->tensors[fc_op->inputs[1]].get();
+  EXPECT_EQ(input_2->name, "tensor_3");
+  EXPECT_EQ(input_2->type, quantization_type);
+  EXPECT_FLOAT_EQ(input_2->quantization->scale[0], 0.12);
+  EXPECT_EQ(input_2->quantization->zero_point[0], 50);
+
+  auto output_1 = model->subgraphs[0]->tensors[fc_op->outputs[0]].get();
+  EXPECT_EQ(output_1->name, "tensor_4");
+  EXPECT_EQ(output_1->type, quantization_type);
+  EXPECT_FLOAT_EQ(output_1->quantization->scale[0], 0.45);
+  EXPECT_EQ(output_1->quantization->zero_point[0], 28);
+
+  auto output_2 = model->subgraphs[0]->tensors[fc_op->outputs[1]].get();
+  EXPECT_EQ(output_2->name, "tensor_5");
+  EXPECT_EQ(output_2->type, quantization_type);
+  EXPECT_FLOAT_EQ(output_2->quantization->scale[0], 0.22);
+  EXPECT_EQ(output_2->quantization->zero_point[0], 50);
+}
+
+INSTANTIATE_TEST_SUITE_P(MultipleInputOutputTests, ModelInterface,
+                         ::testing::Values(TensorType_INT8, TensorType_INT16));
+
+TEST(ModelInterface, MixedTypeSingleInputOutput) {
+  auto model = CreateQuantizedModelSingleInputOutput(TensorType_INT8);
+
+  // Change model type.
+  flatbuffers::FlatBufferBuilder builder;
+  EXPECT_EQ(ModifyModelInterface(&builder, model.get(), TensorType_UINT8,
+                                 TensorType_INT8),
+            kTfLiteOk);
+
+  // Verify results.
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->inputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->inputs[0], 0);
+  EXPECT_EQ(model->subgraphs[0]->outputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->outputs[0], 2);
+  EXPECT_EQ(model->operator_codes.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 2);
+  EXPECT_EQ(model->subgraphs[0]->operators[0]->opcode_index, 0);
+  EXPECT_EQ(model->subgraphs[0]->operators[1]->opcode_index, 1);
+
+  auto quant_op = model->subgraphs[0]->operators[0].get();
+  auto input = model->subgraphs[0]->tensors[quant_op->inputs[0]].get();
+  EXPECT_EQ(input->name, "tensor_0");
+  EXPECT_EQ(input->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(input->quantization->scale[0], 0.35);
+  EXPECT_EQ(input->quantization->zero_point[0], 156);
+
+  auto fc_op = model->subgraphs[0]->operators[1].get();
+  auto output = model->subgraphs[0]->tensors[fc_op->outputs[0]].get();
+  EXPECT_EQ(output->name, "tensor_2");
+  EXPECT_EQ(output->type, TensorType_INT8);
+  EXPECT_FLOAT_EQ(output->quantization->scale[0], 0.12);
+  EXPECT_EQ(output->quantization->zero_point[0], 50);
+}
+
+TEST(ModelInterface, Uint8SingleInputOutput) {
+  auto model = CreateQuantizedModelSingleInputOutput(TensorType_INT8);
+
+  // Change model type.
+  flatbuffers::FlatBufferBuilder builder;
+  EXPECT_EQ(ModifyModelInterface(&builder, model.get(), TensorType_UINT8,
+                                 TensorType_UINT8),
+            kTfLiteOk);
+
+  // Verify results.
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 4);
+  EXPECT_EQ(model->subgraphs[0]->inputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->inputs[0], 0);
+  EXPECT_EQ(model->subgraphs[0]->outputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->outputs[0], 3);
+  EXPECT_EQ(model->operator_codes.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->operators[0]->opcode_index, 0);
+  EXPECT_EQ(model->subgraphs[0]->operators[1]->opcode_index, 1);
+  EXPECT_EQ(model->subgraphs[0]->operators[2]->opcode_index, 0);
+
+  auto input_quant_op = model->subgraphs[0]->operators[0].get();
+  auto input = model->subgraphs[0]->tensors[input_quant_op->inputs[0]].get();
+  EXPECT_EQ(input->name, "tensor_0");
+  EXPECT_EQ(input->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(input->quantization->scale[0], 0.35);
+  EXPECT_EQ(input->quantization->zero_point[0], 156);
+
+  auto output_quant_op = model->subgraphs[0]->operators[2].get();
+  auto output = model->subgraphs[0]->tensors[output_quant_op->outputs[0]].get();
+  EXPECT_EQ(output->name, "tensor_3");
+  EXPECT_EQ(output->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(output->quantization->scale[0], 0.12);
+  EXPECT_EQ(output->quantization->zero_point[0], 178);
+}
+
+TEST(ModelInterface, Uint8MutipleInputOutput) {
+  auto model = CreateQuantizedModelMultipleInputOutput(TensorType_INT8);
+
+  // Change model type.
+  flatbuffers::FlatBufferBuilder builder;
+  EXPECT_EQ(ModifyModelInterface(&builder, model.get(), TensorType_UINT8,
+                                 TensorType_UINT8),
+            kTfLiteOk);
+
+  // Verify results.
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 8);
+  EXPECT_EQ(model->subgraphs[0]->inputs.size(), 2);
+  EXPECT_EQ(model->subgraphs[0]->inputs[0], 0);
+  EXPECT_EQ(model->subgraphs[0]->inputs[1], 1);
+  EXPECT_EQ(model->subgraphs[0]->outputs.size(), 2);
+  EXPECT_EQ(model->subgraphs[0]->outputs[0], 6);
+  EXPECT_EQ(model->subgraphs[0]->outputs[1], 7);
+  EXPECT_EQ(model->operator_codes.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 5);
+  EXPECT_EQ(model->subgraphs[0]->operators[0]->opcode_index, 0);
+  EXPECT_EQ(model->subgraphs[0]->operators[1]->opcode_index, 0);
+  EXPECT_EQ(model->subgraphs[0]->operators[2]->opcode_index, 1);
+  EXPECT_EQ(model->subgraphs[0]->operators[3]->opcode_index, 0);
+  EXPECT_EQ(model->subgraphs[0]->operators[4]->opcode_index, 0);
+
+  auto input_quant_1 = model->subgraphs[0]->operators[0].get();
+  auto input_1 = model->subgraphs[0]->tensors[input_quant_1->inputs[0]].get();
+  EXPECT_EQ(input_1->name, "tensor_0");
+  EXPECT_EQ(input_1->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(input_1->quantization->scale[0], 0.35);
+  EXPECT_EQ(input_1->quantization->zero_point[0], 156);
+
+  auto input_quant_2 = model->subgraphs[0]->operators[1].get();
+  auto input_2 = model->subgraphs[0]->tensors[input_quant_2->inputs[0]].get();
+  EXPECT_EQ(input_2->name, "tensor_1");
+  EXPECT_EQ(input_2->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(input_2->quantization->scale[0], 0.12);
+  EXPECT_EQ(input_2->quantization->zero_point[0], 178);
+
+  auto output_quant_1 = model->subgraphs[0]->operators[3].get();
+  auto output_1 =
+      model->subgraphs[0]->tensors[output_quant_1->outputs[0]].get();
+  EXPECT_EQ(output_1->name, "tensor_6");
+  EXPECT_EQ(output_1->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(output_1->quantization->scale[0], 0.45);
+  EXPECT_EQ(output_1->quantization->zero_point[0], 156);
+
+  auto output_quant_2 = model->subgraphs[0]->operators[4].get();
+  auto output_2 =
+      model->subgraphs[0]->tensors[output_quant_2->outputs[0]].get();
+  EXPECT_EQ(output_2->name, "tensor_7");
+  EXPECT_EQ(output_2->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(output_2->quantization->scale[0], 0.22);
+  EXPECT_EQ(output_2->quantization->zero_point[0], 178);
+}
+
+TEST(ModelInterface, Int8MutipleInputOutput) {
+  auto model = CreateQuantizedModelMultipleInputOutput(TensorType_INT8);
+
+  // Change model type.
+  flatbuffers::FlatBufferBuilder builder;
+  EXPECT_EQ(ModifyModelInterface(&builder, model.get(), TensorType_INT8,
+                                 TensorType_INT8),
+            kTfLiteOk);
+
+  // Verify results.
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 6);
+  EXPECT_EQ(model->subgraphs[0]->inputs.size(), 2);
+  EXPECT_EQ(model->subgraphs[0]->inputs[0], 2);
+  EXPECT_EQ(model->subgraphs[0]->inputs[1], 3);
+  EXPECT_EQ(model->subgraphs[0]->outputs.size(), 2);
+  EXPECT_EQ(model->subgraphs[0]->outputs[0], 4);
+  EXPECT_EQ(model->subgraphs[0]->outputs[1], 5);
+  EXPECT_EQ(model->operator_codes.size(), 3);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->operators[0]->opcode_index, 1);
+
+  auto fc_op = model->subgraphs[0]->operators[0].get();
+
+  auto input_1 = model->subgraphs[0]->tensors[fc_op->inputs[0]].get();
+  EXPECT_EQ(input_1->name, "tensor_2");
+  EXPECT_EQ(input_1->type, TensorType_INT8);
+  EXPECT_FLOAT_EQ(input_1->quantization->scale[0], 0.35);
+  EXPECT_EQ(input_1->quantization->zero_point[0], 28);
+
+  auto input_2 = model->subgraphs[0]->tensors[fc_op->inputs[1]].get();
+  EXPECT_EQ(input_2->name, "tensor_3");
+  EXPECT_EQ(input_2->type, TensorType_INT8);
+  EXPECT_FLOAT_EQ(input_2->quantization->scale[0], 0.12);
+  EXPECT_EQ(input_2->quantization->zero_point[0], 50);
+
+  auto output_1 = model->subgraphs[0]->tensors[fc_op->outputs[0]].get();
+  EXPECT_EQ(output_1->name, "tensor_4");
+  EXPECT_EQ(output_1->type, TensorType_INT8);
+  EXPECT_FLOAT_EQ(output_1->quantization->scale[0], 0.45);
+  EXPECT_EQ(output_1->quantization->zero_point[0], 28);
+
+  auto output_2 = model->subgraphs[0]->tensors[fc_op->outputs[1]].get();
+  EXPECT_EQ(output_2->name, "tensor_5");
+  EXPECT_EQ(output_2->type, TensorType_INT8);
+  EXPECT_FLOAT_EQ(output_2->quantization->scale[0], 0.22);
+  EXPECT_EQ(output_2->quantization->zero_point[0], 50);
+}
+
+TEST(ModelInterface, Float) {
+  // Create the model.
+  std::unique_ptr<ModelT> input_model_t = CreateFloatModel();
+  flatbuffers::FlatBufferBuilder builder_temp;
+  flatbuffers::Offset<Model> output_model_location =
+      Model::Pack(builder_temp, input_model_t.get());
+  FinishModelBuffer(builder_temp, output_model_location);
+  const uint8_t* buffer_temp = builder_temp.GetBufferPointer();
+  const Model* input_model = GetModel(buffer_temp);
+
+  // Change model type.
+  flatbuffers::FlatBufferBuilder builder;
+  EXPECT_EQ(Uint8QuantizeModelInputsOutputs(&builder, input_model,
+                                            {{"tensor_0", {0.4, 2}}},
+                                            {{"tensor_1", {0.5, -5}}}),
+            kTfLiteOk);
+
+  const uint8_t* buffer = builder.GetBufferPointer();
+  const Model* output_model = GetModel(buffer);
+  std::unique_ptr<ModelT> model;
+  model.reset(output_model->UnPack());
+
+  // Verify results.
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 4);
+  EXPECT_EQ(model->subgraphs[0]->inputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->inputs[0], 0);
+  EXPECT_EQ(model->subgraphs[0]->outputs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->outputs[0], 1);
+  EXPECT_EQ(model->operator_codes.size(), 3);
+  EXPECT_EQ(GetBuiltinCode(model->operator_codes[0].get()),
+            BuiltinOperator_FULLY_CONNECTED);
+  EXPECT_EQ(GetBuiltinCode(model->operator_codes[1].get()),
+            BuiltinOperator_DEQUANTIZE);
+  EXPECT_EQ(GetBuiltinCode(model->operator_codes[2].get()),
+            BuiltinOperator_QUANTIZE);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 3);
+
+  auto dequantize_op = model->subgraphs[0]->operators[0].get();
+  auto input = model->subgraphs[0]->tensors[dequantize_op->inputs[0]].get();
+  EXPECT_EQ(input->name, "tensor_0_uint8");
+  EXPECT_EQ(input->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(input->quantization->scale[0], 0.4);
+  EXPECT_EQ(input->quantization->zero_point[0], 2);
+
+  auto quantize_op = model->subgraphs[0]->operators[2].get();
+  auto output = model->subgraphs[0]->tensors[quantize_op->outputs[0]].get();
+  EXPECT_EQ(output->name, "tensor_1_uint8");
+  EXPECT_EQ(output->type, TensorType_UINT8);
+  EXPECT_FLOAT_EQ(output->quantization->scale[0], 0.5);
+  EXPECT_EQ(output->quantization->zero_point[0], -5);
+}
+
+}  // namespace
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/python/BUILD

@@ -0,0 +1,77 @@
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+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"],
+)
+
+py_binary(
+    name = "modify_model_interface",
+    srcs = ["modify_model_interface.py"],
+    strict_deps = True,
+    deps = [
+        ":modify_model_interface_constants",
+        ":modify_model_interface_lib",
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+    ],
+)
+
+py_library(
+    name = "modify_model_interface_lib",
+    srcs = ["modify_model_interface_lib.py"],
+    strict_deps = True,
+    deps = [
+        ":_pywrap_modify_model_interface",
+        ":modify_model_interface_constants",
+        "//tensorflow/lite/python:schema_py",
+    ],
+)
+
+# Use --config=disable_tf_lite_py when running this test under github.
+py_test(
+    name = "modify_model_interface_lib_test",
+    srcs = ["modify_model_interface_lib_test.py"],
+    strict_deps = True,
+    deps = [
+        ":modify_model_interface_lib",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "modify_model_interface_constants",
+    srcs = ["modify_model_interface_constants.py"],
+    strict_deps = True,
+    deps = ["//tensorflow/python/framework:dtypes"],
+)
+
+pybind_extension(
+    name = "_pywrap_modify_model_interface",
+    srcs = ["modify_model_interface.cc"],
+    common_lib_packages = [
+        "litert/python",
+        "tensorflow/lite/python",
+    ],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_modify_model_interface.pyi",
+    ],
+    wrap_py_init = True,
+    deps = [
+        "//tensorflow/lite/schema:schema_fbs",
+        "//tensorflow/lite/tools/optimize:modify_model_interface",
+        "@pybind11",
+    ],
+)

tensorflow/lite/tools/optimize/python/_pywrap_modify_model_interface.pyi

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+
+def modify_model_interface(arg0: str, arg1: str, arg2: int, arg3: int) -> int: ...

tensorflow/lite/tools/optimize/python/modify_model_interface.cc

@@ -0,0 +1,40 @@
+/* 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.
+==============================================================================*/
+
+// Python wrapper to modify model interface.
+
+#include "tensorflow/lite/tools/optimize/modify_model_interface.h"
+
+#include <string>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace pybind11 {
+
+PYBIND11_MODULE(_pywrap_modify_model_interface, m) {
+  // An anonymous function that invokes the C++ function
+  // after applying transformations to the python function arguments
+  m.def("modify_model_interface",
+        [](const std::string& input_file, const std::string& output_file,
+           const int input_type, const int output_type) -> int {
+          return tflite::optimize::ModifyModelInterface(
+              input_file, output_file,
+              static_cast<tflite::TensorType>(input_type),
+              static_cast<tflite::TensorType>(output_type));
+        });
+}
+
+}  // namespace pybind11

tensorflow/lite/tools/optimize/python/modify_model_interface.py

@@ -0,0 +1,54 @@
+# 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.
+# ==============================================================================
+r"""Modify a quantized model's interface from float to integer."""
+
+from absl import app
+from absl import flags
+
+from tensorflow.lite.tools.optimize.python import modify_model_interface_constants as mmi_constants
+from tensorflow.lite.tools.optimize.python import modify_model_interface_lib as mmi_lib
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string('input_tflite_file', None,
+                    'Full path name to the input TFLite file.')
+flags.DEFINE_string('output_tflite_file', None,
+                    'Full path name to the output TFLite file.')
+flags.DEFINE_enum('input_type', mmi_constants.DEFAULT_STR_TYPE,
+                  mmi_constants.STR_TYPES,
+                  'Modified input integer interface type.')
+flags.DEFINE_enum('output_type', mmi_constants.DEFAULT_STR_TYPE,
+                  mmi_constants.STR_TYPES,
+                  'Modified output integer interface type.')
+
+flags.mark_flag_as_required('input_tflite_file')
+flags.mark_flag_as_required('output_tflite_file')
+
+
+def main(_):
+  input_type = mmi_constants.STR_TO_TFLITE_TYPES[FLAGS.input_type]
+  output_type = mmi_constants.STR_TO_TFLITE_TYPES[FLAGS.output_type]
+
+  mmi_lib.modify_model_interface(FLAGS.input_tflite_file,
+                                 FLAGS.output_tflite_file, input_type,
+                                 output_type)
+
+  print('Successfully modified the model input type from FLOAT to '
+        '{input_type} and output type from FLOAT to {output_type}.'.format(
+            input_type=FLAGS.input_type, output_type=FLAGS.output_type))
+
+
+if __name__ == '__main__':
+  app.run(main)

tensorflow/lite/tools/optimize/python/modify_model_interface_constants.py

@@ -0,0 +1,30 @@
+# 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.
+# ==============================================================================
+"""Constants for modify_model_interface."""
+
+from tensorflow.python.framework import dtypes
+
+STR_TO_TFLITE_TYPES = {
+    'INT8': dtypes.int8,
+    'UINT8': dtypes.uint8,
+    'INT16': dtypes.int16,
+}
+TFLITE_TO_STR_TYPES = {v: k for k, v in STR_TO_TFLITE_TYPES.items()}
+
+STR_TYPES = STR_TO_TFLITE_TYPES.keys()
+TFLITE_TYPES = STR_TO_TFLITE_TYPES.values()
+
+DEFAULT_STR_TYPE = 'INT8'
+DEFAULT_TFLITE_TYPE = STR_TO_TFLITE_TYPES[DEFAULT_STR_TYPE]

tensorflow/lite/tools/optimize/python/modify_model_interface_lib.py

@@ -0,0 +1,74 @@
+# 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.
+# ==============================================================================
+"""Library to modify a quantized model's interface from float to integer."""
+
+from tensorflow.lite.python import schema_py_generated as schema_fb
+from tensorflow.lite.tools.optimize.python import _pywrap_modify_model_interface
+from tensorflow.lite.tools.optimize.python import modify_model_interface_constants as mmi_constants
+
+
+def _parse_type_to_int(dtype, flag):
+  """Converts a tflite type to it's integer representation.
+
+  Args:
+    dtype: tf.DType representing the inference type.
+    flag: str representing the flag name.
+
+  Returns:
+     integer, a tflite TensorType enum value.
+
+  Raises:
+    ValueError: Unsupported tflite type.
+  """
+  # Validate if dtype is supported in tflite and is a valid interface type.
+  if dtype not in mmi_constants.TFLITE_TYPES:
+    raise ValueError(
+        "Unsupported value '{0}' for {1}. Only {2} are supported.".format(
+            dtype, flag, mmi_constants.TFLITE_TYPES))
+
+  dtype_str = mmi_constants.TFLITE_TO_STR_TYPES[dtype]
+  dtype_int = schema_fb.TensorType.__dict__[dtype_str]
+
+  return dtype_int
+
+
+def modify_model_interface(input_file, output_file, input_type, output_type):
+  """Modify a quantized model's interface (input/output) from float to integer.
+
+  Args:
+    input_file: Full path name to the input tflite file.
+    output_file: Full path name to the output tflite file.
+    input_type: Final input interface type.
+    output_type: Final output interface type.
+
+  Raises:
+    RuntimeError: If the modification of the model interface was unsuccessful.
+    ValueError: If the input_type or output_type is unsupported.
+
+  """
+  # Map the interface types to integer values
+  input_type_int = _parse_type_to_int(input_type, 'input_type')
+  output_type_int = _parse_type_to_int(output_type, 'output_type')
+
+  # Invoke the function to modify the model interface
+  status = _pywrap_modify_model_interface.modify_model_interface(
+      input_file, output_file, input_type_int, output_type_int)
+
+  # Throw an exception if the return status is an error.
+  if status != 0:
+    raise RuntimeError(
+        'Error occurred when trying to modify the model input type from float '
+        'to {input_type} and output type from float to {output_type}.'.format(
+            input_type=input_type, output_type=output_type))

tensorflow/lite/tools/optimize/python/modify_model_interface_lib_test.py

@@ -0,0 +1,163 @@
+# 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.
+# ==============================================================================
+"""Tests for modify_model_interface_lib.py."""
+
+import os
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.lite.python import lite
+from tensorflow.lite.tools.optimize.python import modify_model_interface_lib
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+def build_tflite_model_with_full_integer_quantization(
+    supported_ops=lite.OpsSet.TFLITE_BUILTINS_INT8):
+  # Define TF model
+  input_size = 3
+  model = tf.keras.Sequential([
+      tf.keras.layers.InputLayer(input_shape=(input_size,), dtype=tf.float32),
+      tf.keras.layers.Dense(units=5, activation=tf.nn.relu),
+      tf.keras.layers.Dense(units=2, activation=tf.nn.softmax)
+  ])
+
+  # Convert TF Model to a Quantized TFLite Model
+  converter = lite.TFLiteConverterV2.from_keras_model(model)
+  converter.optimizations = [lite.Optimize.DEFAULT]
+
+  def representative_dataset_gen():
+    for i in range(10):
+      yield [np.array([i] * input_size, dtype=np.float32)]
+
+  converter.representative_dataset = representative_dataset_gen
+  converter.target_spec.supported_ops = [supported_ops]
+  tflite_model = converter.convert()
+
+  return tflite_model
+
+
+class ModifyModelInterfaceTest(test_util.TensorFlowTestCase):
+
+  def testInt8Interface(self):
+    # 1. SETUP
+    # Define the temporary directory and files
+    temp_dir = self.get_temp_dir()
+    initial_file = os.path.join(temp_dir, 'initial_model.tflite')
+    final_file = os.path.join(temp_dir, 'final_model.tflite')
+    # Define initial model
+    initial_model = build_tflite_model_with_full_integer_quantization()
+    with open(initial_file, 'wb') as model_file:
+      model_file.write(initial_model)
+
+    # 2. INVOKE
+    # Invoke the modify_model_interface function
+    modify_model_interface_lib.modify_model_interface(initial_file, final_file,
+                                                      tf.int8, tf.int8)
+
+    # 3. VALIDATE
+    # Load TFLite model and allocate tensors.
+    initial_interpreter = lite.Interpreter(model_path=initial_file)
+    initial_interpreter.allocate_tensors()
+    final_interpreter = lite.Interpreter(model_path=final_file)
+    final_interpreter.allocate_tensors()
+
+    # Get input and output types.
+    initial_input_dtype = initial_interpreter.get_input_details()[0]['dtype']
+    initial_output_dtype = initial_interpreter.get_output_details()[0]['dtype']
+    final_input_dtype = final_interpreter.get_input_details()[0]['dtype']
+    final_output_dtype = final_interpreter.get_output_details()[0]['dtype']
+
+    # Validate the model interfaces
+    self.assertEqual(initial_input_dtype, np.float32)
+    self.assertEqual(initial_output_dtype, np.float32)
+    self.assertEqual(final_input_dtype, np.int8)
+    self.assertEqual(final_output_dtype, np.int8)
+
+  def testInt16Interface(self):
+    # 1. SETUP
+    # Define the temporary directory and files
+    temp_dir = self.get_temp_dir()
+    initial_file = os.path.join(temp_dir, 'initial_model.tflite')
+    final_file = os.path.join(temp_dir, 'final_model.tflite')
+    # Define initial model
+    initial_model = build_tflite_model_with_full_integer_quantization(
+        supported_ops=lite.OpsSet
+        .EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8)
+    with open(initial_file, 'wb') as model_file:
+      model_file.write(initial_model)
+
+    # 2. INVOKE
+    # Invoke the modify_model_interface function
+    modify_model_interface_lib.modify_model_interface(initial_file, final_file,
+                                                      tf.int16, tf.int16)
+
+    # 3. VALIDATE
+    # Load TFLite model and allocate tensors.
+    initial_interpreter = lite.Interpreter(model_path=initial_file)
+    initial_interpreter.allocate_tensors()
+    final_interpreter = lite.Interpreter(model_path=final_file)
+    final_interpreter.allocate_tensors()
+
+    # Get input and output types.
+    initial_input_dtype = initial_interpreter.get_input_details()[0]['dtype']
+    initial_output_dtype = initial_interpreter.get_output_details()[0]['dtype']
+    final_input_dtype = final_interpreter.get_input_details()[0]['dtype']
+    final_output_dtype = final_interpreter.get_output_details()[0]['dtype']
+
+    # Validate the model interfaces
+    self.assertEqual(initial_input_dtype, np.float32)
+    self.assertEqual(initial_output_dtype, np.float32)
+    self.assertEqual(final_input_dtype, np.int16)
+    self.assertEqual(final_output_dtype, np.int16)
+
+  def testUInt8Interface(self):
+    # 1. SETUP
+    # Define the temporary directory and files
+    temp_dir = self.get_temp_dir()
+    initial_file = os.path.join(temp_dir, 'initial_model.tflite')
+    final_file = os.path.join(temp_dir, 'final_model.tflite')
+    # Define initial model
+    initial_model = build_tflite_model_with_full_integer_quantization()
+    with open(initial_file, 'wb') as model_file:
+      model_file.write(initial_model)
+
+    # 2. INVOKE
+    # Invoke the modify_model_interface function
+    modify_model_interface_lib.modify_model_interface(initial_file, final_file,
+                                                      tf.uint8, tf.uint8)
+
+    # 3. VALIDATE
+    # Load TFLite model and allocate tensors.
+    initial_interpreter = lite.Interpreter(model_path=initial_file)
+    initial_interpreter.allocate_tensors()
+    final_interpreter = lite.Interpreter(model_path=final_file)
+    final_interpreter.allocate_tensors()
+
+    # Get input and output types.
+    initial_input_dtype = initial_interpreter.get_input_details()[0]['dtype']
+    initial_output_dtype = initial_interpreter.get_output_details()[0]['dtype']
+    final_input_dtype = final_interpreter.get_input_details()[0]['dtype']
+    final_output_dtype = final_interpreter.get_output_details()[0]['dtype']
+
+    # Validate the model interfaces
+    self.assertEqual(initial_input_dtype, np.float32)
+    self.assertEqual(initial_output_dtype, np.float32)
+    self.assertEqual(final_input_dtype, np.uint8)
+    self.assertEqual(final_output_dtype, np.uint8)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/lite/tools/optimize/quantization_utils.cc

@@ -0,0 +1,860 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/quantization_utils.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <type_traits>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/logger.h"
+#include "tensorflow/lite/minimal_logging.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/tools/optimize/model_utils.h"
+
+namespace tflite {
+namespace optimize {
+namespace utils {
+
+namespace {
+
+// LINT.IfChange(QuantizationUtilsConstants)
+const int8_t kMinQuantizedValue8bit = -127;
+const int8_t kMaxQuantizedValue8bit = 127;
+const int8_t kMinQuantizedValue4bit = -7;
+const int8_t kMaxQuantizedValue4bit = 7;
+// The maximum number of dimensions supported in per-channel quantization.
+constexpr int kPerChannelMaxDim = 4;
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:QuantizationUtilsConstants)
+const int8_t kMinQuantizedValue2bit = -1;
+const int8_t kMaxQuantizedValue2bit = 1;
+}  // namespace
+
+// LINT.IfChange(NumElements)
+TfLiteStatus NumElements(const TensorT& tensor, uint64_t* num_elements) {
+  *num_elements = 1;
+  for (const int64_t dim : tensor.shape) {
+    if (dim <= 0 || *num_elements > UINT64_MAX / static_cast<uint64_t>(dim)) {
+      return kTfLiteError;
+    }
+    *num_elements *= dim;
+  }
+  return kTfLiteOk;
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:NumElements)
+
+// Nudge min and max so that floating point 0 falls exactly on a quantized
+// value, returning the nudges scale and zero_point.
+//
+// Although this code originates from FakeQuantization in quantized training,
+// we may deviate from that implementation as we please since we do not fine
+// tune the weights with quantized training.
+void GetAsymmetricQuantizationParams(
+    float min, float max, const int quant_min, const int quant_max,
+    QuantizationParametersT* quantization_params) {
+  const float quant_min_float = static_cast<float>(quant_min);
+  const float quant_max_float = static_cast<float>(quant_max);
+  // Adjust the boundaries to guarantee 0 is included.
+  min = std::min(static_cast<float>(min), 0.0f);
+  max = std::max(static_cast<float>(max), 0.0f);
+  const float scale = (max - min) / (quant_max_float - quant_min_float);
+  // Scale can be zero if min and max are exactly 0.0f.
+  float zero_point_from_min = quant_min_float;
+  if (scale != 0) {
+    zero_point_from_min = quant_min_float - min / scale;
+  }
+  int64_t zero_point;
+  if (zero_point_from_min < quant_min_float) {
+    zero_point = static_cast<int64_t>(quant_min);
+  } else if (zero_point_from_min > quant_max_float) {
+    zero_point = static_cast<int64_t>(quant_max);
+  } else {
+    zero_point = static_cast<int64_t>(std::round(zero_point_from_min));
+  }
+  quantization_params->min = std::vector<float>(1, min);
+  quantization_params->max = std::vector<float>(1, max);
+  quantization_params->scale = std::vector<float>(1, scale);
+  quantization_params->zero_point = std::vector<int64_t>(1, zero_point);
+}
+
+void GetSymmetricQuantizationParams(
+    float min, float max, const int half_quant_range,
+    QuantizationParametersT* quantization_params) {
+  // Adjust the boundaries to guarantee 0 is included.
+  min = std::min(min, 0.0f);
+  max = std::max(max, 0.0f);
+  const float scale = std::max(std::abs(max), std::abs(min)) / half_quant_range;
+  quantization_params->min = std::vector<float>(1, min);
+  quantization_params->max = std::vector<float>(1, max);
+  quantization_params->scale = std::vector<float>(1, scale);
+  quantization_params->zero_point = std::vector<int64_t>(1, 0);
+}
+
+TfLiteStatus GetQuantizationParams(TensorT* tensor, TensorType activations_type,
+                                   QuantizationParametersT* quantization_params,
+                                   ErrorReporter* error_reporter) {
+  if (activations_type == TensorType_INT8) {
+    GetAsymmetricQuantizationParams(
+        tensor->quantization->min[0], tensor->quantization->max[0],
+        std::numeric_limits<int8_t>::min(), std::numeric_limits<int8_t>::max(),
+        quantization_params);
+  } else if (activations_type == TensorType_INT16) {
+    const int half_quantized_range = 32767;
+    GetSymmetricQuantizationParams(tensor->quantization->min[0],
+                                   tensor->quantization->max[0],
+                                   half_quantized_range, quantization_params);
+  } else {
+    TF_LITE_REPORT_ERROR(
+        error_reporter,
+        "Unsupported activation type for quantize-activation: %d",
+        activations_type);
+    return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+// Set the max and min quantization parameter for a single tensor given its
+// values.
+void FillSingleMinMax(const float* const input, const uint64_t input_size,
+                      QuantizationParametersT* quantization_params) {
+  const auto minmax = std::minmax_element(input, input + input_size);
+  quantization_params->min.assign(1, *minmax.first);
+  quantization_params->max.assign(1, *minmax.second);
+}
+
+// LINT.IfChange(FillPerChannelMinMax)
+TfLiteStatus FillPerChannelMinMax(const float* const input,
+                                  const std::vector<int32_t>& dimension,
+                                  int32_t channel_dim_index,
+                                  QuantizationParametersT* quantization_params,
+                                  ErrorReporter* error_reporter) {
+  if (!quantization_params->min.empty() || !quantization_params->max.empty()) {
+    TF_LITE_REPORT_ERROR(
+        error_reporter,
+        "Min or max already present in tensor quantization params.");
+    return kTfLiteError;
+  }
+
+  if (dimension.size() > kPerChannelMaxDim) {
+    TF_LITE_REPORT_ERROR(
+        error_reporter,
+        "Expected tensor with less than %d dimensions, but got %d.",
+        kPerChannelMaxDim + 1, dimension.size());
+    return kTfLiteError;
+  }
+  if (channel_dim_index >= dimension.size()) {
+    TF_LITE_REPORT_ERROR(
+        error_reporter,
+        "Expected channel_dim_index to be less than %d, but got %d.",
+        dimension.size(), channel_dim_index);
+    return kTfLiteError;
+  }
+
+  const int32_t channel_dim_size = dimension[channel_dim_index];
+  quantization_params->quantized_dimension = channel_dim_index;
+  quantization_params->min = std::vector<float>(channel_dim_size);
+  quantization_params->max = std::vector<float>(channel_dim_size);
+  std::vector<bool> has_min_max_value(channel_dim_size, false);
+  int indices[kPerChannelMaxDim];
+  RuntimeShape unextended_tensor_dims(dimension.size(), dimension.data());
+  RuntimeShape tensor_dims =
+      RuntimeShape::ExtendedShape(kPerChannelMaxDim, unextended_tensor_dims);
+  channel_dim_index +=
+      kPerChannelMaxDim - unextended_tensor_dims.DimensionsCount();
+
+  // Compute min max ranges per channel
+  for (indices[0] = 0; indices[0] < tensor_dims.Dims(0); indices[0]++) {
+    for (indices[1] = 0; indices[1] < tensor_dims.Dims(1); indices[1]++) {
+      for (indices[2] = 0; indices[2] < tensor_dims.Dims(2); indices[2]++) {
+        for (indices[3] = 0; indices[3] < tensor_dims.Dims(3); indices[3]++) {
+          int channel_idx = indices[channel_dim_index];
+          const float val = input[Offset(tensor_dims, indices)];
+          if (has_min_max_value[channel_idx]) {
+            if (quantization_params->min[channel_idx] > val) {
+              quantization_params->min[channel_idx] = val;
+            } else if (quantization_params->max[channel_idx] < val) {
+              quantization_params->max[channel_idx] = val;
+            }
+          } else {
+            quantization_params->min[channel_idx] = val;
+            quantization_params->max[channel_idx] = val;
+            has_min_max_value[channel_idx] = true;
+          }
+        }
+      }
+    }
+  }
+  return kTfLiteOk;
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:FillPerChannelMinMax)
+
+// Populates the scales vector based on max and min values of quant_params
+TfLiteStatus GetSymmetricScalesFromMaxMin(QuantizationParametersT* quant_params,
+                                          std::vector<float>* scales,
+                                          ErrorReporter* error_reporter) {
+  // Check that max and min values are present and their sizes match.
+  if (quant_params->min.empty() || quant_params->max.empty()) {
+    TF_LITE_REPORT_ERROR(error_reporter,
+                         "Max and min values are not populated.");
+    return kTfLiteError;
+  }
+  if (quant_params->min.size() != quant_params->max.size()) {
+    TF_LITE_REPORT_ERROR(error_reporter,
+                         "Dimensions of max and min values do not match.");
+    return kTfLiteError;
+  }
+  if (scales->size() != quant_params->min.size()) {
+    TF_LITE_REPORT_ERROR(error_reporter,
+                         "Provided scale vector has incorrect size.");
+    return kTfLiteError;
+  }
+
+  // num_channels is calculated from min.size() to infer whether quantization
+  // is per axis.
+  int num_channels = quant_params->min.size();
+  // Calculate scales per channel.
+  for (int channel_idx = 0; channel_idx < num_channels; ++channel_idx) {
+    const float half_range = std::max(std::abs(quant_params->min[channel_idx]),
+                                      std::abs(quant_params->max[channel_idx]));
+    scales->at(channel_idx) = half_range / kMaxQuantizedValue8bit;
+  }
+  return kTfLiteOk;
+}
+
+// Checks that the bias is quantized to within the middle half of the
+// allowable bit range determined by the scales of the input and weight tensors.
+// If this condition is not satisfied, the scale of the weights is increased in
+// order to prevent overflow. The scale of the bias is not set here, only the
+// min/max.
+// The quant_params are the quantization parameters that correspond to the
+// weight tensor.
+TfLiteStatus AdjustWeightsForBiasScale(QuantizationParametersT* quant_params,
+                                       const float* bias_data,
+                                       const size_t bias_size,
+                                       const float input_scale,
+                                       ErrorReporter* error_reporter) {
+  // TODO(dmolitor) Allow adjusting activation scale.
+  // TODO(dmolitor) Tighten scale adjustment.
+  // TODO(dmolitor) Test using a separate strategy for scales of 0.
+  const int32_t kScale = std::numeric_limits<int32_t>::max();
+  if (quant_params == nullptr) {
+    TF_LITE_REPORT_ERROR(error_reporter,
+                         "Missing max and min values for weight tensor.");
+    return kTfLiteError;
+  }
+  // channel_dim_size is calculated from min.size() to infer whether
+  // quantization is per axis
+  int channel_dim_size = quant_params->min.size();
+  if (channel_dim_size == 0) {
+    TF_LITE_REPORT_ERROR(
+        error_reporter,
+        "Missing weight scales. Unable to check compatibility with bias "
+        "scale.");
+    return kTfLiteError;
+  }
+
+  std::vector<float> weight_scales(channel_dim_size);
+  TF_LITE_ENSURE_STATUS(GetSymmetricScalesFromMaxMin(
+      quant_params, &weight_scales, error_reporter));
+
+  // Per channel quantization
+  if (channel_dim_size > 1) {
+    for (int i = 0; i < channel_dim_size; ++i) {
+      // Current scale is not compatible with bias. Adjust max/min values.
+      if (std::abs(bias_data[i]) >=
+          0.5 * input_scale * weight_scales[i] * kScale) {
+        quant_params->max[i] = 2.0 * std::abs(bias_data[i]) / kScale *
+                               (kMaxQuantizedValue8bit / input_scale);
+        quant_params->min[i] = -quant_params->max[i];
+      }
+    }
+    // Per layer quantization
+  } else if (channel_dim_size == 1) {
+    const auto minmax = std::minmax_element(bias_data, bias_data + bias_size);
+    const float bias_half_range =
+        std::max(std::abs(*minmax.first), std::abs(*minmax.second));
+
+    // Need to adjust weight min/max; not compatible with bias.
+    if (bias_half_range / kScale >= 0.5 * input_scale * weight_scales[0]) {
+      quant_params->min[0] = 2.0 * bias_half_range / kScale *
+                             (kMinQuantizedValue8bit / input_scale);
+      quant_params->max[0] = 2.0 * bias_half_range / kScale *
+                             (kMaxQuantizedValue8bit / input_scale);
+    }
+  }
+  return kTfLiteOk;
+}
+
+// LINT.IfChange(SymmetricPerChannelQuantization)
+// Per-channel quantize a tensor at the given index and fills both scales and
+// quantized values.
+TfLiteStatus SymmetricPerChannelQuantization(TensorT* tensor,
+                                             const float* const input,
+                                             int32_t channel_dim_index,
+                                             std::vector<float>* output_scales,
+                                             std::vector<int8_t>* output_value,
+                                             ErrorReporter* error_reporter) {
+  if (tensor == nullptr) {
+    TF_LITE_REPORT_ERROR(error_reporter, "Cannot quantize. Tensor is null.");
+    return kTfLiteError;
+  }
+  const int32_t channel_dim_size = tensor->shape[channel_dim_index];
+  // Fill per channel max and min values if needed
+  if (tensor->quantization == nullptr) {
+    tensor->quantization = std::make_unique<QuantizationParametersT>();
+  }
+  if (!HasMinMax(tensor)) {
+    TF_LITE_ENSURE_STATUS(
+        FillPerChannelMinMax(input, tensor->shape, channel_dim_index,
+                             tensor->quantization.get(), error_reporter));
+  }
+
+  // Calculate scales per channel using max and min values from tensor.
+  std::vector<float> scale_invs(channel_dim_size);
+  const float half_scale = kMaxQuantizedValue8bit;
+  for (int channel_idx = 0; channel_idx < channel_dim_size; channel_idx++) {
+    const float half_range =
+        std::max(std::abs(tensor->quantization->min[channel_idx]),
+                 std::abs(tensor->quantization->max[channel_idx]));
+    output_scales->at(channel_idx) = half_range / half_scale;
+    if (half_range == 0) {
+      scale_invs[channel_idx] = 0;
+    } else {
+      scale_invs[channel_idx] = half_scale / half_range;
+    }
+  }
+
+  // Quantize the input values.
+  SymmetricPerChannelQuantizeValues(input, scale_invs, tensor->shape,
+                                    channel_dim_index, output_value);
+  return kTfLiteOk;
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:SymmetricPerChannelQuantization)
+
+std::vector<int16_t> SymmetricQuantizeFloatsToInt16(const float* data,
+                                                    uint64_t num_elements,
+                                                    float scaling_factor) {
+  // Compute the inverse of scale.
+  const float scaling_factor_inv =
+      (scaling_factor == 0) ? 0 : 1.0 / scaling_factor;
+  std::vector<int16_t> buffer(num_elements);
+  const int32_t kScale = std::numeric_limits<int16_t>::max();
+
+  for (size_t i = 0; i < num_elements; i++) {
+    const int32_t quantized_value =
+        static_cast<int32_t>(TfLiteRound(data[i] * scaling_factor_inv));
+    buffer[i] = std::min(kScale, std::max(-kScale, quantized_value));
+  }
+  return buffer;
+}
+
+TfLiteStatus SymmetricQuantizeFloatsToInt16(ModelT* model, TensorT* tensor,
+                                            float scaling_factor,
+                                            ErrorReporter* error_reporter) {
+  const BufferT* buffer = model->buffers[tensor->buffer].get();
+  const float* float_data = reinterpret_cast<const float*>(buffer->data.data());
+  uint64_t num_elements;
+  TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements));
+
+  auto final_buffer =
+      SymmetricQuantizeFloatsToInt16(float_data, num_elements, scaling_factor);
+  // Set the buffers and output type.
+  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(final_buffer.data());
+  size_t buffer_size = num_elements * sizeof(int16_t);
+  std::vector<float> scales(1, scaling_factor);
+  std::vector<int64_t> zero_points(1, 0);
+  return AddQuantizationParams(scales, zero_points, 0, uint8_buffer,
+                               buffer_size, TensorType_INT16, model, tensor,
+                               error_reporter);
+}
+
+void SymmetricPerBlockQuantizeValues(
+    const float* const input, const float* const scales_inv,
+    const std::vector<int32_t>& input_dimension,
+    const std::vector<int32_t>& scale_dimension, int32_t channel_dim_index,
+    std::vector<int8_t>* output_value, TfLiteType type) {
+  // Quantize the values.
+  int indices[kPerChannelMaxDim];
+  int blocksize =
+      input_dimension[channel_dim_index] / scale_dimension[channel_dim_index];
+  RuntimeShape unextended_tensor_dims(input_dimension.size(),
+                                      input_dimension.data());
+  RuntimeShape tensor_dims =
+      RuntimeShape::ExtendedShape(kPerChannelMaxDim, unextended_tensor_dims);
+  RuntimeShape unextended_scale_dims(scale_dimension.size(),
+                                     scale_dimension.data());
+  RuntimeShape scale_dims =
+      RuntimeShape::ExtendedShape(kPerChannelMaxDim, unextended_scale_dims);
+  channel_dim_index +=
+      kPerChannelMaxDim - unextended_tensor_dims.DimensionsCount();
+  for (indices[0] = 0; indices[0] < tensor_dims.Dims(0); indices[0]++) {
+    for (indices[1] = 0; indices[1] < tensor_dims.Dims(1); indices[1]++) {
+      for (indices[2] = 0; indices[2] < tensor_dims.Dims(2); indices[2]++) {
+        for (indices[3] = 0; indices[3] < tensor_dims.Dims(3); indices[3]++) {
+          int index = Offset(tensor_dims, indices);
+          int scale_indices[4] = {indices[0], indices[1], indices[2],
+                                  indices[3] / blocksize};
+          int scale_index = Offset(scale_dims, scale_indices);
+          const float val = input[index];
+          const int32_t quantized_value =
+              static_cast<int32_t>(TfLiteRound(val * scales_inv[scale_index]));
+          if (type == kTfLiteInt4) {
+            output_value->at(index) = std::min<int8_t>(
+                kMaxQuantizedValue4bit,
+                std::max<int8_t>(kMinQuantizedValue4bit, quantized_value));
+          } else if (type == kTfLiteInt2) {
+            output_value->at(index) = std::min<int8_t>(
+                kMaxQuantizedValue2bit,
+                std::max<int8_t>(kMinQuantizedValue2bit, quantized_value));
+          }
+        }
+      }
+    }
+  }
+}
+
+// LINT.IfChange(SymmetricPerChannelQuantizeValues)
+void SymmetricPerChannelQuantizeValues(const float* const input,
+                                       const std::vector<float>& scales_inv,
+                                       const std::vector<int32_t>& dimension,
+                                       int32_t channel_dim_index,
+                                       std::vector<int8_t>* output_value,
+                                       TfLiteType type) {
+  // Quantize the values.
+  int indices[kPerChannelMaxDim];
+  RuntimeShape unextended_tensor_dims(dimension.size(), dimension.data());
+  RuntimeShape tensor_dims =
+      RuntimeShape::ExtendedShape(kPerChannelMaxDim, unextended_tensor_dims);
+  channel_dim_index +=
+      kPerChannelMaxDim - unextended_tensor_dims.DimensionsCount();
+  for (indices[0] = 0; indices[0] < tensor_dims.Dims(0); indices[0]++) {
+    for (indices[1] = 0; indices[1] < tensor_dims.Dims(1); indices[1]++) {
+      for (indices[2] = 0; indices[2] < tensor_dims.Dims(2); indices[2]++) {
+        for (indices[3] = 0; indices[3] < tensor_dims.Dims(3); indices[3]++) {
+          int channel_idx = indices[channel_dim_index];
+          int index = Offset(tensor_dims, indices);
+          const float val = input[index];
+          const int32_t quantized_value =
+              static_cast<int32_t>(TfLiteRound(val * scales_inv[channel_idx]));
+          if (type == kTfLiteInt4) {
+            output_value->at(index) = std::min<int8_t>(
+                kMaxQuantizedValue4bit,
+                std::max<int8_t>(kMinQuantizedValue4bit, quantized_value));
+          } else {
+            output_value->at(index) = std::min<int8_t>(
+                kMaxQuantizedValue8bit,
+                std::max<int8_t>(kMinQuantizedValue8bit, quantized_value));
+          }
+        }
+      }
+    }
+  }
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:SymmetricPerChannelQuantizeValues)
+
+// Quantize the tensor using the max and min values recorded in its
+// quantization parameters. Applies per-layer quantization.
+TfLiteStatus SymmetricQuantizeTensorFromMinMax(ModelT* model, TensorT* tensor,
+                                               ErrorReporter* error_reporter) {
+  if (model == nullptr || tensor == nullptr) {
+    TF_LITE_REPORT_ERROR(error_reporter, "No tensor to quantize.");
+    return kTfLiteError;
+  }
+
+  BufferT* buffer = model->buffers[tensor->buffer].get();
+  if (buffer == nullptr) {
+    TF_LITE_REPORT_ERROR(error_reporter, "Missing buffer.");
+    return kTfLiteError;
+  }
+
+  if (!HasMinMax(tensor)) {
+    TF_LITE_REPORT_ERROR(error_reporter,
+                         "Missing min or max values for quantization.");
+    return kTfLiteError;
+  }
+  if (tensor->quantization->min.size() != 1 ||
+      tensor->quantization->max.size() != 1) {
+    TF_LITE_REPORT_ERROR(error_reporter,
+                         "Expected single entry in max and min.");
+    return kTfLiteError;
+  }
+
+  const float* float_data = reinterpret_cast<const float*>(buffer->data.data());
+  uint64_t num_elements;
+  TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements));
+
+  std::vector<int8_t> quantized_buffer;
+  quantized_buffer.resize(num_elements);
+
+  // Quantize tensor using recorded min and max values
+  float scaling_factor;
+  tensor_utils::SymmetricQuantizeFloats(
+      float_data, num_elements, quantized_buffer.data(),
+      tensor->quantization->min[0], tensor->quantization->max[0],
+      &scaling_factor);
+  tensor->quantization->scale = std::vector<float>(1, scaling_factor);
+  tensor->quantization->zero_point = std::vector<int64_t>(1, 0);
+
+  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(quantized_buffer.data());
+  model->buffers[tensor->buffer]->data.assign(uint8_buffer,
+                                              uint8_buffer + num_elements);
+  // Update the tensor type.
+  tensor->type = TensorType_INT8;
+
+  return kTfLiteOk;
+}
+
+// LINT.IfChange(SymmetricQuantizeTensor)
+TfLiteStatus SymmetricQuantizeTensor(ModelT* model, TensorT* tensor) {
+  if (model == nullptr || tensor == nullptr) {
+    TFLITE_LOG(TFLITE_LOG_ERROR, "No tensor to quantize.");
+    return kTfLiteError;
+  }
+
+  BufferT* buffer = model->buffers[tensor->buffer].get();
+  if (buffer == nullptr) {
+    TFLITE_LOG(TFLITE_LOG_ERROR, "Missing buffer.");
+    return kTfLiteError;
+  }
+  const float* float_data = reinterpret_cast<const float*>(buffer->data.data());
+  uint64_t num_elements;
+  TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements));
+
+  std::vector<int8_t> quantized_buffer;
+  quantized_buffer.resize(num_elements);
+
+  float min_value, max_value, scaling_factor;
+  tensor_utils::SymmetricQuantizeFloats(float_data, num_elements,
+                                        quantized_buffer.data(), &min_value,
+                                        &max_value, &scaling_factor);
+
+  if (tensor->quantization == nullptr) {
+    tensor->quantization = std::make_unique<QuantizationParametersT>();
+  }
+  tensor->quantization->scale = std::vector<float>(1, scaling_factor);
+  tensor->quantization->zero_point = std::vector<int64_t>(1, 0);
+
+  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(quantized_buffer.data());
+  model->buffers[tensor->buffer]->data.assign(uint8_buffer,
+                                              uint8_buffer + num_elements);
+
+  // Update the tensor type.
+  tensor->type = TensorType_INT8;
+
+  return kTfLiteOk;
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:SymmetricQuantizeTensor)
+
+// LINT.IfChange(QuantizeTensorFloat16)
+TfLiteStatus QuantizeTensorFloat16(ModelT* model, TensorT* tensor) {
+  if (model == nullptr || tensor == nullptr) {
+    TFLITE_LOG(TFLITE_LOG_ERROR, "No tensor to quantize.");
+    return kTfLiteError;
+  }
+
+  BufferT* buffer = model->buffers[tensor->buffer].get();
+  if (buffer == nullptr) {
+    TFLITE_LOG(TFLITE_LOG_ERROR, "Missing buffer.");
+    return kTfLiteError;
+  }
+
+  uint64_t num_elements;
+  TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements));
+
+  // Copy single byte buffer data to float vector to guard against
+  // misalignment.
+  std::vector<float> float_vector(num_elements);
+  uint8_t* first = buffer->data.data();
+  std::copy(first, first + buffer->data.size(),
+            reinterpret_cast<uint8_t*>(float_vector.data()));
+
+  // Transform float data to float16.
+  std::vector<Eigen::half> quantized_buffer;
+  quantized_buffer.resize(num_elements);
+  constexpr float kMaxFloat16Value = 65504.f;
+  constexpr float kMinFloat16Value = -65504.f;
+  std::transform(float_vector.begin(), float_vector.end(),
+                 quantized_buffer.begin(), [=](float a) {
+                   float clamped = std::min(std::max(a, kMinFloat16Value),
+                                            kMaxFloat16Value);
+                   return static_cast<Eigen::half>(clamped);
+                 });
+
+  char* half_buffer = reinterpret_cast<char*>(quantized_buffer.data());
+  model->buffers[tensor->buffer]->data.assign(
+      half_buffer, half_buffer + sizeof(Eigen::half) * num_elements);
+
+  // Update the tensor type.
+  tensor->type = TensorType_FLOAT16;
+
+  return kTfLiteOk;
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:QuantizeTensorFloat16)
+
+// LINT.IfChange(AddQuantizationParams)
+TfLiteStatus AddQuantizationParams(const std::vector<float>& scales,
+                                   const std::vector<int64_t>& zero_point,
+                                   int quantized_dimension,
+                                   const uint8_t* buffer_data,
+                                   size_t buffer_size, TensorType output_type,
+                                   ModelT* model, TensorT* tensor,
+                                   ErrorReporter* error_reporter) {
+  if (tensor->quantization == nullptr) {
+    tensor->quantization = std::make_unique<QuantizationParametersT>();
+  }
+  tensor->quantization->scale.assign(scales.begin(), scales.end());
+  if (zero_point.size() != scales.size()) {
+    TF_LITE_REPORT_ERROR(
+        error_reporter,
+        "Received zero_point of size %d and scales of size %d. "
+        "These sizes should match.",
+        zero_point.size(), scales.size());
+    return kTfLiteError;
+  }
+  tensor->quantization->zero_point.assign(zero_point.begin(), zero_point.end());
+  tensor->quantization->quantized_dimension = quantized_dimension;
+  model->buffers[tensor->buffer]->data.assign(buffer_data,
+                                              buffer_data + buffer_size);
+  // Update the tensor type.
+  tensor->type = output_type;
+  return kTfLiteOk;
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:AddQuantizationParams)
+
+// LINT.IfChange(SymmetricQuantizeTensorPerChannel)
+TfLiteStatus SymmetricQuantizeTensorPerChannel(ModelT* model, TensorT* tensor,
+                                               int32_t channel_dim_index,
+                                               ErrorReporter* error_reporter) {
+  if (tensor->shape.size() > kPerChannelMaxDim) {
+    TF_LITE_REPORT_ERROR(
+        error_reporter,
+        "SymmetricQuantizeTensorPerChannel requires tensor with less than %d "
+        "dimensions, but got %d dimension(s).",
+        kPerChannelMaxDim + 1, tensor->shape.size());
+    return kTfLiteError;
+  }
+
+  // Get dimensions.
+  uint64_t num_elements;
+  TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements));
+  const int32_t channel_dim_size = tensor->shape[channel_dim_index];
+
+  // Get input float data.
+  const BufferT* buffer = model->buffers[tensor->buffer].get();
+  const float* float_input_data =
+      reinterpret_cast<const float*>(buffer->data.data());
+
+  // Create container for output scale and output data.
+  std::vector<float> scales(channel_dim_size);
+  std::vector<int8_t> final_buffer(num_elements);
+
+  // Quantize the input data with respect to channel_dim_index.
+  TF_LITE_ENSURE_STATUS(SymmetricPerChannelQuantization(
+      tensor, float_input_data, channel_dim_index, &scales, &final_buffer,
+      error_reporter));
+
+  // Set the buffers and output type.
+  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(final_buffer.data());
+  const size_t buffer_size = num_elements * sizeof(int8_t);
+  std::vector<int64_t> zero_point(scales.size(), 0);
+  return AddQuantizationParams(scales, zero_point, channel_dim_index,
+                               uint8_buffer, buffer_size, TensorType_INT8,
+                               model, tensor, error_reporter);
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.cc:SymmetricQuantizeTensorPerChannel)
+
+template <class BiasType>
+std::vector<BiasType> SymmetricBiasQuantize(const float* data,
+                                            uint64_t num_elements,
+                                            const std::vector<float>& scales) {
+  std::vector<BiasType> buffer(num_elements);
+  const BiasType kScale = std::numeric_limits<BiasType>::max();
+  float scaling_factor_inv_per_layer = (scales[0] == 0) ? 0 : 1.0 / scales[0];
+
+  for (int32_t idx = 0; idx < num_elements; idx++) {
+    float scaling_factor_inv =
+        scales.size() == 1 ? scaling_factor_inv_per_layer
+                           : ((scales[idx] == 0) ? 0 : 1.0 / scales[idx]);
+    const BiasType quantized_value =
+        tflite::SafeCast<BiasType>(TfLiteRound(data[idx] * scaling_factor_inv));
+    buffer[idx] = std::min(kScale, std::max(-kScale, quantized_value));
+  }
+  return buffer;
+}
+
+template std::vector<std::int32_t> SymmetricBiasQuantize<std::int32_t>(
+    const float* data, uint64_t num_elements, const std::vector<float>& scales);
+
+template <class BiasType>
+TfLiteStatus SymmetricPerLayerBiasQuantize(ModelT* model, TensorT* tensor,
+                                           float scaling_factor,
+                                           ErrorReporter* error_reporter) {
+  const BufferT* buffer = model->buffers[tensor->buffer].get();
+  const float* float_data = reinterpret_cast<const float*>(buffer->data.data());
+  uint64_t num_elements;
+  TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements));
+
+  auto final_buffer = SymmetricBiasQuantize<BiasType>(float_data, num_elements,
+                                                      {scaling_factor});
+
+  // Set the buffers and output type.
+  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(final_buffer.data());
+  size_t buffer_size = num_elements * sizeof(BiasType);
+  std::vector<float> scales(1, scaling_factor);
+  std::vector<int64_t> zero_points(1, 0);
+
+  auto output_type = std::is_same<BiasType, std::int32_t>::value
+                         ? TensorType_INT32
+                         : TensorType_INT64;
+  return AddQuantizationParams(scales, zero_points, 0, uint8_buffer,
+                               buffer_size, output_type, model, tensor,
+                               error_reporter);
+}
+
+template TfLiteStatus SymmetricPerLayerBiasQuantize<std::int32_t>(
+    ModelT* model, TensorT* tensor, float scaling_factor,
+    ErrorReporter* error_reporter);
+
+template TfLiteStatus SymmetricPerLayerBiasQuantize<std::int64_t>(
+    ModelT* model, TensorT* tensor, float scaling_factor,
+    ErrorReporter* error_reporter);
+
+template <class BiasType>
+TfLiteStatus SymmetricPerChannelBiasQuantize(ModelT* model, TensorT* tensor,
+                                             float input_scale,
+                                             const float* weight_scales,
+                                             int number_of_dimension,
+                                             ErrorReporter* error_reporter) {
+  // Compute scales.
+  std::vector<float> scales(number_of_dimension);
+  for (int i = 0; i < number_of_dimension; i++) {
+    scales[i] = input_scale * weight_scales[i];
+  }
+
+  const BufferT* buffer = model->buffers[tensor->buffer].get();
+  const float* float_data = reinterpret_cast<const float*>(buffer->data.data());
+  uint64_t num_elements;
+  TF_LITE_ENSURE_STATUS(NumElements(*tensor, &num_elements));
+
+  auto final_buffer =
+      SymmetricBiasQuantize<BiasType>(float_data, num_elements, scales);
+
+  // Set the buffers and output type.
+  uint8_t* uint8_buffer = reinterpret_cast<uint8_t*>(final_buffer.data());
+  size_t buffer_size = num_elements * sizeof(BiasType);
+  std::vector<int64_t> zero_point(scales.size(), 0);
+
+  auto output_type = std::is_same<BiasType, std::int32_t>::value
+                         ? TensorType_INT32
+                         : TensorType_INT64;
+  return AddQuantizationParams(scales, zero_point, 0, uint8_buffer, buffer_size,
+                               output_type, model, tensor, error_reporter);
+}
+
+template TfLiteStatus SymmetricPerChannelBiasQuantize<std::int64_t>(
+    ModelT* model, TensorT* tensor, float input_scale,
+    const float* weight_scales, int number_of_dimension,
+    ErrorReporter* error_reporter);
+
+template TfLiteStatus SymmetricPerChannelBiasQuantize<std::int32_t>(
+    ModelT* model, TensorT* tensor, float input_scale,
+    const float* weight_scales, int number_of_dimension,
+    ErrorReporter* error_reporter);
+
+TfLiteStatus QuantizeWeight(ModelT* model, TensorT* tensor, bool per_channel,
+                            int per_axis_index, ErrorReporter* error_reporter) {
+  // TODO(suharshs): Currently we conflate quantizing weights and constants.
+  // Its possible that the right thing to do is asymmetric quantize the
+  // weight. Add support for this.
+  if (per_channel) {
+    return SymmetricQuantizeTensorPerChannel(model, tensor, per_axis_index,
+                                             error_reporter);
+  } else if (HasMinMax(tensor) && (tensor->quantization->min.size() == 1) &&
+             (tensor->quantization->max.size() == 1)) {
+    // Quantize using recorded min/max values if per-tensor.
+    return SymmetricQuantizeTensorFromMinMax(model, tensor, error_reporter);
+  } else {
+    // Quantize using min/max from buffer.
+    return SymmetricQuantizeTensor(model, tensor);
+  }
+}
+
+float GetEffectiveScale(ModelT* model, SubGraphT* subgraph, int op_idx,
+                        std::vector<int> input_index,
+                        std::vector<int> intermediate_index,
+                        std::vector<float> factors) {
+  float scale = 1.0f;
+  OperatorT* op = subgraph->operators[op_idx].get();
+  for (int i = 0, end = input_index.size(); i < end; ++i) {
+    const int index_local = input_index[i];
+    const int index_global = op->inputs[index_local];
+    const TensorT* tensor = subgraph->tensors[index_global].get();
+    scale *= tensor->quantization->scale[0];
+  }
+  for (int i = 0, end = intermediate_index.size(); i < end; ++i) {
+    const int index_local = intermediate_index[i];
+    const int index_global = op->intermediates[index_local];
+    const TensorT* tensor = subgraph->tensors[index_global].get();
+    scale *= tensor->quantization->scale[0];
+  }
+  for (int i = 0, end = factors.size(); i < end; ++i) {
+    scale *= factors[i];
+  }
+  return scale;
+}
+
+TfLiteStatus QuantizeActivation(TensorT* tensor, TensorType activations_type,
+                                ErrorReporter* error_reporter) {
+  TF_LITE_ENSURE_STATUS(GetQuantizationParams(
+      tensor, activations_type, tensor->quantization.get(), error_reporter));
+  tensor->type = activations_type;
+  return kTfLiteOk;
+}
+
+TfLiteStatus QuantizeActivationToInt16(TensorT* tensor, float scale) {
+  const int32_t zero_point = 0;
+  tensor->quantization = std::make_unique<QuantizationParametersT>();
+  tensor->quantization->scale.push_back(scale);
+  tensor->quantization->zero_point.push_back(zero_point);
+  tensor->type = TensorType_INT16;
+  return kTfLiteOk;
+}
+
+int GetPowerOfTwoScale(float min, float max) {
+  const float range = std::max(std::abs(min), std::abs(max));
+  int pot = 0;
+  for (int i = 0; i < 10; i++) {
+    // NOTE: use std::pow() for bitwise accuracy.
+    if (std::pow(2, pot) < range) {  // NOLINT
+      pot++;
+    }
+  }
+  return pot;
+}
+
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/quantization_utils.h

@@ -0,0 +1,193 @@
+/* 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_TOOLS_OPTIMIZE_QUANTIZATION_UTILS_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZATION_UTILS_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+#include "tensorflow/lite/c/c_api_types.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace optimize {
+namespace utils {
+
+// LINT.IfChange(num_elements)
+// Returns the number of elements in the given tensor.
+TfLiteStatus NumElements(const TensorT& tensor, uint64_t* num_elements);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:num_elements)
+
+// Populates the scale and zero point for quantization parameters.
+//
+// Nudges min and max so that floating point 0 falls exactly on a quantized
+// value, returning the nudges scale and zero_point.
+void GetAsymmetricQuantizationParams(
+    float min, float max, const int quant_min, const int quant_max,
+    QuantizationParametersT* quantization_params);
+
+// Populates the single total max and min values for a tensor.
+void FillSingleMinMax(const float* const input, const uint64_t input_size,
+                      QuantizationParametersT* quantization_params);
+
+// LINT.IfChange(fill_per_channel_min_max)
+// Populates the max and min values for per channel quantization.
+TfLiteStatus FillPerChannelMinMax(const float* const input,
+                                  const std::vector<int>& dimension,
+                                  int32_t channel_dim_index,
+                                  QuantizationParametersT* quantization_params,
+                                  ErrorReporter* error_reporter);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:fill_per_channel_min_max)
+
+void SymmetricPerBlockQuantizeValues(
+    const float* input, const float* scales_inv,
+    const std::vector<int32_t>& input_dimension,
+    const std::vector<int32_t>& scale_dimension, int32_t channel_dim_index,
+    std::vector<int8_t>* output_value, TfLiteType type);
+
+// LINT.IfChange(symmetric_per_channel_quantization)
+// Per-channel quantize a tensor at the given index and returns both scales and
+// quantized values.
+// Parameters:
+// - tensor is the tensor to be quantized, needed to access associated
+//   quantization parameters
+// - input is the float input data to be quantized.
+// - channel_dim_index is the channel index within "dimension".
+//   dimension[channel_dim_index] gives the number of channels.
+// - output_scale is the output scale, the size of which equals the number of
+//   channels.
+// - output_value is the output data, the size of which equals the number of
+//   inputs.
+TfLiteStatus SymmetricPerChannelQuantization(TensorT* tensor,
+                                             const float* const input,
+                                             int32_t channel_dim_index,
+                                             std::vector<float>* output_scales,
+                                             std::vector<int8_t>* output_value,
+                                             ErrorReporter* error_reporter);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:symmetric_per_channel_quantization)
+
+// LINT.IfChange(symmetric_per_channel_quantize_values)
+// Quantize the values given an array of scales.
+void SymmetricPerChannelQuantizeValues(const float* const input,
+                                       const std::vector<float>& scales_inv,
+                                       const std::vector<int32_t>& dimension,
+                                       int32_t channel_dim_index,
+                                       std::vector<int8_t>* output_value,
+                                       TfLiteType type = kTfLiteNoType);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:symmetric_per_channel_quantize_values)
+
+// LINT.IfChange(symmetric_quantize_tensor)
+// Quantizes tensor using symmetric quantization with the min and max elements
+// of the tensor.
+TfLiteStatus SymmetricQuantizeTensor(ModelT* model, TensorT* tensor);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:symmetric_quantize_tensor)
+
+// LINT.IfChange(quantize_tensor_float16)
+// Quantizes tensor to float16.
+TfLiteStatus QuantizeTensorFloat16(ModelT* model, TensorT* tensor);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:quantize_tensor_float16)
+
+// LINT.IfChange(add_quantization_params)
+// Add quantization parameters.
+TfLiteStatus AddQuantizationParams(const std::vector<float>& scales,
+                                   const std::vector<int64_t>& zero_point,
+                                   int quantized_dimension,
+                                   const uint8_t* buffer_data,
+                                   size_t buffer_size, TensorType output_type,
+                                   ModelT* model, TensorT* tensor,
+                                   ErrorReporter* error_reporter);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:add_quantization_params)
+
+// Populates the scales vector based on max and min values of quant_params
+TfLiteStatus GetSymmetricScalesFromMaxMin(QuantizationParametersT* quant_params,
+                                          std::vector<float>* scales,
+                                          ErrorReporter* error_reporter);
+
+// Adjusts scale of weights if incompatible with bias scale and likely to
+// cause overflow.
+TfLiteStatus AdjustWeightsForBiasScale(QuantizationParametersT* quant_params,
+                                       const float* bias_data,
+                                       const size_t bias_size,
+                                       const float input_scale,
+                                       ErrorReporter* error_reporter);
+
+// LINT.IfChange(symmetric_quantize_tensor_per_channel)
+// Quantizes tensor with per channel.
+TfLiteStatus SymmetricQuantizeTensorPerChannel(ModelT* model, TensorT* tensor,
+                                               int32_t channel_dim_index,
+                                               ErrorReporter* error_reporter);
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/quantization/lite/toco_legacy/quantization_utils.h:symmetric_quantize_tensor_per_channel)
+
+// Symmetrically quantizes float to 16bits.
+TfLiteStatus SymmetricQuantizeFloatsToInt16(ModelT* model, TensorT* tensor,
+                                            float scaling_factor,
+                                            ErrorReporter* error_reporter);
+
+std::vector<int16_t> SymmetricQuantizeFloatsToInt16(const float* data,
+                                                    uint64_t num_elements,
+                                                    float scaling_factor);
+
+// Symmetrically quantizes the bias for per-layer ops (i.e. FullyConnected).
+template <typename BiasType>
+TfLiteStatus SymmetricPerLayerBiasQuantize(ModelT* model, TensorT* tensor,
+                                           float scaling_factor,
+                                           ErrorReporter* error_reporter);
+
+// Symmetrically quantizes the bias for ops like Conv and DepthwiseConv.
+// The scale of bias if weight_per_channel_scale[channel] * input_scale.
+template <typename BiasType>
+TfLiteStatus SymmetricPerChannelBiasQuantize(ModelT* model, TensorT* tensor,
+                                             float input_scale,
+                                             const float* weight_scales,
+                                             int number_of_dimension,
+                                             ErrorReporter* error_reporter);
+
+template <typename BiasType>
+std::vector<BiasType> SymmetricBiasQuantize(const float* data,
+                                            uint64_t num_elements,
+                                            const std::vector<float>& scales);
+
+// Quantize weight with or without per channel.
+TfLiteStatus QuantizeWeight(ModelT* model, TensorT* tensor, bool per_channel,
+                            int per_axis_index, ErrorReporter* error_reporter);
+
+// Get effective scale by combining input scale, intermediate scale and factors.
+float GetEffectiveScale(ModelT* model, SubGraphT* subgraph, int op_idx,
+                        std::vector<int> input_index,
+                        std::vector<int> intermediate_index,
+                        std::vector<float> factors);
+
+// Return quantization parameters depending on activations type.
+TfLiteStatus GetQuantizationParams(TensorT* tensor, TensorType activations_type,
+                                   QuantizationParametersT* quantization_params,
+                                   ErrorReporter* error_reporter);
+
+// Quantize activation.
+TfLiteStatus QuantizeActivation(TensorT* tensor, TensorType activations_type,
+                                ErrorReporter* error_reporter);
+
+// Quantize activation to 16bit.
+TfLiteStatus QuantizeActivationToInt16(TensorT* tensor, float scale);
+
+// Get the power of two scale for min and max for symmetric quantization case.
+int GetPowerOfTwoScale(float min, float max);
+
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZATION_UTILS_H_

tensorflow/lite/tools/optimize/quantization_utils_test.cc

@@ -0,0 +1,919 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/quantization_utils.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <iostream>
+#include <limits>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/compiler/mlir/lite/quantization/lite/test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/util/command_line_flags.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace {
+std::string* g_test_model_dir = nullptr;
+}  // namespace
+
+namespace tflite {
+namespace optimize {
+namespace utils {
+namespace {
+
+std::unique_ptr<FlatBufferModel> ReadModel(const char* model) {
+  auto model_path = tensorflow::io::JoinPath(*g_test_model_dir, model);
+  return FlatBufferModel::BuildFromFile(model_path.c_str());
+}
+
+std::unique_ptr<FlatBufferModel> ReadConvModel() {
+  return ReadModel(mlir::lite::internal::kConvModelWith0Plus10Weights);
+}
+
+using ::testing::ElementsAreArray;
+
+class QuantizationUtilsTest : public testing::Test {
+ protected:
+  tflite::TestErrorReporter error_reporter_;
+};
+
+TEST_F(QuantizationUtilsTest, NumElements) {
+  TensorT tensor;
+  tensor.shape = {1, 2, 3, 4};
+  uint64_t num_elements;
+  EXPECT_EQ(kTfLiteOk, NumElements(tensor, &num_elements));
+  EXPECT_EQ(num_elements, 1 * 2 * 3 * 4);
+
+  tensor.shape = {5};
+  EXPECT_EQ(kTfLiteOk, NumElements(tensor, &num_elements));
+  EXPECT_EQ(num_elements, 5);
+
+  tensor.shape = {};
+  EXPECT_EQ(kTfLiteOk, NumElements(tensor, &num_elements));
+  // Scalars with empty shape have 1 element.
+  EXPECT_EQ(num_elements, 1);
+
+  tensor.shape = {1, 2, 3, -1};
+  EXPECT_EQ(kTfLiteError, NumElements(tensor, &num_elements));
+}
+
+TEST_F(QuantizationUtilsTest, GetAsymmetricQuantizationParamsUnitRange) {
+  const float float_min = -128.0;
+  const float float_max = 127.0;
+  const int quant_min = -128;
+  const int quant_max = 127;
+  QuantizationParametersT params;
+  GetAsymmetricQuantizationParams(float_min, float_max, quant_min, quant_max,
+                                  &params);
+  ASSERT_EQ(params.max.size(), 1);
+  ASSERT_EQ(params.min.size(), 1);
+  ASSERT_EQ(params.scale.size(), 1);
+  ASSERT_EQ(params.zero_point.size(), 1);
+  EXPECT_EQ(params.max[0], float_max);
+  EXPECT_EQ(params.min[0], float_min);
+
+  int64_t zero_point = params.zero_point[0];
+  float scale = params.scale[0];
+  const float eps = 1e-7f;
+  EXPECT_EQ(zero_point, 0);
+  EXPECT_NEAR(scale, 1, eps);
+}
+
+TEST_F(QuantizationUtilsTest,
+       AsymmetricQuantizationParamsWithAllPositiveRange) {
+  // The min should get nudged to include 0, so the effective range is [0, 6].
+  const float float_min = 1.0;
+  const float float_max = 6.0;
+  const int quant_min = -128;
+  const int quant_max = 127;
+  QuantizationParametersT params;
+  GetAsymmetricQuantizationParams(float_min, float_max, quant_min, quant_max,
+                                  &params);
+  ASSERT_EQ(params.max.size(), 1);
+  ASSERT_EQ(params.min.size(), 1);
+  ASSERT_EQ(params.scale.size(), 1);
+  ASSERT_EQ(params.zero_point.size(), 1);
+  EXPECT_EQ(params.max[0], float_max);
+  EXPECT_EQ(params.min[0], 0.0);
+  int64_t zero_point = params.zero_point[0];
+  float scale = params.scale[0];
+  const float eps = 1e-7f;
+  EXPECT_EQ(zero_point, -128);
+  EXPECT_NEAR(scale, 6 / 255.0f, eps);
+}
+
+TEST_F(QuantizationUtilsTest,
+       AsymmetricQuantizationParamsWithAllNegativeRange) {
+  // The max should get nudged to include 0, so the effective range is [-6, 0].
+  const float float_min = -6.0;
+  const float float_max = -1.0;
+  const int quant_min = -128;
+  const int quant_max = 127;
+  QuantizationParametersT params;
+  GetAsymmetricQuantizationParams(float_min, float_max, quant_min, quant_max,
+                                  &params);
+  ASSERT_EQ(params.max.size(), 1);
+  ASSERT_EQ(params.min.size(), 1);
+  ASSERT_EQ(params.scale.size(), 1);
+  ASSERT_EQ(params.zero_point.size(), 1);
+  EXPECT_EQ(params.max[0], 0.0);
+  EXPECT_EQ(params.min[0], float_min);
+  int64_t zero_point = params.zero_point[0];
+  float scale = params.scale[0];
+  const float eps = 1e-7f;
+  EXPECT_EQ(zero_point, 127);
+  EXPECT_NEAR(scale, 6 / 255.0f, eps);
+}
+
+TEST_F(QuantizationUtilsTest, AsymmetricQuantizationParamsWithZeroInRange) {
+  const float float_min = -5.0;
+  const float float_max = 1.0;
+  const int quant_min = -128;
+  const int quant_max = 127;
+  QuantizationParametersT params;
+  GetAsymmetricQuantizationParams(float_min, float_max, quant_min, quant_max,
+                                  &params);
+  ASSERT_EQ(params.max.size(), 1);
+  ASSERT_EQ(params.min.size(), 1);
+  ASSERT_EQ(params.scale.size(), 1);
+  ASSERT_EQ(params.zero_point.size(), 1);
+  EXPECT_EQ(params.max[0], float_max);
+  EXPECT_EQ(params.min[0], float_min);
+  int64_t zero_point = params.zero_point[0];
+  float scale = params.scale[0];
+  const float eps = 1e-7f;
+  EXPECT_NEAR(scale, 6 / 255.0f, eps);
+  EXPECT_GT(zero_point, quant_min);
+  EXPECT_LT(zero_point, quant_max);
+}
+
+TEST_F(QuantizationUtilsTest, AsymmetricQuantizationParamsWithZeroMinMax) {
+  const float float_min = 0;
+  const float float_max = 0;
+  const int quant_min = -128;
+  const int quant_max = 127;
+  QuantizationParametersT params;
+  GetAsymmetricQuantizationParams(float_min, float_max, quant_min, quant_max,
+                                  &params);
+  ASSERT_EQ(params.max.size(), 1);
+  ASSERT_EQ(params.min.size(), 1);
+  ASSERT_EQ(params.scale.size(), 1);
+  ASSERT_EQ(params.zero_point.size(), 1);
+  EXPECT_EQ(params.max[0], float_max);
+  EXPECT_EQ(params.min[0], float_min);
+  int64_t zero_point = params.zero_point[0];
+  float scale = params.scale[0];
+  const float eps = 1e-7f;
+  EXPECT_NEAR(scale, 0, eps);
+  EXPECT_NEAR(zero_point, quant_min, eps);
+  EXPECT_LT(zero_point, quant_max);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantizationWithNullQParams) {
+  // Set up an input with [3, 2, 2, 2] size and 0 is the channel index.
+  const std::vector<float> input = {
+      3.0, 2.0, 5.0,  -2.0, 3.0,  2.0,  5.0,  -2.0,  // Channel 1.
+      1.0, 2.0, 3.0,  4.0,  5.0,  6.0,  7.0,  8.0,   // Channel 2.
+      1.0, 0.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0,  // Channel 3.
+  };
+  const int channel_index = 0;
+
+  // Create holder for output scale and data.
+  std::vector<float> output_scales(3);
+  std::vector<int8_t> output_data(3 * 2 * 2 * 2);
+
+  // Call SymmetricPerChannelQuantization with quant_params as a null pointer
+  // and verify the result.
+  TensorT tensor = TensorT();
+  tensor.quantization = nullptr;
+  tensor.shape = {3, 2, 2, 2};
+  SymmetricPerChannelQuantization(&tensor, input.data(), channel_index,
+                                  &output_scales, &output_data,
+                                  &error_reporter_);
+  const std::vector<float> expected_output_scales = {0.0393700786, 0.0629921257,
+                                                     0.0472440943};
+  const std::vector<int8_t> expected_output_data = {
+      76, 51, 127, -51, 76,  51,  127,  -51,   // Channel 1.
+      16, 32, 48,  64,  79,  95,  111,  127,   // Channel 2.
+      21, 0,  -21, -42, -64, -85, -106, -127,  // Channel 3.
+  };
+  EXPECT_THAT(output_scales, ElementsAreArray(expected_output_scales));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantization) {
+  // Set up an input with [3, 2, 2, 2] size and 0 is the channel index.
+  const std::vector<float> input = {
+      3.0, 2.0, 5.0,  -2.0, 3.0,  2.0,  5.0,  -2.0,  // Channel 1.
+      1.0, 2.0, 3.0,  4.0,  5.0,  6.0,  7.0,  8.0,   // Channel 2.
+      1.0, 0.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0,  // Channel 3.
+  };
+  const int32_t channel_index = 0;
+
+  // Create holder for output scale and data.
+  std::vector<float> output_scales(3);
+  std::vector<int8_t> output_data(3 * 2 * 2 * 2);
+
+  // Initialize pointer to quantization parameters
+  TensorT tensor = TensorT();
+  tensor.quantization = std::make_unique<QuantizationParametersT>();
+  tensor.shape = {3, 2, 2, 2};
+  FillPerChannelMinMax(input.data(), tensor.shape, channel_index,
+                       tensor.quantization.get(), &error_reporter_);
+
+  // Test that FillPerChanneMinMax worked
+  const std::vector<float> expected_mins = {-2.0, 1.0, -6.0};
+  const std::vector<float> expected_maxs = {5.0, 8.0, 1.0};
+  EXPECT_THAT(tensor.quantization->min, ElementsAreArray(expected_mins));
+  EXPECT_THAT(tensor.quantization->max, ElementsAreArray(expected_maxs));
+
+  // Call SymmetricPerChannelQuantization with quant_params as a null pointer
+  // and verify the result.
+  SymmetricPerChannelQuantization(&tensor, input.data(), channel_index,
+                                  &output_scales, &output_data,
+                                  &error_reporter_);
+  const std::vector<float> expected_output_scales = {0.0393700786, 0.0629921257,
+                                                     0.0472440943};
+  const std::vector<int8_t> expected_output_data = {
+      76, 51, 127, -51, 76,  51,  127,  -51,   // Channel 1.
+      16, 32, 48,  64,  79,  95,  111,  127,   // Channel 2.
+      21, 0,  -21, -42, -64, -85, -106, -127,  // Channel 3.
+  };
+  EXPECT_THAT(output_scales, ElementsAreArray(expected_output_scales));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantization2DTensor) {
+  // Set up an input with [3, 8] size and 0 is the channel index.
+  const std::vector<float> input = {
+      3.0, 2.0, 5.0,  -2.0, 3.0,  2.0,  5.0,  -2.0,  // Batch 1.
+      1.0, 2.0, 3.0,  4.0,  5.0,  6.0,  7.0,  8.0,   // Batch 2.
+      1.0, 0.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0,  // Batch 3.
+  };
+  const int32_t channel_index = 1;
+
+  // Create holder for output scale and data.
+  std::vector<float> output_scales(8);
+  std::vector<int8_t> output_data(3 * 8);
+
+  // Initialize pointer to quantization parameters
+  TensorT tensor = TensorT();
+  tensor.quantization = std::make_unique<QuantizationParametersT>();
+  tensor.shape = {3, 8};
+  FillPerChannelMinMax(input.data(), tensor.shape, channel_index,
+                       tensor.quantization.get(), &error_reporter_);
+
+  // Test that FillPerChanneMinMax worked
+  const std::vector<float> expected_mins = {1.0,  0.0,  -1.0, -2.0,
+                                            -3.0, -4.0, -5.0, -6.0};
+  const std::vector<float> expected_maxs = {3.0, 2.0, 5.0, 4.0,
+                                            5.0, 6.0, 7.0, 8.0};
+  EXPECT_THAT(tensor.quantization->min, ElementsAreArray(expected_mins));
+  EXPECT_THAT(tensor.quantization->max, ElementsAreArray(expected_maxs));
+
+  // Call SymmetricPerChannelQuantization with quant_params as a null pointer
+  // and verify the result.
+  SymmetricPerChannelQuantization(&tensor, input.data(), channel_index,
+                                  &output_scales, &output_data,
+                                  &error_reporter_);
+  const std::vector<float> expected_output_scales = {
+      0.02362204724, 0.01574803149, 0.03937007874, 0.03149606299,
+      0.03937007874, 0.04724409448, 0.05511811023, 0.06299212598};
+  const std::vector<int8_t> expected_output_data = {
+      127, 127, 127, -64, 76,  42,  91,  -32,  // Batch 1.
+      42,  127, 76,  127, 127, 127, 127, 127,  // Batch 2.
+      42,  0,   -25, -64, -76, -85, -91, -95,  // Batch 3.
+  };
+  EXPECT_THAT(output_scales, ElementsAreArray(expected_output_scales));
+  EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricPerChannelQuantizeValues) {
+  // Set up an input with [3, 1, 1, 2] size and 0 is the channel index.
+  const std::vector<float> input = {
+      13.0, 21.0,  // Channel 1.
+      21.0, 22.0,  // Channel 2.
+      31.0, 40.0,  // Channel 3.
+  };
+  const std::vector<float> scales_inv = {2, 0.5, 3};
+  const std::vector<int32_t> dimension = {3, 1, 1, 2};
+  const int channel_index = 0;
+
+  // Create holder for output data.
+  std::vector<int8_t> output_data(3 * 1 * 1 * 2);
+
+  // Call SymmetricPerChannelQuantizeValues and verify the result.
+  SymmetricPerChannelQuantizeValues(input.data(), scales_inv, dimension,
+                                    channel_index, &output_data);
+  const std::vector<int8_t> expected_output_data = {
+      26, 42,   // Channel 1.
+      11, 11,   // Channel 2.
+      93, 120,  // Channel 3.
+  };
+  EXPECT_THAT(output_data, ElementsAreArray(expected_output_data));
+}
+
+TEST_F(QuantizationUtilsTest, FillPerChannelMinMax) {
+  // Set up an input with [3, 1, 1, 2] size.
+  const std::vector<float> input = {
+      13.0, 21.0,  // Channel 1.
+      21.0, 22.0,  // Channel 2.
+      31.0, 40.0,  // Channel 3.
+  };
+
+  // Initialize pointer to quantization parameters.
+  QuantizationParametersT quantization_params = QuantizationParametersT();
+  std::vector<int> dimension = {3, 1, 1, 2};
+  int32_t channel_dim_idx = 0;
+  const std::vector<float> expected_mins = {13.0, 21.0, 31.0};
+  const std::vector<float> expected_maxs = {21.0, 22.0, 40.0};
+
+  FillPerChannelMinMax(input.data(), dimension, channel_dim_idx,
+                       &quantization_params, &error_reporter_);
+
+  EXPECT_EQ(quantization_params.min, expected_mins);
+  EXPECT_EQ(quantization_params.max, expected_maxs);
+  EXPECT_EQ(quantization_params.quantized_dimension, channel_dim_idx);
+}
+
+TEST_F(QuantizationUtilsTest, FillPerChannelMinMaxFillDim3) {
+  // Set up an input with [3, 1, 1, 2] size.
+  const std::vector<float> input = {
+      // Channel 1, Channel 2
+      13.0, 21.0, 21.0, 22.0, 31.0, 40.0,
+  };
+
+  // Initialize pointer to quantization parameters.
+  QuantizationParametersT quantization_params = QuantizationParametersT();
+  std::vector<int> dimension = {3, 1, 1, 2};
+  int32_t channel_dim_idx = 3;
+  const std::vector<float> expected_mins = {13.0, 21.0};
+  const std::vector<float> expected_maxs = {31.0, 40.0};
+
+  FillPerChannelMinMax(input.data(), dimension, channel_dim_idx,
+                       &quantization_params, &error_reporter_);
+
+  EXPECT_EQ(quantization_params.min, expected_mins);
+  EXPECT_EQ(quantization_params.max, expected_maxs);
+  EXPECT_EQ(quantization_params.quantized_dimension, channel_dim_idx);
+}
+
+TEST_F(QuantizationUtilsTest, FillPerChannelMinMax2DTensor) {
+  // Set up an input with [3, 2] size.
+  const std::vector<float> input = {
+      // Channel 1, Channel 2
+      13.0, 21.0, 21.0, 22.0, 31.0, 40.0,
+  };
+
+  // Initialize pointer to quantization parameters.
+  QuantizationParametersT quantization_params = QuantizationParametersT();
+  std::vector<int> dimension = {3, 2};
+  int32_t channel_dim_idx = 1;
+  const std::vector<float> expected_mins = {13.0, 21.0};
+  const std::vector<float> expected_maxs = {31.0, 40.0};
+
+  FillPerChannelMinMax(input.data(), dimension, channel_dim_idx,
+                       &quantization_params, &error_reporter_);
+
+  EXPECT_EQ(quantization_params.min, expected_mins);
+  EXPECT_EQ(quantization_params.max, expected_maxs);
+  EXPECT_EQ(quantization_params.quantized_dimension, channel_dim_idx);
+}
+
+TEST_F(QuantizationUtilsTest, FillSingleMinMax) {
+  // Set up an input with [3, 1, 1, 2] size.
+  const std::vector<float> input = {
+      13.0, 21.0,  // Channel 1.
+      21.0, 22.0,  // Channel 2.
+      31.0, 40.0,  // Channel 3.
+  };
+  const uint32_t input_size = input.size();
+
+  // Initialize pointer to quantization parameters.
+  QuantizationParametersT quantization_params = QuantizationParametersT();
+
+  FillSingleMinMax(input.data(), input_size, &quantization_params);
+  const std::vector<float> expected_min_max = {
+      13, 40,  // min max
+  };
+  EXPECT_EQ(quantization_params.min.size(), 1);
+  EXPECT_EQ(quantization_params.max.size(), 1);
+  EXPECT_EQ(quantization_params.min[0], expected_min_max[0]);
+  EXPECT_EQ(quantization_params.max[0], expected_min_max[1]);
+}
+
+// kMaxQuantizedValue should match that of quantization_utils.cc.
+TEST_F(QuantizationUtilsTest, GetSymmetricScalesFromMaxMin) {
+  const int8_t kMaxQuantizedValue = 127;
+  tflite::TestErrorReporter error_reporter_;
+  // Create data.
+  auto quantization = std::make_unique<QuantizationParametersT>();
+  quantization->min = {-0.00001, -7.0, -2.0};
+  quantization->max = {0.00001, 1.0, -1.0};
+  std::vector<float> scales = std::vector<float>(quantization->min.size());
+
+  GetSymmetricScalesFromMaxMin(quantization.get(), &scales, &error_reporter_);
+  const std::vector<float> expected_scales = {0.00001 / kMaxQuantizedValue,
+                                              7.0 / kMaxQuantizedValue,
+                                              2.0 / kMaxQuantizedValue};
+  EXPECT_EQ(scales, expected_scales);
+}
+
+// kMaxQuantizedValue should match that of quantization_utils.cc
+TEST_F(QuantizationUtilsTest, AdjustWeightScaleForBiasPerChannel) {
+  // 2^(31) = 2147483648
+  const int32_t kScale = std::numeric_limits<int32_t>::max();
+  const int8_t kMaxQuantizedValue = 127;
+  tflite::TestErrorReporter error_reporter_;
+
+  // Create data.
+  auto quant_params = std::make_unique<QuantizationParametersT>();
+  const float small_val = 0.0000001;
+  const std::vector<float> orig_mins = {-small_val, -7.0};
+  quant_params->min = orig_mins;
+  const std::vector<float> orig_maxs = {small_val, 1.0};
+  quant_params->max = orig_maxs;
+  std::vector<float> scales = std::vector<float>(quant_params->min.size());
+
+  GetSymmetricScalesFromMaxMin(quant_params.get(), &scales, &error_reporter_);
+  const std::vector<float> expected_scales = {small_val / kMaxQuantizedValue,
+                                              7.0 / kMaxQuantizedValue};
+  EXPECT_EQ(scales, expected_scales);
+
+  const float input_scale = 0.05;
+
+  // Initialize bias.
+  float bias_data[] = {4.0, 4.0};
+  const size_t bias_size = 2;
+  // Quantized bias would be {101600000000, -- }.
+
+  AdjustWeightsForBiasScale(quant_params.get(), bias_data, bias_size,
+                            input_scale, &error_reporter_);
+
+  std::vector<float> new_scales = std::vector<float>(quant_params->min.size());
+  GetSymmetricScalesFromMaxMin(quant_params.get(), &new_scales,
+                               &error_reporter_);
+
+  // Adjust min and max for first channel.
+  EXPECT_TRUE(new_scales[0] > scales[0]);
+  EXPECT_TRUE(std::abs(bias_data[0]) / kScale <=
+              0.6 * input_scale * new_scales[0]);
+  EXPECT_TRUE(orig_mins[0] > quant_params->min[0]);
+  EXPECT_TRUE(orig_maxs[0] < quant_params->max[0]);
+  // No change for second channel.
+  EXPECT_TRUE(std::abs(bias_data[1]) / kScale <=
+              0.6 * input_scale * new_scales[1]);
+  EXPECT_EQ(orig_mins[1], quant_params->min[1]);
+  EXPECT_EQ(orig_maxs[1], quant_params->max[1]);
+}
+
+// kMaxQuantizedValue should match that of quantization_utils.cc.
+TEST_F(QuantizationUtilsTest, AdjustWeightScaleForBiasPerLayer) {
+  // 2^(31) = 2147483648
+  const int32_t kScale = std::numeric_limits<int32_t>::max();
+  const int8_t kMaxQuantizedValue = 127;
+  tflite::TestErrorReporter error_reporter_;
+
+  // Create data.
+  auto quant_params = std::make_unique<QuantizationParametersT>();
+  float small_val = 0.0000001;
+  const std::vector<float> orig_mins = {-small_val};
+  quant_params->min = orig_mins;
+  const std::vector<float> orig_maxs = {small_val};
+  quant_params->max = orig_maxs;
+  std::vector<float> scales = std::vector<float>(quant_params->min.size());
+  GetSymmetricScalesFromMaxMin(quant_params.get(), &scales, &error_reporter_);
+  const std::vector<float> expected_scales = {small_val / kMaxQuantizedValue};
+  EXPECT_EQ(scales, expected_scales);
+
+  const float input_scale = 0.05;
+
+  // Initialize bias.
+  float bias_data[] = {4.0};
+  const size_t bias_size = 1;
+  // Quantized bias would be {101600000000}.
+
+  AdjustWeightsForBiasScale(quant_params.get(), bias_data, bias_size,
+                            input_scale, &error_reporter_);
+
+  std::vector<float> new_scales = std::vector<float>(quant_params->min.size());
+  GetSymmetricScalesFromMaxMin(quant_params.get(), &new_scales,
+                               &error_reporter_);
+
+  // Adjust min and max.
+  EXPECT_TRUE(new_scales[0] > scales[0]);
+  EXPECT_TRUE(std::abs(bias_data[0]) / kScale <=
+              0.6 * input_scale * new_scales[0]);
+  EXPECT_TRUE(orig_mins[0] > quant_params->min[0]);
+  EXPECT_TRUE(orig_maxs[0] < quant_params->max[0]);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricQuantizeTensorFromMinMax) {
+  // Conv model has weights between 0 and 10.
+  // Quantize the weights tensor.
+  ASSERT_TRUE(g_test_model_dir);
+  ASSERT_FALSE(g_test_model_dir->empty());
+  auto test_model = ReadConvModel();
+  ASSERT_TRUE(test_model);
+  auto readonly_model = test_model->GetModel();
+  ASSERT_TRUE(readonly_model);
+  ASSERT_TRUE(readonly_model->subgraphs());
+  ASSERT_GE(readonly_model->subgraphs()->size(), 1);
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+  auto subgraph = model.subgraphs[0].get();
+  auto conv_op = subgraph->operators.at(0).get();
+  ASSERT_EQ(
+      GetBuiltinCode(model.operator_codes.at(conv_op->opcode_index).get()),
+      BuiltinOperator_CONV_2D);
+  int32_t weights_tensor_idx = conv_op->inputs[1];
+  TensorT* weights_tensor = subgraph->tensors.at(weights_tensor_idx).get();
+
+  // The test model has the incorrect number of min and max values for per-layer
+  // quantization.
+  weights_tensor->quantization->min =
+      std::vector<float>(1, weights_tensor->quantization->min[0]);
+  weights_tensor->quantization->max =
+      std::vector<float>(1, weights_tensor->quantization->max[0]);
+
+  EXPECT_EQ(weights_tensor->type, TensorType_FLOAT32);
+  size_t float_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+
+  bool per_channel = false;
+  int per_axis_index = 0;
+  tflite::TestErrorReporter error_reporter_;
+  EXPECT_EQ(QuantizeWeight(&model, weights_tensor, per_channel, per_axis_index,
+                           &error_reporter_),
+            kTfLiteOk);
+
+  size_t quant_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+  EXPECT_EQ(weights_tensor->type, TensorType_INT8);
+  EXPECT_EQ(quant_buffer_size * 4, float_buffer_size);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricQuantizeTensorNullInputs) {
+  tflite::TestErrorReporter error_reporter_;
+  EXPECT_EQ(SymmetricQuantizeTensor(nullptr, nullptr), kTfLiteError);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricQuantizeTensorNullQuantParams) {
+  // Conv model has weights between 0 and 10.
+  // Quantize the weights tensor.
+  ASSERT_TRUE(g_test_model_dir);
+  ASSERT_FALSE(g_test_model_dir->empty());
+  auto test_model = ReadConvModel();
+  ASSERT_TRUE(test_model);
+  auto readonly_model = test_model->GetModel();
+  ASSERT_TRUE(readonly_model);
+  ASSERT_TRUE(readonly_model->subgraphs());
+  ASSERT_GE(readonly_model->subgraphs()->size(), 1);
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+  auto subgraph = model.subgraphs[0].get();
+  auto conv_op = subgraph->operators.at(0).get();
+  ASSERT_EQ(
+      GetBuiltinCode(model.operator_codes.at(conv_op->opcode_index).get()),
+      BuiltinOperator_CONV_2D);
+  int32_t weights_tensor_idx = conv_op->inputs[1];
+  TensorT* weights_tensor = subgraph->tensors.at(weights_tensor_idx).get();
+  // Empty quantization parameters.
+  weights_tensor->quantization = std::make_unique<QuantizationParametersT>();
+
+  EXPECT_EQ(weights_tensor->type, TensorType_FLOAT32);
+  size_t float_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+
+  EXPECT_EQ(SymmetricQuantizeTensor(&model, weights_tensor), kTfLiteOk);
+
+  size_t quant_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+  EXPECT_EQ(weights_tensor->type, TensorType_INT8);
+  EXPECT_EQ(quant_buffer_size * 4, float_buffer_size);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricQuantizeTensor) {
+  // Conv model has weights between 0 and 10.
+  // Quantize the weights tensor.
+  ASSERT_TRUE(g_test_model_dir);
+  ASSERT_FALSE(g_test_model_dir->empty());
+  auto test_model = ReadConvModel();
+  ASSERT_TRUE(test_model);
+  auto readonly_model = test_model->GetModel();
+  ASSERT_TRUE(readonly_model);
+  ASSERT_TRUE(readonly_model->subgraphs());
+  ASSERT_GE(readonly_model->subgraphs()->size(), 1);
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+  auto subgraph = model.subgraphs[0].get();
+  auto conv_op = subgraph->operators.at(0).get();
+  ASSERT_EQ(
+      GetBuiltinCode(model.operator_codes.at(conv_op->opcode_index).get()),
+      BuiltinOperator_CONV_2D);
+  int32_t weights_tensor_idx = conv_op->inputs[1];
+  TensorT* weights_tensor = subgraph->tensors.at(weights_tensor_idx).get();
+
+  EXPECT_EQ(weights_tensor->type, TensorType_FLOAT32);
+  size_t float_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+
+  EXPECT_EQ(SymmetricQuantizeTensor(&model, weights_tensor), kTfLiteOk);
+
+  size_t quant_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+  EXPECT_EQ(weights_tensor->type, TensorType_INT8);
+  EXPECT_EQ(quant_buffer_size * 4, float_buffer_size);
+}
+
+TEST_F(QuantizationUtilsTest, QuantizeFloat16Clamp) {
+  // Create data.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto tensor = std::make_unique<TensorT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  constexpr int kNumElements = 6;
+  const std::vector<float> weights = {2.0, 1.0, 65504., 65505, -65504., -99999};
+  auto weights_reinterpreted_data =
+      reinterpret_cast<const unsigned char*>(weights.data());
+  buffer->data.assign(weights_reinterpreted_data,
+                      weights_reinterpreted_data + weights.size() * 4);
+  tensor->buffer = 0;
+  tensor->shape = {1, kNumElements};
+
+  // Wire the model.
+  model->subgraphs.push_back(std::move(subgraph));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor));
+  model->buffers.push_back(std::move(buffer));
+
+  // Call and verify.
+  EXPECT_EQ(
+      QuantizeTensorFloat16(model.get(), model->subgraphs[0]->tensors[0].get()),
+      kTfLiteOk);
+  auto weightsf16 = reinterpret_cast<Eigen::half*>(
+      model->buffers[model->subgraphs[0]->tensors[0]->buffer]->data.data());
+  std::vector<float> wf32(kNumElements);
+  std::transform(weightsf16, weightsf16 + 6, wf32.begin(),
+                 [](Eigen::half a) { return static_cast<float>(a); });
+
+  EXPECT_THAT(wf32,
+              ElementsAreArray({2.0, 1.0, 65504., 65504., -65504., -65504.}));
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->type, TensorType_FLOAT16);
+}
+
+TEST_F(QuantizationUtilsTest, QuantizeFloat16) {
+  // Conv model has weights between 0 and 10.
+  // Quantize the weights tensor.
+  ASSERT_TRUE(g_test_model_dir != nullptr);
+  ASSERT_FALSE(g_test_model_dir->empty());
+  auto test_model = ReadConvModel();
+  ASSERT_TRUE(test_model);
+  auto readonly_model = test_model->GetModel();
+  ASSERT_TRUE(readonly_model);
+  ASSERT_TRUE(readonly_model->subgraphs());
+  ASSERT_GE(readonly_model->subgraphs()->size(), 1);
+  tflite::ModelT model;
+  readonly_model->UnPackTo(&model);
+  auto subgraph = model.subgraphs[0].get();
+  auto conv_op = subgraph->operators.at(0).get();
+  ASSERT_EQ(
+      GetBuiltinCode(model.operator_codes.at(conv_op->opcode_index).get()),
+      BuiltinOperator_CONV_2D);
+  int32_t weights_tensor_idx = conv_op->inputs[1];
+  TensorT* weights_tensor = subgraph->tensors.at(weights_tensor_idx).get();
+
+  EXPECT_EQ(weights_tensor->type, TensorType_FLOAT32);
+  size_t float_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+
+  EXPECT_EQ(QuantizeTensorFloat16(&model, weights_tensor), kTfLiteOk);
+
+  size_t quant_buffer_size =
+      model.buffers.at(weights_tensor->buffer)->data.size();
+  EXPECT_EQ(weights_tensor->type, TensorType_FLOAT16);
+  EXPECT_EQ(quant_buffer_size * 2, float_buffer_size);
+}
+
+TEST_F(QuantizationUtilsTest, AddQuantizationParams) {
+  // Create data.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto tensor = std::make_unique<TensorT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  const std::vector<float> scales = {0.5, 1.0, 1.5};
+  const std::vector<int64_t> zero_points = {5, 10, 15};
+  const int32_t quantizated_dimension = 3;
+  const std::vector<uint8_t> buffer_data = {1, 2, 3, 4};
+  const int32_t buffer_size = 4;
+  tensor->buffer = 0;
+
+  // Wire the model.
+  model->subgraphs.push_back(std::move(subgraph));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor));
+  model->buffers.push_back(std::move(buffer));
+
+  // Call and verify.
+  EXPECT_EQ(AddQuantizationParams(
+                scales, zero_points, quantizated_dimension, buffer_data.data(),
+                buffer_size, TensorType_INT8, model.get(),
+                model->subgraphs[0]->tensors[0].get(), &error_reporter_),
+            kTfLiteOk);
+  EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->scale,
+              ElementsAreArray(scales));
+  EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->zero_point,
+              ElementsAreArray(zero_points));
+  EXPECT_THAT(model->buffers[model->subgraphs[0]->tensors[0]->buffer]->data,
+              ElementsAreArray(buffer_data));
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->type, TensorType_INT8);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricQuantizeFloatsToInt16Test) {
+  // Create data.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto tensor = std::make_unique<TensorT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  const float weight_scale = 0.5;
+  const float input_scale = 0.5;
+  std::vector<float> layer_norm_data = {4.0, 1.0, -1.0, 8.0};
+  auto layer_norm_reinterpreted_data =
+      reinterpret_cast<const unsigned char*>(layer_norm_data.data());
+  buffer->data.assign(
+      layer_norm_reinterpreted_data,
+      layer_norm_reinterpreted_data + layer_norm_data.size() * 4);
+  tensor->buffer = 0;
+  tensor->shape = {4};
+  tensor->quantization = std::make_unique<QuantizationParametersT>();
+
+  // Wire the model.
+  model->subgraphs.push_back(std::move(subgraph));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor));
+  model->buffers.push_back(std::move(buffer));
+
+  // Call and verify.
+  EXPECT_EQ(SymmetricQuantizeFloatsToInt16(
+                model.get(), model->subgraphs[0]->tensors[0].get(),
+                input_scale * weight_scale, &error_reporter_),
+            kTfLiteOk);
+
+  EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->scale[0],
+              weight_scale * input_scale);
+  EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->zero_point[0], 0);
+
+  auto result = reinterpret_cast<int16_t*>(
+      model->buffers[model->subgraphs[0]->tensors[0]->buffer]->data.data());
+  EXPECT_EQ(result[0], 16);
+  EXPECT_EQ(result[1], 4);
+  EXPECT_EQ(result[2], -4);
+  EXPECT_EQ(result[3], 32);
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->type, TensorType_INT16);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricPerLayerBiasQuantize) {
+  // Create data.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto tensor = std::make_unique<TensorT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  const float weight_scale = 0.5;
+  const float input_scale = 0.5;
+  std::vector<float> bias_data = {4.0, 1.0};
+  auto bias_reinterpreted_data =
+      reinterpret_cast<const unsigned char*>(bias_data.data());
+  buffer->data.assign(bias_reinterpreted_data,
+                      bias_reinterpreted_data + bias_data.size() * 4);
+  tensor->buffer = 0;
+  tensor->shape = {2, 1, 1, 1};
+  tensor->quantization = std::make_unique<QuantizationParametersT>();
+
+  // Wire the model.
+  model->subgraphs.push_back(std::move(subgraph));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor));
+  model->buffers.push_back(std::move(buffer));
+
+  // Call and verify.
+  EXPECT_EQ(SymmetricPerLayerBiasQuantize<int32_t>(
+                model.get(), model->subgraphs[0]->tensors[0].get(),
+                input_scale * weight_scale, &error_reporter_),
+            kTfLiteOk);
+
+  EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->scale[0],
+              weight_scale * input_scale);
+  EXPECT_THAT(model->subgraphs[0]->tensors[0]->quantization->zero_point[0], 0);
+
+  const uint32_t* d1 = reinterpret_cast<const uint32_t*>(
+      model->buffers[model->subgraphs[0]->tensors[0]->buffer]->data.data());
+  EXPECT_EQ(d1[0], 0x00000010);
+  EXPECT_EQ(d1[1], 0x00000004);
+
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->type, TensorType_INT32);
+}
+
+TEST_F(QuantizationUtilsTest, GetEffectiveScale) {
+  // Create data.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<SubGraphT>();
+  auto tensor = std::make_unique<TensorT>();
+  auto op = std::make_unique<OperatorT>();
+  tensor->quantization = std::make_unique<QuantizationParametersT>();
+  tensor->quantization->scale.push_back(3.0);
+  op->inputs.push_back(0);
+
+  // Wire the model.
+  subgraph->operators.push_back(std::move(op));
+  model->subgraphs.push_back(std::move(subgraph));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor));
+
+  // Call and verify.
+  EXPECT_EQ(GetEffectiveScale(model.get(), model->subgraphs[0].get(), 0, {0},
+                              {}, {5.0}),
+            15.0);
+}
+
+TEST_F(QuantizationUtilsTest, SymmetricPerChannelBiasQuantize) {
+  // Create data.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto tensor = std::make_unique<TensorT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  const std::vector<float> weight_scales = {0.5, 1.0};
+  const float input_scale = 0.5;
+  std::vector<float> bias_data = {4.0, 1.0};
+  auto bias_reinterpreted_data =
+      reinterpret_cast<const unsigned char*>(bias_data.data());
+  buffer->data.assign(bias_reinterpreted_data,
+                      bias_reinterpreted_data + bias_data.size() * 4);
+  tensor->buffer = 0;
+  tensor->shape = {2, 1, 1, 1};
+  tensor->quantization = std::make_unique<QuantizationParametersT>();
+
+  // Wire the model.
+  model->subgraphs.push_back(std::move(subgraph));
+  model->subgraphs[0]->tensors.push_back(std::move(tensor));
+  model->buffers.push_back(std::move(buffer));
+
+  // Call and verify.
+  EXPECT_EQ(SymmetricPerChannelBiasQuantize<int32_t>(
+                model.get(), model->subgraphs[0]->tensors[0].get(), input_scale,
+                weight_scales.data(), 2, &error_reporter_),
+            kTfLiteOk);
+
+  const uint32_t* d1 = reinterpret_cast<const uint32_t*>(
+      model->buffers[model->subgraphs[0]->tensors[0]->buffer]->data.data());
+  EXPECT_EQ(d1[0], 0x00000010);
+  EXPECT_EQ(d1[1], 0x00000002);
+
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->type, TensorType_INT32);
+}
+
+TEST_F(QuantizationUtilsTest, ExtendToPowerOfTwo) {
+  EXPECT_EQ(GetPowerOfTwoScale(-1.0, 1.0), 0);
+  EXPECT_EQ(GetPowerOfTwoScale(-10.0, 10.0), 4);
+  EXPECT_EQ(GetPowerOfTwoScale(3.0, 10.0), 4);
+}
+
+}  // namespace
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite
+
+int main(int argc, char** argv) {
+  std::string model_file;
+  const std::vector<tensorflow::Flag> flag_list = {
+      tensorflow::Flag("test_model_file", &model_file,
+                       "Path to test tflite model file."),
+  };
+
+  const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
+  if (!parse_result) {
+    std::cerr << "Required test_model_file\n";
+    std::abort();
+  }
+  g_test_model_dir = new std::string(tensorflow::io::Dirname(model_file));
+  ::tensorflow::port::InitMain(argv[0], &argc, &argv);
+  return RUN_ALL_TESTS();
+}

tensorflow/lite/tools/optimize/quantization_wrapper.cc

@@ -0,0 +1,56 @@
+/* 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/tools/optimize/quantization_wrapper.h"
+
+#include <string>
+
+#include "tensorflow/lite/tools/optimize/quantization_wrapper_utils.h"
+#include "tensorflow/lite/tools/optimize/quantize_model.h"
+
+namespace tflite {
+namespace optimize {
+
+bool CreateModelForCalibration(const std::string& input_path,
+                               const std::string& output_path) {
+  ModelT model;
+  if (LoadModel(input_path, &model) != kTfLiteOk) {
+    return false;
+  }
+  flatbuffers::FlatBufferBuilder builder;
+  if (AddIntermediateTensorsToFusedOp(&builder, &model) != kTfLiteOk) {
+    return false;
+  }
+  return WriteFile(output_path, builder.GetBufferPointer(), builder.GetSize());
+}
+
+bool CreateQuantizedModel(const std::string& path) {
+  ModelT model;
+  if (LoadModel(path, &model) != kTfLiteOk) {
+    return false;
+  }
+  flatbuffers::FlatBufferBuilder builder;
+  tflite::StderrReporter error_reporter;
+  if (tflite::optimize::QuantizeModel(
+          &builder, &model, tflite::TensorType_FLOAT32,
+          tflite::TensorType_FLOAT32,
+          // TODO(b/159351372): Pass required activation type if needed
+          tflite::TensorType_INT8, &error_reporter) != kTfLiteOk) {
+    return false;
+  }
+  return WriteFile(path, builder.GetBufferPointer(), builder.GetSize());
+}
+
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/quantization_wrapper.h

@@ -0,0 +1,39 @@
+/* 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_TOOLS_OPTIMIZE_QUANTIZATION_WRAPPER_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZATION_WRAPPER_H_
+
+#include <string>
+
+namespace tflite {
+namespace optimize {
+
+// Makes an copy of the model at input_path and writes it to output_path, adding
+// tensors to the model needed for calibration.
+// Returns true if it is successful.
+// Example: a/b/c.tflite becomes a/b/c.calibrated.tflite and has
+// intermediate tensors added according to operator properties.
+bool CreateModelForCalibration(const std::string& input_path,
+                               const std::string& output_path);
+
+// Quantize a model in place. This function is only to be called after calling
+// CreateModelForCalibration and running calibration over data.
+// Returns true if it is successful.
+bool CreateQuantizedModel(const std::string& path);
+
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZATION_WRAPPER_H_

tensorflow/lite/tools/optimize/quantization_wrapper_utils.cc

@@ -0,0 +1,136 @@
+/* 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/tools/optimize/quantization_wrapper_utils.h"
+
+#include <fstream>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "tensorflow/compiler/mlir/lite/tools/optimize/operator_property.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace impl {
+class FlatBufferModel;
+}
+namespace optimize {
+namespace {
+
+#ifdef TFLITE_CUSTOM_LSTM
+constexpr bool kUseCustomLSTM = true;
+#else
+constexpr bool kUseCustomLSTM = false;
+#endif
+
+void MakeTensor(const string& name, std::unique_ptr<TensorT>* tensor) {
+  TensorT* tensor_raw = new TensorT;
+  tensor_raw->name = name;
+  tensor_raw->shape = {0};
+  tensor_raw->type = TensorType_FLOAT32;
+
+  tensor->reset(tensor_raw);
+}
+
+string CreateTensorName(int op_index, int tensor_index) {
+  return "intermediate_" + std::to_string(op_index) + "_" +
+         std::to_string(tensor_index);
+}
+
+bool IntermediateTensorExists(ModelT* model) {
+  for (int subgraph_idx = 0; subgraph_idx < model->subgraphs.size();
+       ++subgraph_idx) {
+    SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get();
+    for (size_t op_idx = 0; op_idx < subgraph->operators.size(); op_idx++) {
+      OperatorT* op = subgraph->operators[op_idx].get();
+      if (!op->intermediates.empty()) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+}  // namespace
+
+TfLiteStatus LoadModel(const string& path, ModelT* model) {
+  auto input_model = impl::FlatBufferModel::BuildFromFile(path.c_str());
+  if (!input_model) {
+    return kTfLiteError;
+  }
+  auto readonly_model = input_model->GetModel();
+  if (!readonly_model) {
+    return kTfLiteError;
+  }
+  readonly_model->UnPackTo(model);
+  return kTfLiteOk;
+}
+
+TfLiteStatus AddIntermediateTensorsToFusedOp(
+    flatbuffers::FlatBufferBuilder* builder, ModelT* model) {
+  // Return early when the model has no operator.
+  if (model->subgraphs.size() == 1 && model->subgraphs[0]->operators.empty()) {
+    return kTfLiteOk;
+  }
+  // Return early if the model already has intermediate tensors.
+  if (IntermediateTensorExists(model)) {
+    return kTfLiteOk;
+  }
+  // Process the model.
+  for (int subgraph_idx = 0; subgraph_idx < model->subgraphs.size();
+       ++subgraph_idx) {
+    SubGraphT* subgraph = model->subgraphs.at(subgraph_idx).get();
+    for (size_t op_idx = 0; op_idx < subgraph->operators.size(); op_idx++) {
+      // Find ops that need additional tensor.
+      OperatorT* op = subgraph->operators[op_idx].get();
+      operator_property::OperatorProperty property =
+          operator_property::GetOperatorProperty(model, subgraph_idx, op_idx);
+      if (property.intermediates.empty()) {
+        continue;
+      }
+      // Add tensors.
+      const int next_tensor_index = subgraph->tensors.size();
+      int num_intermediates = property.intermediates.size();
+      if (kUseCustomLSTM) {
+        num_intermediates = 12;
+      }
+      for (int i = 0; i < num_intermediates; ++i) {
+        std::unique_ptr<TensorT> intermediate_tensor;
+        auto name = CreateTensorName(op_idx, i);
+        MakeTensor(name, &intermediate_tensor);
+        subgraph->tensors.push_back(std::move(intermediate_tensor));
+        op->intermediates.push_back(next_tensor_index + i);
+      }
+    }
+  }
+
+  // Export the model.
+  flatbuffers::Offset<Model> output_model_location =
+      Model::Pack(*builder, model);
+  FinishModelBuffer(*builder, output_model_location);
+
+  return kTfLiteOk;
+}
+
+bool WriteFile(const std::string& out_file, const uint8_t* bytes,
+               size_t num_bytes) {
+  std::fstream stream(out_file, std::ios::binary | std::ios::out);
+  for (size_t i = 0; i < num_bytes; i++) {
+    stream << bytes[i];
+  }
+  return (!stream.bad() && !stream.fail());
+}
+
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/quantization_wrapper_utils.h

@@ -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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZATION_WRAPPER_UTILS_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZATION_WRAPPER_UTILS_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+namespace optimize {
+
+// Load a tflite model from path.
+TfLiteStatus LoadModel(const string& path, ModelT* model);
+
+// Going through the model and add intermediates tensors if the ops have any.
+// Returns early if the model has already intermediate tensors. This is to
+// support cases where a model is initialized multiple times.
+TfLiteStatus AddIntermediateTensorsToFusedOp(
+    flatbuffers::FlatBufferBuilder* builder, ModelT* model);
+
+// Write model to a given location.
+bool WriteFile(const std::string& out_file, const uint8_t* bytes,
+               size_t num_bytes);
+
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZATION_WRAPPER_UTILS_H_

tensorflow/lite/tools/optimize/quantization_wrapper_utils_custom_test.cc

@@ -0,0 +1,140 @@
+/* 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 <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
+#include "tensorflow/lite/tools/optimize/quantization_wrapper_utils.h"
+
+namespace tflite {
+namespace optimize {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+TEST(LstmPreprocess, Add2Tensors) {
+  // Create a model with 1 lstm layer.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  auto lstm_op_code = std::make_unique<OperatorCodeT>();
+  auto lstm_op = std::make_unique<OperatorT>();
+
+  lstm_op_code->builtin_code = BuiltinOperator_LSTM;
+  lstm_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_LSTM);
+  lstm_op_code->version = 2;
+  lstm_op->opcode_index = 0;
+  lstm_op->inputs = {0, 1,  2,  3,  4,  5,  6,  7,  8,  -1, -1, -1,
+                     9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20};
+  lstm_op->outputs = {24};
+
+  model->subgraphs.push_back(std::move(subgraph));
+  for (int i = 0; i < lstm_op->inputs.size(); ++i) {
+    const int index = lstm_op->inputs[i];
+    if (index == -1) {
+      continue;
+    }
+    auto tensor = std::make_unique<TensorT>();
+    tensor->name = "lstm_tensor" + std::to_string(index);
+    tensor->shape = {2, 3, 4};
+    tensor->type = TensorType_FLOAT32;
+    model->subgraphs[0]->tensors.push_back(std::move(tensor));
+  }
+  model->subgraphs[0]->operators.push_back(std::move(lstm_op));
+  model->operator_codes.push_back(std::move(lstm_op_code));
+  model->buffers.push_back(std::move(buffer));
+
+  // Add 2 tensors.
+  flatbuffers::FlatBufferBuilder builder;
+  tflite::optimize::AddIntermediateTensorsToFusedOp(&builder, model.get());
+
+  // Verify results.
+  EXPECT_EQ(model->operator_codes.size(), 1);
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 33);
+  EXPECT_EQ(model->buffers.size(), 1);
+
+  EXPECT_EQ(GetBuiltinCode(model->operator_codes[0].get()),
+            BuiltinOperator_LSTM);
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->name, "lstm_tensor0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[21]->name, "intermediate_0_0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[22]->name, "intermediate_0_1");
+  EXPECT_EQ(model->subgraphs[0]->tensors[23]->name, "intermediate_0_2");
+  EXPECT_EQ(model->subgraphs[0]->tensors[24]->name, "intermediate_0_3");
+  EXPECT_EQ(model->subgraphs[0]->tensors[25]->name, "intermediate_0_4");
+  EXPECT_EQ(model->subgraphs[0]->tensors[26]->name, "intermediate_0_5");
+  EXPECT_EQ(model->subgraphs[0]->tensors[27]->name, "intermediate_0_6");
+  EXPECT_EQ(model->subgraphs[0]->tensors[28]->name, "intermediate_0_7");
+  EXPECT_EQ(model->subgraphs[0]->tensors[29]->name, "intermediate_0_8");
+  EXPECT_EQ(model->subgraphs[0]->tensors[30]->name, "intermediate_0_9");
+  EXPECT_EQ(model->subgraphs[0]->tensors[31]->name, "intermediate_0_10");
+  EXPECT_EQ(model->subgraphs[0]->tensors[32]->name, "intermediate_0_11");
+  EXPECT_THAT(
+      model->subgraphs[0]->operators[0]->inputs,
+      ElementsAreArray({0, 1,  2,  3,  4,  5,  6,  7,  8,  -1, -1, -1,
+                        9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}));
+  EXPECT_THAT(model->subgraphs[0]->operators[0]->outputs,
+              ElementsAreArray({24}));
+  EXPECT_THAT(
+      model->subgraphs[0]->operators[0]->intermediates,
+      ElementsAreArray({21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}));
+
+  // Call AddIntermediateTensorsToFusedOp again and expect no change in model.
+  tflite::optimize::AddIntermediateTensorsToFusedOp(&builder, model.get());
+
+  // Verify results.
+  EXPECT_EQ(model->operator_codes.size(), 1);
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 33);
+  EXPECT_EQ(model->buffers.size(), 1);
+
+  EXPECT_EQ(GetBuiltinCode(model->operator_codes[0].get()),
+            BuiltinOperator_LSTM);
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->name, "lstm_tensor0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[21]->name, "intermediate_0_0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[22]->name, "intermediate_0_1");
+  EXPECT_EQ(model->subgraphs[0]->tensors[23]->name, "intermediate_0_2");
+  EXPECT_EQ(model->subgraphs[0]->tensors[24]->name, "intermediate_0_3");
+  EXPECT_EQ(model->subgraphs[0]->tensors[25]->name, "intermediate_0_4");
+  EXPECT_EQ(model->subgraphs[0]->tensors[26]->name, "intermediate_0_5");
+  EXPECT_EQ(model->subgraphs[0]->tensors[27]->name, "intermediate_0_6");
+  EXPECT_EQ(model->subgraphs[0]->tensors[28]->name, "intermediate_0_7");
+  EXPECT_EQ(model->subgraphs[0]->tensors[29]->name, "intermediate_0_8");
+  EXPECT_EQ(model->subgraphs[0]->tensors[30]->name, "intermediate_0_9");
+  EXPECT_EQ(model->subgraphs[0]->tensors[31]->name, "intermediate_0_10");
+  EXPECT_EQ(model->subgraphs[0]->tensors[32]->name, "intermediate_0_11");
+  EXPECT_THAT(
+      model->subgraphs[0]->operators[0]->inputs,
+      ElementsAreArray({0, 1,  2,  3,  4,  5,  6,  7,  8,  -1, -1, -1,
+                        9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}));
+  EXPECT_THAT(model->subgraphs[0]->operators[0]->outputs,
+              ElementsAreArray({24}));
+  EXPECT_THAT(
+      model->subgraphs[0]->operators[0]->intermediates,
+      ElementsAreArray({21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}));
+}
+
+}  // namespace
+}  // namespace optimize
+}  // namespace tflite

tensorflow/lite/tools/optimize/quantization_wrapper_utils_test.cc

@@ -0,0 +1,127 @@
+/* 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/tools/optimize/quantization_wrapper_utils.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
+
+namespace tflite {
+namespace optimize {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+TEST(LstmPreprocess, Add2Tensors) {
+  // Create a model with 1 lstm layer.
+  auto model = std::make_unique<ModelT>();
+  auto subgraph = std::make_unique<tflite::SubGraphT>();
+  auto buffer = std::make_unique<tflite::BufferT>();
+  auto lstm_op_code = std::make_unique<OperatorCodeT>();
+  auto lstm_op = std::make_unique<OperatorT>();
+
+  lstm_op_code->builtin_code = BuiltinOperator_LSTM;
+  lstm_op_code->deprecated_builtin_code =
+      static_cast<int8_t>(BuiltinOperator_LSTM);
+  lstm_op_code->version = 2;
+  lstm_op->opcode_index = 0;
+  lstm_op->inputs = {0, 1,  2,  3,  4,  5,  6,  7,  8,  -1, -1, -1,
+                     9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20};
+  lstm_op->outputs = {24};
+
+  model->subgraphs.push_back(std::move(subgraph));
+  for (int i = 0; i < lstm_op->inputs.size(); ++i) {
+    const int index = lstm_op->inputs[i];
+    if (index == -1) {
+      continue;
+    }
+    auto tensor = std::make_unique<TensorT>();
+    tensor->name = "lstm_tensor" + std::to_string(index);
+    tensor->shape = {2, 3, 4};
+    tensor->type = TensorType_FLOAT32;
+    model->subgraphs[0]->tensors.push_back(std::move(tensor));
+  }
+  model->subgraphs[0]->operators.push_back(std::move(lstm_op));
+  model->operator_codes.push_back(std::move(lstm_op_code));
+  model->buffers.push_back(std::move(buffer));
+
+  // Add 2 tensors.
+  flatbuffers::FlatBufferBuilder builder;
+  tflite::optimize::AddIntermediateTensorsToFusedOp(&builder, model.get());
+
+  // Verify results.
+  EXPECT_EQ(model->operator_codes.size(), 1);
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 26);
+  EXPECT_EQ(model->buffers.size(), 1);
+
+  EXPECT_EQ(GetBuiltinCode(model->operator_codes[0].get()),
+            BuiltinOperator_LSTM);
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->name, "lstm_tensor0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[21]->name, "intermediate_0_0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[22]->name, "intermediate_0_1");
+  EXPECT_EQ(model->subgraphs[0]->tensors[23]->name, "intermediate_0_2");
+  EXPECT_EQ(model->subgraphs[0]->tensors[24]->name, "intermediate_0_3");
+  EXPECT_EQ(model->subgraphs[0]->tensors[25]->name, "intermediate_0_4");
+  EXPECT_THAT(
+      model->subgraphs[0]->operators[0]->inputs,
+      ElementsAreArray({0, 1,  2,  3,  4,  5,  6,  7,  8,  -1, -1, -1,
+                        9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}));
+  EXPECT_THAT(model->subgraphs[0]->operators[0]->outputs,
+              ElementsAreArray({24}));
+  EXPECT_THAT(model->subgraphs[0]->operators[0]->intermediates,
+              ElementsAreArray({21, 22, 23, 24, 25}));
+
+  // Call AddIntermediateTensorsToFusedOp again and expect no change in model.
+  tflite::optimize::AddIntermediateTensorsToFusedOp(&builder, model.get());
+
+  // Verify results.
+  EXPECT_EQ(model->operator_codes.size(), 1);
+  EXPECT_EQ(model->subgraphs.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->operators.size(), 1);
+  EXPECT_EQ(model->subgraphs[0]->tensors.size(), 26);
+  EXPECT_EQ(model->buffers.size(), 1);
+
+  EXPECT_EQ(GetBuiltinCode(model->operator_codes[0].get()),
+            BuiltinOperator_LSTM);
+  EXPECT_EQ(model->subgraphs[0]->tensors[0]->name, "lstm_tensor0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[21]->name, "intermediate_0_0");
+  EXPECT_EQ(model->subgraphs[0]->tensors[22]->name, "intermediate_0_1");
+  EXPECT_EQ(model->subgraphs[0]->tensors[23]->name, "intermediate_0_2");
+  EXPECT_EQ(model->subgraphs[0]->tensors[24]->name, "intermediate_0_3");
+  EXPECT_EQ(model->subgraphs[0]->tensors[25]->name, "intermediate_0_4");
+  EXPECT_THAT(
+      model->subgraphs[0]->operators[0]->inputs,
+      ElementsAreArray({0, 1,  2,  3,  4,  5,  6,  7,  8,  -1, -1, -1,
+                        9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}));
+  EXPECT_THAT(model->subgraphs[0]->operators[0]->outputs,
+              ElementsAreArray({24}));
+  EXPECT_THAT(model->subgraphs[0]->operators[0]->intermediates,
+              ElementsAreArray({21, 22, 23, 24, 25}));
+}
+
+}  // namespace
+}  // namespace optimize
+}  // namespace tflite
+
+int main(int argc, char** argv) { return RUN_ALL_TESTS(); }

tensorflow/lite/tools/optimize/quantize_model.h

@@ -0,0 +1,139 @@
+/* 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_TOOLS_OPTIMIZE_QUANTIZE_MODEL_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZE_MODEL_H_
+
+#include <memory>
+#include <unordered_set>
+
+#include "tensorflow/lite/context.h"
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace optimize {
+
+// Quantizes input_model and populates the provided builder with the new model.
+// input_model is required to have min/max information populated in its
+// quantization params.
+//
+// Inputs and output types default to float instead of a quantized type.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder,
+                           ModelT* input_model, ErrorReporter* error_reporter);
+
+// Same as above, but the types of quantized inputs and outputs are
+// configurable.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder,
+                           ModelT* input_model, const TensorType& input_type,
+                           const TensorType& output_type,
+                           ErrorReporter* error_reporter);
+
+// Same as above, but can enable allowing float intermediate operations for ops
+// that do not yet support quantizable.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder,
+                           ModelT* input_model, const TensorType& input_type,
+                           const TensorType& output_type, bool allow_float,
+                           ErrorReporter* error_reporter);
+
+// Same as above but with added option of disabling per channel quantization
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModel(
+    flatbuffers::FlatBufferBuilder* builder, ModelT* input_model,
+    const TensorType& input_type, const TensorType& output_type,
+    bool allow_float, bool disable_per_channel,
+    bool disable_per_channel_quantization_for_dense_layers,
+    ErrorReporter* error_reporter);
+
+// Same as above but with added option of handling quantization of external
+// state tensors. This assumes first input and output tensors are ouputs and
+// rest are state tensors which are quantized later with type as
+// activation type (hence no fake quant ops).
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModel(
+    flatbuffers::FlatBufferBuilder* builder, ModelT* input_model,
+    const TensorType& input_type, const TensorType& output_type,
+    bool allow_float, bool disable_per_channel,
+    bool disable_per_channel_quantization_for_dense_layers,
+    ErrorReporter* error_reporter, bool handle_external_state);
+
+// Same as above, but enables only quantizing an allowlist of operations,
+// specified by their operator output name.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder,
+                           ModelT* input_model, const TensorType& input_type,
+                           const TensorType& output_type, bool allow_float,
+                           const std::unordered_set<string>& operator_names,
+                           ErrorReporter* error_reporter);
+
+// Same as above, but enables to provide activation type, which
+// could be TensorType_INT16 or TensorType_INT8.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModel(flatbuffers::FlatBufferBuilder* builder,
+                           ModelT* model, const TensorType& input_type,
+                           const TensorType& output_type, bool allow_float,
+                           const std::unordered_set<string>& operator_names,
+                           const TensorType& activations_type,
+                           const TensorType& bias_type,
+                           ErrorReporter* error_reporter);
+
+// Same as above, but all operators supporting quantization are quantized.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModelAllOperators(
+    flatbuffers::FlatBufferBuilder* builder, ModelT* model,
+    const TensorType& input_type, const TensorType& output_type,
+    bool allow_float, const TensorType& activations_type,
+    const TensorType& bias_type, ErrorReporter* error_reporter);
+
+// Same as above, but allows disabling per channel quantization.
+//
+// Note: This is a private API, subject to change.
+TfLiteStatus QuantizeModelAllOperators(
+    flatbuffers::FlatBufferBuilder* builder, ModelT* model,
+    const TensorType& input_type, const TensorType& output_type,
+    bool allow_float, const TensorType& activations_type,
+    const TensorType& bias_type, bool disable_per_channel,
+    bool disable_per_channel_quantization_for_dense_layers,
+    ErrorReporter* error_reporter);
+
+// Quantizes input_model and populates the provided builder with the new model
+// with all possible input parameters including disabling per_channel
+// quantization.
+//
+// All functions above call this function underneath.
+TfLiteStatus QuantizeModel(
+    flatbuffers::FlatBufferBuilder* builder, ModelT* model,
+    const TensorType& input_type, const TensorType& output_type,
+    bool allow_float, const std::unordered_set<string>& operator_names,
+    const TensorType& activations_type, const TensorType& bias_type,
+    bool disable_per_channel,
+    bool disable_per_channel_quantization_for_dense_layers,
+    ErrorReporter* error_reporter, bool handle_external_state);
+
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_QUANTIZE_MODEL_H_

tensorflow/lite/tools/optimize/reduced_precision_support.h

@@ -0,0 +1,91 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_TOOLS_OPTIMIZE_REDUCED_PRECISION_SUPPORT_H_
+#define TENSORFLOW_LITE_TOOLS_OPTIMIZE_REDUCED_PRECISION_SUPPORT_H_
+
+#include <string>
+
+#include "tensorflow/compiler/mlir/lite/tools/optimize/reduced_precision_metadata.h"
+
+namespace tflite {
+namespace optimize {
+
+inline bool ReadInferenceType(const std::string& metadata, size_t* idx,
+                              ReducedPrecisionSupport* mask) {
+  if (metadata.substr(*idx, 4) == kTfLiteFloat16String) {
+    *idx += 4;
+    *mask = *mask | ReducedPrecisionSupport::Float16Inference;
+    return true;
+  } else if (metadata.substr(*idx, 4) == kTfLiteBfloat16String) {
+    *idx += 4;
+    *mask = *mask | ReducedPrecisionSupport::Bfloat16Inference;
+    return true;
+  }
+  return false;
+}
+
+inline bool ReadAccumulationType(const std::string& metadata, size_t* idx,
+                                 ReducedPrecisionSupport* mask) {
+  if (metadata.substr(*idx, 4) == kTfLiteFloat16String) {
+    *idx += 4;
+    *mask = *mask | ReducedPrecisionSupport::Float16Accumulation;
+    return true;
+  } else if (metadata.substr(*idx, 4) == kTfLiteFloat32String) {
+    *idx += 4;
+    *mask = *mask | ReducedPrecisionSupport::Float32Accumulation;
+    return true;
+  }
+  return false;
+}
+
+// If the string is valid, set the given mask to indicate the state in
+// string and return true. If the string is invalid, return false.
+// A valid string is:
+// >= 1 valid inference types + accumulation token + 1 valid accumulation type.
+// Valid examples would be: "fp16accfp16", "bf16accfp32"
+inline bool SetMaskFromReducedPrecisionMetadata(const std::string& metadata,
+                                                ReducedPrecisionSupport* mask) {
+  bool check = true;
+  size_t idx = 0;
+  ReducedPrecisionSupport rsp = ReducedPrecisionSupport::None;
+  do {
+    check = ReadInferenceType(metadata, &idx, &rsp);
+  } while (check);
+  // Ensure we read at least 1 inference type.
+  if (idx == 0) {
+    return false;
+  }
+  // Next read the accumulation token.
+  if (metadata.substr(idx, 3) != kTfLiteAccumulationString) {
+    return false;
+  }
+  idx += std::string(kTfLiteAccumulationString).size();
+  // Next read a valid accumulation type.
+  if (!ReadAccumulationType(metadata, &idx, &rsp)) {
+    return false;
+  }
+  // This should be the end of string.
+  if (idx != metadata.length()) {
+    return false;
+  }
+  // The string is a valid mask description. Set the value and return.
+  *mask = rsp;
+  return true;
+}
+
+}  // namespace optimize
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_OPTIMIZE_REDUCED_PRECISION_SUPPORT_H_

tensorflow/lite/tools/optimize/reduced_precision_support_test.cc

@@ -0,0 +1,142 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/tools/optimize/reduced_precision_support.h"
+
+#include <cstddef>
+#include <string>
+#include <utility>
+
+#include <gtest/gtest.h>
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/lite/core/model.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/schema/schema_utils.h"
+#include "tensorflow/lite/testing/util.h"
+
+namespace tflite {
+namespace optimize {
+namespace utils {
+namespace {
+
+class ReducedPrecisionSupportTest : public testing::Test {
+ protected:
+  tflite::TestErrorReporter error_reporter_;
+};
+
+TEST_F(ReducedPrecisionSupportTest, BitwiseOps) {
+  ReducedPrecisionSupport mask0 = ReducedPrecisionSupport::None;
+  ReducedPrecisionSupport mask1 = ReducedPrecisionSupport::Float16Inference;
+  ReducedPrecisionSupport bf16 = ReducedPrecisionSupport::Bfloat16Inference;
+  ReducedPrecisionSupport fp16 = ReducedPrecisionSupport::Float16Inference;
+  EXPECT_EQ(mask0, mask0 & mask1);
+  EXPECT_EQ(mask1, mask0 | mask1);
+  mask0 |= fp16;
+  EXPECT_EQ(true, SupportsFP16Inference(mask0));
+  mask0 |= bf16;
+  EXPECT_EQ(true, SupportsBfloat16Inference(mask0));
+  ReducedPrecisionSupport mask2 = ReducedPrecisionSupport::Float16Accumulation;
+  mask2 &= fp16;
+  EXPECT_EQ(mask2, ReducedPrecisionSupport::None);
+}
+
+TEST_F(ReducedPrecisionSupportTest, SupportTests) {
+  ReducedPrecisionSupport bf16 = ReducedPrecisionSupport::Bfloat16Inference;
+  ReducedPrecisionSupport fp16 = ReducedPrecisionSupport::Float16Inference;
+  ReducedPrecisionSupport mask = bf16 | fp16;
+  EXPECT_EQ(true, SupportsFP16Inference(mask));
+  EXPECT_EQ(true, SupportsBfloat16Inference(mask));
+  EXPECT_EQ(false, SupportsFP16Accumulation(mask));
+  EXPECT_EQ(false, SupportsFP32Accumulation(mask));
+  EXPECT_EQ(true, SupportsReducedPrecisionInference(mask));
+  EXPECT_EQ(true, SupportsReducedPrecisionInference(mask));
+  EXPECT_EQ(false, SupportsEitherFP16OrFP32Accumulation(mask));
+  mask = mask | ReducedPrecisionSupport::Float16Accumulation;
+  EXPECT_EQ(true, SupportsFP16Accumulation(mask));
+  EXPECT_EQ(false, SupportsFP32Accumulation(mask));
+  EXPECT_EQ(true, SupportsEitherFP16OrFP32Accumulation(mask));
+}
+
+TEST_F(ReducedPrecisionSupportTest, MetadataStrings) {
+  ReducedPrecisionSupport bf16 = ReducedPrecisionSupport::Bfloat16Inference;
+  ReducedPrecisionSupport fp16 = ReducedPrecisionSupport::Float16Inference;
+  ReducedPrecisionSupport accfp32 =
+      ReducedPrecisionSupport::Float32Accumulation;
+  ReducedPrecisionSupport accfp16 =
+      ReducedPrecisionSupport::Float16Accumulation;
+  ReducedPrecisionSupport maskA = bf16 | fp16 | accfp32;
+  std::pair<std::string, std::string> ans =
+      MetadataForReducedPrecisionSupport(maskA);
+  EXPECT_EQ("fp16bf16accfp32", ans.second);
+  EXPECT_EQ("reduced_precision_support", ans.first);
+  ReducedPrecisionSupport maskB = fp16 | accfp16;
+  EXPECT_EQ("fp16accfp16", MetadataForReducedPrecisionSupport(maskB).second);
+}
+
+TEST_F(ReducedPrecisionSupportTest, ReadStringsIntoMasks) {
+  ReducedPrecisionSupport fp16 = ReducedPrecisionSupport::Float16Inference;
+  ReducedPrecisionSupport accfp16 =
+      ReducedPrecisionSupport::Float16Accumulation;
+  ReducedPrecisionSupport maskfp16 = fp16;
+  ReducedPrecisionSupport maskfp16accfp16 = fp16 | accfp16;
+  ReducedPrecisionSupport mask = ReducedPrecisionSupport::None;
+  size_t idx = 0;
+  std::string metadata = "fp16accfp16";
+  EXPECT_EQ(true, ReadInferenceType(metadata, &idx, &mask));
+  EXPECT_EQ(maskfp16, mask);
+  EXPECT_EQ(idx, 4);
+  idx = 7;
+  EXPECT_EQ(true, ReadAccumulationType(metadata, &idx, &mask));
+  EXPECT_EQ(maskfp16accfp16, mask);
+  EXPECT_EQ(idx, 11);
+}
+
+TEST_F(ReducedPrecisionSupportTest, SetMasks) {
+  ReducedPrecisionSupport fp16 = ReducedPrecisionSupport::Float16Inference;
+  ReducedPrecisionSupport bf16 = ReducedPrecisionSupport::Bfloat16Inference;
+  ReducedPrecisionSupport accfp16 =
+      ReducedPrecisionSupport::Float16Accumulation;
+  ReducedPrecisionSupport accfp32 =
+      ReducedPrecisionSupport::Float32Accumulation;
+  ReducedPrecisionSupport mask = ReducedPrecisionSupport::None;
+  EXPECT_EQ(true, SetMaskFromReducedPrecisionMetadata("bf16accfp32", &mask));
+  EXPECT_EQ(mask, bf16 | accfp32);
+  mask = ReducedPrecisionSupport::None;
+  EXPECT_EQ(true, SetMaskFromReducedPrecisionMetadata("fp16accfp16", &mask));
+  EXPECT_EQ(mask, fp16 | accfp16);
+  mask = ReducedPrecisionSupport::None;
+  EXPECT_EQ(true,
+            SetMaskFromReducedPrecisionMetadata("fp16bf16accfp32", &mask));
+  EXPECT_EQ(mask, fp16 | bf16 | accfp32);
+  mask = ReducedPrecisionSupport::None;
+  EXPECT_EQ(false, SetMaskFromReducedPrecisionMetadata("accfp32", &mask));
+  EXPECT_EQ(mask, ReducedPrecisionSupport::None);
+  EXPECT_EQ(false, SetMaskFromReducedPrecisionMetadata("qwerwer", &mask));
+  EXPECT_EQ(mask, ReducedPrecisionSupport::None);
+  EXPECT_EQ(false,
+            SetMaskFromReducedPrecisionMetadata("fp16accfp32fp16", &mask));
+  EXPECT_EQ(mask, ReducedPrecisionSupport::None);
+  EXPECT_EQ(false, SetMaskFromReducedPrecisionMetadata("fp16accbf16", &mask));
+  EXPECT_EQ(mask, ReducedPrecisionSupport::None);
+}
+
+}  // namespace
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite
+
+int main(int argc, char** argv) {
+  ::tensorflow::port::InitMain(argv[0], &argc, &argv);
+  return RUN_ALL_TESTS();
+}

tensorflow/lite/tools/optimize/sparsity/BUILD

@@ -0,0 +1,50 @@
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
+load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+pybind_extension(
+    name = "format_converter_wrapper_pybind11",
+    srcs = ["format_converter_wrapper_pybind11.cc"],
+    common_lib_packages = [
+        "litert/python",
+        "tensorflow/lite/python",
+    ],
+    copts = [
+        "-fexceptions",
+        "-fno-strict-aliasing",
+    ],
+    data = [
+        "format_converter_wrapper_pybind11.pyi",
+    ],
+    enable_stub_generation = True,
+    features = ["-use_header_modules"],
+    wrap_py_init = True,
+    deps = [
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal/utils:sparsity_format_converter",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+py_test(
+    name = "format_converter_wrapper_pybind11_test",
+    srcs = ["format_converter_wrapper_pybind11_test.py"],
+    strict_deps = True,
+    deps = [
+        ":format_converter_wrapper_pybind11",
+        "@absl_py//absl/testing:absltest",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+    ] + if_pywrap(
+        if_true = ["//tensorflow/lite/python:pywrap_tflite"],
+    ),
+)

tensorflow/lite/tools/optimize/sparsity/format_converter_wrapper_pybind11.cc

@@ -0,0 +1,64 @@
+/* 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 <vector>
+
+#include "pybind11/attr.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/utils/sparsity_format_converter.h"
+
+namespace py = pybind11;
+
+using FormatConverterFp32 = tflite::internal::sparsity::FormatConverter<float>;
+
+PYBIND11_MODULE(format_converter_wrapper_pybind11, m) {
+  m.doc() = "Python wrapper for the tflite sparse tensor converter.";
+
+  py::enum_<TfLiteStatus>(m, "TfLiteStatus")
+      .value("TF_LITE_OK", TfLiteStatus::kTfLiteOk)
+      .value("TF_LITE_ERROR", TfLiteStatus::kTfLiteError)
+      .export_values();
+
+  py::enum_<TfLiteDimensionType>(m, "TfLiteDimensionType")
+      .value("TF_LITE_DIM_DENSE", TfLiteDimensionType::kTfLiteDimDense)
+      .value("TF_LITE_DIM_SPARSE_CSR", TfLiteDimensionType::kTfLiteDimSparseCSR)
+      .export_values();
+
+  py::class_<TfLiteSparsity> sparsity_class(m, "TfLiteSparsity",
+                                            py::module_local());
+
+  py::class_<FormatConverterFp32>(m, "FormatConverterFp32")
+      .def(py::init</*shape=*/const std::vector<int>&,
+                    /*traversal_order=*/const std::vector<int>&,
+                    /*format=*/const std::vector<TfLiteDimensionType>&,
+                    /*block_size=*/const std::vector<int>&,
+                    /*block_map=*/const std::vector<int>&>())
+      .def(py::init</*shape=*/const std::vector<int>&,
+                    /*sparsity=*/const TfLiteSparsity&>())
+      .def("GetData", &FormatConverterFp32::GetData)
+      .def("GetDimMetadata", &FormatConverterFp32::GetDimMetadata)
+      .def("DenseToSparse",
+           [](FormatConverterFp32& converter, py::buffer buf) {
+             py::buffer_info buffer_info = buf.request();
+             return converter.DenseToSparse(
+                 static_cast<float*>(buffer_info.ptr));
+           })
+      .def("SparseToDense", [](FormatConverterFp32& converter, py::buffer buf) {
+        py::buffer_info buffer_info = buf.request();
+        return converter.SparseToDense(static_cast<float*>(buffer_info.ptr));
+      });
+}

tensorflow/lite/tools/optimize/sparsity/format_converter_wrapper_pybind11.pyi

@@ -0,0 +1,66 @@
+# 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 ClassVar, overload
+
+TF_LITE_DIM_DENSE: TfLiteDimensionType
+TF_LITE_DIM_SPARSE_CSR: TfLiteDimensionType
+TF_LITE_ERROR: TfLiteStatus
+TF_LITE_OK: TfLiteStatus
+
+class FormatConverterFp32:
+    @overload
+    def __init__(self, arg0: list[int], arg1: list[int], arg2: list[TfLiteDimensionType], arg3: list[int], arg4: list[int]) -> None: ...
+    @overload
+    def __init__(self, arg0: list[int], arg1: TfLiteSparsity) -> None: ...
+    def DenseToSparse(self, arg0) -> TfLiteStatus: ...
+    def GetData(self) -> list[float]: ...
+    def GetDimMetadata(self) -> list[list[int]]: ...
+    def SparseToDense(self, arg0) -> TfLiteStatus: ...
+
+class TfLiteDimensionType:
+    __members__: ClassVar[dict] = ...  # read-only
+    TF_LITE_DIM_DENSE: ClassVar[TfLiteDimensionType] = ...
+    TF_LITE_DIM_SPARSE_CSR: ClassVar[TfLiteDimensionType] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...
+
+class TfLiteSparsity:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TfLiteStatus:
+    __members__: ClassVar[dict] = ...  # read-only
+    TF_LITE_ERROR: ClassVar[TfLiteStatus] = ...
+    TF_LITE_OK: ClassVar[TfLiteStatus] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...

tensorflow/lite/tools/optimize/sparsity/format_converter_wrapper_pybind11_test.py

@@ -0,0 +1,69 @@
+# 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.
+# =============================================================================
+"""Tests for the pybind11 bindings of format_converter."""
+
+from absl.testing import absltest
+import numpy as np
+
+from tensorflow.lite.tools.optimize.sparsity import format_converter_wrapper_pybind11 as format_converter
+
+
+class FormatConverterTest(absltest.TestCase):
+
+  def test_bcsr_fp32(self):
+    """Same as FormatConverterTest::BlockTestD0S1 but via pybind11."""
+    # pyformat: disable
+    dense_matrix = [1.0, 0.0, 2.0, 3.0,
+                    0.0, 4.0, 0.0, 0.0,
+                    0.0, 0.0, 5.0, 0.0,
+                    0.0, 0.0, 0.0, 6.0]
+    # pyformat: enable
+    dense_shape = [4, 4]
+    traversal_order = [0, 1, 2, 3]
+    dim_types = [
+        format_converter.TfLiteDimensionType.TF_LITE_DIM_DENSE,
+        format_converter.TfLiteDimensionType.TF_LITE_DIM_SPARSE_CSR
+    ]
+    block_size = [2, 2]
+    block_map = [0, 1]
+    converter = format_converter.FormatConverterFp32(dense_shape,
+                                                     traversal_order, dim_types,
+                                                     block_size, block_map)
+
+    converter.DenseToSparse(np.asarray(dense_matrix, dtype=np.float32).data)
+
+    dim_metadata = converter.GetDimMetadata()
+    self.assertEqual([2], dim_metadata[0])
+    self.assertEmpty(dim_metadata[1])  # rows are dense.
+
+    self.assertEqual([0, 2, 3], dim_metadata[2])  # array segments.
+    self.assertEqual([0, 1, 1], dim_metadata[3])  # array indices.
+
+    self.assertEqual([2], dim_metadata[4])
+    self.assertEmpty(dim_metadata[5])  # sub block rows are dense.
+
+    self.assertEqual([2], dim_metadata[6])
+    self.assertEmpty(dim_metadata[7])  # sub block columns are dense.
+
+    expected_data = [1.0, 0.0, 0.0, 4.0, 2.0, 3.0, 0.0, 0.0, 5.0, 0.0, 0.0, 6.0]
+    sparse_data = converter.GetData()
+    self.assertTrue(np.allclose(expected_data, sparse_data))
+
+    converter.SparseToDense(np.asarray(sparse_data, dtype=np.float32).data)
+    self.assertTrue(np.allclose(dense_matrix, converter.GetData()))
+
+
+if __name__ == '__main__':
+  absltest.main()