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

ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py

@@ -0,0 +1,664 @@
+# 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.
+# ==============================================================================
+
+import os
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.nonlinear_quantize import float_casting
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+
+_TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile("../../tests/models")
+_TFLOpName = qtyping.TFLOperationName
+_ComputePrecision = qtyping.ComputePrecision
+_TensorQuantConfig = qtyping.TensorQuantizationConfig
+_QuantTransformation = qtyping.QuantTransformation
+
+
+class Fp16QuantizeTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(666)
+    self._test_model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, "conv_fc_mnist.tflite"
+    )
+    self._test_model = tfl_flatbuffer_utils.read_model(self._test_model_path)
+    # The test model has one subgraph for now.
+    self._graph_info = qtyping.GraphInfo(
+        subgraph_tensors=self._test_model.subgraphs[0].tensors,
+        buffers=self._test_model.buffers,
+    )
+    self._tensor_name_to_qsv = {}
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="fc",
+          op_name=_TFLOpName.FULLY_CONNECTED,
+      ),
+      dict(
+          testcase_name="conv2d",
+          op_name=_TFLOpName.CONV_2D,
+      ),
+      dict(
+          testcase_name="embedding_lookup",
+          op_name=_TFLOpName.EMBEDDING_LOOKUP,
+      ),
+  )
+  def test_check_op_quantization_config_succeeds(self, op_name):
+    float_casting.check_op_quantization_config(
+        op_name,
+        qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+            ),
+            compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+            explicit_dequantize=True,
+        ),
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="invalid_fc",
+          op_name=_TFLOpName.FULLY_CONNECTED,
+      ),
+      dict(
+          testcase_name="invalid_conv2d",
+          op_name=_TFLOpName.CONV_2D,
+      ),
+      dict(
+          testcase_name="invalid_embedding_lookup",
+          op_name=_TFLOpName.EMBEDDING_LOOKUP,
+      ),
+  )
+  def test_check_op_quantization_config_invalid_activation_tensor_config_raises_exception(
+      self, op_name
+  ):
+    # With activation tensor config.
+    error_message = (
+        "Activation tensor quantization is not supported for float casting"
+        " quantization."
+    )
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      float_casting.check_op_quantization_config(
+          op_name,
+          qtyping.OpQuantizationConfig(
+              activation_tensor_config=_TensorQuantConfig(
+                  num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+              ),
+              weight_tensor_config=_TensorQuantConfig(
+                  num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+              ),
+              compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+              explicit_dequantize=True,
+          ),
+      )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="invalid_fc",
+          op_name=_TFLOpName.FULLY_CONNECTED,
+      ),
+      dict(
+          testcase_name="invalid_conv2d",
+          op_name=_TFLOpName.CONV_2D,
+      ),
+      dict(
+          testcase_name="invalid_embedding_lookup",
+          op_name=_TFLOpName.EMBEDDING_LOOKUP,
+      ),
+  )
+  def test_check_op_quantization_config_invalid_bit_width_raises_exception(
+      self, op_name
+  ):
+    error_message = (
+        "float casting quantization config requires number of bits to be set as"
+        " 16, dtype as float"
+    )
+    # Wrong bit width.
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      float_casting.check_op_quantization_config(
+          op_name,
+          qtyping.OpQuantizationConfig(
+              activation_tensor_config=None,
+              weight_tensor_config=_TensorQuantConfig(
+                  num_bits=8, dtype=qtyping.TensorDataType.FLOAT
+              ),
+              compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+              explicit_dequantize=True,
+          ),
+      )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="invalid_fc",
+          op_name=_TFLOpName.FULLY_CONNECTED,
+      ),
+      dict(
+          testcase_name="invalid_conv2d",
+          op_name=_TFLOpName.CONV_2D,
+      ),
+      dict(
+          testcase_name="invalid_embedding_lookup",
+          op_name=_TFLOpName.EMBEDDING_LOOKUP,
+      ),
+  )
+  def test_check_op_quantization_config_invalid_dtype_raises_exception(
+      self, op_name
+  ):
+    error_message = (
+        "float casting quantization config requires number of bits to be set as"
+        " 16, dtype as float"
+    )
+    # Wrong dtype.
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      float_casting.check_op_quantization_config(
+          op_name,
+          qtyping.OpQuantizationConfig(
+              activation_tensor_config=None,
+              weight_tensor_config=_TensorQuantConfig(
+                  num_bits=16, dtype=qtyping.TensorDataType.INT
+              ),
+              compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+              explicit_dequantize=True,
+          ),
+      )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="invalid_fc",
+          op_name=_TFLOpName.FULLY_CONNECTED,
+      ),
+      dict(
+          testcase_name="invalid_conv2d",
+          op_name=_TFLOpName.CONV_2D,
+      ),
+      dict(
+          testcase_name="invalid_embedding_lookup",
+          op_name=_TFLOpName.EMBEDDING_LOOKUP,
+      ),
+  )
+  def test_check_op_quantization_config_no_weight_config_raises_exception(
+      self, op_name
+  ):
+    error_message = (
+        "Weight tensor quantization config is required for float casting"
+        " quantization."
+    )
+    # No weight quantization config.
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      float_casting.check_op_quantization_config(
+          op_name,
+          qtyping.OpQuantizationConfig(
+              activation_tensor_config=None,
+              compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+              explicit_dequantize=True,
+          ),
+      )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="averagepool2D",
+          op_name=_TFLOpName.AVERAGE_POOL_2D,
+      ),
+      dict(
+          testcase_name="reshape",
+          op_name=_TFLOpName.RESHAPE,
+      ),
+  )
+  def test_check_op_quantization_config_invalid_ops_raises_exception(
+      self, op_name
+  ):
+    error_message = "Unsupported op"
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      float_casting.check_op_quantization_config(
+          op_name=op_name,
+          op_quant_config=qtyping.OpQuantizationConfig(
+              activation_tensor_config=None,
+              weight_tensor_config=_TensorQuantConfig(
+                  num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+              ),
+              compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+              explicit_dequantize=True,
+          ),
+      )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="fc_with_bias",
+          subgraph_op_id=3,
+          op_tensor_names={
+              "weight": "arith.constant1",
+              "bias": "arith.constant2",
+              "input": "sequential/flatten/Reshape",
+              "output": "sequential/dense/MatMul;sequential/dense/Relu;sequential/dense/BiasAdd",
+          },
+      ),
+      dict(
+          testcase_name="fc_with_no_bias",
+          subgraph_op_id=4,
+          op_tensor_names={
+              "weight": "arith.constant",
+              "input": "sequential/dense/MatMul;sequential/dense/Relu;sequential/dense/BiasAdd",
+              "output": "sequential/dense_1/MatMul",
+          },
+      ),
+  )
+  def test_fully_connected_weight_only_succeeds(
+      self, subgraph_op_id, op_tensor_names
+  ):
+    subgraph0 = self._test_model.subgraphs[0]
+    fc_op = subgraph0.operators[subgraph_op_id]
+    op_info = qtyping.OpInfo(
+        op=fc_op,
+        op_name=_TFLOpName.FULLY_CONNECTED,
+        subgraph_op_index=subgraph_op_id,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+            ),
+            compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+            explicit_dequantize=True,
+        ),
+    )
+
+    self._test_fc_conv(
+        op_info,
+        self._graph_info,
+        op_tensor_names,
+        float_casting.materialize_fc_conv,
+    )
+
+  def test_conv2d_weight_only_succeeds(self):
+    # Read from Model Explorer.
+    subgraph0 = self._test_model.subgraphs[0]
+    subgraph_op_id = 0
+    op = subgraph0.operators[subgraph_op_id]
+
+    op_info = qtyping.OpInfo(
+        op=op,
+        op_name=_TFLOpName.CONV_2D,
+        subgraph_op_index=subgraph_op_id,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+            ),
+            compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+            explicit_dequantize=True,
+        ),
+    )
+
+    op_tensor_names = {}
+    op_tensor_names["weight"] = "sequential/conv2d/Conv2D"
+    op_tensor_names["bias"] = (
+        "sequential/conv2d/Relu;sequential/conv2d/BiasAdd;sequential/conv2d/Conv2D;sequential/conv2d/BiasAdd/ReadVariableOp"
+    )
+    op_tensor_names["input"] = "serving_default_conv2d_input:0"
+    op_tensor_names["output"] = (
+        "sequential/conv2d/Relu;sequential/conv2d/BiasAdd;sequential/conv2d/Conv2D;sequential/conv2d/BiasAdd/ReadVariableOp1"
+    )
+    self._test_fc_conv(
+        op_info,
+        self._graph_info,
+        op_tensor_names,
+        float_casting.materialize_fc_conv,
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="invalid_fc",
+          op_name=_TFLOpName.FULLY_CONNECTED,
+      ),
+      dict(
+          testcase_name="invalid_conv2d",
+          op_name=_TFLOpName.CONV_2D,
+      ),
+      dict(
+          testcase_name="invalid_embedding_lookup",
+          op_name=_TFLOpName.EMBEDDING_LOOKUP,
+      ),
+  )
+  def test_check_op_quantization_config_invalid_execution_mode_raises_exception(
+      self, op_name
+  ):
+    # Use DRQ instead of WEIGHT-ONLY.
+    error_message = (
+        "Currently, only Weight-Only is supported for float casting"
+        " quantization."
+    )
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      float_casting.check_op_quantization_config(
+          op_name,
+          qtyping.OpQuantizationConfig(
+              activation_tensor_config=None,
+              weight_tensor_config=_TensorQuantConfig(
+                  num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+              ),
+              compute_precision=_ComputePrecision.INTEGER,  # DRQ.
+          ),
+      )
+
+  def test_conv2d_transpose_weight_only_succeeds(self):
+    # Read from Model Explorer.
+    test_model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, "single_conv2d_transpose_bias.tflite"
+    )
+
+    test_model = tfl_flatbuffer_utils.read_model(test_model_path)
+    # The test model has one subgraph for now.
+    graph_info = qtyping.GraphInfo(
+        subgraph_tensors=test_model.subgraphs[0].tensors,
+        buffers=test_model.buffers,
+    )
+
+    subgraph0 = test_model.subgraphs[0]
+    subgraph_op_id = 0
+    op = subgraph0.operators[subgraph_op_id]
+
+    op_info = qtyping.OpInfo(
+        op=op,
+        op_name=_TFLOpName.CONV_2D_TRANSPOSE,
+        subgraph_op_index=subgraph_op_id,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+            ),
+            compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+            explicit_dequantize=True,
+        ),
+    )
+
+    op_tensor_names = {}
+    op_tensor_names["weight"] = (
+        "sequential_5/conv2d_transpose_3/conv2d_transpose"
+    )
+    op_tensor_names["bias"] = (
+        "sequential_5/conv2d_transpose_3/BiasAdd;sequential_5/conv2d_transpose_3/conv2d_transpose;sequential_5/conv2d_transpose_3/BiasAdd/ReadVariableOp"
+    )
+    op_tensor_names["input"] = "serving_default_input_6:0"
+    op_tensor_names["output"] = "StatefulPartitionedCall:0"
+
+    tensor_quant_params = float_casting.materialize_conv2d_transpose(
+        op_info, graph_info, self._tensor_name_to_qsv
+    )
+    _, weight_tensor, bias_tensor, _ = (
+        tfl_flatbuffer_utils.parse_fc_bmm_conv_tensors(
+            op_info.op, graph_info.subgraph_tensors
+        )
+    )
+
+    num_configs = 4 if bias_tensor is not None else 3
+    self.assertLen(tensor_quant_params, num_configs)
+
+    # Test input tensor params.
+    self._test_fp16_nonweight_tensor_transformation_params(
+        op_tensor_names["input"],
+        op_info.subgraph_op_index,
+        transformation_params=tensor_quant_params[0],
+        desired_transformations=[_QuantTransformation.NO_QUANTIZE],
+        is_inbounding_tensor=True,
+    )
+
+    # Test weight tensor params.
+    weight_tensor_data = tfl_flatbuffer_utils.get_tensor_data(
+        weight_tensor,
+        graph_info.buffers,
+    )
+    self._test_fp16_weight_tensor_transformation_params(
+        op_tensor_names["weight"],
+        op_info.subgraph_op_index,
+        tensor_quant_config=op_info.op_quant_config.weight_tensor_config,
+        transformation_params=tensor_quant_params[1],
+        desired_transformations=[_QuantTransformation.ADD_DEQUANTIZE],
+        tensor_data=weight_tensor_data,
+    )
+    # Test output tensor params.
+    self._test_fp16_nonweight_tensor_transformation_params(
+        op_tensor_names["output"],
+        op_info.subgraph_op_index,
+        transformation_params=tensor_quant_params[2],
+        desired_transformations=[_QuantTransformation.NO_QUANTIZE],
+        is_inbounding_tensor=False,
+    )
+
+    # Test bias tensor params.
+    if bias_tensor is not None:
+      self._test_fp16_nonweight_tensor_transformation_params(
+          op_tensor_names["bias"],
+          op_info.subgraph_op_index,
+          transformation_params=tensor_quant_params[3],
+          desired_transformations=[_QuantTransformation.NO_QUANTIZE],
+          is_inbounding_tensor=True,
+      )
+
+  def test_depthwise_conv2d_weight_only_succeeds(self):
+    # Read from Model Explorer.
+    test_model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, "single_depthwise_conv2d_bias.tflite"
+    )
+
+    test_model = tfl_flatbuffer_utils.read_model(test_model_path)
+    # The test model has one subgraph for now.
+    graph_info = qtyping.GraphInfo(
+        subgraph_tensors=test_model.subgraphs[0].tensors,
+        buffers=test_model.buffers,
+    )
+
+    subgraph0 = test_model.subgraphs[0]
+    subgraph_op_id = 0
+    op = subgraph0.operators[subgraph_op_id]
+
+    op_info = qtyping.OpInfo(
+        op=op,
+        op_name=_TFLOpName.DEPTHWISE_CONV_2D,
+        subgraph_op_index=subgraph_op_id,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+            ),
+            compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+            explicit_dequantize=True,
+        ),
+    )
+
+    op_tensor_names = {}
+    op_tensor_names["weight"] = "sequential/depthwise_conv2d/depthwise"
+    op_tensor_names["bias"] = (
+        "sequential/depthwise_conv2d/BiasAdd;sequential/depthwise_conv2d/depthwise;sequential/depthwise_conv2d/BiasAdd/ReadVariableOp"
+    )
+    op_tensor_names["input"] = "serving_default_input_1:0"
+    op_tensor_names["output"] = "StatefulPartitionedCall:0"
+    self._test_fc_conv(
+        op_info,
+        graph_info,
+        op_tensor_names,
+        float_casting.materialize_fc_conv,
+    )
+
+  def test_embedding_lookup_weight_only_succeeds(self):
+    test_model_path = os.path.join(
+        _TEST_DATA_PREFIX_PATH, "embedding_lookup.tflite"
+    )
+
+    test_model = tfl_flatbuffer_utils.read_model(test_model_path)
+    graph_info = qtyping.GraphInfo(
+        subgraph_tensors=test_model.subgraphs[0].tensors,
+        buffers=test_model.buffers,
+    )
+
+    subgraph0 = test_model.subgraphs[0]
+    subgraph_op_id = 0
+    op = subgraph0.operators[subgraph_op_id]
+
+    op_info = qtyping.OpInfo(
+        op=op,
+        op_name=_TFLOpName.EMBEDDING_LOOKUP,
+        subgraph_op_index=subgraph_op_id,
+        op_quant_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=16, dtype=qtyping.TensorDataType.FLOAT
+            ),
+            compute_precision=_ComputePrecision.FLOAT,  # WEIGHT_ONLY.
+            explicit_dequantize=True,
+        ),
+    )
+
+    op_tensor_names = {}
+    op_tensor_names["weight"] = (
+        "jax2tf_export_func_/...y_yz-_...z/pjit__einsum_/MatMul;jax2tf_export_func_/pjit__one_hot_/Equal;jax2tf_export_func_/pjit__one_hot_/Cast_1"
+    )
+    op_tensor_names["input"] = "inputs"
+    op_tensor_names["output"] = "Identity_1"
+
+    # TODO: b/335913710 - Rename the test function.
+    self._test_fc_conv(
+        op_info,
+        graph_info,
+        op_tensor_names,
+        float_casting.materialize_fc_conv,
+    )
+
+  def _test_fc_conv(
+      self,
+      op_info,
+      graph_info,
+      op_tensor_names,
+      materialization_func,
+  ):
+
+    tensor_quant_params = materialization_func(
+        op_info, graph_info, self._tensor_name_to_qsv
+    )
+    _, weight_tensor, bias_tensor, _ = (
+        tfl_flatbuffer_utils.parse_fc_bmm_conv_tensors(
+            op_info.op, graph_info.subgraph_tensors
+        )
+    )
+
+    num_configs = 4 if bias_tensor is not None else 3
+    self.assertLen(tensor_quant_params, num_configs)
+
+    # Test input tensor params.
+    self._test_fp16_nonweight_tensor_transformation_params(
+        op_tensor_names["input"],
+        op_info.subgraph_op_index,
+        transformation_params=tensor_quant_params[0],
+        desired_transformations=[_QuantTransformation.NO_QUANTIZE],
+        is_inbounding_tensor=True,
+    )
+
+    # Test weight tensor params.
+    weight_tensor_data = tfl_flatbuffer_utils.get_tensor_data(
+        weight_tensor,
+        graph_info.buffers,
+    )
+
+    self._test_fp16_weight_tensor_transformation_params(
+        op_tensor_names["weight"],
+        op_info.subgraph_op_index,
+        tensor_quant_config=op_info.op_quant_config.weight_tensor_config,
+        transformation_params=tensor_quant_params[1],
+        desired_transformations=[_QuantTransformation.ADD_DEQUANTIZE],
+        tensor_data=weight_tensor_data,
+    )
+    # Test output tensor params.
+    self._test_fp16_nonweight_tensor_transformation_params(
+        op_tensor_names["output"],
+        op_info.subgraph_op_index,
+        transformation_params=tensor_quant_params[2],
+        desired_transformations=[_QuantTransformation.NO_QUANTIZE],
+        is_inbounding_tensor=False,
+    )
+
+    # Test bias tensor params.
+    if bias_tensor is not None:
+      self._test_fp16_nonweight_tensor_transformation_params(
+          op_tensor_names["bias"],
+          op_info.subgraph_op_index,
+          transformation_params=tensor_quant_params[3],
+          desired_transformations=[_QuantTransformation.NO_QUANTIZE],
+          is_inbounding_tensor=True,
+      )
+
+  def _test_fp16_weight_tensor_transformation_params(
+      self,
+      tensor_name,
+      subgraph_op_id,
+      tensor_quant_config,
+      transformation_params,
+      desired_transformations,
+      tensor_data,
+  ):
+    self.assertEqual(transformation_params.tensor_name, tensor_name)
+    # Weight-only means the transformation is added from the consumer.
+    self.assertIsNone(transformation_params.producer)
+    self.assertLen(transformation_params.consumers, 1)
+    # Check op params.
+    op_params = transformation_params.consumers[0]
+    self.assertEqual(op_params.subgraph_op_id, subgraph_op_id)
+    self.assertSequenceEqual(op_params.transformations, desired_transformations)
+    # Check quantization params.
+    quantization_params = op_params.parameters
+    self.assertIsNotNone(quantization_params)
+    self.assertIsNotNone(tensor_quant_config)
+    self.assertEqual(quantization_params.num_bits, tensor_quant_config.num_bits)
+    quantized_data = quantization_params.quantized_data
+    self.assertIsNotNone(quantized_data)
+    self.assertEqual(quantized_data.dtype, "float16")
+    # fp16 quantization implies very small error.
+    self.assertSequenceAlmostEqual(
+        list(tensor_data.flatten()),  # pytype: disable=attribute-error
+        list(quantization_params.quantized_data.flatten()),  # pytype: disable=attribute-error
+        delta=5,
+    )
+
+  def _test_fp16_nonweight_tensor_transformation_params(
+      self,
+      tensor_name,
+      subgraph_op_id,
+      transformation_params,
+      desired_transformations,
+      is_inbounding_tensor,
+  ):
+    self.assertEqual(transformation_params.tensor_name, tensor_name)
+    if is_inbounding_tensor:
+      self.assertIsNone(transformation_params.producer)
+      self.assertLen(transformation_params.consumers, 1)
+      op_params = transformation_params.consumers[0]
+    else:
+      self.assertIsNone(transformation_params.consumers)
+      op_params = transformation_params.producer
+    self.assertIsNotNone(op_params)
+    self.assertEqual(op_params.subgraph_op_id, subgraph_op_id)
+    self.assertSequenceEqual(op_params.transformations, desired_transformations)
+    self.assertIsNone(op_params.parameters)
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/algorithms/uniform_quantize/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+# ==============================================================================
+