Python numpy (einsum) 优化 : 1D to ND outer dot products

Question

如何最有效地计算一维数组的 2 ND 外积的点积？

“i,j,k,l->ij,kl->”将是一个很好的 einsum 替代方案，但失败了。

invalid subscript ',' in einstein sum subscripts string, subscripts must be letters

下面是我试图加速的简单实现(A 和 B 将是其他东西)。理想的情况是具有以下 API 的快速版本，其结果与以下示例相同:nd_outer_from1D(2, A, B)、nd_outer_from1D(3, A, B)。正如您将看到的，随着 |A|、|B| 和 N 增加存储结果并重新插入 einsum，因为参数很快变得不可行。

$  python -m timeit 'import numpy as np; A=np.arange(50); B=np.arange(50); sum(sum(np.outer(A,A) * np.outer(B,B)))'
10000 loops, best of 3: 72.1 usec per loop
$  python -m timeit 'import numpy as np; A=np.arange(50); B=np.arange(50); sum(sum(np.einsum("i,j->ij",A,A) * np.einsum("i,j->ij",B,B)  ))'
10000 loops, best of 3: 61.4 usec per loop
$  python -m timeit 'import numpy as np; A=np.arange(50); B=np.arange(50); sum(sum(sum(np.einsum("i,j,k->ijk",A,A,A) * np.einsum("i,j,k->ijk",B,B,B)  )))'
1000 loops, best of 3: 1.78 msec per loop

编辑(示例):

>>> A
array([0, 1, 2, 3])
>>> B
array([0.58394169, 0.22495002, 0.08322459, 0.05406281])
>>> sum(sum(np.einsum('i,j->ij',A,A) * np.einsum('i,j->ij', B, B)))
0.3064592592321492

显然 sum(sum( 和领先的 einsum 没有按我的预期工作。

对于以下专家来说小菜一碟，比较时间:

$  python -m timeit 'import numpy as np; A=np.arange(50); B=np.arange(50); np.einsum("i,i->",A,B)**3'
100000 loops, best of 3: 6.77 usec per loop
$  python -m timeit 'import numpy as np; A=np.arange(50); B=np.arange(50); np.einsum("i,i->",A,B)**2'
100000 loops, best of 3: 6.63 usec per loop
$  python -m timeit 'import numpy as np; A=np.arange(50); B=np.arange(50); A.dot(B)**3'
100000 loops, best of 3: 3.75 usec per loop
$  python -m timeit 'import numpy as np; A=np.arange(50); B=np.arange(50); A.dot(B)**2'
100000 loops, best of 3: 3.68 usec per loop

哇，这比我预期的要快:

$  python -m timeit 'import numpy as np; A=np.arange(5000); B=np.arange(5000); A.dot(B)**10'
100000 loops, best of 3: 12.1 usec per loop

Answer 1

这可以通过einsum进行优化 -

np.einsum("i,i->",A,B)**2

使用矩阵乘法 -

A.dot(B)**2