nstarman/quax/quax

Using Quax Effectively

Quax runs a JAX program under a custom interpreter, reinterpreting each primitive according to the types passed in. Just as jax.vmap takes a program and reinterprets every operation as its batched version, quax.quaxify takes a program and reinterprets every operation according to multiple-dispatch rules you register. This means custom array-ish objects work with existing JAX code that was never written to accept them.

Checked against quax v0.4.x (jax 0.7.2–0.11.x, equinox>=0.13.3, plum-dispatch>=2.5.7, Python >=3.11). Docs: <https://nstarman.github.io/quax>. Note that patrick-kidger/quax is at 0.2.x while PyPI quax now comes from nstarman/quax — prior knowledge of this library is probably a version behind. See "Version notes" at the end.

Quick start

Take any existing JAX program and pass custom types through it:

import equinox as eqx
import jax.random as jr
import quax
import quax.examples.lora as lora
key1, key2, key3 = jr.split(jr.key(0), 3)
linear = eqx.nn.Linear(10, 12, key=key1)
vector = jr.normal(key2, (10,))
def run(model, x):
    return model(x)
run(linear, vector)  # ordinary JAX, works as normal
# Swap the weight for a LoRA array, then quaxify the *unchanged* function.
lora_weight = lora.LoraArray(linear.weight, rank=2, key=key3)
lora_linear = eqx.tree_at(lambda l: l.weight, linear, lora_weight)
out = quax.quaxify(run)(lora_linear, vector)

quax.examples ships lora, zero, unitful, named, sparse, prng, and structured_matrices — read their _core.py files as reference implementations.

Pair quaxify with jax.jit

Bare quax.quaxify(fn)(*args) is 50–100x slower than the JIT path for small operations. Every call pays Python-level trace setup, jaxpr interpretation, and equinox module overhead — roughly 1–2 µs of fixed cost regardless of array size. This is the single most expensive mistake available here, and it is invisible: the code is correct, just slow.

import jax
import jax.numpy as jnp
jit_fn = jax.jit(quax.quaxify(jnp.add))  # preferred: jit outermost

Put jax.jit at the outermost level. quax.quaxify(jax.jit(fn)) is correct but compiles less efficiently, because each inner jitted call is compiled in isolation before quaxify sees it. With vmap, ordering does not matter for correctness or speed — choose whichever expresses your intent.

Quaxify short-circuits entirely when no argument is a quax.Value, so quaxifying a function that might receive plain arrays costs nothing extra.

How dispatch resolves

For primitive.bind(x, y, ...) under a quaxify, in order:

1. A rule registered for the types of (x, y, ...) matches → use it.
2. Exactly one argument's type overrides Value.default → use that.
3. No type overrides default → materialise every operand and call ordinary

JAX.

1. Two or more types override default → `TypeError: Multiple array-ish types

{...} are specifying default process rules.`

Rule 3 is the one that surprises people. A missing rule is not an error by default; it silently degrades to plain JAX, and your custom type vanishes from the output. If that is never acceptable for your type, make materialise() raise — then rule 3 becomes a loud failure instead of a silent one.

An operand that is neither a quax.Value nor array-like (a str, None) raises TypeError: Primitive ... got an operand of type ... that is neither a quax.Value nor array-like.

Creating a custom ArrayValue

1. Subclass `quax.ArrayValue` (not Value — ArrayValue is for array-ish

things and gives you .shape/.dtype/.ndim/.size for free). It is an equinox.Module, so it is a frozen dataclass and a pytree.

1. Mark metadata `eqx.field(static=True)` — units, axis names, stored shapes,

bool flags. Anything JAX must not trace through.

1. Implement `aval()` returning a jax.core.ShapedArray. It must be pure.
2. Implement `materialise()` — or raise, to forbid the rule-3 fallback.
3. Register rules for the primitives you care about, starting with the ones

your users will actually hit (add_p, mul_p, dot_general_p).

1. Add mixed-type rules so your type interoperates with plain arrays and with

other people's quax types.

1. Test both paths: the registered rule, and what happens when no rule

matches.

import jax
import jax.numpy as jnp
from jax import lax
from jaxtyping import ArrayLike
class Meters(quax.ArrayValue):
    array: ArrayLike
    def aval(self) -> jax.core.ShapedArray:
        return jax.core.ShapedArray(jnp.shape(self.array), jnp.result_type(self.array))
    def materialise(self):
        raise ValueError("Refusing to materialise Meters: it would drop the unit.")
@quax.register(lax.add_p)
def add_meters_meters(x: Meters, y: Meters) -> Meters:
    return Meters(x.array + y.array)
quax.quaxify(jnp.add)(Meters(jnp.arange(3.0)), Meters(jnp.ones(3)))

aval() and materialise()

aval() must be pure: the same instance must return the same AbstractValue every time. Quax caches the result at tracer-construction time and never observes later changes.

Pure does not mean array-free — aval() may bind JAX primitives, and quax evaluates it under the parent trace so that works. Prefer not to: aval runs once per tracer, so a bind there is on the hot path. Read the field directly when the primitive would not change the shape or dtype. (Under jax ≤0.10 this held by luck; jax 0.11 leaves no current trace at that point. Fixed on main after v0.4.3 — on v0.4.3 or earlier with jax 0.11, an aval() that binds anything fails.)

