numpy.einsum 的惰性求值以避免在内存中存储中间大维数组

Question

想象一下，我有整数n,q和具有以下维度的向量/数组：

import numpy as np
n = 100
q = 102
A = np.random.normal(size=(n,n))
B = np.random.normal(size=(q, ))

C = np.einsum("i, jk -> ijk", B, A)
D = np.einsum('ijk, ikj -> k', C, C)

如果所有中间数组都适合内存，则工作正常。

现在假设我可以存储 size 的内存数组(n,n)，(q,n)但不能存储任何三维数组，例如 shape (n,n,q)。我无法存储在C上面的内存数组中。相反，要计算D，

D1 = np.einsum('i, jk, i, kj -> k', B, A, B, A, optimize='optimal')

工作正常，np.einsum通常足够聪明，可以找到一个einsum_path3d 数组。伟大的！

现在让我们稍微复杂一点：

C = np.einsum("i, jk -> ijk", B, A)    # as before

Y2 = np.random.normal(size=(n, ))
Z2 = np.random.normal(size=(q, n))
C2 = np.einsum("j, ik -> ijk", Y2, Z2)

E = np.einsum('ijk, ikj -> k', C+C2, C+C2)

在这里，我找不到一种合理的方法（合理的，如简短/可读的代码）来构造E而不构造中间 3d 数组，例如 C 和 C2。

问题：

1. 是否有np.einsum一个衬垫可以构造E，而无需构造中间 3d 数组 C 和 C2？
   以下内容似乎通过扩展为四个术语来工作，但与问题 2 中的假设 API 相比相当不切实际......

E_CC   = np.einsum('i, jk, i, kj -> k', B,  A,  B,  A, optimize='optimal') # as D before
E_C2C2 = np.einsum('j, ik, k, ij -> k', Y2, Z2, Y2, Z2, optimize='optimal')
E_CC2  = np.einsum('i, jk, k, ij -> k', B,  A,  Y2, Z2, optimize='optimal')
E_C2C  = np.einsum('j, ik, i, kj -> k', Y2, Z2, B,  A, optimize='optimal')

E_new  = E_CC + E_C2C2 + E_CC2 + E_C2C 
np.isclose(E_new, E) # all True!

2. 是否有一个“惰性”版本np.einsum会在最终调用之前等待以einsum_path在几个惰性 einsum 的整个组合中找到最佳值，包括上面示例中的总和？例如，对于一个假设的einsum_lazy，以下将构造E而不在内存中存储 3d 数组（例如 C 或 C2）：

C = np.einsum_lazy("i, jk -> ijk", B, A)  # nothing has been computed yet!
C2 = np.einsum_lazy("j, ik -> ijk", Y2, Z2) # nothing has been computed yet!
E = np.einsum('ijk, ikj -> k', C+C2, C+C2)  # expand the sums and uses optimal einsum_path to compute E 

Answer 1

目标问题2：

不幸的是einsum，没有惰性版本。einsum只是返回一个 numpyndarray对象 - 这正是后续调用所einsum期望的场景中的参数。但是，您可以通过使用生成器来利用 Python 本身。在你的情况下，以下方法可以解决问题：

C1 = (np.einsum_lazy("i, jk -> ijk", b, a) for a, b in ((A, B),))
C2 = (np.einsum_lazy("j, ik -> ijk", y2, z2) for y2, z2 in ((Y2, Z2),))

def _einsum(v, w):
    u = v + w # no need to do this twice
    return np.einsum('ijk, ikj -> k', u, u)

E = (_einsum(c1, c2) for c1, c2 in ((C1, C2),))

for e in E: # only HERE C1, C2 and E are actually computed
    print(e)

上面的示例使用了链式生成器表达式。这是最后一个 for 循环，它触发链的实际评估。或多或少是懒惰的。还有另一个缺点：从内存角度来看，C1 和 C2 实际上是（临时）构造/创建的。

如果内存消耗是您主要关心的问题，并且您正在执行多个类似的操作，则可以查看out的参数einsum。事实上，大多数 numpyufunc恰好都有一个out参数，它允许您指定一个“预先存在的”numpyndarray作为操作结果的目标。因此，不需要分配新的内存，这也会加速计算的副作用。