Optimize argument handling and improve ONNX graph building

[skewballfox]

Pull Request Template

Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

#1840

Changes

This was the alternative design I was talking about the other day. It's definitely a different set of compromises compared to #1795.

1. got rid of graph_io, and move renaming to the end, which mostly works
2. I altered how arguments are added to the map, so that it only runs on inputs (which mostly works).
3. Args are cloned into nodes, so functions that examine inputs or alter outputs don't need a second struct.

Problems

• in order to check if inputs are passed, since input names are altered at the beginning, we need to double map them (new_name->old->name->entry). This is both marginally better if a graph has 50+ inputs, and needed to check if the input is also an initializer(both are the case for web inference example).
• In order for unsqueeze to remap to a reshape(if the shape of the lhs is unknown or symbolic), it needs to check if the rhs is a graph output, which right now feels a little bolted on because nothing else needs access to that.

Testing

testing with existing onnx test, and the minst classification web example. passes the test, both work as expected.

I figured out why the example was failing before. All but the first input are also initializers.

[codecov]

Codecov Report

Attention: Patch coverage is 95.74468% with 6 lines in your changes missing coverage. Please review.

Project coverage is 86.11%. Comparing base (a5af19b) to head (5fe61f6).
Report is 8 commits behind head on main.

Files	Patch %	Lines
crates/burn-import/src/onnx/from_onnx.rs	95.31%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1857      +/-   ##
==========================================
- Coverage   86.28%   86.11%   -0.18%     
==========================================
  Files         773      776       +3     
  Lines       89382    90402    +1020     
==========================================
+ Hits        77125    77847     +722     
- Misses      12257    12555     +298     

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

[antimora]

A few minor comments that I have. But once they're resolved, we can merge them.

I'll copy our offline communication here to convey the benefits this PR will bring. I presume this is one in a series of refactorings to make the ONNX IR stand-alone.

[antimora]

A copy of our offline conversation about the benefits of these changes:

can you please summarize this?

[4:10 PM]dt: Could you remind me again what's the ultimate benefit we get after these changes? More efficiency? Better for future refactoring? So I am and others are onboard?
[6:04 PM]whoami(self,null): A little of both. The separation pr brings the number of copies of args from 3 to 1, eliminates most places were values were cloned, and tries to "enforce" correct use with panics (parsing would fail rather than perform an invalid operation). this one brings it to 2, but avoids some of the complexity of the former.

in both, I guess I'm trying to get the api for burn-import "just right" before starting a PR for turning it into a separate library. I suspect there are edge cases (like remapping unsqueeze), that are going to appear as more of the spec gets implemented, and I'm trying to account for the addition of custom behavior to burn-import(if/when that happens). Ideally, it should be easy to add new fuctions and checks without causing another step to fail (unless they do something stupid like skip a node unintentionally)
[6:10 PM]whoami(self,null): like right now, if someone adds a function to modify the output of a node after dim inference but doesn't explicity sync it to graph_io, that change won't show up in the copy given as an input to the next node.

[skewballfox]

can you please summarize this?

This PR reduces the number of copies of the arguments from 3 to 2, and removes the need to sync changes to arguments between 2 different data structures. It contains a few other minor changes like making the graph_builder actually produce an onnx graph. Part of the goal is just making onnx parsing easier to add to and debug.

[antimora]

LGTM!

Thanks for your contribution!