De-duplicate zero points for per channel quantized tensors when all the zero points are the same.

PiperOrigin-RevId: 763795520

Author: alankelly

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py

@@ -435,7 +435,8 @@ def _is_valid_quantization_params(
   """Checks if the quantization parameters are valid.
 
   A valid quantization params requires:
-    1. scale and zero point have the same shape (TFL Runtime requirement).
+    1. scale and zero point either have the same shape or the zero point is a
+    scalar.
     2. scale and zero point have the same rank as the tensor content (avoid
     ambiguous broadcasting).
 
@@ -446,17 +447,20 @@ def _is_valid_quantization_params(
   Returns:
     True if the quantization parameters are valid.
   """
-  if quantization_params.scale.shape != quantization_params.zero_point.shape:
+  if (
+      quantization_params.scale.shape != quantization_params.zero_point.shape
+      and quantization_params.zero_point.size != 1
+  ):
     raise ValueError(
-        "scale and zero_point must have the same shape. Got"
-        f" {quantization_params.scale.shape} and"
+        "scale and zero_point must have the same shape or zero_point must have"
+        f" only one element. Got {quantization_params.scale.shape} and"
         f" {quantization_params.zero_point.shape}"
     )
 
   tensor_rank = tensor_data.ndim
   scale_rank = quantization_params.scale.ndim
   zero_point_rank = quantization_params.zero_point.ndim
-  if (tensor_rank != scale_rank) or (tensor_rank != zero_point_rank):
+  if tensor_rank != scale_rank or (tensor_rank != zero_point_rank):
     raise ValueError(
         f"Ranks of scales ({scale_rank}) and zps"
         f" ({zero_point_rank}) must be the same as the tensor rank"

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor_test.py

@@ -160,7 +160,9 @@ def test_uniform_quantize(
   def test_uniform_quantize_wrong_shape(self):
     tensor = [-3.0, 1.3, 2.4, 16.0]
 
-    error_message = "scale and zero_point must have the same shape."
+    error_message = (
+        "Ranks of scales (3) and zps (2) must be the same as the tensor rank"
+    )
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: error_message in str(err)
     ):
@@ -233,7 +235,9 @@ def test_uniform_dequantize(
   def test_uniform_dequantize_wrong_shape(self):
     tensor = [-3.0, 1.3, 2.4, 16.0]
 
-    error_message = "scale and zero_point must have the same shape."
+    error_message = (
+        "Ranks of scales (3) and zps (2) must be the same as the tensor rank"
+    )
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: error_message in str(err)
     ):