numpy - 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)

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

现在假设我可以存储大小为 (n,n) 的内存数组。 , (q,n)但不是任何三维数组，例如形状 (n,n,q) .我无法存储在内存数组中 C以上。相反，要计算 D ,

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

工作正常，np.einsum通常足够聪明，可以找到 einsum_path这样就不会构造 3d 数组。伟大的!

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

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。

问题:

有没有np.einsum一个可以 build 的类轮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!

np.einsum 是否有“懒惰”版本这将在最终调用之前等待以找到最佳 einsum_path贯穿几个懒惰einsum的组成，包括上面例子中的sums？例如，假设 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简单地返回一个 numpy ndarray对象 - 这正是对 einsum 的后续调用期望作为您方案中的参数。但是，您可以使用 generators 来利用 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 的参数.事实上，most numpy ufunc s happen to have an out parameter ，它允许您指定“预先存在的”numpy ndarray作为操作结果的目标。因此，不需要分配新的内存，这也加速了您的计算作为副作用。