Python vectorization of matrix-vector operation

Question

I have a Matrix A with shape (2,2,N) and a Matrix V with shape (2,N)

I want to vectorize the following:

F = np.zeros(N)
for k in xrange(N):
    F[k] = np.dot( A[:,:,k], V[:,k] ).sum()

Any way this can be done with either tensordot or any other numpy function without explicit looping?

Answer 1

With np.einsum -

F = np.einsum('ijk,jk->k',A,V)

We can optimize it further with optimize flag (check docs) set as True .