张量的numpy dot积(3d次2d)

Question

目前我用

Na = (3, 2, 4)
Nb = Na[1:]
A = np.arange(np.prod(Na)).reshape(Na)
b = np.arange(np.prod(Nb)).reshape(Nb)

我想计算:

r = np.empty((A.shape[0], A.shape[2])
for i in range(A.shape[2]):
    r[:, i] = np.dot(A[:, :, i], b[:, i])

在文字中:A是"4个矩阵的堆叠"(形状(3,2)),即3d-阵列,b是"4个矢量的堆叠"(形状(3,)),即2d阵列.期望的结果是"4个矩阵 - 矢量 - 产物的堆叠",即矢量堆栈,即同样是2d阵列(形状(3,4)).

我对np.einsum和np.tensordot有一个中等深度的外观,但我用这些构建的任何解决方案至少与我的循环解决方案一样长且不易读.

但是我觉得这个简单问题应该有一个单线程.

Answer 1

我不确定为什么np.einsum不适合你.这应该做的伎俩:

r = np.einsum('ijk,jk->ik', A, b)

我认为这是非常易读的,因为它完全反映了你在问题中给出的数学公式.