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

ai_edge_quantizer/transformations/insert_decomposed_hadamard_rotation.py

@@ -0,0 +1,291 @@
+# 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 decomposed 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 _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 (
+        transformation_utils.get_schema_op_id(transformation, consumer)
+        == schema_py_generated.BuiltinOperator.FULLY_CONNECTED
+    ):
+      transformation.subgraph.operators[consumer].inputs[0] = new_tensor_id
+      updated = True
+  return updated
+
+
+def _make_hadamard_matrix(size: int):
+  """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 insert_decomposed_hadamard_rotation(
+    transformation_input: transformation_utils.TransformationInput,
+) -> qtyping.TransformationInfo:
+  """Inserts a decomposed pattern of Hadamard rotation on this tensor.
+
+  This function works for float32 tensors only. Instead of inserting a single
+  custom op (aeq.hadamard_rotation), this inserts the mathematical equivalent
+  expressed in built-in TFLite ops. The mathematical equivalent is:
+    x' = reshape(x, (-1, hadamard_size))
+    x' = x' @ H(hadamard_size)
+    x' = reshape(x, x.shape)
+  where H(n) is a Hadamard matrix of size n.
+
+  Args:
+    transformation_input: The transformation input to insert the ops on.
+
+  Returns:
+    The transformation info of the inserted ops.
+
+  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.'
+    )
+
+  # Insert x' = tfl.reshape to reshape x to (-1, hadamard_size)
+  hadamard_size = transformation_input.quant_params.hadamard.hadamard_size
+  tensor_size = np.prod(tensor.shape)
+  num_hadamard_blocks = tensor_size // hadamard_size
+  prerotate_shape = [num_hadamard_blocks, hadamard_size]
+  prerotate_shape_tensor_id = transformation_utils.add_new_constant_tensor(
+      tensor.name + b'_prerotate_shape',
+      np.array(prerotate_shape, dtype=np.int32),
+      schema_py_generated.TensorType.INT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+  prerotate_reshape_output_tensor_id = (
+      transformation_utils.add_new_activation_tensor(
+          tensor.name + b'_prerotate_reshaped',
+          prerotate_shape,
+          schema_py_generated.TensorType.FLOAT32,
+          transformation_input.subgraph,
+      )
+  )
+
+  prerotate_reshape_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.RESHAPE,
+      transformation_input.op_codes,
+      'RESHAPE',
+  )
+  prerorate_reshape_op = schema_py_generated.OperatorT()
+  prerorate_reshape_op.opcodeIndex = prerotate_reshape_op_code_idx
+  prerorate_reshape_op.inputs = [
+      transformation_input.tensor_id,
+      prerotate_shape_tensor_id,
+  ]
+  prerorate_reshape_op.outputs = [prerotate_reshape_output_tensor_id]
+
+  # Generate hadamard_matrix(hadamard_size).
+  # We could quantize this to INT4 for better memory efficiency, but for large
+  # models the memory overhead is not significant, and floating point
+  # computation does seem to result in better accuracy.
+  hadamard_matrix = _make_hadamard_matrix(hadamard_size)
+  hadamard_matrix_tensor_id = transformation_utils.add_new_constant_tensor(
+      tensor.name + b'_hadamard_matrix',
+      hadamard_matrix.astype(np.float32),
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+
+  # Insert x' = tfl.fully_connected(x', hadamard_matrix)
+  fc_output_tensor_id = transformation_utils.add_new_activation_tensor(
+      tensor.name + b'_rotated',
+      prerotate_shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+
+  fc_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+      transformation_input.op_codes,
+      'FULLY_CONNECTED',
+  )
+  fc_op = schema_py_generated.OperatorT()
+  fc_op.opcodeIndex = fc_op_code_idx
+  fc_op.inputs = [prerotate_reshape_output_tensor_id, hadamard_matrix_tensor_id]
+  fc_op.outputs = [fc_output_tensor_id]
+
+  # Insert x' = tfl.reshape(x', x.shape)
+  post_reshape_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.RESHAPE,
+      transformation_input.op_codes,
+      'RESHAPE',
+  )
+  post_reshape_op = schema_py_generated.OperatorT()
+  post_reshape_op.opcodeIndex = post_reshape_op_code_idx
+  post_reshape_shape_tensor_id = transformation_utils.add_new_constant_tensor(
+      tensor.name + b'_postrotate_shape',
+      np.array(tensor.shape, dtype=np.int32),
+      schema_py_generated.TensorType.INT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+
+  post_reshape_output_tensor_id = (
+      transformation_utils.add_new_activation_tensor(
+          tensor.name + b'_postrotate_reshaped',
+          tensor.shape,
+          schema_py_generated.TensorType.FLOAT32,
+          transformation_input.subgraph,
+      )
+  )
+  post_reshape_op.inputs = [
+      fc_output_tensor_id,
+      post_reshape_shape_tensor_id,
+  ]
+  post_reshape_op.outputs = [post_reshape_output_tensor_id]
+
+  # Update the users of this tensor to use the new tensor.
+  if (
+      transformation_utils.get_producer_schema_op_id(transformation_input)
+      == schema_py_generated.BuiltinOperator.EMBEDDING_LOOKUP
+  ):
+    _update_embedding_lookup_consumers(
+        transformation_input, post_reshape_output_tensor_id
+    )
+  elif not _update_fully_connected_consumers(
+      transformation_input, post_reshape_output_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] = post_reshape_output_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 new ops in the correct order.
+  transformation_input.subgraph.operators.insert(op_id, prerorate_reshape_op)
+  transformation_input.subgraph.operators.insert(op_id + 1, fc_op)
+  transformation_input.subgraph.operators.insert(op_id + 2, post_reshape_op)
+
+  return qtyping.TransformationInfo(
+      op_id=op_id,
+      num_ops_added=3,
+      output_tensor_id=post_reshape_output_tensor_id,
+  )

ai_edge_quantizer/transformations/insert_decomposed_hadamard_rotation_test.py

@@ -0,0 +1,244 @@
+# 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 Decomposed Hadamard rotation ops."""
+
+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_decomposed_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 InsertDecomposedHadamardRotationFullyConnectedTest(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_decomposed_hadamard_rotation.insert_decomposed_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_decomposed_hadamard_rotation.insert_decomposed_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_decomposed_hadamard_rotation.insert_decomposed_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_decomposed_ops(self):
+    # Insert Decomposed Hadamard ops before fully_connected
+    info = (
+        insert_decomposed_hadamard_rotation.insert_decomposed_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=[0],  # Consumer is the FC op
+                quant_params=self.params,
+            )
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 0)
+    self.assertEqual(info.num_ops_added, 3)
+    # Model had 4 tensors, added 6 tensors (3 activations 3 constants).
+    self.assertEqual(info.output_tensor_id, 9)
+    self.assertLen(subgraph.tensors, 10)
+    # Model had 1 op code, added RESHAPE and FC.
+    self.assertLen(self.model.operatorCodes, 3)
+    self.assertEqual(
+        self.model.operatorCodes[1].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(
+        self.model.operatorCodes[2].builtinCode,
+        schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+    )
+
+    # Op 0: RESHAPE
+    reshape_op = subgraph.operators[0]
+    self.assertEqual(
+        self.model.operatorCodes[reshape_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(reshape_op.inputs[0], 0)  # Graph input
+    self.assertEqual(reshape_op.outputs[0], 5)  # Reshape output
+
+    # Op 1: FULLY_CONNECTED
+    fc_op = subgraph.operators[1]
+    self.assertEqual(
+        self.model.operatorCodes[fc_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+    )
+    self.assertEqual(fc_op.inputs[0], 5)  # Reshape output
+    self.assertEqual(fc_op.inputs[1], 6)  # Hadamard matrix tensor
+    self.assertEqual(fc_op.outputs[0], 7)  # FC output
+
+    # Op 2: RESHAPE (post)
+    post_reshape_op = subgraph.operators[2]
+    self.assertEqual(
+        self.model.operatorCodes[post_reshape_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(post_reshape_op.inputs[0], 7)  # FC output
+    self.assertEqual(post_reshape_op.outputs[0], 9)  # Post Reshape output
+
+    # Op 3: Original FULLY_CONNECTED
+    orig_fc_op = subgraph.operators[3]
+    self.assertEqual(
+        self.model.operatorCodes[orig_fc_op.opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+    )
+    # Input to the original FC is the post reshape output
+    self.assertEqual(orig_fc_op.inputs[0], 9)
+
+
+class InsertDecomposedHadamardRotationEmbeddingLookupTest(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_decomposed_ops(self):
+    # Insert Decomposed Hadamard ops after embedding_lookup
+    info = (
+        insert_decomposed_hadamard_rotation.insert_decomposed_hadamard_rotation(
+            transformation_utils.TransformationInput(
+                tensor_id=2,  # Output of embedding_lookup
+                op_codes=self.model.operatorCodes,
+                buffers=self.model.buffers,
+                subgraph=self.model.subgraphs[0],
+                producer=0,
+                consumers=[-1],  # Output is a graph output
+                quant_params=self.params,
+            )
+        )
+    )
+    subgraph = self.model.subgraphs[0]
+    self.assertEqual(info.op_id, 1)
+    self.assertEqual(info.num_ops_added, 3)
+    # Model had 3 tensors, added 6 (3 activations 3 constants).
+    self.assertEqual(info.output_tensor_id, 8)
+    self.assertLen(subgraph.tensors, 9)
+    # Model had 1 op code, added RESHAPE and FC.
+    self.assertLen(self.model.operatorCodes, 3)
+
+    # Op 0: EMBEDDING_LOOKUP (Original)
+    # Op 1: RESHAPE
+    reshape_op = subgraph.operators[1]
+    self.assertEqual(reshape_op.inputs[0], 2)  # Embedding lookup output
+    self.assertEqual(reshape_op.outputs[0], 4)
+
+    # Op 2: FULLY_CONNECTED
+    fc_op = subgraph.operators[2]
+    self.assertEqual(fc_op.inputs[0], 4)
+    self.assertEqual(fc_op.inputs[1], 5)  # Hadamard matrix
+    self.assertEqual(fc_op.outputs[0], 6)
+
+    # Op 3: RESHAPE (post)
+    post_reshape_op = subgraph.operators[3]
+    self.assertEqual(post_reshape_op.inputs[0], 6)
+    self.assertEqual(post_reshape_op.outputs[0], 8)
+
+    # Check graph output
+    self.assertIn(8, subgraph.outputs)
+    self.assertNotIn(2, subgraph.outputs)
+
+
+if __name__ == '__main__':
+  googletest.main()

ai_edge_quantizer/transformations/insert_hadamard_rotation.py

@@ -34,35 +34,6 @@ def _to_flexbuffer(
   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,
@@ -107,7 +78,10 @@ def _update_fully_connected_consumers(
   """
   updated = False
   for consumer in transformation.consumers:
-    if _is_fully_connected(transformation, consumer):
+    if (
+        transformation_utils.get_schema_op_id(transformation, consumer)
+        == schema_py_generated.BuiltinOperator.FULLY_CONNECTED
+    ):
       transformation.subgraph.operators[consumer].inputs[0] = new_tensor_id
       updated = True
   return updated
@@ -177,7 +151,10 @@ def insert_hadamard_rotation(
   custom_op.outputs = [new_tensor_id]
 
   # Update the users of this tensor to use the new tensor.
-  if _is_producer_embedding_lookup(transformation_input):
+  if (
+      transformation_utils.get_producer_schema_op_id(transformation_input)
+      == schema_py_generated.BuiltinOperator.EMBEDDING_LOOKUP
+  ):
     _update_embedding_lookup_consumers(transformation_input, new_tensor_id)
   elif not _update_fully_connected_consumers(
       transformation_input, new_tensor_id

ai_edge_quantizer/transformations/quantize_tensor.py

@@ -68,29 +68,6 @@ def nonlinear_quant_params_to_tflite_type(
     raise ValueError(f"Unsupported nonlinear params: {bitwidth}")
 
 
-def _pack_data(bitwidth: int, flattened_data: np.ndarray) -> np.ndarray:
-  """Pack the data to the corresponding bit width.
-
-  Currently only support 4 bits. If no packing is needed, the original data is
-  returned.
-
-  Args:
-    bitwidth: Bit width from NonLinearQuantParams.
-    flattened_data: The data to be packed.
-
-  Returns:
-    Packed data.
-  """
-  if bitwidth == 4:
-    even_data = flattened_data[::2] & 0x0F
-    odd_data = np.left_shift(flattened_data[1::2], 4).astype(np.uint8)
-    if odd_data.shape[0] == even_data.shape[0] - 1:
-      odd_data = np.pad(odd_data, (0, 1), constant_values=0)
-    return np.bitwise_or(even_data, odd_data)
-  else:
-    return flattened_data
-
-
 def _perform_channelwise_quantization(
     transformation_input: transformation_utils.TransformationInput,
 ) -> schema_py_generated.QuantizationParametersT():
@@ -180,14 +157,17 @@ def quantize_tensor(
   # is not provided.
   if tensor.buffer:
     if transformation_input.quant_params.quantized_data is not None:
-      transformation_input.buffers[tensor.buffer].data = _pack_data(
-          transformation_input.quant_params.num_bits,
-          np.frombuffer(
-              cast(
-                  np.ndarray, transformation_input.quant_params.quantized_data
-              ).tobytes(),
-              dtype=np.uint8,
-          ).flatten(),
+      transformation_input.buffers[tensor.buffer].data = (
+          transformation_utils.pack_data(
+              transformation_input.quant_params.num_bits,
+              np.frombuffer(
+                  cast(
+                      np.ndarray,
+                      transformation_input.quant_params.quantized_data,
+                  ).tobytes(),
+                  dtype=np.uint8,
+              ).flatten(),
+          )
       )
 
   if isinstance(transformation_input.quant_params, qtyping.UniformQuantParams):