ai-edge-quantizer-nightly 0.1.0.dev20250415__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/transformations/quantize_tensor_test.py

@@ -170,7 +170,7 @@ class QuantizeTensorTest(parameterized.TestCase):
     # Check if the scale and zero point tensors are inserted correctly.
     self.assertEqual(quant_param.details.scales, 9)
     # So far we don't have zero point in blockwise quantization.
-    self.assertEqual(quant_param.details.zeroPoints, 0)
+    self.assertEqual(quant_param.details.zeroPoints, -1)
 
   def test_int4_constant_packed_correctly(self):
     subgraph = self._model.subgraphs[0]

ai_edge_quantizer/transformations/transformation_utils.py

@@ -15,6 +15,7 @@
 
 """Utility functions for graph transformations."""
 
+import copy
 import dataclasses
 from typing import Optional, Union
 
@@ -51,39 +52,78 @@ class TransformationInput:
 def add_op_code(
     op_code: schema_py_generated.OperatorCodeT,
     model_op_codes: list[schema_py_generated.OperatorCodeT],
+    custom_op_name: Optional[str] = None,
 ) -> int:
   """Add an op code into a model if it's not present.
 
   Args:
     op_code: The op code to be added.
     model_op_codes: The op codes of the model.
+    custom_op_name: The custom string of the op code. If None, the op code will
+      be added as a builtin op code.
 
   Returns:
     The index of the op code in the model.
   """
+  if (
+      op_code == schema_py_generated.BuiltinOperator.CUSTOM
+      and custom_op_name is None
+  ):
+    raise ValueError('Custom string is required for custom op code.')
+
   for i, model_op_code in enumerate(model_op_codes):
+    # If the model already has the op code, just return the index.
     if model_op_code.builtinCode == op_code:
-      return i
+      if custom_op_name is not None:
+        if model_op_code.customCode == custom_op_name:
+          return i
+      else:
+        # Built-in op
+        return i
+
   model_op_codes.append(schema_py_generated.OperatorCodeT())
   model_op_codes[-1].builtinCode = op_code
+  if custom_op_name is not None:
+    model_op_codes[-1].customCode = custom_op_name
   return len(model_op_codes) - 1
 
 
