Fix aten_linear rank mismatch for 1D weight (missing squeeze)

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

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

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,