Numpy 3darray matrix multiplication function

Question

Suppose I have an ndarray, W of shape (m,n,n) and a vector C of dimension (m,n). I need to multiply these two in the following way

result = np.empty(m,n)
for i in range(m):
    result[i] = W[i] @ C[i]

How do I do this in a vectorized way without loops and all?

Answer 1

Since, you need to keep the first axis from both W and C aligned, while loosing the last axis from them with the matrix-multiplication, I would suggest using np.einsum for a very efficient approach, like so -

np.einsum('ijk,ik->ij',W,C)

np.tensordot or np.dot doesn't have the feature to keep axes aligned and that's where np.einsum improves upon.