-def add_new_constant_buffer(
+def get_constant_buffer(
     data: np.ndarray,
     buffers: list[schema_py_generated.BufferT],
+    force_duplicate_buffer: bool = False,
 ) -> int:
-  """Add a new constant buffer to the model.
+  """Get the index of the constant buffer that contains the given data.
+
+  creating new buffer if provided data is not found in buffers list.
 
   Args:
     data: The data of the new tensor.
     buffers: The buffers of the model.
+    force_duplicate_buffer: Whether to add a new buffer even if the same buffer
+      already exists.
 
   Returns:
     The index of the new buffer in the model.
   """
+
+  if isinstance(data, np.ndarray):
+    # in the case where the data is passed from quantization_params.
+    new_data = np.frombuffer(data.tobytes(), dtype=np.uint8).flatten()
+  elif isinstance(data, bytes):
+    # in the case where the data is coming from duplicating buffers, we need to
+    # make a copy of the data to avoid having two buffers pointing to the same
+    # data.
+    new_data = copy.deepcopy(data)
+  else:
+    raise ValueError('data passed in must be either np.ndarray or bytes.')
+  # TODO: b/417811116 - we should make this more efficient.
+  if not force_duplicate_buffer:
+    for index, buffer in enumerate(buffers):
+      if np.array_equal(buffer.data, new_data):
+        return index
   new_buffer = schema_py_generated.BufferT()
-  new_buffer.data = np.frombuffer(data.tobytes(), dtype=np.uint8).flatten()
+  new_buffer.data = new_data
   new_buffer.offset = 0
   new_buffer.size = 0
   new_buffer_id = len(buffers)
@@ -99,6 +139,7 @@ def add_new_constant_tensor(
     subgraph: schema_py_generated.SubGraphT,
     buffers: list[schema_py_generated.BufferT],
     tensor_shape: Optional[list[int]] = None,
+    force_duplicate_buffer: bool = False,
 ) -> int:
   """Add a new constant tensor to the model.
 
@@ -110,11 +151,13 @@ def add_new_constant_tensor(
     buffers: The buffers of the model.
     tensor_shape: The shape of the new tensor. If not provided, the shape of the
       data will be used.
+    force_duplicate_buffer: Whether to add a new buffer even if the same buffer
+      already exists.
 
   Returns:
     The index of the new tensor in the subgraph.
   """
-  new_buffer_id = add_new_constant_buffer(data, buffers)
+  new_buffer_id = get_constant_buffer(data, buffers, force_duplicate_buffer)
 
   new_tensor = schema_py_generated.TensorT()
   if tensor_shape is None:
@@ -146,10 +189,90 @@ def add_new_activation_tensor(
     The index of the new tensor in the subgraph.
   """
   new_tensor = schema_py_generated.TensorT()
-  new_tensor.shape = shape
+  # If there's a dynamic shape, we need to read from the shapeSignature field
+  # instead of shape. Shape should contain just 1 for the dynamic dimension but
+  # shapeSignature should contain the true shape.
+  if -1 in shape:
+    new_tensor.shapeSignature = shape
+    new_tensor.shape = [1 if i == -1 else i for i in shape]
+  else:
+    new_tensor.shape = shape
   new_tensor.type = tensor_type
   new_tensor.name = tensor_name
   new_tensor.buffer = 0
   new_tensor_id = len(subgraph.tensors)
   subgraph.tensors.append(new_tensor)
   return new_tensor_id
+
+
+def raise_deprecated_error(_: TransformationInput):
+  raise NotImplementedError(
+      'This transformation is deprecated. Please contact AI Edge Quantizer team'
+      ' if you see this error.'
+  )
+
+
+def pack_data(bitwidth: int, flattened_data: np.ndarray) -> np.ndarray:
+  """Pack the data to the corresponding bit width.
+
+  Currently only support 4 bits. If no packing is needed, the original data is
+  returned.
+
+  Args:
+    bitwidth: Bit width from NonLinearQuantParams.
+    flattened_data: The data to be packed.
+
+  Returns:
+    Packed data.
+  """
+  if bitwidth == 4:
+    even_data = flattened_data[::2] & 0x0F
+    odd_data = np.left_shift(flattened_data[1::2], 4).astype(np.uint8)
+    if odd_data.shape[0] == even_data.shape[0] - 1:
+      odd_data = np.pad(odd_data, (0, 1), constant_values=0)
+    return np.bitwise_or(even_data, odd_data)
+  else:
+    return flattened_data
+
+
+def get_producer_schema_op_id(
+    transformation: TransformationInput,
+) -> int:
+  """Checks if the tensor's producer matches the given op.
+
+  Args:
+    transformation: The transformation input to check the producer of.
+
+  Returns:
+    The schema op id of the producer op. E.g.
+    schema_py_generated.BuiltinOperator.FULLY_CONNECTED.
+  """
+  if transformation.producer == -1:
+    return False
+  else:
+    return (
+        transformation.op_codes[
+            transformation.subgraph.operators[
+                transformation.producer
+            ].opcodeIndex
+        ].builtinCode
+    )
+
+
+def get_schema_op_id(
+    transformation: TransformationInput, op_id: int
+) -> bool:
+  """Returns the schema op id of the given op.
+
+  Args:
+    transformation: The transformation input to check the consumers of.
+    op_id: The op id in the list of operators to check for.
+
+  Returns:
+    The schema op id of the given op.
+  """
+  return (
+      transformation.op_codes[
+          transformation.subgraph.operators[op_id].opcodeIndex
+      ].builtinCode
+  )

ai_edge_quantizer/transformations/transformation_utils_test.py

@@ -41,19 +41,62 @@ class TransformationUtilsTest(parameterized.TestCase):
           testcase_name="add_new_op_code",
           op_code=schema_py_generated.BuiltinOperator.LOGISTIC,
           expected=1,
+          custom_op_name=None,
       ),
       dict(
           testcase_name="add_existing_op_code",
           op_code=schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
           expected=0,
+          custom_op_name=None,
+      ),
+      dict(
+          testcase_name="add_new_custom_op_code",
+          op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+          expected=1,
+          custom_op_name="random_new_custom_op",
       ),
   )
