ai-edge-quantizer-nightly 0.0.1.dev20250302__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py

@@ -31,8 +31,7 @@ _TensorQuantConfig = qtyping.TensorQuantizationConfig
 
 
 class CommonQuantizeTest(parameterized.TestCase):
-  """Tests for general quantize functions.
-  """
+  """Tests for general quantize functions."""
 
   def setUp(self):
     super().setUp()
@@ -69,6 +68,34 @@ class CommonQuantizeTest(parameterized.TestCase):
           default_policy.DEFAULT_CONFIG_CHECK_POLICY,
       )
 
+  def test_reshape_data_for_blockwise_raises_error_when_quantized_dim_not_divisible_by_block_size(
+      self,
+  ):
+    tensor_data = np.ones((24, 128), dtype=np.float32)
+    block_size = 256
+    quantized_dim = 1
+    with self.assertRaisesWithPredicateMatch(
+        ValueError,
+        lambda err: (
+            "Tensor quantization dimension must be divisible by block"
+            " size for blockwise quantization."
+        )
+        in str(err),
+    ):
+      common_quantize._reshape_data_for_blockwise(
+          tensor_data, quantized_dim, block_size
+      )
+
+  def test_reshape_data_for_blockwise_returns_correct_values(self):
+    tensor_data = np.ones((24, 128), dtype=np.float32)
+    block_size = 32
+    quantized_dim = 1
+    new_tensor_data, reduce_dim = common_quantize._reshape_data_for_blockwise(
+        tensor_data, quantized_dim, block_size
+    )
+    self.assertEqual(new_tensor_data.shape, (24, 4, 32))
+    self.assertEqual(reduce_dim, 2)
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py

@@ -70,17 +70,17 @@ def _get_scale(arr: np.ndarray, min_scale: float) -> float:
   return min_scale
 
 
