Interactive online version:

Quick Start

We now showcase some basic functionality present in CoLA. We’ll start by showing how to define different types of Linear ops, then we’ll show how to perform basic arithmetic with Linear ops and, finally, we’ll conclude applying some linear algebra operations (like solves or log determinants) to the Linear ops.

We’ll work with torch.tensors in this example, but the same code can be run using JAX arrays (jnp.ndarrays).

import cola
import torch
import numpy as np
import warnings

warnings.filterwarnings("ignore")
torch.manual_seed(21)

<torch._C.Generator at 0x7f547f30c210>

Creating a Linear Operator

You can find several predefined Linear ops under cola.ops. We’ll illustrate three basic cases: Dense, Diagonal and Tridiagonal.

A Dense Linear Operator is nothing more than a wrapper on a dense matrix, where the wrapper defines a matmat function \(v \mapsto Av\) and holds several attributes such as dtype and shape.

Let’s start by defining a dense matrix and a vector to act upon. Below we show the entries of the matrix \(A\), of the vector \(v\) and the result of \(Av\).

N = 4
A_dense = torch.randn(N, N)
vec = torch.randn(N)
print(f"A: {A_dense}\nAv: {A_dense @ vec}")

A: tensor([[-0.2386, -1.0934,  0.1558,  0.1750],
        [-0.9526, -0.5442,  1.1985,  0.9604],
        [-1.1074, -0.8403, -0.0020,  0.2240],
        [ 0.8766, -0.5379, -0.2994,  0.9785]])
Av: tensor([-0.3842, -1.8328, -0.7573,  0.9795])

To create a Dense operator simply run:

A = cola.ops.Dense(A_dense)
print(type(A))

<class 'cola.ops.operators.Dense'>

The previous operator now has a dtype and a shape attribute. More importantly, it can now act on the vector \(v\) and get the same result as above.

print(f"Dtype: {A.dtype} | Shape: {A.shape}")
print(A @ vec)

Dtype: torch.float32 | Shape: torch.Size([4, 4])
tensor([-0.3842, -1.8328, -0.7573,  0.9795])

In order to make use of the benefits of CoLA, let’s consider using some large structured Linear Operators. First let’s start by creating some structured linear operators, such as a diagonal, a dense matrix, and a permutation matrix.

ones = torch.ones(300)
upper = ones[:-1]
lower = ones[:-1]
diagonal= -2*ones

T = cola.ops.Tridiagonal(lower, diagonal, upper)
print(f"{T[:3,:3].to_dense()}, Type: {T}")

tensor([[-2.,  1.,  0.],
        [ 1., -2.,  1.],
        [ 0.,  1., -2.]]), Type: Tridiagonal

perm = torch.randperm(200)
P = cola.ops.Permutation(perm)
v = torch.randn(200)
print((P @ v)[:5], v[perm[:5]])

tensor([-1.9827, -0.5036,  0.7356, -0.2840,  1.3906]) tensor([-1.9827, -0.5036,  0.7356, -0.2840,  1.3906])

For these Linear Operators, the memory and compute cost for multiplication is only \(O(n)\) rather than the \(O(n^2)\) for dense.

Doing binary operations with Linear ops

CoLA provides a similar interface to combine Linear ops as you would combine matrices. We can add and subtract (+,-), multiply by scalars (*,/), matrix multiply (@), transpose (.T), slice [:,:] and more

D = cola.ops.Diagonal(torch.tensor([1, 2, 3, 4.]))
DT = T[:4, :4] - D
print(DT.to_dense())

tensor([[-3.,  1.,  0.,  0.],
        [ 1., -4.,  1.,  0.],
        [ 0.,  1., -5.,  1.],
        [ 0.,  0.,  1., -6.]])

Let’s use these operations to make \(B= A^T(T_{[:4,:4]}-D) + \mu I\) regularized by \(\mu\) by running the following code.

mu = 1e-6
B = A.T @ DT
B += mu * cola.ops.I_like(B)
print(f"Object: {B}\n {B.to_dense()}")

Object: DenseTridiagonal[slice(None, 4, None),slice(None, 4, None)]+-1.0diag(tensor([1., 2., 3., 4.]))+9.999999974752427e-07I
 tensor([[-0.2368,  2.4642,  5.4611, -6.3671],
        [ 2.7361,  0.2431,  3.1194,  2.3872],
        [ 0.7310, -4.6402,  0.9091,  1.7946],
        [ 0.4354, -3.4425,  0.8188, -5.6472]])

Linear Operators can also be combined in more involved ways, such as with kron for \(\otimes\), block_diag, and concatenate

K = cola.kron(A, D)
print(f"{K[:5,:5].to_dense()}\n shape: {K.shape}, object: {K}")

