Refactor ir.Function builders to use build_function from onnxscript 0.7.0 - #189
Merged
Conversation
….7.0 Replace manual ir.Graph + GraphBuilder + ir.Function construction boilerplate with the new builder.build_function() utility. This eliminates ~25 lines of setup per function and automatically handles initializer lifting to Constant nodes. Updated files: - causal_conv.py: body extracted to trace function - linear_attention.py: body extracted, conditional inputs preserved - packed_multi_head_attention.py: body extracted, raw ir.Node via op._builder._graph.append() for ref_attr_name forwarding - skip_layer_normalization.py: both functions refactored similarly - pyproject.toml: bumped onnxscript requirement to >=0.7.0 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Replace *args unpacking with named parameters and None entries. The function signature is now stable (6 formals) across all update_rule variants. Absent optional inputs (decay, beta) get placeholder formals via build_function; the trace function receives None and branches naturally with 'if decay_v is not None'. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
- skip_layer_normalization: Replace raw ir.Node with op.LayerNormalization and op.RMSNormalization, passing ir.Attr objects with ref_attr_name directly as kwargs. OpBuilder passes them through to ir.node() correctly. Remove outdated docstring note about OpBuilder limitation. - packed_multi_head_attention: Replace raw ir.Node with op.Attention, passing ir.Attr objects with ref_attr_name for num_heads and scale. - linear_attention: Revert to per-variant input arity. The function is specialized per-model, so only inputs used by the update_rule are declared. This matches call sites (Mamba2/Bamba pass 5 args for 'gated', GatedDeltaNet passes 6 for 'gated_delta'). Shape inference does strict arity checking. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
The strict arity check is in the Python onnx-shape-inference package (onnx_shape_inference._functions.infer_function_call_output_shapes), not in ONNX's C++ shape inference. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Performance Comparison
|
justinchuby
approved these changes
Apr 22, 2026
justinchuby
left a comment
justinchuby
left a comment
Member
There was a problem hiding this comment.
Thanks! This reminds me we can also update the call sites now.
Contributor
There was a problem hiding this comment.
Pull request overview
Refactors Mobius’ standard-ONNX ir.Function fallback builders to use onnxscript._internal.builder.build_function() (onnxscript >= 0.7.0), removing manual ir.Graph/GraphBuilder boilerplate while preserving the same function signatures and attribute-forwarding behavior used by InlinePass expansion.
Changes:
- Bump
onnxscriptdependency to>=0.7.0. - Refactor multiple
ir.Functionfactories to define abody(op, ...)trace function and build viabuilder.build_function(...). - Replace prior raw
ir.Nodeconstruction patterns withOpBuildercalls usingir.Attr(..., ref_attr_name=...)where needed.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| pyproject.toml | Updates the minimum onnxscript version to support build_function(). |
| src/mobius/functions/causal_conv.py | Switches to build_function() with a traced body, keeping the same function inputs/outputs and activation attribute. |
| src/mobius/functions/linear_attention.py | Switches to build_function() and preserves per-variant input arity while keeping attribute forwarding. |
| src/mobius/functions/packed_multi_head_attention.py | Switches to build_function() and emits op.Attention(...) with forwarded ir.Attr kwargs. |
| src/mobius/functions/skip_layer_normalization.py | Switches to build_function() and uses op.LayerNormalization/op.RMSNormalization with forwarded epsilon via ir.Attr. |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
|
The author of this PR, gramalingam, is not an activated member of this organization on Codecov. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replaces manual
ir.Graph+GraphBuilder+ir.Functionconstruction boilerplate with the newbuilder.build_function()utility from onnxscript 0.7.0. This eliminates ~25 lines of setup per function and automatically handles initializer lifting to Constant nodes.Changes
>=0.7.0op.Attention(...)withir.Attrkwargs forref_attr_nameforwarding (eliminates rawir.Nodeconstruction)op.LayerNormalization/op.RMSNormalizationwithir.Attrkwargs forref_attr_name(eliminates rawir.Nodeconstruction and outdated docstring note)Key findings
ir.Attrwithref_attr_nameas kwargs — it flows through toir.node()which acceptsir.AttrinSupportedAttrTypes. This eliminates the need for rawir.Nodeconstruction in all three files that previously required it.onnx-shape-inference(Python package) does a strictlen(f.inputs) != len(node.inputs)check. ONNX's own C++ shape inference (onnx.shape_inference.infer_shapes) accepts trailing optional inputs fine, even in strict mode. A TODO is added to investigate relaxing this upstream.Testing
All 2558 tests pass, 39 skipped (unchanged from baseline).