From 6128059bc87bba0e2da3bc08158a698f6c7b5b88 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 30 Jul 2026 23:25:20 -0300 Subject: [PATCH 1/2] Fix aten_linear rank mismatch for 1D weight 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 #2982 --- onnxscript/function_libs/torch_lib/ops/nn.py | 11 ++++--- tests/function_libs/torch_lib/extra_opinfo.py | 33 +++++++++++++++++++ .../function_libs/torch_lib/ops_test_data.py | 5 +++ 3 files changed, 45 insertions(+), 4 deletions(-) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index a37c48f259..e790467d27 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -822,11 +822,14 @@ def aten_linear(input: TFloat, weight: TFloat, bias: Optional[TFloat] = None) -> # 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 + # In rare cases the weight can be 1d. aten::linear does not support a + # bias with a 1d weight (mat2 must be a matrix). eager linear contracts + # the last dim away in this case, so squeeze the dim added for MatMul + # back off. 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]) + 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 diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index d925be6877..43779d59cd 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/tests/function_libs/torch_lib/extra_opinfo.py @@ -37,6 +37,31 @@ def sample_inputs_scalar_tensor(op_info, device, dtype, requires_grad, **kwargs) yield opinfo_core.SampleInput(item, dtype=dtype) +def sample_inputs_linear_1d_weight(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for linear with a 1D weight (in_features,), no out_features dim. + + Regression coverage for https://github.com/microsoft/onnxscript/issues/2982: + the upstream `sample_inputs_linear` in + torch/testing/_internal/common_methods_invocations.py never generates a + 1D-weight case, so this branch of aten_linear went untested. + + Note: aten::linear does not accept a bias together with a 1D weight + ("mat2 must be a matrix, got 1-D tensor"), so no bias samples are + generated here. + """ + del op_info + del kwargs + + make_arg = functools.partial( + torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + + for in_features, batch_shape in itertools.product([3, 8], [(), (2,), (2, 3)]): + input_tensor = make_arg((*batch_shape, in_features)) + weight = make_arg((in_features,)) + yield opinfo_core.SampleInput(input_tensor, args=(weight, None)) + + def sample_inputs_bilinear(op_info, device, dtype, requires_grad, **kwargs): """Sample inputs for bilinear operation.""" del op_info @@ -2506,6 +2531,14 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs sample_inputs_func=sample_inputs_bilinear, supports_out=False, ), + opinfo_core.OpInfo( + "ops.aten.linear.1d_weight", + op=torch.nn.functional.linear, + aten_name="linear", + dtypes=common_dtype.floating_types(), + sample_inputs_func=sample_inputs_linear_1d_weight, + supports_out=False, + ), opinfo_core.OpInfo( "ops.aten.bernoulli.p", aten_name="bernoulli.p", diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index c4f47f2097..62ad59c446 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -1652,6 +1652,11 @@ def _where_input_wrangler( TorchLibOpInfo( "nn.functional.linear", nn_ops.aten_linear, tolerance={torch.float16: (1e-2, 1e-3)} ), + TorchLibOpInfo( + "ops.aten.linear.1d_weight", + nn_ops.aten_linear, + tolerance={torch.float16: (1e-2, 1e-3)}, + ), TorchLibOpInfo( "nn.functional.unfold", nn_ops.aten_im2col, From 88766e9283ce7e1fb3c1960159cc0f787874ed28 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 31 Jul 2026 10:17:41 -0300 Subject: [PATCH 2/2] fix: raise on 1d-weight linear with bias --- onnxscript/function_libs/torch_lib/ops/nn.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index e790467d27..24571d0965 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -826,6 +826,9 @@ def aten_linear(input: TFloat, weight: TFloat, bias: Optional[TFloat] = None) -> # bias with a 1d weight (mat2 must be a matrix). eager linear contracts # the last dim away in this case, so squeeze the dim added for MatMul # back off. + if bias is not None: + raise NotImplementedError("aten::linear with 1d weight and bias is not supported") + weight_transposed = op.Unsqueeze(weight, [1]) return op.Squeeze(op.MatMul(input, weight_transposed), [-1]) assert len(weight.shape) == 2