ai-edge-quantizer-nightly 0.1.0.dev20250327__py3-none-any.whl → 0.1.0.dev20250329__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -170,6 +170,7 @@ register_config_check_policy_func(
 
 DEQUANTIZED_WEIGHT_RECOVERY_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.FULLY_CONNECTED: common_quantize.materialize_fc_conv,
+    _TFLOpName.CONV_2D: common_quantize.materialize_fc_conv,
     _TFLOpName.EMBEDDING_LOOKUP: common_quantize.materialize_embedding_lookup,
 }
 

ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py

@@ -70,17 +70,17 @@ def _get_scale(arr: np.ndarray, min_scale: float) -> float:
   return min_scale
 
 
-def get_zp_scale_from_2d_dequantized_symmetric_weights(
+def get_zp_scale_from_dequantized_symmetric_weights(
     dequant_vals: np.ndarray,
     quantized_dimension: Optional[int] = None,
     min_scale: float = 1e-9,
 ) -> tuple[np.ndarray, np.ndarray]:
-  """Calculates scale and zero point from 2D dequantized, symmetric weights.
+  """Calculates scale and zero point from dequantized and symmetric weights.
 
   Handles both per-tensor and per-channel (axis) quantization.
 
   Args:
-      dequant_vals: The 2D dequantized weight values (numpy array).
+      dequant_vals: The dequantized weight values (numpy array).
       quantized_dimension:  The dimension along which quantization was performed
         (0 or 1), or None for per-tensor quantization.
       min_scale: The minimum allowed scale value.
@@ -91,15 +91,9 @@ def get_zp_scale_from_2d_dequantized_symmetric_weights(
           - scales: Scales (scalar for per-tensor, array for per-channel).
 
   Raises:
-      ValueError: If `dequant_vals` is not 2D, or if
-          `quantized_dimension` is not 0, 1, or None.
+      ValueError: If `quantized_dimension` is not 0, 1, or None.
   """
 
-  if dequant_vals.ndim != 2:
-    raise ValueError(
-        f"Only 2D weights are supported. Got {dequant_vals.ndim} dimensions."
-    )
-
   if quantized_dimension not in (0, 1, None):
     raise ValueError(
         f"quantized_dimension must be 0, 1, or None. Got {quantized_dimension}"
@@ -112,23 +106,26 @@ def get_zp_scale_from_2d_dequantized_symmetric_weights(
     # Per-tensor quantization: One scale for the entire tensor.
     scales = _get_scale(dequant_vals.flatten(), min_scale)
     scales = np.array([[scales]])
-
   else:
     # Per-channel quantization: A scale for each slice along the dimension.
-    scales = []
-    for i in range(dequant_vals.shape[quantized_dimension]):
-      if quantized_dimension == 0:
-        vec = dequant_vals[i, :]
-      else:  # quantized_dimension == 1
-        vec = dequant_vals[:, i]
-      scales.append(_get_scale(vec, min_scale))
-
-    # Reshape for correct broadcasting.
-    scales = (
-        np.array(scales).reshape(-1, 1)
-        if quantized_dimension == 0
-        else np.array(scales).reshape(1, -1)
+    # Create a broadcasted array for per-channel scales. It should have the same
+    # number of dimensions as the input, with 1 in all dimensions except for the
+    # quantized dimension, which retains its original size.
+    scales = np.empty(
+        tuple(
+            [
+                1
+                if i != quantized_dimension
+                else dequant_vals.shape[quantized_dimension]
+                for i in range(dequant_vals.ndim)
+            ]
+        )
     )
+    for i in range(dequant_vals.shape[quantized_dimension]):
+      slices = [slice(None)] * dequant_vals.ndim
+      slices[quantized_dimension] = i
+      vec = dequant_vals[tuple(slices)]
+      scales[tuple(slices)] = _get_scale(vec, min_scale)
 
   zero_points = np.zeros_like(scales, dtype=np.int32)
   return zero_points, scales
@@ -153,7 +150,7 @@ def get_tensor_quant_params(
 
   Raises:
     ValueError: If the quantization granularity is blockwise, or if the tensor
-    is not a 2D symmetric weight tensor.
+    is not a symmetric weight tensor.
   """
   # Fallback to naive_min_max_quantize.py for non-weight tensors.
   if tensor_content is None:
@@ -166,10 +163,9 @@ def get_tensor_quant_params(
         "Blockwise quantization is not supported for dequantized weight"
         " recovery."
     )
-  if tensor_content.ndim != 2 or not tensor_quant_config.symmetric:
+  if not tensor_quant_config.symmetric:
     raise ValueError(
-        "Only 2D symmetric weights are supported for dequantized weight"
-        " recovery."
+        "Only symmetric weights are supported for dequantized weight recovery."
     )
 
   quantized_dim = None
@@ -178,7 +174,7 @@ def get_tensor_quant_params(
         op_info, tensor_content
     )
 
-  zp, scale = get_zp_scale_from_2d_dequantized_symmetric_weights(
+  zp, scale = get_zp_scale_from_dequantized_symmetric_weights(
       dequant_vals=tensor_content,
       quantized_dimension=quantized_dim,
   )

ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py

@@ -62,7 +62,7 @@ class DequantizedWeightRecoveryTest(parameterized.TestCase):
   ):
     dequant_vals = scale * self._dummy_quantized_weights
     zp, recovered_scale = (
-        dequantized_weight_recovery.get_zp_scale_from_2d_dequantized_symmetric_weights(
+        dequantized_weight_recovery.get_zp_scale_from_dequantized_symmetric_weights(
             dequant_vals, quantized_dimension
         )
     )
@@ -72,17 +72,40 @@ class DequantizedWeightRecoveryTest(parameterized.TestCase):
     self.assertEqual(np.sum(zp), 0)
     self.assertEqual(zp.shape, scale.shape)
 
-  def test_tensor_zp_scale_from_2d_dequantized_symmetric_weights_raises_error_for_non_2d_weights(
-      self,
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="per-tensor-recovery",
+          quantized_dimension=None,
+          scale=np.array([0.1875]).reshape(1, 1),
+      ),
+      dict(
+          testcase_name="channel0-recovery",
+          quantized_dimension=0,
+          scale=np.array([0.1875, 1e-4, 12.3]).reshape(3, 1, 1),
+      ),
+      dict(
+          testcase_name="channel1-recovery",
+          quantized_dimension=1,
+          scale=np.array([0.003, 1.234]).reshape(1, 2, 1),
+      ),
+  )
+  def test_tensor_zp_scale_from_3d_dequantized_symmetric_weights_success(
+      self, quantized_dimension, scale
   ):
-    weights_3d = self._dummy_quantized_weights.reshape(1, 3, 4)
-    weights_3d = weights_3d * 1.02
-    with self.assertRaisesRegex(
-        ValueError, "Only 2D weights are supported. Got 3 dimensions."
-    ):
-      dequantized_weight_recovery.get_zp_scale_from_2d_dequantized_symmetric_weights(
-          weights_3d, quantized_dimension=None
-      )
+    dequant_vals = scale * self._dummy_quantized_weights.reshape(3, 2, 2)
+    zp, recovered_scale = (
+        dequantized_weight_recovery.get_zp_scale_from_dequantized_symmetric_weights(
+            dequant_vals, quantized_dimension
+        )
+    )
+    with self.subTest("shapes_match"):
+      self.assertEqual(recovered_scale.shape, scale.shape)
+      self.assertEqual(zp.shape, scale.shape)
+    with self.subTest("scale_value_match"):
+      self.assertSequenceAlmostEqual(recovered_scale.flatten(), scale.flatten())
+    with self.subTest("zp_is_zero"):
+      # Zero point should be zero for symmetric quantization.
+      self.assertEqual(np.sum(zp), 0)
 
   @parameterized.named_parameters(
       dict(testcase_name="negative_dimension", quantized_dimension=-1),
@@ -95,7 +118,7 @@ class DequantizedWeightRecoveryTest(parameterized.TestCase):
     with self.assertRaisesRegex(
         ValueError, "quantized_dimension must be 0, 1, or None. Got"
     ):
-      dequantized_weight_recovery.get_zp_scale_from_2d_dequantized_symmetric_weights(
+      dequantized_weight_recovery.get_zp_scale_from_dequantized_symmetric_weights(
           dequant_vals, quantized_dimension
       )
 

ai_edge_quantizer/params_generator.py

@@ -394,14 +394,6 @@ def _compatible_tensor_params(
   ]
   if _same_tensor_params_except_id(params1, params2):
     return True
-  if (
-      params1.transformations[0] != _QuantTrans.NO_QUANTIZE
-      and params2.transformations[0] != _QuantTrans.NO_QUANTIZE
-  ):
-    # NO_QUANTIZE has no parameters. So only if both params aren't NO_QUANTIZE
-    # do we expect the parameters to be the same.
-    if params1.parameters != params2.parameters:
-      return False
   # We only need to check the first transformation because transformations are
   # applied in order, and as long as the one that's immediately after the tensor
   # is the same, it's compatible.
@@ -413,6 +405,7 @@ def _compatible_tensor_params(
   if (
       params1.transformations[0] in quantized_source_transformations
       and params2.transformations[0] in quantized_source_transformations
+      and params1.parameters == params2.parameters
   ):
     return True
   return False

ai_edge_quantizer/params_generator_test.py

@@ -914,6 +914,58 @@ class ParamsGeneratorTest(parameterized.TestCase):
           ),
           expected=True,
       ),
+      dict(
+          testcase_name='compatible_no_numeric_check',
+          param1=qtyping.TensorTransformationParams(
+              tensor_name='tfl.quantize',
+              producer=None,
+              consumers=[
+                  qtyping.OpToTensorParams(
+                      subgraph_op_id=4,
+                      transformations=[
+                          qtyping.QuantTransformation.ADD_QUANTIZE,
+                      ],
+                      parameters=qtyping.UniformQuantParams(
+                          8, None, np.array([0.00028806]), np.array([0])
+                      ),
+                  ),
+                  qtyping.OpToTensorParams(
+                      subgraph_op_id=5,
+                      transformations=[
+                          qtyping.QuantTransformation.ADD_QUANTIZE,
+                      ],
+                      parameters=qtyping.UniformQuantParams(
+                          8, None, np.array([0.00027501]), np.array([0])
+                      ),
+                  ),
+              ],
+          ),
+          param2=qtyping.TensorTransformationParams(
+              tensor_name='tfl.quantize',
+              producer=None,
+              consumers=[
+                  qtyping.OpToTensorParams(
+                      subgraph_op_id=4,
+                      transformations=[
+                          qtyping.QuantTransformation.ADD_QUANTIZE,
+                      ],
+                      parameters=qtyping.UniformQuantParams(
+                          8, None, np.array([0.00028806]), np.array([0])
+                      ),
+                  ),
+                  qtyping.OpToTensorParams(
+                      subgraph_op_id=5,
+                      transformations=[
+                          qtyping.QuantTransformation.ADD_QUANTIZE,
+                      ],
+                      parameters=qtyping.UniformQuantParams(
+                          8, None, np.array([0.00027501]), np.array([0])
+                      ),
+                  ),
+              ],
+          ),
+          expected=True,
+      ),
   )
   def test_params_compatible(self, param1, param2, expected):
     # adding a test to make production coverage happy.

ai_edge_quantizer/transformations/quantize_tensor.py

@@ -121,6 +121,26 @@ def _perform_channelwise_quantization(
   return flatbuffer_quantization
 
 
+def _downcast_and_truncate_scale(input_scale: np.ndarray) -> np.ndarray:
+  """Given a fp32 scale, downcast it to fp16 and truncate mantissa to 7 bits.
+
+  CPU kernel can only utilize 7 bits of mantissa for fp16, so we want to produce
+  scale this way to unify behaviours across different platforms.
+
+  Args:
+    input_scale: The input scale in fp32.
+
+  Returns:
+    The downcasted & truncated scale in fp16.
+  """
+
+  # A regular fp16 has 10 bits of mantissa, so we need to zero out the 3 least
+  # significant bits.
+  return (
+      input_scale.astype(np.float16).view(dtype=np.uint16) & np.uint16(0xFFF8)
+  ).view(dtype=np.float16)
+
+
 def _perform_blockwise_quantization(
     transformation_input: transformation_utils.TransformationInput,
 ) -> schema_py_generated.QuantizationParametersT():
@@ -142,9 +162,13 @@ def _perform_blockwise_quantization(
   )
   tensor = transformation_input.subgraph.tensors[transformation_input.tensor_id]
   blockwise_details = schema_py_generated.BlockwiseQuantizationT()
+  # Downcast and truncate the scale to fp16.
+  downcasted_scale = _downcast_and_truncate_scale(
+      transformation_input.quant_params.scale
+  )
   scale_tensor_id = transformation_utils.add_new_constant_tensor(
       tensor.name + b"_scales",
-      transformation_input.quant_params.scale.astype(np.float16),
+      downcasted_scale,
       schema_py_generated.TensorType.FLOAT16,
       transformation_input.subgraph,
       transformation_input.buffers,

{ai_edge_quantizer_nightly-0.1.0.dev20250327.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250329.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=sOZ1T8n0YYi_ijDDuzryNJi2HUPggeo9uWNJri3elv0,10431
+ai_edge_quantizer/algorithm_manager.py,sha256=0uootLsVD6h9ph9TrnXZMI-ExkX8UvXSV0lbWxBLybU,10492
 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=n7AD9j7UScR-CieoI6DQRMeiG_fhLBfSLRiM4460xaM,11895
@@ -10,8 +10,8 @@ ai_edge_quantizer/model_modifier.py,sha256=SPt9X-xBzRvcd4xIS24zLHt3aUS2QwsNDqweF
 ai_edge_quantizer/model_modifier_test.py,sha256=cJd04SLOG-fQZZNZPcisoBLx3cLtWEwGqUBbLb-pif4,4751
 ai_edge_quantizer/model_validator.py,sha256=fRNz0jO54cthPTibsCuViUXUuFRHl_fbvEiCukIVy20,13030
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
-ai_edge_quantizer/params_generator.py,sha256=f-KhJMFdRv2oHxfM8tAANPOtfBMw8vD7Vjv0rYQbnF4,16062
-ai_edge_quantizer/params_generator_test.py,sha256=zmDS6jG5zKhHL_hzJw2wlMTx1LLcNCK6S5WlwogWF-A,41122
+ai_edge_quantizer/params_generator.py,sha256=46XDjnP4R3m4xsoXNp7brv0sNQPdQMg217_CbEl-Wgg,15780
+ai_edge_quantizer/params_generator_test.py,sha256=9WTUl87XqbM4NruX5ypLuVRtuhcw-CmxndsMOUzZ92Q,43171
 ai_edge_quantizer/qtyping.py,sha256=UBZ3HgO8IDLY6VJmO05rGtFv_idMD3Os3WWsnriA0NA,15235
 ai_edge_quantizer/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
 ai_edge_quantizer/quantizer_test.py,sha256=K_HBA56JkFI3HL8VLWCqGEfC0ISh5ldMKoNyBdGRAJg,20368
@@ -30,8 +30,8 @@ ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py,sha256=s64
 ai_edge_quantizer/algorithms/uniform_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=SVu1RSX5xOWhuNEi9hHqgIDGe_ywyHBZAczp7KAcl3k,27220
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py,sha256=qMmKbWqxrCoVKbLKHn9WuCrGKPfHkEyU0Nmhokh8Qeo,2597
-ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=OTXjEZ3Ctq3ffYzisX-6HwgK_DuA7uos_aap5PiIUPE,8686
-ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py,sha256=y7BK11fkF63Ex_Jzg3fbIdy0D_Ca6HuvChVZR7Uwggc,8073
+ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=Fk3s9Qy2A_hjUepFOUmTwIZ_wKYVPbdDX4eoP-eoAQU,8726
+ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py,sha256=sT5eX5TLZEHTtPfnSkCPDlS0sQxlTFWbCsbvOuj--yY,8889
 ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py,sha256=cbyyYAoQnEraOYSV00wZ557ElBndHduVGeHikYUEFCE,7995
 ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py,sha256=B30SEISYZ9DPs3suKeG2elgXylR98pCEMWSEGgZo20o,7648
 ai_edge_quantizer/algorithms/uniform_quantize/octav.py,sha256=e5wYtki-vl739gSVAZHAKcs2hA87GvFUjVoSUPlnkyM,6433
@@ -50,7 +50,7 @@ ai_edge_quantizer/transformations/emulated_subchannel.py,sha256=HVaRxoC8PCAvy3xe
 ai_edge_quantizer/transformations/emulated_subchannel_test.py,sha256=gZP6u9NdPXl7s19qB_Un8evou9ZZV6I9Gy0E1rdobHM,7722
 ai_edge_quantizer/transformations/quant_insert.py,sha256=jn6HsJaV-sqBiFPY-Aqbd64t8zgcYVkEkZI375x_FWY,3958
 ai_edge_quantizer/transformations/quant_insert_test.py,sha256=X9ptPDvJCFkR5tejKnD1SlHFGPazQTW-wNNMV9MEAuw,10107
-ai_edge_quantizer/transformations/quantize_tensor.py,sha256=vzKtrXILqVsr1NGlribhdtKEIsXA93o37embLRe9TwQ,7493
+ai_edge_quantizer/transformations/quantize_tensor.py,sha256=y6As38mTzhva50YvNQ7p0SFpuWet3LPqFwE3qIO0gEQ,8231
 ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=mHLO3_MRt36A8-ZN8ADn5tBBJlqjTWa7ZUN8Mmu5Rcw,9116
 ai_edge_quantizer/transformations/transformation_utils.py,sha256=R42OIbzwQ7JYJ-Qt46jsqwb6u4MfDGiIPCRZCUGLVCw,4664
 ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=xH64SF3UHDh84vYbt-WvmXNjM-Jg-mefES1ACO1tkqw,6269
@@ -64,8 +64,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=x2xA2CFPpe_2trcV8v5xGaBE
 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.1.0.dev20250327.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250327.dist-info/METADATA,sha256=KES2W7tXAwTOpymOVvoQXovvZ2eaObKUUUT5L06a2gw,1527
-ai_edge_quantizer_nightly-0.1.0.dev20250327.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250327.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250327.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250329.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250329.dist-info/METADATA,sha256=3GoKN9dRRW0IhcJA2xnqAhU3znEgEzkDec2STODYtLA,1527
+ai_edge_quantizer_nightly-0.1.0.dev20250329.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250329.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250329.dist-info/RECORD,,