Quax

JAX + multiple dispatch + custom array-ish objects

For example, this can be mean overloading matrix multiplication to exploit sparsity or structure, or automatically rewriting a LoRA's matmul (W + AB)v into the more-efficient Wv + ABv.

Applications include:

• LoRA weight matrices
• symbolic zeros
• arrays with named dimensions
• structured (e.g. tridiagonal) matrices
• sparse arrays
• quantised arrays
• arrays with physical units attached
• etc! (See the built-in quax.examples library for most of the above!)

This works via a custom JAX transform. Take an existing JAX program, wrap it in a quax.quaxify, and then pass in the custom array-ish objects. This means it will work even with existing programs, that were not written to accept such array-ish objects!

(Just like how jax.vmap takes a program, but reinterprets each operation as its batched version, so to will quax.quaxify take a program and reinterpret each operation according to what array-ish types are passed.)

Installation

pip install quax

What that looks like

A few examples of quax propagating physical units through JAX code.

Your code doesn't change. kinetic_energy is ordinary JAX with no notion of units. Quax tracks them through it and works out the result's units:

import jax.numpy as jnp
import quax
from quax.examples.unitful import kilograms, meters, seconds, Unitful

def kinetic_energy(m, v):
    return 0.5 * m * v**2

mass = Unitful(jnp.asarray(2.0), kilograms)
velocity = Unitful(jnp.asarray(3.0), {meters: 1, seconds: -1})

energy = quax.quaxify(kinetic_energy)(mass, velocity)
print(energy.array, energy.units)  # 9.0 {kg: 1, m: 2, s: -2}

Units are checked, not just carried along. A Unitful knows that adding metres to seconds is meaningless, and raises instead of silently returning a wrong number:

try:
    quax.quaxify(jnp.add)(
        Unitful(jnp.asarray(1.0), meters), Unitful(jnp.asarray(1.0), seconds)
    )
except ValueError as e:
    print(e)  # Cannot add two arrays with units {m: 1} and {s: 1}.

It also works on code you didn't write. Diffrax was written years before Quax existed and knows nothing about units, but its solver carries them through anyway — unmodified:

import diffrax
import jax.numpy as jnp
import quax
from quax.examples.unitful import meters, Unitful

term = diffrax.ODETerm(lambda t, y, args: -0.5 * y)
solver = diffrax.Euler()

def step(y0):
    state = solver.init(term, 0.0, 0.1, y0, None)
    y1, _, _, _, _ = solver.step(term, 0.0, 0.1, y0, None, state, made_jump=False)
    return y1

y1 = quax.quaxify(step)(Unitful(jnp.asarray([1.0]), meters))
print(y1.array, y1.units)  # [0.95] {m: 1}

The one place this breaks down: if a library allocates its own buffers internally, Quax never sees the allocation and you get plain arrays back. Sharp bits covers that and the other boundaries.

Documentation

https://nstarman.github.io/quax

• Custom rules — build your own array-ish type, start to finish
• API reference
• Sharp bits — where Quax stops being transparent, and why
• FAQ — jit, vmap, and writing aval()

See also: other libraries in the JAX ecosystem

Always useful
Equinox: neural networks and everything not already in core JAX!
jaxtyping: type annotations for shape/dtype of arrays.

Deep learning
Optax: first-order gradient (SGD, Adam, ...) optimisers.
Orbax: checkpointing (async/multi-host/multi-device).
Levanter: scalable+reliable training of foundation models (e.g. LLMs).

Scientific computing
Diffrax: numerical differential equation solvers.
Optimistix: root finding, minimisation, fixed points, and least squares.
Lineax: linear solvers.
BlackJAX: probabilistic+Bayesian sampling.
sympy2jax: SymPy<->JAX conversion; train symbolic expressions via gradient descent.

Built on Quax
Quaxed: a namespace of already-wrapped quaxify(jnp.foo) operations.
quax-blocks: blocks for constructing quax classes.
unxt: unitful quantities.
coordinax: coordinates in JAX.
galax: galactic and gravitational dynamics, with GPU and autodiff.
phasecurvefit: construct paths through phase-space points.

Awesome JAX
Awesome JAX: a longer list of other JAX projects.

Acknowledgements

Significantly inspired by https://github.com/davisyoshida/qax, https://github.com/stanford-crfm/levanter, and jax.experimental.sparse.