与 einsum 和 tensordot 与 Numpy 的相同操作 [英] Same operation with einsum and tensordot with Numpy

问题描述

假设我有两个形状为 (3, 4, N) 和 (4,3、N).

Let's say I have two 3D arrays A and Bof shape (3, 4, N) and (4, 3, N).

我可以用

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

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

Is it possible to perform the same operation with numpy.tensordot ?

推荐答案

With np.einsum('ikl,kjl->ijl', A, B), 有轴对齐要求字符串 - l 与输入和输出保持一致.因此，使用 np.tensordot 可能不一定会导致性能改进，但由于问题已经特别要求它，无论如何我们还是建议它.现在，np.tensordot 会将不参与和约简的轴展开为单独的轴，从而得到 (N,N).因此，为了获得最终输出，我们需要沿着展开轴的对角线提取元素.

With np.einsum('ikl,kjl->ijl', A, B), there is axis alignment requirement with string - l that stays with the inputs and the output. As such, using np.tensordot might not necessarily result in performance improvement, but since the question has specifically asked for it, let's suggest it anyway. Now, np.tensordot would spread out the axes that don't take part in sum-reduction as separate axes, resulting in (N,N). So, to get to the final output, we need to extract the elements along the diagonal of the spread-out axes.

以下是使用 np.tensordot 的解决方案的样子 -

Here's how a solution with np.tensordot would look like -

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

在某些情况下，np.dot/np.tensordot 可能会成为赢家，但这需要求和轴具有合适的长度.有关更详细的分析，请参阅这篇文章.

There would be cases where np.dot/np.tensordot might come out as winner, but that requires that the sum-reduction axes have decent lengths. See this post for a more detailed analysis.

对于给定的问题，情况并非如此，因为求和的长度仅为 4.所以，我认为 einsum 最好在这里！

For the given problem, that's not the case as the length of sum-reduction is just 4. So, I would think einsum would be best here!

这篇关于与 einsum 和 tensordot 与 Numpy 的相同操作的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！