Code design: moving away from kwargs and towards abstracting each algorithm as a class with its methods and settings.

[daniel-dodd]

The default approach is to use kwargs e.g., method="auto", max_iters , etc., to define the algorithm or functions like cola.solve, cola.logdet, and others. It can make it challenging to work out what operations have been executed, especially in a library with extensive dispatch mechanisms. I strongly advocate for moving away from this design.

A more robust approach would involve abstracting each algorithm and allowing them to be passed as arguments. This approach is effectively demonstrated in Lineax for structured matrix solves: Lineax Structured Matrices Documentation.

# Solver:

def solve(linop: LinearOperator, rhs: Array, algorithm: AbstractSolveAlgorithm) -> Array:
      ...

@dataclass
class AbstractSolveAlgorithm:
    ... # Put algorithm settings here.
    
    ... # Put any relevant methods here.

@dataclass 
class CholeskySolver(AbstractSolveAlgorithm):
    ... # Put algorithm settings here.
    
    ... # Put any relevant methods here.

@dataclass 
class CGSolver(AbstractSolveAlgorithm):
    ... # Put algorithm settings here.

    ... # Put any relevant methods here.

# Logdeterminant:

def logdet(linop: LinearOperator algorithm: AbstractLogDetAlgorithm) -> Array:
      ...

@dataclass
class AbstractLogDetAlgorithm:
   ... # Put algorithm settings here.

   ... # Put any relevant methods here.

[mfinzi]

Yeah we're open to some suggestions here. I get the need for some more transparency and control over what algorithm is ultimately run.

I guess it would look like

cola.linalg.inv(A, algorithm=Auto(tol=1e-5,max_iters=100))

and alternatives could be specified e.g. with

cola.linalg.inv(A, algorithm=CG(tol=1e-5,max_iters=100))
cola.linalg.inv(A, algorithm=LU())

A challenge comes from situations created by dispatch. For example what happens if we have
A=cola.kron(B, C) where B is PSD and C is not.

Dispatch splits the inverse into B⁻¹⊗C⁻¹. Using algorithm=CG(...) or algorithm=Cholesky would need to produce an error presumably since C cannot be inverted with CG, but using algorithm=GMRES(...) or algorithm=LU would be very inefficient for B⁻¹.

I suppose in these more complicated cases, the user could specify a mapping

algorithm={cola.PSD: CG(...), other: GMRES(...)}

or better as a function (since then it could also make the choices based on the sizes or types involved):

algorithm=lambda A: CG(...) if A.isa(cola.PSD) else GMRES(...)

[mfinzi]

I do see this being especially helpful as we begin to have more algorithms for the same linear algebra operation
e.g. LOBPCG vs Lanczos for smallest eigenvalue or BICGSTAB vs GMRES for asymmetric solves.

[mfinzi]

We've decided to go ahead and start working on this implementation.

Algorithms will most likely be simple immutable dataclasses (pytrees as well)

@dataclass_pytree
class Algorithm:
     pass

class LU(Algorithm):
     pass

class CG(Algorithm):
     tol=1e-5
     max_iters=1000
     pbar=False
     P = None
     x0 = None

The algorithm will most likely be specified also through dispatch:

@dispatch
def inv(A: LinearOperator, alg: LU):
      # implementation

@dispatch
def inv(A: LinearOperator, alg: CG):
      # implementation

but we will need to set the dispatch precedence so the LinearOperator dispatch rules are used first and these function merely as base cases unless the user wants to change that.

[daniel-dodd]

Nice work - this looks clean! I'm pleased to see you're adopting dataclasses for this (and I would suggest this for your LinearOperators too). Yeah immutability is super important here! A frozen dataclass would be the safest option. :)

[mfinzi]

@daniel-dodd Do you have any opinions on using jax_dataclasses vs tjax dataclasses vs vanilla dataclasses + simple_pytree?

[daniel-dodd]

All of these libraries have quite a lot of common features, and another one that shares similar traits is py🌲class, which has impressive display printing functionality.

For the most robust PyTrees available, I'd recommend Equinox's Module and Flax's struct.dataclass. These options are top-notch.

Regarding performance, I observed better performance when using simple_pytree compared to Equinox and Flax for flatten/unflatten operations on both CPU (M1) and GPU (Tesla v100) during some crude testing. For instance, it appears that py🌲class's findings align with my observations (you can find something similar on their GitHub repository: https://github.com/ASEM000/pytreeclass), and their PyTree implementation also exhibits good performance.

However, it's important to note that these performance differences are diminished when we consider the computational costs associated with linear algebra operations. Therefore, it's advisable not to take these observations at face value without considering the broader context and application-specific requirements.

On the topic of flexibility, simple_pytree is good choice. The codebase is incredibly minimal, which lends itself well to flexibility and easy customisation. You may need such flexibility when it comes to setting values during class initialisation (i.e.,. before declaring the object as immutable), e.g., with your annotations where you set new_obj.annotations = obj.annotations | {self} within the Metaclass. For example something like this (pseudo code) might work if you run into issues:

from simple_pytree import Pytree, static_field
from dataclasses import dataclass
from typing import Set
from typing_extensions import Self
from copy import copy
from jax.tree_util import tree_flatten

class Module(Pytree):
    _pytree__annotations: Set[str] = static_field()

    def __init_subclass__(cls):
        cls._pytree__annotations = {} # Define annotations.
        super().__init_subclass__(mutable=False) # Immutable.
        dataclass(cls, frozen=True)

    def update_annotations(self: Self, new: Set) -> Self:
        pytree = copy(self)
        pytree.__dict__.update(_pytree__annotations=new)
        return pytree

    @property
    def annotations(self):
        return self._pytree__annotations

class AbstractLinearOperator(Module):
    pass

class Example(AbstractLinearOperator):
    a: int = 1
    b: int = 2


# Init example.
e = Example()

# No annotations yet.
print(e.annotations) # {}

# Lets add some.
e = e.update_annotations({"a", "b", "c"})
print(e.annotations) # {'b', 'a', 'c'}

# Lets apply some PyTree transformations.
leaves, tree_def = tree_flatten(e)
e = tree_def.unflatten(leaves)

# We still have our annotations.
print(e.annotations) # {'b', 'a', 'c'}

# Base class is unaffected.
print(Example._pytree__annotations) # {}

[daniel-dodd]

Actually, forgot we are using Python 3.10 -> this simpler option would work here:

import equinox
from dataclasses import field, replace
from typing import Set
from jax.tree_util import tree_flatten

class Module(equinox.Module):
    annotations: Set[str] = field(default_factory=lambda: set(), metadata={"static": True}, kw_only=True)


class AbstractLinearOperator(Module):
    pass

class Example(AbstractLinearOperator):
    a: int = 1
    b: int = 2


# Init example.
e = Example()

# No annotations yet.
print(e.annotations) # {}

# Lets add some.
e = replace(e, annotations={"a", "b", "c"})
print(e.annotations) # {'b', 'a', 'c'}

# Lets apply some PyTree transformations.
leaves, tree_def = tree_flatten(e)
print(leaves) # [1, 2]
e = tree_def.unflatten(leaves)

# We still have our annotations.
print(e.annotations) # {'b', 'a', 'c'}

[mfinzi]

Merged on main with #73 for the Algorithms implementation and assorted changes.
We will have a new release out soon.

For the pytree migration, we are hoping to use py🌲class, but some issues remain to be sorted out. So that is coming, but not ready yet.