Add custom operator names to the script decorator

onnxscript/_internal/irbuilder.py

@@ -77,8 +77,8 @@ def append_parameter(self, parameter: ir.Value | ir.Attr) -> None:
     def add_nested_function(self, fun: IRFunction) -> None:
         self.nested_functions[fun.name] = fun
 
-    def get_called_functions(self) -> dict[str, values.OnnxFunction]:
-        called_functions: dict[str, values.OnnxFunction] = {}
+    def get_called_functions(self) -> dict[ir.OperatorIdentifier, values.OnnxFunction]:
+        called_functions: dict[ir.OperatorIdentifier, values.OnnxFunction] = {}
 
         def visit(function_ir: IRFunction):
             for node in ir.traversal.RecursiveGraphIterator(function_ir.graph):
@@ -87,9 +87,16 @@ def visit(function_ir: IRFunction):
                     add(callee)
 
         def add(f: values.OnnxFunction):
-            if f.name in called_functions:
+            identifier = f.function_ir.identifier()
+            existing = called_functions.get(identifier)
+            if existing is f:
                 return
-            called_functions[f.name] = f
+            if existing is not None:
+                raise ValueError(
+                    "Multiple distinct ONNX functions use the same identifier "
+                    f"{identifier!r}. Give each function a unique domain, name, or overload."
+                )
+            called_functions[identifier] = f
             visit(f.function_ir)
 
         visit(self)

onnxscript/_internal/main.py

