Squeeze Onnx Import #1753
Squeeze Onnx Import #1753
Conversation
Codecov ReportAttention: Patch coverage is
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 ☔ View full report in Codecov by Sentry. |
| match key.as_str() { | ||
| "axes" => return value.clone().into_i64s(), | ||
| _ => {} |
There was a problem hiding this comment.
Great! This will support ONNX OPset 13
Can we, to be sure it works, add a unit test for OPset 13? This is similar to unsqueeze with opset 16 and 13.
There was a problem hiding this comment.
Unsqueeze uses OPset 16 and 11. I think my original one used 16, so I added one for 13. Is that what you meant?
There was a problem hiding this comment.
Yes, both should be good. OpSet 13 has axes attribute so it should work.
| } | ||
| } | ||
| _ => panic!("Arg for squeeze must be tensor or scalar"), | ||
| } |
There was a problem hiding this comment.
axes can contain negative values, which means counting dimensions from the back (see squeeze spec), and Burn squeeze only supports positive values (see doc). So we should account for this.
We are already doing this for one dimension (see code) for gather_config, so you can see the logic. We need to do this for all items in axes.
There was a problem hiding this comment.
Thanks for your contribution! I have the same comments as already mentioned, and one question/comment regarding the squeeze op spec support.
This PR adds support for squeeze op on a single dim, so support is not complete but as long as it is explicit I have no issues adding full support if needed in another PR.
| if let Some(Data::Int64s(axes)) = &node.inputs[1].value { | ||
| if axes.len() != 1 { | ||
| panic!("Squeeze: Only one axis should be specified for squeezing."); | ||
| } | ||
| } else { | ||
| panic!("Squeeze: Axes input must be an integer list."); | ||
| }; |
There was a problem hiding this comment.
If we're checking for axes in the attributes in squeeze_config to support another opset we need to check here too to make sure the output is properly adjusted.
See the unsqueeze_update_output function which is doing something similar (though it captures the axes to support unsqueeze on multiple axes).
There was a problem hiding this comment.
Sounds good, should be a little closer to what you had in mind now!
Pull Request Template
Checklist
run-checks allscript has been executed.Related Issues/PRs
#1714
Changes
Added the squeeze operation for to
burn-importTesting