chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/lite/python/BUILD

@@ -0,0 +1,261 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("@xla//xla/tsl/platform:build_config_root.bzl", "if_pywrap")
+load("//tensorflow:tensorflow.default.bzl", "tf_python_pybind_extension")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+licenses(["notice"])
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        ":friends",
+        "//tensorflow:__pkg__",
+    ],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        "//learning/brain/mobile/lite/tooling/model_analyzer/...",
+        "//learning/brain/mobile/lite/tools/analyzer/...",
+        "//tensorflow/lite/python/...",
+        "//tensorflow/lite/python/converter/...",
+        "//tensorflow/lite/toco/...",
+        "//third_party/odml/litert/litert/...",
+    ],
+)
+
+cc_library(
+    name = "tf_tfl_flatbuffer_helpers",
+    srcs = ["tf_tfl_flatbuffer_helpers.cc"],
+    hdrs = ["tf_tfl_flatbuffer_helpers.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/lite:common",
+        "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:model_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite:tf_to_tfl_flatbuffer",
+        "//tensorflow/compiler/mlir/lite:types_proto_cc",
+        "//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
+        "//tensorflow/compiler/mlir/lite/tools/optimize:reduced_precision_metadata",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "graphdef_to_tfl_flatbuffer",
+    srcs = ["graphdef_to_tfl_flatbuffer.cc"],
+    hdrs = [
+        "graphdef_to_tfl_flatbuffer.h",
+    ],
+    deps = [
+        ":tf_tfl_flatbuffer_helpers",
+        "//tensorflow/compiler/mlir/lite:common",
+        "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:model_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:types_proto_cc",
+        "//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
+        "//tensorflow/compiler/mlir/tensorflow/translate/tools:parsers",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "saved_model_to_tfl_flatbuffer",
+    srcs = ["saved_model_to_tfl_flatbuffer.cc"],
+    hdrs = [
+        "saved_model_to_tfl_flatbuffer.h",
+    ],
+    deps = [
+        ":tf_tfl_flatbuffer_helpers",
+        "//tensorflow/compiler/mlir/lite:common",
+        "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:model_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite:tf_to_tfl_flatbuffer",
+        "//tensorflow/compiler/mlir/lite:types_proto_cc",
+        "//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "jax_to_tfl_flatbuffer",
+    srcs = ["jax_to_tfl_flatbuffer.cc"],
+    hdrs = [
+        "jax_to_tfl_flatbuffer.h",
+    ],
+    deps = [
+        ":tf_tfl_flatbuffer_helpers",
+        "//tensorflow/compiler/mlir/lite:common",
+        "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:model_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite:types_proto_cc",
+        "//tensorflow/compiler/mlir/lite/quantization/common/quantization_lib:quantization_config",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/utility",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/hlo/parser:hlo_parser",
+        "@xla//xla/hlo/translate:stablehlo",
+        "@xla//xla/service:hlo_proto_cc",
+    ],
+)
+
+# Smaller version of flatbuffer_translate which only converts flatbuffer to MLIR.
+cc_library(
+    name = "flatbuffer_to_mlir",
+    srcs = [
+        "flatbuffer_to_mlir.cc",
+    ],
+    hdrs = [
+        "flatbuffer_to_mlir.h",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/lite:flatbuffer_import",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TranslateLib",
+    ],
+)
+
+py_library(
+    name = "wrap_converter",
+    srcs = [
+        "wrap_converter.py",
+    ],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":_pywrap_converter_api",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib_py",
+        "//tensorflow/python:pywrap_tensorflow",
+    ] + if_pywrap(
+        if_true = [] + select({
+            "//conditions:default": [],
+            # This reqirement is for LiteRT only as tensorflow/compiler/mlirt/lite is
+            # duplicated in LiteRT/tflite/converter.
+            "//tensorflow:disable_tf_lite_py": [
+                "//tensorflow/lite/python:pywrap_tflite",
+            ],
+        }),
+    ),
+)
+
+config_setting(
+    name = "tflite_convert_with_select_tf_ops",
+    define_values = {"tflite_convert_with_select_tf_ops": "true"},
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "converter_python_api_hdrs",
+    srcs = [
+        "converter_python_api.h",
+    ],
+    visibility = ["//visibility:private"],
+)
+
+cc_library(
+    name = "converter_python_api",
+    srcs = ["converter_python_api.cc"],
+    hdrs = ["converter_python_api.h"],
+    features = ["-parse_headers"],
+    visibility = [
+        "//visibility:public",
+    ],
+    deps = [
+        ":flatbuffer_to_mlir",
+        ":graphdef_to_tfl_flatbuffer",
+        ":jax_to_tfl_flatbuffer",
+        ":saved_model_to_tfl_flatbuffer",
+        "//tensorflow/c:kernels",
+        "//tensorflow/c:tf_status_headers",
+        "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:model_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:types_proto_cc",
+        "//tensorflow/compiler/mlir/lite/core:absl_error_model_builder",
+        "//tensorflow/compiler/mlir/lite/debug:debug_options_proto_cc",
+        "//tensorflow/compiler/mlir/lite/metrics:error_collector",
+        "//tensorflow/compiler/mlir/lite/python/interpreter_wrapper:python_error_reporter",
+        "//tensorflow/compiler/mlir/lite/python/interpreter_wrapper:python_utils",
+        "//tensorflow/compiler/mlir/lite/quantization/lite:quantize_model",
+        "//tensorflow/compiler/mlir/lite/schema:schema_fbs",
+        "//tensorflow/compiler/mlir/lite/sparsity:sparsify_model",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_protobuf//:protobuf",
+        "@com_google_protobuf//:protobuf_headers",
+        "@flatbuffers//:runtime_cc",
+        "@xla//third_party/python_runtime:headers",  # build_cleaner: keep; DNR: b/35864863
+        "@xla//xla/tsl/platform:status",
+    ] + select({
+        # This is required when running `tflite_convert` from `bazel`.
+        # It requires to link with TensorFlow Ops to get the op definitions.
+        ":tflite_convert_with_select_tf_ops": [
+            "//tensorflow/core:ops",
+        ],
+        "//conditions:default": [],
+    }),
+    alwayslink = True,
+)
+
+tf_python_pybind_extension(
+    name = "_pywrap_converter_api",
+    srcs = [
+        "converter_python_api_wrapper.cc",
+    ],
+    hdrs = [":converter_python_api_hdrs"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "_pywrap_converter_api.pyi",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "//tensorflow/python/lib/core:pybind11_lib",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",
+    ] + if_pywrap([":converter_python_api"]),
+)

tensorflow/compiler/mlir/lite/python/_pywrap_converter_api.pyi

@@ -0,0 +1,21 @@
+# 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 Convert(model_flags_proto_txt_raw: object, converter_flags_proto_txt_raw: object, input_contents_txt_raw: object, extended_return: bool = ..., debug_info_txt_raw: object = ..., quantization_py_function_library=...) -> object: ...
+def ExperimentalMlirQuantizeModel(input_contents_txt_raw: object, disable_per_channel: bool = ..., fully_quantize: bool = ..., inference_type: int = ..., input_data_type: int = ..., output_data_type: int = ..., enable_numeric_verify: bool = ..., enable_whole_model_verify: bool = ..., op_blocklist: object = ..., node_blocklist: object = ..., enable_variable_quantization: bool = ..., disable_per_channel_for_dense_layers: bool = ..., debug_options_proto_txt_raw: object = ...) -> object: ...
+def ExperimentalMlirSparsifyModel(input_contents_txt_raw: object) -> object: ...
+def FlatBufferToMlir(arg0: str, arg1: bool) -> str: ...
+def RegisterCustomOpdefs(custom_opdefs_txt_raw: object) -> object: ...
+def RetrieveCollectedErrors() -> list: ...

tensorflow/compiler/mlir/lite/python/converter_python_api.cc

@@ -0,0 +1,440 @@
+/* Copyright 2017 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/compiler/mlir/lite/python/converter_python_api.h"
+
+#include <Python.h>
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "google/protobuf/text_format.h"
+#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/core/absl_error_model_builder.h"
+#include "tensorflow/compiler/mlir/lite/debug/debug_options.pb.h"
+#include "tensorflow/compiler/mlir/lite/metrics/error_collector.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h"
+#include "tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.h"
+#include "tensorflow/compiler/mlir/lite/python/interpreter_wrapper/python_error_reporter.h"
+#include "tensorflow/compiler/mlir/lite/python/interpreter_wrapper/python_utils.h"
+#include "tensorflow/compiler/mlir/lite/python/jax_to_tfl_flatbuffer.h"
+#include "tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.h"
+#include "tensorflow/compiler/mlir/lite/quantization/lite/quantize_model.h"
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+#include "tensorflow/compiler/mlir/lite/sparsity/sparsify_model.h"
+#include "tensorflow/compiler/mlir/lite/types.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tflite {
+
+// NOTE(aselle): We are using raw PyObject's here because we want to make
+// sure we input and output bytes rather than unicode strings for Python3.
+PyObject* Convert(PyObject* model_flags_proto_txt_raw,
+                  PyObject* converter_flags_proto_txt_raw,
+                  PyObject* input_contents_txt_raw, bool extended_return,
+                  PyObject* debug_info_txt_raw,
+                  const tensorflow::quantization::PyFunctionLibrary*
+                      quantization_py_function_library) {
+  // Use Python C API to validate and convert arguments. In py3 (bytes),
+  // in py2 (str).
+  auto ConvertArg = [&](PyObject* obj, bool* error) {
+    char* buf;
+    Py_ssize_t len;
+    if (mlirlite::python_utils::ConvertFromPyString(obj, &buf, &len) == -1) {
+      *error = true;
+      return std::string();
+    } else {
+      *error = false;
+      return std::string(buf, len);
+    }
+  };
+
+  bool error;
+  std::string model_flags_proto_txt =
+      ConvertArg(model_flags_proto_txt_raw, &error);
+  if (error) {
+    PyErr_SetString(PyExc_ValueError, "Model flags are invalid.");
+    return nullptr;
+  }
+  std::string converter_flags_proto_txt =
+      ConvertArg(converter_flags_proto_txt_raw, &error);
+  if (error) {
+    PyErr_SetString(PyExc_ValueError, "Converter flags are invalid.");
+    return nullptr;
+  }
+
+  // Produce new outputs.
+  tflite::ModelFlags model_flags;
+  if (!model_flags.ParseFromString(model_flags_proto_txt)) {
+    PyErr_SetString(PyExc_ValueError,
+                    "Failed to convert Model to Python String.");
+    return nullptr;
+  }
+  tflite::ConverterFlags converter_flags;
+  if (!converter_flags.ParseFromString(converter_flags_proto_txt)) {
+    PyErr_SetString(PyExc_ValueError,
+                    "Failed to convert ConverterFlags to Python String.");
+    return nullptr;
+  }
+
+  tensorflow::GraphDebugInfo debug_info;
+  if (debug_info_txt_raw && debug_info_txt_raw != Py_None) {
+    std::string debug_info_txt = ConvertArg(debug_info_txt_raw, &error);
+    if (error) {
+      PyErr_SetString(PyExc_ValueError, "Input DebugInfo is invalid.");
+      return nullptr;
+    }
+    if (!debug_info.ParseFromString(debug_info_txt)) {
+      PyErr_SetString(PyExc_ValueError,
+                      "Failed to convert DebugInfo to Python String.");
+      return nullptr;
+    }
+  }
+
+  tensorflow::GraphDef graph_def;
+  std::string input_contents_txt;
+  if (model_flags.saved_model_dir().empty()) {
+    input_contents_txt = ConvertArg(input_contents_txt_raw, &error);
+    if (error) {
+      PyErr_SetString(PyExc_ValueError, "Input GraphDef is invalid.");
+      return nullptr;
+    }
+    if (!model_flags.use_hlo_import() &&
+        !graph_def.ParseFromString(input_contents_txt)) {
+      PyErr_SetString(PyExc_ValueError,
+                      "Failed to convert GraphDef to Python String.");
+      return nullptr;
+    }
+  }
+
+  std::string output_file_contents_txt;
+  absl::Status status;
+
+  // Convert model.
+  if (model_flags.use_hlo_import() && model_flags.has_saved_model_dir()) {
+    PyErr_SetString(PyExc_ValueError,
+                    "Cannot specify both saved_model and hlo import.");
+    return nullptr;
+  }
+
+  if (model_flags.use_hlo_import()) {
+    status = tensorflow::ConvertJaxToTFLiteFlatBuffer(
+        input_contents_txt, model_flags, converter_flags,
+        &output_file_contents_txt);
+  } else if (!model_flags.saved_model_dir().empty()) {
+    status = tensorflow::ConvertSavedModelToTFLiteFlatBuffer(
+        model_flags, converter_flags, &output_file_contents_txt,
+        quantization_py_function_library);
+  } else {
+    tensorflow::GraphDef graph_def;
+    if (!graph_def.ParseFromString(input_contents_txt)) {
+      PyErr_SetString(PyExc_ValueError,
+                      "Failed to convert GraphDef to Python String.");
+      return nullptr;
+    }
+
+    status = tensorflow::ConvertGraphDefToTFLiteFlatBuffer(
+        model_flags, converter_flags, debug_info, graph_def,
+        &output_file_contents_txt);
+  }
+
+  if (!status.ok()) {
+    PyErr_SetString(PyExc_Exception, absl::StatusMessageAsCStr(status));
+    return nullptr;
+  }
+  // Convert arguments back to byte (py3) or str (py2)
+  return mlirlite::python_utils::ConvertToPyString(
+      output_file_contents_txt.data(), output_file_contents_txt.size());
+}
+
+tflite::TensorType FromConverterFlagsToTfLiteDType(int inference_type) {
+  switch (inference_type) {
+    case tflite::IODataType::QUANTIZED_INT16:
+      return tflite::TensorType_INT16;
+    case tflite::IODataType::QUANTIZED_UINT8:
+      return tflite::TensorType_UINT8;
+    case tflite::IODataType::UINT8:
+      return tflite::TensorType_UINT8;
+    case tflite::IODataType::QUANTIZED_INT8:
+      return tflite::TensorType_INT8;
+    case tflite::IODataType::INT8:
+      return tflite::TensorType_INT8;
+    default:
+      return tflite::TensorType_FLOAT32;
+  }
+}
+
+int ToStringSet(PyObject* py_denylist,
+                absl::flat_hash_set<std::string>* string_set) {
+  using mlirlite::python_utils::ConvertFromPyString;
+  // Ensure op_denylist is non null
+  if (!py_denylist) {
+    return 0;
+  }
+  if (PyList_Check(py_denylist)) {
+    for (int i = 0; i < PyList_GET_SIZE(py_denylist); ++i) {
+      PyObject* value = PyList_GetItem(py_denylist, i);
+      char* str_buf;
+      Py_ssize_t length;
+      if (ConvertFromPyString(value, &str_buf, &length) == -1) {
+        return -1;
+      }
+      string_set->emplace(str_buf, length);
+    }
+  }
+  if (PySet_Check(py_denylist)) {
+    auto* tmp = PySet_New(py_denylist);
+    while (PySet_GET_SIZE(tmp)) {
+      PyObject* value = PySet_Pop(tmp);
+      char* str_buf;
+      Py_ssize_t length;
+      if (ConvertFromPyString(value, &str_buf, &length) == -1) {
+        return -1;
+      }
+      string_set->emplace(str_buf, length);
+    }
+  }
+  return 0;
+}
+
+PyObject* MlirQuantizeModel(PyObject* data, bool disable_per_channel,
+                            bool fully_quantize, int inference_type,
+                            int input_data_type, int output_data_type,
+                            bool enable_numeric_verify,
+                            bool enable_whole_model_verify,
+                            PyObject* op_denylist, PyObject* node_denylist,
+                            bool enable_variable_quantization,
+                            bool disable_per_channel_for_dense_layers,
+                            PyObject* debug_options_proto_txt_raw) {
+  using tflite_migration::interpreter_wrapper::PythonErrorReporter;
+  char* buf = nullptr;
+  Py_ssize_t length;
+  std::unique_ptr<PythonErrorReporter> error_reporter =
+      std::make_unique<PythonErrorReporter>();
+
+  if (mlirlite::python_utils::ConvertFromPyString(data, &buf, &length) == -1) {
+    PyErr_Format(PyExc_ValueError, "Failed to convert input PyObject");
+    return nullptr;
+  }
+
+  std::optional<tensorflow::converter::DebugOptions> debug_options =
+      tensorflow::converter::DebugOptions();
+  if (debug_options_proto_txt_raw != nullptr) {
+    auto ConvertArg = [&](PyObject* obj, bool* error) {
+      char* buf;
+      Py_ssize_t len;
+      if (mlirlite::python_utils::ConvertFromPyString(obj, &buf, &len) == -1) {
+        *error = true;
+        return std::string();
+      } else {
+        *error = false;
+        return std::string(buf, len);
+      }
+    };
+
+    bool error;
+    std::string debug_options_proto_txt =
+        ConvertArg(debug_options_proto_txt_raw, &error);
+    if (error) {
+      PyErr_SetString(PyExc_ValueError, "Converter flags are invalid.");
+      return nullptr;
+    }
+
+    if (!debug_options->ParseFromString(debug_options_proto_txt)) {
+      PyErr_SetString(PyExc_ValueError,
+                      "Failed to convert ConverterFlags to Python String.");
+      return nullptr;
+    }
+  } else {
+    debug_options = std::nullopt;
+  }
+
+  absl::flat_hash_set<std::string> denylisted_ops;
+  absl::flat_hash_set<std::string> denylisted_nodes;
+  if (ToStringSet(op_denylist, &denylisted_ops) == -1) {
+    PyErr_Format(PyExc_ValueError, "Failed to convert op denylist PyObject");
+    return nullptr;
+  }
+  if (ToStringSet(node_denylist, &denylisted_nodes) == -1) {
+    PyErr_Format(PyExc_ValueError, "Failed to convert node denylist PyObject");
+    return nullptr;
+  }
+
+  std::unique_ptr<mlir::TFL::FlatBufferModelAbslError> model =
+      mlir::TFL::FlatBufferModelAbslError::BuildFromBuffer(
+          buf, length, error_reporter.get());
+  if (!model) {
+    PyErr_Format(PyExc_ValueError, "Invalid model");
+    return nullptr;
+  }
+  auto tflite_model = std::make_unique<tflite::ModelT>();
+  model->GetModel()->UnPackTo(tflite_model.get(), nullptr);
+
+  const tflite::TensorType inference_tensor_type =
+      FromConverterFlagsToTfLiteDType(inference_type);
+  const tflite::TensorType input_type =
+      FromConverterFlagsToTfLiteDType(input_data_type);
+  const tflite::TensorType output_type =
+      FromConverterFlagsToTfLiteDType(output_data_type);
+
+  std::string output_model;
+  const absl::string_view input_model_buffer(buf, length);
+  auto status = mlir::lite::QuantizeModel(
+      input_model_buffer, input_type, output_type, inference_tensor_type,
+      /*operator_names=*/{}, disable_per_channel, fully_quantize, output_model,
+      enable_numeric_verify, enable_whole_model_verify,
+      /*legacy_float_scale=*/true, denylisted_ops, denylisted_nodes,
+      enable_variable_quantization, disable_per_channel_for_dense_layers,
+      debug_options);
+  if (!status.ok()) {
+    LOG(ERROR) << "Failed to quantize model: " << status;
+    error_reporter->exception();
+    return nullptr;
+  }
+
+  return mlirlite::python_utils::ConvertToPyString(output_model.data(),
+                                                   output_model.size());
+}
+
+PyObject* MlirSparsifyModel(PyObject* data) {
+  using tflite_migration::interpreter_wrapper::PythonErrorReporter;
+  char* buf = nullptr;
+  Py_ssize_t length;
+  std::unique_ptr<PythonErrorReporter> error_reporter =
+      std::make_unique<PythonErrorReporter>();
+
+  if (mlirlite::python_utils::ConvertFromPyString(data, &buf, &length) == -1) {
+    PyErr_Format(PyExc_ValueError, "Failed to convert input PyObject");
+    return nullptr;
+  }
+  std::unique_ptr<mlir::TFL::FlatBufferModelAbslError> model =
+      mlir::TFL::FlatBufferModelAbslError::BuildFromBuffer(
+          buf, length, error_reporter.get());
+  if (!model) {
+    PyErr_Format(PyExc_ValueError, "Invalid model");
+    return nullptr;
+  }
+  auto tflite_model = std::make_unique<tflite::ModelT>();
+  model->GetModel()->UnPackTo(tflite_model.get(), nullptr);
+
+  flatbuffers::FlatBufferBuilder builder;
+  auto status = mlir::lite::SparsifyModel(*tflite_model, &builder);
+
+  if (!status.ok()) {
+    error_reporter->exception();
+    return nullptr;
+  }
+  return mlirlite::python_utils::ConvertToPyString(
+      reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
+      builder.GetSize());
+}
+
+PyObject* RegisterCustomOpdefs(PyObject* list) {
+  if (!PyList_Check(list)) {
+    PyErr_SetString(PyExc_TypeError, "Expected list in argument");
+    return nullptr;
+  }
+
+  int64_t size = PyList_Size(list);
+  for (int i = 0; i < size; ++i) {
+    // Get character array from Python object.
+    char* tf_opdefs;
+    Py_ssize_t len;
+    if (mlirlite::python_utils::ConvertFromPyString(PyList_GetItem(list, i),
+                                                    &tf_opdefs, &len) == -1) {
+      PyErr_Format(PyExc_ValueError,
+                   "Failed to convert Python string at index %d of custom op "
+                   "defs argument",
+                   i);
+      return nullptr;
+    }
+
+    // Parse op def from character array.
+    tensorflow::OpDef opdef;
+    if (!tensorflow::protobuf::TextFormat::ParseFromString(tf_opdefs, &opdef)) {
+      PyErr_Format(
+          PyExc_ValueError,
+          "Failed to parse opdefs at index %d of custom op defs argument: %s",
+          i, tf_opdefs);
+      return nullptr;
+    }
+
+    // Register extra opdefs to TensorFlow global op registry.
+    tensorflow::OpRegistry::Global()->Register(
+        [opdef](tensorflow::OpRegistrationData* op_reg_data) -> absl::Status {
+          *op_reg_data = tensorflow::OpRegistrationData(opdef);
+          return absl::OkStatus();
+        });
+
+    // Register the corresponding fake op kernel.
+    const char* node_name = opdef.name().c_str();
+    const char* op_name = opdef.name().c_str();
+    const char* device_name = "CPU";
+    static auto fake_compute_func = [](void* kernel, TF_OpKernelContext* ctx) {
+    };
+
+    TF_KernelBuilder* builder =
+        TF_NewKernelBuilder(op_name, device_name, /*create_func=*/nullptr,
+                            fake_compute_func, /*delete_func=*/nullptr);
+
+    TF_Status* status = TF_NewStatus();
+    TF_RegisterKernelBuilder(node_name, builder, status);
+    if (TF_GetCode(status) != TF_OK) {
+      TF_DeleteStatus(status);
+      PyErr_Format(PyExc_ValueError,
+                   "Failed to register fake op kernel at index %d of custom op "
+                   "defs argument",
+                   i);
+      return nullptr;
+    }
+    TF_DeleteStatus(status);
+  }
+
+  Py_RETURN_TRUE;
+}
+
+std::vector<std::string> RetrieveCollectedErrors() {
+  mlir::TFL::ErrorCollector* collector =
+      mlir::TFL::ErrorCollector::GetErrorCollector();
+  std::vector<std::string> collected_errors;
+  for (const auto& error_data : collector->CollectedErrors()) {
+    collected_errors.push_back(error_data.SerializeAsString());
+  }
+  collector->Clear();
+  return collected_errors;
+}
+
+std::string FlatBufferFileToMlir(const std::string& model,
+                                 bool input_is_filepath) {
+  return ::tensorflow::FlatBufferFileToMlir(model, input_is_filepath);
+}
+
+}  // namespace tflite

