Numpy 3darray matrix multiplication function

问题

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?

回答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.

回答2:

np.dot should be your friend. http://docs.scipy.org/doc/numpy-1.10.0/reference/generated/numpy.dot.html

回答3:

It's done using np.tensordot

ans=np.tensordot(W,C,axes=[2,1])[np.arange(m),:,np.arange(m)]
assert np.all(result==ans)

来源：https://stackoverflow.com/questions/37738401/numpy-3darray-matrix-multiplication-function