Refactor: Introduce OpBuilderBase/TapeBuilder as unified op-building interface

[gramalingam]

Summary

Introduces OpBuilderBase (ABC) and TapeBuilder (concrete implementation) as the unified interface for building IR nodes across the rewriter, optimizer, and version converter. This is a step towards unifying these builders with the recently introduced GraphBuilder. As of now, there is still some differences internally (as these are oriented towards incrementally modifying an existing graph), but for end-users (eg., writing rewrite-rules) we should be able to move towards same API.

Changes

New: onnxscript/tape_builder.py

• OpBuilderBase — Abstract base class providing the op-building API:
  • Dynamic dispatch: op.Relu(x), op.MatMul(a, b, _domain=...)
  • Explicit creation: op.op("Conv", inputs, attributes, domain=...)
  • Initializer creation: op.initializer(tensor, name=...)
  • Subclasses implement _add_node, _add_initializer, _record_opset
• TapeBuilder — Concrete subclass with list-based storage and harvesting properties (nodes, initializers, used_opsets)
• Both are exported as public names from onnxscript

Rewriter

• rewriter/_context.py slimmed to re-exports + local alias RewriterContext = OpBuilderBase
• _rewrite_rule.py uses TapeBuilder() directly; harvests from the context
• Deleted _node_sink.py (the intermediate NodeSink/TapeSink layer)

Optimizer

• optimizer/_constant_folding.py defines OptimizerContext = OpBuilderBase locally
• Uses TapeBuilder() instead of the old _tape.Builder()

Version Converter

• version_converter/_version_converter.py defines VCContext = OpBuilderBase locally
• Uses TapeBuilder() instead of the old _tape.Builder()

Cleanup

• Deleted onnxscript/ir/_tape.py and _tape_test.py (fully superseded)
• Removed incorrect ir.Value type annotations from pattern() method signatures in rule files (pattern methods work with pattern-value objects, not ir.Value)

Design

OpBuilderBase (ABC)          <- shared interface
  _add_node()                <- abstract
  _add_initializer()         <- abstract
  _record_opset()            <- abstract
  __getattr__()              <- dynamic dispatch (concrete)
  op()                       <- explicit node creation (concrete)
  initializer()              <- initializer creation (concrete)

TapeBuilder(OpBuilderBase)   <- list-backed implementation
  nodes                      <- harvesting property
  initializers               <- harvesting property
  used_opsets                <- harvesting property

Aliases:
  RewriterContext  = OpBuilderBase  (in rewriter)
  OptimizerContext = OpBuilderBase  (in optimizer)
  VCContext        = OpBuilderBase  (in version converter)

This design allows future alternative implementations (e.g., graph-backed builder) by subclassing OpBuilderBase without changing rule/evaluator code.

[Copilot]

Pull request overview

This draft PR refactors the builder used by rewrite rules by introducing a dedicated rewriter context and a pluggable sink abstraction, with the goal of decoupling rule authoring from the current tape-backed storage implementation. It fits into the rewriter subsystem as an internal step toward unifying builder behavior across rewrite, optimization, and graph-construction paths.

Changes:

• Introduces RewriterContext and NodeSink/TapeSink to separate rule-facing node construction from storage/harvesting details.
• Updates rewrite-rule internals to instantiate the new context/sink pair when building replacement subgraphs.
• Adjusts several common rewrite rules to accept the new context type and adds focused unit tests for context/sink behavior.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
onnxscript/rewriter/rules/common/_remove_optional_bias.py	Removes explicit ir.tape.Tape type annotations from rule hooks.
onnxscript/rewriter/rules/common/_fuse_pad_into_conv.py	Updates rule hook signatures to use the new generic rewriter context parameter.
onnxscript/rewriter/rules/common/_fuse_conv_affine.py	Drops concrete builder type annotations from rewrite methods.
onnxscript/rewriter/pattern.py	Re-exports the new RewriterContext from _context instead of aliasing _tape.Builder.
onnxscript/rewriter/_rewrite_rule.py	Switches replacement construction from direct builder instantiation to RewriterContext(TapeSink()).
onnxscript/rewriter/_node_sink.py	Adds sink abstraction for collecting created nodes, initializers, and opset usage.
onnxscript/rewriter/_context_test.py	Adds unit tests for sink harvesting, forbidden attribute access, and dynamic op creation.
onnxscript/rewriter/_context.py	Implements the restricted rewrite-rule context API and sink-backed node/initializer creation.

