TopolearnModel parameters, optimizer state and batches are all pytrees; manipulating them is daily JAX work.
A pytree is any nested structure of Python containers (lists, tuples, dicts, None, and registered classes) whose leaves are arrays or scalars. JAX transformations accept and return pytrees: grad of a loss over a dict of parameters returns a dict of gradients with the same keys.
The jax.tree module works on whole trees:
jax.tree.map(f, tree)appliesfto every leaf; with several trees of the same structure,fgets matching leaves, e.g. an SGD updatejax.tree.map(lambda p, g: p - lr * g, params, grads).jax.tree.leaves(tree)gives a flat list of leaves, useful for counting parameters.jax.tree.structure(tree)gives the "treedef" describing the shape;jax.tree.flatten/unflattenconvert back and forth.
Your own classes become pytrees once registered. The simplest way is a dataclass with @jax.tree_util.register_dataclass, marking non-array configuration as static metadata so it is not traced. Libraries like Flax and Equinox do this for you.
Things to know:
Noneis an empty subtree, not a leaf, sojax.tree.mapskips it.- Tree structure is part of a jitted function's cache key: changing a dict's keys triggers recompilation.
- Trees passed to a multi-tree
mapmust have identical structure, or you get an error naming the mismatch.
Example
import jax
import jax.numpy as jnp
from dataclasses import dataclass, field
params = {"w": jnp.ones((3, 2)), "b": jnp.zeros(2)}
def loss(p, x):
return jnp.sum(jnp.tanh(x @ p["w"] + p["b"]))
grads = jax.grad(loss)(params, jnp.ones((4, 3)))
print(jax.tree.map(jnp.shape, grads)) # {'b': (2,), 'w': (3, 2)}
new = jax.tree.map(lambda p, g: p - 0.1 * g, params, grads)
print(sum(x.size for x in jax.tree.leaves(new))) # 8 parameters
@jax.tree_util.register_dataclass
@dataclass
class Linear:
w: jax.Array
b: jax.Array
name: str = field(default="dense", metadata=dict(static=True))
layer = Linear(jnp.ones(2), jnp.zeros(2))
print(jax.tree.leaves(layer)) # [Array([1., 1.], ...), Array([0., 0.], ...)]Remember
- Transformations take and return arbitrarily nested pytrees of arrays
jax.tree.mapover several same-shaped trees implements updates leaf by leaf- Register custom classes (e.g.
register_dataclass) and mark config as static
Watch out for
- Storing Python config strings or ints as leaves, so they get traced and cause errors or recompiles.