numpy - 沿轴的点积总和

numpy - Sum of dot products along axis

我有一个 A x B 数组和另一个 D x A x A 数组,我正在尝试想出有效的方法来计算两个数组沿 D 轴的点积之和(结果将是 A x B 数组)。最明显的方法是使用 for 循环:

result = np.zeros(first_array.shape)
for d in range(0,second_array.shape[0]):
    result = result + np.dot(second_array[d], first_array)
print result

我想知道在 numpy 中是否有更有效的计算方法。我读了一些 np.einsum 但不幸的是不完全理解它是否能够在这种情况下提供帮助。

In [436]: np.einsum('ijk,km->jm',np.ones((2,3,3)),np.ones((3,4)))
Out[436]: 
array([[ 6.,  6.,  6.,  6.],
       [ 6.,  6.,  6.,  6.],
       [ 6.,  6.,  6.,  6.]])
In [437]: _.shape
Out[437]: (3, 4)

您的 dot 表示为(k 在第一个的最后一个和第二个到第二个的最后一个之间共享):

In [438]: np.einsum('jk,km->jm',np.ones((3,3)),np.ones((3,4)))
Out[438]: 
array([[ 3.,  3.,  3.,  3.],
       [ 3.,  3.,  3.,  3.],
       [ 3.,  3.,  3.,  3.]])

i 添加到第一个数组以匹配其 3d 形状。但是从 result 中省略它会告诉 einsum 对其值求和。没有求和

In [439]: np.einsum('ijk,km->ijm',np.ones((2,3,3)),np.ones((3,4))).shape
Out[439]: (2, 3, 4)