ai-edge-quantizer-nightly 0.4.0.dev20250919__py3-none-any.whl → 0.4.0.dev20250921__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py

@@ -127,23 +127,9 @@ def get_tensor_quant_params(
         "Hadamard rotation is only supported for tensors with rank >= 2."
     )
 
-  if tensor_quant_config.granularity != qtyping.QuantGranularity.CHANNELWISE:
-    raise ValueError(
-        "Hadamard rotation is not supported for"
-        f" {tensor_quant_config.granularity} granularity."
-    )
-
-  quantized_dim = common_utils.get_weight_quantized_dim(
-      op_info, tensor_content, tensor_quant_config.granularity
-  )
-  if quantized_dim != 0:
-    raise ValueError(
-        f"Unsupported quantized dimension: {quantized_dim}. Only 0 is"
-        " supported."
-    )
-
   # Reduction axis is the last non-quantized dimension. Since we only support
-  # quantized_dim of 0, the reduction axis is the last axis.
+  # quantized_dim of 0 (or 1 for blockwise), the reduction axis is the last
+  # axis.
   reduce_axis = tensor_content.ndim - 1
 
   # Rotate the tensor with a Hadamard matrix.

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py

@@ -46,13 +46,13 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
     )
     self._tensor_name_to_qsv = None
     self._subgraph = self._test_model.subgraphs[0]
-    fc_subgraph_op_index = 3
-    self._fc_op = self._subgraph.operators[fc_subgraph_op_index]
+    self._fc_subgraph_op_index = 3
+    self._fc_op = self._subgraph.operators[self._fc_subgraph_op_index]
     self._fc_buffer_id = self._subgraph.tensors[self._fc_op.inputs[1]].buffer
     self._op_info = qtyping.OpInfo(
         op=self._fc_op,
         op_name=_TFLOpName.FULLY_CONNECTED,
-        subgraph_op_index=fc_subgraph_op_index,
+        subgraph_op_index=self._fc_subgraph_op_index,
         op_quant_config=qtyping.OpQuantizationConfig(
             weight_tensor_config=_TensorQuantConfig(
                 num_bits=8,
@@ -98,6 +98,87 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
           [qtyping.QuantTransformation.NO_QUANTIZE],
       )
 
+  def test_fully_connected_tensorwise_supported(self):
+    self._op_info = qtyping.OpInfo(
+        op=self._fc_op,
+        op_name=_TFLOpName.FULLY_CONNECTED,
+        subgraph_op_index=self._fc_subgraph_op_index,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=8,
+                symmetric=True,
+                granularity=qtyping.QuantGranularity.TENSORWISE,
+            ),
+        ),
+    )
+    params = hadamard_rotation.materialize_fully_connected(
+        self._op_info, self._graph_info, self._tensor_name_to_qsv
+    )
+    self.assertLen(params, 4)
+    fc_input = params[0]
+    self.assertIsNotNone(fc_input)
+    self.assertIsNotNone(fc_input.consumers)
+    self.assertIsNotNone(fc_input.consumers[0].parameters)
+    self.assertIsInstance(
+        fc_input.consumers[0].parameters, qtyping.UniformQuantParams
+    )
+    if isinstance(
+        fc_input.consumers[0].parameters, qtyping.UniformQuantParams
+    ):
+      self.assertIsNone(fc_input.consumers[0].parameters.quantized_dimension)
+    weight = params[1]
+    self.assertIsNotNone(weight)
+    self.assertIsNotNone(weight.consumers)
+    self.assertIsNotNone(weight.consumers[0].parameters)
+    self.assertIsInstance(
+        weight.consumers[0].parameters, qtyping.UniformQuantParams
+    )
+    if isinstance(
+        weight.consumers[0].parameters, qtyping.UniformQuantParams
+    ):
+      self.assertIsNone(weight.consumers[0].parameters.quantized_dimension)
+
+  def test_fully_connected_blockwise_supported(self):
+    self._op_info = qtyping.OpInfo(
+        op=self._fc_op,
+        op_name=_TFLOpName.FULLY_CONNECTED,
+        subgraph_op_index=self._fc_subgraph_op_index,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=8,
+                symmetric=True,
+                granularity=qtyping.QuantGranularity.BLOCKWISE,
+                block_size=32,
+            ),
+        ),
+    )
+    params = hadamard_rotation.materialize_fully_connected(
+        self._op_info, self._graph_info, self._tensor_name_to_qsv
+    )
+    self.assertLen(params, 4)
+    fc_input = params[0]
+    self.assertIsNotNone(fc_input)
+    self.assertIsNotNone(fc_input.consumers)
+    self.assertIsNotNone(fc_input.consumers[0].parameters)
+    self.assertIsInstance(
+        fc_input.consumers[0].parameters, qtyping.UniformQuantParams
+    )
+    if isinstance(
+        fc_input.consumers[0].parameters, qtyping.UniformQuantParams
+    ):
+      self.assertEqual(fc_input.consumers[0].parameters.quantized_dimension, 1)
+    weight = params[1]
+    self.assertIsNotNone(weight)
+    self.assertIsNotNone(weight.consumers)
+    self.assertIsNotNone(weight.consumers[0].parameters)
+    self.assertIsInstance(
+        weight.consumers[0].parameters, qtyping.UniformQuantParams
+    )
+    if isinstance(
+        weight.consumers[0].parameters, qtyping.UniformQuantParams
+    ):
+      self.assertEqual(weight.consumers[0].parameters.quantized_dimension, 1)
+
   def test_get_tensor_quant_params_basic(self):
     input_tensor = self._subgraph.tensors[self._fc_op.inputs[1]]
     buffer = self._graph_info.buffers[self._fc_buffer_id]

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py