Which fields need static=True follows from that. Shape and dtype derived from an array field need nothing special — JAX arrays are immutable, so reading self.array.shape is already stable. A separately stored shape, or units, or a flag consulted by aval(), must be eqx.field(static=True); forget it and JAX tries to trace through the integer, which errors long before staleness matters.

import equinox as eqx
class Zeros(quax.ArrayValue):
    _shape: tuple[int, ...] = eqx.field(static=True)  # REQUIRED
    _dtype: jnp.dtype = eqx.field(static=True)  # REQUIRED
    def aval(self) -> jax.core.ShapedArray:
        return jax.core.ShapedArray(self._shape, self._dtype)
    def materialise(self):
        return jnp.zeros(self._shape, self._dtype)

materialise() raising is a legitimate, common design — Unitful, LoraArray, and NamedArray all do it, because materialising would silently discard the very information the type exists to carry. LoraArray and NamedArray expose an allow_materialise flag so users can opt into the fallback. If you see a materialise() that raises, do not "fix" it.

eqx.field(converter=jnp.asarray) on the payload is a cheap way to accept lists/scalars without writing __init__.

Writing dispatch rules

@quax.register(primitive) takes the jax.extend.core.Primitive, and dispatches on the type annotations — plum reads them, so they are load-bearing, not documentation.

Name the rule; never `def _(...)`. Dispatch does not read the name, but everything you debug with does: tracebacks, plum ambiguity and redefinition errors, and profiles all identify a rule by its function name, and a module of rules all called _ makes every one of them indistinguishable. Name it after the primitive and the types it dispatches on — add_meters_meters, mul_meters_array_like, select_n_unitful — matching the existing convert_element_type_zero in quax.examples.zero and cond_quax in quax._primitives.

Find the primitive behind a jnp function by tracing it:

print(jax.make_jaxpr(jnp.square)(jnp.arange(3.0)))  # shows integer_pow

Positional arguments are the primitive's operands; keyword arguments are its params (bind(*args, **params)). Accept `kw and forward it** — params come and go across JAX versions (out_dtype on mul_p, out_sharding on reduce_sum_p, sharding on broadcast_in_dim_p`), and a rule with a rigid signature breaks on upgrade.

For mixed types, register a pair, using ArrayLike | quax.ArrayValue for the operand you do not own:

@quax.register(lax.mul_p)
def mul_meters_array_like(x: Meters, y: ArrayLike, /, **kw) -> Meters:
    return Meters(lax.mul_p.bind(x.array, y, **kw))
@quax.register(lax.mul_p)
def mul_array_like_meters(x: ArrayLike, y: Meters, /, **kw) -> Meters:
    return Meters(lax.mul_p.bind(x, y.array, **kw))

Annotating the other operand as ArrayLike | quax.ArrayValue and then redispatching through a nested quax.quaxify is what lets your type interoperate with types you have never heard of — the other author's rules handle their half. This is the preferred pattern for library types.

A rule may return a plain array rather than your type when that is the honest answer (comparisons returning bools, integer_pow with y=0 returning ones).

If two rules can both match — (Meters, ArrayLike | ArrayValue) and (ArrayLike | ArrayValue, Meters) both match a (Meters, Meters) call — plum raises AmbiguousLookupError. Fix it by adding the specific rule with precedence=1:

@quax.register(lax.mul_p, precedence=1)
def mul_meters_meters(x: Meters, y: Meters, /, **kw) -> Meters:
    return Meters(lax.mul_p.bind(x.array, y.array, **kw))

Prefer resolving ambiguity with nested quaxify(fn, filter_spec=...) at the call site when the clash is between two libraries' types: registering a cross-library rule mutates a global dispatch table and commits you to an implementation for a combination you may not understand. Order matters — the outer quaxify sees the type it selects last.

Overriding Value.default

default is a rule for every primitive at once, keyed only on your type. Use it when the behaviour is uniform across primitives — a tag that propagates through everything (detecting the backward pass, quantisation policy), rather than per-primitive semantics.

class Tagged(quax.ArrayValue):
    array: ArrayLike
    def aval(self) -> jax.core.ShapedArray:
        return jax.core.ShapedArray(jnp.shape(self.array), jnp.result_type(self.array))
    def materialise(self):
        return self.array
    @staticmethod
    def default(primitive, values, params):
        raw = [x.array if isinstance(x, Tagged) else x for x in values]
        out = primitive.bind(*raw, **params)
        return [Tagged(o) for o in out] if primitive.multiple_results else Tagged(out)

Only one type per expression may override default (dispatch rule 4). Two default-overriding types meeting in one computation is a hard error, so this is a strong commitment for a library type — per-primitive rules compose, default does not.

Boundaries

Construct Values outside the quaxify boundary, not inside. Constructing one inside appears to work — it is combining it with values that did cross the boundary that fails, with TypeError: unsupported operand type(s) for +: '_QuaxTracer' and 'YourType'. A value built inside has no trace to associate with (which of two nested quaxifies would own it?), so it never becomes a tracer. Build your values, then cross the boundary.

Inside the boundary your type looks like an array. That is the entire point, and it has a consequence for runtime typechecking: the jaxtyping/beartype import hook installs inside the quaxify, so annotations on the wrapped function must describe arrays, not your Value type. A function annotated q: Shaped[UnitfulArray, "3"] will fail typechecking when quaxify hands it something that presents as f64[3].