Merge branch 'main' into copilot/create-fusion-rule-remove-expand-node

Author: xadupre

.lintrunner.toml

@@ -98,42 +98,6 @@ init_command = [
 ]
 is_formatter = true
 
-[[linter]]
-code = 'PYLINT'
-include_patterns = [
-    '**/*.py',
-]
-exclude_patterns = [
-    'docs/**',
-    'examples/**',
-    'onnxscript/_internal/converter_test.py',
-    'onnxscript/optimizer/**',  # FIXME
-    'onnxscript/rewriter/**',  # FIXME
-    'tests/functions/**',
-    'tests/models/**',
-    'tests/onnx_backend_test_code/**',
-]
-command = [
-    'python',
-    '-m',
-    'lintrunner_adapters',
-    'run',
-    'pylint_linter',
-    '--rcfile=pyproject_pylint.toml',
-    '--show-disable',
-    '--',
-    '@{{PATHSFILE}}'
-]
-init_command = [
-    'python',
-    '-m',
-    'lintrunner_adapters',
-    'run',
-    'pip_init',
-    '--dry-run={{DRYRUN}}',
-    '--requirement=requirements/lintrunner/requirements.txt',
-]
-
 [[linter]]
 code = 'EDITORCONFIG-CHECKER'
 include_patterns = ['**']

docs/tutorial/rewriter/node_value_checkers.md

@@ -179,9 +179,16 @@ This means you should be careful when designing patterns with multiple alternati
 
 ## Error Handling
 
-Checkers can return either:
-- `True`: Check passed, continue matching
-- `False`: Check failed, pattern does not match
-- `MatchResult`: More detailed result with potential failure reasons
+Both check functions (including condition functions and node/value-level checkers) and
+rewrite functions support the same conventions for indicating failure:
 
-If a checker raises an exception, it will be caught and treated as a match failure, allowing patterns to fail gracefully when encountering unexpected conditions.
+- **`MatchResult` with `.fail()`** *(recommended)*: Return `MatchResult().fail("reason", source)` to indicate failure with a descriptive reason and optional source node/value. This provides the most useful diagnostic information for debugging.
+- **Raise `MatchFailureError`** *(recommended)*: Import it as `from onnxscript.rewriter.rewriter import MatchFailureError` and raise `MatchFailureError("reason", source1, source2, ...)` to indicate failure associated with one or more `ir.Node` or `ir.Value` objects. Each source should be passed as a separate positional argument (do not pass a list as a single argument). This is especially convenient in utility functions called from a check or rewrite, since it avoids having to explicitly propagate failure status through the call chain.
+- **Return `None` or `False`**: These indicate failure without providing a reason. They are supported but not recommended, since a failure reason is valuable for debugging why a rule did not apply.
+
+Including a descriptive failure reason is strongly encouraged. The rewriter's tracing infrastructure
+uses these reasons to report why rules failed to match, which is essential for diagnosing
+issues when developing or debugging rewrite rules.
+
+For **check functions**, success is indicated by returning `True` or a truthy `MatchResult`.
+For **rewrite functions**, success is indicated by returning one or more `ir.Value` results.

onnxscript/_internal/builder.py

@@ -31,6 +31,7 @@
     Sequence[float],
     Sequence[bool],
     Sequence[str],
+    None,
 ]
 
 # Mapping from Python scalar types to their default ONNX DataType,
@@ -132,6 +133,7 @@ def build_graph(
     *,
     opset_imports: dict[str, int] | None = None,
     name: str = "subgraph",
+    parent: GraphBuilder | None = None,
 ) -> ir.Graph:
     """Build an :class:`ir.Graph` suitable for use as a graph-valued attribute.
 
@@ -164,6 +166,10 @@ def build_graph(
         opset_imports: Opset version map for the subgraph (e.g.
             ``{"": 23}``).  Defaults to ``{"": 23}`` when *None*.
         name: Name of the resulting :class:`ir.Graph`.
+        parent: Optional parent :class:`GraphBuilder`.  When provided, the
+            sub-builder's ``_root`` points to the root builder of the parent,
+            so that :meth:`Parameter._realize` registers initializers in the
+            root (main) graph rather than the subgraph.
 
     Returns:
         An :class:`ir.Graph` whose inputs and outputs are populated and whose
@@ -187,7 +193,9 @@ def build_graph(
     for input_name, ts in resolved_inputs:
         subgraph.inputs.append(ir.Value(name=input_name, type=ts.type, shape=ts.shape))
 
-    sub_builder = GraphBuilder(subgraph)
+    sub_builder = GraphBuilder(subgraph, parent=parent)
+    if parent is not None:
+        sub_builder._scope_stack = list(parent._scope_stack)
     trace_outputs = trace_function(sub_builder.op, *subgraph.inputs)
     if not isinstance(trace_outputs, Sequence):
         trace_outputs = [trace_outputs]
@@ -208,8 +216,10 @@ def build_graph(
 class GraphBuilder:
     """Imperative builder for constructing ONNX IR graphs with automatic constant promotion, type casting, and shape inference."""
 
-    def __init__(self, graph: ir.Graph) -> None:
+    def __init__(self, graph: ir.Graph, parent: GraphBuilder | None = None) -> None:
         self._graph = graph
+        self._parent = parent
+        self._root: GraphBuilder = parent._root if parent is not None else self
 
         # Get the opset version for "" (default domain) from the graph
         if "" not in graph.opset_imports:
@@ -237,6 +247,16 @@ def opset(self, domain: str, version: int = 1) -> OpBuilder:
     def op(self) -> OpBuilder:
         return self._op_builder
 
+    @property
+    def parent(self) -> GraphBuilder | None:
+        """The parent builder, or None for a top-level builder."""
+        return self._parent
+
+    @property
+    def root(self) -> GraphBuilder:
+        """The root (top-level) builder in the parent chain."""
+        return self._root
+
     @property
     def graph(self) -> ir.Graph:
         return self._graph
@@ -258,7 +278,7 @@ def initializer(
 
     def _input_to_ir_value(
         self, value: VALUE_LIKE, like_type: ir.Value | None = None
-    ) -> ir.Value:
+    ) -> ir.Value | None:
         """Convert a permissible input (for a call to an op) into an ir.Value.
 
         Permissible values include ir.Value as well as python constants that can be converted