tensorflow/compiler/mlir/lite/python/converter_python_api.h

@@ -0,0 +1,74 @@
+/* Copyright 2017 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_COMPILER_MLIR_LITE_PYTHON_CONVERTER_PYTHON_API_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_CONVERTER_PYTHON_API_H_
+
+#include <Python.h>
+
+#include <string>
+#include <vector>
+
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+
+namespace tflite {
+
+// Convert a model represented in `input_contents`. `model_flags_proto`
+// describes model parameters. `flags_proto` describes conversion
+// parameters (see relevant .protos for more information). Returns a string
+// representing the contents of the converted model. When extended_return
+// flag is set to true returns a dictionary that contains string representation
+// of the converted model and some statistics like arithmetic ops count.
+// `debug_info_str` contains the `GraphDebugInfo` proto.
+PyObject* Convert(PyObject* model_flags_proto_txt_raw,
+                  PyObject* converter_flags_proto_txt_raw,
+                  PyObject* input_contents_txt_raw,
+                  bool extended_return = false,
+                  PyObject* debug_info_txt_raw = nullptr,
+                  const tensorflow::quantization::PyFunctionLibrary*
+                      quantization_py_function_library = nullptr);
+
+// Quantize the model with calibration data. Throw errors if `fully_quantize`
+// is specified by the calibration data are not sufficient to quantize the
+// model.
+PyObject* MlirQuantizeModel(PyObject* data, bool disable_per_channel,
+                            bool fully_quantize, int inference_type,
+                            int input_data_type, int output_data_type,
+                            bool enable_numeric_verify = false,
+                            bool enable_whole_model_verify = false,
+                            PyObject* op_denylist = nullptr,
+                            PyObject* node_denylist = nullptr,
+                            bool enable_variable_quantization = false,
+                            bool disable_per_channel_for_dense_layers = false,
+                            PyObject* debug_options_proto_txt_raw = nullptr);
+
+// Sparsifies model to encode sparse tensors with proper format. Throws error if
+// sparsification fails.
+PyObject* MlirSparsifyModel(PyObject* data);
+
+// Registers the given custom opdefs to TensorFlow global op registry.
+PyObject* RegisterCustomOpdefs(PyObject* list);
+
+// Returns the collected TFLite conversion errors.
+std::vector<std::string> RetrieveCollectedErrors();
+
+// Returns MLIR string dump of the given Flatbuffer model.
+std::string FlatBufferFileToMlir(const std::string& model,
+                                 bool input_is_filepath);
+
+// All the exported functions should be listed in
+// tensorflow/tools/def_file_filter/symbols_pybind.txt for the Windows build.
+}  // namespace tflite
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_CONVERTER_PYTHON_API_H_

tensorflow/compiler/mlir/lite/python/converter_python_api_wrapper.cc

@@ -0,0 +1,128 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <string>
+#include <vector>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/compiler/mlir/lite/python/converter_python_api.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_converter_api, m) {
+  m.def(
+      "Convert",
+      [](py::object model_flags_proto_txt_raw,
+         py::object converter_flags_proto_txt_raw,
+         py::object input_contents_txt_raw, bool extended_return,
+         py::object debug_info_txt_raw,
+         const tensorflow::quantization::PyFunctionLibrary*
+             quantization_py_function_library) {
+        return tensorflow::PyoOrThrow(tflite::Convert(
+            model_flags_proto_txt_raw.ptr(),
+            converter_flags_proto_txt_raw.ptr(), input_contents_txt_raw.ptr(),
+            extended_return, debug_info_txt_raw.ptr(),
+            quantization_py_function_library));
+      },
+      py::arg("model_flags_proto_txt_raw"),
+      py::arg("converter_flags_proto_txt_raw"),
+      py::arg("input_contents_txt_raw"), py::arg("extended_return") = false,
+      py::arg("debug_info_txt_raw") = py::none(),
+      py::arg("quantization_py_function_library") = py::none(),
+      R"pbdoc(
+      Convert a model represented in `input_contents`. `model_flags_proto`
+      describes model parameters. `flags_proto` describes conversion
+      parameters (see relevant .protos for more information). Returns a string
+      representing the contents of the converted model. When extended_return
+      flag is set to true returns a dictionary that contains string representation
+      of the converted model and some statistics like arithmetic ops count.
+      `debug_info_str` contains the `GraphDebugInfo` proto.
+    )pbdoc");
+  m.def(
+      "ExperimentalMlirQuantizeModel",
+      [](py::object input_contents_txt_raw, bool disable_per_channel,
+         bool fully_quantize, int inference_type, int input_data_type,
+         int output_data_type, bool enable_numeric_verify,
+         bool enable_whole_model_verify, py::object op_blocklist,
+         py::object node_blocklist, bool enable_variable_quantization,
+         bool disable_per_channel_for_dense_layers,
+         py::object debug_options_proto_txt_raw) {
+        return tensorflow::PyoOrThrow(tflite::MlirQuantizeModel(
+            input_contents_txt_raw.ptr(), disable_per_channel, fully_quantize,
+            inference_type, input_data_type, output_data_type,
+            enable_numeric_verify, enable_whole_model_verify,
+            op_blocklist.ptr(), node_blocklist.ptr(),
+            enable_variable_quantization, disable_per_channel_for_dense_layers,
+            debug_options_proto_txt_raw.ptr()));
+      },
+      py::arg("input_contents_txt_raw"), py::arg("disable_per_channel") = false,
+      py::arg("fully_quantize") = true, py::arg("inference_type") = 9,
+      py::arg("input_data_type") = 0, py::arg("output_data_type") = 0,
+      py::arg("enable_numeric_verify") = false,
+      py::arg("enable_whole_model_verify") = false,
+      py::arg("op_blocklist") = py::none(),
+      py::arg("node_blocklist") = py::none(),
+      py::arg("enable_variable_quantization") = false,
+      py::arg("disable_per_channel_for_dense_layers") = false,
+      py::arg("debug_options_proto_txt_raw") = nullptr,
+      R"pbdoc(
+      Returns a quantized model.
+    )pbdoc");
+  m.def(
+      "ExperimentalMlirSparsifyModel",
+      [](py::object input_contents_txt_raw) {
+        return tensorflow::PyoOrThrow(
+            tflite::MlirSparsifyModel(input_contents_txt_raw.ptr()));
+      },
+      py::arg("input_contents_txt_raw"),
+      R"pbdoc(
+      Returns a sparsified model.
+    )pbdoc");
+  m.def(
+      "RegisterCustomOpdefs",
+      [](py::object custom_opdefs_txt_raw) {
+        return tensorflow::PyoOrThrow(
+            tflite::RegisterCustomOpdefs(custom_opdefs_txt_raw.ptr()));
+      },
+      py::arg("custom_opdefs_txt_raw"),
+      R"pbdoc(
+      Registers the given custom opdefs to the TensorFlow global op registry.
+    )pbdoc");
+  m.def(
+      "RetrieveCollectedErrors",
+      []() {
+        std::vector<std::string> collected_errors =
+            tflite::RetrieveCollectedErrors();
+        pybind11::list serialized_message_list(collected_errors.size());
+        int i = 0;
+        for (const auto& error_data : collected_errors) {
+          serialized_message_list[i++] = pybind11::bytes(error_data);
+        }
+        return serialized_message_list;
+      },
+      R"pbdoc(
+      Returns and clears the list of collected errors in ErrorCollector.
+    )pbdoc");
+  m.def(
+      "FlatBufferToMlir",
+      [](const std::string& model, bool input_is_filepath) {
+        return tflite::FlatBufferFileToMlir(model, input_is_filepath);
+      },
+      R"pbdoc(
+      Returns MLIR dump of the given TFLite model.
+    )pbdoc");
+}

tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.cc

@@ -0,0 +1,94 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/string_view.h"
+#include "llvm/Support/LogicalResult.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SMLoc.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Verifier.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Tools/mlir-translate/Translation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/flatbuffer_import.h"
+
+namespace tensorflow {
+namespace {
+static mlir::OwningOpRef<mlir::ModuleOp> FlatBufferFileToMlirTranslation(
+    llvm::SourceMgr* source_mgr, mlir::MLIRContext* context) {
+  const llvm::MemoryBuffer* input =
+      source_mgr->getMemoryBuffer(source_mgr->getMainFileID());
+  std::string error;
+  auto loc =
+      mlir::FileLineColLoc::get(context, input->getBufferIdentifier(), 0, 0);
+  std::vector<std::string> inputs;
+  std::vector<std::string> outputs;
+  return tflite::FlatBufferToMlir(
+      absl::string_view(input->getBufferStart(), input->getBufferSize()),
+      context, loc, false, inputs, outputs, false);
+}
+
+}  // namespace
+
+std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer,
+                                 bool input_is_filepath) {
+  // referred logic from mlir::mlirTranslateMain().
+
+  std::string errorMessage;
+  std::unique_ptr<llvm::MemoryBuffer> input;
+  if (input_is_filepath) {
+    input = mlir::openInputFile(model_file_or_buffer, &errorMessage);
+    if (!input) {
+      llvm::errs() << errorMessage << "\n";
+      return "";
+    }
+  } else {
+    input = llvm::MemoryBuffer::getMemBuffer(model_file_or_buffer, "flatbuffer",
+                                             false);
+    if (!input) {
+      llvm::errs() << "Can't get llvm::MemoryBuffer\n";
+      return "";
+    }
+  }
+
+  mlir::MLIRContext context;
+  context.printOpOnDiagnostic(true);
+  llvm::SourceMgr sourceMgr;
+  sourceMgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc());
+
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      FlatBufferFileToMlirTranslation(&sourceMgr, &context);
+  if (!module || failed(verify(*module))) return "";
+
+  std::string mlir_output;
+  llvm::raw_string_ostream output_stream(mlir_output);
+  // Dump MLIR with eliding large elements.
+  module->print(
+      output_stream,
+      mlir::OpPrintingFlags().useLocalScope().elideLargeElementsAttrs());
+  return mlir_output;
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/lite/python/flatbuffer_to_mlir.h

@@ -0,0 +1,30 @@
+/* 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_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_
+
+#include <string>
+
+namespace tensorflow {
+
+// Translates the given FlatBuffer filename or buffer into MLIR and returns
+// translated MLIR as string.
+std::string FlatBufferFileToMlir(const std::string& model_file_or_buffer,
+                                 bool input_is_filepath);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_FLATBUFFER_TO_MLIR_H_

tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.cc

@@ -0,0 +1,133 @@
+/* Copyright 2022 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/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h"
+#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
+#include "tensorflow/compiler/mlir/lite/types.pb.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/tools/parsers.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/graph_to_tf_executor.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/common_runtime/graph_constructor.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/graph.h"
+
+namespace tensorflow {
+
+absl::Status ConvertGraphDefToTFLiteFlatBuffer(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags, const GraphDebugInfo& debug_info,
+    const GraphDef& input, std::string* result) {
+  auto context = std::make_unique<mlir::MLIRContext>();
+  GraphImportConfig specs;
+  mlir::TFL::QuantizationSpecs quant_specs;
+
+  // Parse input arrays.
+  std::vector<std::string> node_names;
+  std::vector<std::string> node_dtypes;
+  std::vector<std::optional<std::vector<int>>> node_shapes;
+  std::vector<std::optional<double>> node_mins;
+  std::vector<std::optional<double>> node_maxs;
+
+  // Populate quantization specs.
+  TF_RETURN_IF_ERROR(internal::PopulateQuantizationSpecs(
+      model_flags, converter_flags, &quant_specs, &node_names, &node_dtypes,
+      &node_shapes, &node_mins, &node_maxs));
+
+  TF_RETURN_IF_ERROR(
+      ParseInputArrayInfo(node_names, node_dtypes, node_shapes, &specs.inputs));
+
+  // Parse output arrays.
+  std::vector<std::string> output_arrays(model_flags.output_arrays().begin(),
+                                         model_flags.output_arrays().end());
+  TF_RETURN_IF_ERROR(ParseOutputArrayInfo(output_arrays, &specs.outputs));
+
+  // Parse control output arrays.
+  std::vector<std::string> control_output_arrays(
+      model_flags.control_output_arrays().begin(),
+      model_flags.control_output_arrays().end());
+  TF_RETURN_IF_ERROR(
+      ParseOutputArrayInfo(control_output_arrays, &specs.control_outputs));
+
+  specs.prune_unused_nodes = true;
+  specs.convert_legacy_fed_inputs = true;
+  specs.graph_as_function = false;
+  specs.upgrade_legacy = true;
+  specs.unconditionally_use_set_output_shapes = true;
+  internal::WarningUnusedFlags(model_flags, converter_flags);
+
+  // Register all custom ops, including user-specified custom ops.
+  TF_RETURN_IF_ERROR(internal::RegisterAllCustomOps(converter_flags));
+  GraphConstructorOptions options;
+  options.allow_internal_ops = true;
+  options.upgrade_legacy = specs.upgrade_legacy;
+  Graph graph(OpRegistry::Global());
+  TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(options, input, &graph));
+  TF_ASSIGN_OR_RETURN(
+      auto module,
+      tensorflow::tf2xla::v2::ConvertGraphToTfExecutor(
+          graph, debug_info, graph.flib_def(), specs, context.get()));
+
+  mlir::TFL::PassConfig pass_config(quant_specs);
+  bool emit_builtin_tflite_ops = !converter_flags.force_select_tf_ops();
+  pass_config.emit_builtin_tflite_ops = emit_builtin_tflite_ops;
+  pass_config.unfold_batch_matmul = converter_flags.unfold_batchmatmul();
+  pass_config.lower_tensor_list_ops = converter_flags.lower_tensor_list_ops();
+  pass_config.legalize_custom_tensor_list_ops =
+      converter_flags.legalize_custom_tensor_list_ops();
+  // Disable the unfolding of the 16x16 TF::BatchMatMulOp to avoid the
+  // conversion to an unsupported 16x16 TFL::FullyConnectedOp.
+  if (converter_flags.inference_type() == tflite::IODataType::QUANTIZED_INT16) {
+    pass_config.unfold_batch_matmul = false;
+  }
+  pass_config.unfold_large_splat_constant =
+      converter_flags.unfold_large_splat_constant();
+  pass_config.enable_dynamic_update_slice =
+      converter_flags.enable_dynamic_update_slice();
+  pass_config.preserve_assert_op = converter_flags.preserve_assert_op();
+  pass_config.guarantee_all_funcs_one_use =
+      converter_flags.guarantee_all_funcs_one_use();
+  pass_config.enable_stablehlo_conversion =
+      converter_flags.convert_to_stablehlo();
+  pass_config.canonicalizing_inf_as_min_max_float =
+      converter_flags.canonicalizing_inf_as_min_max_float();
+
+  // StableHLO Quantizer is not supported for GraphDef inputs, so
+  // quantization_py_function_lib is set to nullptr.
+  return internal::ConvertMLIRToTFLiteFlatBuffer(
+      model_flags, converter_flags, std::move(context), std::move(module),
+      pass_config,
+      /*saved_model_tags=*/{}, result,
+      /*quantization_py_function_lib=*/nullptr);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/lite/python/graphdef_to_tfl_flatbuffer.h

