ai-edge-quantizer-nightly 0.4.0.dev20250828__py3-none-any.whl → 0.4.0.dev20250830__py3-none-any.whl

ai_edge_quantizer/default_policy.py

@@ -266,7 +266,11 @@ DEFAULT_JSON_POLICY = """
   }
 }
 """
-QUANTIZABLE_COMPOSITES = ["od" + "ml.npu_call", "od" + "ml.rms_norm"]
+QUANTIZABLE_COMPOSITES = [
+    "od" + "ml.npu_call",
+    "od" + "ml.rms_norm",
+    "od" + "ml.l2_norm",
+]
 
 
 def _unroll_json_config(

ai_edge_quantizer/transformation_instruction_generator.py

@@ -23,10 +23,13 @@ from collections.abc import Iterator
 import dataclasses
 from typing import Optional
 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_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
 
 
+_OpQuantConstraint = common_utils.OpQuantConstraint
 _QuantTransformation = qtyping.QuantTransformation
 
 
@@ -165,6 +168,16 @@ class TransformationInstructionsGenerator:
     else:
       self.flatbuffer_model = tfl_flatbuffer_utils.read_model(float_tflite)
       self._create_tensor_name_to_graph_info_map()
+    self._same_as_input_scale_ops = (
+        constrained_ops_utils.get_constrained_op_list(
+            _OpQuantConstraint.SAME_AS_INPUT_SCALE
+        )
+    )
+    self._same_as_output_scale_ops = (
+        constrained_ops_utils.get_constrained_op_list(
+            _OpQuantConstraint.SAME_AS_OUTPUT_SCALE
+        )
+    )
 
   @dataclasses.dataclass(frozen=True)
   class TensorGraphInfo:
@@ -506,6 +519,89 @@ class TransformationInstructionsGenerator:
       ):
         instructions.pop(i)
 
