ai-edge-quantizer-nightly 0.3.0.dev20250613__py3-none-any.whl → 0.3.0.dev20250615__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -112,6 +112,7 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
         common_quantize.materialize_squared_difference
     ),
     _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
+    _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -250,6 +251,7 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
         common_quantize.materialize_squared_difference
     ),
     _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
+    _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -728,6 +728,23 @@ def materialize_max_pool_2d(
   )
 
 
+def materialize_resize_bilinear(
+    get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: dict[str, Any],
+) -> list[qtyping.TensorTransformationParams]:
+  """Materialize tensors in tfl.resize_bilinear."""
+  return common_utils.materialize_standard_op(
+      op_info,
+      graph_info,
+      tensor_name_to_qsv,
+      get_tensor_quant_params_fn,
+      constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
+      inputs_to_ignore=[1],  # Resize size does not need to be quantized.
+  )
+
+
 def _get_tensor_shape_for_blockwise(
     tensor_shape: Sequence[int], quantized_dim: int, block_size: int
 ) -> list[int]:

ai_edge_quantizer/default_policy.py

@@ -185,7 +185,8 @@ DEFAULT_JSON_POLICY = """
       "SELECT_V2",
       "STABLEHLO_COMPOSITE",
       "PAD",
-      "MAX_POOL_2D"
+      "MAX_POOL_2D",
+      "RESIZE_BILINEAR"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -219,7 +220,8 @@ DEFAULT_JSON_POLICY = """
       "STABLEHLO_COMPOSITE",
       "PAD",
       "SQUARED_DIFFERENCE",
-      "MAX_POOL_2D"
+      "MAX_POOL_2D",
+      "RESIZE_BILINEAR"
     ],
     "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],
     "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],

ai_edge_quantizer/qtyping.py

@@ -65,6 +65,7 @@ class TFLOperationName(str, enum.Enum):
   PAD = 'PAD'
   SQUARED_DIFFERENCE = 'SQUARED_DIFFERENCE'
   MAX_POOL_2D = 'MAX_POOL_2D'
+  RESIZE_BILINEAR = 'RESIZE_BILINEAR'
 
 
 class QuantizeMode(enum.Enum):

ai_edge_quantizer/recipe_manager_test.py

@@ -29,19 +29,6 @@ _AlgorithmName = recipe_manager.AlgorithmName
 _QuantGranularity = qtyping.QuantGranularity
 
 
