ai-edge-quantizer-nightly 0.1.0.dev20250511__py3-none-any.whl → 0.1.0.dev20250513__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -24,6 +24,7 @@ from ai_edge_quantizer import qtyping
 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 naive_min_max_quantize
 from ai_edge_quantizer.algorithms.uniform_quantize import octav
 
@@ -58,6 +59,8 @@ class AlgorithmName(str, enum.Enum):
   FLOAT_CASTING = float_casting.ALGORITHM_KEY
   DEQUANTIZED_WEIGHT_RECOVERY = dequantized_weight_recovery.ALGORITHM_KEY
   OCTAV = octav.ALGORITHM_KEY
+  HADAMARD_ROTATION = hadamard_rotation.ALGORITHM_KEY
+
 
 ### MIN/MAX_UNIFORM_QUANT ###
 
@@ -104,6 +107,7 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
         common_quantize.materialize_dynamic_update_slice
     ),
     _TFLOpName.STABLEHLO_COMPOSITE: common_quantize.materialize_composite,
+    _TFLOpName.PAD: common_quantize.materialize_pad,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -237,6 +241,7 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
         common_quantize.materialize_dynamic_update_slice
     ),
     _TFLOpName.STABLEHLO_COMPOSITE: common_quantize.materialize_composite,
+    _TFLOpName.PAD: common_quantize.materialize_pad,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():
@@ -250,3 +255,32 @@ for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():
           octav.get_tensor_quant_params,
       ),
   )
+
+# Register the Hadamard Rotation algorithm.
+register_op_quant_config_validation_func(
+    AlgorithmName.HADAMARD_ROTATION,
+    common_quantize.check_op_quantization_config,
+)
+
+# Register a config check policy for the Hadamard Rotation algorithm.
+register_config_check_policy_func(
+    AlgorithmName.HADAMARD_ROTATION,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+# Register specialized hadamard rotation materialize functions.
+_HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
+    _TFLOpName.FULLY_CONNECTED: hadamard_rotation.materialize_fully_connected,
+    _TFLOpName.EMBEDDING_LOOKUP: hadamard_rotation.materialize_embedding_lookup,
+})
+for (
+    op_name,
+    materialize_func,
+) in _HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      AlgorithmName.HADAMARD_ROTATION,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=materialize_func,
+  )

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -680,6 +680,23 @@ def materialize_split(
   )
 
 
+def materialize_pad(
+    get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: dict[str, Any],
+) -> list[qtyping.TensorTransformationParams]:
+  """Materialize tensors in tfl.pad."""
+  return common_utils.materialize_standard_op(
+      op_info,
+      graph_info,
+      tensor_name_to_qsv,
+      get_tensor_quant_params_fn,
+      constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
+      inputs_to_ignore=[1],  # Padding value does not need to be quantized.
+  )
+
+
 def _get_tensor_shape_for_blockwise(
     tensor_shape: Sequence[int], quantized_dim: int, block_size: int
 ) -> list[int]:

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py

