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

ai_edge_quantizer/calibrator.py

@@ -165,7 +165,7 @@ class Calibrator:
             )
             if algorithm_name == algorithm_manager.AlgorithmName.NO_QUANTIZE:
               continue
-            if policy.is_conditionally_unquantized(op):
+            if policy.is_non_quantizable_composite_op(op):
               continue
 
             # Step2.2: query algorithm_manager to get/call the related

ai_edge_quantizer/default_policy.py

@@ -245,6 +245,7 @@ DEFAULT_JSON_POLICY = """
   }
 }
 """
+QUANTIZABLE_COMPOSITES = ["od" + "ml.npu_call", "od" + "ml.rms_norm"]
 
 
 def _unroll_json_config(
@@ -322,10 +323,10 @@ def _unroll_json_config(
 
 
 # TODO: b/401024954 - Have a better way to specify recipes based on op options.
-def is_conditionally_unquantized(
+def is_non_quantizable_composite_op(
     op: Union[schema.Operator, schema.OperatorT],
 ) -> bool:
-  """Checks if the operator is conditionally unquantized.
+  """Checks if the operator is a non-quantizable composite op.
 
   We may want to quantize an op only when its has certain options.
   Policies/recipes
@@ -340,10 +341,9 @@ def is_conditionally_unquantized(
   if opts := flatbuffer_utils.get_options_as(
       op, schema.StableHLOCompositeOptionsT
   ):
-    name: bytes = opts.name
-    # Non npu_call composites may have a kernel and as such will not be
-    # quantized.
-    return ("od" + "ml.npu_call") not in name.decode("utf-8")
+    name = opts.name.decode("utf-8")
+    if name not in QUANTIZABLE_COMPOSITES:
+      return True
 
   return False
 

ai_edge_quantizer/params_generator.py

@@ -109,7 +109,7 @@ class ParamsGenerator:
         algorithm_name, op_quant_config = (
             model_recipe_manager.get_quantization_configs(op_key, op_scope)
         )
-        if policy.is_conditionally_unquantized(op):
+        if policy.is_non_quantizable_composite_op(op):
           algorithm_name = algorithm_manager.AlgorithmName.NO_QUANTIZE
 
         if algorithm_name == algorithm_manager.AlgorithmName.NO_QUANTIZE:

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

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

@@ -2,15 +2,15 @@ ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas
 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.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=nKtghUjTQ8QS9CgLRwQb3iB2eZOyQv0FqyISlcgzSH4,11195
+ai_edge_quantizer/default_policy.py,sha256=9CNd5zIk_BA560kOLkoXD5mapDBWj0yXzSFDedLhzYw,11192
 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.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/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
@@ -70,8 +70,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EtOv6cpKM_F0uv2bWuSXylYm
 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.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,,
+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,,