ai-edge-quantizer-nightly 0.3.0.dev20250610__py3-none-any.whl → 0.3.0.dev20250612__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -111,6 +111,7 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.SQUARED_DIFFERENCE: (
         common_quantize.materialize_squared_difference
     ),
+    _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -248,6 +249,7 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     _TFLOpName.SQUARED_DIFFERENCE: (
         common_quantize.materialize_squared_difference
     ),
+    _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -712,6 +712,22 @@ def materialize_squared_difference(
   )
 
 
+def materialize_max_pool_2d(
+    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.max_pool_2d."""
+  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,
+  )
+
+
 def _get_tensor_shape_for_blockwise(
     tensor_shape: Sequence[int], quantized_dim: int, block_size: int
 ) -> list[int]:

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py

@@ -435,7 +435,8 @@ def _is_valid_quantization_params(
   """Checks if the quantization parameters are valid.
 
   A valid quantization params requires:
-    1. scale and zero point have the same shape (TFL Runtime requirement).
+    1. scale and zero point either have the same shape or the zero point is a
+    scalar.
     2. scale and zero point have the same rank as the tensor content (avoid
     ambiguous broadcasting).
 
@@ -446,17 +447,20 @@ def _is_valid_quantization_params(
   Returns:
     True if the quantization parameters are valid.
   """
-  if quantization_params.scale.shape != quantization_params.zero_point.shape:
+  if (
+      quantization_params.scale.shape != quantization_params.zero_point.shape
+      and quantization_params.zero_point.size != 1
+  ):
     raise ValueError(
-        "scale and zero_point must have the same shape. Got"
-        f" {quantization_params.scale.shape} and"
+        "scale and zero_point must have the same shape or zero_point must have"
+        f" only one element. Got {quantization_params.scale.shape} and"
         f" {quantization_params.zero_point.shape}"
     )
 
   tensor_rank = tensor_data.ndim
   scale_rank = quantization_params.scale.ndim
   zero_point_rank = quantization_params.zero_point.ndim
-  if (tensor_rank != scale_rank) or (tensor_rank != zero_point_rank):
+  if tensor_rank != scale_rank or (tensor_rank != zero_point_rank):
     raise ValueError(
         f"Ranks of scales ({scale_rank}) and zps"
         f" ({zero_point_rank}) must be the same as the tensor rank"

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py

@@ -160,7 +160,9 @@ class TensorUtilsTest(parameterized.TestCase):
   def test_uniform_quantize_wrong_shape(self):
     tensor = [-3.0, 1.3, 2.4, 16.0]
 
-    error_message = "scale and zero_point must have the same shape."
+    error_message = (
+        "Ranks of scales (3) and zps (2) must be the same as the tensor rank"
+    )
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: error_message in str(err)
     ):
@@ -233,7 +235,9 @@ class TensorUtilsTest(parameterized.TestCase):
   def test_uniform_dequantize_wrong_shape(self):
     tensor = [-3.0, 1.3, 2.4, 16.0]
 
-    error_message = "scale and zero_point must have the same shape."
+    error_message = (
+        "Ranks of scales (3) and zps (2) must be the same as the tensor rank"
+    )
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: error_message in str(err)
     ):

ai_edge_quantizer/default_policy.py

@@ -185,7 +185,7 @@ DEFAULT_JSON_POLICY = """
       "SELECT_V2",
       "STABLEHLO_COMPOSITE",
       "PAD",
-      "SQUARED_DIFFERENCE"
+      "MAX_POOL_2D"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -218,7 +218,8 @@ DEFAULT_JSON_POLICY = """
       "SELECT_V2",
       "STABLEHLO_COMPOSITE",
       "PAD",
-      "SQUARED_DIFFERENCE"
+      "SQUARED_DIFFERENCE",
+      "MAX_POOL_2D"
     ],
     "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

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

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -58,6 +58,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,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

{ai_edge_quantizer_nightly-0.3.0.dev20250610.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250612.dist-info}/METADATA

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

@@ -1,18 +1,18 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=6Q2emCwebYbU_A6epyLBZ3C_1L5EWP9N-5mGAt-FNiU,11924
+ai_edge_quantizer/algorithm_manager.py,sha256=lfCazb2b0Q4L3of0cTWkF5lMr3AD6LWW1ekmFoEGB_4,12062
 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=-_jX_KkfIepkQAwxxDrZjvPO1JsoSjHXVy1DPc1iFjM,12068
 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=H6Jn5jPNRkDT54GYR2xI-1A4v_R4zfu_6BqgcmRevq0,11181
+ai_edge_quantizer/default_policy.py,sha256=nKtghUjTQ8QS9CgLRwQb3iB2eZOyQv0FqyISlcgzSH4,11195
 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=j1BV2cGFLlQmUY6aoW5uglYqf77b9ytN8oZ1gh6o0mM,20096
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=mJi8B--NYH-ad9rCJpXZIZsua3qkPBI6cm90vD44VHA,16497
+ai_edge_quantizer/qtyping.py,sha256=0Dwz6LHQG8LhZMhVAo_h6ieZ_gcfkJl2yJcsGf17YYs,16527
 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
@@ -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=2vxo8pmEPWD0fklhibNweiah6LgLL8tenSbDvsui5Oo,29128
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=ofDoiZhOKjF7Tm-v0a4xsLSvytjfvMALXLDcuwcKNK0,29634
 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
@@ -38,8 +38,8 @@ ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py,sha256=8
 ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py,sha256=zoF_EHjYqsKkuev8wfuutIITEmp_maa70IpJI_Df3ck,7431
 ai_edge_quantizer/algorithms/uniform_quantize/octav.py,sha256=Umxh4kJyeHddZf-Wd4aXE5MTI1XWFa5KRuM17uYU714,6922
 ai_edge_quantizer/algorithms/uniform_quantize/octav_test.py,sha256=sha1d99Xk87bI87tgz0g5LeDC-EeE4WMfM5rRC98-m4,9140
-ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=W2QbXP96xeleAmA7qFwco1iq_bOtArGDK6Qj_g6kNl8,15986
-ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=MgG7Qh2_z4I6InBqEEDSVlaR0q48aMz4xqAlxeG2EMk,12436
+ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=3zq2AO_PRYKHuNvHzwg0pVDZT7kcpaMgXx6OEyEl6co,16103
+ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=JlX3fLHiknGH1osu6gwWEGUizLrEsE6d8iRpzDODmXo,12510
 ai_edge_quantizer/algorithms/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/utils/common_utils.py,sha256=UoZxeAQmZk3b3hK51KFwq6XfdbeduXVjdYIxAxlAzB8,34982
 ai_edge_quantizer/algorithms/utils/common_utils_test.py,sha256=zqapGEfYhjQWe9cNGPLmdbwtEUUYQRhlO_kNe0cXX6E,18104
@@ -64,14 +64,14 @@ ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V
 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/test_utils.py,sha256=Y2pdMvn1k4gmqDo3noJfzx3fJcDHX_1hcsP6oiIz65Y,8240
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=JmfVgXLEdt0cYKE1At9H8yKyFISE9GC_mKcGSF06o8E,10694
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=Yy1u53FzRBFx-fr1TqoycWMZwAlAl0b2IB4MmGV1xJA,10758
 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.dev20250610.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.3.0.dev20250610.dist-info/METADATA,sha256=Az9C-MKHQfNVxpw2DcQvQVCChViByoruNFGEnA8r_58,1528
-ai_edge_quantizer_nightly-0.3.0.dev20250610.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.3.0.dev20250610.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.3.0.dev20250610.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.3.0.dev20250612.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.3.0.dev20250612.dist-info/METADATA,sha256=C0e46WStdRDjUO05juSVs83Bu_Tm27ZBa9XZoLRmAZo,1528
+ai_edge_quantizer_nightly-0.3.0.dev20250612.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.3.0.dev20250612.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.3.0.dev20250612.dist-info/RECORD,,