[codecov]

❌ 1 Tests Failed:

Tests completed	Failed	Passed	Skipped
14050	1	14049	2690

View the full list of 3 ❄️ flaky test(s)

onnxscript.rewriter.rules.common._basic_rules_test.ReshapeReshapeTest::test_reshape_reshape_dynamic_rule_1

Flake rate in main: 95.75% (Passed 45 times, Failed 1015 times)

Stack Traces | 0.011s run time

..../test_ort_nightly/lib/python3.11.../site-packages/parameterized/parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.../rules/common/_basic_rules_test.py:531: in test_reshape_reshape_dynamic_rule
    testing.assert_numerically_equal(model, updated_model, (inputs,), atol=0, rtol=0)
onnxscript/rewriter/testing.py:102: in assert_numerically_equal
    np.testing.assert_allclose(
E   AssertionError: 
E   Not equal to tolerance rtol=0, atol=0
E   
E   (shapes (1, 0, 6, 3), (1, 0, 0, 3) mismatch)
E    ACTUAL: array([], shape=(1, 0, 6, 3), dtype=float32)
E    DESIRED: array([], shape=(1, 0, 0, 3), dtype=float32)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__addmv_cpu_float32

Flake rate in main: 14.56% (Passed 223 times, Failed 38 times)

Stack Traces | 0.628s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 1 / 5 (20.0%)
E   Greatest absolute difference: 1.52587890625e-05 at index (0,) (up to 1e-05 allowed)
E   Greatest relative difference: 7.915093192423228e-06 at index (0,) (up to 1.3e-06 allowed)

tests.function_libs.torch_lib.ops_test.TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__logsumexp_cpu_float16

Flake rate in main: 14.93% (Passed 114 times, Failed 20 times)

Stack Traces | 0.571s run time

.../function_libs/torch_lib/ops_test.py:243: in run_test_output_match
    torch.testing.assert_close(
E   AssertionError: Tensor-likes are not close!
E   
E   Mismatched elements: 1 / 5 (20.0%)
E   Greatest absolute difference: 2.288818359375e-05 at index (1,) (up to 1e-05 allowed)
E   Greatest relative difference: 0.0022869110107421875 at index (1,) (up to 0.001 allowed)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (2)

onnxscript/version_converter/_version_converter.py:263

• TapeBuilder can now create initializers and track referenced opsets, but process_node() still harvests only context.nodes. If a future adapter uses op.initializer(...) or creates a node in another domain, the converted graph will drop those initializers/opset imports and become invalid even though the new context API advertises that functionality.

        context = TapeBuilder()
        output = adapter(node, context)
        if output is not None:
            if isinstance(output, ir.Value):
                output = [output]
            return Replacement(output, context.nodes)

onnxscript/optimizer/_constant_folding.py:1192

• TapeBuilder now supports initializer() and records opsets, but this replacement path only returns context.nodes. Any partial evaluator that starts using the unified builder to emit an initializer or a node from a non-default domain will silently lose that state when the replacement is applied, producing an invalid rewrite.

            context = TapeBuilder()
            try:
                output = optimizer(node, context, self._state)
            except Exception as e:
                raise RuntimeError(
                    f"Error during constant folding for node {node.name!r} ({node.domain}::{node.op_type})"
                ) from e
            if output is not None:
                if isinstance(output, Replacement):
                    return output
                if isinstance(output, ir.Value):
                    output = [output]
                return Replacement(output, context.nodes)