ai-edge-quantizer-nightly 0.4.0.dev20250827__py3-none-any.whl → 0.4.0.dev20250829__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/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_nightly-0.4.0.dev20250827.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250829.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.4.0.dev20250827
+Version: 0.4.0.dev20250829
 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.dev20250827.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250829.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
@@ -59,8 +59,10 @@ 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_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
@@ -68,8 +70,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EoVjI_hplX_Rml3hfRsGmQOi
 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.dev20250827.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20250827.dist-info/METADATA,sha256=aGdUBHwbY4eeSd9c7pQyQLu2D5LoOEDXmKmM06XN-yI,1535
-ai_edge_quantizer_nightly-0.4.0.dev20250827.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20250827.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20250827.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20250829.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20250829.dist-info/METADATA,sha256=fg5k0J7zQJc0ufSBvuidEZKz57iydiIhRI4teV-7AZI,1535
+ai_edge_quantizer_nightly-0.4.0.dev20250829.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20250829.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20250829.dist-info/RECORD,,