ai-edge-quantizer-nightly 0.1.0.dev20250519__py3-none-any.whl → 0.1.0.dev20250521__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py

@@ -58,19 +58,19 @@ def _rotate_with_diagonal_hadamard(
 
   Args:
     tensor_content: The float array to quantize.
-    axis: The axis of the tensor to quantize.
+    axis: The axis of the tensor to rotate.
 
   Returns:
     A tuple containing the quantized array and the recovered array.
 
   Raises:
-    ValueError: If the axis is not 1. To support other axes, please add
-      support to the matrix multiplication.
+    ValueError: If the axis is not the last axis of tensor_content. To support
+      other axes, please add support to the matrix multiplication.
   """
-  if axis != 1:
+  if axis != tensor_content.ndim - 1:
     raise ValueError(
-        "Hadamard rotation is only supported for 2D tensors with quantized"
-        " dimension 0."
+        "Hadamard rotation is only supported for tensors with quantized"
+        " dimension 0 (rotate last dimension)."
     )
 
   # Use the largest power of 2 that is a factor of the dimension and then
@@ -78,15 +78,15 @@ def _rotate_with_diagonal_hadamard(
   # of 2 to calculate this factor.
   hadamard_size = np.gcd(tensor_content.shape[axis], 2 ** 30)
   diagonal_size = tensor_content.shape[axis] // hadamard_size
-  output_size = tensor_content.shape[1 - axis]
+  # Output size is the product of all dimensions except the one being rotated.
+  output_size = np.prod(np.delete(tensor_content.shape, axis))
   random_vector = np.ones(hadamard_size, dtype=np.int8)
 
   # Use a canonical Hadamard matrix.
   hadamard = _make_hadamard_matrix(hadamard_size)
   reshaped_tensor = tensor_content.reshape(
-      diagonal_size, output_size, hadamard_size
-  )
-  w_rotated = np.einsum("jk,ilk->ilj", hadamard, reshaped_tensor)
+      diagonal_size * output_size, hadamard_size)
+  w_rotated = np.matmul(hadamard, reshaped_tensor.mT).mT
   return w_rotated.reshape(tensor_content.shape), hadamard_size, random_vector
 
 
@@ -122,8 +122,10 @@ def get_tensor_quant_params(
         "Hadamard rotation is not supported for static quantization."
     )
 
-  if tensor_content.ndim != 2:
-    raise ValueError("Hadamard rotation is only supported for 2D tensors.")
+  if tensor_content.ndim < 2:
+    raise ValueError(
+        "Hadamard rotation is only supported for tensors with rank >= 2."
+    )
 
   if tensor_quant_config.granularity != qtyping.QuantGranularity.CHANNELWISE:
     raise ValueError(
@@ -140,9 +142,9 @@ def get_tensor_quant_params(
         " supported."
     )
 
-  # Reduction axis is the non-quantized dimension. Since we only support 2D
-  # tensors and quantized_dim of 0, the reduction axis is 1.
-  reduce_axis = 1
+  # Reduction axis is the last non-quantized dimension. Since we only support
+  # quantized_dim of 0, the reduction axis is the last axis.
+  reduce_axis = tensor_content.ndim - 1
 
   # Rotate the tensor with a Hadamard matrix.
   w_rotated, hadamard_size, random_vector = _rotate_with_diagonal_hadamard(

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py

@@ -168,6 +168,36 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
         np.array(qparams.quantized_data), expected
     )
 
+  def test_get_tensor_quant_params_golden_3(self):
+    # test_data:
+    #   [[[1 2 1 2 1 2]
+    #     [3 4 3 4 3 4]
+    #     [1 2 1 2 1 2]]
+    #    [[3 4 3 4 3 4]
+    #     [1 2 1 2 1 2]
+    #     [3 4 3 4 3 4]]]
+    test_data = np.tile([[1, 2], [3, 4]], [3, 3])
+    test_data = np.reshape(test_data, (2, 3, 6))
+    # expected:
+    #   [[[ 54 -18  54 -18  54 -18]
+    #     [127 -18 127 -18 127 -18]
+    #     [ 54 -18  54 -18  54 -18]]
+    #    [[127 -18 127 -18 127 -18]
+    #     [ 54 -18  54 -18  54 -18]
+    #     [127 -18 127 -18 127 -18]]]
+    expected = np.tile([[54, -18], [127, -18]], [3, 3])
+    expected = np.reshape(expected, (2, 3, 6))
+    qparams = hadamard_rotation.get_tensor_quant_params(
+        self._op_info,
+        self._op_info.op_quant_config.weight_tensor_config,
+        test_data,
+        self._tensor_name_to_qsv,
+    )
+    self.assertIsNotNone(qparams.quantized_data)
+    np.testing.assert_array_equal(
+        np.array(qparams.quantized_data), expected
+    )
+
   def test_raise_missing_tensor_content(self):
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: "weight tensor" in str(err)
@@ -190,9 +220,9 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
           self._graph_info.buffers[self._fc_buffer_id],
       )
 
-  def test_raise_non_2d_constant(self):
+  def test_raise_1d_constant(self):
     with self.assertRaisesWithPredicateMatch(
-        ValueError, lambda err: "2D tensors" in str(err)
+        ValueError, lambda err: "rank >= 2" in str(err)
     ):
       hadamard_rotation.get_tensor_quant_params(
           self._op_info,

{ai_edge_quantizer_nightly-0.1.0.dev20250519.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250521.dist-info}/METADATA

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

@@ -32,8 +32,8 @@ ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=NpZ-JvZt
 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=pN4hwggrdI4eBdqvsdwnFagFxpd4D8LkWK0o4HG_xxk,12536
-ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py,sha256=MajG6DqpP4HvVzcZwgiKojWL3RBxCpkU3u2mKyeB0hA,9191
+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/naive_min_max_quantize.py,sha256=8_tNLTbOWTKId4DfHBjkOR9RvELUyIpxlGxKu7tv5Ko,7556
 ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py,sha256=zoF_EHjYqsKkuev8wfuutIITEmp_maa70IpJI_Df3ck,7431
 ai_edge_quantizer/algorithms/uniform_quantize/octav.py,sha256=Umxh4kJyeHddZf-Wd4aXE5MTI1XWFa5KRuM17uYU714,6922
@@ -70,8 +70,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=WoewyiZpaua80oP0tpgyrw5W
 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.1.0.dev20250519.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250519.dist-info/METADATA,sha256=Miou6XsSuDYRUvHBvwHgxOz3L3X3CWYcmcftB0F5cHw,1528
-ai_edge_quantizer_nightly-0.1.0.dev20250519.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250519.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250519.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250521.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250521.dist-info/METADATA,sha256=RBpATRDheE6-5ohVAasX_rrpKfQiRhFWu2aEsc4bsJg,1528
+ai_edge_quantizer_nightly-0.1.0.dev20250521.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250521.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250521.dist-info/RECORD,,