ai-edge-quantizer-nightly 0.4.0.dev20250930__py3-none-any.whl → 0.4.0.dev20251002__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -61,7 +61,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
+  HADAMARD_ROTATION = hadamard_rotation.CUSTOM_OP_ALGORITHM_KEY
+  DECOMPOSED_HADAMARD_ROTATION = hadamard_rotation.DECOMPOSED_ALGORITHM_KEY
   MSE = mse.ALGORITHM_KEY
 
 
@@ -130,6 +131,7 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.REDUCE_MIN: common_quantize.materialize_reduce_min,
     _TFLOpName.EQUAL: common_quantize.materialize_equal,
     _TFLOpName.NOT_EQUAL: common_quantize.materialize_not_equal,
+    _TFLOpName.MIRROR_PAD: common_quantize.materialize_mirror_pad,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -283,6 +285,7 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     _TFLOpName.REDUCE_MIN: common_quantize.materialize_reduce_min,
     _TFLOpName.EQUAL: common_quantize.materialize_equal,
     _TFLOpName.NOT_EQUAL: common_quantize.materialize_not_equal,
+    _TFLOpName.MIRROR_PAD: common_quantize.materialize_mirror_pad,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():
@@ -311,8 +314,12 @@ register_config_check_policy_func(
 
 # 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,
+    _TFLOpName.FULLY_CONNECTED: (
+        hadamard_rotation.materialize_fully_connected_custom_op
+    ),
+    _TFLOpName.EMBEDDING_LOOKUP: (
+        hadamard_rotation.materialize_embedding_lookup_custom_op
+    ),
 })
 for (
     op_name,
@@ -326,6 +333,36 @@ for (
       materialize_func=materialize_func,
   )
 
+register_op_quant_config_validation_func(
+    AlgorithmName.DECOMPOSED_HADAMARD_ROTATION,
+    common_quantize.check_op_quantization_config,
+)
+
+register_config_check_policy_func(
+    AlgorithmName.DECOMPOSED_HADAMARD_ROTATION,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+_DECOMPOSED_HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
+    _TFLOpName.FULLY_CONNECTED: (
+        hadamard_rotation.materialize_fully_connected_decomposed
+    ),
+    _TFLOpName.EMBEDDING_LOOKUP: (
+        hadamard_rotation.materialize_embedding_lookup_decomposed
+    ),
+})
+for (
+    op_name,
+    materialize_func,
+) in _DECOMPOSED_HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      AlgorithmName.DECOMPOSED_HADAMARD_ROTATION,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=materialize_func,
+  )
+
 
 # Register the MSE algorithm.
 register_op_quant_config_validation_func(

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -748,6 +748,34 @@ def materialize_padv2(
   )
 
 
+def materialize_mirror_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.mirror_pad.
+
+  Args:
+    get_tensor_quant_params_fn: Function to get quantization parameters for the
+      tensor.
+    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:
+    A list of `qtyping.TensorTransformationParams` for the tensors in the op.
+  """
+  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],  # Paddings tensor does not need to be quantized.
+  )
+
+
 def materialize_squared_difference(
     get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
     op_info: qtyping.OpInfo,

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py

@@ -23,16 +23,17 @@ from ai_edge_quantizer.algorithms.utils import common_utils
 from ai_edge_quantizer.utils import tfl_flatbuffer_utils
 
 
-ALGORITHM_KEY = "HADAMARD_ROTATION"
+CUSTOM_OP_ALGORITHM_KEY = "HADAMARD_ROTATION"
+DECOMPOSED_ALGORITHM_KEY = "DECOMPOSED_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.
+    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.
@@ -157,9 +158,10 @@ def get_tensor_quant_params(
   )
 
 
-def materialize_fully_connected(
+def _materialize_fully_connected(
     op_info: qtyping.OpInfo,
     graph_info: qtyping.GraphInfo,
+    is_decomposed: bool = False,
     tensor_name_to_qsv: Optional[dict[str, Any]] = None,  # pylint: disable=unused-argument
 ) -> list[qtyping.TensorTransformationParams]:
   """Materialize the fully_connected op.
@@ -167,12 +169,20 @@ def materialize_fully_connected(
   Args:
     op_info: Aggregated information about the op (e.g., quantization config).
     graph_info: Graph information needed to perform quantization for the op.
+    is_decomposed: Whether to use decomposed Hadamard rotation ops or a custom
+      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).
   """
+  if op_info.op_quant_config.weight_tensor_config is None:
+    raise ValueError(
+        "Weight tensor quantization config is not provided for Hadamard"
+        " Rotation quantization."
+    )
+
   op_tensor_params = []
 
   # Materialize weight.
@@ -209,7 +219,9 @@ def materialize_fully_connected(
       op_info.op.inputs[input_tensor_index]
   ]
   transformations = [
-      qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
+      qtyping.QuantTransformation.INSERT_DECOMPOSED_HADAMARD_ROTATION
+      if is_decomposed
+      else qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
   ]
   op2tensor_params = qtyping.OpToTensorParams(
       subgraph_op_id=op_info.subgraph_op_index,
@@ -258,16 +270,45 @@ def materialize_fully_connected(
   return op_tensor_params
 
 
-def materialize_embedding_lookup(
+def materialize_fully_connected_custom_op(
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: Optional[dict[str, Any]] = None,  # pylint: disable=unused-argument
+) -> list[qtyping.TensorTransformationParams]:
+  return _materialize_fully_connected(
+      op_info,
+      graph_info,
+      is_decomposed=False,
+      tensor_name_to_qsv=tensor_name_to_qsv,
+  )
+
+
+def materialize_fully_connected_decomposed(
     op_info: qtyping.OpInfo,
     graph_info: qtyping.GraphInfo,
     tensor_name_to_qsv: Optional[dict[str, Any]] = None,  # pylint: disable=unused-argument
+) -> list[qtyping.TensorTransformationParams]:
+  return _materialize_fully_connected(
+      op_info,
+      graph_info,
+      is_decomposed=True,
+      tensor_name_to_qsv=tensor_name_to_qsv,
+  )
+
+
+def _materialize_embedding_lookup(
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    is_decomposed: bool = False,
+    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.
+    is_decomposed: Whether to use decomposed Hadamard rotation ops or a custom
+      op.
     tensor_name_to_qsv: A map of tensor name to quantization parameters.
 
   Returns:
@@ -329,7 +370,9 @@ def materialize_embedding_lookup(
       op_info.op.outputs[output_tensor_index]
   ]
   transformations = [
-      qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
+      qtyping.QuantTransformation.INSERT_DECOMPOSED_HADAMARD_ROTATION
+      if is_decomposed
+      else qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
   ]
   op2tensor_params = qtyping.OpToTensorParams(
       subgraph_op_id=op_info.subgraph_op_index,
@@ -343,3 +386,29 @@ def materialize_embedding_lookup(
   op_tensor_params.append(output_transformation_params)
 
   return op_tensor_params
+
+
+def materialize_embedding_lookup_custom_op(
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: Optional[dict[str, Any]] = None,  # pylint: disable=unused-argument
+) -> list[qtyping.TensorTransformationParams]:
+  return _materialize_embedding_lookup(
+      op_info,
+      graph_info,
+      is_decomposed=False,
+      tensor_name_to_qsv=tensor_name_to_qsv,
+  )
+
+
+def materialize_embedding_lookup_decomposed(
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: Optional[dict[str, Any]] = None,  # pylint: disable=unused-argument
+) -> list[qtyping.TensorTransformationParams]:
+  return _materialize_embedding_lookup(
+      op_info,
+      graph_info,
+      is_decomposed=True,
+      tensor_name_to_qsv=tensor_name_to_qsv,
+  )

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation_test.py

@@ -63,7 +63,7 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
     )
 
   def test_materialize_fully_connected_basic(self):
-    params = hadamard_rotation.materialize_fully_connected(
+    params = hadamard_rotation.materialize_fully_connected_custom_op(
         self._op_info, self._graph_info, self._tensor_name_to_qsv
     )
     fc_input = params[0]
@@ -111,7 +111,7 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
             ),
         ),
     )
-    params = hadamard_rotation.materialize_fully_connected(
+    params = hadamard_rotation.materialize_fully_connected_custom_op(
         self._op_info, self._graph_info, self._tensor_name_to_qsv
     )
     self.assertLen(params, 4)
@@ -152,7 +152,7 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
             ),
         ),
     )
-    params = hadamard_rotation.materialize_fully_connected(
+    params = hadamard_rotation.materialize_fully_connected_custom_op(
         self._op_info, self._graph_info, self._tensor_name_to_qsv
     )
     self.assertLen(params, 4)
@@ -179,6 +179,34 @@ class HadamardRotationFullyConnectedTest(parameterized.TestCase):
     ):
       self.assertEqual(weight.consumers[0].parameters.quantized_dimension, 1)
 
+  def test_materialize_fully_connected_decomposed(self):
+    params = hadamard_rotation.materialize_fully_connected_decomposed(
+        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.assertEqual(
+        fc_input.consumers[0].transformations,
+        [qtyping.QuantTransformation.INSERT_DECOMPOSED_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]
@@ -344,7 +372,7 @@ class HadamardRotationEmbeddingLookupTest(parameterized.TestCase):
             ),
         ),
     )
-    params = hadamard_rotation.materialize_embedding_lookup(
+    params = hadamard_rotation.materialize_embedding_lookup_custom_op(
         op_info, self._graph_info, self._tensor_name_to_qsv
     )
     self.assertLen(params, 3)
@@ -371,6 +399,43 @@ class HadamardRotationEmbeddingLookupTest(parameterized.TestCase):
           [qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION],
       )
 
+  def test_materialize_embedding_lookup_decomposed(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_decomposed(
+        op_info, self._graph_info, self._tensor_name_to_qsv
+    )
+    self.assertLen(params, 3)
+    lookup = params[0]
+    value = params[1]
+    output = params[2]
+    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_DECOMPOSED_HADAMARD_ROTATION],
+      )
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/default_policy.py

@@ -199,7 +199,8 @@ DEFAULT_JSON_POLICY = """
       "PADV2",
       "REDUCE_MIN",
       "EQUAL",
-      "NOT_EQUAL"
+      "NOT_EQUAL",
+      "MIRROR_PAD"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -248,7 +249,8 @@ DEFAULT_JSON_POLICY = """
       "PADV2",
       "REDUCE_MIN",
       "EQUAL",
-      "NOT_EQUAL"
+      "NOT_EQUAL",
+      "MIRROR_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

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

ai_edge_quantizer/qtyping.py

@@ -80,6 +80,7 @@ class TFLOperationName(str, enum.Enum):
   REDUCE_MIN = 'REDUCE_MIN'
   EQUAL = 'EQUAL'
   NOT_EQUAL = 'NOT_EQUAL'
+  MIRROR_PAD = 'MIRROR_PAD'
 
 
 class QuantizeMode(enum.Enum):
@@ -133,6 +134,9 @@ class QuantTransformation(enum.Enum):
   DUPLICATE_TENSOR = 6
   # Insert the aeq.hadamard_rotation op.
   INSERT_HADAMARD_ROTATION = 7
+  # Insert decomposed Hadamard rotation ops. This expresses the Hadamard
+  # rotation as matrix multiplication with Hadamard matrices.
+  INSERT_DECOMPOSED_HADAMARD_ROTATION = 8
 
 
 @dataclasses.dataclass(frozen=True)
@@ -305,6 +309,7 @@ class TensorQuantizationConfig:
       quantization.
     dtype: The data type of the tensor.
     block_size: The block size for blockwise quantization, ignored otherwise.
+    algorithm_key: The algorithm key to use for quantization.
   """
 
   num_bits: int

ai_edge_quantizer/transformation_performer.py

@@ -24,6 +24,7 @@ from ai_edge_quantizer import qtyping
 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 insert_decomposed_hadamard_rotation
 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
@@ -83,6 +84,9 @@ class TransformationPerformer:
         qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION: (
             insert_hadamard_rotation.insert_hadamard_rotation
         ),
+        qtyping.QuantTransformation.INSERT_DECOMPOSED_HADAMARD_ROTATION: (
+            insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation
+        ),
     }
     # transformations are seprated in two categories:
     # op_insertion_transformations are transformations that only insert ops
@@ -95,6 +99,7 @@ class TransformationPerformer:
         qtyping.QuantTransformation.DUPLICATE_BUFFER,
         qtyping.QuantTransformation.DUPLICATE_TENSOR,
         qtyping.QuantTransformation.INSERT_HADAMARD_ROTATION,
+        qtyping.QuantTransformation.INSERT_DECOMPOSED_HADAMARD_ROTATION,
     ])
     self._op_replacement_transformations = set(
         [qtyping.QuantTransformation.EMULATED_SUBCHANNEL]