ai-edge-quantizer-nightly 0.0.1.dev20250115__py3-none-any.whl

ai_edge_quantizer/transformations/emulated_subchannel.py

@@ -0,0 +1,363 @@
+# 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.
+# ==============================================================================
+
+"""Transformation pattern for emulated subchannel quantization."""
+
+from typing import cast
+import numpy as np
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import quantize_tensor
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+
+def emulated_subchannel(
+    transformation_input: transformation_utils.TransformationInput,
+) -> qtyping.TransformationInfo:
+  """Emulated subchannel quantization for fully_connected op.
+
+  The input tensor must also be the weight tensor of the fully_connected op.
+
+  after the transformation, the fully connected op will be replaced by:
+  reshape -> batch_matmul -> mul -> sum -> add (if bias is present) ->
+  activation (if fused activation function exist, only support ReLU for now)
+
+  Args:
+    transformation_input: input structure that contains all information needed
+      for the transformation.
+
+  Returns:
+    The transformation info.
+  """
+  # only apply to a single fully_connected op
+  if len(transformation_input.consumers) > 1:
+    raise ValueError('Emulated Subchannel transformation only support one op')
+  if isinstance(
+      transformation_input.quant_params, qtyping.NonLinearQuantParams
+  ):
+    raise ValueError(
+        'Emulated Subchannel transformation only support uniform quantization'
+    )
+  if (
+      transformation_input.op_codes[
+          transformation_input.subgraph.operators[
+              transformation_input.consumers[0]
+          ].opcodeIndex
+      ].builtinCode
+      != schema_py_generated.BuiltinOperator.FULLY_CONNECTED
+  ):
+    raise ValueError(
+        'Emulated Subchannel transformation only support fully_connected op'
+    )
+  if transformation_input.producer != -1:
+    raise ValueError(
+        'Emulated Subchannel transformation only support constant tensor'
+    )
+
+  # insert all tne necessary op codes into the model
+  reshape_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.RESHAPE, transformation_input.op_codes
+  )
+  bmm_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.BATCH_MATMUL,
+      transformation_input.op_codes,
+  )
+  mul_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.MUL, transformation_input.op_codes
+  )
+  sum_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.SUM, transformation_input.op_codes
+  )
+
+  original_fc_op_idx = transformation_input.consumers[0]
+  if cast(
+      schema_py_generated.FullyConnectedOptionsT,
+      transformation_input.subgraph.operators[
+          original_fc_op_idx
+      ].builtinOptions,
+  ).fusedActivationFunction not in (
+      schema_py_generated.ActivationFunctionType.NONE,
+      schema_py_generated.ActivationFunctionType.RELU,
+  ):
+    raise ValueError(
+        'Emulated Subchannel transformation only support'
+        ' fusedActivationFunction NONE and RELU for now'
+    )
+
+  weight_tensor = transformation_input.subgraph.tensors[
+      transformation_input.tensor_id
+  ]
+  weight_tensor.type = quantize_tensor.quant_params_to_tflite_type(
+      transformation_input.quant_params.num_bits
+  )
+
+  # modify the weight tensor with the correct quantization parameters
+  transformation_input.buffers[weight_tensor.buffer].data = np.frombuffer(
+      cast(
+          np.ndarray, transformation_input.quant_params.quantized_data
+      ).tobytes(),
+      dtype=np.uint8,
+  )
+  weight_tensor.shape = cast(
+      np.ndarray, transformation_input.quant_params.quantized_data
+  ).shape
+  weight_tensor.quantization.scale = np.ones(shape=[1], dtype=np.float32)
+  weight_tensor.quantization.zeroPoint = np.zeros(
+      shape=[1], dtype=np.int64
+  ).flatten()
+
+  # assuming zero point is 0, so no need to add a zero point tensor
+  for val in transformation_input.quant_params.zero_point.flatten():
+    if val != 0:
+      raise ValueError(
+          'Emulated Subchannel transformation only support zero point 0 for now'
+      )
+
+  scale_tensor_id = transformation_utils.add_new_constant_tensor(
+      weight_tensor.name + b'_scale',
+      transformation_input.quant_params.scale,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+
+  # for fully connected op, the reduce axis is always 1
+  reduce_axes_data = np.array([1], dtype=np.int32)
+  reduce_axes_tensor_id = transformation_utils.add_new_constant_tensor(
+      weight_tensor.name + b'_reduce_axes',
+      reduce_axes_data,
+      schema_py_generated.TensorType.INT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+
+  # find the input and output tensor of the fully connected op
+  activation_input_id = transformation_input.subgraph.operators[
+      transformation_input.consumers[0]
+  ].inputs[0]
+  activation_output_id = transformation_input.subgraph.operators[
+      transformation_input.consumers[0]
+  ].outputs[0]
+  activation_input = transformation_input.subgraph.tensors[activation_input_id]
+  activation_output = transformation_input.subgraph.tensors[
+      activation_output_id
+  ]
+
+  if len(activation_input.shape) != 3:
+    raise ValueError(
+        'Emulated Subchannel transformation only support 3D input tensor'
+    )
+  bmm_input_shape = [
+      activation_input.shape[0] * activation_input.shape[1],
+      weight_tensor.shape[1],
+      1,
+      weight_tensor.shape[2],
+  ]
+  intermediate_tensor_shape = [
+      activation_input.shape[0] * activation_input.shape[1],
+      weight_tensor.shape[1],
+      1,
+      weight_tensor.shape[3],
+  ]
+  sum_output_shape = [
+      activation_input.shape[0] * activation_input.shape[1],
+      1,
+      1,
+      weight_tensor.shape[3],
+  ]
+
+  # create constant tensors for reshape
+  reshape1_shape_id = transformation_utils.add_new_constant_tensor(
+      activation_output.name + b'_reshape_op1_shape',
+      np.array(bmm_input_shape, dtype=np.int32),
+      schema_py_generated.TensorType.INT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+  reshape2_shape_id = transformation_utils.add_new_constant_tensor(
+      activation_output.name + b'_reshape_op2_shape',
+      np.array(activation_output.shape, dtype=np.int32),
+      schema_py_generated.TensorType.INT32,
+      transformation_input.subgraph,
+      transformation_input.buffers,
+  )
+
+  # create all intermediate tensors
+  bmm_input_id = transformation_utils.add_new_activation_tensor(
+      activation_output.name + b'_bmm_input',
+      bmm_input_shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+  mul_input_id = transformation_utils.add_new_activation_tensor(
+      activation_output.name + b'_mul_input',
+      intermediate_tensor_shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+  sum_input_id = transformation_utils.add_new_activation_tensor(
+      activation_output.name + b'_reduce_sum_input',
+      intermediate_tensor_shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+  reshape_op2_input_id = transformation_utils.add_new_activation_tensor(
+      activation_output.name + b'_reshape_op2_input',
+      sum_output_shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+
+  # reshape
+  reshape_op1 = schema_py_generated.OperatorT()
+  reshape_op1.opcodeIndex = reshape_op_code_idx
+  reshape_op1_option = schema_py_generated.ReshapeOptionsT()
+  reshape_op1_option.newShape = bmm_input_shape
+  reshape_op1.inputs = [activation_input_id, reshape1_shape_id]
+  reshape_op1.outputs = [bmm_input_id]
+  reshape_op1.builtinOptionsType = (
+      schema_py_generated.BuiltinOptions.ReshapeOptions
+  )  # reshape option index
+  reshape_op1.builtinOptions = reshape_op1_option
+
+  # batch_matmul
+  bmm_op = schema_py_generated.OperatorT()
+  bmm_op.opcodeIndex = bmm_op_code_idx
+  bmm_op.inputs = [bmm_input_id, transformation_input.tensor_id]
+  bmm_op.outputs = [mul_input_id]
+  bmm_op.builtinOptionsType = (
+      schema_py_generated.BuiltinOptions.BatchMatMulOptions
+  )
+  bmm_op.builtinOptions = schema_py_generated.BatchMatMulOptionsT()
+
+  # mul
+  mul_op = schema_py_generated.OperatorT()
+  mul_op.opcodeIndex = mul_op_code_idx
+  mul_option = schema_py_generated.MulOptionsT()
+  mul_option.fusedActivationFunction = (
+      schema_py_generated.ActivationFunctionType.NONE
+  )
+  mul_op.inputs = [mul_input_id, scale_tensor_id]
+  mul_op.outputs = [sum_input_id]
+  mul_op.builtinOptionsType = schema_py_generated.BuiltinOptions.MulOptions
+  mul_op.builtinOptions = mul_option
+
+  # sum
+  sum_op = schema_py_generated.OperatorT()
+  sum_op.opcodeIndex = sum_op_code_idx
+  sum_op.inputs = [sum_input_id, reduce_axes_tensor_id]
+  sum_op.outputs = [reshape_op2_input_id]
+  sum_op.builtinOptionsType = schema_py_generated.BuiltinOptions.ReducerOptions
+  sum_op.builtinOptions = schema_py_generated.ReducerOptionsT()
+  sum_op.builtinOptions.keepDims = True
+
+  # reshape
+  reshape_op2 = schema_py_generated.OperatorT()
+  reshape_op2.opcodeIndex = reshape_op_code_idx
+  reshape_op2_option = schema_py_generated.ReshapeOptionsT()
+  reshape_op2_option.newShape = activation_output.shape
+  reshape_op2.inputs = [reshape_op2_input_id, reshape2_shape_id]
+  reshape_op2.outputs = [activation_output_id]
+  reshape_op2.builtinOptionsType = (
+      schema_py_generated.BuiltinOptions.ReshapeOptions
+  )
+  reshape_op2.builtinOptions = reshape_op2_option
+
+  transformation_input.subgraph.operators.insert(
+      original_fc_op_idx, reshape_op1
+  )
+  transformation_input.subgraph.operators.insert(original_fc_op_idx + 1, bmm_op)
+  transformation_input.subgraph.operators.insert(original_fc_op_idx + 2, mul_op)
+  transformation_input.subgraph.operators.insert(original_fc_op_idx + 3, sum_op)
+  transformation_input.subgraph.operators.insert(
+      original_fc_op_idx + 4, reshape_op2
+  )
+  ops_added = 5
+  last_op = reshape_op2
+
+  # If there is a bias tensor (the third input to the original fc op),
+  # we need an add to process it. The current fc op id need to be recalculated
+  # because we added operators in front of it.
+  current_fc_op_id = original_fc_op_idx + ops_added
+  if (
+      len(transformation_input.subgraph.operators[current_fc_op_id].inputs) > 2
+      and transformation_input.subgraph.operators[current_fc_op_id].inputs[2]
+      != -1
+  ):
+    add_op_code_idx = transformation_utils.add_op_code(
+        schema_py_generated.BuiltinOperator.ADD, transformation_input.op_codes
+    )
+    reshape_op2_output_id = transformation_utils.add_new_activation_tensor(
+        activation_output.name + b'_reshape_op2_output',
+        activation_output.shape,
+        schema_py_generated.TensorType.FLOAT32,
+        transformation_input.subgraph,
+    )
+    last_op.outputs = [reshape_op2_output_id]
+    add_op = schema_py_generated.OperatorT()
+    add_op.opcodeIndex = add_op_code_idx
+    add_option = schema_py_generated.AddOptionsT()
+    add_op.builtinOptionsType = schema_py_generated.BuiltinOptions.AddOptions
+    add_op.builtinOptions = add_option
+    add_op.inputs = [
+        reshape_op2_output_id,
+        transformation_input.subgraph.operators[
+            original_fc_op_idx + ops_added
+        ].inputs[2],
+    ]
+    add_op.outputs = [activation_output_id]
+    transformation_input.subgraph.operators.insert(
+        original_fc_op_idx + ops_added, add_op
+    )
+    ops_added += 1
+    last_op = add_op
+
+  # If the fused activation function is RELU, we need to add a relu op.
+  # The current fc op id need to be recalculated because we added operators
+  # in front of it.
+  fc_fused_activation_function = cast(
+      schema_py_generated.FullyConnectedOptionsT,
+      transformation_input.subgraph.operators[
+          original_fc_op_idx + ops_added
+      ].builtinOptions,
+  ).fusedActivationFunction
+  if (
+      fc_fused_activation_function
+      == schema_py_generated.ActivationFunctionType.RELU
+  ):
+    activation_output.name += b'_relu'
+    relu_input_id = transformation_utils.add_new_activation_tensor(
+        activation_output.name + b'_relu_input',
+        activation_output.shape,
+        schema_py_generated.TensorType.FLOAT32,
+        transformation_input.subgraph,
+    )
+    last_op.outputs = [relu_input_id]
+    relu_op = schema_py_generated.OperatorT()
+    relu_op.opcodeIndex = transformation_utils.add_op_code(
+        schema_py_generated.BuiltinOperator.RELU, transformation_input.op_codes
+    )
+    relu_op.inputs = [relu_input_id]
+    relu_op.outputs = [activation_output_id]
+    transformation_input.subgraph.operators.insert(
+        original_fc_op_idx + ops_added, relu_op
+    )
+    ops_added += 1
+    last_op = relu_op
+
+  del transformation_input.subgraph.operators[original_fc_op_idx + ops_added]
+  return qtyping.TransformationInfo(
+      original_fc_op_idx, ops_added - 1, activation_output_id
+  )

ai_edge_quantizer/transformations/emulated_subchannel_test.py

@@ -0,0 +1,212 @@
+# 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.
+# ==============================================================================
+
+"""Tests for emulated_subchannel."""
+
+import os
+import numpy as np
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import emulated_subchannel
+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 EmulatedSubchannelTest(googletest.TestCase):
+  """Tests for emulated_subchannel."""
+
+  def setUp(self):
+    super().setUp()
+    self.params = qtyping.UniformQuantParams(
+        num_bits=8,
+        quantized_dimension=None,
+        scale=np.ones([1, 1, 1, 4], dtype=np.float32),
+        zero_point=np.zeros([1, 1, 1, 4], dtype=np.int64),
+        symmetric=True,
+        quantized_data=np.ones([1, 4, 2, 4], dtype=np.int8),
+    )
+
+  def test_emulate_subchannel_without_bias_succeeds(self):
+    """Tests the emulated_subchannel function."""
+    self._model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_no_bias.tflite"
+    )
+    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
+    subgraph = self._model.subgraphs[0]
+    model = self._model
+    ret = emulated_subchannel.emulated_subchannel(
+        transformation_utils.TransformationInput(
+            tensor_id=1,
+            op_codes=model.operatorCodes,
+            buffers=model.buffers,
+            subgraph=subgraph,
+            producer=-1,
+            consumers=[0],
+            quant_params=self.params,
+        )
+    )
+    self.assertEqual(ret.op_id, 0)
+    self.assertEqual(ret.num_ops_added, 4)
+    self.assertEqual(ret.output_tensor_id, 2)
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[0].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[1].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.BATCH_MATMUL,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[2].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.MUL,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[3].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.SUM,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[4].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(
+        subgraph.tensors[subgraph.operators[2].inputs[1]].name,
+        b"arith.constant_scale",
+    )
+    self.assertListEqual(
+        np.frombuffer(
+            model.buffers[
+                subgraph.tensors[subgraph.operators[2].inputs[1]].buffer
+            ].data,
+            dtype=np.float32,
+        ).tolist(),
+        np.ones([1, 1, 1, 4]).flatten().tolist(),
+    )
+
+  def test_emulate_subchannel_with_bias_succeeds(self):
+    """Tests the emulated_subchannel function."""
+    self._model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias.tflite"
+    )
+    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
+    subgraph = self._model.subgraphs[0]
+    model = self._model
+    ret = emulated_subchannel.emulated_subchannel(
+        transformation_utils.TransformationInput(
+            tensor_id=1,
+            op_codes=model.operatorCodes,
+            buffers=model.buffers,
+            subgraph=subgraph,
+            producer=-1,
+            consumers=[0],
+            quant_params=self.params,
+        )
+    )
+    self.assertEqual(ret.op_id, 0)
+    self.assertEqual(ret.num_ops_added, 5)
+    self.assertEqual(ret.output_tensor_id, 3)
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[0].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[1].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.BATCH_MATMUL,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[2].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.MUL,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[3].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.SUM,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[4].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RESHAPE,
+    )
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[5].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.ADD,
+    )
+    self.assertEqual(
+        subgraph.tensors[subgraph.operators[2].inputs[1]].name,
+        b"arith.constant_scale",
+    )
+    self.assertListEqual(
+        np.frombuffer(
+            model.buffers[
+                subgraph.tensors[subgraph.operators[2].inputs[1]].buffer
+            ].data,
+            dtype=np.float32,
+        ).tolist(),
+        np.ones([1, 1, 1, 4]).flatten().tolist(),
+    )
+
+  def test_emulated_subchannel_with_fused_relu_succeeds(self):
+    """Tests the emulated_subchannel function with fused relu."""
+    self._model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias_relu.tflite"
+    )
+    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
+    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
+    subgraph = self._model.subgraphs[0]
+    model = self._model
+    ret = emulated_subchannel.emulated_subchannel(
+        transformation_utils.TransformationInput(
+            tensor_id=1,
+            op_codes=model.operatorCodes,
+            buffers=model.buffers,
+            subgraph=subgraph,
+            producer=-1,
+            consumers=[0],
+            quant_params=self.params,
+        )
+    )
+    self.assertEqual(ret.op_id, 0)
+    self.assertEqual(ret.num_ops_added, 6)
+    self.assertEqual(ret.output_tensor_id, 3)
+    self.assertEqual(
+        model.operatorCodes[subgraph.operators[6].opcodeIndex].builtinCode,
+        schema_py_generated.BuiltinOperator.RELU,
+    )
+
+  def test_emulated_subchannel_raises_when_unsupported_activation(self):
+    """Tests the emulated_subchannel function with unsupported activation."""
+    self._model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias_relu6.tflite"
+    )
+    self._model = tfl_flatbuffer_utils.read_model(self._model_path)
+    subgraph = self._model.subgraphs[0]
+    model = self._model
+    with self.assertRaises(ValueError):
+      emulated_subchannel.emulated_subchannel(
+          transformation_utils.TransformationInput(
+              tensor_id=1,
+              op_codes=model.operatorCodes,
+              buffers=model.buffers,
+              subgraph=subgraph,
+              producer=-1,
+              consumers=[0],
+              quant_params=self.params,
+          )
+      )
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/transformations/quant_insert.py

@@ -0,0 +1,100 @@
+# 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.
+# ==============================================================================
+
+"""Apply dequantization transformations to the given op/tensor.
+
+Inserts dequantize node after the given tensor to enable float execution of
+the tensor consumer
+"""
+
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import quantize_tensor
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+
+def insert_quant(
+    transformation_input: transformation_utils.TransformationInput,
+) -> qtyping.TransformationInfo:
+  """Insert quant op after the given tensor in the subgraph.
+
+  Args:
+    transformation_input: input structure that contains all information needed
+      for the transformation.
+
+  Returns:
+    TransformationInfo:
+      op_id: the index where the dequant op is added
+      num_ops_added: the total number of ops inserted by this operation, which
+        is 1
+  """
+  quant_op_code_idx = transformation_utils.add_op_code(
+      schema_py_generated.BuiltinOperator.QUANTIZE,
+      transformation_input.op_codes,
+  )
+
+  # create output tensor for the quantize op
+  tensor = transformation_input.subgraph.tensors[transformation_input.tensor_id]
+  new_tensor_id = transformation_utils.add_new_activation_tensor(
+      tensor.name + b'_quantized',
+      tensor.shape,
+      schema_py_generated.TensorType.FLOAT32,
+      transformation_input.subgraph,
+  )
+
+  # quantize the output tensor
+  ## we need a new transformation input because we don't want to modify the
+  ## original input
+  quantize_tensor.quantize_tensor(
+      transformation_utils.TransformationInput(
+          new_tensor_id,
+          transformation_input.op_codes,
+          transformation_input.buffers,
+          transformation_input.subgraph,
+          transformation_input.producer,
+          transformation_input.consumers,
+          transformation_input.quant_params,
+      )
+  )
+
+  # create quantize_op
+  quant_op = schema_py_generated.OperatorT()
+  quant_op.opcodeIndex = quant_op_code_idx
+  quant_op.outputs = [new_tensor_id]
+  quant_op.inputs = [transformation_input.tensor_id]
+
+  # update the original consumers of the op to take the dequant op,
+  # and find the first consumer of the new tensor
+  first_consumer_id = min(transformation_input.consumers)
+  for consumer_id in transformation_input.consumers:
+    op = transformation_input.subgraph.operators[consumer_id]
+    for input_idx in range(len(op.inputs)):
+      if op.inputs[input_idx] == transformation_input.tensor_id:
+        op.inputs[input_idx] = new_tensor_id
+
+  # if the output is also an output to the graph, we need to update that as well
+  for output_idx, output in enumerate(transformation_input.subgraph.outputs):
+    if output == transformation_input.tensor_id:
+      transformation_input.subgraph.outputs[output_idx] = new_tensor_id
+
+  # add dequant into the subgraph op list,
+  # must insert the op right before it's first consumer
+  # in the case of output goes to graph output, we need to ensure the dequant
+  # op is inserted after the producer
+  op_id = max(transformation_input.producer + 1, first_consumer_id)
+  transformation_input.subgraph.operators.insert(op_id, quant_op)
+  return qtyping.TransformationInfo(
+      op_id=op_id, num_ops_added=1, output_tensor_id=new_tensor_id
+  )