ai-edge-quantizer-nightly 0.3.0.dev20250724__py3-none-any.whl → 0.3.0.dev20250726__py3-none-any.whl

ai_edge_quantizer/algorithm_manager.py

@@ -117,6 +117,8 @@ MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT = {
     _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,
 }
 for op_name, materialize_func in MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT.items():
   register_quantized_op(
@@ -260,6 +262,8 @@ _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT = immutabledict({
     _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,
 })
 
 for op_name, materialize_func in _OCTAV_OP_NAME_MATERIALIZE_FUNC_DICT.items():

ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py

@@ -809,6 +809,38 @@ def materialize_div(
   )
 
 
+def materialize_broadcast_to(
+    get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: dict[str, Any],
+) -> list[qtyping.TensorTransformationParams]:
+  """Materialize tensors in tfl.broadcast_to."""
+  return common_utils.materialize_standard_op(
+      op_info,
+      graph_info,
+      tensor_name_to_qsv,
+      get_tensor_quant_params_fn,
+      constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
+      inputs_to_ignore=[1],  # Shape tensor does not need to be quantized.
+  )
+
+
+def materialize_sqrt(
+    get_tensor_quant_params_fn: qtyping.GetTensorQuantParamsFuncSignature,
+    op_info: qtyping.OpInfo,
+    graph_info: qtyping.GraphInfo,
+    tensor_name_to_qsv: dict[str, Any],
+) -> list[qtyping.TensorTransformationParams]:
+  """Materialize tensors in tfl.sqrt."""
+  return common_utils.materialize_standard_op(
+      op_info,
+      graph_info,
+      tensor_name_to_qsv,
+      get_tensor_quant_params_fn,
+  )
+
+
 def _get_tensor_shape_for_blockwise(
     tensor_shape: Sequence[int], quantized_dim: int, block_size: int
 ) -> list[int]:

ai_edge_quantizer/default_policy.py

@@ -190,7 +190,9 @@ DEFAULT_JSON_POLICY = """
       "GATHER_ND",
       "PACK",
       "UNPACK",
-      "DIV"
+      "DIV",
+      "BROADCAST_TO",
+      "SQRT"
     ],
     "static_wi8_ai8": [
       "ADD",
@@ -229,7 +231,9 @@ DEFAULT_JSON_POLICY = """
       "GATHER_ND",
       "PACK",
       "UNPACK",
-      "DIV"
+      "DIV",
+      "BROADCAST_TO",
+      "SQRT"
     ],
     "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],
     "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT", "EMBEDDING_LOOKUP"],

ai_edge_quantizer/qtyping.py

@@ -70,6 +70,8 @@ class TFLOperationName(str, enum.Enum):
   PACK = 'PACK'
   UNPACK = 'UNPACK'
   DIV = 'DIV'
+  BROADCAST_TO = 'BROADCAST_TO'
+  SQRT = 'SQRT'
 
 
 class QuantizeMode(enum.Enum):

ai_edge_quantizer/utils/test_utils.py

@@ -18,7 +18,7 @@
 import inspect as _inspect
 import os.path as _os_path
 import sys as _sys
-from typing import Union
+from typing import Optional, Union
 
 from absl.testing import parameterized
 
@@ -32,6 +32,7 @@ _OpName = qtyping.TFLOperationName
 _TensorQuantConfig = qtyping.TensorQuantizationConfig
 _OpQuantConfig = qtyping.OpQuantizationConfig
 _AlgorithmName = quantizer.AlgorithmName
