[MRG] Add LazyTensor for large scale OT by rflamary · Pull Request #544 · PythonOT/POT
This PR adds a very simple LazyTensor class that can implement large scale matrices and tensors that are computed only lazily when getting slices of the tensor.
This class will be very interesting when implementing large scale OT problems such as factored OT, low rank OT and lazy implementation of sinkhorn. The objective is to be very simple with a storage of all necessary data and a function that compute the values on sliced (without computing the full tensor).
import numpy as np from ot.utils import LazyTensor, n1=100 n2=200 x1 = np.random.randn(n1,2) x2 = np.random.randn(n2,2) # i,j can be integers or slices, x1,x2 have to be passed as keyword arguments def getitem(i,j, x1, x2): return np.dot(x1[i], x2[j].T) # create a lazy tensor (giv data as arguments) T = LazyTensor((n1,n2),getitem, x1=x1, x2=x2) print(T.shape) # (100, 200) print(T) # LazyTensor(shape=(100, 200),attributes=(x1,x2)) # get the full tensor (not lazy) full_T = T[:] # get one component T11 = T[1,1] # get one row T1 = T[1] # get one column with slices Tsliced = T[::10,5] # get the data as attributes x1_0 = T.x1 x2_0 = T.x2 # create a LazyTensor from another one : exp(T) T2 = ot.utils.LazyTensor(T.shape,lambda i,j,T: np.exp(T[i,j]) , T=T) # compute sum s = reduce_lazytensor(T, np.sum) # compute marginals c = reduce_lazytensor(T, np.sum, axis=1) r = reduce_lazytensor(T, np.sum, axis=0, batch_size=50)