ai-edge-quantizer-nightly 0.1.0.dev20250415__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/mse_test.py

@@ -0,0 +1,195 @@
+# 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.
+# ==============================================================================
+
+import os
+from typing import cast
+
+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 mse
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+
+
+class MseQuantizeTest(parameterized.TestCase):
+  """Tests for general functions for MSE."""
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(666)
+    self._test_model_path = os.path.join(
+        test_utils.get_path_to_datafile("../../tests/models"),
+        "conv_fc_mnist.tflite",
+    )
+    self._test_model = tfl_flatbuffer_utils.read_model(self._test_model_path)
+    # The test model has one subgraph for now.
+    self._graph_info = qtyping.GraphInfo(
+        subgraph_tensors=self._test_model.subgraphs[0].tensors,
+        buffers=self._test_model.buffers,
+    )
+    self._tensor_name_to_qsv = {}
+    subgraph0 = self._test_model.subgraphs[0]
+    self._subgraph_op_index = 3
+    self._fc_op = subgraph0.operators[self._subgraph_op_index]
+    self._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=None,
+        ),
+    )
+
+  def test_get_tensor_quant_params_raises_error_with_unsupported_symmetry(self):
+    err_msg = "Unsupported symmetry"
+    test_data = np.array([[-7, 7], [4, -4], [4, -4], [7, 7]])
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: err_msg in str(err)
+    ):
+      _ = mse.get_tensor_quant_params(
+          op_info=self._fc_op_info,
+          tensor_quant_config=qtyping.TensorQuantizationConfig(
+              num_bits=4,
+              symmetric=False,
+              granularity=qtyping.QuantGranularity.CHANNELWISE,
+          ),
+          tensor_content=test_data,
+      )
+
+  def test_get_tensor_quant_params_raises_error_with_unsupported_granularity(
+      self,
+  ):
+    err_msg = "Blockwise quantization is not supported"
+    test_data = np.array([[-7, 7], [4, -4], [4, -4], [7, 7]])
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: err_msg in str(err)
+    ):
+      _ = mse.get_tensor_quant_params(
+          op_info=self._fc_op_info,
+          tensor_quant_config=qtyping.TensorQuantizationConfig(
+              num_bits=4,
+              symmetric=True,
+              granularity=qtyping.QuantGranularity.BLOCKWISE_32,
+          ),
+          tensor_content=test_data,
+      )
+
+  def test_get_tensor_quant_params_succeeds_with_qsv(self):
+    # Fall back to naive_min_max_quantize.py for non-weight tensors.
+    tensor_quant_params = mse.get_tensor_quant_params(
+        op_info=self._fc_op_info,
+        tensor_quant_config=qtyping.TensorQuantizationConfig(
+            num_bits=8,
+            granularity=qtyping.QuantGranularity.TENSORWISE,
+        ),
+        tensor_qsv={
+            "min": np.array([-1]),
+            "max": np.array([1]),
+        },
+    )
+
+    self.assertIsNone(tensor_quant_params.quantized_dimension)
+    scale = tensor_quant_params.scale
+    self.assertEqual(scale.shape, (1,))
+    self.assertSequenceAlmostEqual(scale.flatten(), [1 / 127])
+
+    # Zero point should be zero for symmetric quantization.
+    zp = tensor_quant_params.zero_point
+    self.assertEqual(np.sum(zp), 0)
+    self.assertEqual(zp.shape, (1,))
+
+  def test_get_tensor_quant_params_succeeds_with_tensorwise_granularity(self):
+    test_data = np.array([
+        [-1e5, 25, -50, 75, -100, 125],
+        [25, -30, 50, -75, 1e5, -125],
+        [50, -60, 70, -80, 90, -100],
+    ])
+    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 = mse.get_tensor_quant_params(
+        op_info=fc_op_info,
+        tensor_quant_config=tensor_config,
+        tensor_content=test_data,
+    )
+
+    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)
+      self.assertIsNotNone(quant_params.quantized_data)
+      self.assertTupleEqual(
+          cast(np.ndarray, quant_params.quantized_data).shape, test_data.shape
+      )
+
+    with self.subTest(name="CheckQuantParamsValues"):
+      self.assertTrue(np.all(quant_params.zero_point == 0))
+
+  def test_get_tensor_quant_params_succeeds_with_channelwise_granularity(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],
+    ])
+    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 = mse.get_tensor_quant_params(
+        op_info=fc_op_info,
+        tensor_quant_config=tensor_config,
+        tensor_content=test_data,
+    )
+
+    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)
+      self.assertTupleEqual(
+          cast(np.ndarray, quant_params.quantized_data).shape, test_data.shape
+      )
+
+    with self.subTest(name="CheckQuantParamsValues"):
+      self.assertTrue(np.all(quant_params.zero_point == 0))
+      self.assertEqual(quant_params.quantized_dimension, 0)
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py

@@ -15,8 +15,8 @@
 
 """Performs naive min/max uniform quantization."""
 
