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

ai_edge_quantizer/params_generator_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Tests for params_generator."""
-
 from collections.abc import Generator
 import os
 from typing import Any
@@ -38,8 +36,11 @@ _TensorQuantConfig = qtyping.TensorQuantizationConfig
 _QuantTransformation = qtyping.QuantTransformation
 _AlgorithmName = recipe_manager.AlgorithmName
 _QuantGranularity = qtyping.QuantGranularity
+_QTransf = qtyping.QuantTransformation
+
 
 TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile('')
+_PARAMS_8BIT = qtyping.UniformQuantParams(8, None, np.array([1]), np.array([0]))
 
 
 def _single_fc_model_representative_dataset_gen(num_samples=5):
@@ -64,6 +65,20 @@ def _get_calibration_data(
   return calibration_data
 
 
+def _get_test_consumers(
+    transformations_per_consumer: list[list[_QTransf]],
+    params_per_consumer: list[qtyping.OpToTensorParams],
+) -> list[qtyping.OpToTensorParams]:
+  return [
+      qtyping.OpToTensorParams(
+          subgraph_op_id=i + 1,
+          transformations=transformations_per_consumer[i],
+          parameters=params_per_consumer[i],
+      )
+      for i in range(len(transformations_per_consumer))
+  ]
+
+
 class ParamsGeneratorTest(parameterized.TestCase):
 
   def setUp(self):
@@ -570,9 +585,27 @@ class ParamsGeneratorTest(parameterized.TestCase):
     )
     self.assertLen(quant_params, 6)
 
-  @parameterized.parameters('no_quant', 'execution_mode', 'num_bits')
-  def test_generate_params_buffer_sharing_graphs_fails(
-      self, the_other_fc_difference
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='different_quant_config_fc2_no_quant',
+          fc_2_num_bits=None,
+          expected_tensor_with_buffer_duplication='BatchMatMulV3',
+      ),
+      dict(
+          testcase_name='different_quant_config_fc2_int4',
+          fc_2_num_bits=4,
+          expected_tensor_with_buffer_duplication='BatchMatMulV3',
+      ),
+      dict(
+          testcase_name='same_quant_config',
+          fc_2_num_bits=8,
+          expected_tensor_with_buffer_duplication=None,
+      ),
+  )
+  def test_generate_params_marks_correct_buffers_for_duplication_when_distinct_tensors_share_constant_buffer(
+      self,
+      fc_2_num_bits,
+      expected_tensor_with_buffer_duplication,
   ):
     model_path = os.path.join(
         TEST_DATA_PREFIX_PATH, 'tests/models/weight_sharing_fcs.tflite'
@@ -580,33 +613,204 @@ class ParamsGeneratorTest(parameterized.TestCase):
     # Setup the quantization config for the first FC.
     self._recipe_manager.add_quantization_config(
         regex='PartitionedCall:0',
-        operation_name=qtyping.TFLOperationName.ALL_SUPPORTED,
+        operation_name=qtyping.TFLOperationName.FULLY_CONNECTED,
         op_config=qtyping.OpQuantizationConfig(
-            weight_tensor_config=_TensorQuantConfig(num_bits=8),
+            weight_tensor_config=_TensorQuantConfig(
+                num_bits=8, granularity=qtyping.QuantGranularity.CHANNELWISE
+            ),
             compute_precision=_ComputePrecision.INTEGER,
         ),
     )
     # Setup the quantization config for the second FC (weight shared with the
     # first FC).
-    if the_other_fc_difference == 'no_quant':
-      pass
-    elif the_other_fc_difference == 'num_bits':
+    if fc_2_num_bits is not None:
       self._recipe_manager.add_quantization_config(
           regex='PartitionedCall_1:0',
-          operation_name=qtyping.TFLOperationName.ALL_SUPPORTED,
+          operation_name=qtyping.TFLOperationName.FULLY_CONNECTED,
           op_config=qtyping.OpQuantizationConfig(
-              weight_tensor_config=_TensorQuantConfig(num_bits=4),
+              weight_tensor_config=_TensorQuantConfig(
+                  num_bits=fc_2_num_bits,
+                  granularity=qtyping.QuantGranularity.CHANNELWISE,
+              ),
               compute_precision=_ComputePrecision.INTEGER,
           ),
       )
     pg = params_generator.ParamsGenerator(model_path)
-    error_message = 'do not have the same quantization parameters'
-    with self.assertRaisesWithPredicateMatch(
-        RuntimeError, lambda err: error_message in str(err)
-    ):
-      pg.generate_quantization_parameters(
-          self._recipe_manager,
+    quant_params = pg.generate_quantization_parameters(
+        self._recipe_manager,
+    )
+    self.assertLen(quant_params, 6)
+
+    # Check that the expected tensor has buffer duplication transformation as
+    # the first one to be applied. And no other tensor has buffer duplication
+    # transformation at all.
+    for tensor_name in quant_params:
+      if tensor_name == expected_tensor_with_buffer_duplication:
+        self.assertIsNotNone(quant_params[tensor_name].consumers)
+        for consumer in quant_params[tensor_name].consumers:
+          self.assertNotEmpty(consumer.transformations)
+          self.assertEqual(
+              consumer.transformations[0],
+              _QTransf.DUPLICATE_BUFFER,
+          )
+          self.assertNotIn(
+              _QTransf.DUPLICATE_BUFFER, consumer.transformations[1:]
+          )
+      elif quant_params[tensor_name].consumers is not None:
+        for consumer in quant_params[tensor_name].consumers:
+          self.assertNotIn(_QTransf.DUPLICATE_BUFFER, consumer.transformations)
+
+  def _get_fc_recipe_entry(self, regex: str, num_bits: int):
+    return {
+        'regex': regex,
+        'operation': 'FULLY_CONNECTED',
+        'algorithm_key': 'min_max_uniform_quantize',
+        'op_config': {
+            'weight_tensor_config': {
+                'num_bits': num_bits,
+                'symmetric': True,
+                'granularity': 'CHANNELWISE',
+                'dtype': 'INT',
+                'block_size': 0,
+            },
+            'compute_precision': 'INTEGER',
+            'explicit_dequantize': False,
+            'skip_checks': False,
+            'min_weight_elements': 0,
+        },
+    }
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='fc1_quant_fc2_no_quant',
+          fc1_num_bits=8,
+          fc2_num_bits=None,
+      ),
+      dict(
+          testcase_name='fc1_no_quant_fc2_quant',
+          fc1_num_bits=None,
+          fc2_num_bits=8,
+      ),
+      dict(
+          testcase_name='fc1_quant_fc2_quant_different_params',
+          fc1_num_bits=8,
+          fc2_num_bits=4,
+      ),
+  )
+  def test_generate_params_marks_correct_buffers_tensors_for_duplication(
+      self,
+      fc1_num_bits,
+      fc2_num_bits,
+  ):
+    model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH,
+        'tests/models/constant_tensor_and_buffer_only_sharing_weight_fcs.tflite',
+    )
+    sig1_fc1_regex = 'BatchMatMulV3;'
+    sig1_fc2_regex = 'PartitionedCall:0;'
+    recipe = []
+    if fc1_num_bits is not None:
+      recipe.append(self._get_fc_recipe_entry(sig1_fc1_regex, fc1_num_bits))
+    if fc2_num_bits is not None:
+      recipe.append(self._get_fc_recipe_entry(sig1_fc2_regex, fc2_num_bits))
+    self._recipe_manager.load_quantization_recipe(recipe)
+    pg = params_generator.ParamsGenerator(model_path)
+    quant_params = pg.generate_quantization_parameters(self._recipe_manager)
+
+    expected_tensor = 'arith.constant'
+    consumers = quant_params[expected_tensor].consumers
+    self.assertLen(consumers, 2)
+
+    # Check FC1 transformations.
+    if fc1_num_bits is None:
+      fc1_quant_transformation = _QTransf.NO_QUANTIZE
+    else:
+      fc1_quant_transformation = _QTransf.QUANTIZE_TENSOR
+    self.assertEqual(
+        consumers[0].transformations,
+        [
+            _QTransf.DUPLICATE_TENSOR,
+            _QTransf.DUPLICATE_BUFFER,
+            fc1_quant_transformation,
+        ],
+    )
+    # Check FC2 transformations.
+    if fc2_num_bits is None:
+      fc2_quant_transformation = _QTransf.NO_QUANTIZE
+    else:
+      fc2_quant_transformation = _QTransf.QUANTIZE_TENSOR
+    self.assertEqual(
+        consumers[1].transformations,
+        [
+            _QTransf.DUPLICATE_TENSOR,
+            _QTransf.DUPLICATE_BUFFER,
+            fc2_quant_transformation,
+        ],
+    )
+    # Check that no other tensor has tensor or buffer duplication
+    # transformations.
+    for tensor_name, params in quant_params.items():
+      if tensor_name == expected_tensor:
+        continue
+      for consumer in params.consumers:
+        self.assertNotIn(_QTransf.DUPLICATE_TENSOR, consumer.transformations)
+        self.assertNotIn(_QTransf.DUPLICATE_BUFFER, consumer.transformations)
+
+  def test_generate_params_returns_valid_results_when_multiple_tensor_duplication_for_one_buffer(
+      self,
+  ):
+    model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH,
+        'tests/models/constant_tensor_and_buffer_only_sharing_weight_fcs.tflite',
+    )
+    sig1_fc1_regex = 'BatchMatMulV3;'
+    sig1_fc2_regex = 'PartitionedCall:0;'
+    sig2_fc1_regex = 'BatchMatMulV31;'
+    sig2_fc2_regex = 'PartitionedCall_1:0;'
+    recipe = [
+        self._get_fc_recipe_entry(sig1_fc1_regex, num_bits=8),
+        self._get_fc_recipe_entry(sig1_fc2_regex, num_bits=4),
+        self._get_fc_recipe_entry(sig2_fc1_regex, num_bits=8),
+        self._get_fc_recipe_entry(sig2_fc2_regex, num_bits=4),
+    ]
+    self._recipe_manager.load_quantization_recipe(recipe)
+    pg = params_generator.ParamsGenerator(model_path)
+    quant_params = pg.generate_quantization_parameters(self._recipe_manager)
+    # Check transformations for sig1.
+    sig1_expected_tensor = 'arith.constant'
+    sig1_consumers = quant_params[sig1_expected_tensor].consumers
+    self.assertLen(sig1_consumers, 2)
+    sig1_expected_transformations = [
+        _QTransf.DUPLICATE_TENSOR,
+        _QTransf.DUPLICATE_BUFFER,
+        _QTransf.QUANTIZE_TENSOR,
+    ]
+    for sig1_consumer in sig1_consumers:
+      self.assertEqual(
+          sig1_consumer.transformations,
+          sig1_expected_transformations,
+      )
+    # Check transformations for sig2.
+    sig2_expected_tensor = 'arith.constant1'
+    sig2_consumers = quant_params[sig2_expected_tensor].consumers
+    self.assertLen(sig2_consumers, 2)
+    sig2_expected_transformations = [
+        _QTransf.DUPLICATE_TENSOR,
+        _QTransf.QUANTIZE_TENSOR,
+    ]
+    for sig2_consumer in sig2_consumers:
+      self.assertEqual(
+          sig2_consumer.transformations,
+          sig2_expected_transformations,
       )
+    # Check that no other tensor has tensor or buffer duplication
+    # transformations.
+    for tensor_name, params in quant_params.items():
+      if tensor_name in [sig1_expected_tensor, sig2_expected_tensor]:
+        continue
+      for consumer in params.consumers:
+        self.assertNotIn(_QTransf.DUPLICATE_TENSOR, consumer.transformations)
+        self.assertNotIn(_QTransf.DUPLICATE_BUFFER, consumer.transformations)
 
   @parameterized.named_parameters(
       dict(
@@ -615,279 +819,185 @@ class ParamsGeneratorTest(parameterized.TestCase):
               tensor_name='tfl.quantize',
               producer=qtyping.OpToTensorParams(
                   subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.ADD_DEQUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+                  transformations=[_QTransf.ADD_DEQUANTIZE],
+                  parameters=_PARAMS_8BIT,
+              ),
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.NO_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.NO_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
           param2=qtyping.TensorTransformationParams(
-              'tfl.other_quantize',
-              qtyping.OpToTensorParams(
+              tensor_name='tfl.other_quantize',
+              producer=qtyping.OpToTensorParams(
                   subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.NO_QUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+                  transformations=[_QTransf.NO_QUANTIZE],
+                  parameters=_PARAMS_8BIT,
+              ),
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              [
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
           expected=False,
       ),
       dict(
-          testcase_name='param2_consumer_incompatible',
+          testcase_name='compatible',
           param1=qtyping.TensorTransformationParams(
               tensor_name='tfl.quantize',
-              producer=qtyping.OpToTensorParams(
-                  subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.ADD_QUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE, _QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
           param2=qtyping.TensorTransformationParams(
-              'tfl.other_quantize',
-              qtyping.OpToTensorParams(
-                  subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.NO_QUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+              tensor_name='tfl.other_quantize',
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              [
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.QUANTIZE_TENSOR,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
-          expected=False,
+          expected=True,
       ),
       dict(
-          testcase_name='compatible',
+          testcase_name='compatible_no_numeric_check',
           param1=qtyping.TensorTransformationParams(
               tensor_name='tfl.quantize',
               producer=None,
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00028806]), np.array([0])
                       ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.NO_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00027501]), np.array([0])
                       ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.NO_QUANTIZE,
-                      ],
-                  ),
-              ],
+                  ],
+              ),
           ),
           param2=qtyping.TensorTransformationParams(
-              'tfl.other_quantize',
-              None,
-              [
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+              tensor_name='tfl.quantize',
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00028806]), np.array([0])
                       ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00027501]), np.array([0])
                       ),
-                  ),
-              ],
+                  ],
+              ),
           ),
           expected=True,
       ),
   )
-  def test_params_compatible(self, param1, param2, expected):
-    # adding a test to make production coverage happy.
+  def test__are_self_compatible_tensors_compatible_to_each_other(
+      self, param1, param2, expected
+  ):
     self.assertEqual(
-        params_generator._compatible_tensor_transformation_params(
+        params_generator._are_self_compatible_tensors_compatible_to_each_other(
             param1, param2
         ),
         expected,
     )
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='consumer_incompatible',
+          params=qtyping.TensorTransformationParams(
+              tensor_name='tfl.quantize',
+              producer=qtyping.OpToTensorParams(
+                  subgraph_op_id=0,
+                  transformations=[_QTransf.NO_QUANTIZE],
+                  parameters=_PARAMS_8BIT,
+              ),
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.QUANTIZE_TENSOR],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
+              ),
+          ),
+          expected=False,
+      ),
+      dict(
+          testcase_name='compatible',
+          params=qtyping.TensorTransformationParams(
+              tensor_name='tfl.quantize',
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE, _QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
+              ),
+          ),
+          expected=True,
+      ),
+      dict(
+          testcase_name='compatible_no_numeric_check',
+          params=qtyping.TensorTransformationParams(
+              tensor_name='tfl.quantize',
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00028806]), np.array([0])
+                      ),
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00027501]), np.array([0])
+                      ),
+                  ],
+              ),
+          ),
+          expected=True,
+      ),
+  )
+  def test__are_tensor_consumer_params_compatible(self, params, expected):
+    self.assertEqual(
+        params_generator._are_tensor_consumer_params_compatible(params),
+        expected,
+    )
+
   def test_model_with_duplicated_tensor_names_fails(self):
     model_path = os.path.join(
         TEST_DATA_PREFIX_PATH, 'tests/models/duplicated_tensor_names.tflite'
@@ -1025,16 +1135,11 @@ class ParamsGeneratorAlreadyQuantizedModelTest(googletest.TestCase):
     )
     _ = params_generator.ParamsGenerator(test_model_path)
 
-  def test_check_is_float_model_raises_error_when_model_is_quantized(self):
+  def test_check_is_quantized_model_succeeds_when_model_is_quantized(self):
     test_model_path = os.path.join(
         TEST_DATA_PREFIX_PATH, 'tests/models/mnist_quantized.tflite'
     )
-    with self.assertRaisesRegex(
-        ValueError,
-        'The input model for quantization parameters generation is not a float'
-        ' model.',
-    ):
-      _ = params_generator.ParamsGenerator(test_model_path)
+    _ = params_generator.ParamsGenerator(test_model_path)
 
 
 if __name__ == '__main__':