Lazy evaluation

[ingmarschuster]

What are your thoughts on lazy evaluation? This can make working with extremely large matrices possible. One example of an algorithm that approximates rather large PSD matrices, their cholesky factors and inverses with only few iterations is
https://paperswithcode.com/paper/randomly-pivoted-cholesky-practical
This would probably necessitate to introduce an API that allows interfacing with user-defined compute backends. Or thats at least how the RPCholesky code is currently implemented.

[mfinzi]

Hi @ingmarschuster
Actually CoLA is already built around this lazy evaluation concept, although maybe we have some work to do in making this clear in the landing page and docs.

For example, constructing a product of low rank factors:

U = cola.ops.Dense(torch.randn(10000,5))
V = cola.ops.Dense(torch.randn(10000,5))
A = U@V.T

The matrices are never explicitly multiplied. All of the iterative algorithms in CoLA are defined through matrix vector multiplies, so in this way the matrices (A in this case) are only used lazily and never materialized.

Is this what you had in mind or are you speaking towards something else?

[ingmarschuster]

Hey @mfinzi,

Not exactly what I mean. The part I'm talking about is exemplified in RP choleskys lazy KernelMatrix. It allows to not store matrix entries as explicit numbers but rather when you access a matrix element, they're computed on the fly (and potentially cached) using _getitem_helper _diag_helper, _function_vec, _function and others.
These are then used inside abstracted linear algebra quite similar to CoLa for computing cholesky factors, inverses etc.

[ingmarschuster]

See also this file for low rank eigenvalue/cholesky factor approximation in the RPCholesky code.

[mfinzi]

Hi @ingmarschuster
I think that should be very doable using existing cola constructs. For kernels I think one could implement it with something like this (would need to verify on getitem and the off diagonal shifts). Is this the kind of thing you're thinking about?

import jax.numpy as jnp
import cola
from cola.ops import LinearOperator
from cola import Algorithm

def example_kernel_fn(x,y):
    ls=0.5
    return jnp.exp(-jnp.sum((x-y)**2/ls**2,-1))

class Kernel(LinearOperator):
    def __init__(self, kernel_fn, x, y, bs=100):
        self.kernel_fn = kernel_fn # broadcastable function of two array arguments
        self.x = x # (n, ) array
        self.y = y # (m, ) array
        self.bs = bs
        super().__init__(dtype=x.dtype, shape=(x.shape[0], y.shape[0]))

    def _matmat(self, v):
        outputs = []
        for i in range(0, self.x.shape[0], self.bs):
            outputs.append(self.kernel_fn(self.x[i:i+self.bs, None], self.y) @ v)
        return self.xnp.concatenate(outputs, axis=0)

    def __getitem__(self, idx):
        if isinstance(idx, tuple) and len(idx) == 1:
            idx = (idx[0], slice(None))
        return self.kernel_fn(self.x[idx[0],None],self.y[idx[1]])

@cola.dispatch
def diag(A: Kernel, k:int, alg:Algorithm):
    shifted_x, shifted_y = A.x, A.y
    if k > 0:
        shifted_x = shifted_x[:-k]
        shifted_y = shifted_y[k:]
    elif k < 0:
        shifted_x = shifted_x[-k:]
        shifted_y = shifted_y[:k]
    return A.kernel_fn(shifted_x, shifted_y)

[mfinzi]

It should also be possible to wrap KeOps with LinearOperators if we want a more hardware optimized implementation