+  def _is_valid_quantize_requantize_pair(
+      self,
+      instr_0: qtyping.TransformationInst,
+      instr_1: qtyping.TransformationInst,
+  ) -> bool:
+    """Checks if the two instructions form a valid quantize and requantize pair."""
+    return (
+        instr_0.transformation == _QuantTransformation.QUANTIZE_TENSOR
+        and instr_1.transformation == _QuantTransformation.ADD_QUANTIZE
+        and instr_0.consumers == instr_1.consumers
+    )
+
+  def _is_op_constrained(
+      self, subgraph_id: int, op_index: int
+  ) -> bool:
+    """Checks if the op has same as input or output scale constraints."""
+    op_name = tfl_flatbuffer_utils.get_op_name_by_index(
+        self.flatbuffer_model, subgraph_id, op_index
+    )
+    return (
+        op_name in self._same_as_input_scale_ops
+        or op_name in self._same_as_output_scale_ops
+    )
+
+  def _are_quant_params_compatible(
+      self,
+      params_0: qtyping.UniformQuantParams,
+      params_1: qtyping.UniformQuantParams,
+  ) -> bool:
+    """Checks if quant params are the same except for the scale and zero point."""
+    ignore_set = {"scale", "zero_point"}
+    for field_info in dataclasses.fields(qtyping.UniformQuantParams):
+      field_name = field_info.name
+      if field_name in ignore_set:
+        continue
+      if getattr(params_0, field_name) != getattr(params_1, field_name):
+        return False
+    return True
+
+  def _eliminate_requantization_for_nonconstrained_provider(
+      self, tensor_trans_insts: qtyping.TensorTransformationInsts
+  ) -> None:
+    """Removes requantization for tensors with a non-constrained provider.
+
+    Fuses [QUANTIZE_TENSOR, ADD_QUANTIZE] instructions when a tensor has a
+    provider op without same as input/ouput scale constrains. Quant params from
+    the second instruction are copied to the first one and ADD_QUANTIZE is
+    removed.
+
+    Args:
+      tensor_trans_insts: Transformation instructions for a tensor.
+    """
+    instructions = tensor_trans_insts.instructions
+    if instructions is None or len(instructions) != 2:
+      return
+
+    instr_0, instr_1 = instructions
+    params_0 = instr_0.parameters
+    params_1 = instr_1.parameters
+    producer_op_index = instr_0.producer
+    if (
+        not isinstance(params_0, qtyping.UniformQuantParams)
+        or not isinstance(params_1, qtyping.UniformQuantParams)
+        or not self._is_valid_quantize_requantize_pair(instr_0, instr_1)
+        or not self._are_quant_params_compatible(params_0, params_1)
+        # To avoid fusion when subgraph inputs connected to the main subgraph
+        # (e.g. while_body), we skip all tensors with no producer.
+        or producer_op_index == -1
+        # Can't apply fusion to tensors with a constrained producer since that
+        # will break the constraint.
+        or self._is_op_constrained(
+            tensor_trans_insts.subgraph_id, producer_op_index
+        )
+    ):
+      return
+
+    # Fuse the quantize and requantize.
+    instr_0.parameters = dataclasses.replace(
+        params_0, scale=params_1.scale, zero_point=params_1.zero_point
+    )
+    # Remove the requantize instruction.
+    instructions.pop(1)
+
   def _quant_params_to_transformation_insts(
       self,
       param: qtyping.TensorTransformationParams,
@@ -578,6 +674,12 @@ class TransformationInstructionsGenerator:
     # will raise an error if the instructions are not valid.
     self._check_tensor_transformation_instructions_valid(tensor_trans_insts)
 
+    # Remove unnecessary [QUANTIZE_TENSOR, ADD_QUANTIZE] pairs for tensors with
+    # providers without same as input/output scale constraints.
+    self._eliminate_requantization_for_nonconstrained_provider(
+        tensor_trans_insts
+    )
+
     return tensor_trans_insts
 
   def _split_instructions_by_tensor_duplication(

ai_edge_quantizer/transformation_instruction_generator_test.py

@@ -15,7 +15,9 @@
 
 """Tests for instruction_generator."""
 
+from collections.abc import Sequence
 import os
+from typing import Optional
 
 import numpy as np
 
@@ -1337,5 +1339,166 @@ class InstructionGeneratorTest(parameterized.TestCase):
       )
 
 
+class EliminateUnnecessaryRequantizationTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.ins_gen = instruction_generator.TransformationInstructionsGenerator(
+        os.path.join(TEST_DATA_PREFIX_PATH, "tests/models/conv_fc_mnist.tflite")
+    )
+
+  def _get_test_instruction(
+      self,
+      transformation: qtyping.QuantTransformation,
+      producer: int = -1,
+      consumers: Optional[Sequence[int]] = None,
+      qparams: Optional[qtyping.UniformQuantParams] = None,
+  ) -> qtyping.TransformationInst:
+    if consumers is None:
+      consumers = []
+    if qparams is None:
+      qparams = qtyping.UniformQuantParams(
+          num_bits=8,
+          quantized_dimension=None,
+          scale=np.array([1]),
+          zero_point=np.array([0]),
+      )
+    return qtyping.TransformationInst(
+        transformation=transformation,
+        producer=producer,
+        consumers=consumers,
+        parameters=qparams,
+        # Dummy values below.
+        tensor_id=0,
+    )
+
+  def _create_test_insts(
+      self, instructions: list[qtyping.TransformationInst]
+  ) -> qtyping.TensorTransformationInsts:
+    return qtyping.TensorTransformationInsts(
+        tensor_name="test_tensor", subgraph_id=0, instructions=instructions
+    )
+
+  def test_no_fusion_when_too_few_instructions(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 1)
+
+  def test_no_fusion_when_too_many_instructions(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE),
+        self._get_test_instruction(_QTransf.ADD_DEQUANTIZE),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 3)
+
+  def test_no_fusion_when_invalid_transformation_pair(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.ADD_DEQUANTIZE),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_consumers_mismatch(self):
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR, consumers=[0]),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE, consumers=[1]),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_no_producer(self):
+    producer = -1
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR, producer),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE, producer),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_quant_params_are_incompatible(self):
+    params_8_bits = qtyping.UniformQuantParams(
+        8, None, np.array([1]), np.array([0])
+    )
+    params_16_bits = qtyping.UniformQuantParams(
+        16, None, np.array([1]), np.array([0])
+    )
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(
+            _QTransf.QUANTIZE_TENSOR, qparams=params_8_bits
+        ),
+        self._get_test_instruction(
+            _QTransf.ADD_QUANTIZE, qparams=params_16_bits
+        ),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_no_fusion_when_producer_constrained(self):
+    # Reshape op (op index 2) has same as input scale constraint.
+    tensor_insts = self._create_test_insts([
+        self._get_test_instruction(_QTransf.QUANTIZE_TENSOR, producer=2),
+        self._get_test_instruction(_QTransf.ADD_QUANTIZE, producer=2),
+    ])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+    self.assertLen(tensor_insts.instructions, 2)
+
+  def test_fusion_succeeds(self):
+    producer = 0
+    consumers = [1]
+    params_0 = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.array([1]),
+        zero_point=np.array([0]),
+    )
+    params_1 = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.array([2]),
+        zero_point=np.array([1]),
+    )
+    inst_0 = self._get_test_instruction(
+        _QTransf.QUANTIZE_TENSOR, producer, consumers, params_0
+    )
+    inst_1 = self._get_test_instruction(
+        _QTransf.ADD_QUANTIZE, producer, consumers, params_1
+    )
+    tensor_insts = self._create_test_insts([inst_0, inst_1])
+    self.ins_gen._eliminate_requantization_for_nonconstrained_provider(
+        tensor_insts
+    )
+
+    self.assertLen(tensor_insts.instructions, 1)
+    result_inst = tensor_insts.instructions[0]
+    self.assertEqual(result_inst.transformation, _QTransf.QUANTIZE_TENSOR)
+
+    result_params = result_inst.parameters
+    # Explicitly narrow the type for pytype.
+    if not isinstance(result_params, qtyping.UniformQuantParams):
+      self.fail("Fused instruction parameters are not UniformQuantParams")
+
+    self.assertEqual(result_params.scale, params_1.scale)
+    self.assertEqual(result_params.zero_point, params_1.zero_point)
+
+
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/utils/calibration_utils.py