-  def test_add_op_code(self, op_code, expected):
+  def test_add_op_code(self, op_code, expected, custom_op_name):
     """Tests if the op code is added to the model."""
     got = transformation_utils.add_op_code(
-        op_code=op_code, model_op_codes=self.model.operatorCodes
+        op_code=op_code,
+        model_op_codes=self.model.operatorCodes,
+        custom_op_name=custom_op_name,
     )
     self.assertEqual(expected, got)
+    if custom_op_name is not None:
+      self.assertEqual(self.model.operatorCodes[got].customCode, custom_op_name)
+
+  def test_add_custom_op_code_without_op_string_raises_error(self):
+    with self.assertRaisesRegex(ValueError, "Custom string is required"):
+      transformation_utils.add_op_code(
+          op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+          model_op_codes=self.model.operatorCodes,
+          custom_op_name=None,
+      )
+
+  def test_add_two_custom_op_codes(self):
+    custom_op_name = "random_new_custom_op"
+    added_index = transformation_utils.add_op_code(
+        op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+        model_op_codes=self.model.operatorCodes,
+        custom_op_name=custom_op_name,
+    )
+    self.assertEqual(1, added_index)
+    self.assertEqual(
+        self.model.operatorCodes[added_index].customCode, custom_op_name
+    )
+
+    custom_op_name_2 = "random_new_custom_op_2"
+    added_index = transformation_utils.add_op_code(
+        op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+        model_op_codes=self.model.operatorCodes,
+        custom_op_name=custom_op_name_2,
+    )
+    self.assertEqual(2, added_index)
+    self.assertEqual(
+        self.model.operatorCodes[added_index].customCode, custom_op_name_2
+    )
 
   @parameterized.named_parameters(
       dict(
@@ -68,7 +111,7 @@ class TransformationUtilsTest(parameterized.TestCase):
   def test_add_new_constant_buffer(self, data):
     """Tests if the constant buffer is added to the model."""
     prev_num_buffers = len(self.model.buffers) - 1
-    new_buffer_idx = transformation_utils.add_new_constant_buffer(
+    new_buffer_idx = transformation_utils.get_constant_buffer(
         data=data,
         buffers=self.model.buffers,
     )
@@ -189,6 +232,25 @@ class TransformationUtilsTest(parameterized.TestCase):
         self.model.subgraphs[0].tensors[-1].shape,
     )
 
+  def test_add_new_activation_tensor_with_dynamic_shape(self):
+    """Tests adding an activation tensor with dynamic shape."""
+    subgraph = self.model.subgraphs[0]
+    new_id = transformation_utils.add_new_activation_tensor(
+        tensor_name="test_tensor",
+        shape=[1, -1, -1, 1],
+        tensor_type=schema_py_generated.TensorType.FLOAT32,
+        subgraph=subgraph,
+    )
+    # Originally had 4 tensors, new tensor is added at index 4.
+    self.assertEqual(new_id, 4)
+    self.assertLen(subgraph.tensors, 5)
+    self.assertEqual(subgraph.tensors[-1].name, "test_tensor")
+    self.assertEqual(
+        subgraph.tensors[-1].type, schema_py_generated.TensorType.FLOAT32
+    )
+    self.assertEqual(subgraph.tensors[-1].shape, [1, 1, 1, 1])
+    self.assertEqual(subgraph.tensors[-1].shapeSignature, [1, -1, -1, 1])
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/utils/calibration_utils.py

@@ -15,9 +15,24 @@
 
 """Utilities for model calibration."""
 
-from typing import Union
+import copy
+from typing import Any, Union
+
 import numpy as np
+
+from ai_edge_litert.tools import flatbuffer_utils
 from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import constrained_ops_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+
+_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 +99,250 @@ 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 = (
+        constrained_ops_utils.get_constrained_op_list(
+            _OpQuantConstraint.SAME_AS_INPUT_SCALE
+        )
+    )
+
+    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
+
+  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