Squeeze Onnx Import

[agelas]

Pull Request Template

Checklist

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

Related Issues/PRs

#1714

Changes

Added the squeeze operation for to burn-import

Testing

• onnx_tests
• codegen_nodes test

[codecov[bot]]

Codecov Report

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

Project coverage is 86.41%. Comparing base (1073752) to head (3e60613). Report is 7 commits behind head on main.

Files	Patch %	Lines
crates/burn-import/src/onnx/dim_inference.rs	65.51%	10 Missing :warning:
crates/burn-import/src/onnx/op_configuration.rs	77.14%	8 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1753      +/-   ##
==========================================
- Coverage   86.61%   86.41%   -0.21%     
==========================================
  Files         700      735      +35     
  Lines       83423    85729    +2306     
==========================================
+ Hits        72258    74083    +1825     
- Misses      11165    11646     +481     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[mtobin-tdab]

Hi there, aware this is closed but wanted to ask why Squeeze has to specify only one axis? Is there a plan to make it possible to Squeeze to take > 1 axis? Should I make an issue to track this? In the mean time would it work to split into two layers which each Squeeze a different axis? Sorry I am not an expert in this field

  DEBUG onnx_ir::from_onnx: checking node squeeze1 for constants    
  DEBUG onnx_ir::from_onnx: checking input Argument { name: "const_axes__420", ty: Tensor(TensorType { elem_type: Int64, dim: 1, shape: Some([2]) }), value: Some(Int64s([2, 3])), passed: false } for const    
  ERROR burn_import::logger: PANIC => panicked at crates/onnx-ir/src/dim_inference.rs:463:9:
  Squeeze must specify only 1 axis, found 2

[laggui]

Hi @mtobin-tdab 👋

The ONNX spec is quite large so we don't necessarily aim to cover all the edge cases when adding support for an operator. That's why this PR was merged - we just make sure to panic when there's is a case which is not supported yet (such as yours).

You can definitely open an issue to ask for support for multiple axes. That's what issues are for (bugs, feature requests, etc.).

In the meantime, if you can generate your model with multiple squeeze nodes instead then yes that should work!

[mtobin-tdab]

Thank you @laggui !

[agelas]

@mtobin-tdab I think squeeze_dims added here in https://github.com/tracel-ai/burn/pull/1779 is what you might be looking for.

[agelas]

Ah wait, actually I think I see the issue. It looks like I didn't update the dimension inference, sorry about that!

[laggui]

Oh I didn't even see the follow-up PR for squeeze_dims 😄 looks like you've correctly identified the issue