Suppose that I have a 3d array A and a 2d array B. A has dimension (s,m,m) while B has dimension (m,s).
I want to write code for a 2d array C with dimension (m,s) such that C[:,i] = A[i,:,:] @ B[:,i].
Is there a way to do this elegantly without using a for loop in numpy?
One solution I thought of was to reshape B into a 3d array with dimension (m,s,1), multiply A and B via A@B, then reshape the resulting 3d array into a 2d array. This sounds a bit tedious and was wondering if tensordot or einsum can be applied here.
Suggestions appreciated. Thanks!
Using einsum is straight forward here:
A = np.arange(18).reshape(2,3,3)
B = np.arange(6).reshape(3,2)
C = np.einsum("ijk,ki->ji",A,B)
for i in range(2):
A[i]@B[:,i]==C[:,i]
# array([ True, True, True])
# array([ True, True, True])
Thanks so much! This looks very clean and efficient.