3D Array Multiplication

Question

I have $2$ matrices ($A$ with shape $(5,2,3)$ and $B$ with shape $(5,3,8)$), and I want to perform some kind of multiplication in order to take a new matrix with shape $(5,2,8)$.

Pseudo Code:

 for i= 0 to A.dim[0] do:
       C[i] = A[i].dot(B[i])

Is it possible to do the above operation without using a loop?

Answer 1

The correct Python syntax would be for i in range(A.shape[0]) and would use matmul instead of dot, but you don't want the for loop anyway. You could write C=np.array([a.matmul(b) for a, b in zip(A, B)]), which is a declarative comprehension rather than an imperative for loop. But better still is np.tensordot, C=np.tensordot(A, B, axes=((0, 2), (0, 1))).

Alternatively, another answer at the same link recommends C=np.einsum('nmk,nkj->nmj', A, B).