ai-edge-quantizer-nightly 0.4.0.dev20250925__py3-none-any.whl → 0.4.0.dev20250926__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -25,9 +25,11 @@ from ai_edge_quantizer.algorithms.nonlinear_quantize import float_casting
 from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
 from ai_edge_quantizer.algorithms.uniform_quantize import dequantized_weight_recovery
 from ai_edge_quantizer.algorithms.uniform_quantize import hadamard_rotation
+from ai_edge_quantizer.algorithms.uniform_quantize import mse
 from ai_edge_quantizer.algorithms.uniform_quantize import naive_min_max_quantize
 from ai_edge_quantizer.algorithms.uniform_quantize import octav
 
+
 # TODO: b/399775701 - Clean up this file.
 
 _TFLOpName = qtyping.TFLOperationName
@@ -60,6 +62,7 @@ class AlgorithmName(str, enum.Enum):
   DEQUANTIZED_WEIGHT_RECOVERY = dequantized_weight_recovery.ALGORITHM_KEY
   OCTAV = octav.ALGORITHM_KEY
   HADAMARD_ROTATION = hadamard_rotation.ALGORITHM_KEY
+  MSE = mse.ALGORITHM_KEY
 
 
 ### MIN/MAX_UNIFORM_QUANT ###
@@ -322,3 +325,36 @@ for (
       calibration_func=naive_min_max_quantize.min_max_calibrate,
       materialize_func=materialize_func,
   )
+
+
+# Register the MSE algorithm.
+register_op_quant_config_validation_func(
+    AlgorithmName.MSE,
+    common_quantize.check_op_quantization_config,
+)
+
+# Register a config check policy for the MSE algorithm.
+register_config_check_policy_func(
+    AlgorithmName.MSE,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+# Register specialized MSE materialize functions.
+_MSE_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
+    _TFLOpName.FULLY_CONNECTED: common_quantize.materialize_fc_conv,
+    _TFLOpName.EMBEDDING_LOOKUP: common_quantize.materialize_embedding_lookup,
+})
+for (
+    op_name,
+    materialize_func,
+) in _MSE_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      AlgorithmName.MSE,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=functools.partial(
+          materialize_func,
+          mse.get_tensor_quant_params,
+      ),
+  )

ai_edge_quantizer/algorithms/uniform_quantize/mse.py

@@ -0,0 +1,125 @@
+# 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.
+# ==============================================================================
+
+"""Implements the MSE quantization."""
+
+import dataclasses
+from typing import Any, Optional
+import numpy as np
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
+from ai_edge_quantizer.algorithms.uniform_quantize import naive_min_max_quantize
+from ai_edge_quantizer.algorithms.uniform_quantize import uniform_quantize_tensor
+from ai_edge_quantizer.algorithms.utils import common_utils
+
+ALGORITHM_KEY = "MSE"
+
+# Coefficients from offline numeric analysis.
+_MSE_QUANT_MULS = {
+    8: 0.05408,
+    4: 0.37755,
+}
+
+
+def get_tensor_quant_params(
+    op_info: qtyping.OpInfo,
+    tensor_quant_config: qtyping.TensorQuantizationConfig,
+    tensor_content: Optional[np.ndarray] = None,
+    tensor_qsv: Optional[dict[str, Any]] = None,
+) -> qtyping.UniformQuantParams:
+  """Returns the quantization parameters for a tensor.
+
+  Args:
+    op_info: Aggregated information about the op (e.g., quantization config).
+    tensor_quant_config: The quantization config for the tensor.
+    tensor_content: The content of the tensor. When None, it means the tensor is
+      not a weight tensor (e.g. static quantization) so we fallback to using
+      naive_min_max_quantize.
+    tensor_qsv: A dictionary containing the min/max of the tensor.
+
+  Raises:
+    ValueError: If the blockwise quantization is requested.
+    ValueError: If the asymmetric quantization is requested.
+    ValueError: `tensor_qsv` must contain min/max values, or `tensor_content`
+      must be provided so that they can be inferred.
+  """
+  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
+    raise ValueError(
+        "Blockwise quantization is not supported for MSE quantization."
+    )
+
+  # Fallback to naive_min_max_quantize.py for non-weight tensors.
+  if tensor_content is None:
+    return naive_min_max_quantize.get_tensor_quant_params(
+        op_info, tensor_quant_config, tensor_content, tensor_qsv
+    )
+
+  if not tensor_quant_config.symmetric:
+    raise ValueError(
+        f"Unsupported symmetry: {tensor_quant_config.symmetric}. MSE"
+        " supports symmetric quantization only for now."
+    )
+
+  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
+    # possible that this information is missing here. In that case we
+    # collect min/max on the spot.
+    tensor_min_max = common_quantize.init_tensor_min_max(
+        tensor_content,
+        op_info,
+    )
+  else:
+    tensor_min_max = tensor_qsv
+
+  if "min" not in tensor_min_max or "max" not in tensor_min_max:
+    raise ValueError(
+        "min and max must be provided to produce tensor quantization"
+        " parameters. Check if the correct calibration results are passed into"
+        " the ParamsGenerator."
+    )
+
+  quantized_dim = common_utils.get_weight_quantized_dim(
+      op_info, tensor_content, tensor_quant_config.granularity
+  )
+
+  reshaped_data = tensor_content
+  reduce_dims = common_utils.get_reduce_dims(
+      quantized_dim, tensor_content.shape
+  )
+
+  multiplier = _MSE_QUANT_MULS[tensor_quant_config.num_bits]
+  scale = multiplier * np.sqrt(
+      np.mean(reshaped_data**2, axis=reduce_dims, keepdims=True)
+  )
+  zp = np.zeros_like(scale, dtype=np.int32)
+
+  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,
+  )
+
+  quantized_vars = uniform_quantize_tensor.uniform_quantize(
+      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/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,
+          ),
+          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_nightly-0.4.0.dev20250925.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=Xu2JKUUayFE5O3-nQPKvQ-0xjqtqPM8ITXgR7BDugN0,13916
+ai_edge_quantizer/algorithm_manager.py,sha256=XkLMG_wQqf_X6swp6YBIhJpIbIdRcOt2LJ_6oTZ3GzU,14956
 ai_edge_quantizer/algorithm_manager_api.py,sha256=u903TG0s1uIDhJqfeJne3CFl8A93phZrwgV2-hwdcXU,9247
 ai_edge_quantizer/algorithm_manager_api_test.py,sha256=w6bSONvXkX6bzXAGc0-7b6gNDt9oz9ieq97KP8Sg_JU,7666
 ai_edge_quantizer/calibrator.py,sha256=Sms7_AIHPH9G5xFaz5Ef3a5gPhxuIWQI8d2LUM8C96I,12071
@@ -34,6 +34,8 @@ ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha
 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=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/mse.py,sha256=qiIyzogATGVxjYwxzH0cZvgwPSPBJv_3y8NSumHZXTk,4561
+ai_edge_quantizer/algorithms/uniform_quantize/mse_test.py,sha256=-_P4jQJ7gVo0FNSapP3sIGcnhwfjQHW1AKLfoiAlS_s,7142
 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
@@ -70,8 +72,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.dev20250925.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20250925.dist-info/METADATA,sha256=cNrHy0lFqC22asNZiYjtXRwPC4DdQRhHn4lcwltzKGo,1508
-ai_edge_quantizer_nightly-0.4.0.dev20250925.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20250925.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20250925.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/METADATA,sha256=6ymhTobT9E998G5IZCfmysJbNYMcQr_vrngEtQf5VsE,1508
+ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/RECORD,,