+_Numeric = Union[int, float]
 
 
 DEFAULT_ACTIVATION_QUANT_SETTING = _TensorQuantConfig(
@@ -98,9 +99,9 @@ class BaseOpTestCase(parameterized.TestCase):
       op_name: _OpName,
       op_config: _OpQuantConfig,
       num_validation_samples: int = 4,
-      num_calibration_samples: Union[int, None] = None,
+      num_calibration_samples: Optional[int] = None,
       error_metric: str = 'mse',
-      int_min_max: Union[tuple[int, int], None] = None,
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
   ) -> model_validator.ComparisonResult:
     """Quantizes and validates the given model with the given configurations.
 
@@ -113,7 +114,7 @@ class BaseOpTestCase(parameterized.TestCase):
       num_calibration_samples: The number of samples to use for calibration. If
         None then it will be set to num_validation_samples * 8.
       error_metric: The error error_metric to use for validation.
-      int_min_max: The min and max of the integer input range.
+      min_max_range: The min and max of the input range.
 
     Returns:
       The comparison result of the validation.
@@ -131,7 +132,7 @@ class BaseOpTestCase(parameterized.TestCase):
       calibration_data = tfl_interpreter_utils.create_random_normal_input_data(
           quantizer_instance.float_model,
           num_samples=num_calibration_samples,
-          int_min_max=int_min_max,
+          min_max_range=min_max_range,
       )
       calibration_result = quantizer_instance.calibrate(calibration_data)
       quantization_result = quantizer_instance.quantize(calibration_result)
@@ -140,7 +141,7 @@ class BaseOpTestCase(parameterized.TestCase):
     test_data = tfl_interpreter_utils.create_random_normal_input_data(
         quantization_result.quantized_model,
         num_samples=num_validation_samples,
-        int_min_max=int_min_max,
+        min_max_range=min_max_range,
     )
     return quantizer_instance.validate(test_data, error_metric)
 
@@ -190,7 +191,7 @@ class BaseOpTestCase(parameterized.TestCase):
       expected_model_size_reduction: float,
       weight_tolerance: float = 1e-4,
       output_tolerance: float = 1e-4,
-      int_min_max: Union[tuple[int, int], None] = None,
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
   ):
     """Check if the quantization is successful and the result is valid."""
     validation_result = self.quantize_and_validate(
@@ -198,7 +199,7 @@ class BaseOpTestCase(parameterized.TestCase):
         algorithm_key=algorithm_key,
         op_name=op_name,
         op_config=op_config,
-        int_min_max=int_min_max,
+        min_max_range=min_max_range,
     )
     with self.subTest(name='ModelSizeReduction'):
       self.assert_model_size_reduction_above_min_pct(
@@ -220,9 +221,9 @@ class BaseOpTestCase(parameterized.TestCase):
       op_name: _OpName,
       op_config: _OpQuantConfig,
       num_validation_samples: int = 4,
-      num_calibration_samples: Union[int, None] = None,
+      num_calibration_samples: Optional[int] = None,
       output_tolerance: float = 1e-4,
-      int_min_max: Union[tuple[int, int], None] = None,
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
   ):
     """Checks if the output errors after quantization are within the tolerance."""
     validation_result = self.quantize_and_validate(
@@ -232,7 +233,7 @@ class BaseOpTestCase(parameterized.TestCase):
         num_calibration_samples=num_calibration_samples,
         op_name=op_name,
         op_config=op_config,
-        int_min_max=int_min_max,
+        min_max_range=min_max_range,
     )
     self.assert_output_errors_below_tolerance(
         validation_result, output_tolerance

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -64,6 +64,8 @@ TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
     _TFLOpName.PACK: schema.BuiltinOperator.PACK,
     _TFLOpName.UNPACK: schema.BuiltinOperator.UNPACK,
     _TFLOpName.DIV: schema.BuiltinOperator.DIV,
+    _TFLOpName.BROADCAST_TO: schema.BuiltinOperator.BROADCAST_TO,
+    _TFLOpName.SQRT: schema.BuiltinOperator.SQRT,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(

ai_edge_quantizer/utils/tfl_interpreter_utils.py

@@ -27,6 +27,8 @@ from tensorflow.python.platform import gfile  # pylint: disable=g-direct-tensorf
 
 DEFAULT_SIGNATURE_KEY = "serving_default"
 
+_Numeric = Union[int, float]
+
 
 def create_tfl_interpreter(
     tflite_model: Union[str, bytes],
@@ -329,6 +331,17 @@ def _create_random_normal(
   return rng.normal(size=shape).astype(dtype)
 
 
+def _create_random_uniform(
+    rng: np.random.Generator,
+    shape: tuple[int, ...],
+    dtype: np.dtype,
+    min_value: float = 0.0,
+    max_value: float = 1.0,
+) -> dict[str, Any]:
+  """Creates a random uniform dataset sample for given input details."""
+  return rng.uniform(min_value, max_value, size=shape).astype(dtype)
+
+
 def _create_random_integers(
     rng: np.random.Generator,
     shape: tuple[int, ...],
@@ -353,7 +366,7 @@ def create_random_dataset(
     input_details: dict[str, Any],
     num_samples: int,
     random_seed: Union[int, np._typing.ArrayLike],
-    int_min_max: Union[tuple[int, int], None] = None,
+    min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
 ) -> list[dict[str, Any]]:
   """Creates a random normal dataset for given input details.
 
