ai-edge-quantizer-nightly 0.4.0.dev20251008__py3-none-any.whl → 0.5.0.dev20251121__py3-none-any.whl

ai_edge_quantizer/calibrator.py

@@ -98,9 +98,7 @@ class Calibrator:
       qsv_update_func: The function to update the QSVs.
     """
     op_codes = self._flatbuffer_model.operatorCodes
-    if not self._model_qsvs:
-      self._initialize_model_qsvs(model_recipe_manager)
-    else:
+    if self._model_qsvs:
       logging.warning(
           "Calibrator contains non-empty model qsvs, and the current"
           " calibration process will start on top of this state (i.e., update"
@@ -263,50 +261,3 @@ class Calibrator:
         output_tensor = subgraph_tensors[output_tensor_idx]
         scope += tfl_flatbuffer_utils.get_tensor_name(output_tensor)
     return scope
-
-  # TODO: b/354224138 - Remove code duplication between calibrate and
-  # _initialize_model_qsvs.
-  def _initialize_model_qsvs(
-      self, model_recipe_manager: recipe_manager.RecipeManager
-  ) -> None:
-    """Initialize the model qsvs.
-
-    Args:
-      model_recipe_manager: A RecipeManager object that contains the
-        quantization recipe.
-    """
-    op_codes = self._flatbuffer_model.operatorCodes
-    for subgraph in self._flatbuffer_model.subgraphs:
-      graph_info = qtyping.GraphInfo(
-          subgraph.tensors, self._flatbuffer_model.buffers
-      )
-      for subgraph_op_id, op in enumerate(subgraph.operators):
-        op_code = op_codes[op.opcodeIndex].builtinCode
-        if op_code not in tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME:
-          continue
-        op_key = tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME[op_code]
-        # Step1: query the quantization_recipe to get op quantization
-        # settings.
-        op_scope = self._get_op_scope(op, subgraph.tensors)
-        algorithm_name, op_quant_config = (
-            model_recipe_manager.get_quantization_configs(op_key, op_scope)
-        )
-        if algorithm_name == algorithm_manager.AlgorithmName.NO_QUANTIZE:
-          continue
-        # Step2: query algorithm_manager to get/call the related qsv init
-        # function.
-        qsv_init_func = algorithm_manager.get_init_qsv_func(
-            algorithm_name, op_key
-        )
-        op_info = qtyping.OpInfo(op, op_key, subgraph_op_id, op_quant_config)
-        # Ignore the input tensors where any dimension of the shape is 0.
-        inputs_to_ignore = [
-            opr_idx
-            for opr_idx, tensor_idx in enumerate(op.inputs)
-            if not np.all(graph_info.subgraph_tensors[tensor_idx].shape)
-        ]
-        op_qsvs = qsv_init_func(op_info, graph_info, inputs_to_ignore)
-        # Step3: initialize tensor qsvs.
-        for tensor_name, qsv in op_qsvs.items():
-          if tensor_name not in self._model_qsvs:
-            self._model_qsvs[tensor_name] = qsv

ai_edge_quantizer/calibrator_test.py

@@ -103,58 +103,6 @@ class CalibratorTest(googletest.TestCase):
     model_tensor_qsvs = self._calibrator.get_model_qsvs()
     self.assertEmpty(model_tensor_qsvs)
 
-  def test_calibrator_initialize_qsv(self):
-    _add_default_int8xint8_integer_recipe(self._recipe_manager)
-    # Overwrite the single op to fc
-    self._recipe_manager.add_quantization_config(
-        regex=".*Stateful.*",
-        operation_name=qtyping.TFLOperationName.FULLY_CONNECTED,
-        algorithm_key=_AlgorithmName.MIN_MAX_UNIFORM_QUANT,
-        op_config=qtyping.OpQuantizationConfig(
-            weight_tensor_config=_TENSOR_QUANT_CONFIG(
-                num_bits=4,
-                granularity=qtyping.QuantGranularity.CHANNELWISE,
-            ),
-            compute_precision=_ComputePrecision.INTEGER,
-        ),
-    )
-    self._calibrator._initialize_model_qsvs(self._recipe_manager)
-    model_tensor_qsvs = self._calibrator.get_model_qsvs()
-
-    self.assertLen(model_tensor_qsvs, 4)
-    self.assertIn("serving_default_input_1:0", model_tensor_qsvs)  # input
-    input_qsv = model_tensor_qsvs["serving_default_input_1:0"]
-    self.assertEmpty(input_qsv)
-
-    self.assertIn("sequential/dense/MatMul", model_tensor_qsvs)  # weight
-    weight_tensor_qsv = model_tensor_qsvs["sequential/dense/MatMul"]
-    mins_maxs_shape = (16, 1)
-    self.assertTupleEqual(weight_tensor_qsv["min"].shape, mins_maxs_shape)
-    self.assertAlmostEqual(weight_tensor_qsv["min"][0][0], -0.40436327)
-    self.assertTupleEqual(weight_tensor_qsv["max"].shape, mins_maxs_shape)
-    self.assertAlmostEqual(weight_tensor_qsv["max"][0][0], 0.46138108)
-
-    self.assertIn(
-        "sequential/dense/BiasAdd/ReadVariableOp", model_tensor_qsvs
-    )  # bias
-    bias_tensor_qsv = model_tensor_qsvs[
-        "sequential/dense/BiasAdd/ReadVariableOp"
-    ]
-    mins_maxs_shape = (16,)
-    self.assertTupleEqual(bias_tensor_qsv["min"].shape, mins_maxs_shape)
-    self.assertAlmostEqual(bias_tensor_qsv["min"][0], -0.26978338)
-    self.assertTupleEqual(bias_tensor_qsv["max"].shape, mins_maxs_shape)
-    # Here bias min/max will be the same as each element is a scalar
-    # Bias will be quantized with input_scale * weight_scale.
-    self.assertSequenceEqual(
-        list(bias_tensor_qsv["max"].flatten()),
-        list(bias_tensor_qsv["min"].flatten()),
-    )
-
-    self.assertIn("StatefulPartitionedCall:0", model_tensor_qsvs)  # output
-    output_qsv = model_tensor_qsvs["StatefulPartitionedCall:0"]
-    self.assertEmpty(output_qsv)
-
   def test_calibrate_single_fc_success(self):
     _add_default_int8xint8_integer_recipe(self._recipe_manager)
     self._calibrator.calibrate(
@@ -162,7 +110,7 @@ class CalibratorTest(googletest.TestCase):
     )
     model_tensor_qsvs = self._calibrator.get_model_qsvs()
 
-    self.assertLen(model_tensor_qsvs, 4)
+    self.assertLen(model_tensor_qsvs, 2)
     self.assertIn("serving_default_input_1:0", model_tensor_qsvs)  # input
     input_qsv = model_tensor_qsvs["serving_default_input_1:0"]
     self.assertSequenceAlmostEqual(
@@ -171,19 +119,6 @@ class CalibratorTest(googletest.TestCase):
     self.assertSequenceAlmostEqual(
         input_qsv["max"].flatten(), [TEST_MAX_VAL], delta=1e-5
     )
-
-    self.assertIn("sequential/dense/MatMul", model_tensor_qsvs)  # weight
-    weight_qsv = model_tensor_qsvs["sequential/dense/MatMul"]
-    self.assertSequenceAlmostEqual(weight_qsv["min"].flatten(), [-0.49114203])
-    self.assertSequenceAlmostEqual(weight_qsv["max"].flatten(), [0.4903704])
-
-    self.assertIn(
-        "sequential/dense/BiasAdd/ReadVariableOp", model_tensor_qsvs
-    )  # bias
-    bias_qsv = model_tensor_qsvs["sequential/dense/BiasAdd/ReadVariableOp"]
-    self.assertSequenceAlmostEqual(bias_qsv["min"].flatten(), [-0.38401994])
-    self.assertSequenceAlmostEqual(bias_qsv["max"].flatten(), [0.31727126])
-
     self.assertIn("StatefulPartitionedCall:0", model_tensor_qsvs)  # output
     output_qsv = model_tensor_qsvs["StatefulPartitionedCall:0"]
     # Relu, only check the min
@@ -302,7 +237,7 @@ class CalibratorToyGemma2Test(googletest.TestCase):
         self._toy_gemma2_calibration_dataset,
         model_recipe_manager=recipe_mngr,
     )
-    self.assertLen(calib.get_model_qsvs(), 290)
+    self.assertLen(calib.get_model_qsvs(), 202)
 
 
 if __name__ == "__main__":

ai_edge_quantizer/default_policy.py

@@ -61,9 +61,8 @@ DEFAULT_JSON_POLICY = """
       "weight_tensor_config": {
         "num_bits": 4,
         "symmetric": [true],
-        "granularity": ["BLOCKWISE"],
-        "dtype": "INT",
-        "block_size": [32, 64, 96, 128, 256]
+        "granularity": ["BLOCKWISE_32", "BLOCKWISE_64", "BLOCKWISE_128", "BLOCKWISE_256"],
+        "dtype": "INT"
       },
       "explicit_dequantize": false,
       "compute_precision": "INTEGER"
