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

ai_edge_quantizer/algorithm_manager.py

@@ -17,13 +17,18 @@
 
 import enum
 import functools
+from immutabledict import immutabledict
 from ai_edge_quantizer import algorithm_manager_api
 from ai_edge_quantizer import default_policy
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.algorithms.nonlinear_quantize import float_casting
 from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
 from ai_edge_quantizer.algorithms.uniform_quantize import dequantized_weight_recovery
+from ai_edge_quantizer.algorithms.uniform_quantize import hadamard_rotation
+from ai_edge_quantizer.algorithms.uniform_quantize import mse
 from ai_edge_quantizer.algorithms.uniform_quantize import naive_min_max_quantize
+from ai_edge_quantizer.algorithms.uniform_quantize import octav
+
 
 # TODO: b/399775701 - Clean up this file.
 
@@ -55,6 +60,11 @@ class AlgorithmName(str, enum.Enum):
   MIN_MAX_UNIFORM_QUANT = naive_min_max_quantize.ALGORITHM_KEY
   FLOAT_CASTING = float_casting.ALGORITHM_KEY
   DEQUANTIZED_WEIGHT_RECOVERY = dequantized_weight_recovery.ALGORITHM_KEY
+  OCTAV = octav.ALGORITHM_KEY
+  HADAMARD_ROTATION = hadamard_rotation.CUSTOM_OP_ALGORITHM_KEY
+  DECOMPOSED_HADAMARD_ROTATION = hadamard_rotation.DECOMPOSED_ALGORITHM_KEY
+  MSE = mse.ALGORITHM_KEY
+
 
 ### MIN/MAX_UNIFORM_QUANT ###
 
@@ -96,7 +106,37 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.LOGISTIC: common_quantize.materialize_softmax_and_logistic,
     _TFLOpName.SLICE: common_quantize.materialize_slice,
     _TFLOpName.SUM: common_quantize.materialize_sum,
+    _TFLOpName.SELECT: common_quantize.materialize_select,
     _TFLOpName.SELECT_V2: common_quantize.materialize_select_v2,
+    _TFLOpName.DYNAMIC_UPDATE_SLICE: (
+        common_quantize.materialize_dynamic_update_slice
+    ),
+    _TFLOpName.STABLEHLO_COMPOSITE: common_quantize.materialize_composite,
+    _TFLOpName.PAD: common_quantize.materialize_pad,
+    _TFLOpName.SQUARED_DIFFERENCE: (
+        common_quantize.materialize_squared_difference
+    ),
+    _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
+    _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
+    _TFLOpName.RESIZE_NEAREST_NEIGHBOR: (
+        common_quantize.materialize_resize_nearest_neighbor
+    ),
+    _TFLOpName.GATHER_ND: common_quantize.materialize_gather_nd,
+    _TFLOpName.PACK: common_quantize.materialize_pack,
+    _TFLOpName.UNPACK: common_quantize.materialize_unpack,
+    _TFLOpName.DIV: common_quantize.materialize_div,
+    _TFLOpName.BROADCAST_TO: common_quantize.materialize_broadcast_to,
+    _TFLOpName.SQRT: common_quantize.materialize_sqrt,
+    _TFLOpName.GATHER: common_quantize.materialize_gather,
+    _TFLOpName.HARD_SWISH: common_quantize.materialize_hard_swish,
+    _TFLOpName.MAXIMUM: common_quantize.materialize_maximum,
+    _TFLOpName.PADV2: common_quantize.materialize_padv2,
+    _TFLOpName.REDUCE_MIN: common_quantize.materialize_reduce_min,
+    _TFLOpName.EQUAL: common_quantize.materialize_equal,
+    _TFLOpName.NOT_EQUAL: common_quantize.materialize_not_equal,
+    _TFLOpName.MIRROR_PAD: common_quantize.materialize_mirror_pad,
+    _TFLOpName.SPACE_TO_DEPTH: common_quantize.materialize_space_to_depth,
+    _TFLOpName.RELU: common_quantize.materialize_relu,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -163,6 +203,7 @@ register_config_check_policy_func(
 
 DEQUANTIZED_WEIGHT_RECOVERY_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _TFLOpName.FULLY_CONNECTED: common_quantize.materialize_fc_conv,
+    _TFLOpName.CONV_2D: common_quantize.materialize_fc_conv,
     _TFLOpName.EMBEDDING_LOOKUP: common_quantize.materialize_embedding_lookup,
 }
 
@@ -184,3 +225,186 @@ for (
           dequantized_weight_recovery.get_tensor_quant_params,
       ),
   )
