ai-edge-quantizer-nightly 0.1.0.dev20250512__py3-none-any.whl → 0.1.0.dev20250514__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py

@@ -0,0 +1,265 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Test Hadamard rotation materialization."""
+
+import os
+
+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 hadamard_rotation
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+
+_TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile("../../tests/models")
+_TFLOpName = qtyping.TFLOperationName
+_TensorQuantConfig = qtyping.TensorQuantizationConfig
+
+
+class HadamardRotationFullyConnectedTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(888)
+    self._test_model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, "conv_fc_mnist.tflite"
+    )
+    self._test_model = tfl_flatbuffer_utils.read_model(self._test_model_path)
+    self._graph_info = qtyping.GraphInfo(
+        subgraph_tensors=self._test_model.subgraphs[0].tensors,
+        buffers=self._test_model.buffers,
+    )
+    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_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,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=8,
+                symmetric=True,
+                granularity=qtyping.QuantGranularity.CHANNELWISE,
+            ),
+        ),
+    )
+
+  def test_materialize_fully_connected_basic(self):
+    params = hadamard_rotation.materialize_fully_connected(
+        self._op_info, self._graph_info, self._tensor_name_to_qsv
+    )
+    fc_input = params[0]
+    weight = params[1]
+    bias = params[2]
+    output = params[3]
+
+    self.assertLen(params, 4)
+    self.assertIsNone(fc_input.producer)
+    self.assertIsNotNone(fc_input.consumers)
+    self.assertIsNone(weight.producer)
+    self.assertIsNotNone(weight.consumers)
+    self.assertIsNone(bias.producer)
+    self.assertIsNotNone(bias.consumers)
+    self.assertIsNotNone(output.producer)
+    self.assertIsNone(output.consumers)
+    self.assertEqual(
+        fc_input.consumers[0].transformations,
+        [qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION],
+    )
+    self.assertEqual(
+        weight.consumers[0].transformations,
+        [qtyping.QuantTransformation.QUANTIZE_TENSOR],
+    )
+    self.assertEqual(
+        bias.consumers[0].transformations,
+        [qtyping.QuantTransformation.NO_QUANTIZE],
+    )
+    if output.producer is not None:
+      self.assertEqual(
+          output.producer.transformations,
+          [qtyping.QuantTransformation.NO_QUANTIZE],
+      )
+
+  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]
+    np_buffer = np.frombuffer(buffer.data, dtype=np.float32).reshape(
+        input_tensor.shape
+    )
+    qparams = hadamard_rotation.get_tensor_quant_params(
+        self._op_info,
+        self._op_info.op_quant_config.weight_tensor_config,
+        np_buffer,
+        self._tensor_name_to_qsv,
+    )
+    self.assertEqual(qparams.num_bits, 8)
+    self.assertEqual(qparams.zero_point.all(), 0)
+    self.assertEqual(qparams.symmetric, True)
+    self.assertIsNotNone(qparams.quantized_data)
+    self.assertEqual(qparams.block_size, 0)
+    self.assertIsNotNone(qparams.hadamard)
+    if qparams.hadamard is not None:
+      self.assertEqual(qparams.hadamard.hadamard_size, 32)
+
+  def test_get_tensor_quant_params_golden_1(self):
+    test_data = np.ones((6, 6))
+    # expected:
+    #   [[127   0 127   0 127   0]
+    #    [127   0 127   0 127   0]
+    #    [127   0 127   0 127   0]
+    #    [127   0 127   0 127   0]
+    #    [127   0 127   0 127   0]
+    #    [127   0 127   0 127   0]]
+    expected = np.tile([127, 0], [6, 3])
+    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_get_tensor_quant_params_golden_2(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])
+    # expected:
+    #   [[127 -42 127 -42 127 -42]
+    #    [127 -18 127 -18 127 -18]
+    #    [127 -42 127 -42 127 -42]
+    #    [127 -18 127 -18 127 -18]
+    #    [127 -42 127 -42 127 -42]
+    #    [127 -18 127 -18 127 -18]]
+    expected = np.tile([[127, -42], [127, -18]], [3, 3])
+    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)
+    ):
+      hadamard_rotation.get_tensor_quant_params(
+          self._op_info,
+          self._op_info.op_quant_config.weight_tensor_config,
+          None,
+          self._tensor_name_to_qsv,
+      )
+
+  def test_raise_qsv_set(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: "static quantization" in str(err)
+    ):
+      hadamard_rotation.get_tensor_quant_params(
+          self._op_info,
+          self._op_info.op_quant_config.weight_tensor_config,
+          self._graph_info.buffers[self._fc_buffer_id],
+          self._graph_info.buffers[self._fc_buffer_id],
+      )
+
+  def test_raise_non_2d_constant(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: "2D tensors" in str(err)
+    ):
+      hadamard_rotation.get_tensor_quant_params(
+          self._op_info,
+          self._op_info.op_quant_config.weight_tensor_config,
+          np.array([1.0, 2.0, 3.0]),
+          self._tensor_name_to_qsv,
+      )
+
+
+class HadamardRotationEmbeddingLookupTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(888)
+    self._test_model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, "embedding_lookup.tflite"
+    )
+    self._test_model = tfl_flatbuffer_utils.read_model(self._test_model_path)
+    self._graph_info = qtyping.GraphInfo(
+        subgraph_tensors=self._test_model.subgraphs[0].tensors,
+        buffers=self._test_model.buffers,
+    )
+    self._tensor_name_to_qsv = None
+
+  def test_materialize_embedding_lookup_basic(self):
+    subgraph = self._test_model.subgraphs[0]
+    embedding_subgraph_op_index = 0
+    embedding_op = subgraph.operators[embedding_subgraph_op_index]
+    op_info = qtyping.OpInfo(
+        op=embedding_op,
+        op_name=_TFLOpName.EMBEDDING_LOOKUP,
+        subgraph_op_index=embedding_subgraph_op_index,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=8,
+                symmetric=True,
+                granularity=qtyping.QuantGranularity.CHANNELWISE,
+            ),
+        ),
+    )
+    params = hadamard_rotation.materialize_embedding_lookup(
+        op_info, self._graph_info, self._tensor_name_to_qsv
+    )
+    self.assertLen(params, 3)
+    lookup = params[0]
+    value = params[1]
+    output = params[2]
+    self.assertIsNone(lookup.producer)
+    self.assertIsNotNone(lookup.consumers)
+    self.assertIsNone(value.producer)
+    self.assertIsNotNone(value.consumers)
+    self.assertIsNotNone(output.producer)
+    self.assertIsNone(output.consumers)
+    self.assertEqual(
+        lookup.consumers[0].transformations,
+        [qtyping.QuantTransformation.NO_QUANTIZE],
+    )
+    self.assertEqual(
+        value.consumers[0].transformations,
+        [qtyping.QuantTransformation.QUANTIZE_TENSOR],
+    )
+    if output.producer is not None:
+      self.assertEqual(
+          output.producer.transformations,
+          [qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION],
+      )
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py

@@ -16,7 +16,6 @@
 """Performs naive min/max uniform quantization."""
 
 from typing import Any, Optional
-import ml_dtypes
 import numpy as np
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
@@ -75,35 +74,17 @@ def get_tensor_quant_params(
         " the ParamsGenerator."
     )
   clipping_values = None
-  # Blockwise quantization uses float16 scale, with 7 bit mantissa,
-  # so the maximum representable value is 65280.
-  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
-    clipping_values = np.broadcast_to(
-        np.array(65280), tensor_min_max["min"].shape
-    )
   zp, scale = uniform_quantize_tensor.tensor_zp_scale_from_min_max(
       tensor_min_max["min"],
       tensor_min_max["max"],
       tensor_quant_config.num_bits,
       tensor_quant_config.symmetric,
+      tensor_quant_config.granularity,
       clipping_values,
   )
-  # Round the scale values to 7 bit mantissa.
-  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
-    scale = (
-        scale.astype(ml_dtypes.bfloat16).astype(np.float16).astype(np.float32)
-    )
-  quantized_dim = None
-  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
-        ]
-    )
+  quantized_dim = common_utils.get_weight_quantized_dim(
+      op_info, tensor_content, tensor_quant_config.granularity
+  )
   quant_params = qtyping.UniformQuantParams(
       scale=scale,
       zero_point=zp,
@@ -115,15 +96,10 @@ def get_tensor_quant_params(
   if tensor_content is None:
     return quant_params
 
-  # The reshaping for blockwise quantization is unique hence we do this here
-  # to avoid unexpected broadcast behavior downstream.
-  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
-    quant_params = common_quantize.broadcast_scale_zp_for_blockwise(
-        tensor_content, quant_params
-    )
-
   quantized_vars = uniform_quantize_tensor.uniform_quantize(
-      tensor_content, quant_params
+      tensor_content,
+      quant_params,
+      tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE,
   )
   # Update with quantized values.
   return qtyping.UniformQuantParams(

ai_edge_quantizer/algorithms/uniform_quantize/octav.py

@@ -102,21 +102,13 @@ def get_tensor_quant_params(
         op_info, tensor_quant_config, tensor_content, tensor_qsv
     )
 
-  if (
-      tensor_quant_config.granularity != qtyping.QuantGranularity.CHANNELWISE
-      and tensor_quant_config.granularity != qtyping.QuantGranularity.TENSORWISE
-  ):
-    raise ValueError(
-        f"Unsupported granularity: {tensor_quant_config.granularity}."
-    )
-
   if not tensor_quant_config.symmetric:
     raise ValueError(
         f"Unsupported symmetry: {tensor_quant_config.symmetric}. OCTAV"
         " supports symmetric quantization only for now."
     )
 
-  if tensor_qsv is None:
+  if not tensor_qsv:
     # We need min/max to calculate quantization parameters, which
     # should be collected during the calibration process. However,
     # weight-only and DRQ do not require calibration, thus it is
@@ -136,25 +128,41 @@ def get_tensor_quant_params(
         " the ParamsGenerator."
     )
 
-  quantized_dim = None
-  if tensor_quant_config.granularity == qtyping.QuantGranularity.CHANNELWISE:
-    quantized_dim = common_utils.get_weight_quantized_dim(
-        op_info, tensor_content
+  quantized_dim = common_utils.get_weight_quantized_dim(
+      op_info, tensor_content, tensor_quant_config.granularity
+  )
+  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
+    reshaped_data, reduce_dims = (
+        uniform_quantize_tensor.reshape_data_for_blockwise(
+            tensor_content,
+            op_info.op_name,
+            tensor_quant_config.block_size,
+        )
+    )
+  else:
+    reshaped_data = tensor_content
+    reduce_dims = common_utils.get_reduce_dims(
+        quantized_dim, tensor_content.shape
     )
-
   clipping_constants = _guess_clipping_with_octav(
-      tensor_content,
+      reshaped_data,
       tensor_quant_config.num_bits,
-      common_utils.get_reduce_dims(quantized_dim, tensor_content.shape),
+      reduce_dims,
       max_iterations=10,
       exponent_divisor=3.0 if tensor_quant_config.symmetric else 12.0,
   )
+  # We created a new dimension in order to reduce properly for blockwise
+  # quantization, so we need to reshape the clipping constants back to the
+  # min/max shape for the next step.
+  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
+    clipping_constants = clipping_constants.reshape(tensor_min_max["min"].shape)
 
   zp, scale = uniform_quantize_tensor.tensor_zp_scale_from_min_max(
       tensor_min_max["min"],
       tensor_min_max["max"],
       tensor_quant_config.num_bits,
       tensor_quant_config.symmetric,
+      tensor_quant_config.granularity,
       clipping_constants,
   )
 
@@ -168,7 +176,9 @@ def get_tensor_quant_params(
   )
 
   quantized_vars = uniform_quantize_tensor.uniform_quantize(
-      tensor_content, quant_params
+      tensor_content,
+      quant_params,
+      tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE,
   )
 
   return dataclasses.replace(quant_params, quantized_data=quantized_vars)

ai_edge_quantizer/algorithms/uniform_quantize/octav_test.py

@@ -44,33 +44,17 @@ class OctavQuantizeTest(parameterized.TestCase):
     )
     self._tensor_name_to_qsv = {}
     subgraph0 = self._test_model.subgraphs[0]
-    subgraph_op_index = 3
-    fc_op = subgraph0.operators[subgraph_op_index]
+    self._subgraph_op_index = 3
+    self._fc_op = subgraph0.operators[self._subgraph_op_index]
     self._fc_op_info = qtyping.OpInfo(
-        op=fc_op,
+        op=self._fc_op,
         op_name=qtyping.TFLOperationName.FULLY_CONNECTED,
-        subgraph_op_index=subgraph_op_index,
+        subgraph_op_index=self._subgraph_op_index,
         op_quant_config=qtyping.OpQuantizationConfig(
             weight_tensor_config=None,
         ),
     )
 
-  def test_get_tensor_quant_params_unsupported_granularity_assert(self):
-    err_msg = "Unsupported granularity"
-    test_data = np.array([[-7, 7], [4, -4], [4, -4], [7, 7]])
-    with self.assertRaisesWithPredicateMatch(
-        ValueError, lambda err: err_msg in str(err)
-    ):
-      _ = octav.get_tensor_quant_params(
-          op_info=self._fc_op_info,
-          tensor_quant_config=qtyping.TensorQuantizationConfig(
-              num_bits=4,
-              symmetric=True,
-              granularity=qtyping.QuantGranularity.BLOCKWISE,
-          ),
-          tensor_content=test_data,
-      )
-
   def test_get_tensor_quant_params_unsupported_symmetry(self):
     err_msg = "Unsupported symmetry"
     test_data = np.array([[-7, 7], [4, -4], [4, -4], [7, 7]])
@@ -117,13 +101,22 @@ class OctavQuantizeTest(parameterized.TestCase):
         [25, -30, 50, -75, 1e5, -125],
         [50, -60, 70, -80, 90, -100],
     ])
-    quant_params = octav.get_tensor_quant_params(
-        op_info=self._fc_op_info,
-        tensor_quant_config=qtyping.TensorQuantizationConfig(
-            num_bits=4,
-            symmetric=True,
-            granularity=qtyping.QuantGranularity.TENSORWISE,
+    tensor_config = qtyping.TensorQuantizationConfig(
+        num_bits=4,
+        symmetric=True,
+        granularity=qtyping.QuantGranularity.TENSORWISE,
+    )
+    fc_op_info = qtyping.OpInfo(
+        op=self._fc_op,
+        op_name=qtyping.TFLOperationName.FULLY_CONNECTED,
+        subgraph_op_index=self._subgraph_op_index,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=tensor_config,
         ),
+    )
+    quant_params = octav.get_tensor_quant_params(
+        op_info=fc_op_info,
+        tensor_quant_config=tensor_config,
         tensor_content=test_data,
     )
     adjusted_test_data = quant_params.quantized_data * quant_params.scale
@@ -131,10 +124,10 @@ class OctavQuantizeTest(parameterized.TestCase):
     adjusted_max = np.max(np.abs(adjusted_test_data))
 
     # Check that some clipping occurred.
-    with self.subTest(name="SanityCheckClipping"):
+    with self.subTest(name="CheckClipping"):
       self.assertLess(adjusted_max, real_max)
 
-    with self.subTest(name="SanityCheckQuantParamsShapes"):
+    with self.subTest(name="CheckQuantParamsShapes"):
       self.assertEqual(quant_params.zero_point.shape, (1, 1))
       self.assertEqual(quant_params.scale.shape, (1, 1))
       self.assertIsNone(quant_params.quantized_dimension)
@@ -143,33 +136,47 @@ class OctavQuantizeTest(parameterized.TestCase):
           cast(np.ndarray, quant_params.quantized_data).shape, test_data.shape
       )
 
-    with self.subTest(name="SanityCheckQuantParamsValues"):
+    with self.subTest(name="CheckQuantParamsValues"):
       self.assertTrue(np.all(quant_params.zero_point == 0))
 
   def test_get_tensor_quant_params_sanity_channelwise(self):
+    # Test that the call generates quant params that are appropriately shaped,
+    # have some clipping, and correct config values without checking the
+    # actual values numerically.
     test_data = np.array([
         [-1e5, 25, -50, 75, -100, 125],
         [25, -30, 50, -75, 1e5, -125],
         [50, -60, 70, -80, 90, -100],
     ])
-    quant_params = octav.get_tensor_quant_params(
-        op_info=self._fc_op_info,
-        tensor_quant_config=qtyping.TensorQuantizationConfig(
-            num_bits=4,
-            symmetric=True,
-            granularity=qtyping.QuantGranularity.CHANNELWISE,
+    tensor_config = qtyping.TensorQuantizationConfig(
+        num_bits=4,
+        symmetric=True,
+        granularity=qtyping.QuantGranularity.CHANNELWISE,
+    )
+    fc_op_info = qtyping.OpInfo(
+        op=self._fc_op,
+        op_name=qtyping.TFLOperationName.FULLY_CONNECTED,
+        subgraph_op_index=self._subgraph_op_index,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=tensor_config,
         ),
+    )
+    quant_params = octav.get_tensor_quant_params(
+        op_info=fc_op_info,
+        tensor_quant_config=tensor_config,
         tensor_content=test_data,
     )
+    # Dequantize output to compare with the original test data.
     adjusted_test_data = quant_params.quantized_data * quant_params.scale
+
     for i, row in enumerate(test_data):
       real_max = np.max(np.abs(row))
       adjusted_max = np.max(np.abs(adjusted_test_data[i]))
       # Check that some clipping occurred.
-      with self.subTest(name="SanityCheckClipping"):
+      with self.subTest(name="CheckClipping"):
         self.assertLess(adjusted_max, real_max)
 
-    with self.subTest(name="SanityCheckQuantParamsShapes"):
+    with self.subTest(name="CheckQuantParamsShapes"):
       self.assertEqual(quant_params.zero_point.shape, (test_data.shape[0], 1))
       self.assertEqual(quant_params.scale.shape, (test_data.shape[0], 1))
       self.assertIsNotNone(quant_params.quantized_data)
@@ -177,10 +184,58 @@ class OctavQuantizeTest(parameterized.TestCase):
           cast(np.ndarray, quant_params.quantized_data).shape, test_data.shape
       )
 
-    with self.subTest(name="SanityCheckQuantParamsValues"):
+    with self.subTest(name="CheckQuantParamsValues"):
       self.assertTrue(np.all(quant_params.zero_point == 0))
       self.assertEqual(quant_params.quantized_dimension, 0)
 
+  def test_get_tensor_quant_params_sanity_blockwise(self):
+    # Test that the call generates quant params that are appropriately shaped,
+    # have some clipping, and correct config values without checking the
+    # actual values numerically.
+    test_data = np.random.randint(0, 1024, size=(32, 128))
+    tensor_config = qtyping.TensorQuantizationConfig(
+        num_bits=4,
+        symmetric=True,
+        granularity=qtyping.QuantGranularity.BLOCKWISE,
+        block_size=32,
+    )
+    fc_op_info = qtyping.OpInfo(
+        op=self._fc_op,
+        op_name=qtyping.TFLOperationName.FULLY_CONNECTED,
+        subgraph_op_index=self._subgraph_op_index,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=tensor_config,
+        ),
+    )
+    quant_params = octav.get_tensor_quant_params(
+        op_info=fc_op_info,
+        tensor_quant_config=tensor_config,
+        tensor_content=test_data,
+    )
+
+    with self.subTest(name="CheckQuantParamsShapes"):
+      # Check that quant params have appropriate shapes.
+      self.assertEqual(quant_params.zero_point.shape, (32, 4))
+      self.assertEqual(quant_params.scale.shape, (32, 4))
+      self.assertIsNotNone(quant_params.quantized_data)
+      self.assertTupleEqual(
+          cast(np.ndarray, quant_params.quantized_data).shape, test_data.shape
+      )
+
+    scales = np.repeat(quant_params.scale, 32, axis=1)
+    adjusted_test_data = quant_params.quantized_data * scales
+    for i, row in enumerate(test_data):
+      real_max = np.max(np.abs(row))
+      adjusted_max = np.max(np.abs(adjusted_test_data[i]))
+      # Check that some clipping occurred.
+      with self.subTest(name="CheckClipping"):
+        self.assertLess(adjusted_max, real_max)
+
+    with self.subTest(name="CheckQuantParamsValues"):
+      self.assertTrue(np.all(quant_params.zero_point == 0))
+      # See TFL_OP_TO_BLOCKWISE_WEIGHT_QUANTIZED_DIM.
+      self.assertEqual(quant_params.quantized_dimension, 1)
+
 
 if __name__ == "__main__":
   googletest.main()