ai-edge-quantizer-nightly 0.1.0.dev20250407__py3-none-any.whl → 0.1.0.dev20250410__py3-none-any.whl

ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py

@@ -16,6 +16,7 @@
 """Performs naive min/max uniform quantization."""
 
 from typing import Any, Optional
+import ml_dtypes
 import numpy as np
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
@@ -73,12 +74,25 @@ def get_tensor_quant_params(
         " parameters. Check if the correct calibration results are passed into"
         " the ParamsGenerator."
     )
+  clipping_values = None
+  # Blockwise quantization uses float16 scale, with 7 bit mantissa,
+  # so the maximum representable value is 65280.
+  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
+    clipping_values = np.broadcast_to(
+        np.array(65280), tensor_min_max["min"].shape
+    )
   zp, scale = uniform_quantize_tensor.tensor_zp_scale_from_min_max(
       tensor_min_max["min"],
       tensor_min_max["max"],
       tensor_quant_config.num_bits,
       tensor_quant_config.symmetric,
+      clipping_values,
   )
+  # Round the scale values to 7 bit mantissa.
+  if tensor_quant_config.granularity == qtyping.QuantGranularity.BLOCKWISE:
+    scale = (
+        scale.astype(ml_dtypes.bfloat16).astype(np.float16).astype(np.float32)
+    )
   quantized_dim = None
   if tensor_quant_config.granularity == qtyping.QuantGranularity.CHANNELWISE:
     quantized_dim = common_utils.get_weight_quantized_dim(

ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py

@@ -22,7 +22,6 @@ import numpy as np
 from tensorflow.python.platform import googletest
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.algorithms.uniform_quantize import naive_min_max_quantize
-from ai_edge_quantizer.algorithms.uniform_quantize import uniform_quantize_tensor
 from ai_edge_quantizer.utils import test_utils
 from ai_edge_quantizer.utils import tfl_flatbuffer_utils
 
@@ -185,14 +184,13 @@ class NaiveMinMaxQuantizeTest(parameterized.TestCase):
     )
     scale = quant_params.scale
     zp = quant_params.zero_point
-    expected_zp, expected_scale = (
-        uniform_quantize_tensor.tensor_zp_scale_from_min_max(
-            min_value=np.array([[-7], [-4], [-4], [7]]),
-            max_value=np.array([[7], [4], [4], [7]]),
-            num_bits=4,
-            symmetric=True,
-        )
-    )
+    expected_scale = np.array([
+        [1],
+        [0.5703125],
+        [0.5703125],
+        [1],
+    ])
+    expected_zp = np.zeros([4, 1])
     self.assertTrue(np.array_equal(zp, expected_zp))
     self.assertTrue(np.array_equal(scale, expected_scale))
     self.assertIsNotNone(quant_params.quantized_data)

ai_edge_quantizer/params_generator.py

@@ -344,22 +344,58 @@ def _compatible_tensor_transformation_params(
     params2: qtyping.TensorTransformationParams,
 ) -> bool:
   """Check if two tensor transformation params are compatible."""
+  return (
+      _are_tensor_consumer_params_compatible(params1)
+      and _are_tensor_consumer_params_compatible(params2)
+      and _are_self_compatible_tensors_compatible_to_each_other(
+          params1, params2
+      )
+  )
+
+
+def _are_tensor_consumer_params_compatible(
+    params: qtyping.TensorTransformationParams,
+) -> bool:
+  """Check if all tensor's consumers have the same quantization parameters."""
+  if params.consumers is None or len(params.consumers) < 2:
+    return True
+  consumer_1 = params.consumers[0]
+  for consumer in params.consumers[1:]:
+    if not _compatible_tensor_params(consumer, consumer_1):
+      return False
+  return True
+
+
+def _are_self_compatible_tensors_compatible_to_each_other(
+    params1: qtyping.TensorTransformationParams,
+    params2: qtyping.TensorTransformationParams,
+) -> bool:
+  """Check if two self compatible tensors are compatible to each other.
+
+  Self compatible means that all tensor's consumers have the same quantization
+  parameters.
+
+  Args:
+    params1: The first tensor transformation params.
+    params2: The second tensor transformation params.
+
+  Returns:
+    Whether the two tensors are compatible to each other.
+  """
+  # Check the producer.
   if params1.producer is None or params2.producer is None:
     if params1.producer != params2.producer:
       return False
   elif not _compatible_tensor_params(params1.producer, params2.producer):
     return False
+
+  # Check the consumers.
   if params1.consumers is None or params2.consumers is None:
     if params1.consumers != params2.consumers:
       return False
   else:
