Skip to content

Fix aten_linear: squeeze trailing dim after 1D-weight MatMul - #2983

Draft
justinchuby with Copilot wants to merge 2 commits into
mainfrom
copilot/fix-aten-linear-bias-none
Draft

Fix aten_linear: squeeze trailing dim after 1D-weight MatMul#2983
justinchuby with Copilot wants to merge 2 commits into
mainfrom
copilot/fix-aten-linear-bias-none

Conversation

Copilot AI commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

When weight is 1D, aten_linear unsqueezed it to (in_features, 1) for MatMul but never squeezed the resulting trailing dim back off. This produced output rank input.rank instead of the correct input.rank - 1, mismatching PyTorch eager's dot-product contraction semantics.

Changes

  • onnxscript/function_libs/torch_lib/ops/nn.py: Apply op.Squeeze(..., [-1]) immediately after the 1D-weight MatMul, before the bias is None branch, so both paths return the correctly-ranked result:
if len(weight.shape) == 1:
    weight_transposed = op.Unsqueeze(weight, [1])
    mul = op.Squeeze(op.MatMul(input, weight_transposed), [-1])
else:
    weight_transposed = op.Transpose(weight, perm=[1, 0])
    mul = op.MatMul(input, weight_transposed)
if bias is None:
    return mul
return op.Add(mul, bias)
  • tests/function_libs/torch_lib/extra_opinfo.py / ops_test_data.py: Add nn.functional.linear_1d_weight OpInfo with sample_inputs_linear_1d_weight, covering 1D/2D/3D inputs with and without bias. The 2D-input+bias combination is intentionally excluded — PyTorch eager itself rejects it.

When weight is 1D, the previous code unsqueezed it to (in_features, 1)
for MatMul but never squeezed the resulting trailing dim back off.
This caused the output to have rank input.rank instead of the correct
input.rank - 1 (matching PyTorch eager dot-product semantics).

The fix squeezes axis [-1] immediately after the MatMul, so both the
no-bias and with-bias code paths return the correctly-shaped result.

Also adds a new OpInfo 'nn.functional.linear_1d_weight' with sample
inputs that specifically exercise the 1D-weight path (no-bias and
with-bias for 1-D, 2-D, and 3-D inputs), closing the coverage gap
noted in the issue.

Closes #2982
Copilot AI requested review from Copilot and removed request for Copilot July 30, 2026 22:47
Copilot AI changed the title [WIP] Fix aten_linear behavior when bias is None Fix aten_linear: squeeze trailing dim after 1D-weight MatMul Jul 30, 2026
Copilot AI requested a review from justinchuby July 30, 2026 22:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

aten_linear doesn't squeeze back the dim added for 1D-weight MatMul when bias is None

2 participants