+
+
+# Register OCTAV algorithm.
+register_op_quant_config_validation_func(
+    AlgorithmName.OCTAV,
+    common_quantize.check_op_quantization_config,
+)
+
+# Register a config check policy for OCTAV algorithm.
+register_config_check_policy_func(
+    AlgorithmName.OCTAV,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+_OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
+    _TFLOpName.INPUT: common_quantize.materialize_input,
+    _TFLOpName.OUTPUT: common_quantize.materialize_output,
+    _TFLOpName.FULLY_CONNECTED: common_quantize.materialize_fc_conv,
+    _TFLOpName.BATCH_MATMUL: common_quantize.materialize_batch_matmul,
+    _TFLOpName.CONV_2D: common_quantize.materialize_fc_conv,
+    _TFLOpName.DEPTHWISE_CONV_2D: common_quantize.materialize_fc_conv,
+    _TFLOpName.CONV_2D_TRANSPOSE: common_quantize.materialize_conv2d_transpose,
+    _TFLOpName.RESHAPE: common_quantize.materialize_reshape,
+    _TFLOpName.AVERAGE_POOL_2D: common_quantize.materialize_average_pool_2d,
+    _TFLOpName.EMBEDDING_LOOKUP: common_quantize.materialize_embedding_lookup,
+    _TFLOpName.SOFTMAX: common_quantize.materialize_softmax_and_logistic,
+    _TFLOpName.TANH: common_quantize.materialize_tanh,
+    _TFLOpName.TRANSPOSE: common_quantize.materialize_transpose,
+    _TFLOpName.GELU: common_quantize.materialize_gelu,
+    _TFLOpName.ADD: common_quantize.materialize_add,
+    _TFLOpName.SUB: common_quantize.materialize_sub,
+    _TFLOpName.MUL: common_quantize.materialize_mul,
+    _TFLOpName.MEAN: common_quantize.materialize_mean,
+    _TFLOpName.RSQRT: common_quantize.materialize_rsqrt,
+    _TFLOpName.CONCATENATION: common_quantize.materialize_concatenation,
+    _TFLOpName.STRIDED_SLICE: common_quantize.materialize_strided_slice,
+    _TFLOpName.SPLIT: common_quantize.materialize_split,
+    _TFLOpName.LOGISTIC: common_quantize.materialize_softmax_and_logistic,
+    _TFLOpName.SLICE: common_quantize.materialize_slice,
+    _TFLOpName.SUM: common_quantize.materialize_sum,
+    _TFLOpName.SELECT: common_quantize.materialize_select,
+    _TFLOpName.SELECT_V2: common_quantize.materialize_select_v2,
+    _TFLOpName.DYNAMIC_UPDATE_SLICE: (
+        common_quantize.materialize_dynamic_update_slice
+    ),
+    _TFLOpName.STABLEHLO_COMPOSITE: common_quantize.materialize_composite,
+    _TFLOpName.PAD: common_quantize.materialize_pad,
+    _TFLOpName.SQUARED_DIFFERENCE: (
+        common_quantize.materialize_squared_difference
+    ),
+    _TFLOpName.MAX_POOL_2D: common_quantize.materialize_max_pool_2d,
+    _TFLOpName.RESIZE_BILINEAR: common_quantize.materialize_resize_bilinear,
+    _TFLOpName.RESIZE_NEAREST_NEIGHBOR: (
+        common_quantize.materialize_resize_nearest_neighbor
+    ),
+    _TFLOpName.GATHER_ND: common_quantize.materialize_gather_nd,
+    _TFLOpName.PACK: common_quantize.materialize_pack,
+    _TFLOpName.UNPACK: common_quantize.materialize_unpack,
+    _TFLOpName.DIV: common_quantize.materialize_div,
+    _TFLOpName.BROADCAST_TO: common_quantize.materialize_broadcast_to,
+    _TFLOpName.SQRT: common_quantize.materialize_sqrt,
+    _TFLOpName.GATHER: common_quantize.materialize_gather,
+    _TFLOpName.HARD_SWISH: common_quantize.materialize_hard_swish,
+    _TFLOpName.MAXIMUM: common_quantize.materialize_maximum,
+    _TFLOpName.PADV2: common_quantize.materialize_padv2,
+    _TFLOpName.REDUCE_MIN: common_quantize.materialize_reduce_min,
+    _TFLOpName.EQUAL: common_quantize.materialize_equal,
+    _TFLOpName.NOT_EQUAL: common_quantize.materialize_not_equal,
+    _TFLOpName.MIRROR_PAD: common_quantize.materialize_mirror_pad,
+    _TFLOpName.SPACE_TO_DEPTH: common_quantize.materialize_space_to_depth,
+    _TFLOpName.RELU: common_quantize.materialize_relu,
+})
+
+for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      AlgorithmName.OCTAV,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=functools.partial(
+          materialize_func,
+          octav.get_tensor_quant_params,
+      ),
+  )
+
+# Register the Hadamard Rotation algorithm.
+register_op_quant_config_validation_func(
+    AlgorithmName.HADAMARD_ROTATION,
+    common_quantize.check_op_quantization_config,
+)
+
+# Register a config check policy for the Hadamard Rotation algorithm.
+register_config_check_policy_func(
+    AlgorithmName.HADAMARD_ROTATION,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+# Register specialized hadamard rotation materialize functions.
+_HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
+    _TFLOpName.FULLY_CONNECTED: (
+        hadamard_rotation.materialize_fully_connected_custom_op
+    ),
+    _TFLOpName.EMBEDDING_LOOKUP: (
+        hadamard_rotation.materialize_embedding_lookup_custom_op
+    ),
+})
+for (
+    op_name,
+    materialize_func,
+) in _HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      AlgorithmName.HADAMARD_ROTATION,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=materialize_func,
+  )
+
+register_op_quant_config_validation_func(
+    AlgorithmName.DECOMPOSED_HADAMARD_ROTATION,
+    common_quantize.check_op_quantization_config,
+)
+
+register_config_check_policy_func(
+    AlgorithmName.DECOMPOSED_HADAMARD_ROTATION,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+_DECOMPOSED_HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
+    _TFLOpName.FULLY_CONNECTED: (
+        hadamard_rotation.materialize_fully_connected_decomposed
+    ),
+    _TFLOpName.EMBEDDING_LOOKUP: (
+        hadamard_rotation.materialize_embedding_lookup_decomposed
+    ),
+})
+for (
+    op_name,
+    materialize_func,
+) in _DECOMPOSED_HADAMARD_ROTATION_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      AlgorithmName.DECOMPOSED_HADAMARD_ROTATION,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=materialize_func,
+  )
+
+
+# Register the MSE algorithm.
+register_op_quant_config_validation_func(
+    AlgorithmName.MSE,
+    common_quantize.check_op_quantization_config,
+)
+
+# Register a config check policy for the MSE algorithm.
+register_config_check_policy_func(
+    AlgorithmName.MSE,
+    default_policy.DEFAULT_CONFIG_CHECK_POLICY,
+)
+
+# Register specialized MSE materialize functions.
+_MSE_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
+    _TFLOpName.FULLY_CONNECTED: common_quantize.materialize_fc_conv,
+    _TFLOpName.EMBEDDING_LOOKUP: common_quantize.materialize_embedding_lookup,
+    _TFLOpName.CONV_2D: common_quantize.materialize_fc_conv,
+    _TFLOpName.DEPTHWISE_CONV_2D: common_quantize.materialize_fc_conv,
+    _TFLOpName.CONV_2D_TRANSPOSE: common_quantize.materialize_conv2d_transpose,
+})
+for (
+    op_name,
+    materialize_func,
+) in _MSE_OP_NAME_MATERIALIZE_FUNC_DICT.items():
+  register_quantized_op(
+      AlgorithmName.MSE,
+      op_name,
+      naive_min_max_quantize.init_qsvs,
+      calibration_func=naive_min_max_quantize.min_max_calibrate,
+      materialize_func=functools.partial(
+          materialize_func,
+          mse.get_tensor_quant_params,
+      ),
+  )