-def get_zp_scale_from_2d_dequantized_symmetric_weights(
+def get_zp_scale_from_dequantized_symmetric_weights(
     dequant_vals: np.ndarray,
     quantized_dimension: Optional[int] = None,
     min_scale: float = 1e-9,
 ) -> tuple[np.ndarray, np.ndarray]:
-  """Calculates scale and zero point from 2D dequantized, symmetric weights.
+  """Calculates scale and zero point from dequantized and symmetric weights.
 
   Handles both per-tensor and per-channel (axis) quantization.
 
   Args:
-      dequant_vals: The 2D dequantized weight values (numpy array).
+      dequant_vals: The dequantized weight values (numpy array).
       quantized_dimension:  The dimension along which quantization was performed
         (0 or 1), or None for per-tensor quantization.
       min_scale: The minimum allowed scale value.
@@ -91,15 +91,9 @@ def get_zp_scale_from_2d_dequantized_symmetric_weights(
           - scales: Scales (scalar for per-tensor, array for per-channel).
 
   Raises:
-      ValueError: If `dequant_vals` is not 2D, or if
-          `quantized_dimension` is not 0, 1, or None.
+      ValueError: If `quantized_dimension` is not 0, 1, or None.
   """
 
-  if dequant_vals.ndim != 2:
-    raise ValueError(
-        f"Only 2D weights are supported. Got {dequant_vals.ndim} dimensions."
-    )
-
   if quantized_dimension not in (0, 1, None):
     raise ValueError(
         f"quantized_dimension must be 0, 1, or None. Got {quantized_dimension}"
@@ -112,23 +106,26 @@ def get_zp_scale_from_2d_dequantized_symmetric_weights(
     # Per-tensor quantization: One scale for the entire tensor.
     scales = _get_scale(dequant_vals.flatten(), min_scale)
     scales = np.array([[scales]])
-
   else:
     # Per-channel quantization: A scale for each slice along the dimension.
-    scales = []
-    for i in range(dequant_vals.shape[quantized_dimension]):
-      if quantized_dimension == 0:
-        vec = dequant_vals[i, :]
-      else:  # quantized_dimension == 1
-        vec = dequant_vals[:, i]
-      scales.append(_get_scale(vec, min_scale))
-
-    # Reshape for correct broadcasting.
-    scales = (
-        np.array(scales).reshape(-1, 1)
-        if quantized_dimension == 0
-        else np.array(scales).reshape(1, -1)
+    # Create a broadcasted array for per-channel scales. It should have the same
+    # number of dimensions as the input, with 1 in all dimensions except for the
+    # quantized dimension, which retains its original size.
+    scales = np.empty(
+        tuple(
+            [
+                1
+                if i != quantized_dimension
+                else dequant_vals.shape[quantized_dimension]
+                for i in range(dequant_vals.ndim)
+            ]
+        )
     )
+    for i in range(dequant_vals.shape[quantized_dimension]):
+      slices = [slice(None)] * dequant_vals.ndim
+      slices[quantized_dimension] = i
+      vec = dequant_vals[tuple(slices)]
+      scales[tuple(slices)] = _get_scale(vec, min_scale)
 
   zero_points = np.zeros_like(scales, dtype=np.int32)
   return zero_points, scales
@@ -153,7 +150,7 @@ def get_tensor_quant_params(
 
   Raises:
     ValueError: If the quantization granularity is blockwise, or if the tensor
-    is not a 2D symmetric weight tensor.
+    is not a symmetric weight tensor.
   """
   # Fallback to naive_min_max_quantize.py for non-weight tensors.
   if tensor_content is None:
@@ -161,24 +158,21 @@ def get_tensor_quant_params(
         op_info, tensor_quant_config, tensor_content, tensor_qsv
     )
 
-  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
+  if uniform_quantize_tensor.is_blockwise(tensor_quant_config.granularity):
     raise ValueError(
         "Blockwise quantization is not supported for dequantized weight"
         " recovery."
     )
-  if tensor_content.ndim != 2 or not tensor_quant_config.symmetric:
+  if not tensor_quant_config.symmetric:
     raise ValueError(
-        "Only 2D symmetric weights are supported for dequantized weight"
-        " recovery."
+        "Only symmetric weights are supported for dequantized weight recovery."
     )
 
-  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
+  )
 
-  zp, scale = get_zp_scale_from_2d_dequantized_symmetric_weights(
+  zp, scale = get_zp_scale_from_dequantized_symmetric_weights(
       dequant_vals=tensor_content,
       quantized_dimension=quantized_dim,
   )

ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py

@@ -62,7 +62,7 @@ class DequantizedWeightRecoveryTest(parameterized.TestCase):
   ):
     dequant_vals = scale * self._dummy_quantized_weights
     zp, recovered_scale = (
-        dequantized_weight_recovery.get_zp_scale_from_2d_dequantized_symmetric_weights(
+        dequantized_weight_recovery.get_zp_scale_from_dequantized_symmetric_weights(
             dequant_vals, quantized_dimension
         )
     )
@@ -72,17 +72,40 @@ class DequantizedWeightRecoveryTest(parameterized.TestCase):
     self.assertEqual(np.sum(zp), 0)
     self.assertEqual(zp.shape, scale.shape)
 
-  def test_tensor_zp_scale_from_2d_dequantized_symmetric_weights_raises_error_for_non_2d_weights(
-      self,
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="per-tensor-recovery",
+          quantized_dimension=None,
+          scale=np.array([0.1875]).reshape(1, 1),
+      ),
+      dict(
+          testcase_name="channel0-recovery",
+          quantized_dimension=0,
+          scale=np.array([0.1875, 1e-4, 12.3]).reshape(3, 1, 1),
+      ),
+      dict(
+          testcase_name="channel1-recovery",
+          quantized_dimension=1,
+          scale=np.array([0.003, 1.234]).reshape(1, 2, 1),
+      ),
+  )
+  def test_tensor_zp_scale_from_3d_dequantized_symmetric_weights_success(
+      self, quantized_dimension, scale
   ):
-    weights_3d = self._dummy_quantized_weights.reshape(1, 3, 4)
-    weights_3d = weights_3d * 1.02
-    with self.assertRaisesRegex(
-        ValueError, "Only 2D weights are supported. Got 3 dimensions."
-    ):
-      dequantized_weight_recovery.get_zp_scale_from_2d_dequantized_symmetric_weights(
-          weights_3d, quantized_dimension=None
-      )
+    dequant_vals = scale * self._dummy_quantized_weights.reshape(3, 2, 2)
+    zp, recovered_scale = (
+        dequantized_weight_recovery.get_zp_scale_from_dequantized_symmetric_weights(
+            dequant_vals, quantized_dimension
+        )
+    )
+    with self.subTest("shapes_match"):
+      self.assertEqual(recovered_scale.shape, scale.shape)
+      self.assertEqual(zp.shape, scale.shape)
+    with self.subTest("scale_value_match"):
+      self.assertSequenceAlmostEqual(recovered_scale.flatten(), scale.flatten())
+    with self.subTest("zp_is_zero"):
+      # Zero point should be zero for symmetric quantization.
+      self.assertEqual(np.sum(zp), 0)
 
   @parameterized.named_parameters(
       dict(testcase_name="negative_dimension", quantized_dimension=-1),
@@ -95,7 +118,7 @@ class DequantizedWeightRecoveryTest(parameterized.TestCase):
     with self.assertRaisesRegex(
         ValueError, "quantized_dimension must be 0, 1, or None. Got"
     ):
-      dequantized_weight_recovery.get_zp_scale_from_2d_dequantized_symmetric_weights(
+      dequantized_weight_recovery.get_zp_scale_from_dequantized_symmetric_weights(
           dequant_vals, quantized_dimension
       )
 

ai_edge_quantizer/algorithms/uniform_quantize/hadamard_rotation.py

@@ -0,0 +1,414 @@
+# 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
+
+
+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.
+
+  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 rotate.
+
+  Returns:
+    A tuple containing the quantized array and the recovered array.
+
+  Raises:
+    ValueError: If the axis is not the last axis of tensor_content. To support
+      other axes, please add support to the matrix multiplication.
+  """
+  if axis != tensor_content.ndim - 1:
+    raise ValueError(
+        "Hadamard rotation is only supported for tensors with quantized"
+        " dimension 0 (rotate last dimension)."
+    )
+
+  # 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
+  # Output size is the product of all dimensions except the one being rotated.
+  output_size = np.prod(np.delete(tensor_content.shape, axis))
+  random_vector = np.ones(hadamard_size, dtype=np.int8)
+
+  # Use a canonical Hadamard matrix.
+  hadamard = _make_hadamard_matrix(hadamard_size)
+  reshaped_tensor = tensor_content.reshape(
+      diagonal_size * output_size, hadamard_size)
+  w_rotated = np.matmul(hadamard, reshaped_tensor.mT).mT
+  return w_rotated.reshape(tensor_content.shape), 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 tensors with rank >= 2."
+    )
+
+  # Reduction axis is the last non-quantized dimension. Since we only support
+  # quantized_dim of 0 (or 1 for blockwise), the reduction axis is the last
+  # axis.
+  reduce_axis = tensor_content.ndim - 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,
+    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.
+
+  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.
+  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_DECOMPOSED_HADAMARD_ROTATION
+      if is_decomposed
+      else 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_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:
+    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_DECOMPOSED_HADAMARD_ROTATION
+      if is_decomposed
+      else 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
+
+
+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,
+  )