[IR] Normalize "ai.onnx" domain to "" (#2283)

TODO: Also update opset_imports to handle the "ai.onnx" key.

Fix #2280

Author: justinchuby

onnxscript/ir/_core.py

@@ -1278,6 +1278,11 @@ def _short_tensor_str_for_node(x: Value) -> str:
     return "{...}"
 
 
+def _normalize_domain(domain: str) -> str:
+    """Normalize 'ai.onnx' to ''"""
+    return "" if domain == "ai.onnx" else domain
+
+
 class Node(_protocols.NodeProtocol, _display.PrettyPrintable):
     """IR Node.
 
@@ -1328,6 +1333,7 @@ def __init__(
 
         Args:
             domain: The domain of the operator. For onnx operators, this is an empty string.
+                When it is "ai.onnx", it is normalized to "".
             op_type: The name of the operator.
             inputs: The input values. When an input is ``None``, it is an empty input.
             attributes: The attributes. RefAttr can be used only when the node is defined in a Function.
@@ -1350,7 +1356,7 @@ def __init__(
             ValueError: If an output value has a producer set already, when outputs is specified.
         """
         self._name = name
-        self._domain: str = domain
+        self._domain: str = _normalize_domain(domain)
         self._op_type: str = op_type
         # NOTE: Make inputs immutable with the assumption that they are not mutated
         # very often. This way all mutations can be tracked.
@@ -1482,7 +1488,7 @@ def domain(self) -> str:
 
     @domain.setter
     def domain(self, value: str) -> None:
-        self._domain = value
+        self._domain = _normalize_domain(value)
 
     @property
     def version(self) -> int | None:
@@ -2885,7 +2891,7 @@ def domain(self) -> str:
 
     @domain.setter
     def domain(self, value: str) -> None:
-        self._domain = value
+        self._domain = _normalize_domain(value)
 
     @property
     def overload(self) -> str:

onnxscript/ir/_core_test.py

@@ -850,6 +850,17 @@ def test_successors(self):
     def test_successors_are_unique(self):
         self.assertEqual(self.node.successors(), (self.node_a, self.node_b))
 
+    def test_domain_normalizes_ai_onnx(self):
+        # Node domain is always normalized to "" if it is "ai.onnx"
+        node = _core.Node("ai.onnx", "TestOp", inputs=())
+        self.assertEqual(node.domain, "")
+
+        node.domain = ""
+        self.assertEqual(node.domain, "")
+
+        node.domain = "ai.onnx"
+        self.assertEqual(node.domain, "")
+
     # TODO(justinchuby): Test all methods
 
 

onnxscript/optimizer/_constant_folding.py

@@ -64,7 +64,7 @@ def is_constant_op(node: ir.Node) -> bool:
 
 def _process_constant_node(node: ir.Node) -> None:
     """Sets const_value of output value of a Constant op node."""
-    if node.op_type != "Constant" or node.domain not in {"", "ai.onnx"}:
+    if node.op_type != "Constant" or node.domain != "":
         return
     if len(node.attributes) != 1:
         return

onnxscript/version_converter/_version_converter.py

@@ -229,7 +229,7 @@ def _upgrade_version(self, node: ir.Node, opset_version: int, up_conversion: boo
     def process_node(
         self, node: ir.Node, opset_version: int, up_conversion: bool = True
     ) -> Replacement | None:
-        if node.domain not in {"", "ai.onnx"}:
+        if node.domain != "":
             return None
         adapter = registry.lookup_adapters(
             node.domain, node.op_type, opset_version, up_conversion