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

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)

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",

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,