ai-edge-quantizer-nightly 0.4.0.dev20250907__py3-none-any.whl → 0.4.0.dev20250909__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py

@@ -330,7 +330,8 @@ def symmetric_quantize_bias_tensor(
 
   # symmetric
   bias_zp = np.zeros_like(effective_output_scale, dtype=np.int32)
-  bias_number_bits = 64 if input_tensor_quant_params.num_bits == 16 else 32
+  # Fixed to 32 bits since most of the accelerators use int32 accumulator.
+  bias_number_bits = 32
   symmetric = True
   quantized_dimension = None if len(effective_output_scale) == 1 else 0
   bias_quant_params = qtyping.UniformQuantParams(
@@ -342,6 +343,13 @@ def symmetric_quantize_bias_tensor(
   )
 
   quantized_vars = uniform_quantize(bias_content, bias_quant_params)
+  # Save the int32 quantized bias as int64 if the input tensor is quantized to
+  # 16 bits. This is to assume the matmul is using int64 accumulator (safe from
+  # overflow). For accelerators with int32 accumulator, it is safe to cast int64
+  # back to int32.
+  if input_tensor_quant_params.num_bits == 16:
+    quantized_vars = quantized_vars.astype(np.int64)
+    bias_number_bits = 64
 
   # UniformQuantParams is frozen dataclass, need to recreate.
   return qtyping.UniformQuantParams(

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py

@@ -15,8 +15,11 @@
 
 """Tests for tensor_utils."""
 
+import dataclasses
+
 from absl.testing import parameterized
 import numpy as np
+
 from tensorflow.python.platform import googletest
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.algorithms.uniform_quantize import uniform_quantize_tensor
@@ -276,7 +279,10 @@ class TensorUtilsTest(parameterized.TestCase):
       )
 
   @parameterized.parameters(
-      (8, 8, True, True), (8, 4, False, True), (16, 8, True, False)
+      (8, 8, True, True),
+      (8, 4, False, True),
+      (16, 8, True, False),
+      (16, 8, True, True),
   )
   def test_quantize_bias_tensor(
       self,
@@ -334,6 +340,26 @@ class TensorUtilsTest(parameterized.TestCase):
     self.assertSequenceAlmostEqual(
         list(dequantized_bias.flatten()), list(bias_tensor_data), places=5
     )
+
+    if activation_num_bits == 16:
+      # Check if it is safe to cast int64 bias to int32. We save the int32
+      # quantized bias as int64 if the input tensor is quantized to 16 bits.
+      # This is to assume the matmul is using int64 accumulator (safe from
+      # overflow). For accelerators with int32 accumulator, it is safe to cast
+      # int64 back to int32.
+      quantized_bias = bias_quant_config.quantized_data
+      self.assertIsNotNone(quantized_bias)
+      self.assertEqual(quantized_bias.dtype, np.int64)
+      self.assertSequenceEqual(
+          list(quantized_bias.flatten()),
+          list(quantized_bias.astype(np.int32).flatten()),  # pytype: disable=attribute-error
+      )
+
+      bias_quant_config = dataclasses.replace(
+          bias_quant_config,
+          num_bits=32,
+      )
+
     expected_quantized_data = uniform_quantize_tensor.uniform_quantize(
         bias_tensor_data, bias_quant_config
     )

{ai_edge_quantizer_nightly-0.4.0.dev20250907.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250909.dist-info}/METADATA

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

@@ -38,8 +38,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=uCREMXi0U2ckhXXfgGVzwSgjFZc0IbtnFU-OjlG9IO8,17146
-ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=7kHluzpteMv36hFD6LD_qnwwMoE1GKUP4bGmGMFbOdA,12755
+ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=_WfxguP9dVsWFjUc2l8cvvx6kK_PEiUS5HU1vfC7d5c,17554
+ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=OH1etR0btp2ep7PIP4uc7LCOydrouVBeeBJzkIkMF-k,13625
 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
@@ -70,8 +70,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EoVjI_hplX_Rml3hfRsGmQOi
 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.dev20250907.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20250907.dist-info/METADATA,sha256=ubptLJreX7c1P74IlyA5O73oPqFgCwdpLqNqH9-xick,1508
-ai_edge_quantizer_nightly-0.4.0.dev20250907.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20250907.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20250907.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20250909.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20250909.dist-info/METADATA,sha256=-_Cqm5vEJXyhzkezt-hEOtZOZJ8kOtv8-0oUkMad7aQ,1508
+ai_edge_quantizer_nightly-0.4.0.dev20250909.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20250909.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20250909.dist-info/RECORD,,