ai-edge-quantizer-nightly 0.0.1.dev20250228__py3-none-any.whl → 0.0.1.dev20250301__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -22,8 +22,11 @@ from ai_edge_quantizer import default_policy
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.algorithms.nonlinear_quantize import float_casting
 from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
+from ai_edge_quantizer.algorithms.uniform_quantize import dequantized_weight_recovery
 from ai_edge_quantizer.algorithms.uniform_quantize import naive_min_max_quantize
 
+# TODO: b/399775701 - Clean up this file.
+
 _TFLOpName = qtyping.TFLOperationName
 
 _alg_manager_instance = algorithm_manager_api.AlgorithmManagerApi()
@@ -51,7 +54,9 @@ class AlgorithmName(str, enum.Enum):
   NO_QUANTIZE = "no_quantize"
   MIN_MAX_UNIFORM_QUANT = naive_min_max_quantize.ALGORITHM_KEY
   FLOAT_CASTING = float_casting.ALGORITHM_KEY
+  DEQUANTIZED_WEIGHT_RECOVERY = dequantized_weight_recovery.ALGORITHM_KEY
 
+### MIN/MAX_UNIFORM_QUANT ###
 
 # Register MIN_MAX_UNIFORM_QUANT algorithm.
 register_op_quant_config_validation_func(
@@ -109,7 +114,7 @@ for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
       ),
   )
 
-# Register FLOAT_CASTING algorithm.
+### FLOAT_CASTING ###
 register_op_quant_config_validation_func(
     AlgorithmName.FLOAT_CASTING,
     float_casting.check_op_quantization_config,
@@ -144,3 +149,38 @@ for op_name, materialize_func in zip(
       calibration_func=float_casting.calibrate,
       materialize_func=materialize_func,
   )
+
+### DEQUANTIZED_WEIGHT_RECOVERY ###
+register_op_quant_config_validation_func(
+    AlgorithmName.DEQUANTIZED_WEIGHT_RECOVERY,
+    common_quantize.check_op_quantization_config,
+)
+
+register_config_check_policy_func(
+    AlgorithmName.DEQUANTIZED_WEIGHT_RECOVERY,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+DEQUANTIZED_WEIGHT_RECOVERY_OP_NAME_MATERIALIZE_FUNC_DICT = {
+    _TFLOpName.FULLY_CONNECTED: common_quantize.materialize_fc_conv,
+    _TFLOpName.EMBEDDING_LOOKUP: common_quantize.materialize_embedding_lookup,
+}
+
+for (
+    op_name,
+    materialize_func,
+) in DEQUANTIZED_WEIGHT_RECOVERY_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      algorithm_key=AlgorithmName.DEQUANTIZED_WEIGHT_RECOVERY,
+      tfl_op_name=op_name,
+      init_qsv_func=dequantized_weight_recovery.init_qsvs,
+      calibration_func=dequantized_weight_recovery.calibrate,
+      # Most of the materialize op functions are common for all algorithms
+      # except for the function to get scale and zero point, i.e.,
+      # get_tensor_quant_params. So we use functools.partial here to pass in the
+      # common utility function and the algorithm-specific function.
+      materialize_func=functools.partial(
+          materialize_func,
+          dequantized_weight_recovery.get_tensor_quant_params,
+      ),
+  )

{ai_edge_quantizer_nightly-0.0.1.dev20250228.dist-info → ai_edge_quantizer_nightly-0.0.1.dev20250301.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=dOW_ex_KShoFaib7auISgsx5wRtxpaxhWrwA2nWgito,6003
+ai_edge_quantizer/algorithm_manager.py,sha256=qjpdEei5l-k5wgGDnelMGVDAYb7K7z8Nn9FIuZBWPA0,7534
 ai_edge_quantizer/algorithm_manager_api.py,sha256=u903TG0s1uIDhJqfeJne3CFl8A93phZrwgV2-hwdcXU,9247
 ai_edge_quantizer/algorithm_manager_api_test.py,sha256=tL_ozYFTsOPX8qGcti0KTz37nVsCxf0SSG5C45SyT-g,7319
 ai_edge_quantizer/calibrator.py,sha256=2J-bX0k09A7vZSRnO3eP49YO2uBMUQh6-sk3JRz9fGQ,11363
@@ -60,8 +60,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=SM8H4i7Jq_nfdsJpImopHndN
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=Op3JxtOqlrjzmYF18jnnstL1k9xiY9kKJ8S2vklKGkc,11327
 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.0.1.dev20250228.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.0.1.dev20250228.dist-info/METADATA,sha256=xlHVb_zTpsb3cqr55S8fCH6jp1T143R6tH6ND9j-iWk,1528
-ai_edge_quantizer_nightly-0.0.1.dev20250228.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.0.1.dev20250228.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.0.1.dev20250228.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.0.1.dev20250301.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.0.1.dev20250301.dist-info/METADATA,sha256=xyqDD73R7HgfjIXn-IS_FgqswM8sxMSMMbFwGSyI9sM,1528
+ai_edge_quantizer_nightly-0.0.1.dev20250301.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.0.1.dev20250301.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.0.1.dev20250301.dist-info/RECORD,,