@@ -267,6 +287,8 @@ def _input_to_ir_value(
         """
         if isinstance(value, ir.Value):
             return value
+        if value is None:
+            return value
         dtype = (
             like_type.type.dtype
             if like_type is not None and like_type.type is not None
@@ -356,7 +378,7 @@ def _get_schema(
     def _partition_inputs_attributes(
         self,
         schema: onnx.defs.OpSchema | None,
-        inputs: Sequence[ir.Value | ir.TensorProtocol],
+        inputs: Sequence[ir.Value | ir.TensorProtocol | None],
         kwargs: dict[str, Any],
     ) -> tuple[Sequence[ir.Value | ir.TensorProtocol], dict[str, Any]]:
         if schema is None:
@@ -499,12 +521,13 @@ def subgraph(
             outputs,
             opset_imports=dict(self._graph.opset_imports),
             name=name,
+            parent=self,
         )
 
     def call_op(
         self,
         op_type: str,
-        inputs: Sequence[ir.Value | ir.TensorProtocol],
+        inputs: Sequence[ir.Value | ir.TensorProtocol | None],
         kwargs: dict[str, Any],
     ):
         """Create an ONNX node and add it to the graph, returning its output value(s)."""

onnxscript/_internal/builder_test.py

@@ -848,6 +848,39 @@ def add_mul(X, Y):
 
         self.assertIn("does not match", str(cm.exception))
 
+    def test_none_input_is_passed_through(self):
+        """Test that None inputs are preserved as None in the node's inputs."""
+        op, x, y = _create_builder_with_inputs()
+
+        # Gemm's third input (C) is optional; passing None should work
+        result = op.Gemm(x, y, None, alpha=1.0)
+
+        nodes = list(op.builder.graph)
+        self.assertEqual(len(nodes), 1)
+        node = nodes[0]
+        self.assertEqual(node.op_type, "Gemm")
+        # The third input should be None (optional, omitted)
+        self.assertEqual(len(list(node.inputs)), 3)
+        self.assertIs(node.inputs[0], x)
+        self.assertIs(node.inputs[1], y)
+        self.assertIsNone(node.inputs[2])
+        self.assertIsNotNone(result)
+
+    def test_none_input_with_custom_domain(self):
+        """Test that None inputs work with custom domain ops."""
+        op, x, y = _create_builder_with_inputs()
+
+        result = op.CustomOp(x, None, y, _domain="com.custom")
+
+        nodes = list(op.builder.graph)
+        self.assertEqual(len(nodes), 1)
+        node = nodes[0]
+        self.assertEqual(node.op_type, "CustomOp")
+        self.assertIs(node.inputs[0], x)
+        self.assertIsNone(node.inputs[1])
+        self.assertIs(node.inputs[2], y)
+        self.assertIsNotNone(result)
+
 
 class BuildSubgraphTest(unittest.TestCase):
     """Tests for GraphBuilder.subgraph()."""
@@ -1028,6 +1061,88 @@ def test_build_graph_custom_name(self):
         )
         self.assertEqual(graph.name, "loop_body")
 
