Fix aten_linear rank mismatch for 1D weight (missing squeeze) - #2984
Fix aten_linear rank mismatch for 1D weight (missing squeeze)#2984gabrielfruet wants to merge 2 commits into
Conversation
The 1D-weight branch unsqueezed the weight and ran MatMul but never squeezed the added dim back off, leaving the traced graph's actual output rank one higher than the shape the exporter records for it. onnx.checker.check_model(..., full_check=True) flags this as a rank mismatch (pytorch/pytorch#191332). aten::linear rejects a bias together with a 1D weight at the kernel level, so that combination is unreachable and isn't handled specially. Fixes microsoft#2982
|
@microsoft-github-policy-service agree |
There was a problem hiding this comment.
Pull request overview
This PR fixes an ONNX export rank mismatch for aten::linear when weight is 1D by ensuring the extra dimension introduced for MatMul is removed before returning, aligning traced output rank with eager PyTorch behavior.
Changes:
- Update
aten_lineartoSqueezethe last dimension afterMatMulin the 1D-weight path. - Add dedicated OpInfo-based test coverage that generates 1D-weight
linearsamples. - Wire the new OpInfo into the TorchLib operator test matrix.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxscript/function_libs/torch_lib/ops/nn.py | Fixes 1D-weight aten::linear to drop the extra (…, 1) dim after MatMul. |
| tests/function_libs/torch_lib/extra_opinfo.py | Adds sample generator + OpInfo to ensure 1D-weight linear cases are exercised. |
| tests/function_libs/torch_lib/ops_test_data.py | Registers the new OpInfo name so it runs through the TorchLib test suite. |
| weight_transposed = op.Unsqueeze(weight, [1]) | ||
| else: | ||
| assert len(weight.shape) == 2 | ||
| weight_transposed = op.Transpose(weight, perm=[1, 0]) | ||
| return op.Squeeze(op.MatMul(input, weight_transposed), [-1]) |
There was a problem hiding this comment.
Good catch, fixed — raises NotImplementedError now instead of silently dropping bias.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2984 +/- ##
=======================================
Coverage 72.63% 72.64%
=======================================
Files 265 265
Lines 32204 32205 +1
Branches 3041 3041
=======================================
+ Hits 23391 23394 +3
+ Misses 7779 7778 -1
+ Partials 1034 1033 -1 ☔ View full report in Codecov by Harness. |
|
I'll try to take a look at the CI errors today |
|
Apparently the CI errors are not related to my changes. If I'm mistaken please let me know :) |
Fixes #2982.
aten_linear's 1D-weight branch (added in #2339/#2340) unsqueezes the weight to(in_features, 1)and does aMatMul, but never squeezed the added dim back off before returning. Eageraten::linearwith a 1D weight contracts the last dim away entirely, so the traced graph's actual output rank was one higher than the shape the exporter records for it —onnx.checker.check_model(..., full_check=True)flags it as a rank mismatch (see pytorch/pytorch#191332).Also confirmed
aten::linearrejects a bias together with a 1D weight at the kernel level (mat2 must be a matrix, got 1-D tensor), so that combination is unreachable and isn't handled specially.Added
ops.aten.linear.1d_weighttest coverage inextra_opinfo.py/ops_test_data.py, since the sharednn.functional.linearOpInfo never generates a 1D-weight sample. Verified it reproduces the exactonnx.checkerfailure against the pre-fix code (all 6 subtests fail withShapeInferenceError: ... rank: (2) vs (1)), and passes after the fix, with no regressions in the existing 42linear-related subtests.