diff --git a/onnxscript/function_libs/torch_lib/ops/nn.py b/onnxscript/function_libs/torch_lib/ops/nn.py index a37c48f259..24571d0965 100644 --- a/onnxscript/function_libs/torch_lib/ops/nn.py +++ b/onnxscript/function_libs/torch_lib/ops/nn.py @@ -822,11 +822,17 @@ 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. + if bias is not None: + raise NotImplementedError("aten::linear with 1d weight and bias is not supported") + 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,