Support custom op and dynamic shape in tensor creation

PiperOrigin-RevId: 752821626

Author: paulinesho

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 @@ def setUp(self):
           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 @@ def test_add_new_activation_tensor_to_subgraph(
         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()