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

Question

我有2个矩阵。说出尺寸为6 6 10的 A 和尺寸为6 10的 B 。我想做的是将 A 的10 6 6个矩阵与 B 的10个6元素向量相乘，得到6 * 10的矩阵。有没有一种方法可以不使用循环？

我想要的是

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] 

但没有for循环。

Answer 1

我们可以使用np.einsum-

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

或者，用np.matmul/@-operator-

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