ai-edge-quantizer-nightly 0.3.0.dev20250707__py3-none-any.whl → 0.3.0.dev20250709__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -114,6 +114,8 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
     _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
     _TFLOpName.GATHER_ND: common_quantize.materialize_gather_nd,
+    _TFLOpName.PACK: common_quantize.materialize_pack,
+    _TFLOpName.UNPACK: common_quantize.materialize_unpack,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -254,6 +256,8 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
     _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
     _TFLOpName.GATHER_ND: common_quantize.materialize_gather_nd,
+    _TFLOpName.PACK: common_quantize.materialize_pack,
+    _TFLOpName.UNPACK: common_quantize.materialize_unpack,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -762,6 +762,38 @@ def materialize_gather_nd(
   )
 
 
+def materialize_pack(
+    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.pack."""
+  return common_utils.materialize_standard_op(
+      op_info,
+      graph_info,
+      tensor_name_to_qsv,
+      get_tensor_quant_params_fn,
+      constraint=_OpQuantConstraint.SAME_AS_OUTPUT_SCALE,
+  )
+
+
+def materialize_unpack(
+    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.unpack."""
+  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/default_policy.py

@@ -187,7 +187,9 @@ DEFAULT_JSON_POLICY = """
       "PAD",
       "MAX_POOL_2D",
       "RESIZE_BILINEAR",
-      "GATHER_ND"
+      "GATHER_ND",
+      "PACK",
+      "UNPACK"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -223,7 +225,9 @@ DEFAULT_JSON_POLICY = """
       "SQUARED_DIFFERENCE",
       "MAX_POOL_2D",
       "RESIZE_BILINEAR",
-      "GATHER_ND"
+      "GATHER_ND",
+      "PACK",
+      "UNPACK"
     ],
     "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

@@ -67,6 +67,8 @@ class TFLOperationName(str, enum.Enum):
   MAX_POOL_2D = 'MAX_POOL_2D'
   RESIZE_BILINEAR = 'RESIZE_BILINEAR'
   GATHER_ND = 'GATHER_ND'
+  PACK = 'PACK'
+  UNPACK = 'UNPACK'
 
 
 class QuantizeMode(enum.Enum):

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -61,6 +61,8 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.MAX_POOL_2D: schema.BuiltinOperator.MAX_POOL_2D,
     _TFLOpName.RESIZE_BILINEAR: schema.BuiltinOperator.RESIZE_BILINEAR,
     _TFLOpName.GATHER_ND: schema.BuiltinOperator.GATHER_ND,
+    _TFLOpName.PACK: schema.BuiltinOperator.PACK,
+    _TFLOpName.UNPACK: schema.BuiltinOperator.UNPACK,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

ai_edge_quantizer/utils/tfl_interpreter_utils.py

@@ -346,7 +346,7 @@ def _create_random_bool(
     dtype: np.dtype,
 ) -> dict[str, Any]:
   """Creates a random bool dataset sample for given input details."""
-  return rng.choice([True, False], size=shape, replace=False).astype(dtype)
+  return rng.choice([True, False], size=shape, replace=True).astype(dtype)
 
 
 def create_random_dataset(

{ai_edge_quantizer_nightly-0.3.0.dev20250707.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250709.dist-info}/METADATA

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

@@ -1,18 +1,18 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=UZVS6ZClIAyaX9RzhXvvymbQv_scR0ybMPYl2CgSPVo,12346
+ai_edge_quantizer/algorithm_manager.py,sha256=-a3FMPJ-mLG20v5wO1CbzJryLlbWqqvTzxdyt0CafX0,12574
 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=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=0Am2TrgyV7gNl7dbul07rVp58OKDuPyJW9SIqRTrD2g,11280
+ai_edge_quantizer/default_policy.py,sha256=HHOPkBxuPP5NO-aLfqKCKyrmoZJ4vKT16vwn4d5BOFA,11340
 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=gC7G6Ne4Fumc8RSmIAbx96ZBhszZlHqBKSmE9p6RPTo,20099
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=vq-9jwDViSndHhcC1_RVu2Bk0qu5MgYPGLTRO9z2Naw,16604
+ai_edge_quantizer/qtyping.py,sha256=uCguCckBh9QTnsyCw-EX8IbXROAA5fwZ21_iX5Yu0z8,16640
 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=juXQWnzBTvasxydeDTNxWuE9ag9j6GOmfHMjC4JQu1Y,30799
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=uvPe3QGSMZSeVfHCt2DXhXf_iX7Gu4H-pDmlC_2p34E,31788
 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
@@ -62,14 +62,14 @@ ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V
 ai_edge_quantizer/utils/calibration_utils.py,sha256=e3dG7Nm94Ix0hkTWTWPUhEG6a8QR_cAM3PSwblfJV5g,15106
 ai_edge_quantizer/utils/calibration_utils_test.py,sha256=4BlksXl7b4yptL8xPR67hmJCnjhN9V10a2PunzfHrUE,9372
 ai_edge_quantizer/utils/test_utils.py,sha256=spqUmSNciOKPQHCBkHE7Zo34eMFq_BfBCAnMT3jAulU,8615
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=NnD57Gkx9upNP8Mso-_yp8Z3x1AqlIWb06jPg-hyRkc,10890
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=7r3XL2Gq44yQy6dY6mTr4YcP77nNZr2hHJmjFT74ar0,10994
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
-ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=vTyy6-4PgfFPL3C8uTq_iPFBwdxCjhrWzUiec4DdFPw,14323
+ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=SKSu1nqhGGzVei_DxmzXK-bbOE7G1vKnPDc5skce-yY,14322
 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.dev20250707.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.3.0.dev20250707.dist-info/METADATA,sha256=-L8UUcmsR22WOm48zZvdIhxZ7VSaziiRNOphHbYgDWE,1528
-ai_edge_quantizer_nightly-0.3.0.dev20250707.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.3.0.dev20250707.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.3.0.dev20250707.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.3.0.dev20250709.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.3.0.dev20250709.dist-info/METADATA,sha256=RP7kloAtphdBWt4SPfWUOHmLhNJDd2sI-vS07K4tBrg,1528
+ai_edge_quantizer_nightly-0.3.0.dev20250709.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.3.0.dev20250709.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.3.0.dev20250709.dist-info/RECORD,,