@@ -178,7 +177,6 @@ DEFAULT_JSON_POLICY = """
       "INPUT",
       "OUTPUT",
       "SLICE",
-      "EMBEDDING_LOOKUP",
       "SUM",
       "SELECT",
       "SELECT_V2",
@@ -226,7 +224,6 @@ DEFAULT_JSON_POLICY = """
       "INPUT",
       "OUTPUT",
       "SLICE",
-      "EMBEDDING_LOOKUP",
       "SUM",
       "SELECT",
       "SELECT_V2",
@@ -250,10 +247,11 @@ DEFAULT_JSON_POLICY = """
       "REDUCE_MIN",
       "EQUAL",
       "NOT_EQUAL",
-      "MIRROR_PAD"
+      "MIRROR_PAD",
+      "SPACE_TO_DEPTH"
     ],
-    "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],
-    "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],
+    "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT"],
+    "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT"],
     "dynamic_wi8_afp32": [
       "BATCH_MATMUL",
       "CONV_2D",
@@ -321,16 +319,9 @@ def _unroll_json_config(
           "granularity": granularity,
           "dtype": json_config["weight_tensor_config"]["dtype"],
       }
-      if "block_size" in json_config["weight_tensor_config"]:
-        for block_size in json_config["weight_tensor_config"]["block_size"]:
-          tensor_config["block_size"] = block_size
-          weight_configs.append(
-              qtyping.TensorQuantizationConfig.from_dict(tensor_config)
-          )
-      else:
-        weight_configs.append(
-            qtyping.TensorQuantizationConfig.from_dict(tensor_config)
-        )
+      weight_configs.append(
+          qtyping.TensorQuantizationConfig.from_dict(tensor_config)
+      )
 
       if activation_configs:
         for activation_config in activation_configs:

ai_edge_quantizer/qtyping.py

@@ -81,6 +81,7 @@ class TFLOperationName(str, enum.Enum):
   EQUAL = 'EQUAL'
   NOT_EQUAL = 'NOT_EQUAL'
   MIRROR_PAD = 'MIRROR_PAD'
+  SPACE_TO_DEPTH = 'SPACE_TO_DEPTH'
 
 
 class QuantizeMode(enum.Enum):
@@ -111,7 +112,11 @@ class TensorDataType(str, enum.Enum):
 class QuantGranularity(str, enum.Enum):
   TENSORWISE = 'TENSORWISE'
   CHANNELWISE = 'CHANNELWISE'
-  BLOCKWISE = 'BLOCKWISE'
+  # Blockwise quantization with various block sizes.
+  BLOCKWISE_32 = 'BLOCKWISE_32'
+  BLOCKWISE_64 = 'BLOCKWISE_64'
+  BLOCKWISE_128 = 'BLOCKWISE_128'
+  BLOCKWISE_256 = 'BLOCKWISE_256'
 
 
 class QuantTransformation(enum.Enum):
@@ -309,7 +314,6 @@ class TensorQuantizationConfig:
     granularity: Whether to perform per-tensor, per-channel or per-block
       quantization.
     dtype: The data type of the tensor.
-    block_size: The block size for blockwise quantization, ignored otherwise.
     algorithm_key: The algorithm key to use for quantization.
   """
 