-    # Check all consumers within each params are compatible.
-    for params1_consumer in params1.consumers:
-      if not _compatible_tensor_params(params1_consumer, params1.consumers[0]):
-        return False
-    for params2_consumer in params2.consumers:
-      if not _compatible_tensor_params(params2_consumer, params2.consumers[0]):
-        return False
+    # Since all consumer params within each tensor are the same, it's enough to
+    # check only the first consumers.
     if not _compatible_tensor_params(
         params1.consumers[0], params2.consumers[0]
     ):

ai_edge_quantizer/params_generator_test.py

@@ -37,8 +37,11 @@ _TensorQuantConfig = qtyping.TensorQuantizationConfig
 _QuantTransformation = qtyping.QuantTransformation
 _AlgorithmName = recipe_manager.AlgorithmName
 _QuantGranularity = qtyping.QuantGranularity
+_QTransf = qtyping.QuantTransformation
+
 
 TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile('')
+_PARAMS_8BIT = qtyping.UniformQuantParams(8, None, np.array([1]), np.array([0]))
 
 
 def _single_fc_model_representative_dataset_gen(num_samples=5):
@@ -63,6 +66,20 @@ def _get_calibration_data(
   return calibration_data
 
 
+def _get_test_consumers(
+    transformations_per_consumer: list[list[_QTransf]],
+    params_per_consumer: list[qtyping.OpToTensorParams],
+) -> list[qtyping.OpToTensorParams]:
+  return [
+      qtyping.OpToTensorParams(
+          subgraph_op_id=i + 1,
+          transformations=transformations_per_consumer[i],
+          parameters=params_per_consumer[i],
+      )
+      for i in range(len(transformations_per_consumer))
+  ]
+
+
 class ParamsGeneratorTest(parameterized.TestCase):
 
   def setUp(self):
@@ -635,12 +652,12 @@ class ParamsGeneratorTest(parameterized.TestCase):
           self.assertNotEmpty(consumer.transformations)
           self.assertEqual(
               consumer.transformations[0],
-              qtyping.QuantTransformation.DUPLICATE_BUFFER,
+              _QTransf.DUPLICATE_BUFFER,
           )
       elif quant_params[tensor_name].consumers is not None:
         for consumer in quant_params[tensor_name].consumers:
           self.assertNotIn(
-              qtyping.QuantTransformation.DUPLICATE_BUFFER,
+              _QTransf.DUPLICATE_BUFFER,
               consumer.transformations,
           )
 
@@ -651,328 +668,182 @@ class ParamsGeneratorTest(parameterized.TestCase):
               tensor_name='tfl.quantize',
               producer=qtyping.OpToTensorParams(
                   subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.ADD_DEQUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+                  transformations=[_QTransf.ADD_DEQUANTIZE],
+                  parameters=_PARAMS_8BIT,
+              ),
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.NO_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.NO_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
           param2=qtyping.TensorTransformationParams(
-              'tfl.other_quantize',
-              qtyping.OpToTensorParams(
+              tensor_name='tfl.other_quantize',
+              producer=qtyping.OpToTensorParams(
                   subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.NO_QUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+                  transformations=[_QTransf.NO_QUANTIZE],
+                  parameters=_PARAMS_8BIT,
+              ),
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              [
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
           expected=False,
       ),
       dict(
-          testcase_name='param2_consumer_incompatible',
+          testcase_name='compatible',
           param1=qtyping.TensorTransformationParams(
               tensor_name='tfl.quantize',
-              producer=qtyping.OpToTensorParams(
-                  subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.ADD_QUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE, _QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
           param2=qtyping.TensorTransformationParams(
-              'tfl.other_quantize',
-              qtyping.OpToTensorParams(
-                  subgraph_op_id=0,
-                  transformations=[qtyping.QuantTransformation.NO_QUANTIZE],
-                  parameters=qtyping.UniformQuantParams(
-                      8, None, np.array([1]), np.array([0])
-                  ),
+              tensor_name='tfl.other_quantize',
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
               ),
-              [
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.QUANTIZE_TENSOR,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-              ],
           ),
-          expected=False,
+          expected=True,
       ),
       dict(
-          testcase_name='compatible',
+          testcase_name='compatible_no_numeric_check',
           param1=qtyping.TensorTransformationParams(
               tensor_name='tfl.quantize',
               producer=None,
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00028806]), np.array([0])
                       ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.NO_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00027501]), np.array([0])
                       ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.NO_QUANTIZE,
-                      ],
-                  ),
-              ],
+                  ],
+              ),
           ),
           param2=qtyping.TensorTransformationParams(
-              'tfl.other_quantize',
-              None,
-              [
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=1,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=2,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=3,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                          qtyping.QuantTransformation.ADD_DEQUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+              tensor_name='tfl.quantize',
+              producer=None,
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00028806]), np.array([0])
                       ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([1]), np.array([0])