@@ -0,0 +1,38 @@
+/* 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_COMPILER_MLIR_LITE_PYTHON_GRAPHDEF_TO_TFL_FLATBUFFER_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_GRAPHDEF_TO_TFL_FLATBUFFER_H_
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+
+namespace tensorflow {
+
+// Converts the given GraphDef to a TF Lite FlatBuffer string according to the
+// given model flags, converter flags and debug information. Returns error
+// status if it fails to convert the input.
+absl::Status ConvertGraphDefToTFLiteFlatBuffer(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags, const GraphDebugInfo& debug_info,
+    const GraphDef& input, std::string* result);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_GRAPHDEF_TO_TFL_FLATBUFFER_H_

tensorflow/compiler/mlir/lite/python/interpreter_wrapper/BUILD

@@ -0,0 +1,29 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "python_utils",
+    srcs = ["python_utils.cc"],
+    hdrs = ["python_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+cc_library(
+    name = "python_error_reporter",
+    srcs = ["python_error_reporter.cc"],
+    hdrs = ["python_error_reporter.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/lite:stateful_error_reporter",
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)

tensorflow/compiler/mlir/lite/python/interpreter_wrapper/python_error_reporter.cc

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

tensorflow/compiler/mlir/lite/python/interpreter_wrapper/python_error_reporter.h

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

tensorflow/compiler/mlir/lite/python/interpreter_wrapper/python_utils.cc

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

tensorflow/compiler/mlir/lite/python/interpreter_wrapper/python_utils.h

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

tensorflow/compiler/mlir/lite/python/jax_to_tfl_flatbuffer.cc

@@ -0,0 +1,169 @@
+/* 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/compiler/mlir/lite/python/jax_to_tfl_flatbuffer.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h"
+#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/lite/types.pb.h"
+#include "xla/hlo/parser/hlo_parser.h"
+#include "xla/hlo/translate/stablehlo.h"
+#include "xla/service/hlo.pb.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tsl/platform/protobuf.h"  // IWYU pragma: keep
+
+namespace tensorflow {
+namespace {
+
+// Error collector that simply ignores errors reported.
+class NoOpErrorCollector : public tsl::protobuf::io::ErrorCollector {
+ public:
+  void RecordError(int line, tsl::protobuf::io::ColumnNumber column,
+                   absl::string_view message) override {}
+};
+
+absl::StatusOr<xla::HloProto> LoadHloProto(const std::string& contents) {
+  xla::HloProto hlo_proto;
+  // NOLINTNEXTLINE: Use tsl::protobuf to be compatible with OSS.
+  tsl::protobuf::TextFormat::Parser parser;
+  NoOpErrorCollector collector;
+  parser.RecordErrorsTo(&collector);
+  bool status =
+      hlo_proto.ParseFromString(contents) ||
+      parser.ParseFromString(contents, &hlo_proto) ||
+      hlo_proto.mutable_hlo_module()->ParseFromString(contents) ||
+      parser.ParseFromString(contents, hlo_proto.mutable_hlo_module());
+  if (!status) {
+    return absl::InternalError("Failed to parse HloProto");
+  }
+  return hlo_proto;
+}
+
+}  // namespace
+
+absl::Status ConvertJaxToTFLiteFlatBuffer(
+    const std::string& input, const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags, std::string* result) {
+  auto context = std::make_unique<mlir::MLIRContext>();
+  mlir::TFL::QuantizationSpecs quant_specs;
+
+  // Parse input arrays.
+  std::vector<std::string> node_names;
+  std::vector<std::string> node_dtypes;
+  std::vector<std::optional<std::vector<int>>> node_shapes;
+  std::vector<std::optional<double>> node_mins;
+  std::vector<std::optional<double>> node_maxs;
+
+  // Populate quantization specs.
+  TF_RETURN_IF_ERROR(internal::PopulateQuantizationSpecs(
+      model_flags, converter_flags, &quant_specs, &node_names, &node_dtypes,
+      &node_shapes, &node_mins, &node_maxs));
+
+  internal::WarningUnusedFlags(model_flags, converter_flags);
+
+  // Register all custom ops, including user-specified custom ops.
+  TF_RETURN_IF_ERROR(internal::RegisterAllCustomOps(converter_flags));
+
+  mlir::TFL::PassConfig pass_config(quant_specs);
+  bool emit_builtin_tflite_ops = !converter_flags.force_select_tf_ops();
+  pass_config.emit_builtin_tflite_ops = emit_builtin_tflite_ops;
+  pass_config.enable_tflite_variables =
+      converter_flags.enable_tflite_resource_variables();
+  pass_config.unfold_batch_matmul = converter_flags.unfold_batchmatmul();
+  pass_config.lower_tensor_list_ops = converter_flags.lower_tensor_list_ops();
+  // Disable the unfolding of the 16x16 TF::BatchMatMulOp to avoid the
+  // conversion to an unsupported 16x16 TFL::FullyConnectedOp.
+  if (converter_flags.inference_type() == tflite::IODataType::QUANTIZED_INT16) {
+    pass_config.unfold_batch_matmul = false;
+  }
+  pass_config.unfold_large_splat_constant =
+      converter_flags.unfold_large_splat_constant();
+  pass_config.enable_hlo_to_tf_conversion = true;
+  pass_config.enable_stablehlo_conversion =
+      converter_flags.convert_to_stablehlo();
+
+  mlir::OwningOpRef<mlir::ModuleOp> module;
+  std::string content(input.data(), input.size());
+  if (model_flags.hlo_file_type() == tflite::ModelFlags::HLO_TEXT) {
+    TF_ASSIGN_OR_RETURN(auto hlo_module,
+                        xla::ParseAndReturnUnverifiedModule(content));
+    TF_ASSIGN_OR_RETURN(auto module,
+                        xla::ConvertHloToStablehlo(*context, hlo_module.get()));
+  } else if (model_flags.hlo_file_type() == tflite::ModelFlags::HLO_PROTO) {
+    TF_ASSIGN_OR_RETURN(xla::HloProto hlo_proto, LoadHloProto(content));
+    TF_ASSIGN_OR_RETURN(module, xla::ConvertHloToStablehlo(
+                                    *context, hlo_proto.mutable_hlo_module()));
+  } else {
+    return absl::InvalidArgumentError("Unknown hlo format type");
+  }
+
+  // Set the input names.
+  auto main_func = module->lookupSymbol<mlir::func::FuncOp>("main");
+  if (!main_func)
+    return absl::InternalError("Failed to find the main function.");
+  // Retrieve input names from model flags.
+  std::vector<std::string> input_names;
+  for (const auto& input : model_flags.input_arrays()) {
+    input_names.push_back(input.name());
+  }
+
+  const auto& inputs = absl::StrJoin(input_names, ",");
+  mlir::OpBuilder builder(*module);
+  llvm::SmallVector<mlir::NamedAttribute> attrs;
+  attrs.push_back(
+      builder.getNamedAttr("inputs", builder.getStringAttr(inputs)));
+  // Jax wrapped the output nodes in a tuple, so it's pretty hard to us
+  // to tell the output at this point, we will set the output at the export
+  // phase.
+  main_func->setAttr("tf.entry_function", builder.getDictionaryAttr(attrs));
+
+  // StableHLO Quantizer is not supported for JAX input models, so
+  // quantization_py_function_lib is set to nullptr.
+  auto status = internal::ConvertMLIRToTFLiteFlatBuffer(
+      model_flags, converter_flags, std::move(context), std::move(module),
+      pass_config, /*saved_model_tags=*/{}, result,
+      /*quantization_py_function_lib=*/nullptr);
+  return status;
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/lite/python/jax_to_tfl_flatbuffer.h

