将 3D 矩阵与 2D 矩阵相乘得到 2D 矩阵

Question

I have 2 matrices. Say A of size 6 6 10 and B of size 6 10. What I want to do is multiply the 10 6 6 matrices of A with the 10 6 element vectors of B to give a 6*10 matrix. Is there a way to do this without using a loop?

What I want is

A = np.ones((6,6,10))
B = np.ones((6,10))
mat = np.zeros((6,10))
for i in range(10):
    mat[:,i] = A[:,:,i]@B[:,i] 

but without the for loop.

Answer 1

We can use np.einsum -

mat = np.einsum('ijk,jk->ik',A,B)

Alternatively, with np.matmul/@-operator -

mat = (A.transpose(2,0,1)@B.T[:,:,None])[...,0].T