[Lang] Fixes matrix-vector multiplication

python/taichi/lang/matrix.py

@@ -377,6 +377,10 @@ def _get_entry_to_infer(self, arr):
     return _MatImpl() if is_matrix else _VecImpl()
 
 
+def is_vector(x):
+    return isinstance(x, Vector) or getattr(x, "ndim", None) == 1
+
+
 @_gen_swizzles
 class Matrix(TaichiOperations):
     """The matrix class.
@@ -541,6 +545,11 @@ def __matmul__(self, other):
 
         """
         assert isinstance(other, Matrix), "rhs of `@` is not a matrix / vector"
+        if is_vector(self) and not is_vector(other):
+            # left multiplication
+            assert self.n == other.m, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
+            return other.transpose() @ self
+        # right multiplication
         assert self.m == other.n, f"Dimension mismatch between shapes ({self.n}, {self.m}), ({other.n}, {other.m})"
         entries = []
         for i in range(self.n):
@@ -550,6 +559,9 @@ def __matmul__(self, other):
                 for k in range(1, other.n):
                     acc = acc + self(i, k) * other(k, j)
                 entries[i].append(acc)
+        if is_vector(other) and other.m == 1:
+            # TODO: remove `ndim` when #5783 is merged
+            return Vector(entries, ndim=1)
         return Matrix(entries)
 
     # host access & python scope operation

tests/python/test_matrix.py

@@ -115,7 +115,7 @@ def func(t: ti.i32):
         m += ti.Matrix([[3, 4], [5, t]])
         print(m @ v)
         print(r.x, r.y, r.z, r.w)
-        s = w.transpose() @ m
+        s = w @ m
         print(s)
         print(m)
 
@@ -699,3 +699,21 @@ def bar():
     with pytest.raises(TaichiCompilationError,
                        match=r'Expected 2 indices, got 1'):
         bar()
+
+
+@test_utils.test(arch=get_host_arch_list(), debug=True)
+def test_matrix_vector_multiplication():
+    mat = ti.math.mat3(1)
+    vec = ti.math.vec3(3)
+    r = mat @ vec
+    for i in range(3):
+        assert r[i] == 9
+
+    @ti.kernel
+    def foo():
+        mat = ti.math.mat3(1)
+        vec = ti.math.vec3(3)
+        r = mat @ vec
+        assert r[0] == r[1] == r[2] == 9
+
+    foo()