ai-edge-quantizer-nightly 0.4.0.dev20251013__py3-none-any.whl → 0.4.0.dev20251014__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -132,6 +132,7 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.EQUAL: common_quantize.materialize_equal,
     _TFLOpName.NOT_EQUAL: common_quantize.materialize_not_equal,
     _TFLOpName.MIRROR_PAD: common_quantize.materialize_mirror_pad,
+    _TFLOpName.SPACE_TO_DEPTH: common_quantize.materialize_space_to_depth,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -286,6 +287,7 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     _TFLOpName.EQUAL: common_quantize.materialize_equal,
     _TFLOpName.NOT_EQUAL: common_quantize.materialize_not_equal,
     _TFLOpName.MIRROR_PAD: common_quantize.materialize_mirror_pad,
+    _TFLOpName.SPACE_TO_DEPTH: common_quantize.materialize_space_to_depth,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -776,6 +776,33 @@ def materialize_mirror_pad(
   )
 
 
+def materialize_space_to_depth(
+    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.space_to_depth.
+
+  Args:
+    get_tensor_quant_params_fn: Function to get quantization parameters for the
+      tensor.
+    op_info: Aggregated information about the op (e.g., quantization config).
+    graph_info: Graph information needed to perform quantization for the op.
+    tensor_name_to_qsv: A map of tensor name to quantization parameters.
+
+  Returns:
+    A list of `qtyping.TensorTransformationParams` for the tensors in the op.
+  """
+  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 materialize_squared_difference(
     get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
     op_info: qtyping.OpInfo,

ai_edge_quantizer/default_policy.py

@@ -248,7 +248,8 @@ DEFAULT_JSON_POLICY = """
       "REDUCE_MIN",
       "EQUAL",
       "NOT_EQUAL",
-      "MIRROR_PAD"
+      "MIRROR_PAD",
+      "SPACE_TO_DEPTH"
     ],
     "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT"],
     "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT"],

ai_edge_quantizer/qtyping.py

@@ -81,6 +81,7 @@ class TFLOperationName(str, enum.Enum):
   EQUAL = 'EQUAL'
   NOT_EQUAL = 'NOT_EQUAL'
   MIRROR_PAD = 'MIRROR_PAD'
+  SPACE_TO_DEPTH = 'SPACE_TO_DEPTH'
 
 
 class QuantizeMode(enum.Enum):

ai_edge_quantizer/utils/constrained_ops_utils_test.py

@@ -28,7 +28,7 @@ class ConstrainedOpsUtilsTest(parameterized.TestCase):
       dict(
           testcase_name="same_as_input_scale",
           constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
-          expected_num_ops=16,
+          expected_num_ops=17,
       ),
       dict(
           testcase_name="same_as_output_scale",

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -75,6 +75,7 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.EQUAL: schema.BuiltinOperator.EQUAL,
     _TFLOpName.NOT_EQUAL: schema.BuiltinOperator.NOT_EQUAL,
     _TFLOpName.MIRROR_PAD: schema.BuiltinOperator.MIRROR_PAD,
+    _TFLOpName.SPACE_TO_DEPTH: schema.BuiltinOperator.SPACE_TO_DEPTH,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

{ai_edge_quantizer_nightly-0.4.0.dev20251013.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20251014.dist-info}/METADATA

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

@@ -1,18 +1,18 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=Ri4bNqbSTmtlsYZiJYHtkjNEsl8h5tZ_1uV_stJ3HUY,16156
+ai_edge_quantizer/algorithm_manager.py,sha256=Eev0uVw19oSPrpUJ4sU4CrorSC3h6bsiUEYCOJw1Hsk,16306
 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=ZLzIMWB2FSFU4TOatDioYuwp_kLh8iSCefZ5_Q9FU7s,11900
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
-ai_edge_quantizer/default_policy.py,sha256=DpQUe0ETpRtx2qWwvrj9pFAhfyEwES1mVTTak8kyPzM,11678
+ai_edge_quantizer/default_policy.py,sha256=i_AcnIIElHqoJNc2jyJFEC2tYqfQ4Nvn4mQlTvormzk,11702
 ai_edge_quantizer/model_modifier.py,sha256=U70JByv6CItP8tg4bdyMfX-R3UlwylAGSviZkF_FSAM,10468
 ai_edge_quantizer/model_modifier_test.py,sha256=CV4pgMEQkBJr_qbYR720TO8HBCutbEYLHptDHgdQMUE,7274
 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=0w-sDGk84sVNkXoduon1wDqq30sGOHVgBVbdg44QVF4,20153
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=rp2jdmCuSsP6Ay8rD7NxDCpbFkRNkbYP29Uwe0xBfnA,17196
+ai_edge_quantizer/qtyping.py,sha256=f2NRz4xqM-7gMe0QFpR4x2m5lzTJI3tmsT0cehO5Vsg,17232
 ai_edge_quantizer/quantizer.py,sha256=teYeONdIS31IAY6ubLujCRi1t6lYAd0LkC8dRPxQdbw,18919
 ai_edge_quantizer/quantizer_test.py,sha256=9BVwt7oyM8IsSC7jN73nI0O-4MikBkymm_FigJnSeCM,27117
 ai_edge_quantizer/recipe.py,sha256=MEkfQ2Sg3KAE9LAORHWcbjYNPg06EUbwc1d-VspQA2U,6461
@@ -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=rUHraqNi0iJ0AyUQfAiYkyXG2rSy7GlhnqmwJeoLStg,38952
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=HF7aNccdDmCbZGZ21UxeO5UpSpQOLr3TiOEyLwWOVPQ,39888
 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
@@ -66,16 +66,16 @@ ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V
 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=HNZstSm6-7xgSmM-7ilHRjuOKsq6tivpxayphm9Oghs,1756
+ai_edge_quantizer/utils/constrained_ops_utils_test.py,sha256=i_uERo-KvMj0dvUSuI67kdOBHvRQETg8-qnejs_MgTE,1756
 ai_edge_quantizer/utils/test_utils.py,sha256=a4Nk-wbeB09dFjTDZiA0K67d26j5DD0UDH_GIVmVG_4,8685
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=LN-WonrcJLP9bB4lULd5VIg_8YLTcp891ZuDZ5nDGe8,12006
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=42OWzQsRTXq3XQYmoxlz177_dw2fJfq7mDSJaU--ArQ,12076
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EoVjI_hplX_Rml3hfRsGmQOihexmizeJqt4SQcET9aA,14925
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=6fjkM-rycZ95L4yfvlr0TN6RlrhfPzxNUYrZaYO_F0A,12013
 ai_edge_quantizer/utils/validation_utils.py,sha256=0sOdH4pzk_Pwh1r8O47iaECRng1Xn0ABn9GVc8UPNcY,4994
 ai_edge_quantizer/utils/validation_utils_test.py,sha256=1sblJWHLTYTbn1Qi9rwnrREOSXRy5KwHAWSwgI1e_aU,3697
-ai_edge_quantizer_nightly-0.4.0.dev20251013.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20251013.dist-info/METADATA,sha256=_PaC8xPzMOnl2xLDR39RJ3OC--oC7tdJXwh3yRMA_5U,1508
-ai_edge_quantizer_nightly-0.4.0.dev20251013.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20251013.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20251013.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20251014.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20251014.dist-info/METADATA,sha256=TrlrIOhZNaILuX9hEN1mHXRXKNkHHZE2KubEuGV2oEQ,1508
+ai_edge_quantizer_nightly-0.4.0.dev20251014.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20251014.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20251014.dist-info/RECORD,,