ai-edge-quantizer-nightly 0.1.0.dev20250514__py3-none-any.whl → 0.1.0.dev20250516__py3-none-any.whl

ai_edge_quantizer/transformations/duplicate_buffer.py

@@ -34,9 +34,10 @@ def duplicate_buffer(
         f' Tensor {tensor_name} is not constant.'
     )
 
-  duplicated_buffer_id = transformation_utils.add_new_constant_buffer(
+  duplicated_buffer_id = transformation_utils.get_constant_buffer(
       data=buffer_data,
       buffers=transformation_input.buffers,
+      force_duplicate_buffer=True,
   )
   tensor.buffer = duplicated_buffer_id
 

ai_edge_quantizer/transformations/duplicate_tensor.py

@@ -41,6 +41,7 @@ def duplicate_tensor(
       tensor_shape=tensor.shape,
       subgraph=subgraph,
       buffers=transformation_input.buffers,
+      force_duplicate_buffer=True,
   )
   # Update the tensor name to avoid name collision in case when tensor is
   # duplicated mulitple times.

ai_edge_quantizer/transformations/transformation_utils.py

@@ -87,21 +87,32 @@ def add_op_code(
   return len(model_op_codes) - 1
 
 
-def add_new_constant_buffer(
+def get_constant_buffer(
     data: np.ndarray,
     buffers: list[schema_py_generated.BufferT],
+    force_duplicate_buffer: bool = False,
 ) -> int:
-  """Add a new constant buffer to the model.
+  """Get the index of the constant buffer that contains the given data.
+
+  creating new buffer if provided data is not found in buffers list.
 
   Args:
     data: The data of the new tensor.
     buffers: The buffers of the model.
+    force_duplicate_buffer: Whether to add a new buffer even if the same buffer
+      already exists.
 
   Returns:
     The index of the new buffer in the model.
   """
+  new_data = np.frombuffer(data.tobytes(), dtype=np.uint8).flatten()
+  # TODO: b/417811116 - we should make this more efficient.
+  if not force_duplicate_buffer:
+    for index, buffer in enumerate(buffers):
+      if np.array_equal(buffer.data, new_data):
+        return index
   new_buffer = schema_py_generated.BufferT()
-  new_buffer.data = np.frombuffer(data.tobytes(), dtype=np.uint8).flatten()
+  new_buffer.data = new_data
   new_buffer.offset = 0
   new_buffer.size = 0
   new_buffer_id = len(buffers)
@@ -117,6 +128,7 @@ def add_new_constant_tensor(
     subgraph: schema_py_generated.SubGraphT,
     buffers: list[schema_py_generated.BufferT],
     tensor_shape: Optional[list[int]] = None,
+    force_duplicate_buffer: bool = False,
 ) -> int:
   """Add a new constant tensor to the model.
 
@@ -128,11 +140,13 @@ def add_new_constant_tensor(
     buffers: The buffers of the model.
     tensor_shape: The shape of the new tensor. If not provided, the shape of the
       data will be used.
+    force_duplicate_buffer: Whether to add a new buffer even if the same buffer
+      already exists.
 
   Returns:
     The index of the new tensor in the subgraph.
   """
-  new_buffer_id = add_new_constant_buffer(data, buffers)
+  new_buffer_id = get_constant_buffer(data, buffers, force_duplicate_buffer)
 
   new_tensor = schema_py_generated.TensorT()
   if tensor_shape is None:

ai_edge_quantizer/transformations/transformation_utils_test.py

@@ -111,7 +111,7 @@ class TransformationUtilsTest(parameterized.TestCase):
   def test_add_new_constant_buffer(self, data):
     """Tests if the constant buffer is added to the model."""
     prev_num_buffers = len(self.model.buffers) - 1
-    new_buffer_idx = transformation_utils.add_new_constant_buffer(
+    new_buffer_idx = transformation_utils.get_constant_buffer(
         data=data,
         buffers=self.model.buffers,
     )

{ai_edge_quantizer_nightly-0.1.0.dev20250514.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250516.dist-info}/METADATA

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

@@ -46,9 +46,9 @@ ai_edge_quantizer/algorithms/utils/common_utils_test.py,sha256=zqapGEfYhjQWe9cNG
 ai_edge_quantizer/transformations/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/transformations/dequant_insert.py,sha256=sL1LHFVzBDSd9jgrzlHz38LWU0bwmVX7iBkaNcui0ts,3566
 ai_edge_quantizer/transformations/dequant_insert_test.py,sha256=NJ18PnG71_AvUPz3Cr_TmG6URMeBfa7IiDDyddfTkKQ,10830
-ai_edge_quantizer/transformations/duplicate_buffer.py,sha256=sEod0EtmcHX0VDqBCI4BYCX9CSRyDtx2vmjtOentFiY,1743
+ai_edge_quantizer/transformations/duplicate_buffer.py,sha256=TvTHbm24IiICNkWOlvR2UpJKMU-88puNFycDYc0_ehQ,1774
 ai_edge_quantizer/transformations/duplicate_buffer_test.py,sha256=YYWl3Q5WF60s8T8pLzzA8TCSxz-i7dqc03dJt1LtMw4,3880
-ai_edge_quantizer/transformations/duplicate_tensor.py,sha256=HF1uuKFm5kFF6X0XUpdYlLPoikSRd7pIPK1oxN7TuHY,2455
+ai_edge_quantizer/transformations/duplicate_tensor.py,sha256=WKhf2LIAL0MnZe88b6942A37lvHXe1cFjUDqE5VNmvU,2490
 ai_edge_quantizer/transformations/duplicate_tensor_test.py,sha256=s-RqSxNBMfVJyCunXz2eb7-KA6UiBmbOmL7phLslENQ,5056
 ai_edge_quantizer/transformations/emulated_subchannel.py,sha256=HVaRxoC8PCAvy3xeMv3OIymukUy_yW1zK0xN8Ann6I4,13602
 ai_edge_quantizer/transformations/emulated_subchannel_test.py,sha256=gZP6u9NdPXl7s19qB_Un8evou9ZZV6I9Gy0E1rdobHM,7722
@@ -58,8 +58,8 @@ 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=Hc1jrY3cEUooiTu9qOh4jxyZp58vrokKxzTmzx6V70c,5853
-ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=E90O4PYSjzGdHhaNvm3ii0Xom3cyFfcqQyYjOhYzG-c,8702
+ai_edge_quantizer/transformations/transformation_utils.py,sha256=mYD6gHxSO6JmjC9fSWl0pLh6N_tCW70yTIxHYdmnV3A,6483
+ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=MWgq29t7rvxRQIfi4ny9IoODFCTcbpjnIwoCL40zDKk,8698
 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
@@ -70,8 +70,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=WoewyiZpaua80oP0tpgyrw5W
 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.dev20250514.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250514.dist-info/METADATA,sha256=jGKo0MakT_DI_wLdclRVhNbLFRzJ2GDmQX7UGLDsr4I,1528
-ai_edge_quantizer_nightly-0.1.0.dev20250514.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250514.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250514.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250516.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250516.dist-info/METADATA,sha256=iP17vZdvF02q2glxd2bcRcQLcwle6B9KLQcsH4rFrog,1528
+ai_edge_quantizer_nightly-0.1.0.dev20250516.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250516.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250516.dist-info/RECORD,,