ai-edge-quantizer-nightly 0.4.0.dev20251006__py3-none-any.whl → 0.4.0.dev20251008__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py

@@ -119,6 +119,7 @@ def fix_quantization_params_rank(
       symmetric=quantization_params.symmetric,
       quantized_dimension=quantization_params.quantized_dimension,
       quantized_data=quantization_params.quantized_data,
+      block_size=quantization_params.block_size,
   )
 
 
@@ -209,13 +210,16 @@ def _broadcast_scale_zp_for_blockwise(
       ),
       tensor_content.shape,
   )
-  expanded_zp = np.reshape(
-      np.broadcast_to(
-          np.expand_dims(quant_params.zero_point, quantized_dim + 1),
-          expanded_tensor_shape,
-      ),
-      tensor_content.shape,
-  )
+  if quant_params.zero_point is None or quant_params.zero_point.size == 0:
+    expanded_zp = np.zeros_like(tensor_content, dtype=np.int32)
+  else:
+    expanded_zp = np.reshape(
+        np.broadcast_to(
+            np.expand_dims(quant_params.zero_point, quantized_dim + 1),
+            expanded_tensor_shape,
+        ),
+        tensor_content.shape,
+    )
   return qtyping.UniformQuantParams(
       scale=expanded_scale,
       zero_point=expanded_zp,
@@ -295,6 +299,26 @@ def uniform_dequantize(
   Returns:
     The dequantized tensor.
   """
+  if quantization_params.block_size != 0:
+    # b/443830202: The quantized dimension is currently increased by 1 because
+    # AEQ expects 1 and XNNPack expects 0.
+    quantization_params = dataclasses.replace(
+        quantization_params,
+        quantized_dimension=quantization_params.quantized_dimension + 1,
+    )
+    scale_shape = list(tensor_data.shape)
+    scale_shape[quantization_params.quantized_dimension] = (
+        scale_shape[quantization_params.quantized_dimension]
+        // quantization_params.block_size
+    )
+    quantization_params = dataclasses.replace(
+        quantization_params,
+        scale=quantization_params.scale.reshape(scale_shape),
+    )
+    quantization_params = _broadcast_scale_zp_for_blockwise(
+        tensor_data, quantization_params
+    )
+
   # quant params in flatbuffer is flattened, expand the rank to be the same
   # as the tensor rank to avoid ambiguous broadcasting.
   quantization_params = fix_quantization_params_rank(

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py

@@ -300,6 +300,30 @@ class TensorUtilsTest(parameterized.TestCase):
           ),
       )
 
+  def test_uniform_dequantize_blockwise(self):
+    quantized_tensor = np.array([[-8, -5, -4, 7], [-4, 7, -8, -5]])
+    expected_output_tensor = np.array([
+        [-10.1333336, -6.3333335, -5.0666668, 8.8666669],
+        [-5.0666668, 8.8666669, -10.1333336, -6.3333335],
+    ])
+    quant_params = qtyping.UniformQuantParams(
+        # b/443830202:
+        quantized_dimension=0,
+        num_bits=4,
+        scale=np.array([[[1.2666667, 1.2666667], [1.2666667, 1.2666667]]]),
+        zero_point=np.array([[0]]),
+        symmetric=True,
+        block_size=2,
+    )
+
+    dequantized_tensor = uniform_quantize_tensor.uniform_dequantize(
+        np.array(quantized_tensor), quant_params
+    )
+
+    self.assertSequenceAlmostEqual(
+        expected_output_tensor.flatten(), dequantized_tensor.flatten(), places=4
+    )
+
   @parameterized.parameters(
       (8, 8, True, True),
       (8, 4, False, True),

ai_edge_quantizer/qtyping.py

@@ -219,6 +219,7 @@ class UniformQuantParams:
         scale=quant_params['scales'],
         zero_point=quant_params['zero_points'],
         symmetric=symmetric,
+        block_size=quant_params['block_size'],
     )
 
   def __eq__(self, other):

ai_edge_quantizer/transformations/quantize_tensor.py

@@ -131,9 +131,14 @@ def _perform_blockwise_quantization(
       transformation_input.buffers,
   )
   blockwise_details.scales = scale_tensor_id
-  blockwise_details.blockSize = transformation_input.quant_params.block_size
+  # Blockwise quantization does not support zero point yet, so this points to
+  # a -1 buffer index.
   # TODO: b/404909258 - Add optional zero point to blockwise quantization.
+  blockwise_details.zeroPoints = -1
+  blockwise_details.blockSize = transformation_input.quant_params.block_size
   flatbuffer_quantization.details = blockwise_details
+  # TODO: b/443830202 - Hardcoding to 0 for now.
+  flatbuffer_quantization.quantizedDimension = 0
   return flatbuffer_quantization
 
 

ai_edge_quantizer/transformations/quantize_tensor_test.py

@@ -170,7 +170,7 @@ class QuantizeTensorTest(parameterized.TestCase):
     # Check if the scale and zero point tensors are inserted correctly.
     self.assertEqual(quant_param.details.scales, 9)
     # So far we don't have zero point in blockwise quantization.
-    self.assertEqual(quant_param.details.zeroPoints, 0)
+    self.assertEqual(quant_param.details.zeroPoints, -1)
 
   def test_int4_constant_packed_correctly(self):
     subgraph = self._model.subgraphs[0]

ai_edge_quantizer/utils/validation_utils.py

@@ -102,7 +102,8 @@ def median_diff_ratio(
 
 
 def cosine_similarity(
-    data1: np._typing.ArrayLike, data2: np._typing.ArrayLike
+    data1: np._typing.ArrayLike,
+    data2: np._typing.ArrayLike,
 ) -> float:
   """Calculates the cosine similarity between data1 & data2.
 
@@ -123,9 +124,14 @@ def cosine_similarity(
   # special handling for tensor of size 0
   if data1.size == 0:
     return float(0)
-  return float(
-      np.dot(data1, data2) / (np.linalg.norm(data1) * np.linalg.norm(data2))
-  )
+  norm_data1 = np.linalg.norm(data1)
+  norm_data2 = np.linalg.norm(data2)
+  # special handling for tensor of length 0
+  if norm_data1 == 0 and norm_data2 == 0:
+    return 1.0
+  if norm_data1 == 0 or norm_data2 == 0:
+    return 0.0
+  return np.dot(data1, data2) / (norm_data1 * norm_data2)
 
 
 def _preprocess_same_size_arrays(

{ai_edge_quantizer_nightly-0.4.0.dev20251006.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20251008.dist-info}/METADATA

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

@@ -12,7 +12,7 @@ ai_edge_quantizer/model_validator.py,sha256=Hj0_5o-Oa3dSlJ3ryVjRhvsyelHNyek1GrtG
 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=7aEMPA4qr4CGD3NXtZgG2fDoQX5NzK9jwSv1yWNqQV4,17149
+ai_edge_quantizer/qtyping.py,sha256=rp2jdmCuSsP6Ay8rD7NxDCpbFkRNkbYP29Uwe0xBfnA,17196
 ai_edge_quantizer/quantizer.py,sha256=ckAEOnnBxuCKZuvlzdChevCKPuE-IeDPHCNtFTWr250,17857
 ai_edge_quantizer/quantizer_test.py,sha256=bh4IowxRF249p_XKIKQ0f17PmeDddfcOUzvQ2ht1L0E,26530
 ai_edge_quantizer/recipe.py,sha256=MEkfQ2Sg3KAE9LAORHWcbjYNPg06EUbwc1d-VspQA2U,6461
@@ -40,8 +40,8 @@ ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py,sha256=1
 ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py,sha256=nscKDvNb14ErZdAfG0aXRWyRs6bTvhMqMjKx2vxvUK0,8725
 ai_edge_quantizer/algorithms/uniform_quantize/octav.py,sha256=Umxh4kJyeHddZf-Wd4aXE5MTI1XWFa5KRuM17uYU714,6922
 ai_edge_quantizer/algorithms/uniform_quantize/octav_test.py,sha256=sha1d99Xk87bI87tgz0g5LeDC-EeE4WMfM5rRC98-m4,9140
-ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=so5pMeoclPdXV_5dDiqWaA_cZ0Ud-OWnXxRbzNh9x1E,18576
-ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=eBTi2I12e65_lxVZjGaN2TAiMzvsyyXAhWmEpKEmkLA,15126
+ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=0bLDAjCm5wxasGXKT3XiS4quk-zXlWK6JKb-14FQAd4,19570
+ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=0xOdoIWuEo9JlXvZdX_Gbq5lfwCwEcG7RwOxhXAFOOY,15939
 ai_edge_quantizer/algorithms/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 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
@@ -58,8 +58,8 @@ ai_edge_quantizer/transformations/insert_hadamard_rotation.py,sha256=5D5WwrJCE6h
 ai_edge_quantizer/transformations/insert_hadamard_rotation_test.py,sha256=iV1p3nZfHUATV2YRoBOYurnu3pLy8n3aFppLWGQOPdA,7268
 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=unqInO0we6_cgwPjtHB3tLWIHPajfNuJSLGW-IFnI9E,7029
-ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=mHLO3_MRt36A8-ZN8ADn5tBBJlqjTWa7ZUN8Mmu5Rcw,9116
+ai_edge_quantizer/transformations/quantize_tensor.py,sha256=VGTVpZWla9R-LPfhTzH1NVAp2soOqDF_duIm8ez_z3Y,7264
+ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=CD7OboBcIQxQY8OaRd5ISC1JcwQW726P_vneY4LKVpA,9117
 ai_edge_quantizer/transformations/transformation_utils.py,sha256=IKrtXJNH0msiTcI7KXkCYn2EkzmbZKWMMX_r5PMEx2U,8857
 ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=MWgq29t7rvxRQIfi4ny9IoODFCTcbpjnIwoCL40zDKk,8698
 ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
@@ -72,10 +72,10 @@ ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=LN-WonrcJLP9bB4lULd5VIg_8
 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=yJH9Cvepr_XWn-3Hsh91j7HuC5iLQHAyskyQ48bGNoc,4797
+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.dev20251006.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20251006.dist-info/METADATA,sha256=b2Lbl1imbudi41DzfX1XxHAeWt0timziUlFLpDakO10,1508
-ai_edge_quantizer_nightly-0.4.0.dev20251006.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20251006.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20251006.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20251008.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20251008.dist-info/METADATA,sha256=OhwG-gg_4N_0x8rm_4M_-aCX0mBjdPxj__btANE2Sok,1508
+ai_edge_quantizer_nightly-0.4.0.dev20251008.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20251008.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20251008.dist-info/RECORD,,