ai-edge-quantizer-nightly 0.1.0.dev20250415__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -24,9 +24,12 @@ 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.
 
 _TFLOpName = qtyping.TFLOperationName
@@ -58,6 +61,10 @@ class AlgorithmName(str, enum.Enum):
   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 ###
 
@@ -99,11 +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(
@@ -232,11 +265,37 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     _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():
@@ -250,3 +309,102 @@ for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():
           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/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.