ai-edge-quantizer-nightly 0.0.1.dev20250311__py3-none-any.whl → 0.0.1.dev20250312__py3-none-any.whl

ai_edge_quantizer/qtyping.py

@@ -109,6 +109,8 @@ class QuantTransformation(enum.Enum):
   # Create pattern for emulated subchannel quantization, only support fully
   # connected op.
   EMULATED_SUBCHANNEL = 4
+  # Duplicate the buffer.
+  DUPLICATE_BUFFER = 5
 
 
 @dataclasses.dataclass(frozen=True)

ai_edge_quantizer/transformation_performer.py

@@ -18,6 +18,7 @@
 import numpy as np
 from ai_edge_quantizer import qtyping
 from ai_edge_quantizer.transformations import dequant_insert
+from ai_edge_quantizer.transformations import duplicate_buffer
 from ai_edge_quantizer.transformations import emulated_subchannel
 from ai_edge_quantizer.transformations import quant_insert
 from ai_edge_quantizer.transformations import quantize_tensor
@@ -68,6 +69,9 @@ class TransformationPerformer:
             emulated_subchannel.emulated_subchannel
         ),
         qtyping.QuantTransformation.ADD_QUANTIZE: quant_insert.insert_quant,
+        qtyping.QuantTransformation.DUPLICATE_BUFFER: (
+            duplicate_buffer.duplicate_buffer
+        ),
     }
     # transformations are seprated in two categories:
     # op_insertion_transformations are transformations that only insert ops
@@ -77,6 +81,7 @@ class TransformationPerformer:
         qtyping.QuantTransformation.ADD_DEQUANTIZE,
         qtyping.QuantTransformation.QUANTIZE_TENSOR,
         qtyping.QuantTransformation.ADD_QUANTIZE,
+        qtyping.QuantTransformation.DUPLICATE_BUFFER,
     ])
     self._op_replacement_transformations = set(
         [qtyping.QuantTransformation.EMULATED_SUBCHANNEL]

ai_edge_quantizer/transformations/duplicate_buffer.py

@@ -0,0 +1,45 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Duplicate buffer transformation."""
+
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+
+
+def duplicate_buffer(
+    transformation_input: transformation_utils.TransformationInput,
+) -> qtyping.TransformationInfo:
+  """Duplicates the buffer of the tensor."""
+  tensor_id = transformation_input.tensor_id
+  tensor = transformation_input.subgraph.tensors[tensor_id]
+  buffer_data = transformation_input.buffers[tensor.buffer].data
+  if buffer_data is None:
+    tensor_name = tfl_flatbuffer_utils.get_tensor_name(tensor)
+    raise ValueError(
+        'Duplicate Buffer transformation supports only constant tensors.'
+        f' Tensor {tensor_name} is not constant.'
+    )
+
+  duplicated_buffer_id = transformation_utils.add_new_constant_buffer(
+      data=buffer_data,
+      buffers=transformation_input.buffers,
+  )
+  tensor.buffer = duplicated_buffer_id
+
+  return qtyping.TransformationInfo(
+      op_id=0, num_ops_added=0, output_tensor_id=tensor_id
+  )

ai_edge_quantizer/transformations/duplicate_buffer_test.py

@@ -0,0 +1,106 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import os
+import numpy as np
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.transformations import duplicate_buffer
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+
+TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile('..')
+
+
+class DuplicateBufferTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, 'tests/models/weight_sharing_fcs.tflite'
+    )
+    self.model = tfl_flatbuffer_utils.read_model(model_path)
+
+  def _get_transformation_input(
+      self, subgraph_idx: int, tensor_idx: int
+  ) -> transformation_utils.TransformationInput:
+    return transformation_utils.TransformationInput(
+        tensor_id=tensor_idx,
+        buffers=self.model.buffers,
+        # Dummy params below.
+        op_codes=self.model.operatorCodes,
+        subgraph=self.model.subgraphs[subgraph_idx],
+        producer=-1,
+        consumers=[],
+        quant_params=qtyping.UniformQuantParams(
+            num_bits=8,
+            quantized_dimension=None,
+            scale=np.ones(1),
+            zero_point=np.zeros(1),
+        ),
+    )
+
+  def test_constant_buffer_is_correctly_duplicated(self):
+    # Duplicate the FC weight tensor in the second subgraph.
+    subgraph_idx = 1
+    subgraph = self.model.subgraphs[subgraph_idx]
+    weight_tensor_idx = 1
+    prev_buffer_id = subgraph.tensors[weight_tensor_idx].buffer
+    prev_num_buffers = len(self.model.buffers)
+    transformation_input = self._get_transformation_input(
+        subgraph_idx, weight_tensor_idx
+    )
+    transformation_info = duplicate_buffer.duplicate_buffer(
+        transformation_input
+    )
+    self.assertEqual(transformation_info.op_id, 0)
+    self.assertEqual(transformation_info.num_ops_added, 0)
+    self.assertEqual(transformation_info.output_tensor_id, 1)
+    # Check that a new buffer was added.
+    self.assertLen(self.model.buffers, prev_num_buffers + 1)
+    # Check that the new buffer is used by the weight tensor.
+    new_buffer_id = len(self.model.buffers) - 1
+    self.assertEqual(subgraph.tensors[weight_tensor_idx].buffer, new_buffer_id)
+    # Check that the new buffer has the same data as the original one.
+    self.assertTrue(
+        np.all(
+            np.frombuffer(
+                self.model.buffers[new_buffer_id].data,
+                dtype=np.float32,
+            )
+            == np.frombuffer(
+                self.model.buffers[prev_buffer_id].data,
+                dtype=np.float32,
+            )
+        )
+    )
+
+  def test_duplicate_buffer_raises_error_when_tensor_is_not_constant(self):
+    # Duplicate the FC input tensor in the second subgraph.
+    subgraph_idx = 1
+    weight_tensor_idx = 0
+    transformation_input = self._get_transformation_input(
+        subgraph_idx, weight_tensor_idx
+    )
+    with self.assertRaisesRegex(
+        ValueError,
+        'Duplicate Buffer transformation supports only constant tensors.',
+    ):
+      duplicate_buffer.duplicate_buffer(transformation_input)
+
+
+if __name__ == '__main__':
+  googletest.main()

ai_edge_quantizer/transformations/transformation_utils.py

@@ -69,6 +69,29 @@ def add_op_code(
   return len(model_op_codes) - 1
 
 
+def add_new_constant_buffer(
+    data: np.ndarray,
+    buffers: list[schema_py_generated.BufferT],
+) -> int:
+  """Add a new constant buffer to the model.
+
+  Args:
+    data: The data of the new tensor.
+    buffers: The buffers of the model.
+
+  Returns:
+    The index of the new buffer in the model.
+  """
+  new_buffer = schema_py_generated.BufferT()
+  new_buffer.data = np.frombuffer(data.tobytes(), dtype=np.uint8).flatten()
+  new_buffer.offset = 0
+  new_buffer.size = 0
+  new_buffer_id = len(buffers)
+  buffers.append(new_buffer)
+
+  return new_buffer_id
+
+
 def add_new_constant_tensor(
     tensor_name: str,
     data: np.ndarray,
@@ -88,16 +111,11 @@ def add_new_constant_tensor(
   Returns:
     The index of the new tensor in the subgraph.
   """
-  tensor_buffer = schema_py_generated.BufferT()
-  tensor_buffer.data = np.frombuffer(data.tobytes(), dtype=np.uint8).flatten()
-  tensor_buffer.offset = 0
-  tensor_buffer.size = 0
-  tensor_buffer_id = len(buffers)
-  buffers.append(tensor_buffer)
+  new_buffer_id = add_new_constant_buffer(data, buffers)
 
   new_tensor = schema_py_generated.TensorT()
   new_tensor.shape = data.shape
-  new_tensor.buffer = tensor_buffer_id
+  new_tensor.buffer = new_buffer_id
   new_tensor.type = tensor_type
   new_tensor.name = tensor_name
   new_tensor_id = len(subgraph.tensors)

ai_edge_quantizer/transformations/transformation_utils_test.py

@@ -55,6 +55,38 @@ class TransformationUtilsTest(parameterized.TestCase):
     )
     self.assertEqual(expected, got)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="float32",