+    def test_build_graph_with_parent(self):
+        """build_graph with parent sets root on the sub-builder."""
+        parent_graph = ir.Graph(
+            name="main",
+            inputs=[],
+            outputs=[],
+            nodes=[],
+            opset_imports={"": 23},
+        )
+        parent_builder = builder.GraphBuilder(parent_graph)
+
+        def body(op, x):
+            self.assertIs(op.builder.parent, parent_builder)
+            self.assertIs(op.builder.root, parent_builder)
+            return op.Identity(x)
+
+        builder.build_graph(
+            body,
+            inputs=[FLOAT[3]],
+            outputs=[FLOAT[3]],
+            parent=parent_builder,
+        )
+
+    def test_subgraph_sets_parent_and_root(self):
+        """GraphBuilder.subgraph() sets parent=self on the sub-builder."""
+        parent_graph = ir.Graph(
+            name="main",
+            inputs=[],
+            outputs=[],
+            nodes=[],
+            opset_imports={"": 23},
+        )
+        parent_builder = builder.GraphBuilder(parent_graph)
+
+        def body(op, x):
+            self.assertIs(op.builder.parent, parent_builder)
+            self.assertIs(op.builder.root, parent_builder)
+            return op.Identity(x)
+
+        parent_builder.subgraph(body, inputs=[FLOAT[3]], outputs=[FLOAT[3]])
+
+    def test_build_graph_inherits_parent_scope_stack(self):
+        """build_graph copies the parent's scope stack so nodes in the subgraph carry scoped names."""
+        parent_graph = ir.Graph(
+            name="main",
+            inputs=[],
+            outputs=[],
+            nodes=[],
+            opset_imports={"": 23},
+        )
+        parent_builder = builder.GraphBuilder(parent_graph)
+        parent_builder.push_module("encoder", "Encoder")
+        parent_builder.push_module("layers.0", "TransformerBlock")
+
+        subgraph = builder.build_graph(
+            lambda op, x: op.Relu(x),
+            inputs={"x": FLOAT[3, 4]},
+            outputs={"y": FLOAT[3, 4]},
+            parent=parent_builder,
+        )
+
+        # The single node created inside the subgraph should carry the
+        # parent's scope prefix in its name and metadata.
+        node = subgraph.node(0)
+        self.assertIn("encoder", node.name)
+        self.assertIn("layers.0", node.name)
+        self.assertIn("encoder", node.metadata_props["namespace"])
+        self.assertIn("TransformerBlock", node.metadata_props["namespace"])
+
+    def test_root_graph_builder_is_its_own_root(self):
+        """A top-level GraphBuilder has root == self."""
+        graph = ir.Graph(
+            name="main",
+            inputs=[],
+            outputs=[],
+            nodes=[],
+            opset_imports={"": 23},
+        )
+        gb = builder.GraphBuilder(graph)
+        self.assertIs(gb.root, gb)
+        self.assertIsNone(gb.parent)
+
 
 class PartitionInputsAttributesTest(unittest.TestCase):
     """Tests for GraphBuilder._partition_inputs_attributes."""

onnxscript/nn/_module_test.py

@@ -72,6 +72,77 @@ def test_realize_qualifies_name(self):
         self.assertEqual(value.name, "layer1.bias")
         self.assertIn("layer1.bias", graph.initializers)
 
+    def test_realize_in_subgraph_registers_in_root(self):
+        """Parameter realized inside a subgraph builder is stored in the root graph."""
+        from onnxscript._internal.builder import GraphBuilder
+        from onnxscript.onnx_types import FLOAT
+
+        root_graph = ir.Graph(
+            name="main",
+            inputs=[],
+            outputs=[],
+            nodes=[],
+            opset_imports={"": 23},
+        )
+        root_builder = GraphBuilder(root_graph)
+
+        p = Parameter([3, 4], name="weight")
+
+        def body_fn(op, x):
+            # Realize param inside a sub-builder context
+            p._realize(op.builder)  # pylint: disable=protected-access
+            return op.Add(x, x)
+
+        _sub_graph = root_builder.subgraph(
+            body_fn,
+            inputs=[FLOAT[3, 4]],
+            outputs=[FLOAT[3, 4]],
+        )
+        # Parameter should be in the ROOT graph's initializers, not the subgraph's
+        self.assertIn("weight", root_graph.initializers)
+        self.assertIs(root_graph.initializers["weight"], p)
+        # The subgraph should NOT have the initializer
+        self.assertNotIn("weight", _sub_graph.initializers)
+
+    def test_realize_in_nested_subgraph_registers_in_root(self):
+        """Parameter realized in a doubly-nested subgraph goes to the root graph."""
+        from onnxscript._internal.builder import GraphBuilder, build_graph
+        from onnxscript.onnx_types import FLOAT
+
+        root_graph = ir.Graph(
+            name="main",
+            inputs=[],
+            outputs=[],
+            nodes=[],
+            opset_imports={"": 23},
+        )
+        root_builder = GraphBuilder(root_graph)
+
+        p = Parameter([3], name="bias")
+
+        def inner_fn(op, x):
+            p._realize(op.builder)  # pylint: disable=protected-access
+            return op.Identity(x)
+
+        def outer_fn(op, x):
+            # Build a nested subgraph
+            build_graph(
+                inner_fn,
+                inputs=[FLOAT[3]],
+                outputs=[FLOAT[3]],
+                parent=op.builder,
+            )
+            return op.Identity(x)
+
+        root_builder.subgraph(
+            outer_fn,
+            inputs=[FLOAT[3]],
+            outputs=[FLOAT[3]],
+        )
+        # Even through two levels of nesting, param ends up in root
+        self.assertIn("bias", root_graph.initializers)
+        self.assertIs(root_graph.initializers["bias"], p)
+
 
 class ModuleBasicTest(unittest.TestCase):
     def test_parameter_auto_registration(self):