@@ -317,7 +321,6 @@ class TensorQuantizationConfig:
   symmetric: bool = True
   granularity: QuantGranularity = QuantGranularity.TENSORWISE
   dtype: TensorDataType = TensorDataType.INT
-  block_size: int = 0
 
   def to_dict(self) -> dict[str, Any]:
     """Converts ActivationQuantizationConfig to dict."""
@@ -335,9 +338,28 @@ class TensorQuantizationConfig:
   def from_dict(cls, params: dict[str, Any]) -> 'TensorQuantizationConfig':
     """Converts a given dict to TensorQuantizationConfig."""
     params_copy = copy.deepcopy(params)
+    # Process block_size config from legacy recipe.
+    params_copy = _process_block_size(params_copy)
     return cls(**params_copy)
 
 
+def _process_block_size(params: dict[str, Any]) -> dict[str, Any]:
+  """Processes block size in the params."""
+  block_size = params.pop('block_size', 0)
+  if block_size > 0:
+    if block_size == 32:
+      params['granularity'] = QuantGranularity.BLOCKWISE_32
+    elif block_size == 64:
+      params['granularity'] = QuantGranularity.BLOCKWISE_64
+    elif block_size == 128:
+      params['granularity'] = QuantGranularity.BLOCKWISE_128
+    elif block_size == 256:
+      params['granularity'] = QuantGranularity.BLOCKWISE_256
+    else:
+      raise ValueError(f'Unsupported block size: {block_size}')
+  return params
+
+
 @dataclasses.dataclass(frozen=True)
 class OpQuantizationConfig:
   """Configuration class to control the quantization process behavior.

ai_edge_quantizer/quantizer.py

@@ -126,12 +126,16 @@ class Quantizer:
     float_model: TFLite model file path or bytearray.
     quantization_recipe: Quantization recipe .json filepath or in loaded json
       format.
+    previous_quantized_model: Optional previously quantized TFLite model file
+      path or bytearray. This is useful for validating a quantized model
+      without quantizing it again.
   """
 
   def __init__(
       self,
       float_model: Union[str, bytearray],
       quantization_recipe: Optional[Union[str, _QuantRecipe]] = None,
+      previous_quantized_model: Optional[Union[str, bytearray]] = None,
   ):
     """Initializes the quantizer.
 
@@ -139,6 +143,9 @@ class Quantizer:
       float_model: Path to the float tflite model.
       quantization_recipe: Quantization recipe in .json filepath or loaded json
         format.
+      previous_quantized_model: Path to an optional previously quantized tflite
+        model. This is useful for validating a quantized model without
+        quantizing it again.
     """
     # Use `float model` as bytes for memory efficiency.
     self.float_model: bytes = (
@@ -146,6 +153,14 @@ class Quantizer:
         if isinstance(float_model, str)
         else float_model
     )
