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

[gabrielfruet]

aten_linear in function_libs/torch_lib/ops/nn.py:817-833:

@torch_op("aten::linear", trace_only=True)
def aten_linear(input: TFloat, weight: TFloat, bias: Optional[TFloat] = None) -> TFloat:
    """linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor"""

    if len(input.shape) == 2 and len(weight.shape) == 2:
        # Use Gemm for the rank 2 input
        return op.Gemm(input, weight, bias, transB=True)
    if len(weight.shape) == 1:
        # In rare cases the weight can be 1d
        weight_transposed = op.Unsqueeze(weight, [1])
    else:
        assert len(weight.shape) == 2
        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)

When weight is 1D, Unsqueeze(weight, [1]) gives shape (in_features, 1). MatMul(input, weight_transposed) then produces shape input.shape[:-1] + (1,) — rank equal to input.rank. When bias is None, that mul is returned directly, without squeezing the added dim back off.

aten::linear with a 1D weight is a dot-product contraction in eager PyTorch: it drops the last dim entirely (input.shape[:-1], rank input.rank - 1). So the returned value's actual rank is one higher than what the op is supposed to produce.

Consumers that trust the exporter's recorded shape (e.g. onnx.checker.check_model(model, full_check=True)) then see a rank mismatch on this node, since exporters like torch.onnx.export(..., dynamo=True) stamp the eager (correct, lower) rank onto a value that structurally still has the extra dim. Reported as pytorch/pytorch#191332.

This is a residual gap from #2339 / #2340, which added the Unsqueeze+MatMul 1D-weight handling but didn't add the corresponding squeeze-back for the no-bias case.

Fix: squeeze the added dim back off before returning, in both the bias and no-bias case:

if len(weight.shape) == 1:
    weight_transposed = op.Unsqueeze(weight, [1])
    result = op.Squeeze(op.MatMul(input, weight_transposed), [-1])
    if bias is None:
        return result
    return op.Add(result, bias)

Also, the shared OpInfo used to test this op (nn.functional.linear, via PyTorch's sample_inputs_linear) never generates a 1D-weight sample, so this branch has no test coverage today.

[gabrielfruet]

Justin Chu (@justinchuby) Hey, I would really like to work on this. I'm trying to contribute more to ONNX and PyTorch :)