使用 Numpy 使用 einsum 和 tensordot 进行相同的操作答案

【问题标题】：Same operation with einsum and tensordot with Numpy使用 Numpy 使用 einsum 和 tensordot 进行相同的操作
【发布时间】：2017-01-25 09:48:34
【问题描述】：

假设我有两个 3D 数组 A 和 B，形状为 (3, 4, N) 和 (4, 3, N)。

我可以计算沿第三轴的切片之间的点积

with_einsum = np.eisum('ikl,kjl->ijl', A, B)

是否可以使用numpy.tensordot 执行相同的操作？

【问题讨论】：

标签： python numpy vectorization

【解决方案1】：

对于np.einsum('ikl,kjl->ijl', A, B)，存在带有字符串的轴对齐要求 - l 与输入和输出保持一致。因此，使用np.tensordot 可能不一定会提高性能，但由于问题已明确要求，所以无论如何我们建议它。现在，np.tensordot 会将不参与减和的轴分散为单独的轴，从而产生(N,N)。因此，为了得到最终输出，我们需要沿展开轴的对角线提取元素。

np.tensordot 的解决方案如下所示 -

mask = np.eye(N,dtype=bool)
out = np.tensordot(A,B,axes=((1),(0))).swapaxes(1,2)[:,:,mask]

在某些情况下np.dot/np.tensordot 可能会成为赢家，但这需要减和轴有合适的长度。如需更详细的分析，请参阅this post。

对于给定的问题，情况并非如此，因为 sum-reduction 的长度只是 4。所以，我认为einsum 在这里最好！

【讨论】：

那么，如果一个轴不参与操作（即我们只是沿着这个轴的切片执行相同的操作），np.einsum 是最适合的函数吗？
@floflo29 链接了另一个应该提供更多信息的帖子。也检查一下！
是否有一种通用的方法来代数编写用np.einsum 计算的表达式？我也会看看你的链接！
@floflo29 不知道代数写作是什么意思。你能详细说明一下吗？我认为einsum 使用的字符串表示法非常接近以循环方式实现的方式，因此非常直观，一旦我们了解每个字符串的含义。
例如，如果我们有C=np.einsum('ikl,kjl->ijl', A, B)，我们可以这样写：C(i,j,l) = sum(k)A(i,k,l)*B(k,j, l)