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

Question

我有 2 个矩阵。假设大小为 6610 的 A 和大小为 610 的 B。我想要做的是将 A 的 10 66 个矩阵与 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

Multiply a 3D matrix with 2D matrix to give a 2D matrix