+    if previous_quantized_model is not None:
+      self.previous_quantized_model: bytes = (
+          tfl_flatbuffer_utils.get_model_content(previous_quantized_model)
+          if isinstance(previous_quantized_model, str)
+          else previous_quantized_model
+      )
+    else:
+      self.previous_quantized_model = None
 
     self._recipe_manager: recipe_manager.RecipeManager = (
         recipe_manager.RecipeManager()
@@ -153,6 +168,7 @@ class Quantizer:
     if quantization_recipe is not None:
       self.load_quantization_recipe(quantization_recipe)
     self._result: QuantizationResult = QuantizationResult([{}], None)
+    self._quantize_called = False
 
   def load_quantization_recipe(self, recipe: Union[str, _QuantRecipe]) -> None:
     """Loads a quantization recipe.
@@ -399,7 +415,7 @@ class Quantizer:
     Raises:
       RuntimeError: If quantization recipe is empty.
     """
-
+    self._quantize_called = True
     if calibration_result is not None:
       self._ensure_model_qsv_sufficient(calibration_result)
 
@@ -445,9 +461,16 @@ class Quantizer:
       test_data = tfl_interpreter_utils.create_random_normal_input_data(
           self.float_model, num_samples=1
       )
+    if self._quantize_called:
+      quantized_model = self._result.quantized_model
+    else:
+      quantized_model = self.previous_quantized_model
+
+    if quantized_model is None:
+      raise ValueError('No quantized model available to validate.')
     return model_validator.compare_model(
         self.float_model,
-        self._result.quantized_model,
+        quantized_model,
         test_data,
         error_metrics,
         validation_utils.get_validation_func(error_metrics),

ai_edge_quantizer/quantizer_test.py

@@ -212,7 +212,7 @@ class QuantizerTest(parameterized.TestCase):
     # Calibrate with empty state.
     calib_data = _get_calibration_data()
     calibration_result = self._quantizer.calibrate(calib_data)
-    self.assertLen(calibration_result, 13)
+    self.assertLen(calibration_result, 7)
 
   @parameterized.parameters(
       'recipes/default_a8w8_recipe.json',
@@ -227,7 +227,7 @@ class QuantizerTest(parameterized.TestCase):
     updated_calibration_result = self._quantizer.calibrate(
         calib_data, previous_calibration_result=calibration_result
     )
-    self.assertLen(updated_calibration_result, 13)
+    self.assertLen(updated_calibration_result, 7)
     self.assertNotEqual(
         calibration_result['StatefulPartitionedCall:0'],
         updated_calibration_result['StatefulPartitionedCall:0'],
@@ -309,6 +309,44 @@ class QuantizerTest(parameterized.TestCase):
       saved_recipe = json.load(json_file)
     self.assertEqual(saved_recipe, self._test_recipe)
 
+  def test_saved_legacy_recipe_lacks_block_size(self):
+    model_name = 'test_model'
+    legacy_recipe_path = os.path.join(
+        TEST_DATA_PREFIX_PATH,
+        'recipes/dynamic_legacy_wi8_afp32_recipe.json',
+    )
+    self._quantizer.load_quantization_recipe(legacy_recipe_path)
+    result = self._quantizer.quantize()
+    result.save(self._tmp_save_path, model_name)
+    saved_recipe_path = os.path.join(
+        self._tmp_save_path, model_name + '_recipe.json'
+    )
+    with open(saved_recipe_path) as json_file:
+      saved_recipe = json.load(json_file)
+    with open(legacy_recipe_path) as json_file:
+      legacy_recipe = json.load(json_file)
+
+    self.assertNotEqual(saved_recipe, legacy_recipe)
+
+    # Verify that the default test recipe contains 'block_size'.
+    has_block_size = False
+    for config in legacy_recipe:
+      op_config = config.get('op_config')
+      if op_config:
+        weight_config = op_config.get('weight_tensor_config')
+        if weight_config and 'block_size' in weight_config:
+          has_block_size = True
+          break
+    self.assertTrue(has_block_size)
+
+    # Verify that the saved recipe does not have 'block_size'.
+    for config in saved_recipe:
+      op_config = config.get('op_config')
+      if op_config:
+        weight_config = op_config.get('weight_tensor_config')
+        if weight_config:
+          self.assertNotIn('block_size', weight_config)
+
   def test_save_no_quantize_raise_error(self):
     error_message = 'No quantized model to save.'
     with self.assertRaisesWithPredicateMatch(
@@ -337,6 +375,21 @@ class QuantizerTest(parameterized.TestCase):
         'sequential/dense_1/MatMul', validation_result.intermediate_tensors
     )
 
+  def test_validate_with_quantized_model_arg_succeeds(self):
+    self._quantizer.quantize()
+    quantized_model = self._quantizer._result.quantized_model
+    self.assertIsNotNone(quantized_model)
+
+    new_quantizer = quantizer.Quantizer(
+        self._test_model_path, previous_quantized_model=quantized_model
+    )
+    validation_result = new_quantizer.validate()
+    validation_result = validation_result.get_signature_comparison_result()
+    self.assertIsNotNone(validation_result)
+    self.assertIn(
+        'sequential/dense_1/MatMul', validation_result.intermediate_tensors
+    )
+
   def test_load_custom_policies_succeeds(self):
 
     test_op_config = qtyping.OpQuantizationConfig(
@@ -520,14 +573,12 @@ class QuantizerMultiSignatureModelTest(parameterized.TestCase):
                     'symmetric': False,
                     'granularity': 'TENSORWISE',
                     'dtype': 'INT',
-                    'block_size': 0,
                 },
                 'weight_tensor_config': {
                     'num_bits': 8,
                     'symmetric': True,
                     'granularity': 'CHANNELWISE',
                     'dtype': 'INT',
-                    'block_size': 0,
                 },
                 'compute_precision': 'INTEGER',
                 'explicit_dequantize': False,
@@ -548,8 +599,7 @@ class QuantizerMultiSignatureModelTest(parameterized.TestCase):
 
     # Quantize and expect an error about missing signature in calibration data.
     error_message = (
-        'Missing QSVs (min/max) for tensor multiply_x:0 in Signature'
-        " 'multiply'."
+        'MUL(index: 0) not found in tensor_name_to_qsv'
     )
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: error_message in str(err)

ai_edge_quantizer/recipe_manager_test.py

@@ -569,14 +569,12 @@ class ConfiguratorTest(parameterized.TestCase, googletest.TestCase):
                     'symmetric': False,
                     'granularity': _QuantGranularity.TENSORWISE,
                     'dtype': 'INT',
-                    'block_size': 0,
                 },
                 'weight_tensor_config': {
                     'num_bits': 8,
                     'symmetric': True,
                     'granularity': _QuantGranularity.TENSORWISE,
                     'dtype': 'INT',
-                    'block_size': 0,
                 },
                 # WEIGHT_ONLY.
                 'compute_precision': _ComputePrecision.INTEGER,
@@ -595,7 +593,6 @@ class ConfiguratorTest(parameterized.TestCase, googletest.TestCase):
                     'num_bits': 8,
                     'symmetric': True,
                     'granularity': _QuantGranularity.TENSORWISE,
-                    'block_size': 0,
                 },
                 # WEIGHT_ONLY.
                 'compute_precision': _ComputePrecision.FLOAT,
@@ -614,7 +611,6 @@ class ConfiguratorTest(parameterized.TestCase, googletest.TestCase):
                     'num_bits': 4,
                     'symmetric': True,
                     'granularity': _QuantGranularity.TENSORWISE,
-                    'block_size': 0,
                 },
                 # WEIGHT_ONLY.
                 'compute_precision': _ComputePrecision.FLOAT,
@@ -633,7 +629,6 @@ class ConfiguratorTest(parameterized.TestCase, googletest.TestCase):
                     'num_bits': 6,
                     'symmetric': True,
                     'granularity': _QuantGranularity.TENSORWISE,
-                    'block_size': 0,
                 },
                 # WEIGHT_ONLY.
                 'compute_precision': _ComputePrecision.FLOAT,
