Numpy:这个einsum表达是什么意思,还有另一种选择吗?

时间:2017-11-03 09:01:30

标签: python numpy

我如何理解以下表达式(A是[200,2]数组):

B = numpy.einsum('...i,...j->...ij',A,A)

如何以不使用numpy.einsum的其他方式编写它?

1 个答案:

答案 0 :(得分:2)

对于所有对,基本上在A的最后一个轴的元素之间进行元素相乘。现在,既然einsum表达式没有进行任何总和减少而只是完成broacasting的工作,我们可以通过手动扩展输入数组以获得一个维度来避免它广播做它的工作。这比使用einsum更有效,而且这是寻找einsum的替代方案的一个很好的理由。

2D甚至n-dim A(使用省略号表示法)的实现将是 -

A[...,None,:]*A[...,None]

示例运行 -

In [71]: A = np.random.rand(3,4,5,6)

In [72]: np.allclose(np.einsum('...i,...j->...ij',A,A), A[...,None,:]*A[...,None])
Out[72]: True

当扩展轴具有合适的长度时,前面讨论的性能方面更有意义。因此,A形状(200,2),即2作为最后一个轴的轴长度,使用broadcasting扩展的改进将更少/不在 - 一切都很明显,但只看到不错的长度。让我们验证这些方面 -

In [76]: A = np.random.rand(20000,2)

In [77]: %timeit np.einsum('...i,...j->...ij',A,A)
1000 loops, best of 3: 207 µs per loop

In [78]: %timeit A[...,None,:]*A[...,None]
1000 loops, best of 3: 364 µs per loop

In [79]: A = np.random.rand(200,200)

In [80]: %timeit np.einsum('...i,...j->...ij',A,A)
100 loops, best of 3: 12.1 ms per loop

In [81]: %timeit A[...,None,:]*A[...,None]
100 loops, best of 3: 9.74 ms per loop