diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index a37c48f259..fca5ce1d9d 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -822,12 +822,15 @@ 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. Unsqueeze to (in_features, 1) for + # MatMul, then squeeze the added trailing dim back off so the result has + # rank input.rank - 1, matching PyTorch eager semantics. weight_transposed = op.Unsqueeze(weight, [1]) + mul = op.Squeeze(op.MatMul(input, weight_transposed), [-1]) else: assert len(weight.shape) == 2 weight_transposed = op.Transpose(weight, perm=[1, 0]) - mul = op.MatMul(input, weight_transposed) + mul = op.MatMul(input, weight_transposed) if bias is None: return mul return op.Add(mul, bias) diff --git a/tests/function_libs/torch_lib/extra_opinfo.py b/tests/function_libs/torch_lib/extra_opinfo.py index d925be6877..b12eecf6a6 100644 --- a/tests/function_libs/torch_lib/extra_opinfo.py +++ b/tests/function_libs/torch_lib/extra_opinfo.py @@ -2498,6 +2498,48 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs yield opinfo_core.SampleInput(self_tensor, args=(mask, source)) +def sample_inputs_linear_1d_weight(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for nn.functional.linear with a 1D weight tensor. + + PyTorch's built-in sample_inputs_linear never generates 1D-weight cases. + When weight is 1D (shape ``(in_features,)``), ``aten::linear`` performs a + dot-product contraction that drops the last dimension of *input*, producing + output of rank ``input.rank - 1``. This function exercises that path so + that the squeeze-back behaviour is covered by the op correctness tests. + + Note: the 2D-input + 1D-weight + bias combination is rejected by PyTorch + eager even today (``mat2 must be a matrix``), so bias samples are only + generated for 1-D and 3-D inputs where PyTorch itself succeeds. + """ + del op_info + del kwargs + + make_arg = functools.partial( + torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + + # Without-bias cases: all input ranks work. + for input_shape, in_features in [ + ((3,), 3), # 1-D input -> scalar output + ((2, 4), 4), # 2-D input -> 1-D output + ((2, 3, 5), 5), # 3-D input -> 2-D output + ]: + yield opinfo_core.SampleInput( + make_arg(input_shape), args=(make_arg((in_features,)),) + ) + + # With-bias cases: only 1-D and 3-D inputs are accepted by PyTorch eager + # when weight is 1-D (2-D input + bias raises a mat2-must-be-matrix error). + for input_shape, in_features in [ + ((3,), 3), # 1-D input -> scalar output + ((2, 3, 5), 5), # 3-D input -> 2-D output + ]: + bias = make_arg(()) # scalar bias, broadcast-compatible with any output + yield opinfo_core.SampleInput( + make_arg(input_shape), args=(make_arg((in_features,)), bias) + ) + + OP_DB: List[opinfo_core.OpInfo] = [ opinfo_core.OpInfo( "bilinear", @@ -2506,6 +2548,13 @@ def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs sample_inputs_func=sample_inputs_bilinear, supports_out=False, ), + opinfo_core.OpInfo( + "nn.functional.linear_1d_weight", + op=torch.nn.functional.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..ae137a3748 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( + "nn.functional.linear_1d_weight", + nn_ops.aten_linear, + tolerance={torch.float16: (1e-2, 1e-3)}, + ), TorchLibOpInfo( "nn.functional.unfold", nn_ops.aten_im2col,