@@ -652,7 +647,6 @@ class ConfiguratorTest(parameterized.TestCase, googletest.TestCase):
                     'num_bits': 3,
                     'symmetric': True,
                     'granularity': _QuantGranularity.TENSORWISE,
-                    'block_size': 0,
                 },
                 # WEIGHT_ONLY.
                 'compute_precision': _ComputePrecision.FLOAT,

ai_edge_quantizer/transformations/insert_decomposed_hadamard_rotation.py

@@ -220,6 +220,14 @@ def insert_decomposed_hadamard_rotation(
   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]
+  fc_options = schema_py_generated.FullyConnectedOptionsT()
+  fc_options.fusedActivationFunction = (
+      schema_py_generated.ActivationFunctionType.NONE
+  )
+  fc_op.builtinOptionsType = (
+      schema_py_generated.BuiltinOptions.FullyConnectedOptions
+  )
+  fc_op.builtinOptions = fc_options
 
   # Insert x' = tfl.reshape(x', x.shape)
   post_reshape_op_code_idx = transformation_utils.add_op_code(

ai_edge_quantizer/utils/constrained_ops_utils_test.py

@@ -28,7 +28,7 @@ class ConstrainedOpsUtilsTest(parameterized.TestCase):
       dict(
           testcase_name="same_as_input_scale",
           constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
-          expected_num_ops=16,
+          expected_num_ops=17,
       ),
       dict(
           testcase_name="same_as_output_scale",

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -75,6 +75,7 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.EQUAL: schema.BuiltinOperator.EQUAL,
     _TFLOpName.NOT_EQUAL: schema.BuiltinOperator.NOT_EQUAL,
     _TFLOpName.MIRROR_PAD: schema.BuiltinOperator.MIRROR_PAD,
+    _TFLOpName.SPACE_TO_DEPTH: schema.BuiltinOperator.SPACE_TO_DEPTH,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

ai_edge_quantizer/utils/validation_utils.py

@@ -32,7 +32,7 @@ def get_validation_func(
     a validation function
 
   Raises:
-    Value error if the function name is not supported
+    ValueError: if the function name is not supported
   """
   if func_name == "mse":
     return mean_squared_difference
@@ -40,6 +40,10 @@ def get_validation_func(
     return median_diff_ratio
   elif func_name == "cosine_similarity":
     return cosine_similarity
+  elif func_name == "kl_divergence":
+    return kl_divergence
+  elif func_name == "snr":
+    return signal_to_noise_ratio
   else:
     raise ValueError(f"Validation function {func_name} not supported")
 
@@ -60,7 +64,7 @@ def mean_squared_difference(
     a float value representing the MSD between data1 & 2
 
   Raises:
-    Value error if the two inputs don't have the same number of elements
+    ValueError: if the two inputs don't have the same number of elements
   """
   data1, data2 = _preprocess_same_size_arrays(data1, data2)
   # special handling for tensor of size 0
@@ -89,7 +93,7 @@ def median_diff_ratio(
     a float value representing the median diff ratio between data1 & 2
 
   Raises:
-    Value error if the two inputs don't have the same number of elements
+    ValueError: if the two inputs don't have the same number of elements
   """
   data1, data2 = _preprocess_same_size_arrays(data1, data2)
   # special handling for tensor of size 0
@@ -118,7 +122,7 @@ def cosine_similarity(
     a float value representing the cosine similarity between data1 & 2
 
   Raises:
-    Value error if the two inputs don't have the same number of elements
+    ValueError: if the two inputs don't have the same number of elements
   """
   data1, data2 = _preprocess_same_size_arrays(data1, data2)
   # special handling for tensor of size 0
@@ -134,6 +138,77 @@ def cosine_similarity(
   return np.dot(data1, data2) / (norm_data1 * norm_data2)
 
 
+def kl_divergence(
+    data1: np._typing.ArrayLike,
+    data2: np._typing.ArrayLike,
+    epsilon: float = 1e-9,
+) -> float:
+  """Calculates the KL divergence between data1 & data2.
+
+  KL(data2 || data1) = sum(data2 * log(data2 / data1)).
+  data2 is treated as the true distribution P, and data1 as the
+  approximated distribution Q.
+  Non-positive values in data1 and data2 are clipped to 0 before
+  KL divergence calculation. Epsilon is added to avoid log(0) and
+  division by zero.
+
+  Args:
+    data1: input data to be used for comparison (distribution Q)
+    data2: input data to be used for comparison (distribution P), data1 & 2 must
+      be of the same shape
+    epsilon: small value to avoid log(0) and division by zero.
+
+  Returns:
+    A float value representing the KL divergence between data1 & 2.
+
+  Raises:
+    ValueError: if the two inputs don't have the same number of elements.
+  """
+  data1, data2 = _preprocess_same_size_arrays(data1, data2)
+  # special handling for tensor of size 0
+  if data1.size == 0:
+    return float(0)
+
+  p = np.maximum(0, data2)
+  q = np.maximum(0, data1)
+
+  return float(np.sum(p * np.log((p + epsilon) / (q + epsilon))))
+
+
+def signal_to_noise_ratio(
+    noisy_signal: np._typing.ArrayLike,
+    signal: np._typing.ArrayLike,
+    epsilon: float = 1e-9,
+) -> float:
+  """Calculates the signal to noise ratio between noisy_signal & signal.
+
+  SNR = P_signal / P_noise, where signal is treated as the clean signal and
+  noisy_signal-signal is treated as the noise samples.
+  P_signal = mean(signal^2)
+  P_noise = mean((noisy_signal-signal)^2) = mse(noisy_signal, signal)
+
+  Args:
+    noisy_signal: Input data to be used for comparison (e.g. noisy signal).
+    signal: Input data to be used for comparison (e.g. clean signal),
+      noisy_signal & signal must be of the same shape.
+    epsilon: Small value to avoid division by zero.
+
+  Returns:
+    A float value representing the SNR between noisy_signal & signal.
+
+  Raises:
+    ValueError: If the two inputs don't have the same number of elements.
+  """
+  noisy_signal, signal = _preprocess_same_size_arrays(noisy_signal, signal)
+  if signal.size == 0:
+    return float(0)
+
+  mse = mean_squared_difference(noisy_signal, signal)
+  signal_power = float(np.square(signal).mean())
+  snr = signal_power / (mse + epsilon)
+  return snr
+
+
 def _preprocess_same_size_arrays(
     data1: np._typing.ArrayLike, data2: np._typing.ArrayLike
 ) -> Tuple[np.ndarray, np.ndarray]:
@@ -148,7 +223,7 @@ def _preprocess_same_size_arrays(
     a tuple of the preprocessed data1 & 2
 
   Raises:
-    Value error if the two inputs don't have the same number of elements
+    ValueError: if the two inputs don't have the same number of elements
   """
   data1 = np.array(data1, dtype=np.float32).flatten()
   data2 = np.array(data2, dtype=np.float32).flatten()