+          data=np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
+      ),
+      dict(
+          testcase_name="int8",
+          data=np.array([[1, 2], [3, 4]], dtype=np.int8),
+      ),
+  )
+  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(
+        data=data,
+        buffers=self.model.buffers,
+    )
+    self.assertEqual(new_buffer_idx, prev_num_buffers + 1)
+
+    expected_buffer_data = (
+        np.frombuffer(
+            data.tobytes(),
+            dtype=np.uint8,
+        )
+        .flatten()
+        .tolist()
+    )
+    self.assertEqual(
+        self.model.buffers[new_buffer_idx].data.tolist(),
+        expected_buffer_data,
+    )
+
   @parameterized.named_parameters(
       dict(
           testcase_name="float32",

ai_edge_quantizer/utils/tfl_interpreter_utils.py

@@ -216,8 +216,10 @@ def get_tensor_name_to_details_map(
   """
   tensor_name_to_detail = {}
   for tensor_detail in tflite_interpreter.get_tensor_details(subgraph_index):
-    # Don't return temporary, unnamed tensors
-    if not tensor_detail["name"]:
+    # Don't return temporary, unnamed tensors or scratch tensors.
+    # tensor_detail doesn't include the allocation size (bytes) or an
+    # indicator of scratch tensors, so use the name to filter them out.
+    if not tensor_detail["name"] or "scratch" in tensor_detail["name"]:
       continue
     tensor_name_to_detail[tensor_detail["name"]] = tensor_detail
   return tensor_name_to_detail

{ai_edge_quantizer_nightly-0.0.1.dev20250311.dist-info → ai_edge_quantizer_nightly-0.0.1.dev20250312.dist-info}/METADATA

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

@@ -12,7 +12,7 @@ ai_edge_quantizer/model_validator.py,sha256=fRNz0jO54cthPTibsCuViUXUuFRHl_fbvEiC
 ai_edge_quantizer/model_validator_test.py,sha256=EeqOP_mrZsnZ3rug756s0ryDDqd2KgIDld5Lm_gDuWY,13020
 ai_edge_quantizer/params_generator.py,sha256=wT_TyW9jIaOfvLr9a1JR_kxwZtPXCvDFF0n_2QcrDZg,14087
 ai_edge_quantizer/params_generator_test.py,sha256=d9JwR-yxNJgg1SW-m8sFFPkIRdhgsDwMpVKsBQFL0gg,37658
-ai_edge_quantizer/qtyping.py,sha256=i9d4kG7bqR_5VTImfGY1JkK6M9J-DRfc_OGOc0qB5cQ,15186
+ai_edge_quantizer/qtyping.py,sha256=UBZ3HgO8IDLY6VJmO05rGtFv_idMD3Os3WWsnriA0NA,15235
 ai_edge_quantizer/quantizer.py,sha256=g3DMqFMrMpt9jQttCE0WcdNbMtk0JZnmN5MmCHrNdyM,13202
 ai_edge_quantizer/quantizer_test.py,sha256=38oTMJwMmxwPDeqT3eaVbazjtuIUIzMQ3mJNKh_eNQY,20493
 ai_edge_quantizer/recipe.py,sha256=r5tJiUs-ihZFzeK_jP2sUIUgTqZsL5SWvbUokuIUPDo,2251
@@ -21,7 +21,7 @@ ai_edge_quantizer/recipe_manager_test.py,sha256=LulVxsYp6TBGFI2PLCUCd4VsFq8ELpC7
 ai_edge_quantizer/recipe_test.py,sha256=Fg_sfxovI2fRjk5qdu18ghOvXdUvhDR1TxbE0GHDczc,3381
 ai_edge_quantizer/transformation_instruction_generator.py,sha256=WkECCO85lLs4cEnjZF5eVGbtuul4P8N77gUxUCK9ESY,21605
 ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=23MfOBiXv5Wq9FKJen7DrJ66T58qxf4ECriIY7V013k,39113
-ai_edge_quantizer/transformation_performer.py,sha256=3ZuaN5xwAu1B9zoGZHqi0__FOSzjGWOcydpCy1_Scac,10951
+ai_edge_quantizer/transformation_performer.py,sha256=4GCCOw5WBhO0UUfFoLlbfUh3RfnffF9qI_yEpdWlyps,11181
 ai_edge_quantizer/transformation_performer_test.py,sha256=m3V6nd6jsjd6jVId5wTBNuyDB2h2p4tHlMWhlnomlJo,13341
 ai_edge_quantizer/algorithms/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/nonlinear_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
@@ -42,26 +42,28 @@ 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_test.py,sha256=YYWl3Q5WF60s8T8pLzzA8TCSxz-i7dqc03dJt1LtMw4,3880
 ai_edge_quantizer/transformations/emulated_subchannel.py,sha256=HVaRxoC8PCAvy3xeMv3OIymukUy_yW1zK0xN8Ann6I4,13602
 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=6CyUFR7fGmzbS-mSuDlSSCJJGxY9X_WnCmEuKqL4LzQ,7864
 ai_edge_quantizer/transformations/quantize_tensor_test.py,sha256=QnJmQ_-XN5X0oR57FoY9bWGTp7migf11psbdO9R2pLg,9050
-ai_edge_quantizer/transformations/transformation_utils.py,sha256=BaKy5LYWgqli62XGo3AGRDNtHjwpBNp5VF5XgFbfVmg,4298
-ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=ks81nNvruOC88Tjdk3_qwku0V8p54p3gOqfObzNhWMM,5371
+ai_edge_quantizer/transformations/transformation_utils.py,sha256=R42OIbzwQ7JYJ-Qt46jsqwb6u4MfDGiIPCRZCUGLVCw,4664
+ai_edge_quantizer/transformations/transformation_utils_test.py,sha256=xH64SF3UHDh84vYbt-WvmXNjM-Jg-mefES1ACO1tkqw,6269
 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=_A-h_MqwElzjgkLDmXTZ1iAIWtTRcLjSFGfjNT8fuHU,10480
 ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py,sha256=AbyDxoM62k4ojD8gPdkWo--xe5hlX3t0kobQSA80kuk,7740
-ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=ZI16i25bAOpnJUBgRg38EH2CuZ55wxyHFScM12RsOwc,12487
+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.0.1.dev20250311.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.0.1.dev20250311.dist-info/METADATA,sha256=HQQRge7GnS1pcLK4Mbi0WMqmqXYSBliDlGu6Fg2Ihps,1528
-ai_edge_quantizer_nightly-0.0.1.dev20250311.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.0.1.dev20250311.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.0.1.dev20250311.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.0.1.dev20250312.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.0.1.dev20250312.dist-info/METADATA,sha256=ZRyM3EF8KAmJpo0QjIpL6Lv1qPQ_2tGLP4o_ARspijs,1528
+ai_edge_quantizer_nightly-0.0.1.dev20250312.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.0.1.dev20250312.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.0.1.dev20250312.dist-info/RECORD,,