@@ -41,13 +41,17 @@ def script_check(
 def script(
     opset: Optional[values.Opset] = None,
     default_opset: Optional[values.Opset] = None,
+    *,
+    op_type: str | None = None,
     **kwargs: Any,
 ) -> Callable[[Callable[_P, _R]], onnxscript.OnnxFunction[_P, _R]]:
     """Main decorator. Declares a function as an onnx function.
 
     Args:
         opset: Opset the function belongs to (see :ref:`l-api-opsets`).
         default_opset: Opset to use for operators not in the function's opset.
+        op_type: Optional custom ONNX operator name. It must be a non-empty
+            string. When omitted or ``None``, the Python function name is used.
         kwargs: Additional keyword arguments.
 
     Returns:
@@ -77,6 +81,10 @@ def log2(x):
         raise TypeError(
             "Script parameter must be an opset. Did you use @script instead of @script()?"
         )
+    if op_type is not None and not isinstance(op_type, str):
+        raise TypeError(f"op_type must be a string or None, got {type(op_type).__name__}")
+    if isinstance(op_type, str) and not op_type.strip():
+        raise ValueError("op_type must be a non-empty string")
 
     def transform(f: Callable[_P, _R]) -> onnxscript.OnnxFunction[_P, _R]:
         if not inspect.isfunction(f):
@@ -92,6 +100,9 @@ def transform(f: Callable[_P, _R]) -> onnxscript.OnnxFunction[_P, _R]:
         env = module.__dict__.copy()
         env.update(closure.nonlocals)
         result = script_check(f_ast, opset, env, src, default_opset=default_opset)
+        if isinstance(op_type, str):
+            result.name = op_type
+            result.graph.name = op_type
         # TODO: add transformations.
         return onnxscript.OnnxFunction(opset, f, result, src, kwargs)
 

onnxscript/_internal/main_test.py

@@ -0,0 +1,153 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import inspect
+import unittest
+
+import onnx
+
+from onnxscript import FLOAT, script
+from onnxscript.onnx_opset import opset18 as op
+from onnxscript.values import Opset
+
+
+class ScriptCustomOpTypeTest(unittest.TestCase):
+    def test_custom_op_type_is_used_consistently(self):
+        custom_opset = Opset("com.example.custom", 1)
+
+        @script(custom_opset, op_type="MY_NEW_NAME_OP", producer_name="onnxscript-test")
+        def custom_op(x: FLOAT) -> FLOAT:
+            return op.Abs(x)
+
+        @script()
+        def caller(x: FLOAT) -> FLOAT:
+            return custom_op(x)
+
+        self.assertEqual(custom_op.name, "MY_NEW_NAME_OP")
+        self.assertEqual(custom_op.__name__, "custom_op")
+        self.assertEqual(custom_op.op_signature.name, "MY_NEW_NAME_OP")
+        self.assertEqual(custom_op.function_ir.graph.name, "MY_NEW_NAME_OP")
+        self.assertEqual(custom_op.to_function_proto().name, "MY_NEW_NAME_OP")
+
+        model = caller.to_model_proto()
+        call_node = model.graph.node[0]
+        self.assertEqual(call_node.op_type, "MY_NEW_NAME_OP")
+        self.assertEqual(call_node.domain, "com.example.custom")
+        self.assertEqual(
+            {(function.domain, function.name) for function in model.functions},
+            {("com.example.custom", "MY_NEW_NAME_OP")},
+        )
+        self.assertNotIn("op_type", custom_op.kwargs)
+        self.assertEqual(
+            custom_op.to_model_proto(io_types=FLOAT).producer_name, "onnxscript-test"
+        )
+        onnx.checker.check_model(model)
+
+    def test_issue_example_uses_default_local_opset(self):
+        @script(op_type="MY_NEW_NAME_OP")
+        def custom_op(x: FLOAT) -> FLOAT:
+            return op.Abs(x)
+
+        self.assertEqual(custom_op.name, "MY_NEW_NAME_OP")
+        self.assertEqual(custom_op.to_function_proto().name, "MY_NEW_NAME_OP")
+
+    def test_default_op_type_is_python_function_name(self):
+        @script()
+        def default_name(x: FLOAT) -> FLOAT:
+            return op.Abs(x)
+
+        self.assertEqual(default_name.name, "default_name")
+        self.assertEqual(default_name.to_function_proto().name, "default_name")
+
+    def test_none_uses_default_name_and_public_signature_is_stable(self):
+        @script(op_type=None)
+        def default_name(x: FLOAT) -> FLOAT:
+            return op.Abs(x)
+
+        self.assertIsNone(inspect.signature(script).parameters["op_type"].default)
+        self.assertEqual(default_name.name, "default_name")
+
+    def test_custom_name_can_be_reused_in_different_domains(self):
+        first_opset = Opset("com.example.first", 1)
+        second_opset = Opset("com.example.second", 1)
+
+        @script(first_opset, op_type="SharedName")
+        def first(x: FLOAT) -> FLOAT:
+            return op.Abs(x)
+
+        @script(second_opset, op_type="SharedName")
+        def second(x: FLOAT) -> FLOAT:
+            return op.Neg(x)
+
+        @script()
+        def intermediate(x: FLOAT) -> FLOAT:
+            return first(x)
+
+        @script()
+        def caller(x: FLOAT) -> tuple[FLOAT, FLOAT]:
+            return intermediate(x), second(x)
+
+        model = caller.to_model_proto()
+        self.assertTrue(
+            {
+                ("com.example.first", "SharedName"),
+                ("com.example.second", "SharedName"),
+            }.issubset({(function.domain, function.name) for function in model.functions})
+        )
+        onnx.checker.check_model(model)
+
+    def test_distinct_functions_cannot_share_an_identifier(self):
+        custom_opset = Opset("com.example.custom", 1)
+
+        @script(custom_opset, op_type="Duplicate")
+        def first(x: FLOAT) -> FLOAT:
+            return op.Abs(x)
+
+        @script(custom_opset, op_type="Duplicate")
+        def second(x: FLOAT) -> FLOAT:
+            return op.Neg(x)
+
+        @script()
+        def caller(x: FLOAT) -> tuple[FLOAT, FLOAT]:
+            return first(x), second(x)
+
+        with self.assertRaisesRegex(ValueError, "same identifier"):
+            caller.to_model_proto()
+
+    def test_repeated_calls_to_same_function_are_deduplicated(self):
+        custom_opset = Opset("com.example.custom", 1)
+
+        @script(custom_opset, op_type="CalledTwice")
+        def custom_op(x: FLOAT) -> FLOAT:
+            return op.Abs(x)
+
+        @script()
+        def caller(x: FLOAT) -> tuple[FLOAT, FLOAT]:
+            return custom_op(x), custom_op(x)
+
+        model = caller.to_model_proto()
+        self.assertEqual(len(model.functions), 1)
+
+    def test_invalid_custom_op_types_are_rejected(self):
+        invalid_values = ("", "   ", 42)
+
+        for value in invalid_values:
+            with self.subTest(value=value), self.assertRaises((TypeError, ValueError)):
+
+                @script(op_type=value)
+                def invalid(x: FLOAT) -> FLOAT:
+                    return op.Abs(x)
+
+    def test_non_c_identifier_names_are_supported(self):
+        for value in ("1LeadingDigit", "Name-With-Punctuation", "Name With Spaces"):
+            with self.subTest(value=value):
+
+                @script(op_type=value)
+                def custom_op(x: FLOAT) -> FLOAT:
+                    return op.Abs(x)
+
+                self.assertEqual(custom_op.name, value)
+
+
+if __name__ == "__main__":
+    unittest.main()