tensor([[-0.2386,  0.0000,  0.0000,  0.0000, -1.0934],
        [ 0.0000, -0.4772,  0.0000,  0.0000,  0.0000],
        [ 0.0000,  0.0000, -0.7157,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.0000, -0.9543,  0.0000],
        [-0.9526,  0.0000,  0.0000,  0.0000, -0.5442]])
 shape: (16, 16), object: Dense⊗diag(tensor([1., 2., 3., 4.]))

d = 1 / (1 + torch.linspace(.2, 10, 1000))
X = cola.block_diag(cola.kron(A, P), T, cola.ops.Diagonal(d))
X

<2100x2100 BlockDiag[cola.ops.operators.Kronecker[cola.ops.operators.Dense, cola.ops.operators.Permutation], cola.ops.operators.Tridiagonal, cola.ops.operators.Diagonal] with dtype=torch.float32>

Under the hood the operator \(B\) is lazily defined and would know how to apply the Linear ops to any vector \(v\).

Computing linear algebra operations (e.g. solves, logdet, diag, exp)

To solve the linear system \(Bx=v\) we use the inverse function. This inverse function lazily defines \(B^{-1}\) and hence applying it to \(v\) yields the solution \(x=B^{-1}v\). The inverse of \(B\) is never computed, using \(B^{-1}\) is simply how in CoLA we call linear solves.

X[:, :5].to_dense().shape

torch.Size([2100, 5])

X_inv = cola.inv(X)
v = torch.randn(X.shape[-1])
soln = X_inv @ v
print(f"Solution error: {torch.linalg.norm(X @ soln - v)/torch.linalg.norm(v):.2e}")

Solution error: 9.79e-06

print(f"CoLA logdet: {cola.logdet(X):.3f}, dense logdet {torch.logdet(X.to_dense()):.3f}")

CoLA logdet: -1564.690, dense logdet -1564.690

print(f"CoLA diag: {cola.diag(X)[-5:]}\ndense diag: {torch.diag(X.to_dense())[-5:]}")

CoLA diag: tensor([0.0912, 0.0912, 0.0911, 0.0910, 0.0909])
dense diag: tensor([0.0912, 0.0912, 0.0911, 0.0910, 0.0909])

from scipy.linalg import expm
scipy_expm = expm((T.to_dense()).numpy())
print(f"CoLA exp(X): {cola.exp(T)[-5:,-5:].to_dense().real}\nscipy expm: {scipy_expm[-5:,-5:]}")

CoLA exp(X): tensor([[0.3085, 0.2153, 0.0932, 0.0288, 0.0066],
        [0.2153, 0.3085, 0.2152, 0.0930, 0.0275],
        [0.0932, 0.2152, 0.3083, 0.2139, 0.0864],
        [0.0288, 0.0930, 0.2139, 0.3016, 0.1865],
        [0.0066, 0.0275, 0.0864, 0.1865, 0.2153]])
scipy expm: [[0.30850828 0.21526888 0.09323523 0.02876081 0.0066488 ]
 [0.21526891 0.30850452 0.2152388  0.09302244 0.02746145]
 [0.09323531 0.2152388  0.30829173 0.21393944 0.08637362]
 [0.02876082 0.09302243 0.2139395  0.30164295 0.18647799]
 [0.0066488  0.02746145 0.08637366 0.18647805 0.21526927]]

from cola import Auto, Arnoldi
print(f"CoLA exp(X): {cola.exp(T, alg=Arnoldi())[-5:,-5:].to_dense().real}\nscipy expm: {scipy_expm[-5:,-5:]}")

CoLA exp(X): tensor([[0.3085, 0.2153, 0.0932, 0.0288, 0.0066],
        [0.2153, 0.3085, 0.2152, 0.0930, 0.0275],
        [0.0932, 0.2152, 0.3083, 0.2139, 0.0864],
        [0.0288, 0.0930, 0.2139, 0.3016, 0.1865],
        [0.0066, 0.0275, 0.0864, 0.1865, 0.2153]])
scipy expm: [[0.30850828 0.21526888 0.09323523 0.02876081 0.0066488 ]
 [0.21526891 0.30850452 0.2152388  0.09302244 0.02746145]
 [0.09323531 0.2152388  0.30829173 0.21393944 0.08637362]
 [0.02876082 0.09302243 0.2139395  0.30164295 0.18647799]
 [0.0066488  0.02746145 0.08637366 0.18647805 0.21526927]]

C = X.T @ X
# let's annotate this matrix as positive definite to speed up the computation
C = cola.PSD(C)
e, v = cola.eig(C, k=3, which="SM", alg=Auto(max_iters=500))
print(e, torch.linalg.eigh(C.to_dense())[0][:3])

tensor([2.5006e-06, 5.9026e-05, 1.5939e-04]) tensor([-1.1555e-07,  2.6563e-07,  1.3119e-06])

e, v = cola.eig(C, k=3, which="LM", alg=Auto(max_iters=500))
print(e, torch.linalg.eigh(C.to_dense())[0][-3:])

tensor([15.9922, 15.9965, 15.9991]) tensor([15.9922, 15.9965, 15.9991])