chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/quantization/stablehlo/python/BUILD

@@ -0,0 +1,171 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@xla//xla/tsl/platform:build_config_root.bzl", "if_pywrap")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load(
+    "//tensorflow:tensorflow.default.bzl",
+    "get_compatible_with_portable",
+    "tf_py_strict_test",
+    "tf_python_pybind_extension",
+)
+load("//tensorflow/compiler/mlir/quantization/stablehlo:internal_visibility_allowlist.bzl", "internal_visibility_allowlist")
+
+package_group(
+    name = "internal_visibility_allowlist_package",
+    packages = [
+        "//tensorflow/compiler/mlir/lite/...",
+        "//tensorflow/compiler/mlir/quantization/...",
+        "//tensorflow/compiler/mlir/tf2xla/transforms/...",
+        "//tensorflow/lite/...",
+        "//third_party/cloud_tpu/inference_converter/...",  # TPU Inference Converter V1
+    ] + internal_visibility_allowlist(),
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["@stablehlo//:license"],
+    default_visibility = [
+        ":internal_visibility_allowlist_package",
+        "//tensorflow:__pkg__",
+    ],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "quantization",
+    srcs = ["quantization.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":pywrap_quantization",
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib_py",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:save_model",
+        "//tensorflow/core:protos_all_py",
+    ],
+)
+
+# copybara:uncomment_begin(google-only)
+# pytype_strict_library(
+#     name = "quantize_model_test_base",
+#     testonly = 1,
+#     srcs = ["integration_test/quantize_model_test_base.py"],
+#     tags = ["no_pip"],
+#     visibility = ["//visibility:private"],
+#     deps = [
+#         "//third_party/py/mlir:ir",
+#         "//third_party/py/mlir:stablehlo_dialect",
+#         "//third_party/py/mlir/_mlir_libs:_mlirRegisterEverything",
+#         "//third_party/py/numpy",
+#         "//tensorflow:tensorflow_py",
+#         "//tensorflow/compiler/mlir/stablehlo",
+#         "//tensorflow/python/eager:def_function",
+#         "//tensorflow/python/framework:dtypes",
+#         "//tensorflow/python/framework:ops",
+#         "//tensorflow/python/framework:tensor_spec",
+#         "//tensorflow/python/module",
+#         "//tensorflow/python/ops:array_ops",
+#         "//tensorflow/python/ops:math_ops",
+#         "//tensorflow/python/ops:nn_ops",
+#         "//tensorflow/python/ops:variables",
+#         "//tensorflow/python/platform:client_testlib",
+#         "//tensorflow/python/platform:tf_logging",
+#         "//tensorflow/python/saved_model:load",
+#         "//tensorflow/python/saved_model:loader",
+#         "//tensorflow/python/saved_model:save",
+#         "//tensorflow/python/types:core",
+#         "@absl_py//absl/testing:parameterized",
+#     ],
+# )
+#
+# tf_py_strict_test(
+#     name = "quantize_model_test",
+#     srcs = ["integration_test/quantize_model_test.py"],
+#     shard_count = 50,  # Parallelize the test to avoid timeouts.
+#     deps = [
+#         ":quantization",
+#         ":quantize_model_test_base",
+#         "//tensorflow/compiler/mlir/quantization/common/python:testing",
+#         "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
+#         "//tensorflow/compiler/mlir/quantization/tensorflow/python:representative_dataset",
+#         "//tensorflow/python/eager:def_function",
+#         "//tensorflow/python/framework:dtypes",
+#         "//tensorflow/python/framework:ops",
+#         "//tensorflow/python/framework:tensor_spec",
+#         "//tensorflow/python/framework:test_lib",
+#         "//tensorflow/python/module",
+#         "//tensorflow/python/ops:math_ops",
+#         "//tensorflow/python/ops:nn_ops",
+#         "//tensorflow/python/platform:client_testlib",
+#         "//tensorflow/python/saved_model:load",
+#         "//tensorflow/python/saved_model:save",
+#         "//tensorflow/python/saved_model:tag_constants",
+#         "//tensorflow/python/types:core",
+#         "@absl_py//absl/testing:parameterized",
+#     ],
+# )
+# copybara:uncomment_end
+
+# This is a header-only target. The purpose of `pywrap_quantization_lib_*` targets is to expose only
+# the symbols that are required by `pywrap_quantization` that translates them to python functions.
+# The only intended use case of this library is by `pywrap_quantization`. Not letting
+# `pywrap_quantization` directly depend on sub-libraries like `static_range_srq` and instead haiving
+# a consolidated impl library `pywrap_quantization_lib_impl` allows the maintainers to avoid
+# declaring multiple impl libraries to `libtensorflow_cc` and `lib_pywrap_tensorflow_internal`,
+# which is required to avoid ODR violations.
+cc_library(
+    name = "pywrap_quantization_lib_header_only",
+    srcs = [],
+    hdrs = ["pywrap_quantization_lib.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:private"],  # ONLY for `pywrap_quantization`.
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+# See the comments for `pywrap_quantization_lib_header_only`.
+cc_library(
+    name = "pywrap_quantization_lib_impl",
+    srcs = ["pywrap_quantization_lib.cc"],
+    hdrs = ["pywrap_quantization_lib.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = [
+        "//tensorflow:__pkg__",  # For libtensorflow_cc.so.
+        "//tensorflow/python:__pkg__",  # For lib_pywrap_tensorflow_internal.so.
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_cc",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:config",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:static_range_ptq",
+        "//tensorflow/compiler/mlir/quantization/stablehlo/cc:weight_only_ptq",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:py_function_lib",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "pywrap_quantization",
+    srcs = ["pywrap_quantization.cc"],
+    pytype_srcs = ["pywrap_quantization.pyi"],
+    starlark_only = True,
+    visibility = [
+        "//tensorflow/python:__pkg__",
+    ],
+    # Each dependency MUST be either header-only or exclusive.
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/tensorflow/python:type_casters",
+        "@pybind11",
+        "@pybind11_abseil//pybind11_abseil:absl_casters",
+        "@pybind11_abseil//pybind11_abseil:import_status_module",
+        "@pybind11_abseil//pybind11_abseil:status_casters",
+    ] + if_pywrap(
+        if_false = [":pywrap_quantization_lib_header_only"],
+        if_true = [":pywrap_quantization_lib_impl"],
+    ),
+)

tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/quantize_model_test_base.py

@@ -0,0 +1,551 @@
+# 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.
+# ==============================================================================
+"""Base test class for quantize_model Tests."""
+
+from typing import List, Mapping, Optional, Sequence, Tuple
+
+from absl.testing import parameterized
+from mlir import ir
+from mlir.dialects import stablehlo as stablehlo_dialect
+import numpy as np
+import tensorflow  # pylint: disable=unused-import
+
+from tensorflow.compiler.mlir.stablehlo import stablehlo
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.module import module
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import loader_impl
+from tensorflow.python.saved_model import save as saved_model_save
+from tensorflow.python.types import core
+
+
+FUNC_ALIAS = 'some_alias'
+
+
+class QuantizedModelTest(test.TestCase, parameterized.TestCase):
+  """Base test class for StableHLO quant tests."""
+
+  def setUp(self) -> None:
+    super().setUp()
+
+    # Many test cases for quantization involve creating and saving the input
+    # model and saving the output quantized model. These two member
+    # attributes can be used to specify the paths for such models,
+    # respectively. These paths will be cleaned up after each test case.
+    self._input_saved_model_path = self.create_tempdir('input').full_path
+    self._output_saved_model_path = self.create_tempdir('output').full_path
+    # Extra output path occasionally used for comparing two different
+    # quantized models.
+    self._output_saved_model_path_2 = self.create_tempdir('output2').full_path
+
+  def _extract_first_xla_call_module_op(
+      self, output_saved_model_path: str
+  ) -> str:
+    """Extracts the first XlaCallModule op from output saved model to string."""
+    root = load.load(output_saved_model_path)
+    tf_graph_def = root.signatures['serving_default'].graph.as_graph_def()
+    for function in tf_graph_def.library.function:
+      for node_def in function.node_def:
+        if node_def.op == 'XlaCallModule':
+          with ir.Context() as context:
+            stablehlo_dialect.register_dialect(context)
+            # Serialization in VHLO dialect.
+            serialized = node_def.attr.get('module').s
+            # MLIR bytecode matching StableHLO version.
+            mlir_bytecode = stablehlo.deserialize_portable_artifact_str(
+                serialized)
+            stablehlo_module = ir.Module.parse(mlir_bytecode, context=context)
+            return str(stablehlo_module)
+    raise ValueError('No XlaCallModule found in saved model.')
+
+  def _get_num_xla_call_module_op(self, output_saved_model_path: str) -> int:
+    """Gets the number of XlaCallModule ops in the output saved model."""
+    root = load.load(output_saved_model_path)
+    tf_graph_def = root.signatures['serving_default'].graph.as_graph_def()
+    count = 0
+    for node_def in tf_graph_def.node:
+      if node_def.op == 'XlaCallModule':
+        count += 1
+    for function in tf_graph_def.library.function:
+      for node_def in function.node_def:
+        if node_def.op == 'XlaCallModule':
+          count += 1
+    return count
+
+  def _get_function_aliases(
+      self, output_saved_model_path: str, tags: List[str]
+  ) -> dict[str, str]:
+    """Gets the function aliases in the output saved model."""
+    loader = loader_impl.SavedModelLoader(output_saved_model_path)
+    return loader.get_meta_graph_def_from_tags(
+        tags
+    ).meta_info_def.function_aliases
+
+  def _create_matmul_model(
+      self,
+      input_shape: Sequence[int],
+      weight_shape: Sequence[int],
+      saved_model_path: str,
+      bias_fn: Optional[ops.Operation] = None,
+      activation_fn: Optional[ops.Operation] = None,
+  ) -> module.Module:
+    class MatmulModel(module.Module):
+      """A simple model with a single matmul.
+
+      Bias and activation function are optional.
+      """
+
+      def __init__(
+          self,
+          weight_shape: Sequence[int],
+      ) -> None:
+        """Initializes a MatmulModel.
+
+        Args:
+          weight_shape: Shape of the weight tensor.
+        """
+        self.filters = np.random.uniform(low=-1.0, high=1.0, size=weight_shape)
+
+        if bias_fn is not None:
+          self.bias = np.random.uniform(
+              low=-1.0, high=1.0, size=weight_shape[-1]
+          )
+
+      def has_reshape(self) -> bool:
+        return self.bias_fn() and self.bias_size != self.filters.shape[-1]
+
+      @def_function.function
+      def matmul(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
+        """Performs a matrix multiplication.
+
+        Depending on self.bias_fn and self.activation_fn, it may add a bias
+        term or go through the activaction function.
+
+        Args:
+          input_tensor: Input tensor to matmul with the filter.
+
+        Returns:
+          A map of: output key -> output result.
+        """
+        out = math_ops.matmul(input_tensor, self.filters, name='sample/matmul')
+        if bias_fn is not None:
+          out = bias_fn(out, self.bias)
+        if activation_fn is not None:
+          out = activation_fn(out)
+        return {'output': out}
+
+    model = MatmulModel(weight_shape)
+    saved_model_save.save(
+        model,
+        saved_model_path,
+        signatures=model.matmul.get_concrete_function(
+            tensor_spec.TensorSpec(
+                shape=input_shape, dtype=dtypes.float32, name='input_tensor'
+            )
+        ),
+    )
+    return model
+
+  def _any_log_contains(
+      self, substring: str, log_record_list: List['logging.LogRecord']
+  ) -> bool:
+    """Returns True if any of the log contains a given substring.
+
+    Args:
+      substring: A piece of string to check whether it exists in the log
+        message.
+      log_record_list: A list of `absl.logging.LogRecord`s.
+
+    Returns:
+      True if and only if the substring exists in any of the log in
+      `log_record_list`.
+    """
+    return any(
+        map(
+            lambda log_record: substring in str(log_record.message),
+            log_record_list,
+        )
+    )
+
+  def _create_matmul_and_same_scale_model(
+      self,
+      input_shape: Sequence[int],
+      weight_shape: Sequence[int],
+      saved_model_path: str,
+      same_scale_op: str,
+  ) -> module.Module:
+    class MatmulAndSameScaleModel(module.Module):
+      """A simple model with a same-scale op.
+
+      Op name in StableHLO dialect is given as a string.
+      """
+
+      def __init__(
+          self,
+          weight_shape: Sequence[int],
+          same_scale_op: str,
+      ) -> None:
+        """Initializes a MatmulModel.
+
+        Args:
+          weight_shape: Shape of the weight tensor.
+          same_scale_op: Name of the same-scale op to be tested. Raises error
+            when an unknown name is given.
+        """
+        self.filters = np.random.uniform(low=-1.0, high=1.0, size=weight_shape)
+        self.same_scale_op = same_scale_op
+
+      @def_function.function
+      def matmul_and_same_scale(
+          self, input_tensor: core.Tensor
+      ) -> Mapping[str, core.Tensor]:
+        """Performs a matrix multiplication.
+
+        Args:
+          input_tensor: Input tensor to matmul with the filter.
+
+        Returns:
+          A map of: output key -> output result.
+        """
+        out = math_ops.matmul(input_tensor, self.filters, name='sample/matmul')
+
+        if self.same_scale_op == 'concatenate':
+          ones = array_ops.ones_like(out)
+          out = array_ops.concat([out, ones], 0)
+        elif self.same_scale_op == 'gather':
+          out = array_ops.gather(out, indices=[0], axis=0)
+        elif self.same_scale_op == 'max_pool':
+          out = nn_ops.max_pool(out, ksize=3, strides=1, padding='SAME')
+        elif self.same_scale_op == 'pad':
+          paddings = array_ops.ones(
+              (array_ops.rank(out), 2), dtype=dtypes.int32
+          )
+          out = array_ops.pad(out, paddings, 'CONSTANT')
+        elif self.same_scale_op == 'reshape':
+          out = array_ops.reshape(out, [-1])
+        elif self.same_scale_op == 'select':
+          rng = np.random.default_rng(seed=1234)
+          condition = ops.convert_to_tensor(
+              rng.uniform(low=0.0, high=1.0, size=out.shape) < 0.5
+          )
+          ones = array_ops.ones_like(out)
+          out = math_ops.select(condition, out, ones)
+        elif self.same_scale_op == 'slice':
+          begin = array_ops.zeros((array_ops.rank(out)), dtype=dtypes.int32)
+          size = array_ops.ones((array_ops.rank(out)), dtype=dtypes.int32)
+          out = array_ops.slice(out, begin, size)
+        elif self.same_scale_op == 'transpose':
+          out = array_ops.transpose(out)
+        else:
+          raise NotImplementedError(
+              '{} is not implemented for integration test.'.format(
+                  self.same_scale_op
+              )
+          )
+
+        return {'output': out}
+
+    model = MatmulAndSameScaleModel(weight_shape, same_scale_op)
+    saved_model_save.save(
+        model,
+        saved_model_path,
+        signatures=model.matmul_and_same_scale.get_concrete_function(
+            tensor_spec.TensorSpec(
+                shape=input_shape, dtype=dtypes.float32, name='input_tensor'
+            )
+        ),
+    )
+    return model
+
+  def _create_conv2d_model(
+      self,
+      input_shape: Sequence[int],
+      filter_shape: Sequence[int],
+      saved_model_path: str,
+      bias_fn: Optional[ops.Operation] = None,
+      activation_fn: Optional[ops.Operation] = None,
+      has_batch_norm: bool = False,
+      strides: Sequence[int] = (1, 1, 1, 1),
+      dilations: Sequence[int] = (1, 1, 1, 1),
+      padding: str = 'SAME',
+      has_func_alias: bool = False,
+  ) -> module.Module:
+    class ConvModel(module.Module):
+      """A simple model with a single conv2d, bias and relu."""
+
+      def __init__(self):
+        self.out_channel_size = filter_shape[-1]
+
+        # This ensures filters will have different value range per out channel
+        self.filters = np.stack(
+            [
+                np.random.uniform(
+                    low=-(i + 1), high=(i + 1), size=filter_shape[:-1]
+                ).astype('f4')
+                for i in range(self.out_channel_size)
+            ],
+            axis=-1,
+        )
+
+        self.bias = np.random.uniform(
+            low=0, high=10, size=(self.out_channel_size)
+        ).astype('f4')
+
+      @def_function.function
+      def conv2d(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
+        """Performs a 2D convolution operation.
+
+        Args:
+          input_tensor: Input tensor to perform convolution on.
+
+        Returns:
+          A map of: output key -> output result.
+        """
+        scale = [1.0] * self.out_channel_size
+        offset = [0.5] * self.out_channel_size
+        mean, variance = scale, offset
+        out = nn_ops.conv2d(
+            input_tensor,
+            self.filters,
+            strides=strides,
+            dilations=dilations,
+            padding=padding,
+            data_format='NHWC',
+            name='sample/conv',
+        )
+        if bias_fn is not None:
+          out = nn_ops.bias_add(out, self.bias)
+        if has_batch_norm:
+          # Fusing is supported for non-training case.
+          out, _, _, _, _, _ = nn_ops.fused_batch_norm_v3(
+              out, scale, offset, mean, variance, is_training=False
+          )
+        if activation_fn is not None:
+          out = activation_fn(out)
+        return {'output': out}
+
+    model = ConvModel()
+    save_options = None
+    if has_func_alias:
+      save_options = tensorflow.saved_model.SaveOptions(
+          function_aliases={FUNC_ALIAS: model.conv2d}
+      )
+    saved_model_save.save(
+        model,
+        saved_model_path,
+        signatures=model.conv2d.get_concrete_function(
+            tensor_spec.TensorSpec(
+                shape=input_shape, dtype=dtypes.float32, name='input_tensor'
+            )
+        ),
+        options=save_options,
+    )
+    return model
+
+  def _create_gather_model(self, input_type, use_variable) -> module.Module:
+    class GatherModel(module.Module):
+      """A simple model with a single gather."""
+
+      def __init__(self, use_variable):
+        """Initializes a GatherModel.
+
+        Args:
+          use_variable: If True, creates a variable for weight.
+        """
+        super().__init__()
+        w_val = np.random.randn(128, 32).astype('f4')
+        if use_variable:
+          self.w = variables.Variable(w_val)
+        else:
+          self.w = w_val
+
+      @def_function.function(
+          input_signature=[
+              tensor_spec.TensorSpec(
+                  shape=[6], dtype=input_type, name='input_tensor'
+              )
+          ]
+      )
+      def __call__(
+          self, input_tensor: core.Tensor
+      ) -> Mapping[str, core.Tensor]:
+        """Performs a gather operation."""
+        out = array_ops.gather_v2(self.w, input_tensor)
+        return {'output': out}
+
+    return GatherModel(use_variable)
+
+  def _create_add_model(
+      self,
+      shape: Sequence[int],
+      saved_model_path: str,
+  ) -> module.Module:
+    class AddModel(module.Module):
+      """A simple model with a single add."""
+
+      def __init__(self):
+        pass
+
+      @def_function.function
+      def add(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
+        """Performs an add operation.
+
+        Args:
+          input_tensor: Input tensor to perform add on.
+
+        Returns:
+          A map of: output key -> output result.
+        """
+        out = math_ops.add(input_tensor, input_tensor)
+        return {'output': out}
+
+    model = AddModel()
+    saved_model_save.save(
+        model,
+        saved_model_path,
+        signatures=model.add.get_concrete_function(
+            tensor_spec.TensorSpec(
+                shape=shape, dtype=dtypes.float32, name='input_tensor'
+            )
+        ),
+    )
+    return model
+
+  # Prepares sample einsum input data shapes.
+  # This function returns:
+  # 1. Shape for input 1
+  # 2. Shape for input 2
+  # 3. Shape for bias
+  # 4. Signature for input 1 (Could contain None dimension)
+  # 5. Signature for input 2 (Could contain None dimension)
+  def _prepare_sample_einsum_datashapes(
+      self,
+      equation: str,
+      generate_unknown_shape_signature: bool = False,
+      use_bias: bool = False,
+  ) -> Tuple[
+      List[Optional[int]],
+      List[Optional[int]],
+      Optional[List[Optional[int]]],
+      List[Optional[int]],
+      List[Optional[int]],
+  ]:
+    # 1. Parse equation.
+    comma_pos = equation.find(',')
+    arrow_pos = equation.find('->')
+    x_labels = equation[0:comma_pos]
+    y_labels = equation[comma_pos + 1 : arrow_pos]
+    out_labels = equation[arrow_pos + 1 :]
+
+    # 2. Create sample shapes.
+    label_to_size = {'a': 4, 'b': 32, 'c': 64, 'd': 128, 'e': 8}
+    x_shape = [label_to_size.get(x_label) for x_label in x_labels]
+    y_shape = [label_to_size.get(y_label) for y_label in y_labels]
+    bias_shape = None
+    if use_bias:
+      bias_shape = [label_to_size.get(out_label) for out_label in out_labels]
+      bias_shape = bias_shape[-1:]
+    contracting_dims = set()
+
+    x_signature = list(x_shape)
+    y_signature = list(y_shape)
+    if generate_unknown_shape_signature:
+      for c in x_labels:
+        if c in y_labels:
+          contracting_dims.add(c)
+      x_signature = [
+          None if c not in contracting_dims else x_shape[cidx]
+          for cidx, c in enumerate(x_labels)
+      ]
+      y_signature = [
+          None if c not in contracting_dims else y_shape[cidx]
+          for cidx, c in enumerate(y_labels)
+      ]
+    return x_shape, y_shape, bias_shape, x_signature, y_signature
+
+  def _create_einsum_model(
+      self,
+      saved_model_path: str,
+      equation: str,
+      y_shape: Sequence[int],
+      x_signature: Sequence[Optional[int]],
+      y_signature: Sequence[Optional[int]],
+      bias_shape: Optional[Sequence[int]] = None,
+  ) -> module.Module:
+    class EinsumModel(module.Module):
+      """Einsum class."""
+
+      def __init__(self):
+        self._bias = None
+        if bias_shape is not None:
+          self._bias = array_ops.constant(
+              np.random.uniform(size=bias_shape), dtype=dtypes.float32
+          )
+
+        self._kernel = np.random.uniform(size=y_shape).astype('f4')
+        self._min = (-0.8, -0.8, -0.9)
+        self._max = (0.9, 0.9, 1.0)
+
+      @def_function.function(
+          input_signature=[
+              tensor_spec.TensorSpec(
+                  name='x', shape=x_signature, dtype=dtypes.float32
+              )
+          ]
+      )
+      def einsum_with_kernel(self, x: core.Tensor) -> Mapping[str, core.Tensor]:
+        return self._einsum(x, self._kernel)
+
+      @def_function.function(
+          input_signature=[
+              tensor_spec.TensorSpec(
+                  name='x', shape=x_signature, dtype=dtypes.float32
+              ),
+              tensor_spec.TensorSpec(
+                  name='y', shape=y_signature, dtype=dtypes.float32
+              ),
+          ]
+      )
+      def einsum_without_kernel(
+          self, x: core.Tensor, y: core.Tensor
+      ) -> Mapping[str, core.Tensor]:
+        return self._einsum(x, y)
+
+      def _einsum(self, x, y):
+
+        out = tensorflow.einsum(equation, x, y)
+        if self._bias is not None:
+          out = nn_ops.bias_add(out, self._bias)
+        return {'output': out}
+
+    model = EinsumModel()
+    signatures = {
+        'serving_default': model.einsum_with_kernel.get_concrete_function(
+            tensor_spec.TensorSpec(
+                name='x', shape=x_signature, dtype=dtypes.float32
+            )
+        ),
+    }
+    saved_model_save.save(model, saved_model_path, signatures=signatures)
+    return model

tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/stablehlo_quantizer_odml_oss.ipynb

@@ -0,0 +1,192 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tWhm0JFMPJ5I"
+      },
+      "source": [
+        "Copyright 2024 Google LLC.\n",
+        "\n",
+        "Licensed under the Apache License, Version 2.0 (the \"License\");"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RJcqTAlfPQjk"
+      },
+      "source": [
+        "# [OSS] JAX to TFLite with StableHLO Quantization Demonstration for ODML."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cqeGmbO6PPNd"
+      },
+      "source": [
+        "This example shows a JAX Keras reference model converted into a StableHLO module and via `jax2tf`, then quantized in the ODML Converter via the StableHLO Quantizer.\n",
+        "\n",
+        "Note: This API is experimental and will likely have breakages with other models. Please reach out to [scalable-opt-team@google.com](mailto:scalable-opt-team@google.com) and we will support your use case."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-S0P42BpPSeJ"
+      },
+      "source": [
+        "## StableHLO Quantizer\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FacwMD9MPUew"
+      },
+      "source": [
+        "StableHLO Quantizer is a quantization API to enable ML framework optionality and hardware retargetability."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "RXZUHZQoQZOo"
+      },
+      "outputs": [],
+      "source": [
+        "!pip uninstall tensorflow --yes"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aYz36YEKPYRk"
+      },
+      "outputs": [],
+      "source": [
+        "!pip3 install tf-nightly\n",
+        "!pip3 install keras-core"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "duab6P-nPZzF"
+      },
+      "outputs": [],
+      "source": [
+        "import tensorflow as tf\n",
+        "print(\"TensorFlow version:\", tf.__version__)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c9JX9RJTPaoW"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "os.environ['KERAS_BACKEND'] = 'jax'\n",
+        "import jax.numpy as jnp\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "from keras_core.applications import ResNet50\n",
+        "from jax.experimental import jax2tf"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "rTcHwDPBPchd"
+      },
+      "outputs": [],
+      "source": [
+        "input_shape = (1, 224, 224, 3)\n",
+        "\n",
+        "jax_callable = jax2tf.convert(\n",
+        "    ResNet50(\n",
+        "      input_shape=input_shape[1:],\n",
+        "      pooling='avg',\n",
+        "  ).call,\n",
+        "    with_gradient=False,\n",
+        "    native_serialization=True,\n",
+        "    native_serialization_platforms=('cpu',))\n",
+        "\n",
+        "tf_module = tf.Module()\n",
+        "tf_module.f = tf.function(\n",
+        "    jax_callable,\n",
+        "    autograph=False,\n",
+        "    input_signature=[\n",
+        "        tf.TensorSpec(input_shape, jnp.float32, 'lhs_operand')\n",
+        "    ],\n",
+        ")\n",
+        "\n",
+        "saved_model_dir = '/tmp/saved_model'\n",
+        "tf.saved_model.save(tf_module, saved_model_dir)\n",
+        "\n",
+        "def calibration_dataset():\n",
+        "  rng = np.random.default_rng(seed=1235)\n",
+        "  for _ in range(2):\n",
+        "    yield {\n",
+        "        'lhs_operand': rng.uniform(low=-1.0, high=1.0, size=input_shape).astype(\n",
+        "            np.float32\n",
+        "        )\n",
+        "    }\n",
+        "converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)\n",
+        "converter.target_spec.supported_ops = [\n",
+        "    tf.lite.OpsSet.SELECT_TF_OPS,  # enable TensorFlow ops.\n",
+        "    tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TFL ops.\n",
+        "]\n",
+        "converter.representative_dataset = calibration_dataset\n",
+        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
+        "# Below flag controls whether to use StableHLO Quantizer or TFLite quantizer.\n",
+        "converter.experimental_use_stablehlo_quantizer = True\n",
+        "\n",
+        "quantized_model = converter.convert()\n",
+        "\n",
+        "with open('/tmp/resnet50_quantized.tflite', 'wb') as f:\n",
+        "  f.write(quantized_model)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "u3b9Xj8dPdXo"
+      },
+      "outputs": [],
+      "source": [
+        "print(str(os.path.getsize('/tmp/resnet50_quantized.tflite') >> 20) + 'MB')"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "private_outputs": true,
+      "provenance": [
+        {
+          "file_id": "https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/mlir/quantization/stablehlo/python/integration_test/stablehlo_quantizer_odml_oss.ipynb",
+          "timestamp": 1712841250910
+        }
+      ]
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization.cc

@@ -0,0 +1,112 @@
+/* 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.
+==============================================================================*/
+#include "pybind11/cast.h"  // from @pybind11
+#include "pybind11/detail/common.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11  // IWYU pragma: keep
+#include "pybind11_abseil/absl_casters.h"  // from @pybind11_abseil   // IWYU pragma: keep
+#include "pybind11_abseil/import_status_module.h"  // from @pybind11_abseil
+#include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/type_casters.h"  // IWYU pragma: keep
+
+namespace py = pybind11;
+
+namespace {
+
+using ::stablehlo::quantization::pywrap::PywrapExpandPresets;
+using ::stablehlo::quantization::pywrap::PywrapPopulateDefaults;
+using ::stablehlo::quantization::pywrap::PywrapQuantizeStaticRangePtq;
+using ::stablehlo::quantization::pywrap::PywrapQuantizeWeightOnlyPtq;
+
+}  // namespace
+
+PYBIND11_MODULE(pywrap_quantization, m) {
+  // Supports absl::Status type conversions.
+  pybind11::google::ImportStatusModule();
+
+  m.doc() = "StableHLO Quantization APIs.";
+
+  // If the function signature changes, likely its corresponding .pyi type
+  // hinting should also change.
+  // LINT.IfChange(static_range_ptq)
+  m.def("static_range_ptq", &PywrapQuantizeStaticRangePtq,
+        R"pbdoc(
+        Runs static-range post-training quantization (PTQ) on a SavedModel at
+        `src_saved_model_path` and saves the resulting model to
+        `dst_saved_model_path`.
+
+        The user should pass a serialized `QuantizationConfig` for the
+        `quantization_config_serialized` argument, and a signature key ->
+        serialized `SignatureDef` mapping for the `signature_def_map_serialized`
+        argument.
+
+        Raises `StatusNotOk` exception if when the run was unsuccessful.
+        )pbdoc",
+        py::arg("src_saved_model_path"), py::arg("dst_saved_model_path"),
+        py::arg("quantization_config_serialized"), py::kw_only(),
+        py::arg("signature_keys"), py::arg("signature_def_map_serialized"),
+        py::arg("py_function_library"));
+  // LINT.ThenChange(pywrap_quantization.pyi:static_range_ptq)
+
+  // If the function signature changes, likely its corresponding .pyi type
+  // hinting should also change.
+  // LINT.IfChange(weight_only_ptq)
+  m.def("weight_only_ptq", &PywrapQuantizeWeightOnlyPtq,
+        R"pbdoc(
+        Runs weight-only Quantization on a SavedModel at `src_saved_model_path`
+        and saves the resulting model to `dst_saved_model_path`.
+
+        The user should pass a serialized `QuantizationConfig` for the
+        `quantization_config_serialized` argument, and a signature key ->
+        serialized `SignatureDef` mapping for the `signature_def_map_serialized`
+        argument.
+
+        Raises `StatusNotOk` exception if when the run was unsuccessful.
+        )pbdoc",
+        py::arg("src_saved_model_path"), py::arg("dst_saved_model_path"),
+        py::arg("quantization_config_serialized"), py::kw_only(),
+        py::arg("signature_keys"), py::arg("signature_def_map_serialized"),
+        py::arg("py_function_library"));
+  // LINT.ThenChange(pywrap_quantization.pyi:weight_only_ptq)
+
+  // If the function signature changes, likely its corresponding .pyi type
+  // hinting should also change.
+  // LINT.IfChange(populate_default_configs)
+  m.def("populate_default_configs", &PywrapPopulateDefaults,
+        R"pbdoc(
+        Populates `QuantizationConfig` with default values.
+
+        Returns an updated `QuantizationConfig` (serialized) after populating
+        default values to fields that the user did not explicitly specify.
+        )pbdoc",
+        py::arg("user_provided_config_serialized"));
+  // LINT.ThenChange(pywrap_quantization.pyi:populate_default_configs)
+
+  // If the function signature changes, likely its corresponding .pyi type
+  // hinting should also change.
+  // LINT.IfChange(expand_preset_configs)
+  m.def("expand_preset_configs", &PywrapExpandPresets, R"pbdoc(
+        Expands presets to other fields in `QuantizationConfig`.
+
+        Each preset is expressible by other fields in `QuantizationConfig`.
+        Returns a copy of `QuantizationConfig` (serialized) where the fields are
+        expanded from presets. If no preset has been set, it is a no-op and
+        returns the same copy of the input.
+        )pbdoc",
+        py::arg("quantization_config_serialized"));
+  // LINT.ThenChange(pywrap_quantization.pyi:expand_preset_configs)
+}

tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization.pyi

@@ -0,0 +1,58 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from typing import Any
+
+from tensorflow.compiler.mlir.quantization.tensorflow.python import py_function_lib
+from tensorflow.compiler.mlir.quantization.tensorflow.python import representative_dataset as rd
+
+# LINT.IfChange(static_range_ptq)
+def static_range_ptq(
+    src_saved_model_path: str,
+    dst_saved_model_path: str,
+    quantization_config_serialized: bytes,
+    *,
+    signature_keys: list[str],
+    signature_def_map_serialized: dict[str, bytes],
+    py_function_library: py_function_lib.PyFunctionLibrary,
+) -> Any: ...  # Status
+
+# LINT.ThenChange()
+
+# LINT.IfChange(weight_only_ptq)
+def weight_only_ptq(
+    src_saved_model_path: str,
+    dst_saved_model_path: str,
+    quantization_config_serialized: bytes,
+    *,
+    signature_keys: list[str],
+    signature_def_map_serialized: dict[str, bytes],
+    py_function_library: py_function_lib.PyFunctionLibrary,
+) -> Any: ...  # Status
+
+# LINT.ThenChange()
+
+# LINT.IfChange(populate_default_configs)
+def populate_default_configs(
+    user_provided_quantization_config_serialized: bytes,
+) -> bytes: ...  # QuantizationConfig
+
+# LINT.ThenChange()
+
+# LINT.IfChange(expand_preset_configs)
+def expand_preset_configs(
+    quantization_config_serialized: bytes,
+) -> bytes: ...  # QuantizationConfig
+
+# LINT.ThenChange()

tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.cc

@@ -0,0 +1,72 @@
+/* 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/quantization/stablehlo/python/pywrap_quantization_lib.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/config.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/static_range_ptq.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/cc/weight_only_ptq.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+
+namespace stablehlo::quantization::pywrap {
+
+using ::mlir::quant::stablehlo::QuantizeStaticRangePtq;
+using ::mlir::quant::stablehlo::QuantizeWeightOnlyPtq;
+using ::tensorflow::SignatureDef;
+using ::tensorflow::quantization::PyFunctionLibrary;
+
+// Note for maintainers: the definitions should ONLY mirror existing functions
+// defined in different targets. Do not include any extra business logic that
+// causes divergence from the semantics of mirrored functions.
+
+absl::Status PywrapQuantizeStaticRangePtq(
+    absl::string_view src_saved_model_path,
+    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
+    const std::vector<std::string>& signature_keys,
+    const absl::flat_hash_map<std::string, SignatureDef>& signature_def_map,
+    const PyFunctionLibrary& py_function_library) {
+  return QuantizeStaticRangePtq(src_saved_model_path, dst_saved_model_path,
+                                config, signature_keys, signature_def_map,
+                                py_function_library);
+}
+
+absl::Status PywrapQuantizeWeightOnlyPtq(
+    absl::string_view src_saved_model_path,
+    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
+    const std::vector<std::string>& signature_keys,
+    const absl::flat_hash_map<std::string, SignatureDef>& signature_def_map,
+    const PyFunctionLibrary& py_function_library) {
+  return QuantizeWeightOnlyPtq(src_saved_model_path, dst_saved_model_path,
+                               config, signature_keys, signature_def_map,
+                               py_function_library);
+}
+
+QuantizationConfig PywrapPopulateDefaults(
+    const QuantizationConfig& user_provided_config) {
+  return PopulateDefaults(user_provided_config);
+}
+
+QuantizationConfig PywrapExpandPresets(const QuantizationConfig& config) {
+  return ExpandPresets(config);
+}
+
+}  // namespace stablehlo::quantization::pywrap

tensorflow/compiler/mlir/quantization/stablehlo/python/pywrap_quantization_lib.h

@@ -0,0 +1,64 @@
+/* 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_QUANTIZATION_STABLEHLO_PYTHON_PYWRAP_QUANTIZATION_LIB_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PYTHON_PYWRAP_QUANTIZATION_LIB_H_
+
+// Contains mirror functions from StableHLO Quantizer to be exposed to python
+// via `pywrap_quantization`.
+
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/quantization/stablehlo/quantization_config.pb.h"
+#include "tensorflow/compiler/mlir/quantization/tensorflow/python/py_function_lib.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+
+namespace stablehlo::quantization::pywrap {
+
+// Function used by the pywrap_quantization module to mirror
+// `::mlir::quant::stablehlo::QuantizeStaticRangePtq`.
+absl::Status PywrapQuantizeStaticRangePtq(
+    absl::string_view src_saved_model_path,
+    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
+    const std::vector<std::string>& signature_keys,
+    const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
+        signature_def_map,
+    const tensorflow::quantization::PyFunctionLibrary& py_function_library);
+
+// Function used by the pywrap_quantization module to mirror
+// `::mlir::quant::stablehlo::QuantizeWeightOnlyPtq`.
+absl::Status PywrapQuantizeWeightOnlyPtq(
+    absl::string_view src_saved_model_path,
+    absl::string_view dst_saved_model_path, const QuantizationConfig& config,
+    const std::vector<std::string>& signature_keys,
+    const absl::flat_hash_map<std::string, tensorflow::SignatureDef>&
+        signature_def_map,
+    const tensorflow::quantization::PyFunctionLibrary& py_function_library);
+
+// Function used by the pywrap_quantization module to mirror
+// `::stablehlo::quantization::PopulateDefaults`.
+QuantizationConfig PywrapPopulateDefaults(
+    const QuantizationConfig& user_provided_config);
+
+// Function used by the pywrap_quantization module to mirror
+// `::stablehlo::quantization::ExpandPresets`.
+QuantizationConfig PywrapExpandPresets(const QuantizationConfig& config);
+
+}  // namespace stablehlo::quantization::pywrap
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_PYTHON_PYWRAP_QUANTIZATION_LIB_H_

tensorflow/compiler/mlir/quantization/stablehlo/python/quantization.py

@@ -0,0 +1,121 @@
+# 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.
+# ==============================================================================
+"""StableHLO Quantizer."""
+from typing import Mapping
+
+from tensorflow.compiler.mlir.quantization.stablehlo import quantization_config_pb2 as qc
+from tensorflow.compiler.mlir.quantization.stablehlo.python import pywrap_quantization
+from tensorflow.compiler.mlir.quantization.tensorflow.python import py_function_lib
+from tensorflow.compiler.mlir.quantization.tensorflow.python import save_model
+from tensorflow.core.protobuf import meta_graph_pb2
+
+# Mapping of signature def key -> SignatureDef.
+_SignatureDefMap = Mapping[str, meta_graph_pb2.SignatureDef]
+
+
+def _serialize_signature_def_map(
+    signature_def_map: _SignatureDefMap,
+) -> dict[str, bytes]:
+  """Serializes SignatureDef values in `signature_def_map`.
+
+  Args:
+    signature_def_map: Signature key -> SignatureDef mapping.
+
+  Returns:
+    Signature def map where the values (`SignatureDef`) are serialized.
+  """
+  signature_def_map_serialized = {}
+  for key, signature_def in signature_def_map.items():
+    signature_def_map_serialized[key] = signature_def.SerializeToString()
+
+  return signature_def_map_serialized
+
+
+def _has_quantization_method(
+    quantization_specs: qc.QuantizationSpecs, method: str
+) -> bool:
+  """Returns whether a given QuantizationSpecs has the given quantization method."""
+  for spec in quantization_specs.specs:
+    if spec.method.HasField(method):
+      return True
+  return False
+
+
+# TODO: b/310594193 - Export API to pip package.
+def quantize_saved_model(
+    src_saved_model_path: str,
+    dst_saved_model_path: str,
+    config: qc.QuantizationConfig,
+) -> None:
+  """Quantizes a saved model.
+
+  Args:
+    src_saved_model_path: Path to the directory for the source SavedModel.
+    dst_saved_model_path: Path to the directory for the destination SavedModel.
+    config: Quantization configuration.
+
+  Raises:
+    ValueError: When `config` was not configured for static-range PTQ
+    single representative dataset.
+  """
+  # Updates user-provided `QuantizationConfig`s for the internal quantization
+  # pipeline to work with.
+  print('=== User-provided QuantizationConfig ===')
+  print(config)
+  config = qc.QuantizationConfig.FromString(
+      pywrap_quantization.populate_default_configs(config.SerializeToString())
+  )
+  config = qc.QuantizationConfig.FromString(
+      pywrap_quantization.expand_preset_configs(config.SerializeToString())
+  )
+  print('=== Updated QuantizationConfig ===')
+  print(config)
+
+  if not (
+      _has_quantization_method(config.specs, 'static_range_ptq')
+      and len(config.calibration_options.representative_datasets) == 1
+  ) and not _has_quantization_method(config.specs, 'weight_only_ptq'):
+    raise ValueError(
+        '`quantize_saved_model` currently only supports static-range PTQ with a'
+        ' single signature or weight-only quantization.'
+    )
+
+  signature_def_map = save_model.get_signatures_from_saved_model(
+      src_saved_model_path,
+      signature_keys=None,
+      tags=set(config.tf_saved_model.tags),
+  )
+
+  signature_def_map_serialized = _serialize_signature_def_map(signature_def_map)
+  # Currently, only StaticRangePtq or WeightOnlyPtq is supported.
+  # Consider merging the pipelines to address mixed algorithm models.
+  if _has_quantization_method(config.specs, 'static_range_ptq'):
+    pywrap_quantization.static_range_ptq(
+        src_saved_model_path,
+        dst_saved_model_path,
+        quantization_config_serialized=config.SerializeToString(),
+        signature_keys=list(signature_def_map.keys()),
+        signature_def_map_serialized=signature_def_map_serialized,
+        py_function_library=py_function_lib.PyFunctionLibrary(),
+    )
+  elif _has_quantization_method(config.specs, 'weight_only_ptq'):
+    pywrap_quantization.weight_only_ptq(
+        src_saved_model_path,
+        dst_saved_model_path,
+        quantization_config_serialized=config.SerializeToString(),
+        signature_keys=list(signature_def_map.keys()),
+        signature_def_map_serialized=signature_def_map_serialized,
+        py_function_library=py_function_lib.PyFunctionLibrary(),
+    )