+                      qtyping.UniformQuantParams(
+                          8, None, np.array([0.00027501]), np.array([0])
                       ),
-                  ),
-              ],
+                  ],
+              ),
           ),
           expected=True,
       ),
+  )
+  def test__are_self_compatible_tensors_compatible_to_each_other(
+      self, param1, param2, expected
+  ):
+    self.assertEqual(
+        params_generator._are_self_compatible_tensors_compatible_to_each_other(
+            param1, param2
+        ),
+        expected,
+    )
+
+  @parameterized.named_parameters(
       dict(
-          testcase_name='compatible_no_numeric_check',
-          param1=qtyping.TensorTransformationParams(
+          testcase_name='consumer_incompatible',
+          params=qtyping.TensorTransformationParams(
+              tensor_name='tfl.quantize',
+              producer=qtyping.OpToTensorParams(
+                  subgraph_op_id=0,
+                  transformations=[_QTransf.NO_QUANTIZE],
+                  parameters=_PARAMS_8BIT,
+              ),
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.QUANTIZE_TENSOR],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
+              ),
+          ),
+          expected=False,
+      ),
+      dict(
+          testcase_name='compatible',
+          params=qtyping.TensorTransformationParams(
               tensor_name='tfl.quantize',
               producer=None,
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([0.00028806]), np.array([0])
-                      ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=5,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
-                          8, None, np.array([0.00027501]), np.array([0])
-                      ),
-                  ),
-              ],
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE, _QTransf.ADD_DEQUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE, _QTransf.ADD_QUANTIZE],
+                      [_QTransf.NO_QUANTIZE],
+                  ],
+                  params_per_consumer=[_PARAMS_8BIT] * 4,
+              ),
           ),
-          param2=qtyping.TensorTransformationParams(
+          expected=True,
+      ),
+      dict(
+          testcase_name='compatible_no_numeric_check',
+          params=qtyping.TensorTransformationParams(
               tensor_name='tfl.quantize',
               producer=None,
-              consumers=[
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=4,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
+              consumers=_get_test_consumers(
+                  transformations_per_consumer=[
+                      [_QTransf.ADD_QUANTIZE],
+                      [_QTransf.ADD_QUANTIZE],
+                  ],
+                  params_per_consumer=[
+                      qtyping.UniformQuantParams(
                           8, None, np.array([0.00028806]), np.array([0])
                       ),
-                  ),
-                  qtyping.OpToTensorParams(
-                      subgraph_op_id=5,
-                      transformations=[
-                          qtyping.QuantTransformation.ADD_QUANTIZE,
-                      ],
-                      parameters=qtyping.UniformQuantParams(
+                      qtyping.UniformQuantParams(
                           8, None, np.array([0.00027501]), np.array([0])
                       ),
-                  ),
-              ],
+                  ],
+              ),
           ),
           expected=True,
       ),
   )
-  def test_params_compatible(self, param1, param2, expected):
-    # adding a test to make production coverage happy.
+  def test__are_tensor_consumer_params_compatible(self, params, expected):
     self.assertEqual(
-        params_generator._compatible_tensor_transformation_params(
-            param1, param2
-        ),
+        params_generator._are_tensor_consumer_params_compatible(params),
         expected,
     )
 

ai_edge_quantizer/transformations/quantize_tensor.py

@@ -16,6 +16,7 @@
 """quantize a given tensor."""
 
 from typing import Optional, cast
