Refactor: Move op_configuration.rs from burn-import to onnx-ir #3126

antimora · 2025-04-30T17:08:34Z

Pull Request Template

Checklist

Confirmed that the run-checks all script has been executed.
Ensured the book is updated with changes in this PR.

Related Issues/PRs

Fixes #3091

Changes

Moved op_configuration.rs from burn-import to onnx-ir. The onnx-ir module does not depend on any data types from burn-import or burn. Configurations are now organized by node type.
Added unit tests for the moved node configurations.
Relocated rank inference functions to node folders in onnx-ir and added corresponding unit tests.
Updated the developer book and README files.

Testing

Existing integration tests are passing.
New unit tests have been added.

codecov · 2025-04-30T17:47:41Z

Codecov Report

Attention: Patch coverage is 93.16003% with 421 lines in your changes missing coverage. Please review.

Project coverage is 82.15%. Comparing base (78a2f24) to head (86d8c72).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/onnx-ir/src/node/test_utils.rs	65.01%	141 Missing ⚠️
crates/onnx-ir/src/node/conv_transpose1d.rs	82.85%	24 Missing ⚠️
crates/onnx-ir/src/node/padding.rs	78.49%	20 Missing ⚠️
crates/onnx-ir/src/node/constant.rs	80.45%	17 Missing ⚠️
crates/onnx-ir/src/node/reshape.rs	82.55%	15 Missing ⚠️
crates/onnx-ir/src/node/unsqueeze.rs	92.78%	14 Missing ⚠️
crates/onnx-ir/src/node/split.rs	95.59%	13 Missing ⚠️
crates/onnx-ir/src/node/resize.rs	93.20%	11 Missing ⚠️
crates/onnx-ir/src/node/expand.rs	94.28%	10 Missing ⚠️
crates/onnx-ir/src/node/pad.rs	96.60%	9 Missing ⚠️
... and 44 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3126      +/-   ##
==========================================
+ Coverage   81.57%   82.15%   +0.57%     
==========================================
  Files         904      956      +52     
  Lines      118064   121759    +3695     
==========================================
+ Hits        96314   100026    +3712     
+ Misses      21750    21733      -17

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Still remaining: expand_config tile_config top_k_config trilu_config pad_config unsqueeze_config split_config

Copilot

Pull Request Overview

This PR refactors configuration usage by moving op_configuration.rs from burn-import to onnx-ir and updates the associated node implementations and tests accordingly. Key changes include replacing direct imports from burn::nn with onnx_ir equivalents, updating test configurations for various convolution and pooling nodes, and revising documentation in both the onnx-tests and contributor guide files.

Reviewed Changes

Copilot reviewed 88 out of 88 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
crates/burn-import/src/burn/node/linear.rs	Replaces LinearConfig import from burn::nn with onnx_ir::node::linear::LinearConfig.
crates/burn-import/src/burn/node/layer_norm.rs	Updates LayerNorm configuration by importing from onnx_ir instead of burn::nn.
crates/burn-import/src/burn/node/expand.rs	Adjusts ExpandShape handling and test configuration to use Argument-based parameters.
crates/burn-import/src/burn/node/dropout.rs	Swaps out DropoutConfig import from burn::nn to onnx_ir::node::dropout::DropoutConfig.
crates/burn-import/src/burn/node/conv_transpose_3d.rs	Modifies ConvTranspose3d configuration to use a new parameter ordering and values.
crates/burn-import/src/burn/node/conv_transpose_2d.rs	Updates ConvTranspose2d configuration with revised parameters and stride values.
...	Other files include similar updates for conv1d/conv2d/conv3d, pooling nodes, and documentation updates.

Comments suppressed due to low confidence (2)

crates/burn-import/src/burn/node/conv_transpose_3d.rs:186

The stride values are set to [0, 0, 0] which may result in invalid behavior during execution. Verify whether a zero stride is intentional or if non-zero values should be used.

.with_stride([0, 0, 0])

crates/burn-import/src/burn/node/conv_transpose_2d.rs:186

Stride values are set to [0, 0] which might cause incorrect behavior. Confirm that a zero stride is intended for this configuration.

.with_stride([0, 0])

laggui

That's a lot of lines to go over to essentially 1) move the node configs to onnx-ir and 2) move the node rank inference functions from rank_inference.rs to a file per node 😅

I understand the objective to get onnx-ir standalone (no dependency to burn), but I have some comments below.

crates/burn-import/README.md

crates/burn-import/src/burn/node/avg_pool1d.rs

crates/burn-import/src/burn/node/avg_pool2d.rs

crates/burn-import/src/burn/node/base.rs

crates/burn-import/src/burn/node/expand.rs

crates/onnx-ir/README.md

crates/onnx-ir/src/node/argmax.rs

laggui

The node builder pattern is a big improvement, thanks! This addresses my main concern.

