ai-edge-quantizer-nightly 0.4.0.dev20250926__py3-none-any.whl → 0.4.0.dev20250928__py3-none-any.whl

ai_edge_quantizer/recipe_manager.py

@@ -82,7 +82,6 @@ class RecipeManager:
         str, list[OpQuantizationRecipe]
     ] = collections.OrderedDict()
 
-  # TODO: b/335254997 - Check if an op quantization config is supported.
   def add_quantization_config(
       self,
       regex: str,
@@ -272,7 +271,8 @@ class RecipeManager:
     """
     weight_config = qtyping.TensorQuantizationConfig(
         num_bits=num_bits,
-        symmetric=True,  # LiteRT kernels only support symmetric quantized weights.
+        symmetric=True,  # LiteRT kernels only support symmetric quantized
+        # weights.
         granularity=granularity,
     )
     self.add_quantization_config(
@@ -316,10 +316,18 @@ class RecipeManager:
       granularity: Granularity of quantization.
       algorithm_key: Algorithm key to be applied.
     """
+    # Default to integer quantization but allow float quantization for
+    # FLOAT_CASTING algorithm. This is to support weight-only quantization with
+    # fp16 weights.
+    weight_dtype = qtyping.TensorDataType.INT
+    if algorithm_key == AlgorithmName.FLOAT_CASTING:
+      weight_dtype = qtyping.TensorDataType.FLOAT
+
     weight_config = qtyping.TensorQuantizationConfig(
         num_bits=num_bits,
         symmetric=True,  # TFL kernels only support symmetric quantized weights.
         granularity=granularity,
+        dtype=weight_dtype,
     )
     self.add_quantization_config(
         regex,
@@ -365,7 +373,8 @@ class RecipeManager:
       raise ValueError(
           'Activation quantization is only supported for 16 or 8 bits.'
       )
-    # INT16 is symmetric and INT8 is asymmetric due to LiteRT kernel limitations.
+    # INT16 is symmetric and INT8 is asymmetric due to LiteRT kernel
+    # limitations.
     activation_symmetric = activation_num_bits == 16
     activation_config = qtyping.TensorQuantizationConfig(
         num_bits=activation_num_bits, symmetric=activation_symmetric

ai_edge_quantizer/recipe_manager_test.py

@@ -315,11 +315,12 @@ class ConfiguratorTest(parameterized.TestCase, googletest.TestCase):
         _QuantGranularity.CHANNELWISE,
     )
 
-  def test_add_weight_only_config(self):
+  @parameterized.parameters(4, 8)
+  def test_add_weight_only_config_int(self, num_bits):
     self._recipe_manager.add_weight_only_config(
         regex='.*/Dense/.*',
         operation_name=_TFLOpName.FULLY_CONNECTED,
-        num_bits=4,
+        num_bits=num_bits,
     )
     alg_key, op_config = self._recipe_manager.get_quantization_configs(
         _TFLOpName.FULLY_CONNECTED, 'model/Dense/op'
@@ -330,6 +331,72 @@ class ConfiguratorTest(parameterized.TestCase, googletest.TestCase):
     self.assertIsNone(op_config.activation_tensor_config)
     weight_tensor_config = op_config.weight_tensor_config
     self.assertIsNotNone(weight_tensor_config)
+    self.assertEqual(weight_tensor_config.num_bits, num_bits)
+    self.assertTrue(weight_tensor_config.symmetric)
+    self.assertEqual(
+        weight_tensor_config.granularity,
+        _QuantGranularity.CHANNELWISE,
+    )
+    self.assertEqual(weight_tensor_config.dtype, _TensorDataType.INT)
+
+  def test_add_weight_only_config_fp16(self):
+    self._recipe_manager.add_weight_only_config(
+        regex='.*/Dense2/.*',
+        operation_name=_TFLOpName.FULLY_CONNECTED,
+        num_bits=16,
+        algorithm_key=_AlgorithmName.FLOAT_CASTING,
+    )
+    alg_key, op_config = self._recipe_manager.get_quantization_configs(
+        _TFLOpName.FULLY_CONNECTED, 'model/Dense2/op'
+    )
+    self.assertEqual(alg_key, _AlgorithmName.FLOAT_CASTING)
+    self.assertEqual(op_config.compute_precision, _ComputePrecision.FLOAT)
+    self.assertTrue(op_config.explicit_dequantize)
+    self.assertIsNone(op_config.activation_tensor_config)
+    weight_tensor_config = op_config.weight_tensor_config
+    self.assertIsNotNone(weight_tensor_config)
+    self.assertEqual(weight_tensor_config.num_bits, 16)
+    self.assertTrue(weight_tensor_config.symmetric)
+    self.assertEqual(
+        weight_tensor_config.granularity,
+        _QuantGranularity.CHANNELWISE,
+    )
+    self.assertEqual(weight_tensor_config.dtype, _TensorDataType.FLOAT)
+
+  def test_add_weight_only_config_fp8_raise_error(self):
+    error_message = (
+        'float casting quantization config requires number of bits to be set'
+        ' as 16'
+    )
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, lambda err: error_message in str(err)
+    ):
+      self._recipe_manager.add_weight_only_config(
+          regex='.*/Dense2/.*',
+          operation_name=_TFLOpName.FULLY_CONNECTED,
+          num_bits=8,
+          algorithm_key=_AlgorithmName.FLOAT_CASTING,
+      )
+
+  def test_add_static_config(self):
+    self._recipe_manager.add_static_config(
+        regex='.*/Dense/.*',
+        operation_name=_TFLOpName.FULLY_CONNECTED,
+        activation_num_bits=8,
+        weight_num_bits=4,
+    )
+    alg_key, op_config = self._recipe_manager.get_quantization_configs(
+        _TFLOpName.FULLY_CONNECTED, 'model/Dense/op'
+    )
+    self.assertEqual(alg_key, _AlgorithmName.MIN_MAX_UNIFORM_QUANT)
+    self.assertEqual(op_config.compute_precision, _ComputePrecision.INTEGER)
+    self.assertFalse(op_config.explicit_dequantize)
+    activation_tensor_config = op_config.activation_tensor_config
+    self.assertIsNotNone(activation_tensor_config)
+    self.assertEqual(activation_tensor_config.num_bits, 8)
+    self.assertFalse(activation_tensor_config.symmetric)
+    weight_tensor_config = op_config.weight_tensor_config
+    self.assertIsNotNone(weight_tensor_config)
     self.assertEqual(weight_tensor_config.num_bits, 4)
     self.assertTrue(weight_tensor_config.symmetric)
     self.assertEqual(

{ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250928.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.4.0.dev20250926
+Version: 0.4.0.dev20250928
 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.4.0.dev20250926.dist-info → ai_edge_quantizer_nightly-0.4.0.dev20250928.dist-info}/RECORD

@@ -16,8 +16,8 @@ ai_edge_quantizer/qtyping.py,sha256=tfrPip-uzJuF_PASgUExx5Oy9gghWUbQaApR0XaBpNw,
 ai_edge_quantizer/quantizer.py,sha256=ckAEOnnBxuCKZuvlzdChevCKPuE-IeDPHCNtFTWr250,17857
 ai_edge_quantizer/quantizer_test.py,sha256=m6f4ayyaF3yQb9i4V0aFAbmGw0OKZ2Zam1RoTPh-u24,22917
 ai_edge_quantizer/recipe.py,sha256=MEkfQ2Sg3KAE9LAORHWcbjYNPg06EUbwc1d-VspQA2U,6461
-ai_edge_quantizer/recipe_manager.py,sha256=6dgbE-IZfEetzXH3p3Qm_9eQutNDOpZnMpiaLTbP-ZQ,14744
-ai_edge_quantizer/recipe_manager_test.py,sha256=H-B75vwPN5ND-nUa3pOXizeHTv4mufPiC5cL_OlDIYU,34040
+ai_edge_quantizer/recipe_manager.py,sha256=6l2uq8KL23KLu9OQDmPGkxrFiwHrdDB9xnn-ni8WdEM,15036
+ai_edge_quantizer/recipe_manager_test.py,sha256=qjgGUF-wggXnSXqZ5khmqrDMIQI5CShk52IVWTahq6s,36817
 ai_edge_quantizer/recipe_test.py,sha256=QisyaTol8JRZFcGOGyee7QRCvqj5VbF4guKWdIoMUOE,6213
 ai_edge_quantizer/transformation_instruction_generator.py,sha256=O0U2aZcB8aXQgOV8r9g1rGNzDUiuI5Ta53XnxZbVffE,31576
 ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=KW5-WoTTo9IqLEVnWxVC8ut8eWLi_91xfKgGqVQ9QDk,54635
@@ -72,8 +72,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EoVjI_hplX_Rml3hfRsGmQOi
 ai_edge_quantizer/utils/tfl_interpreter_utils_test.py,sha256=6fjkM-rycZ95L4yfvlr0TN6RlrhfPzxNUYrZaYO_F0A,12013
 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.4.0.dev20250926.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/METADATA,sha256=6ymhTobT9E998G5IZCfmysJbNYMcQr_vrngEtQf5VsE,1508
-ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.4.0.dev20250926.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.4.0.dev20250928.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.4.0.dev20250928.dist-info/METADATA,sha256=Vi0g_M6-Fk7O8Xk5h5AtzbeTEFzse_8cYiBZ1S-H330,1508
+ai_edge_quantizer_nightly-0.4.0.dev20250928.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.4.0.dev20250928.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.4.0.dev20250928.dist-info/RECORD,,