ai-edge-quantizer-nightly 0.1.0.dev20250424__py3-none-any.whl → 0.1.0.dev20250426__py3-none-any.whl

ai_edge_quantizer/transformation_performer.py

@@ -180,6 +180,38 @@ class TransformationPerformer:
             )
           transformation.tensor_id = trans_info.output_tensor_id
 
+  def _get_updated_producer_id(
+      self, original_producer_id: int, subgraph_id: int
+  ) -> int:
+    """Update the producer of a transformation instruction."""
+    if original_producer_id is None or original_producer_id < 0:
+      producer = -1
+    elif original_producer_id < len(self._original_op_id_map[subgraph_id]):
+      producer = self._original_op_id_map[subgraph_id][original_producer_id]
+    else:
+      # If the producer id is not in the original op map, it's an added op,
+      # go the added op map to find the producer.
+      producer = self._added_op_id_map[subgraph_id][
+          original_producer_id - len(self._original_op_id_map[subgraph_id])
+      ]
+    return producer
+
+  def _get_updated_consumer_ids(
+      self,
+      original_consumer_ids: list[int],
+      subgraph_id: int,
+  ) -> list[int]:
+    """Update the consumers of a transformation instruction."""
+    consumers = []
+    for original_op_id in original_consumer_ids:
+      new_consumer_id = (
+          -1
+          if original_op_id == -1
+          else self._original_op_id_map[subgraph_id][original_op_id]
+      )
+      consumers.append(new_consumer_id)
+    return consumers
+
   def _apply_single_transformation(
       self,
       transformation_inst: qtyping.TensorTransformationInsts,
@@ -198,28 +230,12 @@ class TransformationPerformer:
       None, update the transformation_inst & tflite_model in place
     """
     instruction = transformation_inst.instructions[transformation_index]
-    if not instruction.producer or instruction.producer < 0:
-      producer = -1
-    elif instruction.producer < len(
-        self._original_op_id_map[transformation_inst.subgraph_id]
-    ):
-      producer = self._original_op_id_map[transformation_inst.subgraph_id][
-          instruction.producer
-      ]
-    else:
-      # if the producer id is not in the original op map, it's an added op,
-      # go the corresponding new maps
-      producer = self._added_op_id_map[transformation_inst.subgraph_id][
-          instruction.producer
-          - len(self._original_op_id_map[transformation_inst.subgraph_id])
-      ]
-    consumers = []
-    for original_op_id in instruction.consumers:
-      consumers.append(
-          self._original_op_id_map[transformation_inst.subgraph_id][
-              original_op_id
-          ]
-      )
+    producer = self._get_updated_producer_id(
+        instruction.producer, transformation_inst.subgraph_id
+    )
+    consumers = self._get_updated_consumer_ids(
+        instruction.consumers, transformation_inst.subgraph_id
+    )
     trans_info = self._transformation_registration[instruction.transformation](
         transformation_utils.TransformationInput(
             instruction.tensor_id,
@@ -239,7 +255,12 @@ class TransformationPerformer:
     )
     self._update_op_id_map(
         transformation_inst.subgraph_id,
-        min(instruction.consumers),
+        # The added op must be right before the most immediate consumer, unless
+        # the consumer is the graph output (id=-1), then use the producer's
+        # index instead.
+        min(instruction.consumers)
+        if min(instruction.consumers) >= 0
+        else instruction.producer + 1,
         trans_info.num_ops_added,
     )
 

ai_edge_quantizer/transformation_performer_test.py

@@ -112,6 +112,32 @@ class TransformationPerformerTest(parameterized.TestCase):
     for index, op_id in enumerate(op_id_map[0]):
       self.assertEqual(op_id, index)
 
+  def test_update_op_id_map_not_changing_value_single_op_model(self):
+    """test for _update_op_id_map."""
+    model = tfl_flatbuffer_utils.read_model(
+        os.path.join(
+            TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias.tflite"
+        )
+    )
+    self._transformation_performer._create_op_id_map(model)
+    instruction = qtyping.TransformationInst(
+        transformation=qtyping.QuantTransformation.QUANTIZE_TENSOR,
+        tensor_id=0,
+        producer=0,
+        consumers=[-1],
+        parameters=qtyping.UniformQuantParams(
+            8, None, np.array([1]), np.array([0])
+        ),
+    )
+    producer = self._transformation_performer._get_updated_producer_id(
+        instruction.producer, 0
+    )
+    consumers = self._transformation_performer._get_updated_consumer_ids(
+        instruction.consumers, 0
+    )
+    self.assertEqual(producer, 0)
+    self.assertEqual(consumers, [-1])
+
   @parameterized.named_parameters(
       dict(
           testcase_name="test_no_update",
@@ -271,7 +297,7 @@ class TransformationPerformerTest(parameterized.TestCase):
         expected_added_op_id_map,
     )
 
-  def test__update_instructions_updates_tensor_id_after_duplicate_tensor(self):
+  def test_update_instructions_updates_tensor_id_after_duplicate_tensor(self):
     def get_test_instruction(transformation, consumers):
       return qtyping.TransformationInst(
           transformation=transformation,
@@ -325,6 +351,8 @@ class TransformationPerformerTest(parameterized.TestCase):
             tensor_name="sequential/conv2d/Relu;sequential/conv2d/BiasAdd;"
             + "sequential/conv2d/Conv2D;sequential/conv2d/BiasAdd/ReadVariableOp1",
             subgraph_id=0,
+            # Conv2d: op_id=0, output_tensor_id=7.
+            # This should add two sequential dequants after the conv2d.
             instructions=[
                 qtyping.TransformationInst(
                     transformation=qtyping.QuantTransformation.ADD_DEQUANTIZE,
@@ -349,6 +377,8 @@ class TransformationPerformerTest(parameterized.TestCase):
         "sequential/average_pooling2d/AvgPool": qtyping.TensorTransformationInsts(
             tensor_name="sequential/average_pooling2d/AvgPool",
             subgraph_id=0,
+            # Avg_pool: op_id=1, output_tensor_id=8.
+            # This should add two sequential dequants after the avg_pool.
             instructions=[
                 qtyping.TransformationInst(
                     transformation=qtyping.QuantTransformation.ADD_DEQUANTIZE,
@@ -376,19 +406,111 @@ class TransformationPerformerTest(parameterized.TestCase):
     )
     self.assertLen(self._test_model.subgraphs, 1)
     self.assertLen(self._test_model.subgraphs[0].operators, 10)
+    # The original model has 13 tensors, each dequant adds 1 tensor.
     self.assertLen(self._test_model.subgraphs[0].tensors, 17)
+    # Check that the dequant opcode is added to the model.
     self.assertEqual(
         self._test_model.subgraphs[0].operators[1].opcodeIndex,
         len(self._test_model.operatorCodes) - 1,
     )
+    # Conv2d, dequant, dequant, avgpool, dequant, dequant, etc.
+    expected_builtin_op_order = [3, 6, 6, 1, 6, 6, 22, 9, 9, 25]
+    for i, op in enumerate(self._test_model.subgraphs[0].operators):
+      op_code = self._test_model.operatorCodes[op.opcodeIndex].builtinCode
+      self.assertEqual(op_code, expected_builtin_op_order[i])
+    # Check that the first dequant input is connected to the conv2d output.
+    self.assertEqual(self._test_model.subgraphs[0].operators[1].inputs[0], 7)
+    # Output is a new tensor just added.
+    self.assertEqual(self._test_model.subgraphs[0].operators[1].outputs[0], 13)
+    # Second dequant has new tensors.
     self.assertEqual(self._test_model.subgraphs[0].operators[2].inputs[0], 13)
     self.assertEqual(self._test_model.subgraphs[0].operators[2].outputs[0], 14)
-    self.assertEqual(
-        self._test_model.subgraphs[0].operators[2].outputs[0],
-        self._test_model.subgraphs[0].operators[3].inputs[0],
-    )
+    # Avgpool's input is second dequant's output.
+    self.assertEqual(self._test_model.subgraphs[0].operators[3].inputs[0], 14)
+    # Avgpool's output remains the same.
     self.assertEqual(self._test_model.subgraphs[0].operators[3].outputs[0], 8)
+    # Third dequant's output is a new tensor.
     self.assertEqual(self._test_model.subgraphs[0].operators[4].outputs[0], 15)
+    # Fourth dequant.
+    self.assertEqual(self._test_model.subgraphs[0].operators[5].inputs[0], 15)
+    self.assertEqual(self._test_model.subgraphs[0].operators[5].outputs[0], 16)
+
+    # Avgpool (op_id=1) and reshape (op_id=2) are bumped by 2 due to the two
+    # dequants added after it.
+    expected_op_id_map = [0, 3, 6, 7, 8, 9]
+    self.assertEqual(
+        self._transformation_performer._original_op_id_map[0],
+        expected_op_id_map,
+    )
+    # New dequants are added at these indices.
+    expected_added_op_id_map = [1, 2, 4, 5]
+    self.assertEqual(
+        self._transformation_performer._added_op_id_map[0],
+        expected_added_op_id_map,
+    )
+
+  def test_op_insertion_at_input_and_output(self):
+    """test for _update_op_id_map."""
+    model = tfl_flatbuffer_utils.read_model(
+        os.path.join(
+            TEST_DATA_PREFIX_PATH, "tests/models/single_fc_bias.tflite"
+        )
+    )
+    self._transformation_performer._create_op_id_map(model)
+    instructions = {
+        # Fully_connected: op_id=0, input_tensor_id=0, output_tensor_id=3.
+        # Add a new quantize op to the input of the fully_connected.
+        "serving_default_input_2:0": qtyping.TensorTransformationInsts(
+            tensor_name="serving_default_input_2:0",
+            subgraph_id=0,
+            instructions=[
+                qtyping.TransformationInst(
+                    transformation=qtyping.QuantTransformation.ADD_QUANTIZE,
+                    tensor_id=0,
+                    producer=-1,
+                    consumers=[0],
+                    parameters=qtyping.UniformQuantParams(
+                        8, None, np.array([1]), np.array([0])
+                    ),
+                ),
+            ],
+        ),
+        # Add a new dequantize op to the output of the fully_connected.
+        "StatefulPartitionedCall:0": qtyping.TensorTransformationInsts(
+            tensor_name="StatefulPartitionedCall:0",
+            subgraph_id=0,
+            instructions=[
+                qtyping.TransformationInst(
+                    transformation=qtyping.QuantTransformation.ADD_DEQUANTIZE,
+                    tensor_id=3,
+                    producer=0,
+                    consumers=[-1],
+                    parameters=qtyping.UniformQuantParams(
+                        8, None, np.array([1]), np.array([0])
+                    ),
+                ),
+            ],
+        ),
+    }
+    self._transformation_performer.transform_graph(instructions, model)
+
+    # Original fc (op_id=0) should be bumped to op_id=1.
+    self.assertEqual(
+        self._transformation_performer._original_op_id_map[0],
+        [1],
+    )
+    # New quantize added at op_id=0, dequantize added at op_id=1.
+    expected_added_op_id_map = [0, 2]
+    self.assertEqual(
+        self._transformation_performer._added_op_id_map[0],
+        expected_added_op_id_map,
+    )
+    # Quantize, fully_connected, dequantize.
+    expected_builtin_op_order = [114, 9, 6]
+    for i, op in enumerate(model.subgraphs[0].operators):
+      op_code = model.operatorCodes[op.opcodeIndex].builtinCode
+      self.assertEqual(op_code, expected_builtin_op_order[i])
+
 
 if __name__ == "__main__":
   googletest.main()

{ai_edge_quantizer_nightly-0.1.0.dev20250424.dist-info → ai_edge_quantizer_nightly-0.1.0.dev20250426.dist-info}/METADATA

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

@@ -21,8 +21,8 @@ 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=R7A90Qj6iQQROrznXmXLJd-5yXq0PRHbLOdNY51dEu4,27913
 ai_edge_quantizer/transformation_instruction_generator_test.py,sha256=E0QSDCav6N6izlJ-a1ZJOsb2VEUxuxBmTbt0-EgDdxY,49890
-ai_edge_quantizer/transformation_performer.py,sha256=PIrylVhuWZCpnXEl7qSw2BlxRrY7lqj6aQvagJVCVts,11989
-ai_edge_quantizer/transformation_performer_test.py,sha256=n9xI6QMqvrj9KUul2LuObIsF7YdLSqgMg4X6d4BkFP8,15219
+ai_edge_quantizer/transformation_performer.py,sha256=zAzrQOb2n2IpB3qopmKV59e5E99HmTOL60QTCn9-7kA,12821
+ai_edge_quantizer/transformation_performer_test.py,sha256=xk6A3LStCyPclN51--9uO7XjSxNfZmpdfvrzOL0maNM,20349
 ai_edge_quantizer/algorithms/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/nonlinear_quantize/__init__.py,sha256=lpq1g2ayg3lCPLy79t2VicYcnGKw64FfYIj1V7J-4m8,676
 ai_edge_quantizer/algorithms/nonlinear_quantize/float_casting.py,sha256=Bs9CK7wZAw6jNaZ8xEtbwO2vM34VYXNZSMVWvxJo9nw,9297
@@ -66,8 +66,8 @@ ai_edge_quantizer/utils/tfl_interpreter_utils.py,sha256=x2xA2CFPpe_2trcV8v5xGaBE
 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.1.0.dev20250424.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ai_edge_quantizer_nightly-0.1.0.dev20250424.dist-info/METADATA,sha256=ymVF3awwYNfrNqBMSN903Tnc_Catt8qj7xvktKDsnoU,1527
-ai_edge_quantizer_nightly-0.1.0.dev20250424.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_quantizer_nightly-0.1.0.dev20250424.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
-ai_edge_quantizer_nightly-0.1.0.dev20250424.dist-info/RECORD,,
+ai_edge_quantizer_nightly-0.1.0.dev20250426.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ai_edge_quantizer_nightly-0.1.0.dev20250426.dist-info/METADATA,sha256=CIV-3K_joKQSZc9qpwHgYbFHYwaAtPCZWE6yEYWgDkc,1527
+ai_edge_quantizer_nightly-0.1.0.dev20250426.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_quantizer_nightly-0.1.0.dev20250426.dist-info/top_level.txt,sha256=8QTfPnFXNVUhScFLaa-NWZMFWMn72M50DVPubpwWB1g,18
+ai_edge_quantizer_nightly-0.1.0.dev20250426.dist-info/RECORD,,