-# Sample functions for test cases.
-def _sample_init_qsvs(*_, **__):
-  return 1.0, dict()
-
-
-def _sample_calibration_func(*_, **__):
-  return 2.0, dict()
-
-
-def _sample_materialize_func(*_, **__):
-  return 3.0, dict()
-
-
 def _sample_check_op_config_func(op_name, op_config, _):
   if (
       op_config.weight_tensor_config is not None
@@ -67,6 +54,16 @@ def _add_default_int8xint8_integer_recipe(recipe_manager_object):
 
 # register some currently unsupported ops for testing purposes
 def _register_testing_op(algorithm_key, tfl_op):
+  # Sample functions for test cases.
+  def _sample_init_qsvs(*_, **__):
+    return {'name': dict()}
+
+  def _sample_calibration_func(*_, **__):
+    return {'name2': dict()}
+
+  def _sample_materialize_func(*_, **__):
+    return []
+
   algorithm_manager.register_op_quant_config_validation_func(
       algorithm_key, _sample_check_op_config_func
   )

ai_edge_quantizer/utils/calibration_utils.py

@@ -15,9 +15,26 @@
 
 """Utilities for model calibration."""
 
-from typing import Union
+import copy
+from typing import Any, Union
+
 import numpy as np
+
+from ai_edge_quantizer import algorithm_manager
 from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+from ai_edge_litert import schema_py_generated as schema_fb  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.lite.tools import flatbuffer_utils  # pylint: disable=g-direct-tensorflow-import
+
+
+_SignatureInput = dict[str, Any]
+_OpQuantConstraint = common_utils.OpQuantConstraint
+_SignatureData = dict[
+    str, list[str]
+]  # signature_key -> list of signature_names.
 
 
 def _update_moving_average(
@@ -84,3 +101,327 @@ def min_max_update(qsv: qtyping.QSV, new_qsv: qtyping.QSV) -> qtyping.QSV:
   updated_qsv["min"] = np.minimum(qsv["min"], new_qsv["min"])
   updated_qsv["max"] = np.maximum(qsv["max"], new_qsv["max"])
   return updated_qsv
+
+
+def _find_overall_min_max(
+    qsv: qtyping.QSV, tensor_names: list[str]
+) -> tuple[np.ndarray, np.ndarray]:
+  """Finds the overall minimum and maximum values for the given tensors.
+
+  Args:
+    qsv: The quantization statistical value of the tensor (min/max).
+    tensor_names: The list of tensor names to find the minimum and maximum
+      values.
+
+  Returns:
+    The minimum and maximum values for the given tensors.
+  """
+  min_value = np.inf
+  max_value = -np.inf
+  for tensor_name in tensor_names:
+    min_value = min(min_value, qsv[tensor_name]["min"])
+    max_value = max(max_value, qsv[tensor_name]["max"])
+  return min_value, max_value
+
+
+class CalibrationQsvAlignmentUtils:
+  """Calibration utils for alignment of QSVs.
+
+  This class is used to align QSVs for a given model. It builds a list of ops
+  that need to be constrained to the same scale as the input. Based on this
+  list, it finds the corresponding tensor names for a given signature data.
+  """
+
+  def __init__(self, model_path: str):
+    self._same_as_input_scale_ops = []
+
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(model_path)
+    self._flatbuffer_object = tfl_flatbuffer_utils.read_model(model_path)
+
+    signature_keys = list(tfl_interpreter.get_signature_list().keys())
+
+    # Build a dict of signature runners.
+    self._signature_runners = {}
+    for signature_key in signature_keys:
+      signature_runner = tfl_interpreter.get_signature_runner(signature_key)
+      self._signature_runners[signature_key] = signature_runner
+
+    # Make a list of `SAME_AS_INPUT_SCALE` operators. This is used to identify
+    # the operators that need to be constrained to the same scale as the input.
+    self._build_same_as_input_scale_op_list()
+
+  def _build_same_as_input_scale_op_list(self, verbose: bool = False):
+    """Constructs a list of SAME_AS_INPUT_SCALE operators.
+
+    This is achieved by invoking all materialization functions and extracting
+    the constraint argument, using monkey patching to redirect logic to wrapper
+    functions.
+
+    Args:
+      verbose: Flag to enable verbose output.
+    """
+
+    def materialize_standard_op_wrapper(
+        op_info: qtyping.OpInfo,
+        *_args,
+        constraint: _OpQuantConstraint = _OpQuantConstraint.NO_CONSTRAIN,
+        **_kwargs,
+    ) -> list[qtyping.TensorTransformationParams]:
+      if constraint == _OpQuantConstraint.SAME_AS_INPUT_SCALE:
+        self._same_as_input_scale_ops.append(op_info.op_name)
+      # Return dummy values to avoid exceptions.
+      dummy_value = [qtyping.TensorTransformationParams("")] * 2
+      return dummy_value
+
+    # Dummy implementation of the `_are_weights_too_small` function to support
+    # `materialize_standard_op_wrapper` above.
+    def are_weights_too_small_wrapper(*_args, **_kwargs) -> bool:
+      return False
+
+    # Dummy implementation of the `_materialize_bias_for_conv_ops` function to
+    # support `materialize_standard_op_wrapper` above.
+    def materialize_bias_for_conv_ops_wrapper(*_args, **_kwargs):
+      return
+
+    # Do monkey patch to intercept the `materialize_standard_op` function to
+    # support `materialize_standard_op_wrapper` above.
+    original_materialize_standard_op = common_utils.materialize_standard_op
+    original_are_weights_too_small = common_quantize._are_weights_too_small  # pylint: disable=protected-access
+    original_materialize_bias_for_conv_ops = (
+        common_quantize._materialize_bias_for_conv_ops  # pylint: disable=protected-access
+    )
+    common_utils.materialize_standard_op = materialize_standard_op_wrapper
+    common_quantize._are_weights_too_small = are_weights_too_small_wrapper  # pylint: disable=protected-access
+    common_quantize._materialize_bias_for_conv_ops = (  # pylint: disable=protected-access
+        materialize_bias_for_conv_ops_wrapper
+    )
+    minmax_func_dict = algorithm_manager.MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT
+
+    # Loop over all available materialization functions to build up a list of
+    # `SAME_AS_INPUT_SCALE` constrained ops.
+    for op, materialize_fn in minmax_func_dict.items():
+      # Create a dummy op info to trigger the materialization.
+      mock_op = schema_fb.OperatorT()
+      mock_op.inputs = [0]
+      mock_op.outputs = [0]
+      op_info = qtyping.OpInfo(
+          op=mock_op,
+          op_name=op,
+          subgraph_op_index=0,
+          op_quant_config=qtyping.OpQuantizationConfig(),
+      )
+      materialize_fn(
+          get_tensor_quant_params_fn=None,
+          op_info=op_info,
+          graph_info=None,
+          tensor_name_to_qsv=None,
+      )
+
+    if verbose:
+      print(f"  Constrained op list: {self._same_as_input_scale_ops}")
+
+    # Restore the original functions.
+    common_utils.materialize_standard_op = original_materialize_standard_op
+    common_quantize._are_weights_too_small = original_are_weights_too_small  # pylint: disable=protected-access
+    common_quantize._materialize_bias_for_conv_ops = (  # pylint: disable=protected-access
+        original_materialize_bias_for_conv_ops
+    )
+
+  def _search_tensor_by_signature_name(
+      self, signature_key: str, signature_input_output_name: str, verbose=False
+  ) -> list[str]:
+    """Searches for a tensor name for a given signature by signature input or output name.
+
+    Args:
+      signature_key: Name of the signature.
+      signature_input_output_name: Name of the input or output in the signature.
+      verbose: Flag to enable verbose output.
+
+    Returns:
+      The list with one or two tensor names. The first one is the input tensor
+      name, and the second one is the output tensor name.
+    """
+
+    if verbose:
+      print("Searching tensor by signature name.")
+
+    tensor_names = []
+
+    # Search among inputs.
+    input_details = self._signature_runners[signature_key].get_input_details()
+    if signature_input_output_name in input_details.keys():
+      tensor_names.append(input_details[signature_input_output_name]["name"])
+
+    # Search among outputs.
+    output_details = self._signature_runners[signature_key].get_output_details()
+    if signature_input_output_name not in output_details:
+      if not tensor_names:
+        raise ValueError(
+            f"Signature {signature_key} does not have input or output"
+            f" `{signature_input_output_name}`"
+        )
+      return tensor_names
+
+    output_tensor_name = output_details[signature_input_output_name]["name"]
+    if verbose:
+      print(
+          ">> Starting recursive search for the output tensor name:"
+          f" {output_tensor_name}"
+      )
+
+    idx = self._signature_runners[signature_key]._subgraph_index  # pylint: disable=protected-access
+    subgraph = self._flatbuffer_object.subgraphs[idx]
+    graph_info = qtyping.GraphInfo(
+        subgraph.tensors, self._flatbuffer_object.buffers
+    )
+
+    # Recursively search the graph for the output tensor name since it may be
+    # `SAME_AS_INPUT` constrainted.
+    operators = copy.deepcopy(subgraph.operators)
+    tensor_name = self._search_reverse_order_recursively(
+        graph_info, operators, output_tensor_name, indent="  ", verbose=verbose
+    )
+    tensor_names.append(tensor_name)
+
+    if verbose:
+      print(f"\n\nFound tensor name: {tensor_name}")
+
+    return tensor_names
+
+  def _search_reverse_order_recursively(
+      self,
+      graph_info: qtyping.GraphInfo,
+      operators: list[Any],
+      output_tensor_name: str,
+      indent: str,
+      verbose: bool = False,
+  ):
+    """Searches for a tensor name in reverse order recursively.
+
+    Stop criteria is when the tensor belongs to an operator that is not
+    `SAME_AS_INPUT` constrainted.
+
+    Args:
+      graph_info: Graph information.
+      operators: List of operators.
+      output_tensor_name: Name of the output tensor to search for.
+      indent: Indentation string for debug output.
+      verbose: Flag to enable verbose output.
+
+    Returns:
+      The name of the tensor found, or None if not found.
+    """
+    op_codes = self._flatbuffer_object.operatorCodes
+
+    while operators:
+      op = operators.pop()
+      op_code = op_codes[op.opcodeIndex].builtinCode
+      op_name = flatbuffer_utils.opcode_to_name(
+          self._flatbuffer_object, op.opcodeIndex
+      )
+      if op_code not in tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME:
+        continue
+      for output_idx in op.outputs:
+        if output_tensor_name == tfl_flatbuffer_utils.get_tensor_name(
+            graph_info.subgraph_tensors[output_idx]
+        ):
+          dbg_str = (
+              f"{indent}>> Found `{op_name}`, output tensor"
+              f" '{output_tensor_name}'"
+          )
+
+          if op_name not in self._same_as_input_scale_ops:
+            if verbose:
+              print(f"{dbg_str}, returning...")
+            return output_tensor_name
+
+          if verbose:
+            print(f"{dbg_str}, with SAME_AS_INPUT, search recursively among:")
+          for input_idx in op.inputs:
+            input_tensor_name = graph_info.subgraph_tensors[
+                input_idx
+            ].name.decode("utf-8")
+
+            if verbose:
+              print(f"{indent}    Input: {input_tensor_name}")
+
+            return self._search_reverse_order_recursively(
+                graph_info,
+                operators,
+                input_tensor_name,
+                indent=f"{indent}  ",
+                verbose=verbose,
+            )
+    return output_tensor_name
+
+  def align_quant_stats(
+      self, qsv: dict[str, Any], signature_data: _SignatureData
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Aligns quantization statistics for a given signature data.
+
+    This function takes quantization statistics and signature data as input,
+    identifies the tensors associated with the signature data, and aligns
+    the quantization statistics of these tensors by setting their minimum
+    and maximum values to the same value. This ensures that the tensors
+    have the same quantization parameters.
+
+    Args:
+      qsv: Quantization statistics.
+      signature_data: Signature data.
+
+    Returns:
+      Tuple of min and max values.
+    """
+    # Go over all signature info and find the corresponding tensor names.
+    tensor_names = []
+    for signature_key, signature_names in signature_data.items():
+      for signature_name in signature_names:
+        tensor_name = self._search_tensor_by_signature_name(
+            signature_key, signature_name
+        )
+        tensor_names.extend(tensor_name)
+
+    # Find min and max values accross all tensors.
+    min_value, max_value = _find_overall_min_max(qsv, tensor_names)
+
+    # Overwrite the min and max values in the QSV.
+    for tensor_name in tensor_names:
+      qsv[tensor_name]["min"] = min_value
+      qsv[tensor_name]["max"] = max_value
+
+    return min_value, max_value
+
+  def update_quant_stats(
+      self,
+      qsv: dict[str, Any],
+      signature_data: _SignatureData,
+      min_value: np.ndarray,
+      max_value: np.ndarray,
+  ):
+    """Updates quantization statistics for a given signature data.
+
+    This function updates the quantization statistics with the provided min, max
+    values for the tensors specified in the signature data.
+
+    Args:
+      qsv: Quantization statistics.
+      signature_data: Signature data.
+      min_value: Minimum value to update.
+      max_value: Maximum value to update.
+
+    Returns:
+      Updated quantization statistics.
+    """
+    # Go over all signature info and find the corresponding tensor names.
+    tensor_names = []
+    for signature_key, signature_names in signature_data.items():
+      for signature_name in signature_names:
+        tensor_name = self._search_tensor_by_signature_name(
+            signature_key, signature_name
+        )
+        tensor_names.extend(tensor_name)
+
+    # Overwrite the min and max values in the QSV.
+    for tensor_name in tensor_names:
+      qsv[tensor_name]["min"] = min_value
+      qsv[tensor_name]["max"] = max_value

ai_edge_quantizer/utils/calibration_utils_test.py

@@ -14,11 +14,68 @@
 # ==============================================================================
 
 from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
 from tensorflow.python.platform import googletest
+from ai_edge_quantizer import quantizer
 from ai_edge_quantizer.utils import calibration_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+_RNG = np.random.default_rng(66)
+
+_CALIBRATION_DATASET = {
+    "signature_1": [{
+        "cache_0": np.zeros(shape=(1, 100, 4, 4), dtype=np.float32),
+        "cache_1": np.zeros(shape=(1, 100, 4, 4), dtype=np.float32),
+        "positions": np.zeros(shape=(1, 100), dtype=np.int32),
+        "tokens": np.zeros(shape=(1, 100), dtype=np.int32),
+    }],
+    "signature_2": [{
+        "cache_0": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+        "cache_1": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+        "positions": (
+            _RNG.integers(low=0, high=10, size=(1, 100)).astype(np.int32)
+        ),
+        "tokens": _RNG.integers(low=0, high=10, size=(1, 100)).astype(np.int32),
+    }],
+}
+
 
+def _get_quant_parameters(
+    quantized_model: bytes, signature_data: dict[str, list[str]]
+) -> list[np.ndarray]:
+  """Returns the quantization parameters from the quantized model."""
+  quant_params = []
+  tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+      quantized_model
+  )
+  for signature_key, signature_names in signature_data.items():
+    signature_runner = tfl_interpreter.get_signature_runner(signature_key)
+
+    for signature_name in signature_names:
+      input_details = signature_runner.get_input_details()
+      output_details = signature_runner.get_output_details()
+      if signature_name in input_details.keys():
+        quant_param = input_details[signature_name]["quantization_parameters"][
+            "scales"
+        ].squeeze()
+        quant_params.append(quant_param)
+      elif signature_name in output_details.keys():
+        output_details = signature_runner.get_output_details()
+        quant_param = output_details[signature_name]["quantization_parameters"][
+            "scales"
+        ].squeeze()
+        quant_params.append(quant_param)
+      else:
+        raise ValueError(
+            f"Signature name {signature_name} not found in the model."
+        )
+  return quant_params
 
-class CalibrationUtilsTest(parameterized.TestCase):
+
+class CalibrationQsvAlignmentUtilsTest(parameterized.TestCase):
 
   @parameterized.named_parameters(
       dict(
@@ -66,12 +123,126 @@ class CalibrationUtilsTest(parameterized.TestCase):
   def test_update_tensor_qsv_min_max(self, old_qsv, new_qsv, expected_qsv):
     updated_qsv = calibration_utils.min_max_update(old_qsv, new_qsv)
     if isinstance(expected_qsv["min"], list):
-      self.assertListEqual(list(updated_qsv["min"]), expected_qsv["min"])
-      self.assertListEqual(list(updated_qsv["max"]), expected_qsv["max"])
+      self.assertEqual(list(updated_qsv["min"]), expected_qsv["min"])
+      self.assertEqual(list(updated_qsv["max"]), expected_qsv["max"])
     else:
       self.assertEqual(updated_qsv["min"], expected_qsv["min"])
       self.assertEqual(updated_qsv["max"], expected_qsv["max"])
 
+  def test_calibration_utils_init_fails(self):
+    model_path = "non_existent_model.tflite"
+    with self.assertRaisesWithPredicateMatch(
+        tf.errors.NotFoundError, lambda err: f"{model_path}" in str(err)
+    ):
+      calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+
+  def test_calibration_utils_init_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_add.tflite"
+    )
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    self.assertNotEmpty(calib_utils._signature_runners)
+    self.assertNotEmpty(calib_utils._same_as_input_scale_ops)
+
+  def test_search_tensor_by_signature_name_succeeds_on_unconstrained_op(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_add.tflite"
+    )
+    expected_tensor_name = "PartitionedCall:0"
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    tensor_name = calib_utils._search_tensor_by_signature_name(
+        "serving_default", "add"
+    )
+    self.assertEqual(tensor_name, [expected_tensor_name])
+
+  def test_search_tensor_by_signature_name_succeeds_on_constrained_op(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_slice.tflite"
+    )
+    expected_tensor_name = "slice_input_tensor:0"
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    tensor_name = calib_utils._search_tensor_by_signature_name(
+        "slice", "output_0"
+    )
+    self.assertEqual(tensor_name, [expected_tensor_name])
+
+  def test_align_quant_stats_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/toy_model_with_kv_cache_multi_signature.tflite"
+    )
+    recipe_path = test_utils.get_path_to_datafile(
+        "../recipes/default_a8w8_recipe.json"
+    )
+    signature_data = {
+        "signature_1": ["output_1_1"],
+        "signature_2": ["output_1_1"],
+    }
+
+    # Obtain the calibration results.
+    qt = quantizer.Quantizer(model_path, recipe_path)
+    qsv = qt.calibrate(_CALIBRATION_DATASET)
+
+    # First quantize the model without aligning the quantization parameters.
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertFalse(
+        all(x == quant_params[0] for x in quant_params)
+    )  # not equal quantization params.
+
+    # Align the quantization parameters and quantize again.
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    calib_utils.align_quant_stats(qsv, signature_data)
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertTrue(
+        all(x == quant_params[0] for x in quant_params)
+    )  # equal quantization params.
+
+  def test_update_quant_stats_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/toy_model_with_kv_cache_multi_signature.tflite"
+    )
+    recipe_path = test_utils.get_path_to_datafile(
+        "../recipes/default_a8w8_recipe.json"
+    )
+    signature_data = {
+        "signature_1": ["output_1_1"],
+        "signature_2": ["output_1_1"],
+    }
+
+    # Obtain the calibration results.
+    qt = quantizer.Quantizer(model_path, recipe_path)
+    qsv = qt.calibrate(_CALIBRATION_DATASET)
+
+    # First quantize the model without updating the `signature_1`.
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertFalse(
+        all(x == quant_params[0] for x in quant_params)
+    )  # not equal quantization params.
+
+    # Update the `signature_1` with stats from `signature_2`.
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    min_val, max_val = calib_utils.align_quant_stats(  # for min and max only.
+        qsv,
+        {
+            "signature_2": ["output_1_1"],
+        },
+    )
+    calib_utils.update_quant_stats(
+        qsv,
+        {
+            "signature_1": ["output_1_1"],
+        },
+        min_val,
+        max_val,
+    )
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertTrue(
+        all(x == quant_params[0] for x in quant_params)
+    )  # equal quantization params.
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -59,6 +59,7 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.PAD: schema.BuiltinOperator.PAD,
     _TFLOpName.SQUARED_DIFFERENCE: schema.BuiltinOperator.SQUARED_DIFFERENCE,
     _TFLOpName.MAX_POOL_2D: schema.BuiltinOperator.MAX_POOL_2D,
+    _TFLOpName.RESIZE_BILINEAR: schema.BuiltinOperator.RESIZE_BILINEAR,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

{ai_edge_quantizer_nightly-0.3.0.dev20250613.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250615.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.3.0.dev20250613
+Version: 0.3.0.dev20250615
 Summary: A quantizer for advanced developers to quantize converted AI Edge models.
 Home-page: https://github.com/google-ai-edge/ai-edge-quantizer
 Keywords: On-Device ML,AI,Google,TFLite,Quantization,LLMs,GenAI

{ai_edge_quantizer_nightly-0.3.0.dev20250613.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250615.dist-info}/RECORD

@@ -1,23 +1,23 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=lfCazb2b0Q4L3of0cTWkF5lMr3AD6LWW1ekmFoEGB_4,12062
+ai_edge_quantizer/algorithm_manager.py,sha256=rMTM89YDPkmLKlUQV_Rjr7B2KpcvldAHzfpgUqaOqdU,12216
 ai_edge_quantizer/algorithm_manager_api.py,sha256=u903TG0s1uIDhJqfeJne3CFl8A93phZrwgV2-hwdcXU,9247
 ai_edge_quantizer/algorithm_manager_api_test.py,sha256=w6bSONvXkX6bzXAGc0-7b6gNDt9oz9ieq97KP8Sg_JU,7666
 ai_edge_quantizer/calibrator.py,sha256=Sms7_AIHPH9G5xFaz5Ef3a5gPhxuIWQI8d2LUM8C96I,12071
 ai_edge_quantizer/calibrator_test.py,sha256=C_oWOaRugPKYX74jF-eRFH-k6nGOdA8I9_uPiocaOuE,11900
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
-ai_edge_quantizer/default_policy.py,sha256=9CNd5zIk_BA560kOLkoXD5mapDBWj0yXzSFDedLhzYw,11192
+ai_edge_quantizer/default_policy.py,sha256=zghBh9dTB-ouPFumV-0siBSnEbp0WxF6tGOsn3TLirg,11242
 ai_edge_quantizer/model_modifier.py,sha256=teGa8I6kGvn6TQY6Xv53YFIc_pQEhNvM9Zb4bvhezyw,7110
 ai_edge_quantizer/model_modifier_test.py,sha256=cJd04SLOG-fQZZNZPcisoBLx3cLtWEwGqUBbLb-pif4,4751
 ai_edge_quantizer/model_validator.py,sha256=Hj0_5o-Oa3dSlJ3ryVjRhvsyelHNyek1GrtG9buMczg,13153
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
 ai_edge_quantizer/params_generator.py,sha256=gC7G6Ne4Fumc8RSmIAbx96ZBhszZlHqBKSmE9p6RPTo,20099
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=0Dwz6LHQG8LhZMhVAo_h6ieZ_gcfkJl2yJcsGf17YYs,16527
+ai_edge_quantizer/qtyping.py,sha256=kX1AoD-YlHYbDI1RfGVXIbPn-CYT7HUF2x77-hPtKBM,16565
 ai_edge_quantizer/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
 ai_edge_quantizer/quantizer_test.py,sha256=K_HBA56JkFI3HL8VLWCqGEfC0ISh5ldMKoNyBdGRAJg,20368
 ai_edge_quantizer/recipe.py,sha256=FR0uJceumZrnle2VRSOQZ1uXup4S1cTYKRH-N53mWRo,2919
 ai_edge_quantizer/recipe_manager.py,sha256=qcGUD7e7BISKdsY9WH2rdaRR3acmzSA5qMezGNbzlpo,8931
-ai_edge_quantizer/recipe_manager_test.py,sha256=LulVxsYp6TBGFI2PLCUCd4VsFq8ELpC7kMNkUjsLgbo,32230
+ai_edge_quantizer/recipe_manager_test.py,sha256=GVOfGFZPRciUb4EF4GkSi6d96LdjS6PbUkAJ0ayy0k8,32243
 ai_edge_quantizer/recipe_test.py,sha256=Fg_sfxovI2fRjk5qdu18ghOvXdUvhDR1TxbE0GHDczc,3381
 ai_edge_quantizer/transformation_instruction_generator.py,sha256=B_TQQe9_Qs7UKXLjMMuz5lORUvXyZOxBS2SpntTnkI8,28077
 ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=E0QSDCav6N6izlJ-a1ZJOsb2VEUxuxBmTbt0-EgDdxY,49890
@@ -28,7 +28,7 @@ ai_edge_quantizer/algorithms/nonlinear_quantize/__init__.py,sha256=lpq1g2ayg3lCP
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting.py,sha256=Bs9CK7wZAw6jNaZ8xEtbwO2vM34VYXNZSMVWvxJo9nw,9297
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py,sha256=EqIHGEZ1LgUrTN7zf880RuAzEv3Qy7kgh5ivObJGHSo,22646
 ai_edge_quantizer/algorithms/uniform_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=ofDoiZhOKjF7Tm-v0a4xsLSvytjfvMALXLDcuwcKNK0,29634
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=rImKK2ax7LrRx6XurSdvRTk0h6WtFGtQn9sYNJcn-uw,30222
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py,sha256=GGf_n3wIeg3GB_eGsmyNJ0fTcxgpeMMbugTMRONK6TQ,3553
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=BDdn_uBZakfHyzdMJPKadsOqxqyC-s6W2ZzFH99L4fE,8652
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py,sha256=sT5eX5TLZEHTtPfnSkCPDlS0sQxlTFWbCsbvOuj--yY,8889
@@ -61,17 +61,17 @@ ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=mHLO3_MRt36A8-Z
 ai_edge_quantizer/transformations/transformation_utils.py,sha256=GwIaKVsePZYgVG2lSanOswcaZYMjvgyqstDVwXl9DGY,6923
 ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=MWgq29t7rvxRQIfi4ny9IoODFCTcbpjnIwoCL40zDKk,8698
 ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/utils/calibration_utils.py,sha256=1Fj9MIO6aLZIRgyd4axvZN4S_O64nB_-Miu1WP664js,2536
-ai_edge_quantizer/utils/calibration_utils_test.py,sha256=Z-AcdTieesWFKyKBb08ZXm4Mgu6cvJ4bg2-MJ7hLD10,2856
+ai_edge_quantizer/utils/calibration_utils.py,sha256=e3dG7Nm94Ix0hkTWTWPUhEG6a8QR_cAM3PSwblfJV5g,15106
+ai_edge_quantizer/utils/calibration_utils_test.py,sha256=4BlksXl7b4yptL8xPR67hmJCnjhN9V10a2PunzfHrUE,9372
 ai_edge_quantizer/utils/test_utils.py,sha256=Y2pdMvn1k4gmqDo3noJfzx3fJcDHX_1hcsP6oiIz65Y,8240
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=Yy1u53FzRBFx-fr1TqoycWMZwAlAl0b2IB4MmGV1xJA,10758
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=pZv8FMWyjBSLN5MGJ2K_dZ6oqkJGbp9RI4CfnlPuPII,10830
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EtOv6cpKM_F0uv2bWuSXylYmTeXT6zUc182pw4sdYSI,13889
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=6fjkM-rycZ95L4yfvlr0TN6RlrhfPzxNUYrZaYO_F0A,12013
 ai_edge_quantizer/utils/validation_utils.py,sha256=oYw33Sg547AqtGw-choPUJmp9SAKkV46J_ddqSsum2Q,3950
 ai_edge_quantizer/utils/validation_utils_test.py,sha256=V_qNDikPD4OPB-siOLQCWNVWTAu87h2IgNYt7teFd-o,2934
-ai_edge_quantizer_nightly-0.3.0.dev20250613.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.3.0.dev20250613.dist-info/METADATA,sha256=YePqVOXDrP8t8VahK71fzW8xL0X14-coPWeR7eq45Jg,1528
-ai_edge_quantizer_nightly-0.3.0.dev20250613.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.3.0.dev20250613.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.3.0.dev20250613.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.3.0.dev20250615.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.3.0.dev20250615.dist-info/METADATA,sha256=IklxnJKNI7_fJW9CmL-QfF9EWmzzn8DRoGjwtpDZ8Wg,1528
+ai_edge_quantizer_nightly-0.3.0.dev20250615.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.3.0.dev20250615.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.3.0.dev20250615.dist-info/RECORD,,