+import dataclasses
 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,65 +75,43 @@ 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,
       num_bits=tensor_quant_config.num_bits,
       symmetric=tensor_quant_config.symmetric,
       quantized_dimension=quantized_dim,
-      block_size=tensor_quant_config.block_size,
+      block_size=uniform_quantize_tensor.extract_block_size_from_granularity(
+          tensor_quant_config.granularity
+      ),
   )
   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,
+      uniform_quantize_tensor.is_blockwise(tensor_quant_config.granularity),
   )
   # Update with quantized values.
-  return qtyping.UniformQuantParams(
-      scale=scale,
-      zero_point=zp,
-      num_bits=tensor_quant_config.num_bits,
-      symmetric=tensor_quant_config.symmetric,
-      quantized_dimension=quantized_dim,
-      quantized_data=quantized_vars,
-      block_size=tensor_quant_config.block_size,
+  return dataclasses.replace(quant_params, quantized_data=quantized_vars)
+
+
+def check_if_quantized(tensor: Any) -> bool:
+  """Checks if the tensor is quantized."""
+  return (
+      tensor.quantization is not None and tensor.quantization.scale is not None
   )
 
 
@@ -158,6 +136,13 @@ def init_qsvs(
   op_qsvs = {}
 
   inputs_to_ignore = inputs_to_ignore or []
+  quantized_inputs_to_ignore = [
+      opr_idx
+      for opr_idx, tensor_idx in enumerate(op_info.op.inputs)
+      if check_if_quantized(graph_info.subgraph_tensors[tensor_idx])
+  ]
+  inputs_to_ignore.extend(quantized_inputs_to_ignore)
+
   outputs_to_ignore = outputs_to_ignore or []
   for opr_idx, tensor_idx in enumerate(op_info.op.inputs):
     if tensor_idx != -1 and opr_idx not in inputs_to_ignore:
@@ -190,6 +175,7 @@ def min_max_calibrate(
     tensor_content_map: dict[str, np.ndarray],
     inputs_to_ignore: Optional[list[int]] = None,
     outputs_to_ignore: Optional[list[int]] = None,
+    valid_range: tuple[float, float] = (-3e38, 3e38),
 ) -> dict[str, qtyping.QSV]:
   """Collect quantization statistics variable (QSV, e.g., min/max) for the op.
 
@@ -199,11 +185,18 @@ def min_max_calibrate(
     tensor_content_map: A map of tensor name to tensor content.
     inputs_to_ignore: Input tensor indices to ignore.
     outputs_to_ignore: Output tensor indices to ignore.
+    valid_range: The valid range for tensor content, excluding the boundaries.
+      Tensor values outside this range are ignored during calibration. Defaults
+      to an approximate bfloat16 range. This range is chosen to address issues
+      with `padv2` where a bfloat16 -inf padding constant can cause problems.
+      Values exceeding this range can lead to quantization issues and are
+      therefore excluded from min/max calibration.
 
   Returns:
     A dictionary with key as tensor name and value as the collected QSV.
   """
   op_qsvs = {}
+  min_val, max_val = valid_range
 
   def _collect_activation_tensor_min_max(tensor_idx):
     tensor = graph_info.subgraph_tensors[tensor_idx]
@@ -215,12 +208,25 @@ def min_max_calibrate(
       return
     tensor_name = tfl_flatbuffer_utils.get_tensor_name(tensor)
     tensor_content = tensor_content_map[tensor_name]
+    qsv_shape = (1,) * tensor_content.ndim
+    filter_mask = (tensor_content > min_val) & (tensor_content < max_val)
+    if np.any(filter_mask):
+      tensor_content = tensor_content[filter_mask]
+    # Reshape is needed to ensure the scalar min/max have the same number of
+    # dimensions as the input tensor array, for compatibility with subsequent
+    # operations.
     op_qsvs[tensor_name] = {
-        "min": np.min(tensor_content, axis=None, keepdims=True),
-        "max": np.max(tensor_content, axis=None, keepdims=True),
+        "min": np.min(tensor_content, axis=None).reshape(qsv_shape),
+        "max": np.max(tensor_content, axis=None).reshape(qsv_shape),
     }
 
   inputs_to_ignore = inputs_to_ignore or []
+  quantized_inputs_to_ignore = [
+      opr_idx
+      for opr_idx, tensor_idx in enumerate(tfl_op.inputs)
+      if check_if_quantized(graph_info.subgraph_tensors[tensor_idx])
+  ]
+  inputs_to_ignore.extend(quantized_inputs_to_ignore)
   outputs_to_ignore = outputs_to_ignore or []
   for i, tensor_idx in enumerate(tfl_op.inputs):
     if tensor_idx != -1 and i not in inputs_to_ignore:

ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py

@@ -17,6 +17,7 @@ import os
 from typing import cast
 
 from absl.testing import parameterized
+import ml_dtypes
 import numpy as np
 
 from tensorflow.python.platform import googletest
@@ -165,8 +166,7 @@ class NaiveMinMaxQuantizeTest(parameterized.TestCase):
     weight_tensor_config = _TensorQuantConfig(
         num_bits=4,
         symmetric=True,
-        granularity=qtyping.QuantGranularity.BLOCKWISE,
-        block_size=2,
+        granularity=qtyping.QuantGranularity.BLOCKWISE_32,
     )
     op_info = qtyping.OpInfo(
         op=fc_op,
@@ -176,30 +176,69 @@ class NaiveMinMaxQuantizeTest(parameterized.TestCase):
             weight_tensor_config=weight_tensor_config,
         ),
     )
-    test_data = np.array([[-7, 7], [4, -4], [4, -4], [7, 7]])
+    test_data = np.random.uniform(low=-10, high=10, size=(4, 32)).astype(
+        np.float32
+    )
     quant_params = naive_min_max_quantize.get_tensor_quant_params(
         op_info=op_info,
         tensor_quant_config=weight_tensor_config,
         tensor_content=test_data,
     )
-    scale = quant_params.scale
     zp = quant_params.zero_point
-    expected_scale = np.array([
-        [1],
-        [0.5703125],
-        [0.5703125],
-        [1],
-    ])
-    expected_zp = np.zeros([4, 1])
-    self.assertTrue(np.array_equal(zp, expected_zp))
-    self.assertTrue(np.array_equal(scale, expected_scale))
+    self.assertEqual(zp.shape, (4, 1))
+    self.assertTrue(np.array_equal(zp, np.zeros([4, 1])))
+
+    self.assertEqual(quant_params.scale.shape, (4, 1))
+    expected_scales = np.max(np.abs(test_data), axis=1, keepdims=True) / 7.0
+    expected_scales = (
+        expected_scales.astype(ml_dtypes.bfloat16)
+        .astype(np.float16)
+        .astype(np.float32)
+    )
+    self.assertTrue(np.allclose(quant_params.scale, expected_scales, atol=1e-5))
+
     self.assertIsNotNone(quant_params.quantized_data)
     self.assertTupleEqual(
         cast(np.ndarray, quant_params.quantized_data).shape, test_data.shape
     )
-    self.assertEqual(quant_params.block_size, 2)
+    self.assertEqual(quant_params.block_size, 32)
     self.assertEqual(quant_params.quantized_dimension, 1)
 
+  def test_calibrate_ignores_inf_min_max(self):
+    """Tests that calibration ignores infinity values."""
+    # Sample input/output data for the fc op.
+    input_tensor_name = "sequential/flatten/Reshape"
+    output_tensor_name = (
+        "sequential/dense/MatMul;sequential/dense/Relu;sequential/dense/BiasAdd"
+    )
+    bloat16_inf = 3.39e38
+    tensor_content_map = {
+        input_tensor_name: np.array(
+            [[-np.inf, 1.0, 5.0, np.inf, bloat16_inf]], dtype=np.float32
+        ),
+        output_tensor_name: np.array(
+            [[6.0, 7.0, -bloat16_inf, 9.0, np.inf]], dtype=np.float32
+        ),
+    }
+    # Read from Model Explorer.
+    subgraph0 = self._test_model.subgraphs[0]
+    fc_op = subgraph0.operators[3]
+    op_qsvs = naive_min_max_quantize.min_max_calibrate(
+        fc_op,
+        self._graph_info,
+        tensor_content_map,
+        inputs_to_ignore=[1, 2],  # Ignore weight and bias.
+        outputs_to_ignore=[],
+    )
+
+    self.assertIn(input_tensor_name, op_qsvs)
+    self.assertEqual(op_qsvs[input_tensor_name]["min"], 1.0)
+    self.assertEqual(op_qsvs[input_tensor_name]["max"], 5.0)
+
+    self.assertIn(output_tensor_name, op_qsvs)
+    self.assertEqual(op_qsvs[output_tensor_name]["min"], 6.0)
+    self.assertEqual(op_qsvs[output_tensor_name]["max"], 9.0)
+
 
 if __name__ == "__main__":
   googletest.main()

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 uniform_quantize_tensor.is_blockwise(tensor_quant_config.granularity):
+    reshaped_data, reduce_dims = (
+        uniform_quantize_tensor.reshape_data_for_blockwise(
+            tensor_content,
+            op_info.op_name,
+            tensor_quant_config.granularity,
+        )
+    )
+  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 uniform_quantize_tensor.is_blockwise(tensor_quant_config.granularity):
+    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,
   )
 
@@ -164,11 +172,17 @@ def get_tensor_quant_params(
       num_bits=tensor_quant_config.num_bits,
       symmetric=tensor_quant_config.symmetric,
       quantized_dimension=quantized_dim,
-      block_size=tensor_quant_config.block_size,
+      block_size=uniform_quantize_tensor.extract_block_size_from_granularity(
+          tensor_quant_config.granularity
+      ),
   )
 
   quantized_vars = uniform_quantize_tensor.uniform_quantize(
-      tensor_content, quant_params
+      tensor_content,
+      quant_params,
+      is_blockwise_quant=uniform_quantize_tensor.is_blockwise(
+          tensor_quant_config.granularity
+      ),
   )
 
   return dataclasses.replace(quant_params, quantized_data=quantized_vars)