ai_edge_quantizer/algorithm_manager_api_test.py

@@ -18,6 +18,7 @@
 from absl.testing import parameterized
 from tensorflow.python.platform import googletest
 from ai_edge_quantizer import algorithm_manager_api
+from ai_edge_quantizer import default_policy
 from ai_edge_quantizer import qtyping
 
 _TFLOpName = qtyping.TFLOperationName
@@ -205,6 +206,12 @@ class AlgorithmManagerApiTest(parameterized.TestCase):
         self._alg_manager._config_check_policy_registry[test_algorithm_name]
     )
 
+  def test_default_policy_not_empty(self):
+    """Tests that the default policy is not empty & no empty policy is generated."""
+    self.assertNotEmpty(default_policy.DEFAULT_CONFIG_CHECK_POLICY)
+    for policy in default_policy.DEFAULT_CONFIG_CHECK_POLICY.values():
+      self.assertNotEmpty(policy)
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py

@@ -531,9 +531,9 @@ class Fp16QuantizeTest(parameterized.TestCase):
 
     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"
+        "jit(export_func)/jit(main)/...y,yz->...z/dot_general;jit(export_func)/jit(main)/jit(_one_hot)/eq;jit(export_func)/jit(main)/jit(_one_hot)/convert_element_type"
     )
-    op_tensor_names["input"] = "inputs"
+    op_tensor_names["input"] = "lookup"
     op_tensor_names["output"] = "Identity_1"
 
     # TODO: b/335913710 - Rename the test function.