ai-edge-quantizer-nightly 0.1.0.dev20250428__py3-none-any.whl → 0.1.0.dev20250430__py3-none-any.whl

ai_edge_quantizer/model_modifier.py

@@ -104,8 +104,8 @@ class ModelModifier:
         instructions, quantized_model, tensor_processing_order
     )
     constant_buffer_size = self._process_constant_map(quantized_model)
-    # we leave 64MB for the model architecture.
-    if constant_buffer_size > 2**31 - 2**26:
+    # we leave 256MB for the model architecture.
+    if constant_buffer_size > 2**31 - 2**28:
       return self._serialize_large_model(quantized_model)
     else:
       return self._serialize_small_model(quantized_model)

ai_edge_quantizer/transformations/transformation_utils.py

@@ -51,21 +51,39 @@ class TransformationInput:
 def add_op_code(
     op_code: schema_py_generated.OperatorCodeT,
     model_op_codes: list[schema_py_generated.OperatorCodeT],
+    custom_op_name: Optional[str] = None,
 ) -> int:
   """Add an op code into a model if it's not present.
 
   Args:
     op_code: The op code to be added.
     model_op_codes: The op codes of the model.
+    custom_op_name: The custom string of the op code. If None, the op code will
+      be added as a builtin op code.
 
   Returns:
     The index of the op code in the model.
   """
+  if (
+      op_code == schema_py_generated.BuiltinOperator.CUSTOM
+      and custom_op_name is None
+  ):
+    raise ValueError('Custom string is required for custom op code.')
+
   for i, model_op_code in enumerate(model_op_codes):
+    # If the model already has the op code, just return the index.
     if model_op_code.builtinCode == op_code:
-      return i
+      if custom_op_name is not None:
+        if model_op_code.customCode == custom_op_name:
+          return i
+      else:
+        # Built-in op
+        return i
+
   model_op_codes.append(schema_py_generated.OperatorCodeT())
   model_op_codes[-1].builtinCode = op_code
+  if custom_op_name is not None:
+    model_op_codes[-1].customCode = custom_op_name
   return len(model_op_codes) - 1
 
 
