ai-edge-quantizer-nightly 0.4.0.dev20250904__py3-none-any.whl → 0.4.0.dev20250906__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -123,6 +123,7 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.GATHER: common_quantize.materialize_gather,
     _TFLOpName.HARD_SWISH: common_quantize.materialize_hard_swish,
     _TFLOpName.MAXIMUM: common_quantize.materialize_maximum,
+    _TFLOpName.PADV2: common_quantize.materialize_padv2,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -272,6 +273,7 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     _TFLOpName.GATHER: common_quantize.materialize_gather,
     _TFLOpName.HARD_SWISH: common_quantize.materialize_hard_swish,
     _TFLOpName.MAXIMUM: common_quantize.materialize_maximum,
+    _TFLOpName.PADV2: common_quantize.materialize_padv2,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -712,7 +712,24 @@ def materialize_pad(
       tensor_name_to_qsv,
       get_tensor_quant_params_fn,
       constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
-      inputs_to_ignore=[1],  # Padding value does not need to be quantized.
+      inputs_to_ignore=[1],  # Paddings tensor does not need to be quantized.
+  )
+
+
+def materialize_padv2(
+    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.padv2."""
+  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,
+      inputs_to_ignore=[1],  # Paddings tensor does not need to be quantized.
   )
 
 

ai_edge_quantizer/algorithms/utils/common_utils.py

@@ -41,6 +41,7 @@ _DRQ_OR_WEIGHT_ONLY_OPS = frozenset([
 
 _SUPPORTED_SUBCHANNEL_OPS = frozenset([
     _TFLOpName.FULLY_CONNECTED,
+    _TFLOpName.EMBEDDING_LOOKUP,
 ])
 
 
@@ -259,6 +260,60 @@ def _get_single_tensor_params(
   )
 
 
+def _materialize_tensors_with_quantized_data_update(
+    op_tensor_params: list[qtyping.TensorTransformationParams],
+    tensors: Sequence[Any],
+    quant_params: Optional[qtyping.UniformQuantParams],
+    is_inbounding_tensor: bool,
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: dict[str, Any],
+    get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
+) -> None:
+  """Materialize a list of tensors with `quantized_data` updated when needed.
+
+  Args:
+    op_tensor_params: Tensor transformation parameters for the op. Will be
+      modified to include new tensor parameters.
+    tensors: Tensors to be materialized.
+    quant_params: The quantization parameters to be used for materialization.
+    is_inbounding_tensor: Whether the tensor is an inbounding tensor for the op.
+    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.
+    get_tensor_quant_params_fn: Function to get quantization parameters for the
+      tensor.
+  """
+  if quant_params is not None and quant_params.quantized_data is not None:
+    quant_params = dataclasses.replace(quant_params, quantized_data=None)
+
+  for tensor in tensors:
+    tensor_data = tfl_flatbuffer_utils.get_tensor_data(
+        tensor, graph_info.buffers
+    )
+    if quant_params is None or tensor_data is None:
+      tensor_quant_params = quant_params
+    else:
+      # Constant tensors require updating `quantized_data`.
+      quantized_data = uniform_quantize_tensor.uniform_quantize(
+          tensor_data, quant_params
+      )
+      tensor_quant_params = dataclasses.replace(
+          quant_params,
+          quantized_data=quantized_data,
+      )
+    _materialize_op_tensors(
+        op_tensor_params,
+        [tensor],
+        is_inbounding_tensor=is_inbounding_tensor,
+        op_info=op_info,
+        graph_info=graph_info,
+        tensor_name_to_qsv=tensor_name_to_qsv,
+        get_tensor_quant_params_fn=get_tensor_quant_params_fn,
+        quant_params=tensor_quant_params,
+    )
+
+
 def _materialize_standard_op_with_same_as_input_scale(
     input_tensors: Sequence[Any],
     output_tensors: Sequence[Any],
@@ -293,45 +348,24 @@ def _materialize_standard_op_with_same_as_input_scale(
       get_tensor_quant_params_fn=get_tensor_quant_params_fn,
   )
   op_tensor_params.append(input_tensor_params)
-  # Use input quantization params for all output tensors but without
-  # quantized_data in case the input is a constant tensor.
-  input_quant_params = dataclasses.replace(
-      input_tensor_params.consumers[0].parameters,
-      quantized_data=None,
-  )
+  # Use input quantization params for all output tensors.
+  input_quant_params = input_tensor_params.consumers[0].parameters
   if not isinstance(input_quant_params, qtyping.UniformQuantParams):
     raise ValueError(
         "_materialize_standard_op_with_same_as_input_scale only supports"
         f" UniformQuantParams. For tensor {input_tensor_params.tensor_name},"
         f" got {type(input_quant_params)}"
     )
-  # Materialize each of the output tensors separately in case there are
-  # constants among them, requiring updating `quantized_data` first.
-  for output_tensor in output_tensors:
-    output_tensor_data = tfl_flatbuffer_utils.get_tensor_data(
-        output_tensor, graph_info.buffers
-    )
-    # Quantize constant inputs' data with the output quantization params.
-    if output_tensor_data is None:
-      quant_params = input_quant_params
-    else:
-      quantized_data = uniform_quantize_tensor.uniform_quantize(
-          output_tensor_data, input_quant_params
-      )
-      quant_params = dataclasses.replace(
-          input_quant_params,
-          quantized_data=quantized_data,
-      )
-    _materialize_op_tensors(
-        op_tensor_params,
-        [output_tensor],
-        is_inbounding_tensor=False,
-        op_info=op_info,
-        graph_info=graph_info,
-        tensor_name_to_qsv=tensor_name_to_qsv,
-        get_tensor_quant_params_fn=get_tensor_quant_params_fn,
-        quant_params=quant_params,
-    )
+  _materialize_tensors_with_quantized_data_update(
+      op_tensor_params,
+      output_tensors,
+      input_quant_params,
+      is_inbounding_tensor=False,
+      op_info=op_info,
+      graph_info=graph_info,
+      tensor_name_to_qsv=tensor_name_to_qsv,
+      get_tensor_quant_params_fn=get_tensor_quant_params_fn,
+  )
 
   # Change output qsv to be the same as input qsv. This is safe since TFL
   # subgraph is acyclic.
@@ -379,19 +413,26 @@ def _materialize_standard_op_with_same_as_output_scale(
   )
   # Use output quantization params for all input tensors.
   if output_tensor_params.producer is None:
-    quant_params = None
+    output_quant_params = None
   else:
-    quant_params = output_tensor_params.producer.parameters
-  _materialize_op_tensors(
+    output_quant_params = output_tensor_params.producer.parameters
+    if not isinstance(output_quant_params, qtyping.UniformQuantParams):
+      raise ValueError(
+          "_materialize_standard_op_with_same_as_output_scale only supports"
+          f" UniformQuantParams. For tensor {output_tensor_params.tensor_name},"
+          f" got {type(output_quant_params)}"
+      )
+  _materialize_tensors_with_quantized_data_update(
       op_tensor_params,
       input_tensors,
+      output_quant_params,
       is_inbounding_tensor=True,
       op_info=op_info,
       graph_info=graph_info,
       tensor_name_to_qsv=tensor_name_to_qsv,
       get_tensor_quant_params_fn=get_tensor_quant_params_fn,
-      quant_params=quant_params,
   )
+
   op_tensor_params.append(output_tensor_params)
 
   return op_tensor_params

ai_edge_quantizer/default_policy.py

@@ -195,7 +195,8 @@ DEFAULT_JSON_POLICY = """
       "BROADCAST_TO",
       "SQRT",
       "GATHER",
-      "MAXIMUM"
+      "MAXIMUM",
+      "PADV2"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -240,7 +241,8 @@ DEFAULT_JSON_POLICY = """
       "SQRT",
       "GATHER",
       "HARD_SWISH",
-      "MAXIMUM"
+      "MAXIMUM",
+      "PADV2"
     ],
     "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

@@ -76,6 +76,7 @@ class TFLOperationName(str, enum.Enum):
   GATHER = 'GATHER'
   HARD_SWISH = 'HARD_SWISH'
   MAXIMUM = 'MAXIMUM'
+  PADV2 = 'PADV2'
 
 
 class QuantizeMode(enum.Enum):

ai_edge_quantizer/utils/constrained_ops_utils_test.py

@@ -33,7 +33,7 @@ class ConstrainedOpsUtilsTest(parameterized.TestCase):
       dict(
           testcase_name="same_as_output_scale",
           constraint=_OpQuantConstraint.SAME_AS_OUTPUT_SCALE,
-          expected_num_ops=6,
+          expected_num_ops=7,
       ),
       dict(
           testcase_name="no_constrain",

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -70,6 +70,7 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.GATHER: schema.BuiltinOperator.GATHER,
     _TFLOpName.HARD_SWISH: schema.BuiltinOperator.HARD_SWISH,
     _TFLOpName.MAXIMUM: schema.BuiltinOperator.MAXIMUM,
+    _TFLOpName.PADV2: schema.BuiltinOperator.PADV2,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

{ai_edge_quantizer_nightly-0.4.0.dev20250904.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250906.dist-info}/METADATA

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

@@ -1,18 +1,18 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=O_psY-4R0ARmgTQHwfH2px81AJY8PmfamHtE7xJDRjQ,13424
+ai_edge_quantizer/algorithm_manager.py,sha256=3kmn-hTLEhHOfAQTkUoN8xXymFtoljzLU-ADpd7uBrE,13538
 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=G_JZtZaQAnrWyfCusDWXwO27iLysk27RS91GlS61m_Q,11592
+ai_edge_quantizer/default_policy.py,sha256=6kEYu0nOQqBKpclzgmxuzvatiVR0BF_ce6zoKCoudW4,11622
 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=hcgMHJlERZERUyIAEi6AHJcLJ8gsKIBAEojzFFz-tqk,20098
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=zXXmLBZUT-cfjnQrqDkytDZaGg3z_yy1wWhKr34_XVg,16792
+ai_edge_quantizer/qtyping.py,sha256=ygLmj_PPTYM1yAs3oCJ649q75cZPQYjJ8hXtFLVqfv8,16810
 ai_edge_quantizer/quantizer.py,sha256=ckAEOnnBxuCKZuvlzdChevCKPuE-IeDPHCNtFTWr250,17857
 ai_edge_quantizer/quantizer_test.py,sha256=m6f4ayyaF3yQb9i4V0aFAbmGw0OKZ2Zam1RoTPh-u24,22917
 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=rkf7jLPVDKpx2ju1LyyP7bxc6n34cLD2E3w2mxLd6qE,35344
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=TQQxkxeAngrZO6ro6RjOeJAieWHIgK4hrACtbU0-Buk,35919
 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
@@ -41,7 +41,7 @@ ai_edge_quantizer/algorithms/uniform_quantize/octav_test.py,sha256=sha1d99Xk87bI
 ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=uCREMXi0U2ckhXXfgGVzwSgjFZc0IbtnFU-OjlG9IO8,17146
 ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=7kHluzpteMv36hFD6LD_qnwwMoE1GKUP4bGmGMFbOdA,12755
 ai_edge_quantizer/algorithms/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/algorithms/utils/common_utils.py,sha256=QrEeCuvA7gY_vK1nbKtqassNDClyAjN1ClZIiw63k5U,35895
+ai_edge_quantizer/algorithms/utils/common_utils.py,sha256=4eAlGph6DDW18bUdoY0XcUoOXEr3P_3_W1ptidD8qK4,37611
 ai_edge_quantizer/algorithms/utils/common_utils_test.py,sha256=zqapGEfYhjQWe9cNGPLmdbwtEUUYQRhlO_kNe0cXX6E,18104
 ai_edge_quantizer/transformations/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/transformations/dequant_insert.py,sha256=sL1LHFVzBDSd9jgrzlHz38LWU0bwmVX7iBkaNcui0ts,3566
@@ -62,16 +62,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=6k_AqfB-NmiLkW5WwEV5NSuswFWky2sL0xBGmV6Fdwk,1756
+ai_edge_quantizer/utils/constrained_ops_utils_test.py,sha256=xWujKhNR_OFXReFM-njFbiaC_4W7kMNr7lmFFRlGNLw,1756
 ai_edge_quantizer/utils/test_utils.py,sha256=a4Nk-wbeB09dFjTDZiA0K67d26j5DD0UDH_GIVmVG_4,8685
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=aNtL4dpWH5uGGGlaygnMDkh5llTstbgs5ZxO0JkH5VQ,11718
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=1NnRPdvqdvZ5sKbIdePcBv8SaCS2LqZXX_B51oDRXrQ,11770
 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=oYw33Sg547AqtGw-choPUJmp9SAKkV46J_ddqSsum2Q,3950
 ai_edge_quantizer/utils/validation_utils_test.py,sha256=V_qNDikPD4OPB-siOLQCWNVWTAu87h2IgNYt7teFd-o,2934
-ai_edge_quantizer_nightly-0.4.0.dev20250904.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20250904.dist-info/METADATA,sha256=AqVFUY7NzeZmDKSjDGACfkR1HIsitVwnnJ91NlaktX0,1508
-ai_edge_quantizer_nightly-0.4.0.dev20250904.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20250904.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20250904.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20250906.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20250906.dist-info/METADATA,sha256=hywZ1CdX4MVCVLEbrx_plaGJL9tEXziGga_psNOiUsc,1508
+ai_edge_quantizer_nightly-0.4.0.dev20250906.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20250906.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20250906.dist-info/RECORD,,