Add shape ONNX op support

[laggui]

Related Issues/PRs

Progress towards #1544

Changes

Added shape ONNX op import support

Testing

Added unit tests for codegen and expected outputs

[laggui]

Didn't realize torch always exports the shape op with a cast.. that's causing the CI to fail.

PR #1634 will fix that, but will fix the included graph for this PR.

[codecov]

Codecov Report

Attention: Patch coverage is 88.54167% with 11 lines in your changes are missing coverage. Please review.

Project coverage is 86.34%. Comparing base (6d96e8d) to head (3df2b78).

❗ Current head 3df2b78 differs from pull request most recent head 6a8ed3e. Consider uploading reports for the commit 6a8ed3e to get more accurate results

Files	Patch %	Lines
crates/burn-import/src/onnx/op_configuration.rs	56.00%	11 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1639      +/-   ##
==========================================
+ Coverage   86.30%   86.34%   +0.04%     
==========================================
  Files         692      691       -1     
  Lines       79998    79523     -475     
==========================================
- Hits        69040    68666     -374     
+ Misses      10958    10857     -101     

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

[antimora]

One left over comment change. Approving ahead.

We might have to revisit the output type. I know in the standard arrays are treated as 1d tensors but it would worth if we use shape output type because when we will have to feed into a reshape operation, it would be convenient to receive in rust array type.

For now, if your solution works for your use case, we should use it till we encounter other use cases, such as reshape. With ONNX, I am taking a conservative approach because ONNX spec is too complex to implement at once.

[antimora]

Flatten => Shape

[laggui]

We might have to revisit the output type. I know in the standard arrays are treated as 1d tensors but it would worth if we use shape output type because when we will have to feed into a reshape operation, it would be convenient to receive in rust array type.

For now, if your solution works for your use case, we should use it till we encounter other use cases, such as reshape. With ONNX, I am taking a conservative approach because ONNX spec is too complex to implement at once.

I initially thought about returning a shape array but then when I started testing I noticed that most of the time in an ONNX model the shape for a dimension is retrieved with gather to use the resulting scalar on other operations. So by returning a 1d tensor as specified in the standard our gather import should work for this sequence too.

But I agree, importing ONNX is not straightforward.