Merge branch 'main' into copilot/fix-squeeze-layer-axes

Author: justinchuby

.github/workflows/main.yaml

@@ -83,7 +83,7 @@ jobs:
           token: ${{ secrets.CODECOV_TOKEN }}
       - name: Upload torchlib error reports
         if: always()
-        uses: actions/upload-artifact@v6
+        uses: actions/upload-artifact@v7
         with:
           name: Error reports (${{ matrix.name }}-${{ matrix.os }})
           path: error_reports

.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/builder/graph_builder.md

@@ -339,9 +339,11 @@ The subgraph automatically inherits the opset version from the parent
 
 ### Type annotations for subgraph inputs and outputs
 
-`subgraph()` accepts `input_types` and `output_types` lists that describe
-the types and shapes of each input and output. Each element can be either an
-`ir.TypeAndShape` object or — more conveniently — an
+`subgraph()` accepts `inputs` and `outputs` that describe
+the types and shapes of each input and output. They can be provided as a
+:class:`list` of type specs (names are auto-generated) **or** as a
+:class:`dict` mapping explicit names to type specs. Each type spec can be
+either an `ir.TypeAndShape` object or — more conveniently — an
 `onnxscript` tensor-type expression:
 
 | Expression           | Meaning                                 |
@@ -408,8 +410,8 @@ def cumsum_body(op, state, x_i):
 
 body = builder.subgraph(
     cumsum_body,
-    input_types=[FLOAT[D], FLOAT[D]],   # state, x_i
-    output_types=[FLOAT[D], FLOAT[D]],  # new_state, scan_out_i
+    inputs=[FLOAT[D], FLOAT[D]],   # state, x_i
+    outputs=[FLOAT[D], FLOAT[D]],  # new_state, scan_out_i
     name="cumsum_body",
 )
 
@@ -430,7 +432,7 @@ model = ir.Model(graph=graph, ir_version=10)
 
 Key points:
 
-- `builder.subgraph(fn, input_types, output_types)` creates a fresh
+- `builder.subgraph(fn, inputs, outputs)` creates a fresh
   `ir.Graph`, calls `fn(op, *inputs)` to trace the body, and wires up the
   declared input/output types.
 - The `fn` receives an `OpBuilder` as its first argument — exactly the same
@@ -450,8 +452,8 @@ def outer_body(op, state, x_i):
     # Build a nested subgraph inside the scan body
     inner = op.builder.subgraph(
         lambda iop, v: iop.Relu(v),
-        input_types=[FLOAT[D]],
-        output_types=[FLOAT[D]],
+        inputs=[FLOAT[D]],
+        outputs=[FLOAT[D]],
         name="relu_body",
     )
     # ... use inner as a graph attribute of a nested op ...

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/__init__.py

@@ -13,6 +13,7 @@
     "OnnxFunction",
     "TracedOnnxFunction",
     "GraphBuilder",
+    "OpBuilder",
     "proto2python",
     "external_tensor",
     "BFLOAT16",
@@ -129,7 +130,7 @@
 # isort: on
 
 from . import ir, nn, optimizer, rewriter, version_converter
-from ._internal.builder import GraphBuilder
+from ._internal.builder import GraphBuilder, OpBuilder
 from ._internal.utils import external_tensor
 from ._internal.values import OnnxFunction, TracedOnnxFunction
 

onnxscript/_internal/_inliner.py

@@ -10,29 +10,29 @@
 
 
 def instantiate(
-    function: ir.Function,
+    graph: ir.Graph,
     inputs: Sequence[ir.Value | None],
     attributes: Mapping[str, ir.Attr],
     *,
     prefix: str = "",
 ) -> tuple[list[ir.Node], list[ir.Value | None]]:
-    """Instantiate (inline) a function, substituting inputs and attributes.
+    """Instantiate (inline) a graph, substituting inputs and attributes.
 
     Args:
-        function: The function to instantiate.
-        inputs: Actual input values to bind to the function's formal parameters.
+        graph: The graph to instantiate.
+        inputs: Actual input values to bind to the graph's formal parameters.
         attributes: Attribute values to substitute for reference attributes.
         prefix: Optional prefix to prepend to node and output names.
 
     Returns:
