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

PiperOrigin-RevId: 764435889

Author: ai-edge-bot

ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py

@@ -435,8 +435,7 @@ def _is_valid_quantization_params(
   """Checks if the quantization parameters are valid.
 
   A valid quantization params requires:
-    1. scale and zero point either have the same shape or the zero point is a
-    scalar.
+    1. scale and zero point have the same shape (TFL Runtime requirement).
     2. scale and zero point have the same rank as the tensor content (avoid
     ambiguous broadcasting).
 
@@ -447,20 +446,17 @@ def _is_valid_quantization_params(
   Returns:
     True if the quantization parameters are valid.
   """
-  if (
-      quantization_params.scale.shape != quantization_params.zero_point.shape
-      and quantization_params.zero_point.size != 1
-  ):
+  if quantization_params.scale.shape != quantization_params.zero_point.shape:
     raise ValueError(
-        "scale and zero_point must have the same shape or zero_point must have"
-        f" only one element. Got {quantization_params.scale.shape} and"
+        "scale and zero_point must have the same shape. Got"
+        f" {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,9 +160,7 @@ def test_uniform_quantize(
   def test_uniform_quantize_wrong_shape(self):
     tensor = [-3.0, 1.3, 2.4, 16.0]
 
-    error_message = (
-        "Ranks of scales (3) and zps (2) must be the same as the tensor rank"
-    )
+    error_message = "scale and zero_point must have the same shape."
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: error_message in str(err)
     ):
@@ -235,9 +233,7 @@ def test_uniform_dequantize(
   def test_uniform_dequantize_wrong_shape(self):
     tensor = [-3.0, 1.3, 2.4, 16.0]
 
-    error_message = (
-        "Ranks of scales (3) and zps (2) must be the same as the tensor rank"
-    )
+    error_message = "scale and zero_point must have the same shape."
     with self.assertRaisesWithPredicateMatch(
         ValueError, lambda err: error_message in str(err)
     ):