+import ml_dtypes
 import numpy as np
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.transformations import transformation_utils
@@ -121,26 +122,6 @@ def _perform_channelwise_quantization(
   return flatbuffer_quantization
 
 
-def _downcast_and_truncate_scale(input_scale: np.ndarray) -> np.ndarray:
-  """Given a fp32 scale, downcast it to fp16 and truncate mantissa to 7 bits.
-
-  CPU kernel can only utilize 7 bits of mantissa for fp16, so we want to produce
-  scale this way to unify behaviours across different platforms.
-
-  Args:
-    input_scale: The input scale in fp32.
-
-  Returns:
-    The downcasted & truncated scale in fp16.
-  """
-
-  # A regular fp16 has 10 bits of mantissa, so we need to zero out the 3 least
-  # significant bits.
-  return (
-      input_scale.astype(np.float16).view(dtype=np.uint16) & np.uint16(0xFFF8)
-  ).view(dtype=np.float16)
-
-
 def _perform_blockwise_quantization(
     transformation_input: transformation_utils.TransformationInput,
 ) -> schema_py_generated.QuantizationParametersT():
@@ -162,13 +143,12 @@ def _perform_blockwise_quantization(
   )
   tensor = transformation_input.subgraph.tensors[transformation_input.tensor_id]
   blockwise_details = schema_py_generated.BlockwiseQuantizationT()
-  # Downcast and truncate the scale to fp16.
-  downcasted_scale = _downcast_and_truncate_scale(
-      transformation_input.quant_params.scale
-  )
+  # Downcast and round the scale to fp16 with 7 bit mantissa.
   scale_tensor_id = transformation_utils.add_new_constant_tensor(
       tensor.name + b"_scales",
-      downcasted_scale,
+      transformation_input.quant_params.scale.astype(ml_dtypes.bfloat16).astype(
+          np.float16
+      ),
       schema_py_generated.TensorType.FLOAT16,
       transformation_input.subgraph,
       transformation_input.buffers,

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -231,7 +231,8 @@ def get_tensor_data(tensor: Any, buffers: list[Any]) -> Optional[np.ndarray]:
   data = np.frombuffer(
       buffer_data, dtype=TENSOR_CODE_TO_TYPE[tensor.type].lower()
   )
-  data = np.reshape(data, tensor.shape)
+  if tensor.shape is not None:
+    data = np.reshape(data, tensor.shape)
   return data
 
 

{ai_edge_quantizer_nightly-0.1.0.dev20250407.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250410.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.1.0.dev20250407
+Version: 0.1.0.dev20250410
 Summary: A quantizer for advanced developers to quantize converted AI Edge models.
 Home-page: https://github.com/google-ai-edge/ai-edge-quantizer
 Keywords: On-Device ML,AI,Google,TFLite,Quantization,LLMs,GenAI

{ai_edge_quantizer_nightly-0.1.0.dev20250407.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250410.dist-info}/RECORD

@@ -10,8 +10,8 @@ ai_edge_quantizer/model_modifier.py,sha256=SPt9X-xBzRvcd4xIS24zLHt3aUS2QwsNDqweF
 ai_edge_quantizer/model_modifier_test.py,sha256=cJd04SLOG-fQZZNZPcisoBLx3cLtWEwGqUBbLb-pif4,4751
 ai_edge_quantizer/model_validator.py,sha256=fRNz0jO54cthPTibsCuViUXUuFRHl_fbvEiCukIVy20,13030
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
-ai_edge_quantizer/params_generator.py,sha256=46XDjnP4R3m4xsoXNp7brv0sNQPdQMg217_CbEl-Wgg,15780
-ai_edge_quantizer/params_generator_test.py,sha256=9WTUl87XqbM4NruX5ypLuVRtuhcw-CmxndsMOUzZ92Q,43171
+ai_edge_quantizer/params_generator.py,sha256=PeIwoNYg4kJq0cMPucTvyxXTqD0I1Sr8vm5xHZCQ518,16774
+ai_edge_quantizer/params_generator_test.py,sha256=DhULRWs1-UuO55zuuxocMjWDClcjcaKaue6mOcoHq9E,37186
 ai_edge_quantizer/qtyping.py,sha256=FqelZu7j0fGBRSCv_VVsuf3VmbfVlYJGgsjvdMXGgaw,15284
 ai_edge_quantizer/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
 ai_edge_quantizer/quantizer_test.py,sha256=K_HBA56JkFI3HL8VLWCqGEfC0ISh5ldMKoNyBdGRAJg,20368
@@ -32,8 +32,8 @@ ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=SVu1RSX5
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py,sha256=qMmKbWqxrCoVKbLKHn9WuCrGKPfHkEyU0Nmhokh8Qeo,2597
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=Fk3s9Qy2A_hjUepFOUmTwIZ_wKYVPbdDX4eoP-eoAQU,8726
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py,sha256=sT5eX5TLZEHTtPfnSkCPDlS0sQxlTFWbCsbvOuj--yY,8889
-ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py,sha256=cbyyYAoQnEraOYSV00wZ557ElBndHduVGeHikYUEFCE,7995
-ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py,sha256=B30SEISYZ9DPs3suKeG2elgXylR98pCEMWSEGgZo20o,7648
+ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize.py,sha256=Divlsn3NjNGtH0vlvE91wxL-VHb4q1nUE0JTDGiEtYc,8572
+ai_edge_quantizer/algorithms/uniform_quantize/naive_min_max_quantize_test.py,sha256=zoF_EHjYqsKkuev8wfuutIITEmp_maa70IpJI_Df3ck,7431
 ai_edge_quantizer/algorithms/uniform_quantize/octav.py,sha256=e5wYtki-vl739gSVAZHAKcs2hA87GvFUjVoSUPlnkyM,6433
 ai_edge_quantizer/algorithms/uniform_quantize/octav_test.py,sha256=IcTOaJ1pxtqsitqxOEP9LROVEP_19VFutHalqNied4I,6940
 ai_edge_quantizer/algorithms/uniform_quantize/uniform_quantize_tensor.py,sha256=WmZzKQlzfu9gFr9SbUDoPY3rFqTl363om8-0rTLwotw,11629
@@ -52,7 +52,7 @@ ai_edge_quantizer/transformations/emulated_subchannel.py,sha256=HVaRxoC8PCAvy3xe
 ai_edge_quantizer/transformations/emulated_subchannel_test.py,sha256=gZP6u9NdPXl7s19qB_Un8evou9ZZV6I9Gy0E1rdobHM,7722
 ai_edge_quantizer/transformations/quant_insert.py,sha256=jn6HsJaV-sqBiFPY-Aqbd64t8zgcYVkEkZI375x_FWY,3958
 ai_edge_quantizer/transformations/quant_insert_test.py,sha256=X9ptPDvJCFkR5tejKnD1SlHFGPazQTW-wNNMV9MEAuw,10107
-ai_edge_quantizer/transformations/quantize_tensor.py,sha256=y6As38mTzhva50YvNQ7p0SFpuWet3LPqFwE3qIO0gEQ,8231
+ai_edge_quantizer/transformations/quantize_tensor.py,sha256=kjaNrw9mnrn0t8u0vey9S_uPz3iVUicwy4rluxVqV3E,7617
 ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=mHLO3_MRt36A8-ZN8ADn5tBBJlqjTWa7ZUN8Mmu5Rcw,9116
 ai_edge_quantizer/transformations/transformation_utils.py,sha256=5w0fG6TP362elTHs-JZokl24fuK4Gv6DGyIpybQYb3g,4885
 ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=xH64SF3UHDh84vYbt-WvmXNjM-Jg-mefES1ACO1tkqw,6269
@@ -60,14 +60,14 @@ ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V
 ai_edge_quantizer/utils/calibration_utils.py,sha256=1Fj9MIO6aLZIRgyd4axvZN4S_O64nB_-Miu1WP664js,2536
 ai_edge_quantizer/utils/calibration_utils_test.py,sha256=Z-AcdTieesWFKyKBb08ZXm4Mgu6cvJ4bg2-MJ7hLD10,2856
 ai_edge_quantizer/utils/test_utils.py,sha256=HwZCIpO9fJRAhuN6t6voXKOYQtcioFtt_tpkAlDsAYk,6205
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=irrGbbOt14PLFcS4538II0dB4Q7YJMgGvpBERVHevXM,10535
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=NKtw60BJAjIE6Yww8B1vJpxXwp4MSERmpKajXJWm5rI,10568
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
 ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=x2xA2CFPpe_2trcV8v5xGaBETvVCfwAcJuq6yieGJ0Y,12687
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=Op3JxtOqlrjzmYF18jnnstL1k9xiY9kKJ8S2vklKGkc,11327
 ai_edge_quantizer/utils/validation_utils.py,sha256=oYw33Sg547AqtGw-choPUJmp9SAKkV46J_ddqSsum2Q,3950
 ai_edge_quantizer/utils/validation_utils_test.py,sha256=V_qNDikPD4OPB-siOLQCWNVWTAu87h2IgNYt7teFd-o,2934
-ai_edge_quantizer_nightly-0.1.0.dev20250407.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250407.dist-info/METADATA,sha256=u28A5AptBWTj0kkQD0QfUKvfwYUGbzTmi1NzqiItobc,1527
-ai_edge_quantizer_nightly-0.1.0.dev20250407.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250407.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250407.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250410.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250410.dist-info/METADATA,sha256=PeTpfQxSLhSk9QyphqCPCHKAytqdjDAf9FzubQofMmM,1527
+ai_edge_quantizer_nightly-0.1.0.dev20250410.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250410.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250410.dist-info/RECORD,,