Say I have an vector v(ij) and matrix M(jk).

v = np.array([1.0, 1.0, 0.0, 0.0], dtype=np.float32)

M = np.array([[0.0, 0.0, 0.2, 0.0],
              [0.0, 0.0, 0.0, -0.1],
              [0.0, 0.0, 0.0, 0.5],
              [0.0, 0.0, 0.0, 0.0]], dtype=np.float32)

How can I implement the following operation without for loop?

def matmul_in_place(vector, matrix):
    for i in range(vector.shape[0]):
        for j in range(matrix.shape[1]):
            vector[i] += vector[j] * matrix[j][i]

Result of matmul_in_place(v, M) should be [1.0, 1.0, 0.2, 0.0]

In place is not a requirement, instead, it means that when calculating vector[1], I should use updated vector[0] for vector[j] when j is 0. This also means that I need to calculate the vector from left to right.

Vector v can also be a matrix, but I want to keep the problem simple. As long as two matrices can be multiplied together(M1 @ M2), this operation should also work.

I've been looking into it for a while but I'm not familiar with einsum or custom cuda kernel so I don't even know if it's possible.