-        A tuple of (nodes, outputs) where nodes are the cloned function body
-        and outputs are the values corresponding to the function's outputs.
+        A tuple of (nodes, outputs) where nodes are the cloned graph body
+        and outputs are the values corresponding to the graph's outputs.
     """
-    formal_inputs = function.inputs
+    formal_inputs = graph.inputs
     if len(inputs) > len(formal_inputs):
         raise ValueError(
             f"Too many inputs: got {len(inputs)}, "
-            f"but function has {len(formal_inputs)} parameters."
+            f"but graph has {len(formal_inputs)} parameters."
         )
     value_map: dict[ir.Value, ir.Value | None] = dict(zip(formal_inputs, inputs))
 
@@ -50,7 +50,8 @@ def rename(node: ir.Node) -> None:
         metadata_props={},
         post_process=rename,
         resolve_ref_attrs=True,
+        allow_outer_scope_values=True,
     )
-    nodes = [cloner.clone_node(n) for n in function]
-    outputs = [value_map.get(v) for v in function.outputs]
+    nodes = [cloner.clone_node(n) for n in graph]
+    outputs = [value_map.get(v) for v in graph.outputs]
     return nodes, outputs

onnxscript/_internal/analysis.py

@@ -153,7 +153,7 @@ def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
             return live
 
         def do_visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
-            def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
+            def visit_block(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
                 for s in reversed(block):
                     live_out = visit(s, live_out)
                 return live_out
@@ -167,28 +167,28 @@ def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
             if isinstance(stmt, ast.If):
                 constant_cond = self.constant_if_condition(stmt)
                 if constant_cond is None:
-                    live1 = visitBlock(stmt.body, live_out)
-                    live2 = visitBlock(stmt.orelse, live_out)
+                    live1 = visit_block(stmt.body, live_out)
+                    live2 = visit_block(stmt.orelse, live_out)
                     return live1 | live2 | _used_vars(stmt.test)
                 elif constant_cond:
-                    return visitBlock(stmt.body, live_out)
+                    return visit_block(stmt.body, live_out)
                 else:
-                    return visitBlock(stmt.orelse, live_out)
+                    return visit_block(stmt.orelse, live_out)
             if isinstance(stmt, ast.For):
                 p_loop_var = _get_loop_var(stmt, self._formatter)
                 prev = None
                 curr = live_out
                 while curr != prev:
                     prev = curr
-                    curr = visitBlock(stmt.body, prev).difference({p_loop_var})
+                    curr = visit_block(stmt.body, prev).difference({p_loop_var})
                 return curr
             if isinstance(stmt, ast.While):
                 cond_vars = _used_vars(stmt.test)
                 prev = None
                 curr = live_out | cond_vars
                 while curr != prev:
                     prev = curr
-                    curr = visitBlock(stmt.body, prev) | cond_vars
+                    curr = visit_block(stmt.body, prev) | cond_vars
                 return curr
             if isinstance(stmt, ast.Break):
                 # The following is sufficient for the current restricted usage, where
@@ -228,7 +228,7 @@ def exposed_uses(self, stmts: Sequence[ast.stmt]) -> set[str]:
         (in the first statement). Hence x is included in the exposed_uses.
         """
 
-        def visitBlock(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
+        def visit_block(block: Sequence[ast.stmt], live_out: Set[str]) -> Set[str]:
             for stmt in reversed(block):
                 live_out = visit(stmt, live_out)
             return live_out
@@ -243,13 +243,13 @@ def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
             if isinstance(stmt, ast.If):
                 constant_cond = self.constant_if_condition(stmt)
                 if constant_cond is None:
-                    live1 = visitBlock(stmt.body, live_out)
-                    live2 = visitBlock(stmt.orelse, live_out)
+                    live1 = visit_block(stmt.body, live_out)
+                    live2 = visit_block(stmt.orelse, live_out)
                     return (live1 | live2) | _used_vars(stmt.test)
                 elif constant_cond:
-                    return visitBlock(stmt.body, live_out)
+                    return visit_block(stmt.body, live_out)
                 else:
-                    return visitBlock(stmt.orelse, live_out)
+                    return visit_block(stmt.orelse, live_out)
             if ast_utils.is_print_call(stmt):
                 return live_out
             if ast_utils.is_doc_string(stmt):
@@ -259,13 +259,13 @@ def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
                 # for loops that execute at least once.
                 loop_var_set = {_get_loop_var(stmt, self._formatter)}
                 used_after_loop = live_out.difference(loop_var_set)
-                used_inside_loop = visitBlock(stmt.body, set()).difference(loop_var_set)
+                used_inside_loop = visit_block(stmt.body, set()).difference(loop_var_set)
                 used_in_loop_header = _used_vars(stmt.iter)
                 return used_inside_loop | used_in_loop_header | used_after_loop
             if isinstance(stmt, ast.While):
                 # Analysis assumes loop may execute zero times. Results can be improved
                 # for loops that execute at least once.
-                used_inside_loop = visitBlock(stmt.body, set())
+                used_inside_loop = visit_block(stmt.body, set())
                 used_in_loop_header = _used_vars(stmt.test)
                 return used_inside_loop | used_in_loop_header | live_out
             if isinstance(stmt, ast.Break):
@@ -281,7 +281,7 @@ def visit(stmt: ast.stmt, live_out: Set[str]) -> Set[str]:
                 self._formatter(stmt, f"Unsupported statement type {type(stmt)!r}.")
             )
 
-        return visitBlock(stmts, set())
+        return visit_block(stmts, set())
 
     def outer_scope_variables(self, fun: ast.FunctionDef) -> set[str]:
         """Return the set of outer-scope variables used in a nested function.

onnxscript/_internal/autocast.py

@@ -187,15 +187,15 @@ def get_type_info(x: Optional[ir.Value]) -> Optional[ir.Value]:
         argument of CastLike) and None otherwise. In the expression "Add(X, 1), 1 is
         castable, while X can serve as the target-type.
         """
-        return None if x is None or converter_.is_castable(x.name) else x
+        return None if x is None or converter_._is_castable(x.name) else x  # pylint: disable=protected-access
 
     def cast_like(x: Optional[ir.Value], y: Optional[ir.Value]) -> Optional[str]:
         if x is None:
             return None
-        if converter_.is_castable(x.name) and y is not None:
+        if converter_._is_castable(x.name) and y is not None:  # pylint: disable=protected-access
             # Polymorphic constant x is cast to the type of y:
-            x_cast = converter_.generate_unique_name(f"{x.name}_cast")
-            return converter_.emit1([x_cast], "CastLike", [x, y])
+            x_cast = converter_._generate_unique_name(f"{x.name}_cast")  # pylint: disable=protected-access
+            return converter_._emit1([x_cast], "CastLike", [x, y])  # pylint: disable=protected-access
         return x
 
     return cast_inputs(get_type_info, cast_like, op_signature, args)