Ленивые вычисления numpy .einsum, чтобы избежать хранения промежуточных больших размерных массивов в памяти - PullRequest
Новая
       100

Ленивые вычисления numpy .einsum, чтобы избежать хранения промежуточных больших размерных массивов в памяти

6 голосов
/ 01 мая 2020

Представьте, что у меня есть целые числа, 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, так что никакой трехмерный массив никогда не создается. Отлично!

Теперь давайте немного усложним ситуацию: 

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 без построения промежуточных трехмерных массивов. такие как C и C2.

Вопросы:

  1. существует ли np.einsum один вкладыш, который будет составлять E, без создания промежуточных 3d-массивов C и C2?
    Появляется следующее работать, расширившись до четырех терминов, но довольно непрактично по сравнению с гипотетическим API в вопросе 2 ... 
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_lazy следующее будет создавать E без сохранения трехмерного массива (например, 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
...