@@ -146,7 +164,14 @@ def add_new_activation_tensor(
     The index of the new tensor in the subgraph.
   """
   new_tensor = schema_py_generated.TensorT()
-  new_tensor.shape = shape
+  # If there's a dynamic shape, we need to read from the shapeSignature field
+  # instead of shape. Shape should contain just 1 for the dynamic dimension but
+  # shapeSignature should contain the true shape.
+  if -1 in shape:
+    new_tensor.shapeSignature = shape
+    new_tensor.shape = [1 if i == -1 else i for i in shape]
+  else:
+    new_tensor.shape = shape
   new_tensor.type = tensor_type
   new_tensor.name = tensor_name
   new_tensor.buffer = 0

ai_edge_quantizer/transformations/transformation_utils_test.py

@@ -41,19 +41,62 @@ class TransformationUtilsTest(parameterized.TestCase):
           testcase_name="add_new_op_code",
           op_code=schema_py_generated.BuiltinOperator.LOGISTIC,
           expected=1,
+          custom_op_name=None,
       ),
       dict(
           testcase_name="add_existing_op_code",
           op_code=schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
           expected=0,
+          custom_op_name=None,
+      ),
+      dict(
+          testcase_name="add_new_custom_op_code",
+          op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+          expected=1,
+          custom_op_name="random_new_custom_op",
       ),
   )
-  def test_add_op_code(self, op_code, expected):
+  def test_add_op_code(self, op_code, expected, custom_op_name):
     """Tests if the op code is added to the model."""
     got = transformation_utils.add_op_code(
-        op_code=op_code, model_op_codes=self.model.operatorCodes
+        op_code=op_code,
+        model_op_codes=self.model.operatorCodes,
+        custom_op_name=custom_op_name,
     )
     self.assertEqual(expected, got)
+    if custom_op_name is not None:
+      self.assertEqual(self.model.operatorCodes[got].customCode, custom_op_name)
+
+  def test_add_custom_op_code_without_op_string_raises_error(self):
+    with self.assertRaisesRegex(ValueError, "Custom string is required"):
+      transformation_utils.add_op_code(
+          op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+          model_op_codes=self.model.operatorCodes,
+          custom_op_name=None,
+      )
+
+  def test_add_two_custom_op_codes(self):
+    custom_op_name = "random_new_custom_op"
+    added_index = transformation_utils.add_op_code(
+        op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+        model_op_codes=self.model.operatorCodes,
+        custom_op_name=custom_op_name,
+    )
+    self.assertEqual(1, added_index)
+    self.assertEqual(
+        self.model.operatorCodes[added_index].customCode, custom_op_name
+    )
+
+    custom_op_name_2 = "random_new_custom_op_2"
+    added_index = transformation_utils.add_op_code(
+        op_code=schema_py_generated.BuiltinOperator.CUSTOM,
+        model_op_codes=self.model.operatorCodes,
+        custom_op_name=custom_op_name_2,
+    )
+    self.assertEqual(2, added_index)
+    self.assertEqual(
+        self.model.operatorCodes[added_index].customCode, custom_op_name_2
+    )
 
   @parameterized.named_parameters(
       dict(
@@ -189,6 +232,25 @@ class TransformationUtilsTest(parameterized.TestCase):
         self.model.subgraphs[0].tensors[-1].shape,
     )
 
+  def test_add_new_activation_tensor_with_dynamic_shape(self):
+    """Tests adding an activation tensor with dynamic shape."""
+    subgraph = self.model.subgraphs[0]
+    new_id = transformation_utils.add_new_activation_tensor(
+        tensor_name="test_tensor",
+        shape=[1, -1, -1, 1],
+        tensor_type=schema_py_generated.TensorType.FLOAT32,
+        subgraph=subgraph,
+    )
+    # Originally had 4 tensors, new tensor is added at index 4.
+    self.assertEqual(new_id, 4)
+    self.assertLen(subgraph.tensors, 5)
+    self.assertEqual(subgraph.tensors[-1].name, "test_tensor")
+    self.assertEqual(
+        subgraph.tensors[-1].type, schema_py_generated.TensorType.FLOAT32
+    )
+    self.assertEqual(subgraph.tensors[-1].shape, [1, 1, 1, 1])
+    self.assertEqual(subgraph.tensors[-1].shapeSignature, [1, -1, -1, 1])
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/utils/tfl_interpreter_utils.py

@@ -319,7 +319,27 @@ def get_signature_main_subgraph_index(
   return signature_runner._subgraph_index  # pylint:disable=protected-access
 
 
-def create_random_normal_dataset(
+def _create_random_normal(
+    rng: np.random.Generator,
+    shape: tuple[int, ...],
+    dtype: np.dtype,
+) -> dict[str, Any]:
+  """Creates a random normal dataset sample for given input details."""
+  return rng.normal(size=shape).astype(dtype)
+
+
+def _create_random_integers(
+    rng: np.random.Generator,
+    shape: tuple[int, ...],
+    dtype: np.dtype,
+    min_value: int = 0,
+    max_value: int = 1024,
+) -> dict[str, Any]:
+  """Creates a random integer dataset sample for given input details."""
+  return rng.integers(min_value, max_value, size=shape, dtype=dtype)
+
+
+def create_random_dataset(
     input_details: dict[str, Any],
     num_samples: int,
     random_seed: Union[int, np._typing.ArrayLike],
@@ -340,9 +360,14 @@ def create_random_normal_dataset(
   for _ in range(num_samples):
     input_data = {}
     for arg_name, input_tensor in input_details.items():
-      new_data = rng.normal(size=input_tensor["shape"]).astype(
-          input_tensor["dtype"]
-      )
+      dtype = input_tensor["dtype"]
+      shape = input_tensor["shape"]
+      if dtype in (np.int32, np.int64):
+        new_data = _create_random_integers(rng, shape, dtype)
+      elif dtype == np.float32:
+        new_data = _create_random_normal(rng, shape, dtype)
+      else:
+        raise ValueError(f"Unsupported dtype: {input_tensor['dtype']}")
       input_data[arg_name] = new_data
     dataset.append(input_data)
   return dataset
@@ -372,7 +397,7 @@ def create_random_normal_input_data(
   for signature_key in signature_keys:
     signature_runner = tfl_interpreter.get_signature_runner(signature_key)
     input_details = signature_runner.get_input_details()
-    test_data[signature_key] = create_random_normal_dataset(
+    test_data[signature_key] = create_random_dataset(
         input_details, num_samples, random_seed
     )
   return test_data

ai_edge_quantizer/utils/tfl_interpreter_utils_test.py

@@ -19,7 +19,6 @@ from tensorflow.python.platform import googletest
 from ai_edge_quantizer.utils import test_utils
 from ai_edge_quantizer.utils import tfl_interpreter_utils
 
-
 TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile("../tests/models")
 
 
@@ -159,7 +158,6 @@ class TflUtilsQuantizedModelTest(googletest.TestCase):
     signature_output = tfl_interpreter_utils.invoke_interpreter_signature(
         tfl_interpreter, self._signature_input_data
     )
-    print(signature_output)
     self.assertEqual(tuple(signature_output["dense_1"].shape), (1, 10))
 
     # Assert the input data is not modified in-place b/353340272.
@@ -328,5 +326,24 @@ class TflUtilsMultiSignatureModelTest(googletest.TestCase):
     self.assertEqual(multiply_output_content, [20.0])
 
 
+class TflUtilsIntegerInputModelTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(0)
+    self._test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "toy_model_with_kv_cache_multi_signature.tflite"
+    )
+
+  def test_random_integer_input_data(self):
+    test_data = tfl_interpreter_utils.create_random_normal_input_data(
+        self._test_model_path
+    )
+    self.assertEqual(test_data["signature_1"][0]["cache_0"].dtype, np.float32)
+    self.assertEqual(test_data["signature_1"][0]["cache_1"].dtype, np.float32)
+    self.assertEqual(test_data["signature_1"][0]["positions"].dtype, np.int32)
+    self.assertEqual(test_data["signature_1"][0]["tokens"].dtype, np.int32)
+
+
 if __name__ == "__main__":
   googletest.main()

{ai_edge_quantizer_nightly-0.1.0.dev20250428.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250430.dist-info}/METADATA

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

@@ -6,7 +6,7 @@ ai_edge_quantizer/calibrator.py,sha256=n7AD9j7UScR-CieoI6DQRMeiG_fhLBfSLRiM4460x
 ai_edge_quantizer/calibrator_test.py,sha256=C_oWOaRugPKYX74jF-eRFH-k6nGOdA8I9_uPiocaOuE,11900
 ai_edge_quantizer/conftest.py,sha256=SxCz-5LlRD_lQm4hQc4c6IGG7DS8d7IyEWY9gnscPN0,794
 ai_edge_quantizer/default_policy.py,sha256=81z4cruBK7mGFt8xFRZK5LKya65axuZwo2zpbcYSicc,11099
-ai_edge_quantizer/model_modifier.py,sha256=SPt9X-xBzRvcd4xIS24zLHt3aUS2QwsNDqweFqitCAo,7109
+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=fRNz0jO54cthPTibsCuViUXUuFRHl_fbvEiCukIVy20,13030
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
@@ -54,20 +54,20 @@ ai_edge_quantizer/transformations/quant_insert.py,sha256=jn6HsJaV-sqBiFPY-Aqbd64
 ai_edge_quantizer/transformations/quant_insert_test.py,sha256=X9ptPDvJCFkR5tejKnD1SlHFGPazQTW-wNNMV9MEAuw,10107
 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
+ai_edge_quantizer/transformations/transformation_utils.py,sha256=Hc1jrY3cEUooiTu9qOh4jxyZp58vrokKxzTmzx6V70c,5853
+ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=E90O4PYSjzGdHhaNvm3ii0Xom3cyFfcqQyYjOhYzG-c,8702
 ai_edge_quantizer/utils/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 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=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/tfl_interpreter_utils.py,sha256=WoewyiZpaua80oP0tpgyrw5Ws1v7f4vl88vdzS0UjDE,13490
+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.1.0.dev20250428.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250428.dist-info/METADATA,sha256=yyhIbc-7ZiZ6-UFWCpYx1LgbsoYfTxl7pnqkCiTGbA8,1527
-ai_edge_quantizer_nightly-0.1.0.dev20250428.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250428.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250428.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250430.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250430.dist-info/METADATA,sha256=_q1njPlZxzBVNALKPM-yvI1dmjSspbXKq8wWHLYitR4,1527
+ai_edge_quantizer_nightly-0.1.0.dev20250430.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250430.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250430.dist-info/RECORD,,