Pytree states

[jlperla]

v1 states are single arrays (scalar or vector), documented as a deliberate limitation. Supporting pytree states (dicts, nnx params, tuples of arrays) would let users integrate structured models without manual packing.

Implementation sketch (grounded in the v1 code)

solvers.py — every stage combination like

k3 = g(x + dt * (A31 * k1 + A32 * k2), t + C3 * dt)

becomes a tree_map over the state/stage trees:

k3 = g(jax.tree.map(lambda x_, k1_, k2_: x_ + dt * (A31 * k1_ + A32 * k2_), x, k1, k2), t + C3 * dt)

A small tree_axpy(coeffs, x, ks) helper keeps the tableau readable. project already maps tree→tree if the user writes it that way.

controllers.py — the error norm must avoid ravel_pytree (which allocates and re-concatenates every step). Tree-reduced max-norm:

E = jax.tree.reduce(
    jnp.maximum,
    jax.tree.map(
        lambda e, a, b: jnp.max(jnp.abs(e) / (atol + rtol * jnp.maximum(jnp.abs(a), jnp.abs(b)))),
        err, x0, x1,
    ),
)

wrapped in the existing stop_gradient.

ode.py — the carry entries x, f_cur and the masked advance (jnp.where(advance, x1, x)) become per-leaf tree_maps; f_init = jnp.zeros_like maps per leaf; row stacking (jnp.concatenate([x0[None], xs_s])) and the fill="inf" broadcast map per leaf too. num_accepted, done, t are unchanged.

interpolation.py — hermite_interpolate maps per leaf with a shared idx/s (the knot times are shared): compute idx, s, width_safe, degenerate once, then tree_map the Hermite combination over knot_xs/knot_fs leaves.

sde.py — dW becomes a tree of increments (jax.tree.map over split keys), EM step maps per leaf.

Tests to add

• dict-state solve vs the equivalent stacked-array solve (exact equality)
• AD (jvp/vjp/reverse-over-forward) with a dict state through all SaveAt modes
• vmap over a pytree-state batch
• recompile: pytree structure is static, leaf values are data