ai-edge-quantizer-nightly 0.1.0.dev20250328__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_nightly-0.1.0.dev20250328.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.dev20250328
+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.dev20250328.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
@@ -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
@@ -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.dev20250328.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250328.dist-info/METADATA,sha256=ajE8NjYUjSmr3VZQWBzARMc9MLqLkeHooAmK-fCg7Ms,1527
-ai_edge_quantizer_nightly-0.1.0.dev20250328.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250328.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250328.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,,