Fix matmul type_inference for 1-D inputs

coremltools/converters/mil/mil/ops/defs/iOS15/linear.py

@@ -188,19 +188,29 @@ def type_inference(self):
         x_shape = list(self.x.shape)
         y_shape = list(self.y.shape)
         x_rank = len(x_shape)
+        y_rank = len(y_shape)
 
         if x_rank == 1 and self.transpose_x.val:
             msg = "Op {} (matmul): x is rank 1, but transpose_x is True, which is not allowed."
             raise ValueError(msg.format(self.name))
+        if y_rank == 1 and self.transpose_y.val:
+            msg = "Op {} (matmul): y is rank 1, but transpose_y is True, which is not allowed."
+            raise ValueError(msg.format(self.name))
 
+        # Transpose only swaps the last two dims and is a no-op for rank-1 inputs.
         if self.transpose_x.val:
-            x_shape = list(x_shape)
             x_shape[-1], x_shape[-2] = x_shape[-2], x_shape[-1]
-            x_shape = tuple(x_shape)
         if self.transpose_y.val:
-            y_shape = list(y_shape)
             y_shape[-1], y_shape[-2] = y_shape[-2], y_shape[-1]
-            y_shape = tuple(y_shape)
+
+        # Follow numpy.matmul: a 1-D x is promoted to a matrix by prepending a 1,
+        # and a 1-D y is promoted by appending a 1. Both inserted dimensions are
+        # removed from the result. We track them so we drop the correct axes.
+        if x_rank == 1:
+            x_shape = [1] + x_shape
+        if y_rank == 1:
+            y_shape = y_shape + [1]
+
         if not (
             x_shape[-1] == y_shape[-2]
             or is_symbolic(x_shape[-1])
@@ -209,16 +219,18 @@ def type_inference(self):
             msg = "Op {} (matmul): x {}, y {} are not broadcastable"
             raise ValueError(msg.format(self.name, self.x.shape, self.y.shape))
 
-        if x_rank == 1:
-            # promote shape of x to rank 2
-            x_shape = list((1,) + tuple(x_shape))
         ret_shape = list(broadcast_shapes(x_shape[:-2], y_shape[:-2]))
         ret_shape += [x_shape[-2], y_shape[-1]]
+
+        # Remove the dimensions inserted above. The prepended dim of a 1-D x lands
+        # at index -2 of ret_shape and the appended dim of a 1-D y at index -1, so
+        # drop the higher index first to keep the lower one valid.
+        if y_rank == 1:
+            del ret_shape[-1]
         if x_rank == 1:
-            # remove the first dimension of the returned shape
-            return types.tensor(x_type, tuple(ret_shape[1:]))
-        else:
-            return types.tensor(x_type, tuple(ret_shape))
+            del ret_shape[-2 if y_rank != 1 else -1]
+
+        return types.tensor(x_type, tuple(ret_shape))
 
     @precondition(allow=VALUE)
     def value_inference(self):

coremltools/converters/mil/mil/ops/tests/iOS14/test_linear.py

@@ -283,6 +283,46 @@ def build(x):
             backend=backend,
         )
 
+    @pytest.mark.parametrize(
+        "shapes_and_transposes",
+        [
+            # A 1-D input follows numpy.matmul: a 1-D x is treated as a row vector
+            # and a 1-D y as a column vector, and the inserted dimension is removed
+            # from the result.
+            ((4,), (10, 4, 3), False, False),  # 1-D x, batched y -> (10, 3)
+            ((4,), (2, 5, 4, 3), False, False),  # 1-D x, higher-rank batched y
+            ((10, 3, 4), (4,), False, False),  # batched x, 1-D y -> (10, 3)
+            ((2, 5, 3, 4), (4,), False, False),
+            ((4,), (4,), False, False),  # both 1-D -> scalar
+            ((4,), (4, 3), False, False),  # 1-D x, 2-D y -> (3,)
+            ((3, 4), (4,), False, False),  # 2-D x, 1-D y -> (3,)
+            ((4,), (3, 4), False, True),  # transpose_y on the 2-D operand
+            ((4, 3), (4,), True, False),  # transpose_x on the 2-D operand
+            ((4,), (5, 3, 4), False, True),
+        ],
+    )
+    def test_builder_1d_input_type_inference(self, shapes_and_transposes):
+        # Type-inference regression test for https://github.com/apple/coremltools/issues/2263.
+        # Only the inferred output shape is checked, not execution: the Core ML runtime does not
+        # execute a rank-1 matmul (predict raises at the framework level, see issue #2263), so a
+        # 1-D matmul cannot go through run_compare_builder. The bug was purely in type_inference,
+        # which previously dropped batch dims (e.g. inferred (1, 3) instead of (10, 3)) or raised
+        # IndexError when y was 1-D.
+        shape_x, shape_y, transpose_x, transpose_y = shapes_and_transposes
+
+        @mb.program(
+            input_specs=[mb.TensorSpec(shape=shape_x), mb.TensorSpec(shape=shape_y)],
+            opset_version=ct.target.iOS15,
+        )
+        def prog(x, y):
+            return mb.matmul(x=x, y=y, transpose_x=transpose_x, transpose_y=transpose_y)
+
+        np_x = np.swapaxes(np.empty(shape_x), -1, -2) if transpose_x else np.empty(shape_x)
+        np_y = np.swapaxes(np.empty(shape_y), -1, -2) if transpose_y else np.empty(shape_y)
+        expected_shape = np.matmul(np_x, np_y).shape
+
+        assert tuple(prog.functions["main"].outputs[0].shape) == expected_shape
+
 
 class TestEinsum:
     @pytest.mark.parametrize(