@@ -0,0 +1,352 @@
+# 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 Hadamard Rotation quantization."""
+
+from typing import Any, Optional
+import numpy as np
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.uniform_quantize import octav
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+
+
+ALGORITHM_KEY = "HADAMARD_ROTATION"
+
+
+def _make_hadamard_matrix(size: int) -> np.ndarray:
+  """Generates a Hadamard matrix of the given size.
+
+  Args:
+    size: The size of the Hadamard matrix. Must be a power of 2. This
+      represents a single dimension. E.g. if size is 4, then the Hadamard matrix
+      is a 4x4 matrix.
+
+  Returns:
+    The Hadamard matrix.
+
+  Raises:
+    ValueError: If the size is not a power of 2.
+  """
+  if size <= 0 or (size & (size - 1)) != 0:
+    raise ValueError("Hadamard matrix size must be a power of 2. ")
+  h = h2 = np.array([[1, 1], [1, -1]])
+  current_size = 2
+  while current_size < size:
+    h = np.kron(h, h2)
+    current_size *= 2
+  return h / np.sqrt(size)
+
+
+def _rotate_with_diagonal_hadamard(
+    tensor_content: np.ndarray,
+    axis: int,
+):
+  """Quantizes the given float array using the diagonal Hadamard algorithm.
+
+  Args:
+    tensor_content: The float array to quantize.
+    axis: The axis of the tensor to quantize.
+
+  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.
+  """
+  if axis != 1:
+    raise ValueError(
+        "Hadamard rotation is only supported for 2D tensors with quantized"
+        " dimension 0."
+    )
+
+  # Use the largest power of 2 that is a factor of the dimension and then
+  # tile this Hadamard matrix along the diagonal. 2**30 is just a large power
+  # of 2 to calculate this factor.
+  hadamard_size = np.gcd(tensor_content.shape[axis], 2 ** 30)
+  diagonal_size = tensor_content.shape[axis] // hadamard_size
+  random_vector = np.ones(hadamard_size, dtype=np.int8)
+
+  # Use a canonical Hadamard matrix.
+  hadamard = _make_hadamard_matrix(hadamard_size)
+  hadamard_diagonal = np.kron(np.eye(diagonal_size), hadamard)
+  w_rotated = np.einsum("ij,aj->ai", hadamard_diagonal, tensor_content)
+  return w_rotated, hadamard_size, random_vector
+
+
+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.
+
+  This function will rotate the tensor with a Hadamard matrix and then
+  quantize it with OCTAV.
+
+  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).
+    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_content is None:
+    raise ValueError("Hadamard rotation is only supported for weight tensors.")
+
+  if tensor_qsv is not None:
+    raise ValueError(
+        "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_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)
+  if quantized_dim != 0:
+    raise ValueError(
+        f"Unsupported quantized dimension: {quantized_dim}. Only 0 is"
+        " 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
+
+  # Rotate the tensor with a Hadamard matrix.
+  w_rotated, hadamard_size, random_vector = _rotate_with_diagonal_hadamard(
+      tensor_content, axis=reduce_axis
+  )
+
+  # Get the quantized values of the rotated tensor.
+  qparams = octav.get_tensor_quant_params(
+      op_info, tensor_quant_config, w_rotated, tensor_qsv
+  )
+
+  return qtyping.UniformQuantParams(
+      quantized_dimension=qparams.quantized_dimension,
+      num_bits=qparams.num_bits,
+      scale=qparams.scale,
+      zero_point=qparams.zero_point,
+      symmetric=qparams.symmetric,
+      quantized_data=qparams.quantized_data,
+      block_size=qparams.block_size,
+      hadamard=qtyping.UniformQuantParams.HadamardRotationParams(
+          random_binary_vector=random_vector,
+          hadamard_size=hadamard_size,
+      ),
+  )
+
+
+def materialize_fully_connected(
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: Optional[dict[str, Any]] = None,  # pylint: disable=unused-argument
+) -> list[qtyping.TensorTransformationParams]:
+  """Materialize the fully_connected op.
+
+  Args:
+    op_info: Aggregated information about the op (e.g., quantization config).
+    graph_info: Graph information needed to perform quantization for the op.
+    tensor_name_to_qsv: A map of tensor name to quantization parameters.
+
+  Returns:
+    Quantization configuration for the tensors associated with the op (e.g.,
+    weights, bias).
+  """
+  op_tensor_params = []
+
+  # Materialize weight.
+  weight_tensor_index = 1
+  weight_tensor = graph_info.subgraph_tensors[
+      op_info.op.inputs[weight_tensor_index]
+  ]
+  tensor_data = tfl_flatbuffer_utils.get_tensor_data(
+      weight_tensor, graph_info.buffers
+  )
+  # quant_params contains the rotated and quantized weights done by
+  # get_tensor_quant_params().
+  quant_params = get_tensor_quant_params(
+      op_info,
+      op_info.op_quant_config.weight_tensor_config,
+      tensor_data,
+      None,
+  )
+  transformations = [qtyping.QuantTransformation.QUANTIZE_TENSOR]
+  op2tensor_params = qtyping.OpToTensorParams(
+      subgraph_op_id=op_info.subgraph_op_index,
+      parameters=quant_params,
+      transformations=transformations,
+  )
+  weight_transformation_params = qtyping.TensorTransformationParams(
+      tensor_name=tfl_flatbuffer_utils.get_tensor_name(weight_tensor),
+      consumers=[op2tensor_params],
+  )
+
+  # Materialize input. A hadamard rotation op should be inserted on the input
+  # tensor to do the inverse of the weight's transformation.
+  input_tensor_index = 0
+  input_tensor = graph_info.subgraph_tensors[
+      op_info.op.inputs[input_tensor_index]
+  ]
+  transformations = [
+      qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
+  ]
+  op2tensor_params = qtyping.OpToTensorParams(
+      subgraph_op_id=op_info.subgraph_op_index,
+      parameters=quant_params,
+      transformations=transformations,
+  )
+  input_transformation_params = qtyping.TensorTransformationParams(
+      tensor_name=tfl_flatbuffer_utils.get_tensor_name(input_tensor),
+      consumers=[op2tensor_params],
+  )
+  op_tensor_params.append(input_transformation_params)
+  op_tensor_params.append(weight_transformation_params)
+
+  # Materialize bias. Since static quantization is not supported, we do not
+  # quantize the bias tensor.
+  bias_tensor_index = 2
+  bias_tensor = graph_info.subgraph_tensors[
+      op_info.op.inputs[bias_tensor_index]
+  ]
+  no_quant_tensor_params = qtyping.OpToTensorParams(
+      subgraph_op_id=op_info.subgraph_op_index,
+      transformations=[qtyping.QuantTransformation.NO_QUANTIZE],
+  )
+  bias_transformation_params = qtyping.TensorTransformationParams(
+      tensor_name=tfl_flatbuffer_utils.get_tensor_name(bias_tensor),
+      consumers=[no_quant_tensor_params],
+  )
+  op_tensor_params.append(bias_transformation_params)
+
+  # Materialize output. Since static quantization is not supported, we do not
+  # quantize the output tensor.
+  output_tensor_index = 0
+  output_tensor = graph_info.subgraph_tensors[
+      op_info.op.outputs[output_tensor_index]
+  ]
+  no_quant_tensor_params = qtyping.OpToTensorParams(
+      subgraph_op_id=op_info.subgraph_op_index,
+      transformations=[qtyping.QuantTransformation.NO_QUANTIZE],
+  )
+  output_transformation_params = qtyping.TensorTransformationParams(
+      tensor_name=tfl_flatbuffer_utils.get_tensor_name(output_tensor),
+      producer=no_quant_tensor_params,
+  )
+  op_tensor_params.append(output_transformation_params)
+
+  return op_tensor_params
+
+
+def materialize_embedding_lookup(
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: Optional[dict[str, Any]] = None,  # pylint: disable=unused-argument
+) -> list[qtyping.TensorTransformationParams]:
+  """Materialize the embedding_lookup op.
+
+  Args:
+    op_info: Aggregated information about the op (e.g., quantization config).
+    graph_info: Graph information needed to perform quantization for the op.
+    tensor_name_to_qsv: A map of tensor name to quantization parameters.
+
+  Returns:
+    Quantization configuration for the tensors associated with the op (e.g.,
+    weights, bias).
+  """
+  op_tensor_params = []
+
+  # Materialize lookup.
+  lookup_tensor_index = 0
+  lookup_tensor = graph_info.subgraph_tensors[
+      op_info.op.inputs[lookup_tensor_index]
+  ]
+  transformations = [
+      qtyping.QuantTransformation.NO_QUANTIZE,
+  ]
+  op2tensor_params = qtyping.OpToTensorParams(
+      subgraph_op_id=op_info.subgraph_op_index,
+      parameters=None,
+      transformations=transformations,
+  )
+  lookup_transformation_params = qtyping.TensorTransformationParams(
+      tensor_name=tfl_flatbuffer_utils.get_tensor_name(lookup_tensor),
+      consumers=[op2tensor_params],
+  )
+  op_tensor_params.append(lookup_transformation_params)
+
+  # Materialize embedding. The embedding table should be rotated and then
+  # quantized.
+  embedding_tensor_index = 1
+  embedding_tensor = graph_info.subgraph_tensors[
+      op_info.op.inputs[embedding_tensor_index]
+  ]
+  tensor_data = tfl_flatbuffer_utils.get_tensor_data(
+      embedding_tensor, graph_info.buffers
+  )
+  quant_params = get_tensor_quant_params(
+      op_info,
+      op_info.op_quant_config.weight_tensor_config,
+      tensor_data,
+      None,
+  )
+  transformations = [qtyping.QuantTransformation.QUANTIZE_TENSOR]
+  op2tensor_params = qtyping.OpToTensorParams(
+      subgraph_op_id=op_info.subgraph_op_index,
+      parameters=quant_params,
+      transformations=transformations,
+  )
+  weight_transformation_params = qtyping.TensorTransformationParams(
+      tensor_name=tfl_flatbuffer_utils.get_tensor_name(embedding_tensor),
+      consumers=[op2tensor_params],
+  )
+  op_tensor_params.append(weight_transformation_params)
+
+  # Materialize output. A hadamard rotation op should be inserted on the output
+  # tensor to do the inverse of the embedding's transformation.
+  output_tensor_index = 0
+  output_tensor = graph_info.subgraph_tensors[
+      op_info.op.outputs[output_tensor_index]
+  ]
+  transformations = [
+      qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
+  ]
+  op2tensor_params = qtyping.OpToTensorParams(
+      subgraph_op_id=op_info.subgraph_op_index,
+      parameters=quant_params,
+      transformations=transformations,
+  )
+  output_transformation_params = qtyping.TensorTransformationParams(
+      tensor_name=tfl_flatbuffer_utils.get_tensor_name(output_tensor),
+      producer=op2tensor_params,
+  )
+  op_tensor_params.append(output_transformation_params)
+
+  return op_tensor_params

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py

@@ -0,0 +1,216 @@
+# 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_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/default_policy.py

@@ -183,7 +183,8 @@ DEFAULT_JSON_POLICY = """
       "SELECT_V2",
       "DYNAMIC_UPDATE_SLICE",
       "SELECT_V2",
-      "STABLEHLO_COMPOSITE"
+      "STABLEHLO_COMPOSITE",
+      "PAD"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -214,7 +215,8 @@ DEFAULT_JSON_POLICY = """
       "SELECT_V2",
       "DYNAMIC_UPDATE_SLICE",
       "SELECT_V2",
-      "STABLEHLO_COMPOSITE"
+      "STABLEHLO_COMPOSITE",
+      "PAD"
     ],
     "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],
     "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],

ai_edge_quantizer/params_generator.py

@@ -508,6 +508,7 @@ def _compatible_tensor_params(
   float_source_transformations = [
       _QuantTrans.ADD_QUANTIZE,
       _QuantTrans.NO_QUANTIZE,
+      _QuantTrans.INSERT_HADAMARD_ROTATION,
   ]
   quantized_source_transformations = [
       _QuantTrans.QUANTIZE_TENSOR,

ai_edge_quantizer/qtyping.py

@@ -20,7 +20,7 @@ from collections.abc import MutableMapping
 import copy
 import dataclasses
 import enum
-from typing import Any, Optional, Union, Callable
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 from typing_extensions import TypeAlias
@@ -62,6 +62,7 @@ class TFLOperationName(str, enum.Enum):
   SELECT_V2 = 'SELECT_V2'
   DYNAMIC_UPDATE_SLICE = 'DYNAMIC_UPDATE_SLICE'
   STABLEHLO_COMPOSITE = 'STABLEHLO_COMPOSITE'
+  PAD = 'PAD'
 
 
 class QuantizeMode(enum.Enum):
@@ -113,6 +114,8 @@ class QuantTransformation(enum.Enum):
   DUPLICATE_BUFFER = 5
   # Duplicate the tensor.
   DUPLICATE_TENSOR = 6
+  # Insert the aeq.hadamard_rotation op.
+  INSERT_HADAMARD_ROTATION = 7
 
 
 @dataclasses.dataclass(frozen=True)
@@ -128,8 +131,35 @@ class UniformQuantParams:
     quantized_data: The quantized data.
     block_size: The block size for blockwise quantization, block_size=0 meaning
       no blockwise quantization.
+    hadamard: The Hadamard rotation parameters, if set.
   """
 
+  class HadamardRotationParams:
+    """Parameters for the Hadamard rotation.
+
+    Attributes:
+      random_binary_vector: The random binary vector for the Hadamard rotation.
+        TODO(b/415392354): Randomization is an experimental feature that's
+        currently not implemented yet hence this is always 1. We will add
+        support or remove in the future.
+      hadamard_size: The size of the Hadamard matrix.
+    """
+
+    random_binary_vector: np.ndarray
+    hadamard_size: int
+
+    def __init__(self, random_binary_vector: np.ndarray, hadamard_size: int):
+      self.random_binary_vector = random_binary_vector
+      self.hadamard_size = hadamard_size
+
+    def __eq__(self, other):
+      if other.__class__ is not self.__class__:
+        return NotImplemented
+      return (
+          np.array_equal(self.random_binary_vector, other.random_binary_vector)
+          and self.hadamard_size == other.hadamard_size
+      )
+
   num_bits: int
   quantized_dimension: Optional[int]
   scale: np.ndarray
@@ -137,6 +167,7 @@ class UniformQuantParams:
   symmetric: bool = True
   quantized_data: Optional[np.ndarray] = None
   block_size: int = 0
+  hadamard: Optional[HadamardRotationParams] = None
 
   @classmethod
   def from_tfl_tensor_details(cls, tensor_detail) -> 'UniformQuantParams':
@@ -180,6 +211,7 @@ class UniformQuantParams:
         and self.symmetric == other.symmetric
         and _compare_array_or_none(self.quantized_data, other.quantized_data)
         and self.block_size == other.block_size
+        and self.hadamard == other.hadamard
     )
 
 
@@ -492,6 +524,7 @@ class IOOperator:
   outputs: list[int]
   op_key: TFLOperationName
 
+
 # The function signature for `get_tensor_quant_params_fn`.
 GetTensorQuantParamsFuncSignature = Callable[
     [

ai_edge_quantizer/transformation_performer.py

@@ -25,6 +25,7 @@ from ai_edge_quantizer.transformations import dequant_insert
 from ai_edge_quantizer.transformations import duplicate_buffer
 from ai_edge_quantizer.transformations import duplicate_tensor
 from ai_edge_quantizer.transformations import emulated_subchannel
+from ai_edge_quantizer.transformations import insert_hadamard_rotation
 from ai_edge_quantizer.transformations import quant_insert
 from ai_edge_quantizer.transformations import quantize_tensor
 from ai_edge_quantizer.transformations import transformation_utils
@@ -80,6 +81,9 @@ class TransformationPerformer:
         qtyping.QuantTransformation.DUPLICATE_TENSOR: (
             duplicate_tensor.duplicate_tensor
         ),
+        qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION: (
+            insert_hadamard_rotation.insert_hadamard_rotation
+        ),
     }
     # transformations are seprated in two categories:
     # op_insertion_transformations are transformations that only insert ops
@@ -91,6 +95,7 @@ class TransformationPerformer:
         qtyping.QuantTransformation.ADD_QUANTIZE,
         qtyping.QuantTransformation.DUPLICATE_BUFFER,
         qtyping.QuantTransformation.DUPLICATE_TENSOR,
+        qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
     ])
     self._op_replacement_transformations = set(
         [qtyping.QuantTransformation.EMULATED_SUBCHANNEL]

ai_edge_quantizer/transformations/insert_hadamard_rotation.py

@@ -0,0 +1,209 @@
+# 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.
+# ==============================================================================
+
+"""Hadamard rotation pattern transformation."""
+
+from flatbuffers import flexbuffers
+import numpy as np
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+
+def _to_flexbuffer(
+    hadamard_size: int,
+    random_binary_vector: list[np.int8],
+) -> bytes:
+  """Converts hadamard_size to flexbuffer."""
+  fbb = flexbuffers.Builder()
+  with fbb.Map():
+    fbb.Int('hadamard_size', hadamard_size)
+    fbb.VectorFromElements('random_binary_vector', random_binary_vector)
+  return fbb.Finish()
+
+
+def _is_producer_embedding_lookup(
+    transformation: transformation_utils.TransformationInput,
+) -> bool:
+  """Checks if the tensor's producer is an embedding lookup op."""
+  if transformation.producer == -1:
+    return False
+  else:
+    return (
+        transformation.op_codes[
+            transformation.subgraph.operators[
+                transformation.producer
+            ].opcodeIndex
+        ].builtinCode
+        == schema_py_generated.BuiltinOperator.EMBEDDING_LOOKUP
+    )
+
+
+def _is_fully_connected(
+    transformation: transformation_utils.TransformationInput, op_id: int
+) -> bool:
+  """Checks if the any of the tensor's consumers is a fully connected op."""
+  return (
+      transformation.op_codes[
+          transformation.subgraph.operators[op_id].opcodeIndex
+      ].builtinCode
+      == schema_py_generated.BuiltinOperator.FULLY_CONNECTED
+  )
+
+
+def _update_embedding_lookup_consumers(
+    transformation: transformation_utils.TransformationInput,
+    new_tensor_id: int,
+) -> bool:
+  """Updates the consumers of the embedding lookup op to use the new tensor.
+
+  Args:
+    transformation: The transformation input to update the consumers of.
+    new_tensor_id: The new tensor id to use as the input to the embedding lookup
+      consumers.
+  """
+  for consumer in transformation.consumers:
+    # If the consumer is a graph output and not an op, we can ignore it here
+    # since the graph output will be updated later.
+    if consumer == -1:
+      continue
+    consumer_op = transformation.subgraph.operators[consumer]
+    # Find the input that was attached to the insertion point, and replace it
+    # with the new tensor.
+    for i in range(len(consumer_op.inputs)):
+      if consumer_op.inputs[i] == transformation.tensor_id:
+        consumer_op.inputs[i] = new_tensor_id
+
+
+def _update_fully_connected_consumers(
+    transformation: transformation_utils.TransformationInput,
+    new_tensor_id: int,
+) -> bool:
+  """Updates the fully connected op(s) to use the new tensor.
+
+  Since the new tensor is inserted to the fully_connected's input, we need to
+  scan each consumer (in case of multiple fully_connected ops), and update
+  the input tensor to the new tensor.
+
+  Args:
+    transformation: The transformation input to update the consumers of.
+    new_tensor_id: The new tensor id to use as the input to the fully connected
+      consumers.
+
+  Returns:
+    True if the fully connected op(s) were updated to use the new tensor.
+  """
+  updated = False
+  for consumer in transformation.consumers:
+    if _is_fully_connected(transformation, consumer):
+      transformation.subgraph.operators[consumer].inputs[0] = new_tensor_id
+      updated = True
+  return updated
+
+
+def insert_hadamard_rotation(
+    transformation_input: transformation_utils.TransformationInput,
+) -> qtyping.TransformationInfo:
+  """Inserts a custom aeq.hadamard_rotation op on this tensor.
+
+  This function works for float32 tensors only.
+
+  Args:
+    transformation_input: The transformation input to insert the custom op on.
+
+  Returns:
+    The transformation info of the inserted custom op.
+
+  Raises:
+    ValueError: If the transformation input is not a uniform quantization
+    transformation.
+    ValueError: If the Hadamard quantization params are not set.
+    ValueError: If the tensor is not a float32 tensor.
+    ValueError: If no supported ops were found as the tensor's producer or
+    consumers.
+  """
+  if not isinstance(
+      transformation_input.quant_params, qtyping.UniformQuantParams
+  ):
+    raise ValueError('Hadamard rotation supports uniform quantization only')
+
+  if transformation_input.quant_params.hadamard is None:
+    raise ValueError(
+        'Hadamard rotation quantization params are not set but op insertion is'
+        ' requested.'
+    )
+
+  tensor = transformation_input.subgraph.tensors[transformation_input.tensor_id]
+  if tensor.type != schema_py_generated.TensorType.FLOAT32:
+    raise ValueError(
+        'The Hadamard rotation op supports float32 tensors only. Got'
+        f' {tensor.type} tensor.'
+    )
+
+  # Create new custom op with the current tensor as input and a new activation
+  # tensor as output.
+  custom_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.CUSTOM,
+      transformation_input.op_codes,
+      'aeq.hadamard_rotation',
+  )
+  custom_op = schema_py_generated.OperatorT()
+  custom_op.opcodeIndex = custom_op_code_idx
+  custom_op.inputs = [transformation_input.tensor_id]
+  custom_op.customOptions = _to_flexbuffer(
+      transformation_input.quant_params.hadamard.hadamard_size,
+      transformation_input.quant_params.hadamard.random_binary_vector.tolist(),
+  )
+  new_tensor_id = transformation_utils.add_new_activation_tensor(
+      tensor.name + b'_rotated',
+      tensor.shapeSignature
+      if tensor.shapeSignature is not None
+      else tensor.shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+  custom_op.outputs = [new_tensor_id]
+
+  # Update the users of this tensor to use the new tensor.
+  if _is_producer_embedding_lookup(transformation_input):
+    _update_embedding_lookup_consumers(transformation_input, new_tensor_id)
+  elif not _update_fully_connected_consumers(
+      transformation_input, new_tensor_id
+  ):
+    raise ValueError(
+        'The Hadamard rotation op supports embedding lookup and fully connected'
+        ' ops only, but no such ops were found.'
+    )
+
+  # If the tensor is a graph output, we need to replace the tensor with the
+  # new tensor.
+  for i, output in enumerate(transformation_input.subgraph.outputs):
+    if output == transformation_input.tensor_id:
+      transformation_input.subgraph.outputs[i] = new_tensor_id
+
+  # Find the actual insertion point. The insertion point should be after the
+  # producer op and before the first consumer op. The max() operation ensures
+  # that we're not using -1 as the insertion point.
+  first_consumer_id = min(transformation_input.consumers)
+  op_id = max(transformation_input.producer + 1, first_consumer_id)
+
+  # Insert the custom op.
+  transformation_input.subgraph.operators.insert(op_id, custom_op)
+
+  return qtyping.TransformationInfo(
+      op_id=op_id,
+      num_ops_added=1,
+      output_tensor_id=new_tensor_id,
+  )

ai_edge_quantizer/transformations/insert_hadamard_rotation_test.py

@@ -0,0 +1,200 @@
+# 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 insertion of the Hadamard rotation custom op."""
+
+import os
+import numpy as np
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import insert_hadamard_rotation
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+_TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile('..')
+
+
+class InsertHadamardRotationFullyConnectedTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, 'tests/models/single_fc_bias.tflite'
+    )
+    self.model = tfl_flatbuffer_utils.read_model(model_path)
+    self.params = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.ones(1),
+        zero_point=np.zeros(1),
+        hadamard=qtyping.UniformQuantParams.HadamardRotationParams(
+            random_binary_vector=np.ones(1),
+            hadamard_size=2,
+        ),
+    )
+
+  def test_raise_unsupported_qparams(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'uniform quantization' in str(err)
+    ):
+      insert_hadamard_rotation.insert_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=qtyping.NonLinearQuantParams(
+                  num_bits=16, quantized_data=None
+              ),
+          )
+      )
+
+  def test_raise_missing_hadamard_data(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'quantization params are not set' in str(err)
+    ):
+      insert_hadamard_rotation.insert_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=qtyping.UniformQuantParams(
+                  num_bits=8,
+                  quantized_dimension=None,
+                  scale=np.ones(1),
+                  zero_point=np.zeros(1),
+              ),
+          )
+      )
+
+  def test_raise_non_float32_tensor(self):
+    self.model.subgraphs[0].tensors[
+        0
+    ].type = schema_py_generated.TensorType.INT32
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: 'float32 tensors' in str(err)
+    ):
+      insert_hadamard_rotation.insert_hadamard_rotation(
+          transformation_utils.TransformationInput(
+              tensor_id=0,
+              op_codes=self.model.operatorCodes,
+              buffers=self.model.buffers,
+              subgraph=self.model.subgraphs[0],
+              producer=-1,
+              consumers=[-1],
+              quant_params=self.params,
+          ),
+      )
+
+  def test_insert_single_custom_op(self):
+    # Insert aeq.hadamard_rotation before fully_connected
+    info = insert_hadamard_rotation.insert_hadamard_rotation(
+        transformation_utils.TransformationInput(
+            tensor_id=0,
+            op_codes=self.model.operatorCodes,
+            buffers=self.model.buffers,
+            subgraph=self.model.subgraphs[0],
+            producer=-1,
+            consumers=[-1],
+            quant_params=self.params,
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 0)
+    self.assertEqual(info.num_ops_added, 1)
+    # Model had 4 tensors, added 1.
+    self.assertEqual(info.output_tensor_id, 4)
+    self.assertLen(subgraph.tensors, 5)
+    # Model had 1 op, added a new one.
+    self.assertLen(self.model.operatorCodes, 2)
+    self.assertEqual(
+        self.model.operatorCodes[1].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # First op is now the custom op, precedes fully_connected.
+    self.assertEqual(
+        self.model.operatorCodes[subgraph.operators[0].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # Input to the custom op is graph input
+    self.assertEqual(subgraph.operators[0].inputs[0], 0)
+    # Input to the FC is the custom op output
+    self.assertEqual(subgraph.operators[1].inputs[0], 4)
+
+
+class InsertHadamardRotationEmbeddingLookupTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, 'tests/models/embedding_lookup.tflite'
+    )
+    self.model = tfl_flatbuffer_utils.read_model(model_path)
+    self.params = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.ones(1),
+        zero_point=np.zeros(1),
+        hadamard=qtyping.UniformQuantParams.HadamardRotationParams(
+            random_binary_vector=np.ones(1),
+            hadamard_size=2,
+        ),
+    )
+
+  def test_insert_single_custom_op(self):
+    # Insert aeq.hadamard_rotation after embedding_lookup
+    info = insert_hadamard_rotation.insert_hadamard_rotation(
+        transformation_utils.TransformationInput(
+            tensor_id=2,
+            op_codes=self.model.operatorCodes,
+            buffers=self.model.buffers,
+            subgraph=self.model.subgraphs[0],
+            producer=0,
+            consumers=[-1],
+            quant_params=self.params,
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 1)
+    self.assertEqual(info.num_ops_added, 1)
+    # Model had 3 tensors, added 1.
+    self.assertEqual(info.output_tensor_id, 3)
+    self.assertLen(subgraph.tensors, 4)
+    # Model had 1 op, added a new one.
+    self.assertLen(self.model.operatorCodes, 2)
+    self.assertEqual(
+        self.model.operatorCodes[1].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # Second op is now the custom op, after embedding_lookup.
+    self.assertEqual(
+        self.model.operatorCodes[subgraph.operators[1].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.CUSTOM,
+    )
+    # Input to the custom op is embedding's output
+    self.assertEqual(subgraph.operators[1].inputs[0], 2)
+    # Custom op's output is the new tensor
+    self.assertEqual(subgraph.operators[1].outputs[0], 3)
+
+
+if __name__ == '__main__':
+  googletest.main()

ai_edge_quantizer/utils/test_utils.py

@@ -33,6 +33,39 @@ _OpQuantConfig = qtyping.OpQuantizationConfig
 _AlgorithmName = quantizer.AlgorithmName
 
 
+DEFAULT_ACTIVATION_QUANT_SETTING = _TensorQuantConfig(
+    num_bits=8,
+    symmetric=False,
+    granularity=qtyping.QuantGranularity.TENSORWISE,
+)
+DEFAULT_WEIGHT_QUANT_SETTING = _TensorQuantConfig(
+    num_bits=8,
+    symmetric=True,
+    granularity=qtyping.QuantGranularity.CHANNELWISE,
+)
+
+
+def get_static_activation_quant_setting(
+    num_bits: int, symmetric: bool
+) -> _TensorQuantConfig:
+  return _TensorQuantConfig(
+      num_bits=num_bits,
+      symmetric=symmetric,
+      granularity=qtyping.QuantGranularity.TENSORWISE,
+  )
+
+
+def get_static_op_quant_config(
+    activation_config: _TensorQuantConfig = DEFAULT_ACTIVATION_QUANT_SETTING,
+    weight_config: _TensorQuantConfig = DEFAULT_WEIGHT_QUANT_SETTING,
+) -> _OpQuantConfig:
+  return qtyping.OpQuantizationConfig(
+      activation_tensor_config=activation_config,
+      weight_tensor_config=weight_config,
+      compute_precision=_ComputePrecision.INTEGER,
+  )
+
+
 def get_path_to_datafile(path):
   """Get the path to the specified file in the data dependencies.
 

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -56,6 +56,7 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.DYNAMIC_UPDATE_SLICE: (
         schema.BuiltinOperator.DYNAMIC_UPDATE_SLICE
     ),
+    _TFLOpName.PAD: schema.BuiltinOperator.PAD,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

{ai_edge_quantizer_nightly-0.1.0.dev20250511.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250513.dist-info}/METADATA

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

@@ -1,18 +1,18 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=0uootLsVD6h9ph9TrnXZMI-ExkX8UvXSV0lbWxBLybU,10492
+ai_edge_quantizer/algorithm_manager.py,sha256=p-wX2ksIV1hbWEQz-uUnbNMVgDJrsIiIOU2ZYX2ZrTM,11726
 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=n7AD9j7UScR-CieoI6DQRMeiG_fhLBfSLRiM4460xaM,11895
 ai_edge_quantizer/calibrator_test.py,sha256=C_oWOaRugPKYX74jF-eRFH-k6nGOdA8I9_uPiocaOuE,11900
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
-ai_edge_quantizer/default_policy.py,sha256=81z4cruBK7mGFt8xFRZK5LKya65axuZwo2zpbcYSicc,11099
+ai_edge_quantizer/default_policy.py,sha256=zNTeiI_eP5-dLL3P_VWIQB3RzXBrb06peJKngLnSSFY,11125
 ai_edge_quantizer/model_modifier.py,sha256=teGa8I6kGvn6TQY6Xv53YFIc_pQEhNvM9Zb4bvhezyw,7110
 ai_edge_quantizer/model_modifier_test.py,sha256=cJd04SLOG-fQZZNZPcisoBLx3cLtWEwGqUBbLb-pif4,4751
 ai_edge_quantizer/model_validator.py,sha256=Hj0_5o-Oa3dSlJ3ryVjRhvsyelHNyek1GrtG9buMczg,13153
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
-ai_edge_quantizer/params_generator.py,sha256=NEZeHVVIeynmhRzPjl9o-acvWfauFgCS4i45pWFw3V8,20052
+ai_edge_quantizer/params_generator.py,sha256=j1BV2cGFLlQmUY6aoW5uglYqf77b9ytN8oZ1gh6o0mM,20096
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=FqelZu7j0fGBRSCv_VVsuf3VmbfVlYJGgsjvdMXGgaw,15284
+ai_edge_quantizer/qtyping.py,sha256=LKn9w53wmw3gPO0E4DKOhj8gkx9efjXMoipGnsJyGiU,16453
 ai_edge_quantizer/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
 ai_edge_quantizer/quantizer_test.py,sha256=K_HBA56JkFI3HL8VLWCqGEfC0ISh5ldMKoNyBdGRAJg,20368
 ai_edge_quantizer/recipe.py,sha256=FR0uJceumZrnle2VRSOQZ1uXup4S1cTYKRH-N53mWRo,2919
@@ -21,17 +21,19 @@ ai_edge_quantizer/recipe_manager_test.py,sha256=LulVxsYp6TBGFI2PLCUCd4VsFq8ELpC7
 ai_edge_quantizer/recipe_test.py,sha256=Fg_sfxovI2fRjk5qdu18ghOvXdUvhDR1TxbE0GHDczc,3381
 ai_edge_quantizer/transformation_instruction_generator.py,sha256=R7A90Qj6iQQROrznXmXLJd-5yXq0PRHbLOdNY51dEu4,27913
 ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=E0QSDCav6N6izlJ-a1ZJOsb2VEUxuxBmTbt0-EgDdxY,49890
-ai_edge_quantizer/transformation_performer.py,sha256=zAzrQOb2n2IpB3qopmKV59e5E99HmTOL60QTCn9-7kA,12821
+ai_edge_quantizer/transformation_performer.py,sha256=nkkqbs81ITB5u2FoWeG9z5d8EtLtCiltOxcQ34okN8E,13091
 ai_edge_quantizer/transformation_performer_test.py,sha256=xk6A3LStCyPclN51--9uO7XjSxNfZmpdfvrzOL0maNM,20349
 ai_edge_quantizer/algorithms/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/nonlinear_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting.py,sha256=Bs9CK7wZAw6jNaZ8xEtbwO2vM34VYXNZSMVWvxJo9nw,9297
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py,sha256=EqIHGEZ1LgUrTN7zf880RuAzEv3Qy7kgh5ivObJGHSo,22646
 ai_edge_quantizer/algorithms/uniform_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=JyMuqVns3meKOA1DomXYopxqhZE65yamlVZjBF6yOmY,27731
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=-ugXQ4cZoVMrgOVs4m73ozI-49CRyT0YuKrLS5begW8,28297
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py,sha256=qMmKbWqxrCoVKbLKHn9WuCrGKPfHkEyU0Nmhokh8Qeo,2597
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=Fk3s9Qy2A_hjUepFOUmTwIZ_wKYVPbdDX4eoP-eoAQU,8726
 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=f9HhFCAavbrdYkQQH37ivbKRuRXC1g1TO2FmILMApN8,12389
+ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py,sha256=kN9aCPt1yTleiDBiH4g2RZ1vMBm7WAf5pmVFjmYCH-0,7617
 ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py,sha256=Divlsn3NjNGtH0vlvE91wxL-VHb4q1nUE0JTDGiEtYc,8572
 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=e5wYtki-vl739gSVAZHAKcs2hA87GvFUjVoSUPlnkyM,6433
@@ -50,6 +52,8 @@ ai_edge_quantizer/transformations/duplicate_tensor.py,sha256=HF1uuKFm5kFF6X0XUpd
 ai_edge_quantizer/transformations/duplicate_tensor_test.py,sha256=s-RqSxNBMfVJyCunXz2eb7-KA6UiBmbOmL7phLslENQ,5056
 ai_edge_quantizer/transformations/emulated_subchannel.py,sha256=HVaRxoC8PCAvy3xeMv3OIymukUy_yW1zK0xN8Ann6I4,13602
 ai_edge_quantizer/transformations/emulated_subchannel_test.py,sha256=gZP6u9NdPXl7s19qB_Un8evou9ZZV6I9Gy0E1rdobHM,7722
+ai_edge_quantizer/transformations/insert_hadamard_rotation.py,sha256=rBbKgcVKHie38NT2UQ7KQ1xCb2tRu_rVl0yFloOAW_A,7562
+ai_edge_quantizer/transformations/insert_hadamard_rotation_test.py,sha256=iV1p3nZfHUATV2YRoBOYurnu3pLy8n3aFppLWGQOPdA,7268
 ai_edge_quantizer/transformations/quant_insert.py,sha256=jn6HsJaV-sqBiFPY-Aqbd64t8zgcYVkEkZI375x_FWY,3958
 ai_edge_quantizer/transformations/quant_insert_test.py,sha256=X9ptPDvJCFkR5tejKnD1SlHFGPazQTW-wNNMV9MEAuw,10107
 ai_edge_quantizer/transformations/quantize_tensor.py,sha256=kjaNrw9mnrn0t8u0vey9S_uPz3iVUicwy4rluxVqV3E,7617
@@ -59,15 +63,15 @@ ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=E90O4PYSjz
 ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/utils/calibration_utils.py,sha256=1Fj9MIO6aLZIRgyd4axvZN4S_O64nB_-Miu1WP664js,2536
 ai_edge_quantizer/utils/calibration_utils_test.py,sha256=Z-AcdTieesWFKyKBb08ZXm4Mgu6cvJ4bg2-MJ7hLD10,2856
-ai_edge_quantizer/utils/test_utils.py,sha256=HwZCIpO9fJRAhuN6t6voXKOYQtcioFtt_tpkAlDsAYk,6205
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=NKtw60BJAjIE6Yww8B1vJpxXwp4MSERmpKajXJWm5rI,10568
+ai_edge_quantizer/utils/test_utils.py,sha256=fXwQ353P7tSy7W4Hs6YskIbCLLaBYGA724hMMbcqCUk,7129
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=zNlR_SJAkDi-EX63O3pNpFLVqSktysScZKgKk1XT3c8,10616
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=WoewyiZpaua80oP0tpgyrw5Ws1v7f4vl88vdzS0UjDE,13490
 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.dev20250511.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250511.dist-info/METADATA,sha256=ENLYvQfun3PLDrzvKxWBsjTrbyWjM7KRBETP2EKs8Kk,1528
-ai_edge_quantizer_nightly-0.1.0.dev20250511.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250511.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250511.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250513.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250513.dist-info/METADATA,sha256=zL_JxmjzCHEwIUmLkDGzI6B7IACt6YnVQSpaxaNUujY,1528
+ai_edge_quantizer_nightly-0.1.0.dev20250513.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250513.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250513.dist-info/RECORD,,