numpy dot product for tensors (3d times 2d)

Question

Currently I use

Na = (3, 2, 4)
Nb = Na[1:]
A = np.arange(np.prod(Na)).reshape(Na)
b = np.arange(np.prod(Nb)).reshape(Nb)

I want to calculate:

r = np.empty((A.shape[0], A.shape[2])
for i in range(A.shape[2]):
    r[:, i] = np.dot(A[:, :, i], b[:, i])

In words: A is a "stack of 4 matrices" (shaped (3,2)), i.e. a 3d-array, b is a "stack of 4 vectors" (shaped (3,)), i.e. a 2d-array. The desired result is a "stack of 4 matrix-vector-products", i.e. a stack of vectors, i.e. again a 2d-array (shaped (3, 4)).

I had a medium-depth look to np.einsum and np.tensordot but any solution that I constructed with these was at least as long and less readable than my loop-solution.

However I think there should be a one-liner for that simple problem.

Answer 1

I'm not sure why np.einsum didn't work for you. This should do the trick:

r = np.einsum('ijk,jk->ik', A, b)

I consider this to be very readable, because it exactly reflects the mathematical formula you've given in your question.