@@ -0,0 +1,38 @@
+/* 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_COMPILER_MLIR_LITE_PYTHON_JAX_TO_TFL_FLATBUFFER_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_JAX_TO_TFL_FLATBUFFER_H_
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// Converts the given Jax model to a TF Lite FlatBuffer
+// string according to the given model flags, converter flags and tags. Returns
+// error status if it fails to convert the input.
+absl::Status ConvertJaxToTFLiteFlatBuffer(
+    const std::string& input, const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags, std::string* result);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_JAX_TO_TFL_FLATBUFFER_H_

tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.cc

@@ -0,0 +1,262 @@
+/* Copyright 2022 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/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "llvm/ADT/StringSet.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h"
+#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
+#include "tensorflow/compiler/mlir/lite/tf_to_tfl_flatbuffer.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/lite/types.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+using tensorflow::quantization::PyFunctionLibrary;
+
+absl::Status HandleInputOutputArraysWithModule(
+    const tflite::ModelFlags& model_flags,
+    mlir::OwningOpRef<mlir::ModuleOp>* module) {
+  mlir::func::FuncOp entry_function = nullptr;
+  for (auto func : module->get().getOps<mlir::func::FuncOp>()) {
+    if (auto tf_attrs =
+            func->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function")) {
+      // TODO(b/184697652): There could be multiple entry functions. Let's
+      // handle such cases if there are any needs for that.
+      if (entry_function != nullptr) {
+        return absl::InvalidArgumentError(
+            "There should be only one tf.entry_function");
+      }
+      entry_function = func;
+    }
+  }
+  if (entry_function == nullptr) {
+    return absl::InvalidArgumentError("no tf.entry_function found");
+  }
+
+  // Get the list of input Op names from the function attribute.
+  mlir::DictionaryAttr tf_attrs =
+      entry_function->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function");
+  llvm::SmallVector<llvm::StringRef, 4> function_input_names;
+  function_input_names.reserve(model_flags.input_arrays().size());
+  auto input_attr = tf_attrs.get("inputs");
+  if (!input_attr) {
+    return absl::InvalidArgumentError("no inputs attribute found");
+  }
+  auto input_names = mlir::cast<mlir::StringAttr>(input_attr).getValue();
+  input_names.split(function_input_names, ",", /*MaxSplit=*/-1,
+                    /*KeepEmpty=*/false);
+  const int function_input_names_size = function_input_names.size();
+  if (function_input_names_size != model_flags.input_arrays().size()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "input array size mismatch: got ", function_input_names.size(),
+        ", expected: ", model_flags.input_arrays().size()));
+  }
+  llvm::StringSet<> function_input_names_set;
+  function_input_names_set.insert(function_input_names.begin(),
+                                  function_input_names.end());
+  for (const auto& input_array : model_flags.input_arrays()) {
+    if (function_input_names_set.count(input_array.name()) == 0) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("input array name (", input_array.name(),
+                       ") does not exist in the given graph"));
+    }
+  }
+
+  // Get the list of output Op names from the function attribute.
+  llvm::SmallVector<llvm::StringRef, 4> function_output_names;
+  function_output_names.reserve(model_flags.output_arrays().size());
+  auto output_attr = tf_attrs.get("outputs");
+  if (!output_attr) {
+    return absl::InvalidArgumentError("no outputs attribute found");
+  }
+  auto output_names = mlir::cast<mlir::StringAttr>(output_attr).getValue();
+  output_names.split(function_output_names, ",", /*MaxSplit=*/-1,
+                     /*KeepEmpty=*/false);
+  const int function_output_names_size = function_output_names.size();
+  if (function_output_names_size != model_flags.output_arrays().size()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "output array size mismatch: got ", function_output_names.size(),
+        ", expected: ", model_flags.output_arrays().size()));
+  }
+  llvm::StringSet<> function_output_names_set;
+  function_output_names_set.insert(function_output_names.begin(),
+                                   function_output_names.end());
+  for (const auto& output_array : model_flags.output_arrays()) {
+    if (function_output_names_set.count(output_array) == 0) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("output array name (", output_array,
+                       ") does not exist in the given graph"));
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ConvertSavedModelToTFLiteFlatBuffer(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags, std::string* result,
+    const PyFunctionLibrary* quantization_py_function_lib) {
+  auto context = std::make_unique<mlir::MLIRContext>();
+  mlir::TFL::QuantizationSpecs quant_specs;
+
+  // Parse input arrays.
+  std::vector<std::string> node_names;
+  std::vector<std::string> node_dtypes;
+  std::vector<std::optional<std::vector<int>>> node_shapes;
+  std::vector<std::optional<double>> node_mins;
+  std::vector<std::optional<double>> node_maxs;
+
+  // Populate quantization specs.
+  TF_RETURN_IF_ERROR(internal::PopulateQuantizationSpecs(
+      model_flags, converter_flags, &quant_specs, &node_names, &node_dtypes,
+      &node_shapes, &node_mins, &node_maxs));
+
+  internal::WarningUnusedFlags(model_flags, converter_flags);
+
+  // Register all custom ops, including user-specified custom ops.
+  TF_RETURN_IF_ERROR(internal::RegisterAllCustomOps(converter_flags));
+
+  auto& saved_model_tags = model_flags.saved_model_tags();
+  auto& saved_model_exported_names = model_flags.saved_model_exported_names();
+  std::unordered_set<std::string> tags(saved_model_tags.begin(),
+                                       saved_model_tags.end());
+  auto exported_names_in_vector = std::vector<std::string>(
+      saved_model_exported_names.begin(), saved_model_exported_names.end());
+  absl::Span<std::string> exported_names(exported_names_in_vector);
+
+  if (exported_names.empty()) {
+    return absl::UnimplementedError("Need at least one exported name.");
+  }
+
+  tensorflow::GraphImportConfig specs;
+  specs.upgrade_legacy = true;
+
+  std::vector<std::string> custom_opdefs(
+      converter_flags.custom_opdefs().begin(),
+      converter_flags.custom_opdefs().end());
+  TF_ASSIGN_OR_RETURN(
+      auto module,
+      ImportSavedModel(model_flags.saved_model_dir(),
+                       model_flags.saved_model_version(), tags,
+                       absl::MakeSpan(custom_opdefs), exported_names, specs,
+                       /*enable_variable_lifting=*/true, context.get(),
+                       /*saved_model_bundle=*/nullptr));
+
+  if (!model_flags.input_arrays().empty() ||
+      !model_flags.output_arrays().empty()) {
+    TF_RETURN_IF_ERROR(HandleInputOutputArraysWithModule(model_flags, &module));
+  }
+
+  mlir::TFL::PassConfig pass_config(quant_specs);
+  bool emit_builtin_tflite_ops = !converter_flags.force_select_tf_ops();
+  pass_config.emit_builtin_tflite_ops = emit_builtin_tflite_ops;
+  pass_config.enable_tflite_variables =
+      converter_flags.enable_tflite_resource_variables();
+  pass_config.unfold_batch_matmul = converter_flags.unfold_batchmatmul();
+  pass_config.lower_tensor_list_ops = converter_flags.lower_tensor_list_ops();
+  // Disable the unfolding of the 16x16 TF::BatchMatMulOp to avoid the
+  // conversion to an unsupported 16x16 TFL::FullyConnectedOp.
+  if (converter_flags.inference_type() == tflite::IODataType::QUANTIZED_INT16) {
+    pass_config.unfold_batch_matmul = false;
+  }
+  pass_config.unfold_large_splat_constant =
+      converter_flags.unfold_large_splat_constant();
+  pass_config.enable_dynamic_update_slice =
+      converter_flags.enable_dynamic_update_slice();
+  pass_config.preserve_assert_op = converter_flags.preserve_assert_op();
+  pass_config.guarantee_all_funcs_one_use =
+      converter_flags.guarantee_all_funcs_one_use();
+  pass_config.enable_stablehlo_conversion =
+      converter_flags.convert_to_stablehlo();
+  pass_config.legalize_custom_tensor_list_ops =
+      converter_flags.legalize_custom_tensor_list_ops();
+  pass_config.enable_composite_direct_lowering =
+      converter_flags.enable_composite_direct_lowering();
+  pass_config.model_origin_framework = converter_flags.model_origin_framework();
+  pass_config.canonicalizing_inf_as_min_max_float =
+      converter_flags.canonicalizing_inf_as_min_max_float();
+  pass_config.unsafe_fuse_dynamic_shaped_broadcast =
+      converter_flags.unsafe_fuse_dynamic_shaped_broadcast();
+
+  pass_config.unsafe_single_batch_rank_reduction =
+      converter_flags.unsafe_single_batch_rank_reduction();
+
+  if (converter_flags.strict_qdq_mode()) {
+    pass_config.quant_specs.qdq_conversion_mode =
+        mlir::TFL::QDQConversionMode::kQDQStrict;
+  } else if (converter_flags.qdq_conversion_mode() == "STATIC") {
+    pass_config.quant_specs.qdq_conversion_mode =
+        mlir::TFL::QDQConversionMode::kQDQStatic;
+  } else if (converter_flags.qdq_conversion_mode() == "DYNAMIC") {
+    pass_config.quant_specs.qdq_conversion_mode =
+        mlir::TFL::QDQConversionMode::kQDQDynamic;
+    // Need to set this or else the ops will still use floating point kernels
+    pass_config.quant_specs.inference_type = tensorflow::DT_QINT8;
+  } else if (converter_flags.qdq_conversion_mode() == "NONE") {
+    pass_config.quant_specs.qdq_conversion_mode =
+        mlir::TFL::QDQConversionMode::kQDQNone;
+  } else {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Unknown QDQ conversion mode: ",
+                     converter_flags.qdq_conversion_mode()));
+  }
+
+  if (converter_flags.has_qdq_conversion_mode() &&
+      converter_flags.qdq_conversion_mode() != "NONE") {
+    // Setting this flag causes
+    // PrepareQuantize::SetInputNodesQuantizationParams() to be false and allows
+    // PrepareQuantizePass to complete. For the most part this step is
+    // unnecessary for non-TF QDQ models.
+    pass_config.quant_specs.disable_set_input_nodes_quantization_params = true;
+  }
+
+  // TODO(b/153507667): Pass the session object when importing logic is removed.
+  auto status = internal::ConvertMLIRToTFLiteFlatBuffer(
+      model_flags, converter_flags, std::move(context), std::move(module),
+      pass_config, tags, result, quantization_py_function_lib);
+  return status;
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/lite/python/saved_model_to_tfl_flatbuffer.h