…l-ai#3126) * Move conv1d, avg_pool1d, max_pool1d and conv_transpose1d to onnx-ir * Fix format * Fix format * Move conv1d rank update * Move conv_transpose1d_update_outputs * Removed and replaced as same_as_input * Move op config from burn-import * Moved 2d and 3d config functions * Move some config functions to node * Moved op configs to individual node modules Still remaining: expand_config tile_config top_k_config trilu_config pad_config unsqueeze_config split_config * Break down reduce module into individual modules * Move rank inference functions * Move rank updates to node module * Add documentation * Repoint config function from onnx-ir * Remove burn types from onnx-ir * Fixed left over tests * No default init for config structs * Fix format * remove op_configuration.rs * Decouple burn-import types from op_configuration * Moved remaining configs from burn-import to onnx-ir * Update the documentation * Remove "Features" section * Converted links to resources * Remove deadcode * Shorten burn-import readme * Remove inline comments * Add NodeBuilder and refactor test code * Refactor tests to use NodeBuilder * Refactored tests to use NodeBuilder * Refactored tests to use NodeBuilder * Refactored tests to use NodeBuilder * Refactored tests to use NodeBuilder * Refactored tests to use NodeBuilder * Refactored tests to use NodeBuilder * Remove inline comments

antimora added 2 commits April 29, 2025 20:12

Move conv1d, avg_pool1d, max_pool1d and conv_transpose1d to onnx-ir

1ede17c

Merge remote-tracking branch 'upstream/main' into refactor-opconfig

519ddd3

antimora marked this pull request as draft April 30, 2025 17:08

antimora added 4 commits April 30, 2025 12:12

Fix format

1b1287b

Fix format

e523498

Move conv1d rank update

e0083c0

Move conv_transpose1d_update_outputs

9f39454

antimora added 19 commits April 30, 2025 12:55

Removed and replaced as same_as_input

4fa7636

Move op config from burn-import

4e41de9

Moved 2d and 3d config functions

8e81541

Move some config functions to node

81b38d9

Moved op configs to individual node modules

b19ce42

Still remaining: expand_config tile_config top_k_config trilu_config pad_config unsqueeze_config split_config

Break down reduce module into individual modules

d38a7e9

Move rank inference functions

b21369b

Move rank updates to node module

03ea300

Add documentation

c60cde3

Repoint config function from onnx-ir

a66a43b

Remove burn types from onnx-ir

4834dff

Fixed left over tests

e1d9682

No default init for config structs

29a0896

Fix format

995b7eb

remove op_configuration.rs

fa67a32

Decouple burn-import types from op_configuration

8e8c5a3

Moved remaining configs from burn-import to onnx-ir

5fbfbec

Merge remote-tracking branch 'upstream/main' into refactor-opconfig

1c5f098

Update the documentation

9efe0a0

antimora self-assigned this May 2, 2025

antimora added the onnx label May 2, 2025

antimora marked this pull request as ready for review May 2, 2025 03:28

antimora requested review from laggui and Copilot May 2, 2025 03:28

Copilot AI reviewed May 2, 2025

View reviewed changes

antimora mentioned this pull request May 3, 2025

break down test_onnx.rs into test subdirectories #3144

Merged

2 tasks

Merge remote-tracking branch 'upstream/main' into refactor-opconfig

781d75c

laggui reviewed May 8, 2025

View reviewed changes

antimora added 5 commits May 8, 2025 13:41

Remove "Features" section

d571dd6

Converted links to resources

67d0067

Remove deadcode

2ee9317

Shorten burn-import readme

f27645e

Remove inline comments

5ef9401

antimora requested a review from laggui May 8, 2025 19:29

antimora added 9 commits May 8, 2025 15:05

Add NodeBuilder and refactor test code

92890e7

Refactor tests to use NodeBuilder

83112de

Refactored tests to use NodeBuilder

6dfe2c0

Refactored tests to use NodeBuilder

a4592a3

Refactored tests to use NodeBuilder

52ed266

Refactored tests to use NodeBuilder

eca4a8a

Refactored tests to use NodeBuilder

689ba26

Refactored tests to use NodeBuilder

a898404

Remove inline comments

86d8c72

laggui approved these changes May 9, 2025

View reviewed changes

antimora merged commit f18cb7f into tracel-ai:main May 9, 2025
11 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Refactor: Move op_configuration.rs from burn-import to onnx-ir #3126

Refactor: Move op_configuration.rs from burn-import to onnx-ir #3126

Uh oh!

antimora commented Apr 30, 2025 •

edited

Loading

Uh oh!

codecov bot commented Apr 30, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

laggui left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

laggui left a comment

Uh oh!

Uh oh!

Uh oh!

Refactor: Move op_configuration.rs from burn-import to onnx-ir #3126

Refactor: Move op_configuration.rs from burn-import to onnx-ir #3126

Uh oh!

Conversation

antimora commented Apr 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Template

Checklist

Related Issues/PRs

Changes

Testing

Uh oh!

codecov bot commented Apr 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

laggui left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

laggui left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

antimora commented Apr 30, 2025 •

edited

Loading

codecov bot commented Apr 30, 2025 •

edited

Loading