@@ -361,7 +374,7 @@ def create_random_dataset(
     input_details: A dictionary of input details.
     num_samples: The number of samples to generate.
     random_seed: The random seed to use.
-    int_min_max: The min and max of the integer input range.
+    min_max_range: The min and max of the input range.
 
   Returns:
     A list of dictionaries, each containing a sample of input data (for all
@@ -375,15 +388,21 @@ def create_random_dataset(
       dtype = input_tensor["dtype"]
       shape = input_tensor["shape"]
       if dtype in (np.int32, np.int64):
-        if int_min_max is None:
+        if min_max_range is None:
           new_data = _create_random_integers(rng, shape, dtype)
         else:
-          min_value, max_value = int_min_max
+          min_value, max_value = min_max_range
           new_data = _create_random_integers(
               rng, shape, dtype, min_value, max_value
           )
       elif dtype in (np.float32, ml_dtypes.bfloat16):
-        new_data = _create_random_normal(rng, shape, dtype)
+        if min_max_range is None:
+          new_data = _create_random_normal(rng, shape, dtype)
+        else:
+          min_value, max_value = min_max_range
+          new_data = _create_random_uniform(
+              rng, shape, dtype, min_value, max_value
+          )
       elif dtype == np.bool:
         new_data = _create_random_bool(rng, shape, dtype)
       else:
@@ -397,7 +416,7 @@ def create_random_normal_input_data(
     tflite_model: Union[str, bytes],
     num_samples: int = 4,
     random_seed: int = 666,
-    int_min_max: Union[tuple[int, int], None] = None,
+    min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
 ) -> dict[str, list[dict[str, Any]]]:
   """Creates a random normal dataset for a signature runner.
 
@@ -405,7 +424,7 @@ def create_random_normal_input_data(
     tflite_model: TFLite model path or bytearray.
     num_samples: Number of input samples to be generated.
     random_seed: Random seed to be used for function.
-    int_min_max: The min and max of the integer input range.
+    min_max_range: The min and max of the input range.
 
   Returns:
     A list of inputs to the given interpreter, for a single interpreter we may
@@ -420,6 +439,9 @@ def create_random_normal_input_data(
     signature_runner = tfl_interpreter.get_signature_runner(signature_key)
     input_details = signature_runner.get_input_details()
     test_data[signature_key] = create_random_dataset(
-        input_details, num_samples, random_seed, int_min_max
+        input_details,
+        num_samples,
+        random_seed,
+        min_max_range,
     )
   return test_data

{ai_edge_quantizer_nightly-0.3.0.dev20250724.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250726.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-quantizer-nightly
-Version: 0.3.0.dev20250724
+Version: 0.3.0.dev20250726
 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.3.0.dev20250724.dist-info → ai_edge_quantizer_nightly-0.3.0.dev20250726.dist-info}/RECORD

@@ -1,18 +1,18 @@
 ai_edge_quantizer/__init__.py,sha256=4pFSkukSwahYyzwqia0yPRyz8TnFQfGRthVJhYpMWas,793
-ai_edge_quantizer/algorithm_manager.py,sha256=776VBmJerYII1TYwauxjg5zgUuY38lcVMQy3joR1FG0,12680
+ai_edge_quantizer/algorithm_manager.py,sha256=wgC3g7hHvEM1fXARQsT3UgR5YLJqdQ4BLIPicn_bTvM,12932
 ai_edge_quantizer/algorithm_manager_api.py,sha256=u903TG0s1uIDhJqfeJne3CFl8A93phZrwgV2-hwdcXU,9247
 ai_edge_quantizer/algorithm_manager_api_test.py,sha256=w6bSONvXkX6bzXAGc0-7b6gNDt9oz9ieq97KP8Sg_JU,7666
 ai_edge_quantizer/calibrator.py,sha256=Sms7_AIHPH9G5xFaz5Ef3a5gPhxuIWQI8d2LUM8C96I,12071
 ai_edge_quantizer/calibrator_test.py,sha256=ejKc5YC7id8J1Ll9HAYCzMnKzxd0FUENSD06zkSSV0c,11900
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
-ai_edge_quantizer/default_policy.py,sha256=JkvwaQEqvBh1BKBNwS179wbgJp7C5M-bkhSJoN_wDZQ,11366
+ai_edge_quantizer/default_policy.py,sha256=djOEFPStjcDLoqNwK4RH_lfWJmdCLLixhCLwa3mN8pQ,11438
 ai_edge_quantizer/model_modifier.py,sha256=teGa8I6kGvn6TQY6Xv53YFIc_pQEhNvM9Zb4bvhezyw,7110
 ai_edge_quantizer/model_modifier_test.py,sha256=cJd04SLOG-fQZZNZPcisoBLx3cLtWEwGqUBbLb-pif4,4751
 ai_edge_quantizer/model_validator.py,sha256=Hj0_5o-Oa3dSlJ3ryVjRhvsyelHNyek1GrtG9buMczg,13153
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
 ai_edge_quantizer/params_generator.py,sha256=hcgMHJlERZERUyIAEi6AHJcLJ8gsKIBAEojzFFz-tqk,20098
 ai_edge_quantizer/params_generator_test.py,sha256=RDYoRZDJfEZRtjlTAU2kZ_4t3JHOqEHxfJX9V4ETAhg,40597
-ai_edge_quantizer/qtyping.py,sha256=J81aZXCe5jUDsQLP4T6slCmVIfN1R-sduGBXOLGuaRA,16654
+ai_edge_quantizer/qtyping.py,sha256=8sCBPI3IuIHaT4NrMZrLH6Hp_fkrVP8NEQC5zvHG-UU,16702
 ai_edge_quantizer/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
 ai_edge_quantizer/quantizer_test.py,sha256=K_HBA56JkFI3HL8VLWCqGEfC0ISh5ldMKoNyBdGRAJg,20368
 ai_edge_quantizer/recipe.py,sha256=FR0uJceumZrnle2VRSOQZ1uXup4S1cTYKRH-N53mWRo,2919
@@ -28,7 +28,7 @@ ai_edge_quantizer/algorithms/nonlinear_quantize/__init__.py,sha256=lpq1g2ayg3lCP
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting.py,sha256=Bs9CK7wZAw6jNaZ8xEtbwO2vM34VYXNZSMVWvxJo9nw,9297
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting_test.py,sha256=EqIHGEZ1LgUrTN7zf880RuAzEv3Qy7kgh5ivObJGHSo,22646
 ai_edge_quantizer/algorithms/uniform_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
-ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=LH6tBhgWeZWwhSL3hG8kUgh5bmtHAs_issHFsEwagkM,32221
+ai_edge_quantizer/algorithms/uniform_quantize/common_quantize.py,sha256=BUrGoC1TU6hD0QzqjblIs56Il7PKRfaz6s1G6nTXoio,33239
 ai_edge_quantizer/algorithms/uniform_quantize/common_quantize_test.py,sha256=GGf_n3wIeg3GB_eGsmyNJ0fTcxgpeMMbugTMRONK6TQ,3553
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery.py,sha256=BDdn_uBZakfHyzdMJPKadsOqxqyC-s6W2ZzFH99L4fE,8652
 ai_edge_quantizer/algorithms/uniform_quantize/dequantized_weight_recovery_test.py,sha256=sT5eX5TLZEHTtPfnSkCPDlS0sQxlTFWbCsbvOuj--yY,8889
@@ -61,15 +61,15 @@ ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=MWgq29t7rv
 ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/utils/calibration_utils.py,sha256=e3dG7Nm94Ix0hkTWTWPUhEG6a8QR_cAM3PSwblfJV5g,15106
 ai_edge_quantizer/utils/calibration_utils_test.py,sha256=4BlksXl7b4yptL8xPR67hmJCnjhN9V10a2PunzfHrUE,9372
-ai_edge_quantizer/utils/test_utils.py,sha256=spqUmSNciOKPQHCBkHE7Zo34eMFq_BfBCAnMT3jAulU,8615
-ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=vCrFYE3lbnn4eP_L-P6ymjehV8KYpQyqQjoqn2OxOGo,11042
+ai_edge_quantizer/utils/test_utils.py,sha256=a4Nk-wbeB09dFjTDZiA0K67d26j5DD0UDH_GIVmVG_4,8685
+ai_edge_quantizer/utils/tfl_flatbuffer_utils.py,sha256=LPk8yWBjLt_saKobjAvtBR9q_Ets6-3HrfMxPt064Ig,11158
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=K1SbK8q92qYVtiVj0I0GtugsPTkpIpEKv9zakvFV_Sc,8555
-ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=SKSu1nqhGGzVei_DxmzXK-bbOE7G1vKnPDc5skce-yY,14322
+ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=EoVjI_hplX_Rml3hfRsGmQOihexmizeJqt4SQcET9aA,14925
 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.3.0.dev20250724.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.3.0.dev20250724.dist-info/METADATA,sha256=N9YhqRIdEJsB_5yhwxdqUa2mBk88wiILrZqjykffVCc,1528
-ai_edge_quantizer_nightly-0.3.0.dev20250724.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.3.0.dev20250724.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.3.0.dev20250724.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.3.0.dev20250726.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.3.0.dev20250726.dist-info/METADATA,sha256=e6TL7vVEzY0RsL4EadFawcW8PMwPQPQO6_16mmkqqjw,1528
+ai_edge_quantizer_nightly-0.3.0.dev20250726.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.3.0.dev20250726.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.3.0.dev20250726.dist-info/RECORD,,