@@ -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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_SAVED_MODEL_TO_TFL_FLATBUFFER_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_SAVED_MODEL_TO_TFL_FLATBUFFER_H_
+
+#include "absl/status/status.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// Converts the given saved_model(either v1 or v2) to a TF Lite FlatBuffer
+// string according to the given model flags, converter flags and tags. Returns
+// error status if it fails to convert the input.
+absl::Status ConvertSavedModelToTFLiteFlatBuffer(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags, std::string* result,
+    const quantization::PyFunctionLibrary* quantization_py_function_lib);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_SAVED_MODEL_TO_TFL_FLATBUFFER_H_

tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc

@@ -0,0 +1,384 @@
+/* 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/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Support/FileUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
+#include "tensorflow/compiler/mlir/lite/tf_to_tfl_flatbuffer.h"
+#include "tensorflow/compiler/mlir/lite/tools/optimize/reduced_precision_metadata.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/lite/types.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/graph_debug_info.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tsl/platform/protobuf.h"  // IWYU pragma: keep
+
+namespace tensorflow {
+namespace internal {
+namespace {
+
+using ::tensorflow::quantization::PyFunctionLibrary;
+using ::tflite::optimize::ReducedPrecisionSupport;
+
+// Op def string for TFLite_Detection_PostProcess Op.
+constexpr mlir::StringRef kDetectionPostProcessOp =
+    "name: 'TFLite_Detection_PostProcess' input_arg: { name: "
+    "'raw_outputs/box_encodings' type: DT_FLOAT } input_arg: { name: "
+    "'raw_outputs/class_predictions' type: DT_FLOAT } input_arg: { name: "
+    "'anchors' type: DT_FLOAT } output_arg: { name: "
+    "'TFLite_Detection_PostProcess' type: DT_FLOAT } output_arg: { name: "
+    "'TFLite_Detection_PostProcess:1' type: DT_FLOAT } output_arg: { name: "
+    "'TFLite_Detection_PostProcess:2' type: DT_FLOAT } output_arg: { name: "
+    "'TFLite_Detection_PostProcess:3' type: DT_FLOAT } attr : { name: "
+    "'h_scale' type: 'float'} attr : { name: 'max_classes_per_detection' "
+    "type: 'int'} attr : { name: 'max_detections' type: 'int'} attr : { "
+    "name: 'nms_iou_threshold' type: 'float'} attr : { name: "
+    "'nms_score_threshold' type: 'float'} attr : { name: 'num_classes' type: "
+    "'int'} attr : { name: 'w_scale' type: 'float'} attr : { name: 'x_scale' "
+    "type: 'float'} attr : { name: 'y_scale' type: 'float'} attr { name: "
+    "'detections_per_class' type: 'int' default_value { i : 100 }} attr { "
+    "name: 'use_regular_nms' type: 'bool' default_value { b : false }}";
+
+constexpr mlir::StringRef kUnidirectionalSequenceLstmOp =
+    "name: 'UnidirectionalSequenceLstm' input_arg: {name: 'Input' type: "
+    "DT_FLOAT} input_arg: { name: 'InputToInputWeights' type: DT_FLOAT } "
+    "input_arg: { name: 'InputToForgetWeights' type: DT_FLOAT } input_arg: { "
+    "name: 'InputToCellWeights' type: DT_FLOAT} input_arg: { name: "
+    "'InputToOutputWeights' type: DT_FLOAT } input_arg: { name: "
+    "'RecurrentToInputWeights' type: DT_FLOAT} input_arg: { name: "
+    "'RecurrentToForgetWeights' type: DT_FLOAT} input_arg: { name: "
+    "'RecurrentToCellWeights' type: DT_FLOAT } input_arg: { name: "
+    "'RecurrentToOutputWeights' type: DT_FLOAT } input_arg: { name: "
+    "'CellToInputWeights' type: DT_FLOAT} input_arg: { name: "
+    "'CellToForgetWeights' type: DT_FLOAT } input_arg: { name: "
+    "'CellToOutputWeights' type: DT_FLOAT } input_arg: { name: 'InputGateBias' "
+    "type: DT_FLOAT } input_arg: { name: 'ForgetGateBias' type: DT_FLOAT } "
+    "input_arg: { name: 'kCellGateBias' type: DT_FLOAT } input_arg: { name: "
+    "'OutputGateBias' type: DT_FLOAT } input_arg: { name: 'ProjectionWeights' "
+    "type: DT_FLOAT } input_arg: { name: 'ProjectionBias' type: DT_FLOAT } "
+    "input_arg: { name: 'InputActivationState' type: DT_FLOAT} input_arg: { "
+    "name: 'InputCellStateTensor' type: DT_FLOAT } "
+    "output_arg: { name: 'Concat' type: DT_FLOAT} "
+    "output_arg: { name: "
+    "'LastState' type: DT_FLOAT } output_arg: { name: 'Output' type: DT_FLOAT} "
+    "attr : { name: '_tflite_input_indices' type: 'list(int)'}";
+
+constexpr mlir::StringRef kUnidirectionalSequenceRnnOp =
+    "name: 'UnidirectionalSequenceRnn' input_arg: {name: 'Input' type: "
+    "DT_FLOAT} input_arg: { name: 'Weights' type: DT_FLOAT } "
+    "input_arg: { name: 'RecurrentWeights' type: DT_FLOAT } input_arg: { "
+    "name: 'Bias' type: DT_FLOAT} "
+    "input_arg: { name: 'HiddenState' type: DT_FLOAT} "
+    "output_arg: { name: "
+    "'LastState' type: DT_FLOAT } output_arg: { name: 'Output' type: "
+    "DT_FLOAT} "
+    "attr : { name: '_tflite_input_indices' type: 'list(int)'}";
+
+// Converts the tflite::IODataType to tensorflow::DataType. Only contains the
+// conversion mapping for constants defined in TFLite Python API.
+DataType ConvertIODataTypeToDataType(tflite::IODataType dtype) {
+  switch (dtype) {
+    case tflite::IODataType::FLOAT:
+      return DT_FLOAT;
+    case tflite::IODataType::FLOAT16:
+      return DT_HALF;
+    case tflite::IODataType::FLOAT64:
+      return DT_DOUBLE;
+    case tflite::IODataType::QUANTIZED_UINT8:
+      return DT_QUINT8;
+    case tflite::IODataType::QUANTIZED_INT8:
+      return DT_QINT8;
+    case tflite::IODataType::QUANTIZED_INT16:
+      return DT_QINT16;
+    case tflite::IODataType::INT8:
+      return DT_INT8;
+    case tflite::IODataType::INT16:
+      return DT_INT16;
+    case tflite::IODataType::UINT16:
+      return DT_UINT16;
+    case tflite::IODataType::INT32:
+      return DT_INT32;
+    case tflite::IODataType::UINT32:
+      return DT_UINT32;
+    case tflite::IODataType::INT64:
+      return DT_INT64;
+    case tflite::IODataType::UINT8:
+      return DT_UINT8;
+    case tflite::IODataType::UINT64:
+      return DT_UINT64;
+    case tflite::IODataType::STRING:
+      return DT_STRING;
+    case tflite::IODataType::BOOL:
+      return DT_BOOL;
+    case tflite::IODataType::COMPLEX64:
+      return DT_COMPLEX64;
+    case tflite::IODataType::COMPLEX128:
+      return DT_COMPLEX128;
+    case tflite::IODataType::RESOURCE:
+      return DT_RESOURCE;
+    case tflite::IODataType::VARIANT:
+      return DT_VARIANT;
+    default:
+      return DT_INVALID;
+  }
+}
+
+absl::StatusOr<std::pair<double, double>> InputStatsToMinMax(double mean,
+                                                             double std,
+                                                             DataType type) {
+  // Only qint8 and quint8 are considered here.
+  double qmin, qmax;
+  if (type == DT_QUINT8) {
+    qmin = 0.0;
+    qmax = 255.0;
+  } else if (type == DT_QINT8) {
+    qmin = -128.0;
+    qmax = 127.0;
+  } else {
+    return absl::InvalidArgumentError("Only int8 and uint8 are considered.");
+  }
+  return std::make_pair((qmin - mean) / std, (qmax - mean) / std);
+}
+
+absl::Status RegisterCustomBuiltinOps(
+    const std::vector<std::string> extra_tf_opdefs) {
+  for (const auto& tf_opdefs_string : extra_tf_opdefs) {
+    OpDef opdef;
+    // NOLINTNEXTLINE: Use tsl::protobuf to be compatible with OSS.
+    if (!tsl::protobuf::TextFormat::ParseFromString(tf_opdefs_string, &opdef)) {
+      return absl::InvalidArgumentError("fail to parse extra OpDef");
+    }
+    // Make sure the op is not already registered. If registered continue.
+    const OpRegistrationData* op_reg =
+        OpRegistry::Global()->LookUp(opdef.name());
+    if (op_reg) continue;
+
+    OpRegistry::Global()->Register(
+        [opdef](OpRegistrationData* op_reg_data) -> absl::Status {
+          *op_reg_data = OpRegistrationData(opdef);
+          return absl::OkStatus();
+        });
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status RegisterAllCustomOps(
+    const tflite::ConverterFlags& converter_flags) {
+  // Register any custom OpDefs.
+  std::vector<std::string> extra_tf_opdefs(
+      converter_flags.custom_opdefs().begin(),
+      converter_flags.custom_opdefs().end());
+  extra_tf_opdefs.push_back(kDetectionPostProcessOp.str());
+  extra_tf_opdefs.push_back(kUnidirectionalSequenceLstmOp.str());
+  extra_tf_opdefs.push_back(kUnidirectionalSequenceRnnOp.str());
+  return RegisterCustomBuiltinOps(extra_tf_opdefs);
+}
+
+absl::Status PopulateQuantizationSpecs(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags,
+    mlir::TFL::QuantizationSpecs* quant_specs,
+    std::vector<std::string>* node_names, std::vector<std::string>* node_dtypes,
+    std::vector<std::optional<std::vector<int>>>* node_shapes,
+    std::vector<std::optional<double>>* node_mins,
+    std::vector<std::optional<double>>* node_maxs) {
+  quant_specs->inference_input_type =
+      ConvertIODataTypeToDataType(converter_flags.inference_input_type());
+  DataType inference_type =
+      ConvertIODataTypeToDataType(converter_flags.inference_type());
+  // Use non-float flag `inference_input_type` to override the `inference_type`
+  // because we have to apply quantization to satisfy that.
+  if (quant_specs->inference_input_type != DT_FLOAT) {
+    inference_type = quant_specs->inference_input_type;
+  }
+
+  for (auto& flag : model_flags.input_arrays()) {
+    node_names->push_back(flag.name());
+    // TOCO doesn't required `data_type` to be filled for every input.
+    // If it's not filled, make it an empty string so the importer will use
+    // the data type in the NodeDef.
+    auto tflite_data_type = flag.data_type();
+    if (tflite_data_type == tflite::IODataType::IO_DATA_TYPE_UNKNOWN) {
+      node_dtypes->push_back("");
+    } else {
+      node_dtypes->push_back(
+          DataType_Name(ConvertIODataTypeToDataType(tflite_data_type)));
+    }
+    if (flag.shape().unknown_rank()) {
+      node_shapes->push_back(std::nullopt);
+    } else {
+      node_shapes->push_back(std::vector<int>(flag.shape().dims().begin(),
+                                              flag.shape().dims().end()));
+    }
+    // Currently, only UINT8 and INT8 require inputs stats
+    if (inference_type == DT_QINT8 || inference_type == DT_QUINT8) {
+      if (flag.has_mean_value() && flag.has_std_value()) {
+        TF_ASSIGN_OR_RETURN(
+            auto min_max, InputStatsToMinMax(flag.mean_value(),
+                                             flag.std_value(), inference_type));
+        node_mins->push_back(min_max.first);
+        node_maxs->push_back(min_max.second);
+      } else {
+        node_mins->push_back(std::nullopt);
+        node_maxs->push_back(std::nullopt);
+      }
+    }
+  }
+
+  if (mlir::TFL::GetInputNodeQuantSpecs(*node_names, *node_mins, *node_maxs,
+                                        inference_type, quant_specs)) {
+    return absl::InvalidArgumentError("Failed to get input quant spec.");
+  }
+
+  // Some extra flag related to post training quantization. If post-training
+  // quantization is enabled, `inference_type` and `inference_input_type` are
+  // not used by MLIR passes.
+  if (converter_flags.post_training_quantize()) {
+    quant_specs->weight_quantization = true;
+    quant_specs->disable_per_channel =
+        converter_flags.disable_per_channel_quantization();
+    if (converter_flags.quantize_to_float16()) {
+      quant_specs->inference_type = DT_HALF;
+      quant_specs->inference_input_type = DT_HALF;
+    } else {
+      quant_specs->inference_type = DT_QINT8;
+      quant_specs->inference_input_type = DT_QINT8;
+    }
+  } else {
+    // These flags are incompatible with post_training_quantize() as only
+    // QAT models can provide required ranges.
+    quant_specs->disable_infer_tensor_range =
+        converter_flags.disable_infer_tensor_range();
+    quant_specs->use_fake_quant_num_bits =
+        converter_flags.use_fake_quant_num_bits();
+  }
+
+  // Add information about half-precision support if fp16 quantization applies.
+  // TODO(b/195945955): Add e2e test for this.
+  if (converter_flags.quantize_to_float16() ||
+      converter_flags.allow_bfloat16()) {
+    ReducedPrecisionSupport mask = ReducedPrecisionSupport::None;
+    if (converter_flags.quantize_to_float16()) {
+      mask |= ReducedPrecisionSupport::Float16Inference;
+    }
+    if (converter_flags.allow_bfloat16()) {
+      mask |= ReducedPrecisionSupport::Bfloat16Inference;
+    }
+    if (converter_flags.accumulation_type() == tflite::IODataType::FLOAT16) {
+      mask |= ReducedPrecisionSupport::Float16Accumulation;
+    } else {
+      mask |= ReducedPrecisionSupport::Float32Accumulation;
+    }
+    quant_specs->support_mask = mask;
+  }
+
+  // Other flags.
+  if (converter_flags.has_default_ranges_min()) {
+    quant_specs->default_ranges.first = converter_flags.default_ranges_min();
+  }
+  if (converter_flags.has_default_ranges_max()) {
+    quant_specs->default_ranges.second = converter_flags.default_ranges_max();
+  }
+  quant_specs->enable_mlir_dynamic_range_quantizer =
+      converter_flags.enable_mlir_dynamic_range_quantizer();
+  quant_specs->enable_mlir_variable_quantization =
+      converter_flags.enable_mlir_variable_quantization();
+  quant_specs->disable_per_channel_for_dense_layers =
+      converter_flags.disable_per_channel_quantization_for_dense_layers();
+  return absl::OkStatus();
+}
+
+absl::Status ConvertMLIRToTFLiteFlatBuffer(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags,
+    std::unique_ptr<mlir::MLIRContext>&& context,
+    mlir::OwningOpRef<mlir::ModuleOp> module,
+    const mlir::TFL::PassConfig& pass_config,
+    const std::unordered_set<std::string>& saved_model_tags,
+    std::string* result,
+    const PyFunctionLibrary* quantization_py_function_lib) {
+  mlir::TFL::PassConfig pass_config_copy = pass_config;
+  pass_config_copy.outline_tf_while = true;
+
+  // Checks whether the model contains an `XlaCallModuleOp` operation which
+  // is a wrapper around StableHLO.
+  // This option is mutually exclusive to `enable_stablehlo_conversion`, the
+  // latter of which takes precedence.
+  // TODO(b/290109282): explore removing the enable_hlo_to_tf_conversion flag
+  // entirely, such that the added passes are no-ops in the non-shlo case.
+  module->walk([&](mlir::TF::XlaCallModuleOp xla_call_module_op) {
+    pass_config_copy.enable_hlo_to_tf_conversion = true;
+    mlir::WalkResult::interrupt();
+  });
+
+  auto status = ConvertTFExecutorToTFLOrFlatbuffer(
+      std::move(context), std::move(module), converter_flags, pass_config_copy,
+      saved_model_tags, model_flags.saved_model_dir(), result,
+      /*export_to_mlir=*/false, quantization_py_function_lib);
+
+  return status;
+}
+
+void WarningUnusedFlags(const tflite::ModelFlags& model_flags,
+                        const tflite::ConverterFlags& converter_flags) {
+  if (converter_flags.output_format()) {
+    LOG(WARNING) << "Ignored output_format.";
+  }
+  if (converter_flags.drop_control_dependency()) {
+    LOG(WARNING) << "Ignored drop_control_dependency.";
+  }
+  if (converter_flags.reorder_across_fake_quant()) {
+    LOG(WARNING) << "Ignored reorder_across_fake_quant.";
+  }
+  if (model_flags.change_concat_input_ranges()) {
+    LOG(WARNING) << "Ignored change_concat_input_ranges.";
+  }
+  if (converter_flags.dump_graphviz_include_video()) {
+    LOG(WARNING) << "Ignored dump_graphviz_video.";
+  }
+  if (model_flags.allow_nonexistent_arrays()) {
+    LOG(WARNING) << "Allow allow_nonexistent_arrays.";
+  }
+}
+
+}  // namespace internal
+}  // namespace tensorflow

tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.h

@@ -0,0 +1,73 @@
+/* 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_COMPILER_MLIR_LITE_PYTHON_TF_TFL_FLATBUFFER_HELPERS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_TF_TFL_FLATBUFFER_HELPERS_H_
+
+#include <optional>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/common/tfl_pass_config.h"
+#include "tensorflow/compiler/mlir/lite/converter_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/model_flags.pb.h"
+#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h"
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"
+#include "tensorflow/compiler/mlir/lite/types.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace internal {
+
+// Register all custom ops including user specified custom ops.
+absl::Status RegisterAllCustomOps(
+    const tflite::ConverterFlags& converter_flags);
+
+// Populate quantization specs (or not) given user specified ranges for each
+// input arrays.
+absl::Status PopulateQuantizationSpecs(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags,
+    mlir::TFL::QuantizationSpecs* quant_specs,
+    std::vector<std::string>* node_names, std::vector<std::string>* node_dtypes,
+    std::vector<std::optional<std::vector<int>>>* node_shapes,
+    std::vector<std::optional<double>>* node_mins,
+    std::vector<std::optional<double>>* node_maxs);
+
+// Convert imported MLIR file to TfLite flatbuffer.
+// This will also run relevant passes as well.
+absl::Status ConvertMLIRToTFLiteFlatBuffer(
+    const tflite::ModelFlags& model_flags,
+    tflite::ConverterFlags& converter_flags,
+    std::unique_ptr<mlir::MLIRContext>&& context,
+    mlir::OwningOpRef<mlir::ModuleOp> module,
+    const mlir::TFL::PassConfig& pass_config,
+    const std::unordered_set<std::string>& saved_model_tags,
+    std::string* result,
+    const quantization::PyFunctionLibrary* quantization_py_function_lib);
+
+// Give a warning for any unused flags that have been specified.
+void WarningUnusedFlags(const tflite::ModelFlags& model_flags,
+                        const tflite::ConverterFlags& converter_flags);
+}  // namespace internal
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_PYTHON_TF_TFL_FLATBUFFER_HELPERS_H_

tensorflow/compiler/mlir/lite/python/wrap_converter.py

@@ -0,0 +1,90 @@
+# 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.
+# ==============================================================================
+"""Wraps TFLite Converter interface with python lazy loader."""
+# We need to import pywrap_tensorflow prior to the converter wrapper.
+# pylint: disable=invalid-import-order,g-bad-import-order
+from tensorflow.python import pywrap_tensorflow  # pylint: disable=unused-import
+from tensorflow.compiler.mlir.lite.python import _pywrap_converter_api
+from tensorflow.compiler.mlir.quantization.tensorflow.python import py_function_lib
+
+
+def wrapped_convert(
+    model_flags_str,
+    converter_flags_str,
+    input_data_str,
+    debug_info_str,
+):
+  """Wraps Convert with lazy loader."""
+  return _pywrap_converter_api.Convert(
+      model_flags_str,
+      converter_flags_str,
+      input_data_str,
+      False,  # extended_return
+      debug_info_str,
+      py_function_lib.PyFunctionLibrary(),
+  )
+
+
+def wrapped_experimental_mlir_quantize(
+    input_data_str,
+    disable_per_channel,
+    fully_quantize,
+    inference_type,
+    input_data_type,
+    output_data_type,
+    enable_numeric_verify,
+    enable_whole_model_verify,
+    denylisted_ops,
+    denylisted_nodes,
+    enable_variable_quantization,
+    disable_per_channel_for_dense_layers,
+    debug_options_str,
+):
+  """Wraps experimental mlir quantize model."""
+  return _pywrap_converter_api.ExperimentalMlirQuantizeModel(
+      input_data_str,
+      disable_per_channel,
+      fully_quantize,
+      inference_type,
+      input_data_type,
+      output_data_type,
+      enable_numeric_verify,
+      enable_whole_model_verify,
+      denylisted_ops,
+      denylisted_nodes,
+      enable_variable_quantization,
+      disable_per_channel_for_dense_layers,
+      debug_options_str,
+  )
+
+
+def wrapped_experimental_mlir_sparsify(input_data_str):
+  """Wraps experimental mlir sparsify model."""
+  return _pywrap_converter_api.ExperimentalMlirSparsifyModel(input_data_str)
+
+
+def wrapped_register_custom_opdefs(custom_opdefs_list):
+  """Wraps RegisterCustomOpdefs with lazy loader."""
+  return _pywrap_converter_api.RegisterCustomOpdefs(custom_opdefs_list)
+
+
+def wrapped_retrieve_collected_errors():
+  """Wraps RetrieveCollectedErrors with lazy loader."""
+  return _pywrap_converter_api.RetrieveCollectedErrors()
+
+
+def wrapped_flat_buffer_file_to_mlir(model, input_is_filepath):
+  """Wraps FlatBufferFileToMlir with lazy loader."""
+  return _pywrap_converter_api.FlatBufferToMlir(model, input_is_filepath)