ai-edge-quantizer-nightly 0.1.0.dev20250319__py3-none-any.whl → 0.1.0.dev20250321__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -790,8 +790,6 @@ def init_tensor_min_max(
     quantized_dim = None
     if weight_tensor_config is not None and (
         weight_tensor_config.granularity == qtyping.QuantGranularity.CHANNELWISE
-        or weight_tensor_config.granularity
-        == qtyping.QuantGranularity.BLOCKWISE
     ):
       quantized_dim = common_utils.get_weight_quantized_dim(
           op_info, tensor_data
@@ -801,6 +799,11 @@ def init_tensor_min_max(
         and weight_tensor_config.granularity
         == qtyping.QuantGranularity.BLOCKWISE
     ):
+      quantized_dim = (
+          tfl_flatbuffer_utils.TFL_OP_TO_BLOCKWISE_WEIGHT_QUANTIZED_DIM[
+              op_info.op_name
+          ]
+      )
       reshaped_data, reduce_dims = _reshape_data_for_blockwise(
           tensor_data,
           quantized_dim,

ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py

@@ -80,13 +80,16 @@ def get_tensor_quant_params(
       tensor_quant_config.symmetric,
   )
   quantized_dim = None
-  if (
-      tensor_quant_config.granularity == qtyping.QuantGranularity.CHANNELWISE
-      or tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE
-  ):
+  if tensor_quant_config.granularity == qtyping.QuantGranularity.CHANNELWISE:
     quantized_dim = common_utils.get_weight_quantized_dim(
         op_info, tensor_content
     )
+  elif tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
+    quantized_dim = (
+        tfl_flatbuffer_utils.TFL_OP_TO_BLOCKWISE_WEIGHT_QUANTIZED_DIM[
+            op_info.op_name
+        ]
+    )
   quant_params = qtyping.UniformQuantParams(
       scale=scale,
       zero_point=zp,

ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py

@@ -187,8 +187,8 @@ class NaiveMinMaxQuantizeTest(parameterized.TestCase):
     zp = quant_params.zero_point
     expected_zp, expected_scale = (
         uniform_quantize_tensor.tensor_zp_scale_from_min_max(
-            min_value=np.array([[-7, 4], [-4, -4]]),
-            max_value=np.array([[4, 7], [7, 7]]),
+            min_value=np.array([[-7], [-4], [-4], [7]]),
+            max_value=np.array([[7], [4], [4], [7]]),
             num_bits=4,
             symmetric=True,
         )
@@ -200,7 +200,7 @@ class NaiveMinMaxQuantizeTest(parameterized.TestCase):
         cast(np.ndarray, quant_params.quantized_data).shape, test_data.shape
     )
     self.assertEqual(quant_params.block_size, 2)
-    self.assertEqual(quant_params.quantized_dimension, 0)
+    self.assertEqual(quant_params.quantized_dimension, 1)
 
 
 if __name__ == "__main__":

ai_edge_quantizer/transformations/quantize_tensor.py

@@ -143,24 +143,15 @@ def _perform_blockwise_quantization(
   tensor = transformation_input.subgraph.tensors[transformation_input.tensor_id]
   blockwise_details = schema_py_generated.BlockwiseQuantizationT()
   scale_tensor_id = transformation_utils.add_new_constant_tensor(
-      tensor.name + b"_scale",
-      transformation_input.quant_params.scale,
+      tensor.name + b"_scales",
+      transformation_input.quant_params.scale.astype(np.float16),
       schema_py_generated.TensorType.FLOAT16,
       transformation_input.subgraph,
       transformation_input.buffers,
   )
   blockwise_details.scales = scale_tensor_id
   blockwise_details.blockSize = transformation_input.quant_params.block_size
-  # blockwise quantization allows optional zero point.
-  if transformation_input.quant_params.zero_point is not None:
-    zero_point_tensor_id = transformation_utils.add_new_constant_tensor(
-        tensor.name + b"_zero_point",
-        transformation_input.quant_params.zero_point,
-        schema_py_generated.TensorType.INT32,
-        transformation_input.subgraph,
-        transformation_input.buffers,
-    )
-    blockwise_details.zeroPoints = zero_point_tensor_id
+  # TODO: b/404909258 - Add optional zero point to blockwise quantization.
   flatbuffer_quantization.details = blockwise_details
   return flatbuffer_quantization
 

ai_edge_quantizer/transformations/quantize_tensor_test.py

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

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -72,6 +72,11 @@ TFL_OP_TO_WEIGHT_QUANTIZED_DIM = immutabledict.immutabledict({
     _TFLOpName.CONV_2D_TRANSPOSE: 0,
 })
 
+TFL_OP_TO_BLOCKWISE_WEIGHT_QUANTIZED_DIM = immutabledict.immutabledict({
+    _TFLOpName.FULLY_CONNECTED: 1,
+    _TFLOpName.EMBEDDING_LOOKUP: 1,
+})
+
 NUM_TFL_DATATYPES = 18
 TENSOR_CODE_TO_TYPE = {}
 for dtype_code in range(NUM_TFL_DATATYPES):

{ai_edge_quantizer_nightly-0.1.0.dev20250319.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250321.dist-info}/METADATA

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

@@ -28,12 +28,12 @@ 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=s64eDDH9bmRWy6Bl1peHnhGewLnFJjvnhYOdjo1zYOA,22625
 ai_edge_quantizer/algorithms/uniform_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=LnItMEsR47qe8T5pg9UI5NGfhi4cOxt0vAU35IkWnaY,27163
+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/naive_min_max_quantize.py,sha256=aWHU4rneBv7ErufEWKQGAWTK-pgfn-rG9mAkC0d9V6Q,7871
-ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py,sha256=Hok09dloSyBfD0oDM5VABdSZjM9JWSQhm_hDHNbFujA,7640
+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
 ai_edge_quantizer/algorithms/uniform_quantize/octav_test.py,sha256=IcTOaJ1pxtqsitqxOEP9LROVEP_19VFutHalqNied4I,6940
 ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=WmZzKQlzfu9gFr9SbUDoPY3rFqTl363om8-0rTLwotw,11629
@@ -50,22 +50,22 @@ 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=9YaaWR6osxZoyUM8DUPJr_AjWO0QuNhFc65OFnSGzY4,7866
-ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=XZOollD1jnpCb78gMZx7yocF7RDBSf9HIf-XdG-y_io,9052
+ai_edge_quantizer/transformations/quantize_tensor.py,sha256=vzKtrXILqVsr1NGlribhdtKEIsXA93o37embLRe9TwQ,7493
+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
 ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/utils/calibration_utils.py,sha256=1Fj9MIO6aLZIRgyd4axvZN4S_O64nB_-Miu1WP664js,2536
 ai_edge_quantizer/utils/calibration_utils_test.py,sha256=Z-AcdTieesWFKyKBb08ZXm4Mgu6cvJ4bg2-MJ7hLD10,2856
 ai_edge_quantizer/utils/test_utils.py,sha256=HwZCIpO9fJRAhuN6t6voXKOYQtcioFtt_tpkAlDsAYk,6205
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=_A-h_MqwElzjgkLDmXTZ1iAIWtTRcLjSFGfjNT8fuHU,10480
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=51GRkwj7PK0XvAqohdv6mAepOWRk1AnW2y-9ne6LzWo,10628
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=AbyDxoM62k4ojD8gPdkWo--xe5hlX3t0kobQSA80kuk,7740
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=x2xA2CFPpe_2trcV8v5xGaBETvVCfwAcJuq6yieGJ0Y,12687
 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.dev20250319.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250319.dist-info/METADATA,sha256=WTz-_FHdUgNLhVPcpu4VW9rw2drBw92tUqa35_OsDWg,1527
-ai_edge_quantizer_nightly-0.1.0.dev20250319.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250319.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250319.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250321.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250321.dist-info/METADATA,sha256=2IxDz8eWG3N0LQ7Qm9Y3_WH7hckqK_IwwzVc60-rkW4,1527
+ai_edge_quantizer_nightly-0.1.0.dev20250321.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250321.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250321.dist-info/RECORD,,