@@ -143,6 +143,11 @@ def _get_tensor_shape_for_blockwise(
   new_shape = []
   for index, val in enumerate(tensor_shape):
     if index == quantized_dim:
+      if val % block_size != 0:
+        raise ValueError(
+            f"Quantized dimension {val} in tensor shape {tensor_shape} is not"
+            f" divisible by block size {block_size}."
+        )
       new_shape.append(int(val / block_size))
       new_shape.append(block_size)
     else:

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py

@@ -203,6 +203,28 @@ class TensorUtilsTest(parameterized.TestCase):
           ),
       )
 
+  def test_uniform_quantize_quant_dim_not_divisible_by_block_size_raise(self):
+    tensor = np.random.rand(34, 2)
+    error_message = (
+        "Quantized dimension 34 in tensor shape (34, 2) is not divisible by"
+        " block size 32."
+    )
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      uniform_quantize_tensor.uniform_quantize(
+          np.array(tensor),
+          qtyping.UniformQuantParams(
+              quantized_dimension=0,
+              block_size=32,
+              num_bits=4,
+              scale=np.array([1.2666667]),
+              zero_point=np.array([-6]),
+              symmetric=True,
+          ),
+          is_blockwise=True,
+      )
+
   @parameterized.parameters(
       (
           8,

{ai_edge_quantizer_nightly-0.4.0.dev20250919.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250921.dist-info}/METADATA

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

@@ -32,14 +32,14 @@ ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=E17cSR-M
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py,sha256=GGf_n3wIeg3GB_eGsmyNJ0fTcxgpeMMbugTMRONK6TQ,3553
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=BDdn_uBZakfHyzdMJPKadsOqxqyC-s6W2ZzFH99L4fE,8652
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py,sha256=sT5eX5TLZEHTtPfnSkCPDlS0sQxlTFWbCsbvOuj--yY,8889
-ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py,sha256=U3h5scCHSOdqHA-pb1C3pNgwumT4ydGbtkCSM0ORhrs,12740
-ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py,sha256=5VUxlaKP1jz4HV-LcKxXMMtmb6eWamq0A6qWJd63cR4,10179
+ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py,sha256=otKRiZn_C0QH0891pxLsIPIBT1mLDwbKYYP7bI-MXAA,12279
+ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py,sha256=_SpP12aDLujv_7tWf_mCt89WknNXTSGE-JpZWO1bYSE,13238
 ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py,sha256=1sB2j1vlvvWDKyjcGvA_JLCpN2KbCmMslGCBUc4--V4,8461
 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=vsvBGEGFEEUP4kXRUh9hMpVXjsMBpfs6UDk8m4BNGTs,18375
-ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py,sha256=Ympigz0BGcaO5x3OozxNxrRAGiF0to6V_HXAcxNNEpI,14399
+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/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.dev20250919.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20250919.dist-info/METADATA,sha256=xhmp6LEFOIxQNTiacffhdU7GGs_8YE2MfdxUGWXqzEo,1508
-ai_edge_quantizer_nightly-0.4.0.dev20250919.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20250919.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20250919.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20250921.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20250921.dist-info/METADATA,sha256=XYKz5uSw06qVDCUQf-DqGJdgaT-Y5jxREjbdp_g42LY,1508
+ai_edge_quantizer_nightly-0.4.0.dev20250921.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20250921.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20250921.dist-info/RECORD,,