合并Einsum表达式

Question

我想评估一下

E = np.einsum('ij,jk,kl->ijkl',A,A,A)
F = np.einsum('ijki->ijk',E)

其中A是矩阵（大小不超过1000乘1000）。计算速度很慢。我只想通过计算存储在F中的“对角线”元素来加速这个过程。有可能把这两个表达式结合起来吗？/有没有更好的方法来加速这个计算？你知道吗

Answer 1

In [86]: A=np.random.randint(0,100,(100,100))

In [88]: E1=np.einsum('ijki->ijk',np.einsum('ij,jk,kl->ijkl',A,A,A))

In [89]: E2=np.einsum('ij,jk,ki->ijk',A,A,A)

In [90]: np.allclose(E1,E2)
Out[90]: True

良好的时间改进-100x，对应于保存的维度（l）

In [91]: timeit np.einsum('ijki->ijk',np.einsum('ij,jk,kl->ijkl',A,A,A))
1 loops, best of 3: 1.1 s per loop

In [92]: timeit np.einsum('ij,jk,ki->ijk',A,A,A)
100 loops, best of 3: 10.9 ms per loop

einsum对所有索引执行组合迭代，尽管是在Cython代码中。因此，减少指数的数量可以节省大量的时间。看起来这样做i...i组合在初始计算中起作用

由于只有2g内存，（10001000）太大，E1情况下是“迭代器太大”，E2情况下是“内存错误”。你知道吗

Answer 2

我不确定是否有一种自动的方法，但你可以自己做数学，然后给出einsum最后的表达式：

F = np.einsum('ij,jk,ki->ijk', A, A, A)