@@ -20,13 +20,11 @@ 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 constrained_ops_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
 
 
@@ -133,7 +131,11 @@ class CalibrationQsvAlignmentUtils:
   """
 
   def __init__(self, model_path: str):
-    self._same_as_input_scale_ops = []
+    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)
@@ -146,87 +148,6 @@ class CalibrationQsvAlignmentUtils:
       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]:

ai_edge_quantizer/utils/constrained_ops_utils.py

@@ -0,0 +1,111 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# 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.
+# ==============================================================================
+
+"""Utils for handling operators with quantization constraints."""
+
+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_litert import schema_py_generated as schema_fb  # pylint: disable=g-direct-tensorflow-import
+
+
+_OpQuantConstraint = common_utils.OpQuantConstraint
+
+
+def get_constrained_op_list(
+    quant_constraint: _OpQuantConstraint, verbose: bool = False
+) -> list[str]:
+  """Constructs and returns a list of constrained operators.
+
+  This is achieved by invoking all materialization functions and extracting
+  the constraint argument, using monkey patching to redirect logic to wrapper
+  functions.
+
+  Args:
+    quant_constraint: The quantization constraint to filter operators by.
+    verbose: Flag to enable verbose output.
+
+  Returns:
+    A list containing operators with the specified constraint.
+  """
+  constrained_ops = []
+
+  def materialize_standard_op_wrapper(
+      op_info: qtyping.OpInfo,
+      *_args,
+      constraint: _OpQuantConstraint = _OpQuantConstraint.NO_CONSTRAIN,
+      **_kwargs,
+  ) -> list[qtyping.TensorTransformationParams]:
+    if constraint == quant_constraint:
+      constrained_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
+  # ops with the given constraint.
+  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"  {quant_constraint} op list: {constrained_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
+  )
+  return constrained_ops

ai_edge_quantizer/utils/constrained_ops_utils_test.py

@@ -0,0 +1,50 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# 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 tensorflow.python.platform import googletest
+from absl.testing import parameterized
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import constrained_ops_utils
+
+
+_OpQuantConstraint = common_utils.OpQuantConstraint
+
+
+class ConstrainedOpsUtilsTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="same_as_input_scale",
+          constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
+          expected_num_ops=14,
+      ),
+      dict(
+          testcase_name="same_as_output_scale",
+          constraint=_OpQuantConstraint.SAME_AS_OUTPUT_SCALE,
+          expected_num_ops=6,
+      ),
+      dict(
+          testcase_name="no_constrain",
+          constraint=_OpQuantConstraint.NO_CONSTRAIN,
+          expected_num_ops=22,
+      ),
+  )
+  def test_get_constrained_op_list(self, constraint, expected_num_ops):
+    constrained_ops = constrained_ops_utils.get_constrained_op_list(constraint)
+    self.assertLen(constrained_ops, expected_num_ops)
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -342,3 +342,12 @@ def get_op_side_effect_subgraphs(
     return [opts.decompositionSubgraphIndex]
   # Can add other nested ops here (control flow ops, etc).
   return []
+
+
+def get_op_name_by_index(
+    flatbuffer_model: Any, subgraph_id: int, op_index: int
+) -> str:
+  """Get the op name from the flatbuffer model."""
+  op = flatbuffer_model.subgraphs[subgraph_id].operators[op_index]
+  builtin_code = flatbuffer_model.operatorCodes[op.opcodeIndex].builtinCode
+  return TFL_OP_CODE_TO_NAME[builtin_code]

{ai_edge_quantizer_nightly-0.4.0.dev20250828.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250830.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.4.0.dev20250828
+Version: 0.4.0.dev20250830
 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.4.0.dev20250828.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250830.dist-info}/RECORD

@@ -5,7 +5,7 @@ ai_edge_quantizer/algorithm_manager_api_test.py,sha256=w6bSONvXkX6bzXAGc0-7b6gND
 ai_edge_quantizer/calibrator.py,sha256=Sms7_AIHPH9G5xFaz5Ef3a5gPhxuIWQI8d2LUM8C96I,12071
 ai_edge_quantizer/calibrator_test.py,sha256=ZLzIMWB2FSFU4TOatDioYuwp_kLh8iSCefZ5_Q9FU7s,11900
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
-ai_edge_quantizer/default_policy.py,sha256=LXEdwdr0SiCfWo6ZwbHQ8ykoqA40GV6fGAT1aofry3o,11556
+ai_edge_quantizer/default_policy.py,sha256=G_JZtZaQAnrWyfCusDWXwO27iLysk27RS91GlS61m_Q,11592
 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
@@ -19,8 +19,8 @@ ai_edge_quantizer/recipe.py,sha256=MEkfQ2Sg3KAE9LAORHWcbjYNPg06EUbwc1d-VspQA2U,6
 ai_edge_quantizer/recipe_manager.py,sha256=6dgbE-IZfEetzXH3p3Qm_9eQutNDOpZnMpiaLTbP-ZQ,14744
 ai_edge_quantizer/recipe_manager_test.py,sha256=H-B75vwPN5ND-nUa3pOXizeHTv4mufPiC5cL_OlDIYU,34040
 ai_edge_quantizer/recipe_test.py,sha256=GKuo6N65wKLS2xwSpjd-BWWeVRpF1zc7Yt7phSMYSxA,5905
-ai_edge_quantizer/transformation_instruction_generator.py,sha256=iMGXy7_ufqgQRzu4drAfO31VGdze35peEFh1BMZlVHk,27714
-ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=Zw3EOSnvzjuB4NWeo129eJZxK_EHno9oF9OtEQ-0dnM,48905
+ai_edge_quantizer/transformation_instruction_generator.py,sha256=O0U2aZcB8aXQgOV8r9g1rGNzDUiuI5Ta53XnxZbVffE,31576
+ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=KW5-WoTTo9IqLEVnWxVC8ut8eWLi_91xfKgGqVQ9QDk,54635
 ai_edge_quantizer/transformation_performer.py,sha256=o4J6OUbI0dLoobVYjkOFw5Po3yH0gZJXrfuTIYais4o,13029
 ai_edge_quantizer/transformation_performer_test.py,sha256=xk6A3LStCyPclN51--9uO7XjSxNfZmpdfvrzOL0maNM,20349
 ai_edge_quantizer/algorithms/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
@@ -59,17 +59,19 @@ ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=mHLO3_MRt36A8-Z
 ai_edge_quantizer/transformations/transformation_utils.py,sha256=efJdAkA24wlg6Vj5NFO7_7MDuvQLSNn-l11Vs_JPktI,7123
 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=e3dG7Nm94Ix0hkTWTWPUhEG6a8QR_cAM3PSwblfJV5g,15106
+ai_edge_quantizer/utils/calibration_utils.py,sha256=iMf_bSCf-O86MzDt5D9hLKqbTydqLwirluaC6BJ9yHo,11553
 ai_edge_quantizer/utils/calibration_utils_test.py,sha256=4BlksXl7b4yptL8xPR67hmJCnjhN9V10a2PunzfHrUE,9372
+ai_edge_quantizer/utils/constrained_ops_utils.py,sha256=EAITCf7Ku_PFZcw3K-wd-8hGbyuRd5W5UtNdGvalwAE,4478
+ai_edge_quantizer/utils/constrained_ops_utils_test.py,sha256=6k_AqfB-NmiLkW5WwEV5NSuswFWky2sL0xBGmV6Fdwk,1756
 ai_edge_quantizer/utils/test_utils.py,sha256=a4Nk-wbeB09dFjTDZiA0K67d26j5DD0UDH_GIVmVG_4,8685
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=RL6oq6FzZj-xV0Zgh0UBn7-fOQaRXSxZ-PPG_LmtyUY,11384
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=aNtL4dpWH5uGGGlaygnMDkh5llTstbgs5ZxO0JkH5VQ,11718
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EoVjI_hplX_Rml3hfRsGmQOihexmizeJqt4SQcET9aA,14925
 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.4.0.dev20250828.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20250828.dist-info/METADATA,sha256=08u7TZ16Y_SHk1eBL3q6pxP4U79rQEXyXbqUEBWuXFo,1535
-ai_edge_quantizer_nightly-0.4.0.dev20250828.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20250828.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20250828.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20250830.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20250830.dist-info/METADATA,sha256=B1PQw7561EH7WCLLBTwYIoXP8WH-0vJeTyPuQPAiX_M,1535
+ai_edge_quantizer_nightly-0.4.0.dev20250830.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20250830.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20250830.dist-info/RECORD,,