Custom Pytree Nodes

By default, JAX only knows how to traverse dicts, lists, tuples, and namedtuples. If you create a custom class to hold parameters, JAX treats it as a leaf — it won't look inside it. To make JAX understand your custom class, you register it as a pytree node.

import jax
import jax.numpy as jnp
from functools import partial

class LinearParams:
    """A simple container for linear layer parameters."""
    def __init__(self, weights, bias):
        self.weights = weights
        self.bias = bias

    def __repr__(self):
        return f"LinearParams(w={self.weights.shape}, b={self.bias.shape})"

# Without registration, JAX sees this as an opaque leaf
layer = LinearParams(jnp.ones((3, 4)), jnp.zeros(4))
print(jax.tree.leaves(layer))  # [LinearParams(...)] — the object itself!

To register it, you define two functions: one to flatten (extract children + metadata) and one to unflatten (reconstruct from children + metadata):

def linear_flatten(obj):
    """Returns (children, aux_data)."""
    children = (obj.weights, obj.bias)  # the arrays JAX should trace
    aux_data = None                      # any static metadata (none here)
    return children, aux_data

def linear_unflatten(aux_data, children):
    """Reconstructs the object from (aux_data, children)."""
    weights, bias = children
    return LinearParams(weights, bias)

# Register with JAX
jax.tree_util.register_pytree_node(
    LinearParams,
    linear_flatten,
    linear_unflatten,
)

# Now JAX sees inside it
layer = LinearParams(jnp.ones((3, 4)), jnp.zeros(4))
print(jax.tree.leaves(layer))
# [Array([[1., ...]], dtype=float32), Array([0., ...], dtype=float32)]

# tree.map works too
doubled = jax.tree.map(lambda x: x * 2, layer)
print(doubled)  # LinearParams(w=(3, 4), b=(4,)) with doubled values

# Example with aux_data: a layer that stores its activation name
class DenseParams:
    def __init__(self, weights, bias, activation='relu'):
        self.weights = weights
        self.bias = bias
        self.activation = activation  # static config

def dense_flatten(obj):
    children = (obj.weights, obj.bias)
    aux_data = obj.activation  # stored as static metadata
    return children, aux_data

def dense_unflatten(aux_data, children):
    weights, bias = children
    return DenseParams(weights, bias, activation=aux_data)

jax.tree_util.register_pytree_node(DenseParams, dense_flatten, dense_unflatten)

# Now grad works through DenseParams
layer = DenseParams(jnp.ones((3, 4)), jnp.zeros(4), activation='relu')
grads = jax.grad(lambda l: jnp.sum(l.weights))(layer)
print(grads.weights)  # all ones
print(grads.activation)  # 'relu' — static data passed through unchanged