From 5f577456d42a2a9382229dabbc083427f295e010 Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Tue, 4 Aug 2026 17:39:54 -0700
Subject: [PATCH 1/5] feat: SRA1 SDAEs, unroll=, drop compat shims

solve_semi_explicit_sdae accepts SRA1: noise comes from
solver.sample_noise, and each step adds a root-restored drift stage at
t + 3/4 dt before the endpoint root. Milstein stays excluded because its
commutativity contract does not survive the implicit reduction.

solve_sde and fixed-step solve_ode gain a static unroll= argument
forwarded to their scans: identical values, fewer GPU dispatches. On
L40S NN-policy ensembles it cut reverse-mode solve time 2-3x.

Removes the max_steps_is_success field (never broadened root
acceptance) and the private brownian-increment alias, shortens the
public docstrings to match the reorganized docs, and adds float32
parametrization plus GPU tests for vmapped ensembles with explicit
noise and reverse-mode over vmap.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01C8SWZyLFpzVjEX5WCRLJnD
---
 docs/adaptive_ad.md              | 261 --------------------
 docs/aux.md                      | 158 ------------
 docs/rodas5p.md                  | 147 -----------
 docs/static_shapes.md            | 156 ------------
 src/tinydiffeq/__init__.py       |  50 ++--
 src/tinydiffeq/controllers.py    |  29 +--
 src/tinydiffeq/dae.py            | 118 +++------
 src/tinydiffeq/exponential.py    |  17 +-
 src/tinydiffeq/ode.py            |  66 ++---
 src/tinydiffeq/save_at.py        |  30 +--
 src/tinydiffeq/sdae.py           | 135 +++++++----
 src/tinydiffeq/sde.py            | 111 +++++----
 src/tinydiffeq/solution.py       |  56 ++---
 src/tinydiffeq/solvers.py        | 145 ++++++++++-
 tests/test_aux.py                |   2 +-
 tests/test_dae.py                |  40 +--
 tests/test_float64_subprocess.py |  54 +++++
 tests/test_gpu.py                | 190 +++++++++++++++
 tests/test_sdae.py               |  87 ++++++-
 tests/test_sde.py                | 405 ++++++++++++++++++++++++++++++-
 tests/test_solvers_fixed.py      |  87 ++++++-
 21 files changed, 1228 insertions(+), 1116 deletions(-)
 delete mode 100644 docs/adaptive_ad.md
 delete mode 100644 docs/aux.md
 delete mode 100644 docs/rodas5p.md
 delete mode 100644 docs/static_shapes.md

diff --git a/docs/adaptive_ad.md b/docs/adaptive_ad.md
deleted file mode 100644
index 591228e..0000000
--- a/docs/adaptive_ad.md
+++ /dev/null
@@ -1,261 +0,0 @@
-# Adaptive Stepping and AD
-
-The default `adaptive_loop="bounded"` supports forward mode, reverse mode, and
-reverse-over-forward through one adaptive solve. The alternative
-`adaptive_loop="forward"` executes a data-dependent number of attempts and
-supports primal evaluation, JVP, and nested forward mode, but not ordinary
-reverse mode. Both paths use the same deliberate frozen-controller derivative
-convention described below.
-
-## Default tolerances follow the state precision
-
-`IController()` and `PIController()` resolve omitted tolerances from the
-state dtype at trace time:
-
-| State/time dtype | `rtol` | `atol` | unit-scale `dt_min` |
-|---|---:|---:|---:|
-| float32 | `1e-4` | `1e-6` | `10 * eps` ≈ `1.19e-6` |
-| float64 | `1e-7` | `1e-9` | `10 * eps` ≈ `2.22e-15` |
-
-The policy never changes `jax_enable_x64`: a float32 `x_0` stays float32 even
-when x64 is enabled globally. Explicit tolerances override the defaults and
-are cast to the state dtype. Use explicit values whenever tolerances are part
-of a reproducibility or accuracy contract.
-
-Automatic `dt_min` is `10 * eps * max(1, abs(t_1))`; set it explicitly when
-the relevant time scale differs materially from the absolute horizon. The
-controller also floors exact-zero error ratios at machine epsilon before
-applying a negative power; the growth-factor clip then selects maximal step
-growth without introducing infinities.
-
-## The controller is stop-gradiented
-
-The adaptive step-size controllers (`IController` and `PIController`) compute
-their scaled error norms, decisions, and next-step factors inside
-`stop_gradient`. Accept/reject is a discrete branch, and differentiating
-`E**(-1/order)` is singular at the exact-zero error of a flat-start policy.
-
-States remain differentiable through solver stages, but the derivative holds
-the realized step sizes and accept/reject pattern fixed. It is the derivative
-of the discrete flow on a **frozen mesh**, not the total derivative of a mesh
-that moves with parameters. In particular, for parameter-only differentiation,
-the returned adaptive knot times in `SaveAt(steps=True)` have exactly zero
-tangent.
-
-The `E**(-1/5)` blow-up is not hypothetical: a policy initialized flat gives
-an exactly-zero error estimate on the first step, and without the
-`stop_gradient` (plus the machine-epsilon error floor) the backward pass is NaN
-from iteration one. `tests/test_ad.py::test_grad_finite_on_flat_field` pins
-this.
-
-`PIController` additionally carries the previous accepted error ratio. Its
-step-size factor is
-
-```text
-safety * E_n**(-(p_coeff + i_coeff) / order)
-       * E_prev**(p_coeff / order)
-```
-
-and `E_prev` changes only after acceptance. The whole recurrence is
-controller-internal and stop-gradiented. Setting `p_coeff=0, i_coeff=1`
-reproduces `IController` bit for bit.
-
-`IController` remains the default. On the smooth ODE and DAE screening
-problems, the default PI coefficients did not reduce rejections, and
-matched-work accuracy was comparable or favored I, except that PI sometimes
-improved common-grid ODE interpolation. PI remains useful as an opt-in for
-genuinely rejection-prone or step-size-oscillatory problems; those cases were
-not covered by this screen.
-
-## The horizon clip is the growth guard
-
-Every attempt is clipped so it cannot step past `t_1`, and — deliberately
-unlike diffrax — the controller's next-step proposal is computed from the
-**clipped** step:
-
-> the horizon clip doubles as the guard on step growth: without it, a
-> near-flat vector field lets steps quintuple into quarter-horizon leaps
-> whose Gauss–Newton linearization stalls a trust-region optimizer
-> differentiating through the rollout.
-
-With `factor_max = 5`, a flat field would otherwise reach `dt ≈ t_1/4` within
-a few accepted steps; residuals sampled from three or four giant steps make
-the optimizer's linear model useless. Clipping first means the proposal can
-never exceed `factor_max × remaining horizon`.
-
-Adaptive steps use `min(dt, remaining)` and are never enlarged to consume a
-floating-point sliver. Fixed-step times are instead formed arithmetically as
-`t_0 + i * dt_0`, avoiding drift from repeated addition, with a small local
-endpoint snap capped below a quarter step. No time tolerance scales with
-`max_steps`: changing a nonbinding attempt budget must not stretch a step or
-change the numerical method.
-
-## Frozen mesh versus a moving-mesh derivative
-
-Let a parameter-dependent controller choose an internal knot
-\(t_k(\theta)\), and let \(q(x,t,\theta)\) be a saved quantity. Along the exact
-trajectory, its total derivative contains
-
-$$
-\frac{d q_k}{d\theta}
-=q_x\left(x_\theta\rvert_t+f\,t_{k,\theta}\right)
- +q_t t_{k,\theta}+q_\theta.
-$$
-
-tinydiffeq's frozen-mesh rule returns the terms holding \(t_k\) fixed. The
-omitted mesh-motion term is
-
-$$
-\left(q_t+q_x f\right)t_{k,\theta}
-=\frac{d q}{dt}\,t_{k,\theta}.
-$$
-
-For a collocation residual \(r(x_k,t_k,\theta)\), this omission is the total
-trajectory derivative \((d r/dt)t_{k,\theta}\). If the learned policy approaches
-a smooth root for which the residual vanishes along the trajectory, then
-\(d r/dt\) vanishes with it locally; under that regularity, the omitted term is
-\(O(\lVert r\rVert)\), so the frozen-mesh residual Jacobian is asymptotically
-exact near the root. Pointwise cancellation at a few knots alone is not enough
-for that conclusion.
-
-Far from a root the omitted term can be material, so a Gauss--Newton or LM
-predicted reduction and trust ratio may be misleading. A targeted
-finite-difference audit found that adaptive `SaveAt(steps=True)` state
-derivatives need not approach the moving-mesh derivative under tolerance
-refinement. Endpoint output at fixed `t_1` and fixed requested-grid output do
-not have this particular moving-output-time ambiguity and did converge in the
-same audit.
-
-## Interpolation knots are non-differentiable
-
-`SaveAt(ts=...)` brackets each query with `searchsorted` — integer indices,
-no gradient. Values differentiate through the bracketing states and
-derivatives (`xL`, `xR`, `fL`, `fR`) while internal knot locations remain
-frozen. The requested times themselves are fixed outputs, unlike
-`SaveAt(steps=True)`'s adaptive internal times.
-
-Zero-width brackets (duplicate rows from rejections and the frozen tail) use
-the **double-where trick**: the divisor is replaced by 1 *before* dividing,
-
-```python
-width_safe = jnp.where(degenerate, 1.0, width)
-s = jnp.clip((tau - t_left) / width_safe, 0.0, 1.0)
-value = jnp.where(degenerate, x_left, hermite(s, ...))
-```
-
-so neither the primal nor its jvp/vjp ever evaluates `0/0`. A single `where`
-on the output is not enough — reverse mode differentiates both branches, and
-`NaN * 0 = NaN`.
-
-## What this buys you
-
-- With the default bounded loop, `jax.grad`, `jax.jvp`, and
-  `jax.grad(jax.jvp(...))` (the
-  Levenberg–Marquardt geodesic-acceleration pattern) all work through
-  adaptive solves and interpolated output, verified against closed forms in
-  `tests/test_ad.py`.
-- With `adaptive_loop="forward"`, primal evaluation, `jax.jvp`, and nested
-  forward mode work; `jax.vjp`, `jax.grad`, and reverse-over-forward do not.
-- `jax.vmap` over `x_0` or `p` gives genuinely per-lane adaptivity: each lane
-  accepts/rejects independently. Execution continues until the slowest lane
-  finishes.
-
-The dynamic loop is most useful when actual attempts are much fewer than
-`max_steps`. When an integration uses most of its budget, primal performance
-can be similar or slower than the bounded scan; compilation and forward-mode
-paths may still improve. Choose from representative end-to-end measurements,
-not from the loop form alone.
-
-On float32 GPUs, different loop lowerings can also fuse neural-network matrix
-products differently. With accelerator-default reduced-precision matrix
-multiplication, that roundoff was large enough in one width-32 policy test to
-move the adaptive mesh even though both loop contracts were correct. Setting
-`jax_default_matmul_precision="highest"` restored matching paths. tinydiffeq
-does not change this application-global JAX setting; set it in reproducibility
-entry points when adaptive decisions depend on neural-network outputs.
-
-## Reusing a linearization
-
-A single `jax.jvp` must evaluate the primal trajectory and propagate its
-tangent together. If several tangent directions are needed at the same
-`x_0`, `p`, and other primal inputs, cache the linearization instead:
-
-```python
-value, pushforward = jax.linearize(endpoint, x_0)
-tangent_batch = jax.jit(jax.vmap(pushforward))(directions)
-
-value, pullback = jax.vjp(endpoint, x_0)
-cotangent_batch = jax.jit(
-    jax.vmap(lambda cotangent: pullback(cotangent)[0])
-)(cotangents)
-```
-
-The setup computes and stores residuals for that one primal point. Reuse the
-pushforward or pullback only while every primal input is unchanged; otherwise
-linearize again. Rodas5P's custom linear-solve rule explicitly retains and
-reuses its primal pivoted LU factors. Semi-explicit DAE roots have a different
-boundary: nlls-gram supplies both the primal LM root and its implicit derivative,
-selecting direct `LU()` by default for a square residual system. The primal LM
-Cholesky factors are not reused because they represent damped normal equations,
-not the accepted root's nonsymmetric Jacobian. `jax.linearize` still avoids
-repeating the primal trajectory, but the API makes no promise that the DAE
-root's direct factorization is cached across pushforward directions.
-
-On the 256-state fixed Tsit5 benchmark, cached pushforwards were about 8–18%
-faster than a fused `vmap(jvp)` on CPU after excluding setup. On the RTX 3090,
-cached pushforwards and pullbacks were roughly 2.2–2.5x faster for 1–16
-directions because the primal trajectory was not replicated across mapped
-lanes. For one direction at a new primal point, ordinary `jax.jvp` or
-`jax.vjp` remains the right interface.
-
-## What to watch
-
-- Reductions over the valid prefix of `SaveAt(steps=True)` are
-  **discontinuous** in the inputs: its length changes when an accept flips to
-  a reject. Use `sol.accepted` when padding must not contribute. Because the
-  default tail repeats the endpoint, `xs[-1]` remains the reached final state.
-- Adaptive `SaveAt(steps=True)` derivatives hold the internal times fixed and
-  omit mesh motion. Use a fixed requested grid when the residual definition
-  requires fixed sample times, or treat the frozen-mesh Jacobian explicitly as
-  a near-root approximation.
-- Finite-difference checks of adaptive solves are noisy for the same reason;
-  compare AD against closed forms or use fixed-step solvers for FD tests.
-
-## Custom-rule audit
-
-tinydiffeq uses hand-coded derivative boundaries where they enforce the intended
-mathematical derivative or avoid differentiating iteration/factorization work:
-
-There is no custom reverse rule that replays an adaptive solve on a recorded
-mesh. The bounded path uses ordinary traced AD through its scan, subject to the
-controller `stop_gradient`; the dynamic forward path deliberately exposes
-JAX's no-reverse boundary. A replay/custom-VJP design would be a distinct future
-API and would not restore the omitted moving-mesh term by itself.
-
-- Rodas5P's factored linear solve has a custom JVP
-  \(\delta x=A^{-1}(\delta b-\delta A\,x)\). Every stage tangent and the
-  transposed VJP reuse the attempt's pivoted LU factors; pivot selection is not
-  differentiated.
-- Semi-explicit DAE roots delegate both the primal solve and implicit AD to
-  nlls-gram. Its implicit-function rule is
-  \(\delta z=-g_z^{-1}(g_y\delta y+g_t\delta t+g_p\delta p)\); the default
-  square `LU()` implementation differentiates the defining constraint rather
-  than the optimizer iterations, and reverse mode transposes that rule.
-  `LMRootSolver` requires residual-only convergence (`gtol=xtol=0`) and never
-  accepts `MAX_STEPS` as a differentiable root. `ad_solver` remains an
-  nlls-owned option forwarded through `solver_options`.
-- Dense linear exponential actions use a Fréchet custom JVP for active matrix
-  or time tangents and reuse the matrix exponential when only the initial state
-  varies.
-- Matrix-free terminal exponential sensitivities expose
-  `jvp_linear_ode`/`vjp_linear_ode`, applying the forward or transposed
-  exponential directly instead of differentiating Arnoldi orthogonalization.
-- `AdaptiveKrylovExponential` uses a bounded scan and residual-controlled
-  internal slices. Ordinary AD follows the realized controller path; the
-  hand-coded initial-state rules apply independent adaptive forward or
-  transposed actions and avoid differentiating the Arnoldi basis.
-
-Explicit Runge--Kutta, Euler--Maruyama, Hermite/Rodas interpolation, and
-cumulative trapezoids remain ordinary JAX programs. Their recurrences and
-polynomials already transpose efficiently, and a custom rule would either
-duplicate JAX's work or introduce a different continuous-adjoint derivative.
diff --git a/docs/aux.md b/docs/aux.md
deleted file mode 100644
index a0cfc1c..0000000
--- a/docs/aux.md
+++ /dev/null
@@ -1,158 +0,0 @@
-# Auxiliary Outputs
-
-Auxiliary outputs let a model expose quantities it already computes without
-adding them to the dynamical state. They are ordinary JAX pytrees, are saved
-with the solution, and participate in JVP and VJP. The no-aux code path remains
-available explicitly with `has_aux=False`.
-
-## ODE contract
-
-An ODE field may return only its derivative,
-
-```python
-def f(x, t, args, p):
-    return dx
-```
-
-or pair the derivative with saved output:
-
-```python
-def f(x, t, args, p):
-    return dx, {"flow": flow, "moment": moment}
-```
-
-The first item must have the state pytree structure and dtype. Saved aux must
-be a nonempty pytree of nonempty real floating arrays; its leaves may have
-different floating dtypes. The result is available as `sol.aux` with the same
-leading saved-time axis as `sol.xs`.
-
-`has_aux=None` (the default) performs one `jax.eval_shape` trace to detect the
-contract. A complete field output matching the state pytree takes precedence,
-so a two-item tuple state is not misclassified as `(dx, aux)`. Set
-`has_aux=True` to require aux, or `has_aux=False` to skip detection and select
-the minimal value-only path.
-
-## DAE contracts
-
-For a semi-explicit DAE, saved output belongs to the differential field:
-
-```python
-def f(y, z, t, args, p):
-    return dy, saved_aux
-
-
-def g(y, z, t, args, p):
-    return residual
-```
-
-The algebraic function may additionally expose internal context that avoids
-recomputing quantities already formed while evaluating the residual:
-
-```python
-def g(y, z, t, args, p):
-    return residual, algebraic_aux
-
-
-def f(y, z, t, args, p, algebraic_aux):
-    return dy, saved_aux
-```
-
-All four combinations are supported: neither aux, saved aux only, algebraic
-aux only, or both. When algebraic aux is present, the differential function
-must use the full six-argument signature. `has_algebraic_aux=None` detects the
-form; explicit `False` avoids that trace.
-
-Algebraic aux is internal context. It is passed to `f`, included in the
-implicit derivative path, and ignored by the nonlinear solver's residual
-interface. It is not independently stored or interpolated. It may be a
-nonempty pytree of bool, integer, real, or complex arrays; every inexact leaf
-must be finite. This broader dtype contract is useful for cached masks and
-indices that should not become solution output.
-
-Only the differential field's real-floating `saved_aux` becomes `sol.aux`.
-This separation keeps one unambiguous output to differentiate and interpolate:
-the algebraic function shares work with the dynamics, while the differential
-function decides what users retain.
-
-## SDE and SDAE contracts
-
-For an SDE, only the drift owns saved aux:
-
-```python
-def drift(x, t, args, p):
-    return drift_value, saved_aux
-
-
-def diffusion(x, t, args, p):
-    return diffusion_value
-```
-
-For an SDAE, algebraic aux is passed to both stochastic fields, and the drift
-still owns saved aux:
-
-```python
-def g(y, z, t, args, p):
-    return residual, algebraic_aux
-
-
-def drift(y, z, t, args, p, algebraic_aux):
-    return drift_value, saved_aux
-
-
-def diffusion(y, z, t, args, p, algebraic_aux):
-    return diffusion_value
-```
-
-The fixed random key defines the path for both primal evaluation and AD.
-Aux tangents and cotangents are therefore pathwise derivatives under common
-random numbers. SDE/SDAE aux is saved only at actual grid nodes;
-`SaveAt(ts=...)` remains unsupported because deterministic interpolation is
-not valid for rough paths.
-
-## Deterministic interpolation and AD
-
-`SaveAt(t_1=True)` evaluates saved aux only at the final state.
-`SaveAt(steps=True)` stores it at the initial and accepted nodes and applies
-the same prefix/padding mask as the state.
-
-By default, `SaveAt(ts=grid)` uses normalized cubic Hermite interpolation for
-ODE and root-restored DAE aux. Endpoint aux slopes are JVPs along the full
-solution velocity, so they include direct parameter dependence and indirect
-dependence through the state and any implicit algebraic root. Rodas5P obtains
-endpoint state velocities from its published stiff-aware continuous extension,
-then uses the same JVP construction for aux. No requested-time algebraic root
-or aux recomputation is performed on this dense-output path. See
-[Rodas5P](rodas5p.md) and [Semi-Explicit DAEs](dae.md) for the dense-output
-details and links to SciML's implementation.
-
-For an explicit fixed-step ODE with `SaveAt(ts=grid, exact=True)`, every
-requested time must instead coincide with a realized solver knot. The solver
-selects those states without Hermite interpolation and evaluates aux directly
-at the requested knots. Saved-aux validity is therefore checked only at those
-requested knots on this exact path.
-
-Ordinary JAX transformations compose through every saved or interpolated aux
-leaf. This includes `jax.jvp`, reverse-mode VJP/`jax.grad`, `vmap`, and
-reverse-over-forward. Adaptive accept/reject decisions and mesh selection
-retain the package's frozen-controller derivative convention. Consequently,
-adaptive `SaveAt(steps=True)` aux derivatives hold internal times fixed and
-omit mesh motion; fixed requested-grid aux does not have a moving output-time
-axis. See [Adaptive Stepping and AD](adaptive_ad.md#frozen-mesh-versus-a-moving-mesh-derivative).
-
-## Failure behavior
-
-Every saved aux leaf and every inexact algebraic-aux leaf must remain finite.
-Invalid algebraic context at initialization fails before time stepping.
-Invalid algebraic context at a later required stage/node, or invalid saved aux
-in a prefix-saving mode, terminates at the previous accepted node.
-Endpoint-only saved aux is evaluated after integration; if it is invalid, the
-endpoint state is retained, aux is zero-filled, and `sol.ok` is false.
-
-For batched differentiation where inactive lanes may leave the model domain,
-`failure_ad_reference` supplies a finite point used only for safe
-linearization. ODE/SDE references have the form `(x, t, p)`; DAE/SDAE
-references have `(y, z, t, p)`. Without one, tinydiffeq uses an all-ones
-best-effort reference. For DAE/SDAE roots, it is substituted before nlls only
-after a lane is already inactive. A newly attempted root and its model context
-must be JVP-safe at the actual initial point; the reference is not a post-hoc
-replacement for an active failure.
diff --git a/docs/rodas5p.md b/docs/rodas5p.md
deleted file mode 100644
index 2b02514..0000000
--- a/docs/rodas5p.md
+++ /dev/null
@@ -1,147 +0,0 @@
-# Rodas5P: Stiff ODEs and Index-1 DAEs
-
-`Rodas5P()` is an eight-stage, fifth-order Rosenbrock–Wanner method with an
-embedded adaptive estimator and a stiff-aware fourth-order continuous
-extension. It is A-stable and stiffly accurate. In tinydiffeq it supports
-ordinary ODEs and semi-explicit index-1 DAEs through the existing solve APIs.
-
-## Credit and implementation lineage
-
-Rodas5P was constructed by Gerd Steinebach and introduced and benchmarked in
-[Steinebach (2023)](https://doi.org/10.1007/s10543-023-00967-x). The method was
-developed in the Julia Differential Equations ecosystem, and tinydiffeq's JAX
-implementation deliberately follows SciML's authoritative implementation:
-
-- [SciML `OrdinaryDiffEqRosenbrock` package](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock)
-- [SciML Rodas5P tableau and dense coefficients](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_tableaus.jl)
-- [SciML consolidated Rosenbrock step implementation](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_perform_step.jl)
-- [SciML Rosenbrock dense interpolants](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_interpolants.jl)
-- [SciML solver documentation and selection guidance](https://docs.sciml.ai/DiffEqDocs/stable/api/ordinarydiffeq/massmatrixdae/Rosenbrock/)
-
-SciML's implementation is MIT-licensed. Its source is the reference for the
-port and receives explicit attribution here, in the kernel source, and in the
-cross-library regression tests and benchmarks.
-
-The tableau, stage equations, embedded estimate, and continuous extension are
-ported rather than redesigned. Regression tests compare fixed Rodas5P steps
-against values produced by SciML to prevent the implementations from silently
-diverging.
-
-## ODE use
-
-```python
-import jax.numpy as jnp
-
-from tinydiffeq import IController, Rodas5P, solve_ode
-
-
-def stiff_field(x, t):
-    # Exact solution x(t) = cos(t), with a fast transient eigenvalue -1000.
-    return -1000.0 * (x - jnp.cos(t)) - jnp.sin(t)
-
-
-sol = solve_ode(
-    stiff_field,
-    Rodas5P(),
-    0.0,
-    1.0,
-    jnp.asarray(1.0),
-    dt_0=0.01,
-    controller=IController(),
-    max_steps=512,
-)
-```
-
-At the beginning of an attempted step, Rodas5P forms the exact JAX Jacobian
-(J=\partial f/\partial x), the time derivative (f_t), and
-
-$$
-W = \frac{I}{\gamma h} - J.
-$$
-
-It computes one pivoted LU factorization of (W) and reuses it for all eight
-stage right-hand sides. It therefore solves linear systems but performs no
-Newton or Levenberg–Marquardt iteration.
-
-## Semi-explicit DAE use
-
-For
-
-$$
-\dot y=f(y,z,t), \qquad 0=g(y,z,t),
-$$
-
-tinydiffeq constructs the flattened mass-matrix system internally:
-
-$$
-M\dot u=F(u,t), \qquad
-u=(y,z), \quad M=\operatorname{diag}(I_y,0_z), \quad F=(f,g).
-$$
-
-```python
-from tinydiffeq import Rodas5P, solve_semi_explicit_dae
-
-
-sol = solve_semi_explicit_dae(
-    lambda y, z, t, args, p: p * z,
-    lambda y, z: z**2 - y - 2.0,
-    Rodas5P(),
-    0.0,
-    1.0,
-    jnp.asarray(1.0),
-    jnp.sqrt(jnp.asarray(3.0)),
-    p=jnp.asarray(-0.2),
-    dt_0=0.1,
-    controller=IController(),
-    max_steps=128,
-)
-```
-
-`LMRootSolver` is used once to make the initial algebraic state consistent.
-Every later Rodas5P stage solves
-
-$$
-\left(\frac{M}{\gamma h}-F_u\right)k_i=r_i
-$$
-
-with the reused LU factors. There is no nonlinear endpoint restoration.
-Consequently, returned internal `z` values satisfy the constraint to the
-method's integration accuracy, not necessarily to `LMRootSolver.atol`.
-
-This differs intentionally from the RK4 and Tsit5 DAE paths, which perform a
-nonlinear algebraic solve at every stage and accepted endpoint. Rodas5P is
-especially useful when those algebraic solves dominate runtime or when the
-coupled dynamics are stiff.
-
-## Dense output, aux, and AD
-
-`SaveAt(ts=...)` uses Rodas5P's SciML/paper continuous extension for the ODE
-state or combined `(y, z)` state. The three stored coefficient pytrees form a
-fourth-order stiff-aware polynomial; no query-time algebraic solve occurs.
-
-When the differential field returns `(value, saved_aux)`, aux is evaluated and
-stored at the required initial and accepted states. Requested-grid aux uses
-normalized cubic Hermite interpolation with endpoint tangents obtained from
-the Rodas5P polynomial and a JVP through the aux map. It is not recalculated
-or root-solved at each requested time. DAE algebraic aux is separate internal
-context and is never interpolated; see [Auxiliary Outputs](aux.md).
-
-The public derivative is the derivative of the discrete Rodas5P method.
-JVP/VJP propagate through Jacobian construction, stage solves, state output,
-and dense output. Each factored solve uses the mathematical linear-solve rule
-$dx=A^{-1}(db-dA\,x)$, so discrete LU pivot choices are not differentiated.
-As with the other adaptive methods, accept/reject decisions and step-size
-selection are stop-gradiented.
-
-## Deliberate limits
-
-- Dense exact Jacobians and dense pivoted LU only; no sparse, Krylov, or
-  preconditioner interface.
-- Constant identity or internally constructed block-diagonal mass matrices;
-  no public general mass-matrix API.
-- Semi-explicit index-1 DAEs with a square, locally nonsingular $g_z$.
-- No fully implicit residual-form DAE interface or automatic index reduction.
-
-For those broader capabilities, use SciML's
-[`OrdinaryDiffEqRosenbrock`](https://docs.sciml.ai/DiffEqDocs/stable/api/ordinarydiffeq/massmatrixdae/Rosenbrock/)
-and related DAE solvers.
diff --git a/docs/static_shapes.md b/docs/static_shapes.md
deleted file mode 100644
index b83f60f..0000000
--- a/docs/static_shapes.md
+++ /dev/null
@@ -1,156 +0,0 @@
-# Static Shapes
-
-JAX jits fixed-shape programs. An adaptive integrator is naturally dynamic —
-the number of steps depends on the data — so tinydiffeq offers two execution
-contracts. The default `adaptive_loop="bounded"` provides exactly `max_steps`
-static attempt slots in a bounded scan. Adaptive attempts are grouped into
-small nested-scan chunks so a completed solve skips whole padded chunks.
-`adaptive_loop="forward"` uses a dynamic `lax.while_loop` and executes only
-actual attempts, while retaining the same static public output shapes and a
-`max_steps` row buffer when the selected `SaveAt` mode requires one.
-
-Each iteration attempts one step:
-
-- an **accepted** attempt advances `(t, x)` and (for FSAL solvers) reuses the
-  last stage as the next first stage;
-- a **rejected** attempt leaves the state in place and retries with a smaller
-  step;
-- once `t_1` is reached, the remaining iterations **freeze**.
-
-The raw internal scan buffer contains repeated rows for rejected and frozen
-iterations. That buffer is an implementation detail used by interpolation;
-step output compacts it into accepted rows plus tail padding.
-
-The bounded loop preserves reverse-mode AD. Chunk-level and attempt-level
-`lax.cond` branches keep expensive field, stage, and controller computations
-out of its frozen tail. The forward loop supports primal evaluation, JVP, and
-nested forward mode, but JAX cannot transpose its data-dependent while loop.
-Under `vmap`, both strategies advance lanes together until the slowest lane
-finishes; the forward strategy stops at that lane's actual attempt count rather
-than always reaching `max_steps`.
-
-Fixed-step integration uses a smaller specialized scan without adaptive
-controller or embedded-error work. `ConstantStepSize` accepts every attempt,
-so `dt_0 = (t_1 - t_0)/n` with `max_steps = n` reproduces a fixed grid
-exactly. Times are formed arithmetically as `t_0 + i * dt_0`, rather than by
-repeatedly accumulating rounded steps. A small local tolerance snaps the
-nominal last point to `t_1`; it is capped relative to `dt_0` and never scales
-with `max_steps`. Consequently, increasing a nonbinding attempt budget cannot
-stretch an earlier step or otherwise change the numerical method.
-
-If the budget runs out before `t_1`, `sol.ok` is `False` and the outputs hold
-the reached prefix. The package never poisons values; the caller decides:
-
-```python
-xs = jnp.where(sol.ok, sol.xs, jnp.inf)  # kernels-style rejection
-```
-
-`sol.num_steps` counts attempts actually made, including rejections;
-`sol.num_accepted` counts advances. These scalar diagnostics do not change the
-fixed output shape.
-
-## SaveAt is the shape contract
-
-Exactly one of three modes:
-
-### `SaveAt(t_1=True)` — endpoint only (default)
-
-`sol.ts` is the reached time (equals `t_1` when `ok`), `sol.xs` the final
-state.
-
-### `SaveAt(ts=grid)` — interpolation onto a fixed grid
-
-This is the answer to "adaptive steps vs static shapes". Internal steps
-adapt freely; the output is dense interpolation onto your fixed
-query grid, so each output leaf has shape
-`(len(grid),) + corresponding_input_leaf.shape` **regardless of how many
-steps the controller took**. Changing tolerances, initial conditions,
-or curvature changes the internal knots but never the output shape — no
-recompilation. A one-dimensional JAX/NumPy array or Python sequence is
-accepted; times must be nondecreasing and within `[t_0, t_1]`, while repeated
-times and omitted endpoints are allowed. Changing values without changing
-the grid length does not recompile.
-
-By default these are observation times, not mandatory internal stops. The
-adaptive controller chooses the same mesh regardless of the requested grid,
-then the solver evaluates each point through dense interpolation.
-
-For an explicit ODE with `ConstantStepSize`,
-`SaveAt(ts=grid, exact=True)` instead requires every query to coincide with a
-realized internal endpoint and gathers the stored states directly. It avoids
-interpolation and the endpoint-slope work needed by Hermite output. A
-misaligned or unreached query makes `sol.ok` false. Exact mode is unavailable
-for adaptive ODEs, Rodas5P, DAEs, SDEs, and SDAEs.
-
-The interpolation runs directly over the raw padded rows: duplicate knots
-from rejections or the frozen tail form zero-width brackets, and the
-bracketing `searchsorted` lands on the last duplicate at-or-before each
-query, so no compaction pass is needed. Queries outside the knot span clamp
-to the boundary values — in particular, when `ok` is `False`, queries beyond
-the reached time return the last state (flat extrapolation) rather than
-evaluating a cubic outside its bracket.
-
-The interpolant is 4th-order accurate between 5th-order-accurate knots:
-expect grid values slightly less accurate than the knots themselves, which
-is the standard dense-output trade-off.
-
-### `SaveAt(steps=True)` — accepted steps with padding
-
-`max_steps + 1` rows including the initial state. Accepted internal steps are
-gathered chronologically into a contiguous prefix; rejected attempts are not
-returned. `sol.accepted` is the validity mask (`accepted[0]` is always
-`True`, so `accepted.sum() == num_accepted + 1`). On a successful solve, the
-endpoint is at index `num_accepted`.
-
-- `fill="last"` (default) repeats the last valid time and state through the
-  padded tail.
-- `fill="inf"` fills only the invalid tail with `inf`.
-
-If the attempt budget is exhausted, the same contract holds for the reached
-accepted prefix, the tail repeats or masks its last state, and `sol.ok` is
-`False`. No fake endpoint is inserted. Compaction is a stable fixed-size
-gather; it performs no sorting.
-
-## Pytree states
-
-ODE, SDE, and DAE states may be arbitrary JAX pytrees, including registered
-dataclasses. Each state contains at least one nonempty real floating array,
-and all leaves use one dtype. Explicit and stochastic solver arithmetic maps
-directly over leaves, so a single-array state retains its array execution
-path. Rodas5P temporarily ravels the state because its dense Jacobian and LU
-factorization couple all coordinates, then reconstructs the original pytree
-at every public boundary. In every case the structure is resolved while
-tracing. Changing leaf values with unchanged shapes and structure reuses a
-compilation; changing the treedef or a leaf shape requires a new compilation.
-
-For every multi-row `SaveAt` mode the leading row dimension is added to each
-leaf independently. `sol.accepted` is one shared mask for all leaves.
-
-## Why one compilation, precisely
-
-- The attempt budget `max_steps` is static; nothing else about the loop depends
-  on data shapes.
-- Tolerances and PI coefficients (`IController(...)` / `PIController(...)`),
-  `dt_0`, `t_0`, `t_1`, `x_0`,
-  `args`, `p`, and `SaveAt.ts` are pytree **data leaves**. Only genuine
-  structure — the solver type, `SaveAt` mode, `fill`, `exact`, `max_steps`,
-  `adaptive_loop`, and the functions themselves — is static.
-
-An omitted tolerance or `dt_min` is represented by `None`, so switching a
-jitted call between automatic and explicit values changes the controller
-pytree structure and compiles once for each policy. Changing the numeric
-values of already-explicit controller fields does not recompile.
-
-So this compiles once:
-
-```python
-@jax.jit
-def run(x_0, dt_0, controller, args):
-    return solve_ode(f, Tsit5(), 0.0, 1.0, x_0, args=args, dt_0=dt_0,
-                     controller=controller, max_steps=128,
-                     save_at=SaveAt(steps=True))
-```
-
-across different curvatures (different accepted counts), tolerances, initial
-steps, and initial conditions — pinned by `tests/test_recompile.py` with
-`_cache_size() == 1` assertions.
diff --git a/src/tinydiffeq/__init__.py b/src/tinydiffeq/__init__.py
index f16f738..c347abc 100644
--- a/src/tinydiffeq/__init__.py
+++ b/src/tinydiffeq/__init__.py
@@ -1,31 +1,17 @@
 """Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX.
 
 solve_ode integrates dx/dt = f(x, t, args, p) with fixed-step (Euler, RK4),
-adaptive explicit Tsit5, or linearly implicit Rodas5P methods. Fixed stepping
-and the default adaptive path use bounded lax.scan loops with exactly max_steps
-attempt slots and support forward/reverse AD, including reverse-over-forward.
-Adaptive ODE/DAE solves may instead use an actual-work lax.while_loop for
-primal, JVP, and nested forward AD; reverse mode is unsupported on that path.
-SaveAt picks the endpoint, dense interpolation onto a fixed grid, accepted
-internal steps with padding, or exact knot gathering for explicit fixed-step
-ODEs. max_steps bounds attempted steps; Solution.num_steps reports actual work,
-and arithmetic fixed-step times are independent of a nonbinding budget.
-solve_sde is fixed-step
-Euler-Maruyama with presampled
-diagonal noise. solve_semi_explicit_dae handles index-1 systems with either
-root-restored explicit methods or the stiff Rodas5P mass-matrix formulation,
-plus differentiable saved differential-field aux and internal algebraic
-context. DAESolution also reports nonlinear-root calls and LM update counts;
-LMRootSolver delegates primal roots and square implicit derivatives to
-nlls-gram, requires residual-only stopping, and offers previous-root and secant
-stage predictors.
-solve_semi_explicit_sdae applies fixed-step Euler-Maruyama to the reduced
-index-1 stochastic system. solve_linear_ode applies dense, fixed-Krylov, or
-adaptive matrix-free exponential actions to fixed homogeneous linear systems.
-States may be arrays or pytrees of same-dtype real floating arrays. Fully
-implicit solvers, general mass
-matrices, full derivative-term PID control, events, continuous interpolation
-objects, and adjoint methods are non-goals.
+adaptive Tsit5, or linearly implicit Rodas5P methods on bounded lax.scan loops
+with static shapes and composable forward/reverse AD. solve_sde integrates
+diagonal-noise Ito SDEs (EulerMaruyama, Milstein, SRA1) from a PRNG key or an
+explicit, differentiable noise pytree. solve_semi_explicit_dae and
+solve_semi_explicit_sdae handle index-1 systems, delegating algebraic roots
+and their implicit derivatives to nlls-gram. solve_linear_ode applies dense or
+matrix-free Krylov exponential actions to fixed homogeneous linear systems,
+and the Markov tools simulate and forecast finite-state chains. States may be
+arrays or pytrees of same-dtype real floating arrays. Fully implicit solvers,
+general mass matrices, events, continuous interpolation objects, and adjoint
+methods are non-goals.
 """
 
 from tinydiffeq.controllers import ConstantStepSize, IController, PIController
@@ -58,7 +44,16 @@
 from tinydiffeq.sdae import solve_semi_explicit_sdae
 from tinydiffeq.sde import solve_sde
 from tinydiffeq.solution import DAESolution, Solution
-from tinydiffeq.solvers import RK4, Euler, EulerMaruyama, Rodas5P, Tsit5
+from tinydiffeq.solvers import (
+    RK4,
+    SRA1,
+    Euler,
+    EulerMaruyama,
+    Milstein,
+    Rodas5P,
+    Tsit5,
+    diagonal_brownian_increments,
+)
 
 __all__ = [
     "solve_ode",
@@ -73,6 +68,9 @@
     "Tsit5",
     "Rodas5P",
     "EulerMaruyama",
+    "Milstein",
+    "SRA1",
+    "diagonal_brownian_increments",
     "ConstantStepSize",
     "IController",
     "PIController",
diff --git a/src/tinydiffeq/controllers.py b/src/tinydiffeq/controllers.py
index e5914ec..2f074f2 100644
--- a/src/tinydiffeq/controllers.py
+++ b/src/tinydiffeq/controllers.py
@@ -65,14 +65,10 @@ class IController:
     Accept iff ``E = max(|err| / (atol + rtol * max(|x_0|, |x_1|))) <= 1``
     (forced accept once the step reaches ``dt_min``), and propose
     ``dt_next = dt_used * clip(safety * E**(-1/order), factor_min, factor_max)``
-    clipped to ``[dt_min, dt_max]``. This is the classic integral controller —
-    equal to diffrax's ``PIDController`` at its default coefficients
-    (p_coeff=0, d_coeff=0); there is no proportional term, hence the name.
-
-    If omitted, ``rtol``/``atol`` default to ``1e-4``/``1e-6`` for float32
-    states and ``1e-7``/``1e-9`` for float64 states. Explicit tolerances are
-    cast to the state dtype. ``dt_min`` defaults to ten machine epsilons in the
-    time dtype, scaled by ``max(1, |t_1|)``.
+    clipped to ``[dt_min, dt_max]``. Omitted ``rtol``/``atol`` default to
+    ``1e-4``/``1e-6`` for float32 states and ``1e-7``/``1e-9`` for float64;
+    ``dt_min`` defaults to ten machine epsilons in the time dtype, scaled by
+    ``max(1, |t_1|)``.
     """
 
     rtol: float | None = None
@@ -126,19 +122,12 @@ class PIController:
     """Proportional-integral step-size controller with max-norm error.
 
     In addition to the current scaled error ``E``, this controller carries the
-    previous accepted step's error ``E_prev`` and proposes
-
+    previous accepted step's error ``E_prev`` (starting at one) and proposes
     ``dt_next = dt_used * clip(safety * E**(-(p_coeff+i_coeff)/order)``
-    ``* E_prev**(p_coeff/order), factor_min, factor_max)``.
-
-    ``E_prev`` starts at one and changes only after an accepted step. The
-    defaults ``p_coeff=0.4`` and ``i_coeff=0.3`` damp step-size oscillations on
-    harder problems; ``p_coeff=0, i_coeff=1`` reproduces :class:`IController`.
-    The error ratios and step-size update are stop-gradiented, while states
-    remain fully differentiable through the solver stages. If omitted,
-    ``rtol``/``atol`` default to ``1e-4``/``1e-6`` for float32 states and
-    ``1e-7``/``1e-9`` for float64 states. ``dt_min`` defaults to ten machine
-    epsilons in the time dtype, scaled by ``max(1, |t_1|)``.
+    ``* E_prev**(p_coeff/order), factor_min, factor_max)``. The defaults
+    ``p_coeff=0.4`` and ``i_coeff=0.3`` damp step-size oscillations;
+    ``p_coeff=0, i_coeff=1`` reproduces :class:`IController`, and the
+    tolerance and ``dt_min`` defaults match it.
     """
 
     rtol: float | None = None
diff --git a/src/tinydiffeq/dae.py b/src/tinydiffeq/dae.py
index d26a952..dbf15ca 100644
--- a/src/tinydiffeq/dae.py
+++ b/src/tinydiffeq/dae.py
@@ -86,54 +86,24 @@
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True)
 class LMRootSolver:
-    """Configuration for algebraic solves in a semi-explicit DAE.
+    """Configuration for algebraic root solves in semi-explicit DAEs/SDAEs.
 
     The implementation is :class:`nlls_gram.LevenbergMarquardt` at its
-    defaults: dense ``Cholesky()`` for primal LM updates and the direct
-    nonsymmetric ``LU()`` implicit derivative that nlls selects automatically
-    for a square residual system.
-
-    The fields here are the ones this package owns; all of them reach the nlls
-    ``solve`` rather than its constructor. ``max_steps`` counts nonlinear
-    iterations for one algebraic root and is independent of the integration's
-    time-step ``max_steps``. Every accepted algebraic root must have Euclidean
-    residual norm strictly below the root ``atol``. Root solves therefore use
-    residual stopping only: ``gtol`` and ``xtol`` must remain zero, and a
-    ``MAX_STEPS`` iterate is never treated as a differentiable root.
-    ``max_steps_is_success`` remains in the configuration for source
-    compatibility but does not broaden root acceptance. ``atol=None`` selects
-    ``1e-6`` in float32 and ``1e-10`` in float64. Root tolerances are
-    deliberately independent of the outer integration tolerances.
-
-    ``solver_options`` is the escape hatch for the rare root that needs a
-    non-default algorithm: a mapping (or pairs) forwarded verbatim to the
+    defaults. ``max_steps`` bounds one algebraic root's nonlinear iterations,
+    independently of the integration's time-step budget. Roots use residual
+    stopping only: ``gtol`` and ``xtol`` must remain zero, and every accepted
+    root must report ``CONVERGED`` with Euclidean residual norm strictly below
+    ``atol`` (``None`` selects ``1e-6`` in float32, ``1e-10`` in float64).
+    ``solver_options`` is a mapping (or pairs) forwarded verbatim to the
     ``LevenbergMarquardt`` constructor, e.g.
-    ``solver_options={"linear_solver": QR()}``. Names and semantics are
-    nlls-gram's, not this package's, so they track it across versions instead
-    of being mirrored here. It is normalized to a sorted tuple so equal
-    configurations stay hashable and share one compiled solver -- an
-    unhashable config silently rebuilds the solver per call, and nlls-gram
-    keys its compiled loop on solver identity, so that would retrace every
-    step. ``cache_jacobian`` and ``geodesic_acceleration`` are fixed to
-    ``False`` and rejected here: each DAE stage changes the root problem, and
-    the intended path is the ordinary dense LM step. Algebraic residuals do
-    not expose nlls aux. ``ad_solver`` remains an nlls-owned option; its default
-    is the direct square ``LU()`` rule.
-
-    ``predictor="previous"`` warm-starts every explicit RK stage from the most
-    recent successful root. ``predictor="secant"`` instead extrapolates from
-    the accepted-step root through the most recent successful stage at a later
-    time. The secant is used only for a strictly later target; duplicate RK4
-    stage times and failed stages fall back to the previous root. Predictor
-    values are differentiation-inert, so successful roots retain the same
-    implicit derivative. With multiple algebraic roots, however, a different
-    warm start can select a different branch; the secant mode is intended only
-    when the continued root is locally unique. A finite root-iteration budget
-    can also make predictor choice affect values or success status.
+    ``solver_options={"linear_solver": QR()}``; ``cache_jacobian`` and
+    ``geodesic_acceleration`` are fixed to ``False``. ``predictor`` selects
+    the explicit-stage warm start: ``"previous"`` reuses the most recent
+    successful root, ``"secant"`` extrapolates it to a strictly later stage
+    time and assumes the continued root is locally unique.
     """
 
     max_steps: int = field(default=8, metadata=dict(static=True))
-    max_steps_is_success: bool = field(default=False, metadata=dict(static=True))
     atol: float | None = field(default=None, metadata=dict(static=True))
     gtol: float = field(default=0.0, metadata=dict(static=True))
     xtol: float = field(default=0.0, metadata=dict(static=True))
@@ -145,8 +115,6 @@ def __post_init__(self):
             raise ValueError("LMRootSolver.max_steps must be a positive int")
         if self.max_steps <= 0:
             raise ValueError("LMRootSolver.max_steps must be a positive int")
-        if not isinstance(self.max_steps_is_success, bool):
-            raise TypeError("LMRootSolver.max_steps_is_success must be a bool")
         if self.atol is not None and self.atol <= 0:
             raise ValueError("LMRootSolver.atol must be positive or None")
         if self.gtol != 0:
@@ -815,53 +783,21 @@ def solve_semi_explicit_dae(
 ):
     """Integrate a semi-explicit index-1 DAE.
 
-    The system is ``dy/dt = f(y, z, t, args, p)`` and
-    ``0 = g(y, z, t, args, p)``, with a square nonsingular algebraic Jacobian
-    ``dg/dz``. ``z_0`` is a root-finding guess: the initial algebraic state is
-    made consistent automatically, and its derivative is determined by the
-    constraint rather than by the guess.
-
-    RK4 with fixed control, Tsit5 with fixed or adaptive control, and the
-    linearly implicit Rodas5P method with fixed or adaptive control are
-    supported. The outer ``max_steps`` bounds attempted time steps;
-    :class:`LMRootSolver.max_steps` separately bounds algebraic solves. RK4
-    and Tsit5 restore ``g=0`` at every stage. Rodas5P uses the root solver only
-    for initial consistency, then advances the block mass-matrix system with
-    reused linear solves; later ``z`` values satisfy the constraint to the
-    integration accuracy rather than the root tolerance.
-    ``sol.num_root_solves`` counts logical active nonlinear root calls and
-    ``sol.num_root_steps`` sums their LM update steps. Rodas5P therefore reports
-    one root call regardless of its number of integration attempts.
-
-    ``args`` is fixed data by convention. All differentiated model parameters
-    belong in ``p``. Initial consistency and explicit-method roots
-    differentiate implicitly with respect to ``(y, t, p)``; Rodas5P then
-    differentiates its discrete Jacobians and uses implicit derivatives for
-    its linear solves.
-
-    ``f`` may return ``dy`` or ``(dy, saved_aux)``. If ``g`` returns
-    ``(residual, algebraic_aux)``, that second value is internal context and
-    ``f`` must take the full six-argument form
-    ``f(y, z, t, args, p, algebraic_aux)``. Only ``saved_aux`` is returned as
-    ``sol.aux``. It is stored at accepted nodes and interpolated on requested
-    deterministic grids; algebraic aux is never stored or interpolated.
-    ``has_aux`` and ``has_algebraic_aux`` default to abstract auto-detection;
-    explicit ``False`` selects the minimal paths without those traces.
-
-    ``failure_ad_reference=(y, z, t, p)`` may provide a domain-safe point for
-    already-inactive ``vmap`` lanes and model aux/field evaluation. A newly
-    attempted root still requires its actual ``(y, z_guess, t, p)`` to be
-    JVP-safe; the reference does not replace an active root after it fails.
-    A nonfinite inexact algebraic-aux leaf at initialization prevents all
-    time-step work. Saved aux is checked at the initial and accepted nodes in
-    prefix/grid modes; endpoint mode checks it only after integration and
-    retains the endpoint state with zero aux if that check fails.
-
-    ``adaptive_loop="bounded"`` keeps the reverse-mode-capable static scan.
-    ``adaptive_loop="forward"`` executes only actual adaptive attempts and,
-    under ``vmap``, stops after the slowest lane. It supports JVP and nested
-    forward mode but not ordinary reverse mode, matching JAX's dynamic-while
-    differentiation boundary.
+    The system is ``dy/dt = f(y, z, t, args, p)`` with
+    ``0 = g(y, z, t, args, p)`` and a square nonsingular ``dg/dz``. ``z_0``
+    is a root guess: initial consistency is solved automatically and its
+    derivative comes from the constraint. RK4 and Tsit5 restore ``g = 0`` at
+    every stage through :class:`LMRootSolver`; Rodas5P advances the block
+    mass-matrix system with one reused LU factorization per attempted step,
+    so its later ``z`` values satisfy the constraint to integration accuracy.
+    Roots and their implicit derivatives are delegated to nlls-gram. ``f``
+    may return ``(dy, saved_aux)``; ``g`` may return
+    ``(residual, algebraic_aux)``, in which case ``f`` takes
+    ``(y, z, t, args, p, algebraic_aux)``.
+    ``failure_ad_reference=(y, z, t, p)`` provides a domain-safe point for
+    inactive ``vmap`` lanes. ``adaptive_loop`` follows
+    :func:`tinydiffeq.solve_ode`. Returns a :class:`DAESolution` with
+    root-solve diagnostics.
     """
     if dt_0 is None:
         raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)")
diff --git a/src/tinydiffeq/exponential.py b/src/tinydiffeq/exponential.py
index 15de6dd..ac13829 100644
--- a/src/tinydiffeq/exponential.py
+++ b/src/tinydiffeq/exponential.py
@@ -420,18 +420,13 @@ def solve_linear_ode(operator, method, t_0, t_1, x_0, *, save_at=None):
     """Solve ``dx/dt = A(x)`` for a fixed homogeneous linear operator.
 
     ``operator`` is either a square matrix using the column convention
-    ``A @ x`` or a callable that maps the state pytree to an identically
+    ``A @ x`` or a callable mapping the state pytree to an identically
     structured pytree. ``DenseExponential`` materializes a callable operator
-    with forward-mode Jacobian columns before applying a dense matrix
-    exponential. ``KrylovExponential`` and ``AdaptiveKrylovExponential`` only
-    evaluate operator actions and are therefore suitable for structured
-    matrix-free pytrees.
-
-    Endpoint output is the default. ``SaveAt(ts=...)`` evaluates independent
-    exponential actions at the requested times. The solve supports ordinary
-    JVPs and VJPs through the initial state and through differentiable arrays
-    used by the operator. The operator must be autonomous, homogeneous, and
-    linear; affine and nonlinear exponential methods are a separate problem.
+    before a dense matrix exponential; the Krylov methods evaluate only
+    operator actions. Endpoint output is the default, and ``SaveAt(ts=...)``
+    evaluates independent exponential actions at the requested times. JVPs
+    and VJPs flow through the initial state and differentiable operator
+    arrays. The operator must be autonomous, homogeneous, and linear.
     """
     if not isinstance(method, _EXPONENTIAL_METHODS):
         raise TypeError(
diff --git a/src/tinydiffeq/ode.py b/src/tinydiffeq/ode.py
index 64045bc..abba2e9 100644
--- a/src/tinydiffeq/ode.py
+++ b/src/tinydiffeq/ode.py
@@ -132,47 +132,25 @@ def solve_ode(
     has_aux=None,
     failure_ad_reference=None,
     adaptive_loop="bounded",
+    unroll=1,
 ):
     """Integrate ``dx/dt = f(x, t, args, p)`` from ``t_0`` to ``t_1 > t_0``.
 
-    The vector field may be declared ``f(x)``, ``f(x, t)``, ``f(x, t, args)``,
-    or ``f(x, t, args, p)`` — always in that order. ``x`` is a pytree whose
-    nonempty leaves share one real floating dtype. ``args`` is pass-through
-    data that is by convention not an AD target, and ``p`` holds differentiable
-    parameters (any pytree);
-    jvp/vjp with respect to ``p`` and ``x_0`` are first-class.
-
-    Fixed stepping and the default ``adaptive_loop="bounded"`` use bounded
-    ``lax.scan`` loops with exactly ``max_steps`` attempt slots, so shapes are
-    static and curvature-dependent step counts never retrace. Bounded adaptive
-    attempts are grouped into static chunks so one ``lax.cond`` skips an entire
-    padded chunk after completion. ``adaptive_loop="forward"`` instead uses a
-    true ``lax.while_loop`` and stops after the actual attempts. It supports
-    primal evaluation, JVP, and nested forward-mode AD; ordinary reverse mode
-    is unavailable because JAX cannot transpose a dynamic ``while_loop``.
-    Under ``vmap``, the forward loop advances lanes together until the slowest
-    live trajectory finishes rather than running every lane to ``max_steps``.
-    ``dt_0`` is required (no auto-initial-step heuristic).
-    Each attempt is clipped to the remaining horizon; the clipped step also
-    feeds the controller's next-step proposal, which doubles as the growth
-    guard — near-flat fields otherwise grow steps into quarter-horizon leaps.
-    ``project`` (e.g. a positivity clamp, assumed idempotent) is applied at
-    every point where ``f`` is evaluated and to every accepted state. Returns
-    a :class:`Solution`; ``sol.ok`` is False if the budget ran out before
-    ``t_1`` or a required saved output was invalid (outputs remain a finite
-    reached prefix or endpoint, never poisoned).
-
-    The field may return either ``dx`` or ``(dx, aux)``. ``has_aux=None``
-    detects the form with an abstract trace; ``has_aux=False`` selects the
-    minimal no-aux path without that trace. Saved aux is a nonempty pytree of
-    real floating arrays. It follows ``SaveAt`` and participates in JVP/VJP.
-    By default, requested-grid aux uses cubic Hermite interpolation with
-    endpoint slopes obtained by JVP, including for Rodas5P's dense state path.
-    With ``SaveAt(ts=..., exact=True)``, an explicit fixed-step solve instead
-    selects realized knots and evaluates aux directly at those requested knots.
-
-    The time dtype follows the state dtype; the library never
-    sets ``jax_enable_x64`` — do that in your application.
+    The field may be declared ``f(x)``, ``f(x, t)``, ``f(x, t, args)``, or
+    ``f(x, t, args, p)``; ``x`` is an array or pytree with one real floating
+    dtype, ``args`` is inert data, and ``p`` holds differentiable parameters.
+    ``dt_0`` is required. Fixed stepping and the default
+    ``adaptive_loop="bounded"`` use bounded ``lax.scan`` loops with exactly
+    ``max_steps`` attempt slots and support forward and reverse AD;
+    ``adaptive_loop="forward"`` runs a dynamic ``lax.while_loop`` (primal,
+    JVP, and nested forward mode only). ``project`` (an idempotent clamp) is
+    applied at every field evaluation and accepted state. The field may
+    return ``(dx, aux)``; saved aux follows ``SaveAt`` and participates in
+    AD. ``unroll`` (a static int, fixed stepping only) unrolls that many
+    steps per iteration of the integration scan — identical values,
+    fewer/larger GPU dispatches, more compile time. Returns a
+    :class:`Solution`; ``sol.ok`` reports whether ``t_1`` was reached with
+    every requested output valid — outputs are never poisoned.
     """
     if dt_0 is None:
         raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)")
@@ -196,6 +174,10 @@ def solve_ode(
         raise ValueError(
             'adaptive_loop="forward" requires an adaptive error controller'
         )
+    if not isinstance(unroll, int) or isinstance(unroll, bool) or unroll < 1:
+        raise ValueError("unroll must be a static int of at least 1")
+    if unroll != 1 and controller.uses_error_estimate:
+        raise ValueError("unroll requires fixed stepping (ConstantStepSize)")
     f = canonicalize_field(f)
     is_rodas = isinstance(solver, Rodas5P)
     is_fixed = isinstance(controller, ConstantStepSize)
@@ -589,11 +571,11 @@ def outer(chunk_carry, chunk_valid):
         if static_uniform_horizon:
             step_indices = jnp.arange(max_steps, dtype=jnp.int32)
             fixed_final, rows = jax.lax.scan(
-                uniform_fixed_body, fixed_carry_0, step_indices
+                uniform_fixed_body, fixed_carry_0, step_indices, unroll=unroll
             )
         else:
             fixed_final, rows = jax.lax.scan(
-                fixed_body, fixed_carry_0, None, length=max_steps
+                fixed_body, fixed_carry_0, None, length=max_steps, unroll=unroll
             )
         t_final, x_final, _, done, num_accepted = fixed_final
         num_steps = num_accepted
@@ -635,7 +617,9 @@ def outer(chunk_carry, chunk_valid):
             _,
         ) = final_carry
     else:
-        final_carry, rows = jax.lax.scan(body, carry_0, None, length=max_steps)
+        final_carry, rows = jax.lax.scan(
+            body, carry_0, None, length=max_steps, unroll=unroll
+        )
         (
             t_final,
             x_final,
diff --git a/src/tinydiffeq/save_at.py b/src/tinydiffeq/save_at.py
index 722665a..a200fe7 100644
--- a/src/tinydiffeq/save_at.py
+++ b/src/tinydiffeq/save_at.py
@@ -7,27 +7,17 @@
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True)
 class SaveAt:
-    """What ``solve_ode``/``solve_sde`` return. Exactly one mode must be set.
+    """What the solve functions return. Exactly one mode must be set.
 
-    - ``t_1=True``: the endpoint only (the default in the solve functions).
-    - ``ts=grid``: dense interpolation of the internal steps onto a fixed,
-      sorted query grid in ``[t_0, t_1]``. Explicit methods use cubic Hermite;
-      Rodas5P uses its stiff-aware fourth-order extension. Output shape is
-      ``(len(ts), ...)`` regardless of how many internal steps the controller
-      takes, so changing curvature never changes shapes or recompiles.
-      ``ts`` is a data leaf; a different grid of the same length retraces
-      nothing. ODE, deterministic DAE, and linear exponential solves only.
-      For an explicit constant-step ODE, ``exact=True`` instead requires every
-      query to coincide with an internal step endpoint. It gathers those states
-      directly and avoids both interpolation and its extra endpoint-slope field
-      evaluation.
-    - ``steps=True``: the initial state and accepted internal steps as a
-      chronological prefix of a ``max_steps + 1`` buffer. Rejected attempts
-      are omitted. ``fill="last"`` (default) pads the tail with the final
-      valid row; ``fill="inf"`` pads the tail with ``inf``. The returned
-      ``Solution.accepted`` mask distinguishes the valid prefix from padding.
-
-    ``fill`` only applies to ``steps=True``.
+    ``t_1=True`` (the solver default) returns the endpoint only. ``ts=grid``
+    interpolates the internal steps onto a fixed query grid — output shape is
+    ``(len(ts), ...)`` however many steps the controller takes, and ``ts`` is
+    a data leaf; with ``exact=True`` an explicit constant-step ODE instead
+    gathers states at queries that must coincide with realized knots.
+    ``steps=True`` returns the initial state and accepted steps as the valid
+    prefix of a ``max_steps + 1`` buffer, padded with the last valid row
+    (``fill="last"``) or ``inf`` (``fill="inf"``) and masked by
+    ``Solution.accepted``.
     """
 
     t_1: bool = field(default=False, metadata=dict(static=True))
diff --git a/src/tinydiffeq/sdae.py b/src/tinydiffeq/sdae.py
index d510c97..da5d028 100644
--- a/src/tinydiffeq/sdae.py
+++ b/src/tinydiffeq/sdae.py
@@ -1,4 +1,4 @@
-"""Fixed-step Euler--Maruyama for semi-explicit index-1 SDAEs."""
+"""Fixed-step stochastic integration for semi-explicit index-1 SDAEs."""
 
 import jax
 import jax.numpy as jnp
@@ -29,9 +29,8 @@
     _prepare_failure_ad_reference,
 )
 from tinydiffeq.save_at import SaveAt
-from tinydiffeq.sde import _diagonal_brownian_increments
 from tinydiffeq.solution import DAESolution
-from tinydiffeq.solvers import EulerMaruyama
+from tinydiffeq.solvers import INV_SQRT_3, SRA1, EulerMaruyama
 
 
 def solve_semi_explicit_sdae(
@@ -56,41 +55,23 @@ def solve_semi_explicit_sdae(
 ):
     """Integrate a semi-explicit index-1 Ito SDAE with diagonal noise.
 
-    The system is ``dy = drift(y, z, t) dt + diffusion(y, z, t) dW`` and
-    ``0 = g(y, z, t)``. Euler--Maruyama advances the differential state on a
-    fixed uniform grid, then an algebraic root solve restores consistency at
-    every node. This is Euler--Maruyama applied to the reduced SDE obtained
-    from the locally unique root ``z = Z(y, t)``.
-    ``sol.num_root_solves`` counts logical active nonlinear root calls and
-    ``sol.num_root_steps`` sums their LM update steps.
-
-    ``drift`` may return ``value`` or ``(value, saved_aux)``. If ``g`` returns
-    ``(residual, algebraic_aux)``, that internal context is passed to both
-    ``drift(y, z, t, args, p, algebraic_aux)`` and the corresponding
-    six-argument ``diffusion``. Only drift-owned ``saved_aux`` is exposed as
-    ``sol.aux`` and it is stored at consistent stochastic nodes; stochastic
-    interpolation is deliberately unsupported. ``has_aux`` and
-    ``has_algebraic_aux`` default to abstract auto-detection, while explicit
-    ``False`` selects the minimal no-aux paths.
-
-    A fixed key defines
-    one common-random-number path for JVP/VJP with respect to ``y_0`` and
-    ``p``. ``z_0`` is only a root guess and has zero tangent by contract.
-    ``failure_ad_reference=(y, z, t, p)`` may provide a domain-safe point for
-    already-inactive ``vmap`` lanes and model aux/field evaluation. A newly
-    attempted root still requires its actual ``(y, z_guess, t, p)`` to be
-    JVP-safe; the reference does not replace an active root after it fails.
-    A nonfinite inexact algebraic-aux leaf at initialization prevents all
-    stochastic time-step work. Saved aux is checked at every node in steps
-    mode; endpoint mode checks it only after integration and retains the
-    endpoint state with zero aux if that check fails.
+    The system is ``dy = drift(y, z, t) dt + diffusion(y, z, t) dW`` with
+    ``0 = g(y, z, t)``. ``solver`` is ``EulerMaruyama`` or ``SRA1``, applied
+    to the reduced SDE obtained from the locally unique root ``z = Z(y, t)``:
+    the differential state advances on a fixed uniform grid of ``n_steps``
+    steps, and a root solve restores consistency at every node and at SRA1's
+    drift stage. SRA1's strong order 1.5 requires a diffusion that depends
+    only on time. A fixed ``key`` defines one common-random-numbers path;
+    JVP/VJP with respect to ``y_0`` and ``p`` are pathwise, and ``z_0`` is a
+    root guess with zero tangent. Aux contracts, ``failure_ad_reference``,
+    and failure behavior follow ``solve_semi_explicit_dae``.
     """
     if not isinstance(n_steps, int) or isinstance(n_steps, bool):
         raise TypeError("n_steps must be a static Python int")
     if n_steps < 1:
         raise ValueError("n_steps must be at least 1")
-    if not isinstance(solver, EulerMaruyama):
-        raise TypeError("semi-explicit SDAEs currently support EulerMaruyama")
+    if not isinstance(solver, (EulerMaruyama, SRA1)):
+        raise TypeError("semi-explicit SDAEs support EulerMaruyama and SRA1")
     if save_at is None:
         save_at = SaveAt(t_1=True)
     if save_at.ts is not None:
@@ -112,7 +93,8 @@ def solve_semi_explicit_sdae(
     )
     dt = (t_1 - t_0) / n_steps
     time_grid = jnp.linspace(t_0, t_1, n_steps + 1)
-    d_w = _diagonal_brownian_increments(y_0, key, n_steps, dt, time_dtype)
+    noise = solver.sample_noise(y_0, key, n_steps, dt, time_dtype)
+    is_sra1 = isinstance(solver, SRA1)
 
     raw_drift = drift
     raw_diffusion = diffusion
@@ -243,7 +225,7 @@ def attempt_step(carry, inputs):
             num_root_solves,
             num_root_steps,
         ) = carry
-        t_step, t_next, d_w_step = inputs
+        t_step, t_next, noise_step = inputs
         active = ~failed
         y_ref, z_ref, t_ref, p_ref = failure_ad_reference
         y_eval = where(active, y, y_ref)
@@ -261,14 +243,79 @@ def attempt_step(carry, inputs):
             diffusion_output(y_eval, z_eval, t_eval, p_eval, context_eval),
             "diffusion(y, z, t)",
         )
-        y_candidate = add_scaled(
-            y,
-            (dt, drift_value),
-            (1.0, multiply(diffusion_value, d_w_step)),
-        )
-        z_candidate, root_ok, root_solves, root_steps = solve_root(
-            y_candidate, t_next, z, active
-        )
+        if is_sra1:
+            d_w_step, d_z_step = noise_step
+            chi = jax.tree.map(
+                lambda w, v: 0.5 * (w + v * INV_SQRT_3), d_w_step, d_z_step
+            )
+            # Additive-noise contract: the diffusion may depend only on time,
+            # so its endpoint evaluation reuses the node's (y, z, context).
+            t_next_eval = jnp.where(active, t_next, t_ref)
+            diffusion_next = checked_value(
+                diffusion_output(y_eval, z_eval, t_next_eval, p_eval, context_eval),
+                "diffusion(y, z, t)",
+            )
+            t_stage = t_step + 0.75 * dt
+            y_stage = add_scaled(
+                y,
+                (0.75 * dt, drift_value),
+                (1.5, multiply(diffusion_next, chi)),
+            )
+            z_stage, stage_root_ok, stage_solves, stage_steps = solve_root(
+                y_stage, t_stage, z, active
+            )
+            if has_algebraic_aux:
+                context_stage, stage_context_ok = evaluate_context(
+                    y_stage, z_stage, t_stage, stage_root_ok
+                )
+                stage_ok = stage_root_ok & stage_context_ok
+            else:
+                context_stage = None
+                stage_ok = stage_root_ok
+            y_stage_eval = where(stage_ok, y_stage, y_ref)
+            z_stage_eval = where(stage_ok, z_stage, z_ref)
+            t_stage_eval = jnp.where(stage_ok, t_stage, t_ref)
+            context_stage_eval = (
+                where(stage_ok, context_stage, context_reference)
+                if has_algebraic_aux
+                else None
+            )
+            stage_drift_raw, _ = split_field_output(
+                drift_output(
+                    y_stage_eval,
+                    z_stage_eval,
+                    t_stage_eval,
+                    p_eval,
+                    context_stage_eval,
+                ),
+                has_aux,
+            )
+            stage_drift_value = checked_value(stage_drift_raw, "drift(y, z, t)")
+            y_candidate = add_scaled(
+                y,
+                (dt / 3.0, drift_value),
+                (2.0 * dt / 3.0, stage_drift_value),
+                (1.0, multiply(diffusion_next, d_w_step)),
+                (
+                    1.0,
+                    multiply(add_scaled(diffusion_value, (-1.0, diffusion_next)), chi),
+                ),
+            )
+            z_candidate, endpoint_ok, endpoint_solves, endpoint_steps = solve_root(
+                y_candidate, t_next, z_stage, stage_ok
+            )
+            root_ok = stage_ok & endpoint_ok
+            root_solves = stage_solves + endpoint_solves
+            root_steps = stage_steps + endpoint_steps
+        else:
+            y_candidate = add_scaled(
+                y,
+                (dt, drift_value),
+                (1.0, multiply(diffusion_value, noise_step)),
+            )
+            z_candidate, root_ok, root_solves, root_steps = solve_root(
+                y_candidate, t_next, z, active
+            )
         if has_algebraic_aux:
             context_candidate, context_ok = evaluate_context(
                 y_candidate, z_candidate, t_next, root_ok & active
@@ -373,7 +420,7 @@ def body(carry, inputs):
     ) = jax.lax.scan(
         body,
         carry_0,
-        (time_grid[:-1], time_grid[1:], d_w),
+        (time_grid[:-1], time_grid[1:], noise),
     )
     ok = ~failed & (num_accepted == n_steps)
     if save_at.t_1:
diff --git a/src/tinydiffeq/sde.py b/src/tinydiffeq/sde.py
index 781be75..a8224cd 100644
--- a/src/tinydiffeq/sde.py
+++ b/src/tinydiffeq/sde.py
@@ -20,27 +20,6 @@
 from tinydiffeq.solution import Solution
 
 
-def _diagonal_brownian_increments(x_0, key, n_steps, dt, dtype):
-    """Generate one diagonal-noise draw for an array or pytree state."""
-    leaves, treedef = jax.tree.flatten(x_0)
-    if treedef == jax.tree.structure(0):
-        return jnp.sqrt(dt) * jax.random.normal(
-            key, (n_steps,) + x_0.shape, dtype=dtype
-        )
-    sizes = [leaf.size for leaf in leaves]
-    flat_noise = jnp.sqrt(dt) * jax.random.normal(
-        key, (n_steps, sum(sizes)), dtype=dtype
-    )
-    noise_leaves = []
-    start = 0
-    for leaf, size in zip(leaves, sizes, strict=True):
-        noise_leaves.append(
-            flat_noise[:, start : start + size].reshape((n_steps,) + leaf.shape)
-        )
-        start += size
-    return jax.tree.unflatten(treedef, noise_leaves)
-
-
 def solve_sde(
     drift,
     diffusion,
@@ -49,43 +28,42 @@ def solve_sde(
     t_1,
     x_0,
     *,
-    key,
+    key=None,
     n_steps,
+    noise=None,
     p=None,
     args=None,
     save_at=None,
     project=None,
     has_aux=None,
     failure_ad_reference=None,
+    unroll=1,
 ):
-    """Integrate the Ito SDE ``dx = drift dt + diffusion d_w`` (diagonal noise)
-    on the fixed grid of ``n_steps`` uniform steps from ``t_0`` to ``t_1 > t_0``.
+    """Integrate the Ito SDE ``dx = drift dt + diffusion d_w`` with diagonal
+    noise on a fixed grid of ``n_steps`` uniform steps from ``t_0`` to
+    ``t_1 > t_0``.
 
     ``drift`` and ``diffusion`` follow the same signature convention as
-    ``solve_ode`` — ``(x)``, ``(x, t)``, ``(x, t, args)``, or
-    ``(x, t, args, p)``. ``n_steps`` must be a static Python int (the honest
-    static-shape contract; there is currently no adaptive SDE stepping). The
-    Brownian increments are presampled from ``key``. Arrays retain the exact
-    ``(n_steps,) + x_0.shape`` random draw. Pytree states use one shared flat
-    draw, partitioned into leaves in JAX's deterministic pytree leaf order.
-    Thus a fixed key gives a fixed noise process: reproducible across calls
-    and differentiable with respect to ``x_0`` and ``p`` (not ``key``).
-
-    ``SaveAt(ts=...)`` raises — cubic Hermite interpolation is wrong for
-    rough paths; use ``t_1`` (default) or ``steps`` (here ``n_steps + 1``
-    rows unless a saved-aux failure terminates the accepted prefix).
-
-    ``drift`` may return either its value or ``(value, aux)``. The optional
-    real-floating aux pytree is stored at the same fixed nodes as ``xs`` and
-    is differentiated pathwise under the fixed random key. ``diffusion`` is
-    value-only. ``has_aux=None`` auto-detects the drift form;
-    ``has_aux=False`` avoids the abstract detection trace and selects the
-    original no-aux scan.
+    ``solve_ode``. ``solver`` is ``EulerMaruyama``, ``Milstein``, or ``SRA1``,
+    each declaring its per-step noise through
+    ``solver.sample_noise(x_0, key, n_steps, dt, dtype)``. Exactly one of
+    ``key`` and ``noise`` must be provided: a fixed ``key`` presamples a
+    fixed, reproducible noise process, differentiable with respect to ``x_0``
+    and ``p``; an explicit ``noise`` pytree (validated against the solver's
+    spec) is additionally differentiable as data. ``SaveAt(ts=...)`` raises —
+    interpolation is wrong for rough paths. ``drift`` may return
+    ``(value, aux)``; ``diffusion`` is value-only. ``unroll`` (a static int)
+    unrolls that many steps per iteration of the underlying ``lax.scan`` —
+    identical values, fewer/larger GPU dispatches, more compile time.
     """
     if not isinstance(n_steps, int):
         raise TypeError("n_steps must be a static Python int")
     if n_steps < 1:
         raise ValueError("n_steps must be at least 1")
+    if not isinstance(unroll, int) or isinstance(unroll, bool) or unroll < 1:
+        raise ValueError("unroll must be a static int of at least 1")
+    if (key is None) == (noise is None):
+        raise ValueError("solve_sde requires exactly one of key or noise")
     if save_at is None:
         save_at = SaveAt(t_1=True)
     if save_at.ts is not None:
@@ -102,7 +80,36 @@ def solve_sde(
     t_0 = jnp.asarray(t_0, time_dtype)
     t_1 = jnp.asarray(t_1, time_dtype)
     dt = (t_1 - t_0) / n_steps
-    d_w = _diagonal_brownian_increments(x_0, key, n_steps, dt, time_dtype)
+    if noise is None:
+        noise = solver.sample_noise(x_0, key, n_steps, dt, time_dtype)
+    else:
+        # sample_noise only reads shapes from x_0, so the abstract trace is
+        # the solver's authoritative noise spec.
+        reference = jax.eval_shape(
+            lambda noise_key: solver.sample_noise(
+                x_0, noise_key, n_steps, dt, time_dtype
+            ),
+            jax.random.key(0),
+        )
+        noise = jax.tree.map(jnp.asarray, noise)
+        if jax.tree.structure(noise) != jax.tree.structure(reference):
+            raise ValueError(
+                "noise must match the pytree structure of "
+                "solver.sample_noise(x_0, key, n_steps, dt, dtype)"
+            )
+        for leaf, ref in zip(
+            jax.tree.leaves(noise), jax.tree.leaves(reference), strict=True
+        ):
+            if leaf.shape != ref.shape:
+                raise ValueError(
+                    f"noise leaf shape {leaf.shape} does not match the "
+                    f"solver's expected {ref.shape}"
+                )
+            if leaf.dtype != ref.dtype:
+                raise TypeError(
+                    f"noise leaf dtype {leaf.dtype} must match the state "
+                    f"dtype {ref.dtype}"
+                )
     time_grid = jnp.linspace(t_0, t_1, n_steps + 1)
 
     def project_state(x):
@@ -142,12 +149,14 @@ def g_diffusion(x, t):
         return value
 
     def body(x, inputs):
-        t, d_w_step = inputs
-        x_1 = solver.step(g_drift, g_diffusion, t, x, dt, d_w_step, project_state)
+        t, noise_step = inputs
+        x_1 = solver.step(g_drift, g_diffusion, t, x, dt, noise_step, project_state)
         return x_1, x_1 if save_at.steps else None
 
     if save_at.t_1 or not has_aux:
-        x_final, step_states = jax.lax.scan(body, x_0, (time_grid[:-1], d_w))
+        x_final, step_states = jax.lax.scan(
+            body, x_0, (time_grid[:-1], noise), unroll=unroll
+        )
         num_accepted = jnp.asarray(n_steps, jnp.int32)
         num_steps = num_accepted
         ok = jnp.asarray(True)
@@ -194,14 +203,14 @@ def auxiliary(inputs):
 
         def aux_attempt(carry, inputs):
             x, aux, t, failed, count, num_steps = carry
-            t_step, t_next, d_w_step = inputs
+            t_step, t_next, noise_step = inputs
             x_candidate = solver.step(
                 g_drift,
                 g_diffusion,
                 t_step,
                 x,
                 dt,
-                d_w_step,
+                noise_step,
                 project_state,
             )
             aux_candidate, aux_ok = evaluate_aux(
@@ -255,7 +264,9 @@ def aux_body(carry, inputs):
                 num_steps,
             ),
             rows,
-        ) = jax.lax.scan(aux_body, carry_0, (time_grid[:-1], time_grid[1:], d_w))
+        ) = jax.lax.scan(
+            aux_body, carry_0, (time_grid[:-1], time_grid[1:], noise), unroll=unroll
+        )
         ts_s, xs_s, aux_s, advance_s = rows
         all_times = jnp.concatenate([t_0[None], ts_s])
         all_states = prepend(x_0, xs_s)
diff --git a/src/tinydiffeq/solution.py b/src/tinydiffeq/solution.py
index 4b61c7d..7a7a6dd 100644
--- a/src/tinydiffeq/solution.py
+++ b/src/tinydiffeq/solution.py
@@ -9,24 +9,14 @@
 class Solution:
     """Result of ``solve_ode``/``solve_sde``/``solve_linear_ode``.
 
-    - ``ts``/``xs``: times and states in the shape dictated by ``SaveAt`` —
-      scalar/endpoint for ``t_1``; for multi-row modes, every state pytree leaf
-      receives a leading ``len(ts)`` or ``max_steps + 1`` axis.
-    - ``ok``: scalar bool, True iff the integration reached ``t_1`` within the
-      attempt budget and every required saved output was valid. The package
-      never poisons outputs; callers that want diverging residuals map
-      ``jnp.where(sol.ok, x, jnp.inf)`` over leaves.
-    - ``num_accepted``: number of accepted steps (excluding the initial
-      state).
-    - ``num_steps``: scalar integer array counting logical attempted steps,
-      including rejected attempts. Public solver outputs always populate it;
-      the optional default only preserves direct-construction compatibility.
-    - ``accepted``: ``steps`` mode only (otherwise None): validity mask for
-      the contiguous accepted-step prefix. Row 0 (the initial state) is
-      always True, so ``accepted.sum() == num_accepted + 1``.
-    - ``aux``: optional floating pytree returned by the differential field
-      (or stochastic drift), with the same leading saved-time axis as
-      ``xs``. It is ``None`` when the field has no auxiliary output.
+    ``ts``/``xs`` hold times and states in the shape dictated by ``SaveAt``.
+    ``ok`` is a scalar bool: the integration reached ``t_1`` and every
+    required saved output was valid. Outputs are never poisoned; callers that
+    want diverging values map ``jnp.where(sol.ok, x, jnp.inf)`` over leaves.
+    ``num_accepted`` counts accepted steps, ``num_steps`` counts logical
+    attempts including rejections, ``accepted`` masks the valid prefix in
+    ``steps`` mode (row 0 is always True), and ``aux`` holds the field's
+    saved auxiliary pytree with the same leading saved-time axis as ``xs``.
     """
 
     ts: jax.Array
@@ -41,26 +31,16 @@ class Solution:
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True)
 class DAESolution:
-    """Result of the deterministic or stochastic semi-explicit DAE solver.
-
-    ``ts`` and ``ys`` follow the same :class:`tinydiffeq.SaveAt` shape
-    contract as :class:`Solution`. ``zs`` holds algebraic states and ``aux``
-    holds the optional differential-field auxiliary-output pytree. An
-    algebraic auxiliary output is internal context passed into the
-    differential field; it is not stored directly. RK4/Tsit5
-    internal-step and endpoint values are evaluated at converged roots.
-    Rodas5P instead satisfies the algebraic equation to integration accuracy
-    after its initial consistency root. Requested values are dense
-    interpolants and need not satisfy the constraint exactly. ``ok`` is true
-    only when initialization, all required stage and saved-output operations
-    succeeded, and the integration reached ``t_1``. ``num_root_solves`` counts
-    logical active algebraic root calls, including failed calls and the initial
-    consistency solve. ``num_root_steps`` sums their nonlinear LM update steps.
-    ``num_steps`` counts logical attempted integration steps, including rejected
-    attempts, with the same semantics as :class:`Solution`.
-    All three are path diagnostics with exact-zero tangents. Under ``vmap``, masked
-    lanes may still execute physically even though they do not increment these
-    logical per-lane counters.
+    """Result of the deterministic or stochastic semi-explicit DAE solvers.
+
+    ``ts``/``ys``/``zs`` follow the :class:`Solution` shape contract, with
+    ``zs`` the algebraic states. Explicit-method saved values sit at
+    converged roots; Rodas5P satisfies the constraint to integration accuracy
+    after its initial consistency root, and requested-grid interpolants need
+    not satisfy it exactly. ``num_root_solves`` counts logical active root
+    calls (including failures and the initial consistency solve) and
+    ``num_root_steps`` sums their LM update steps; like ``num_steps``, both
+    are path diagnostics with exact-zero tangents.
     """
 
     ts: jax.Array
diff --git a/src/tinydiffeq/solvers.py b/src/tinydiffeq/solvers.py
index dd83f8a..f1cc9f7 100644
--- a/src/tinydiffeq/solvers.py
+++ b/src/tinydiffeq/solvers.py
@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 
 import jax
+import jax.numpy as jnp
 
 from tinydiffeq._tree import add_scaled, multiply, weighted_sum
 
@@ -185,14 +186,13 @@ def step_fixed(self, g, t, x, dt, f_0, project):
 class Rodas5P:
     """Fifth-order Rodas5P Rosenbrock--Wanner method.
 
-    Rodas5P is an eight-stage, linearly implicit method with an embedded
-    error estimate and a stiff-aware fourth-order continuous extension. It is
-    supported by :func:`tinydiffeq.solve_ode` and
-    :func:`tinydiffeq.solve_semi_explicit_dae`; both use one dense LU
-    factorization per attempted step and reuse it across all stages.
-
-    The implementation follows Steinebach (2023) and SciML's
-    ``OrdinaryDiffEqRosenbrock.Rodas5P`` implementation:
+    An eight-stage, linearly implicit method with an embedded error estimate
+    and a stiff-aware fourth-order continuous extension, supported by
+    :func:`tinydiffeq.solve_ode` and
+    :func:`tinydiffeq.solve_semi_explicit_dae` with one dense LU
+    factorization reused across the stages of each attempted step. The
+    implementation follows Steinebach (2023) and SciML's
+    ``OrdinaryDiffEqRosenbrock.Rodas5P``:
 
     - https://doi.org/10.1007/s10543-023-00967-x
     - https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock
@@ -203,14 +203,137 @@ class Rodas5P:
     has_error_estimate = True
 
 
+def diagonal_brownian_increments(x_0, key, n_steps, dt, dtype):
+    """Draw ``n_steps`` diagonal Brownian increments ``sqrt(dt) * N(0, 1)``.
+
+    Arrays retain the exact ``(n_steps,) + x_0.shape`` draw. Pytree states use
+    one shared flat draw, partitioned into leaves in JAX's deterministic
+    pytree leaf order.
+    """
+    leaves, treedef = jax.tree.flatten(x_0)
+    if treedef == jax.tree.structure(0):
+        return jnp.sqrt(dt) * jax.random.normal(
+            key, (n_steps,) + x_0.shape, dtype=dtype
+        )
+    sizes = [leaf.size for leaf in leaves]
+    flat_noise = jnp.sqrt(dt) * jax.random.normal(
+        key, (n_steps, sum(sizes)), dtype=dtype
+    )
+    noise_leaves = []
+    start = 0
+    for leaf, size in zip(leaves, sizes, strict=True):
+        noise_leaves.append(
+            flat_noise[:, start : start + size].reshape((n_steps,) + leaf.shape)
+        )
+        start += size
+    return jax.tree.unflatten(treedef, noise_leaves)
+
+
+# SDE steppers share the contract
+# `step(g_drift, g_diffusion, t, x, dt, noise, project) -> x_1` where `noise`
+# is one per-step slice of the pytree produced by the solver's own
+# `sample_noise(x_0, key, n_steps, dt, dtype)`, so explicit noise handed to
+# `solve_sde` is validated against exactly what the solver expects.
+
+
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True)
 class EulerMaruyama:
-    """Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only."""
+    """Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only.
+
+    Strong order 0.5 for multiplicative noise. ``sample_noise`` returns the
+    Brownian increments with the same pytree structure as the state and a
+    leading ``n_steps`` axis.
+    """
+
+    order = 1
+    strong_order = 0.5
+
+    def sample_noise(self, x_0, key, n_steps, dt, dtype):
+        return diagonal_brownian_increments(x_0, key, n_steps, dt, dtype)
+
+    def step(self, g_drift, g_diffusion, t, x, dt, noise, project):
+        return project(
+            add_scaled(
+                x, (dt, g_drift(x, t)), (1.0, multiply(g_diffusion(x, t), noise))
+            )
+        )
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class Milstein:
+    """Milstein for Ito SDEs with diagonal noise. Fixed-step only.
+
+    Strong order 1.0 under the diagonal commutativity condition: each
+    diffusion component may depend only on its own state component. The
+    correction ``(1/2) g g' (d_w^2 - dt)`` evaluates ``g g'`` as the
+    forward-mode derivative of the diffusion field in the direction of its
+    own value, which equals the diagonal term exactly in that case.
+    ``sample_noise`` matches ``EulerMaruyama``.
+    """
 
     order = 1
+    strong_order = 1.0
+
+    def sample_noise(self, x_0, key, n_steps, dt, dtype):
+        return diagonal_brownian_increments(x_0, key, n_steps, dt, dtype)
 
-    def step(self, g_drift, g_diffusion, t, x, dt, d_w, project):
+    def step(self, g_drift, g_diffusion, t, x, dt, noise, project):
+        g_value = g_diffusion(x, t)
+        _, dg_g = jax.jvp(lambda state: g_diffusion(state, t), (x,), (g_value,))
+        correction = jax.tree.map(lambda dg, w: 0.5 * dg * (w * w - dt), dg_g, noise)
         return project(
-            add_scaled(x, (dt, g_drift(x, t)), (1.0, multiply(g_diffusion(x, t), d_w)))
+            add_scaled(
+                x,
+                (dt, g_drift(x, t)),
+                (1.0, multiply(g_value, noise)),
+                (1.0, correction),
+            )
+        )
+
+
+INV_SQRT_3 = 3.0**-0.5
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class SRA1:
+    """Rossler SRA1 stochastic Runge-Kutta for Ito SDEs with additive
+    diagonal noise. Fixed-step only.
+
+    Strong order 1.5 when the diffusion is independent of the state (it may
+    depend on time). ``sample_noise`` returns ``(d_w, d_z)``: two independent
+    ``sqrt(dt) * N(0, 1)`` draws per step. The time-Wiener integral
+    ``I_10 / dt`` is realized internally as ``(d_w + d_z / sqrt(3)) / 2``,
+    reproducing its variance ``dt^3 / 3`` and covariance ``dt^2 / 2`` with
+    the increment.
+    """
+
+    order = 2
+    strong_order = 1.5
+
+    def sample_noise(self, x_0, key, n_steps, dt, dtype):
+        key_w, key_z = jax.random.split(key)
+        return (
+            diagonal_brownian_increments(x_0, key_w, n_steps, dt, dtype),
+            diagonal_brownian_increments(x_0, key_z, n_steps, dt, dtype),
+        )
+
+    def step(self, g_drift, g_diffusion, t, x, dt, noise, project):
+        d_w, d_z = noise
+        g_0 = g_diffusion(x, t)
+        g_1 = g_diffusion(x, t + dt)
+        chi = jax.tree.map(lambda w, z: 0.5 * (w + z * INV_SQRT_3), d_w, d_z)
+        k_1 = g_drift(x, t)
+        stage = add_scaled(x, (0.75 * dt, k_1), (1.5, multiply(g_1, chi)))
+        k_2 = g_drift(stage, t + 0.75 * dt)
+        return project(
+            add_scaled(
+                x,
+                (dt / 3.0, k_1),
+                (2.0 * dt / 3.0, k_2),
+                (1.0, multiply(g_1, d_w)),
+                (1.0, multiply(add_scaled(g_0, (-1.0, g_1)), chi)),
+            )
         )
diff --git a/tests/test_aux.py b/tests/test_aux.py
index aee7ef5..6e1edeb 100644
--- a/tests/test_aux.py
+++ b/tests/test_aux.py
@@ -308,7 +308,7 @@ def differential(y, z, t, args, p, algebraic_aux):
             max_steps=1,
             save_at=save_at,
             has_aux=True,
-            root_solver=LMRootSolver(max_steps_is_success=False),
+            root_solver=LMRootSolver(),
             failure_ad_reference=(1.0, 1.0, 0.0, 0.0),
         )
 
diff --git a/tests/test_dae.py b/tests/test_dae.py
index 8337797..21ced4b 100644
--- a/tests/test_dae.py
+++ b/tests/test_dae.py
@@ -46,7 +46,6 @@ def constraint(y, z, t, args, p):
         return z - y
 
     defaults = _build_algebraic_solver(constraint, LMRootSolver(), False)
-    assert not LMRootSolver().max_steps_is_success
     # Everything algorithmic is nlls-gram's default; only the two invariants
     # this package owns are pinned.
     assert isinstance(defaults.linear_solver, Cholesky)
@@ -109,7 +108,7 @@ def test_lm_root_solver_options_normalize_and_reject_fixed_keys():
 
 
 def test_max_steps_policy_requires_root_residual_and_strict_batched_derivative():
-    def one_lane(z_0, p, max_steps_is_success):
+    def one_lane(z_0, p):
         return solve_semi_explicit_dae(
             lambda y, z: jnp.zeros_like(y),
             lambda y, z, t, args, p: z - p,
@@ -121,26 +120,20 @@ def one_lane(z_0, p, max_steps_is_success):
             p=p,
             dt_0=0.1,
             max_steps=1,
-            root_solver=LMRootSolver(
-                max_steps=1,
-                max_steps_is_success=max_steps_is_success,
-                atol=1e-12,
-            ),
+            root_solver=LMRootSolver(max_steps=1, atol=1e-12),
         )
 
     p = jnp.asarray(1.0)
-    forgiving = one_lane(jnp.asarray(-1e6), p, True)
-    strict = one_lane(jnp.asarray(0.0), p, False)
-    assert not bool(forgiving.ok)
+    starved = one_lane(jnp.asarray(-1e6), p)
+    strict = one_lane(jnp.asarray(0.0), p)
+    assert not bool(starved.ok)
     assert not bool(strict.ok)
-    assert int(forgiving.num_root_solves) == 1
-    assert int(forgiving.num_root_steps) == 1
-    assert forgiving.zs == jnp.asarray(-1e6)
+    assert int(starved.num_root_solves) == 1
+    assert int(starved.num_root_steps) == 1
+    assert starved.zs == jnp.asarray(-1e6)
 
     def strict_batch(parameter):
-        return jax.vmap(lambda z_0: one_lane(z_0, parameter, False))(
-            jnp.asarray([1.0, 0.0])
-        )
+        return jax.vmap(lambda z_0: one_lane(z_0, parameter))(jnp.asarray([1.0, 0.0]))
 
     def endpoint(parameter):
         result = strict_batch(parameter)
@@ -162,7 +155,7 @@ def test_dae_roots_reject_nonresidual_stopping_rules():
         LMRootSolver(xtol=1e-6)
 
 
-def test_compat_max_steps_success_does_not_broaden_nlls_root_ad():
+def test_root_budget_exhaustion_is_never_a_differentiable_root():
     guesses = jnp.asarray([1.0, 0.0])
 
     def solve_one(parameter, guess):
@@ -177,11 +170,7 @@ def solve_one(parameter, guess):
             p=parameter,
             dt_0=0.1,
             max_steps=1,
-            root_solver=LMRootSolver(
-                max_steps=1,
-                max_steps_is_success=True,
-                atol=1e-12,
-            ),
+            root_solver=LMRootSolver(max_steps=1, atol=1e-12),
         )
 
     def endpoints(parameters):
@@ -803,7 +792,7 @@ def test_initial_root_failure_and_time_budget_failure():
         jnp.asarray(0.0),
         dt_0=0.1,
         max_steps=10,
-        root_solver=LMRootSolver(max_steps_is_success=False),
+        root_solver=LMRootSolver(),
     )
     assert not bool(failed_root.ok)
     assert int(failed_root.num_accepted) == 0
@@ -842,7 +831,6 @@ def test_adaptive_stage_root_failure_retries_with_smaller_step(predictor):
         controller=IController(),
         root_solver=LMRootSolver(
             max_steps=1,
-            max_steps_is_success=False,
             atol=1e-8,
             predictor=predictor,
         ),
@@ -872,7 +860,7 @@ def one_lane(y, z, p):
             p=p,
             dt_0=0.1,
             max_steps=1,
-            root_solver=LMRootSolver(max_steps_is_success=False),
+            root_solver=LMRootSolver(),
             failure_ad_reference=(1.0, 1.0, 0.0, 0.0),
         )
 
@@ -949,8 +937,6 @@ def test_validation():
         LMRootSolver(max_steps=0)
     with pytest.raises(ValueError, match="atol must be positive or None"):
         LMRootSolver(atol=0.0)
-    with pytest.raises(TypeError, match="max_steps_is_success must be a bool"):
-        LMRootSolver(max_steps_is_success=1)
     with pytest.raises(ValueError, match="gtol must be zero"):
         LMRootSolver(gtol=-1.0)
     with pytest.raises(ValueError, match="xtol must be zero"):
diff --git a/tests/test_float64_subprocess.py b/tests/test_float64_subprocess.py
index 186cdac..1006a8d 100644
--- a/tests/test_float64_subprocess.py
+++ b/tests/test_float64_subprocess.py
@@ -431,6 +431,60 @@ def continuous_forecast(value):
 """)
 
 
+def test_sra1_with_x64_disabled_stays_float32_and_differentiable():
+    run_script(r"""
+import jax
+import jax.numpy as jnp
+
+from tinydiffeq import SRA1, solve_sde
+
+
+assert not jax.config.x64_enabled
+n = 32
+x_0 = jnp.asarray(1.0)
+key = jax.random.key(0)
+noise = SRA1().sample_noise(x_0, key, n, jnp.asarray(1.0 / n), x_0.dtype)
+
+
+def endpoint(x_0, noise):
+    return solve_sde(
+        lambda x: -0.5 * x,
+        lambda x: 0.1 * jnp.ones_like(x),
+        SRA1(),
+        0.0,
+        1.0,
+        x_0,
+        noise=noise,
+        n_steps=n,
+    ).xs
+
+
+jaxpr = jax.make_jaxpr(endpoint)(x_0, noise)
+assert "f64" not in str(jaxpr), jaxpr
+
+value = jax.jit(endpoint)(x_0, noise)
+keyed = solve_sde(
+    lambda x: -0.5 * x,
+    lambda x: 0.1 * jnp.ones_like(x),
+    SRA1(),
+    0.0,
+    1.0,
+    x_0,
+    key=key,
+    n_steps=n,
+).xs
+assert value.dtype == jnp.float32
+assert jnp.array_equal(value, keyed)
+
+grad_x, grad_noise = jax.grad(endpoint, argnums=(0, 1))(x_0, noise)
+assert grad_x.dtype == jnp.float32
+assert jnp.isfinite(grad_x)
+for leaf in jax.tree.leaves(grad_noise):
+    assert leaf.dtype == jnp.float32
+    assert jnp.all(jnp.isfinite(leaf))
+""")
+
+
 def test_pytree_states_preserve_float32_and_float64():
     run_script(r"""
 import jax
diff --git a/tests/test_gpu.py b/tests/test_gpu.py
index e8a36f0..74a55f1 100644
--- a/tests/test_gpu.py
+++ b/tests/test_gpu.py
@@ -3,15 +3,19 @@
 import pytest
 
 from tinydiffeq import (
+    RK4,
+    SRA1,
     AdaptiveKrylovExponential,
     AssociativeMarkov,
     ContinuousTimeMarkovChain,
     DenseExponential,
     DiscreteMarkovChain,
+    Euler,
     EulerMaruyama,
     IController,
     KrylovExponential,
     MatrixFreeContinuousTimeMarkovChain,
+    Milstein,
     Rodas5P,
     SaveAt,
     Tsit5,
@@ -625,6 +629,192 @@ def endpoint(p):
         assert bool(jnp.isfinite(leaf))
 
 
+def test_float32_fixed_euler_rk4_ode_on_gpu():
+    gpu = gpu_devices()[0]
+    n = 128
+
+    @jax.jit
+    def run(x_0):
+        euler = solve_ode(
+            lambda x: -F32_DECAY * x, Euler(), 0.0, 1.0, x_0, dt_0=1.0 / n, max_steps=n
+        ).xs
+        rk4 = solve_ode(
+            lambda x: -F32_DECAY * x, RK4(), 0.0, 1.0, x_0, dt_0=1.0 / n, max_steps=n
+        ).xs
+        return euler, rk4
+
+    with jax.default_device(gpu):
+        euler, rk4 = jax.block_until_ready(run(jnp.asarray(1.0, jnp.float32)))
+
+    exact = float(jnp.exp(-F32_DECAY))
+    for out, tol in ((euler, 1e-3), (rk4, 1e-5)):
+        assert out.dtype == jnp.float32
+        assert out.devices().pop().platform == "gpu"
+        assert abs(float(out) - exact) < tol
+
+
+@pytest.mark.parametrize(
+    ("solver", "diffusion"),
+    [
+        (Milstein(), lambda x: 0.1 * x),
+        (SRA1(), lambda x: 0.1 * jnp.ones_like(x)),
+    ],
+    ids=["milstein", "sra1"],
+)
+def test_float32_milstein_and_sra1_run_on_gpu(solver, diffusion):
+    gpu = gpu_devices()[0]
+
+    @jax.jit
+    def run(x_0):
+        return solve_sde(
+            lambda x: -F32_DECAY * x,
+            diffusion,
+            solver,
+            0.0,
+            1.0,
+            x_0,
+            key=jax.random.key(0),
+            n_steps=64,
+        ).xs
+
+    with jax.default_device(gpu):
+        out = jax.block_until_ready(run(jnp.asarray(1.0, jnp.float32)))
+
+    assert out.dtype == jnp.float32
+    assert out.devices().pop().platform == "gpu"
+    assert bool(jnp.isfinite(out))
+    assert abs(float(out)) < 10.0
+
+
+# The kernels use case: vmap over B trajectories with per-trajectory x_0 and
+# explicit noise, jit the whole batch, and take reverse-mode gradients of a
+# scalar residual with respect to (x_0, p, noise). n_steps covers 31 and 127
+# so both a scan XLA may unroll and one it will not take the same paths.
+
+
+def sra1_batch_noise(batch, n_steps, dtype, seed):
+    keys = jax.random.split(jax.random.key(seed), batch)
+    dt = jnp.asarray(1.0 / n_steps, dtype)
+    return jax.vmap(
+        lambda k: SRA1().sample_noise(jnp.zeros((), dtype), k, n_steps, dt, dtype)
+    )(keys)
+
+
+@pytest.mark.parametrize("n_steps", [31, 127])
+def test_float32_vmapped_sra1_ensemble_matches_per_trajectory_on_gpu(n_steps):
+    gpu = gpu_devices()[0]
+    dtype = jnp.float32
+    batch = 8
+    x_0s = jnp.linspace(0.5, 2.0, batch, dtype=dtype)
+    noise = sra1_batch_noise(batch, n_steps, dtype, seed=3)
+
+    def one(x_0, w, z):
+        return solve_sde(
+            lambda x: -F32_DECAY * x,
+            lambda x: jnp.asarray(0.1, x.dtype) * jnp.ones_like(x),
+            SRA1(),
+            0.0,
+            1.0,
+            x_0,
+            noise=(w, z),
+            n_steps=n_steps,
+        ).xs
+
+    with jax.default_device(gpu):
+        batched = jax.block_until_ready(jax.jit(jax.vmap(one))(x_0s, *noise))
+        looped = jnp.stack(
+            [one(x_0s[i], noise[0][i], noise[1][i]) for i in range(batch)]
+        )
+
+    assert batched.dtype == dtype
+    assert batched.devices().pop().platform == "gpu"
+    assert jnp.allclose(batched, looped, atol=1e-5)
+
+
+@pytest.mark.parametrize("n_steps", [31, 127])
+def test_float32_grad_of_vmapped_sra1_ensemble_on_gpu(n_steps):
+    gpu = gpu_devices()[0]
+    dtype = jnp.float32
+    batch = 8
+    x_0s = jnp.linspace(0.5, 2.0, batch, dtype=dtype)
+    theta = jnp.asarray(0.7, dtype)
+    noise = sra1_batch_noise(batch, n_steps, dtype, seed=11)
+
+    def endpoint(x_0, p, w, z):
+        return solve_sde(
+            lambda x, t, args, p: -p * x,
+            lambda x, t, args, p: jnp.asarray(0.1, x.dtype) * jnp.ones_like(x),
+            SRA1(),
+            0.0,
+            1.0,
+            x_0,
+            p=p,
+            noise=(w, z),
+            n_steps=n_steps,
+        ).xs
+
+    def loss(x_0s, p, noise):
+        endpoints = jax.vmap(lambda x_0, w, z: endpoint(x_0, p, w, z))(x_0s, *noise)
+        return jnp.sum((endpoints - 1.0) ** 2)
+
+    with jax.default_device(gpu):
+        value, (grad_x, grad_p, grad_noise) = jax.block_until_ready(
+            jax.jit(jax.value_and_grad(loss, argnums=(0, 1, 2)))(x_0s, theta, noise)
+        )
+
+    for leaf in jax.tree.leaves((value, grad_x, grad_p, grad_noise)):
+        assert leaf.dtype == dtype
+        assert leaf.devices().pop().platform == "gpu"
+        assert bool(jnp.all(jnp.isfinite(leaf)))
+
+    def one_loss(x_0, w, z):
+        return (endpoint(x_0, theta, w, z) - 1.0) ** 2
+
+    per_trajectory = jnp.stack(
+        [jax.grad(one_loss)(x_0s[i], noise[0][i], noise[1][i]) for i in range(batch)]
+    )
+    assert jnp.allclose(grad_x, per_trajectory, atol=1e-4)
+
+
+@pytest.mark.parametrize("n_steps", [31, 127])
+def test_float32_grad_of_vmapped_rk4_ensemble_on_gpu(n_steps):
+    gpu = gpu_devices()[0]
+    dtype = jnp.float32
+    batch = 8
+    x_0s = jnp.linspace(0.5, 2.0, batch, dtype=dtype)
+    theta = jnp.asarray(0.7, dtype)
+
+    def endpoint(x_0, p):
+        return solve_ode(
+            lambda x, t, args, p: -p * x,
+            RK4(),
+            0.0,
+            1.0,
+            x_0,
+            p=p,
+            dt_0=1.0 / n_steps,
+            max_steps=n_steps,
+        ).xs
+
+    def loss(x_0s, p):
+        endpoints = jax.vmap(lambda x_0: endpoint(x_0, p))(x_0s)
+        return jnp.sum((endpoints - 1.0) ** 2)
+
+    with jax.default_device(gpu):
+        grad_x, grad_p = jax.block_until_ready(
+            jax.jit(jax.grad(loss, argnums=(0, 1)))(x_0s, theta)
+        )
+
+    decay = jnp.exp(-theta)
+    expected_x = 2.0 * (x_0s * decay - 1.0) * decay
+    expected_p = jnp.sum(-2.0 * (x_0s * decay - 1.0) * x_0s * decay)
+    for leaf in (grad_x, grad_p):
+        assert leaf.dtype == dtype
+        assert leaf.devices().pop().platform == "gpu"
+    assert jnp.allclose(grad_x, expected_x, atol=1e-3)
+    assert jnp.allclose(grad_p, expected_p, atol=1e-2)
+
+
 def test_float32_y_with_float64_z_keeps_the_tight_root_tolerance_on_gpu():
     # dae.py picks root_atol from z_dtype alone: 1e-10 when z has more than 32
     # bits, 1e-6 otherwise. A float32 y with a float64 z therefore gets the
diff --git a/tests/test_sdae.py b/tests/test_sdae.py
index ed822e9..e545e7d 100644
--- a/tests/test_sdae.py
+++ b/tests/test_sdae.py
@@ -4,8 +4,10 @@
 import pytest
 
 from tinydiffeq import (
+    SRA1,
     EulerMaruyama,
     LMRootSolver,
+    Milstein,
     SaveAt,
     solve_sde,
     solve_semi_explicit_sdae,
@@ -26,7 +28,7 @@ def constraint(y, z, t, args, p):
     return z - y, {"scaled": p["scale"] * z, "square": z**2}
 
 
-def sdae(key, n_steps, *, save_at=None, y_0=Y_0, p=None):
+def sdae(key, n_steps, *, save_at=None, y_0=Y_0, p=None, solver=None):
     if p is None:
         p = {
             "mu": jnp.asarray(MU),
@@ -37,7 +39,7 @@ def sdae(key, n_steps, *, save_at=None, y_0=Y_0, p=None):
         drift,
         diffusion,
         constraint,
-        EulerMaruyama(),
+        EulerMaruyama() if solver is None else solver,
         0.0,
         T,
         jnp.asarray(y_0),
@@ -120,6 +122,83 @@ def exact(key):
     assert 0.35 < slope < 0.7, (slope, errors)
 
 
+def additive_sigma(t, p):
+    # Time-dependent but state-independent, so SRA1's additive contract holds
+    # for the reduced SDE and its endpoint diffusion evaluation is exercised.
+    return p["sigma"] * (1.0 + 0.5 * t)
+
+
+def test_sra1_sdae_matches_reduced_sde_on_identical_noise_path():
+    key = jax.random.key(12)
+    n = 32
+    p = {"mu": jnp.asarray(MU), "sigma": jnp.asarray(SIGMA)}
+    full = solve_semi_explicit_sdae(
+        lambda y, z, t, args, p: p["mu"] * z,
+        lambda y, z, t, args, p: additive_sigma(t, p) * jnp.ones_like(y),
+        lambda y, z: z - y,
+        SRA1(),
+        0.0,
+        T,
+        jnp.asarray(Y_0),
+        jnp.asarray(0.7),
+        key=key,
+        n_steps=n,
+        p=p,
+        save_at=SaveAt(steps=True),
+    )
+    reduced = solve_sde(
+        lambda y, t, args, p: p["mu"] * y,
+        lambda y, t, args, p: additive_sigma(t, p) * jnp.ones_like(y),
+        SRA1(),
+        0.0,
+        T,
+        jnp.asarray(Y_0),
+        key=key,
+        n_steps=n,
+        p=p,
+        save_at=SaveAt(steps=True),
+    )
+    assert bool(full.ok)
+    assert int(full.num_steps) == n
+    # One initial consistency root, then a stage and an endpoint root per step.
+    assert int(full.num_root_solves) == 2 * n + 1
+    assert jnp.allclose(full.ys, reduced.xs, atol=1e-7, rtol=1e-7)
+    assert jnp.allclose(full.zs, full.ys, atol=1e-7)
+
+
+def test_sra1_sdae_aux_and_state_jvp_vjp_under_fixed_key():
+    key = jax.random.key(21)
+
+    def output(mu):
+        p = {"mu": mu, "sigma": jnp.asarray(SIGMA), "scale": mu}
+        sol = sdae(key, 64, p=p, save_at=SaveAt(steps=True), solver=SRA1())
+        return jnp.sum(sol.ys + sol.zs + sol.aux["scaled"])
+
+    mu = jnp.asarray(MU)
+    tangent = jax.jvp(output, (mu,), (jnp.ones_like(mu),))[1]
+    cotangent = jax.grad(output)(mu)
+    eps = 1e-6
+    finite_difference = (output(mu + eps) - output(mu - eps)) / (2 * eps)
+    assert jnp.abs(tangent - finite_difference) < 2e-5
+    assert jnp.abs(cotangent - finite_difference) < 2e-5
+
+
+def test_sdae_rejects_milstein():
+    with pytest.raises(TypeError, match="EulerMaruyama and SRA1"):
+        solve_semi_explicit_sdae(
+            lambda y, z: z,
+            lambda y, z: 0.0 * y,
+            lambda y, z: z - y,
+            Milstein(),
+            0.0,
+            1.0,
+            jnp.asarray(1.0),
+            jnp.asarray(1.0),
+            key=jax.random.key(0),
+            n_steps=4,
+        )
+
+
 def test_sdae_root_guess_has_zero_tangent_and_ts_raises():
     key = jax.random.key(1)
 
@@ -203,7 +282,7 @@ def stochastic(y, z, t, args, p, algebraic_aux):
         save_at=SaveAt(steps=True),
         has_aux=True,
         has_algebraic_aux=True,
-        root_solver=LMRootSolver(max_steps_is_success=False),
+        root_solver=LMRootSolver(),
     )
     assert not bool(sol.ok)
     assert int(sol.accepted.sum()) == int(sol.num_accepted) + 1
@@ -244,7 +323,7 @@ def zero_diffusion(y, z, t, args, p, algebraic_aux):
             save_at=save_at,
             has_aux=True,
             has_algebraic_aux=True,
-            root_solver=LMRootSolver(max_steps_is_success=False),
+            root_solver=LMRootSolver(),
             failure_ad_reference=(1.0, 1.0, 0.0, 0.0),
         )
 
diff --git a/tests/test_sde.py b/tests/test_sde.py
index f32c455..6aa5af8 100644
--- a/tests/test_sde.py
+++ b/tests/test_sde.py
@@ -3,7 +3,15 @@
 import numpy as np
 import pytest
 
-from tinydiffeq import Euler, EulerMaruyama, SaveAt, solve_ode, solve_sde
+from tinydiffeq import (
+    SRA1,
+    Euler,
+    EulerMaruyama,
+    Milstein,
+    SaveAt,
+    solve_ode,
+    solve_sde,
+)
 
 # Geometric Brownian motion dX = mu X dt + sigma X d_w has the exact solution
 # X_T = X_0 exp((mu - sigma^2/2) T + sigma W_T). solve_sde presamples its
@@ -202,3 +210,398 @@ def test_steps_mode_shapes_and_flags():
     assert int(sol.num_steps) == n
     assert bool(jnp.all(sol.accepted))
     assert sol.ts[0] == 0.0 and sol.ts[-1] == T
+
+
+# Additive-noise Ornstein-Uhlenbeck dX = -theta X dt + sigma d_w for the SRA1
+# tests: SRA1's strong order 1.5 holds only for state-independent diffusion.
+
+THETA, OU_SIGMA = 1.0, 0.5
+
+
+def ou_drift(x, t, args, p):
+    return -THETA * x
+
+
+def ou_diffusion(x, t, args, p):
+    return OU_SIGMA * jnp.ones_like(x)
+
+
+@pytest.mark.parametrize("solver", [EulerMaruyama(), Milstein(), SRA1()])
+@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64])
+def test_explicit_noise_matches_key(solver, dtype):
+    key = jax.random.PRNGKey(11)
+    n = 32
+    x_0 = jnp.asarray(X_0, dtype)
+    noise = solver.sample_noise(x_0, key, n, jnp.asarray(T / n, dtype), dtype)
+    a = solve_sde(
+        ou_drift,
+        ou_diffusion,
+        solver,
+        0.0,
+        T,
+        x_0,
+        key=key,
+        n_steps=n,
+        save_at=SaveAt(steps=True),
+    )
+    b = solve_sde(
+        ou_drift,
+        ou_diffusion,
+        solver,
+        0.0,
+        T,
+        x_0,
+        noise=noise,
+        n_steps=n,
+        save_at=SaveAt(steps=True),
+    )
+    assert a.xs.dtype == dtype
+    assert a.ts.dtype == dtype
+    assert jnp.array_equal(a.xs, b.xs)
+    assert jnp.array_equal(a.ts, b.ts)
+
+
+def test_key_noise_exclusivity_raises():
+    x_0 = jnp.asarray(X_0)
+    noise = EulerMaruyama().sample_noise(
+        x_0, jax.random.PRNGKey(0), 8, jnp.asarray(T / 8), x_0.dtype
+    )
+    with pytest.raises(ValueError, match="exactly one"):
+        solve_sde(
+            drift, diffusion, EulerMaruyama(), 0.0, T, x_0, n_steps=8, p=(MU, SIGMA)
+        )
+    with pytest.raises(ValueError, match="exactly one"):
+        solve_sde(
+            drift,
+            diffusion,
+            EulerMaruyama(),
+            0.0,
+            T,
+            x_0,
+            key=jax.random.PRNGKey(0),
+            noise=noise,
+            n_steps=8,
+            p=(MU, SIGMA),
+        )
+
+
+def test_explicit_noise_wrong_shape_raises():
+    with pytest.raises(ValueError, match="shape"):
+        solve_sde(
+            drift,
+            diffusion,
+            EulerMaruyama(),
+            0.0,
+            T,
+            jnp.asarray(X_0),
+            noise=jnp.zeros((7,)),
+            n_steps=8,
+            p=(MU, SIGMA),
+        )
+
+
+def test_explicit_noise_wrong_structure_raises():
+    # SRA1 expects the (d_w, d_z) pair, not a bare increment array.
+    with pytest.raises(ValueError, match="structure"):
+        solve_sde(
+            ou_drift,
+            ou_diffusion,
+            SRA1(),
+            0.0,
+            T,
+            jnp.asarray(X_0),
+            noise=jnp.zeros((8,)),
+            n_steps=8,
+        )
+
+
+def test_explicit_noise_wrong_dtype_raises():
+    with pytest.raises(TypeError, match="dtype"):
+        solve_sde(
+            drift,
+            diffusion,
+            EulerMaruyama(),
+            0.0,
+            T,
+            jnp.asarray(X_0),
+            noise=jnp.zeros((8,), dtype=jnp.float32),
+            n_steps=8,
+            p=(MU, SIGMA),
+        )
+
+
+def test_grad_wrt_explicit_noise_vs_finite_differences():
+    n = 16
+    x_0 = jnp.asarray(X_0)
+    noise = EulerMaruyama().sample_noise(
+        x_0, jax.random.PRNGKey(5), n, jnp.asarray(T / n), x_0.dtype
+    )
+
+    def endpoint(noise):
+        return solve_sde(
+            drift,
+            diffusion,
+            EulerMaruyama(),
+            0.0,
+            T,
+            x_0,
+            noise=noise,
+            n_steps=n,
+            p=(MU, SIGMA),
+        ).xs
+
+    grad = jax.grad(endpoint)(noise)
+    eps = 1e-6
+    for i in (0, n // 2, n - 1):
+        bump = jnp.zeros_like(noise).at[i].set(eps)
+        fd = (endpoint(noise + bump) - endpoint(noise - bump)) / (2 * eps)
+        assert jnp.abs(grad[i] - fd) < 1e-6, i
+
+
+def test_grad_wrt_sra1_noise_pair_vs_finite_differences():
+    n = 16
+    x_0 = jnp.asarray(X_0)
+    noise = SRA1().sample_noise(
+        x_0, jax.random.PRNGKey(6), n, jnp.asarray(T / n), x_0.dtype
+    )
+
+    def endpoint(noise):
+        return solve_sde(
+            ou_drift, ou_diffusion, SRA1(), 0.0, T, x_0, noise=noise, n_steps=n
+        ).xs
+
+    grad_w, grad_z = jax.grad(endpoint)(noise)
+    eps = 1e-6
+    for leaf, grad_leaf in ((0, grad_w), (1, grad_z)):
+        for i in (0, n - 1):
+            bump = jnp.zeros((n,)).at[i].set(eps)
+            if leaf == 0:
+                plus = (noise[0] + bump, noise[1])
+                minus = (noise[0] - bump, noise[1])
+            else:
+                plus = (noise[0], noise[1] + bump)
+                minus = (noise[0], noise[1] - bump)
+            fd = (endpoint(plus) - endpoint(minus)) / (2 * eps)
+            assert jnp.abs(grad_leaf[i] - fd) < 1e-6, (leaf, i)
+
+
+@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64])
+def test_vmap_explicit_noise_matches_loop(dtype):
+    n = 16
+    x_0s = jnp.asarray([0.5, 1.0, 2.0], dtype)
+    keys = jax.random.split(jax.random.PRNGKey(9), 3)
+    d_w, d_z = jax.vmap(
+        lambda k: SRA1().sample_noise(
+            jnp.asarray(X_0, dtype), k, n, jnp.asarray(T / n, dtype), dtype
+        )
+    )(keys)
+
+    def one(x_0, w, z):
+        return solve_sde(
+            ou_drift, ou_diffusion, SRA1(), 0.0, T, x_0, noise=(w, z), n_steps=n
+        ).xs
+
+    batched = jax.vmap(one)(x_0s, d_w, d_z)
+    looped = jnp.stack([one(x_0s[i], d_w[i], d_z[i]) for i in range(3)])
+    assert batched.dtype == dtype
+    assert jnp.max(jnp.abs(batched - looped)) < 200 * jnp.finfo(dtype).eps
+
+
+def test_unroll_matches_rolled_values_and_gradients():
+    key = jax.random.PRNGKey(13)
+    n = 32
+    x_0 = jnp.asarray(X_0)
+
+    def endpoint(x_0, unroll):
+        return solve_sde(
+            ou_drift,
+            ou_diffusion,
+            SRA1(),
+            0.0,
+            T,
+            x_0,
+            key=key,
+            n_steps=n,
+            unroll=unroll,
+        ).xs
+
+    rolled = solve_sde(
+        ou_drift,
+        ou_diffusion,
+        SRA1(),
+        0.0,
+        T,
+        x_0,
+        key=key,
+        n_steps=n,
+        save_at=SaveAt(steps=True),
+        unroll=1,
+    )
+    unrolled = solve_sde(
+        ou_drift,
+        ou_diffusion,
+        SRA1(),
+        0.0,
+        T,
+        x_0,
+        key=key,
+        n_steps=n,
+        save_at=SaveAt(steps=True),
+        unroll=4,
+    )
+    assert jnp.array_equal(rolled.xs, unrolled.xs)
+    grad_rolled = jax.grad(lambda x: endpoint(x, 1))(x_0)
+    grad_unrolled = jax.grad(lambda x: endpoint(x, 4))(x_0)
+    assert jnp.allclose(grad_rolled, grad_unrolled, rtol=1e-12, atol=1e-12)
+    with pytest.raises(ValueError, match="unroll"):
+        endpoint(x_0, 0)
+
+
+def test_milstein_strong_order_gbm():
+    keys = jax.random.split(jax.random.PRNGKey(2), 400)
+    levels = (16, 32, 64, 128)
+    errors = []
+    for n in levels:
+        endpoints = jax.vmap(
+            lambda k, n=n: (
+                solve_sde(
+                    drift,
+                    diffusion,
+                    Milstein(),
+                    0.0,
+                    T,
+                    jnp.asarray(X_0),
+                    key=k,
+                    n_steps=n,
+                    p=(MU, SIGMA),
+                ).xs
+            )
+        )(keys)
+        exact = jax.vmap(lambda k, n=n: exact_endpoint(k, n))(keys)
+        errors.append(float(jnp.mean(jnp.abs(endpoints - exact))))
+    slope = np.polyfit(np.log([T / n for n in levels]), np.log(errors), 1)[0]
+    assert 0.75 < slope < 1.25, (slope, errors)
+
+
+def coarsen_sra1_noise(d_w, d_z, dt):
+    # Pasting adjacent steps: dW_c = dW_1 + dW_2 and the time-Wiener integral
+    # I_c = I_1 + I_2 + dt dW_1, then dZ back from I = (dt/2)(dW + dZ/sqrt(3)).
+    i_10 = 0.5 * dt * (d_w + d_z / np.sqrt(3.0))
+    d_w_c = d_w[..., 0::2] + d_w[..., 1::2]
+    i_c = i_10[..., 0::2] + i_10[..., 1::2] + dt * d_w[..., 0::2]
+    d_z_c = np.sqrt(3.0) * (i_c / dt - d_w_c)
+    return d_w_c, d_z_c
+
+
+def test_sra1_strong_order_additive_nonlinear():
+    # No closed form conditions on (d_w, d_z) alone, so measure
+    # self-convergence against the finest grid on consistently coarsened
+    # noise; additive-noise EM on the same paths is strong order 1.0, so a
+    # slope well above 1 and the error ratio both separate the schemes. The
+    # observed slope may exceed the guaranteed 1.5 when the h^2 error terms
+    # dominate on the tested grids (exactly 2 for linear drift).
+    def cubic_drift(x):
+        return -x - x**3
+
+    n_fine = 256
+    keys = jax.random.split(jax.random.PRNGKey(4), 200)
+    x_0 = jnp.asarray(X_0)
+    d_w, d_z = jax.vmap(
+        lambda k: SRA1().sample_noise(
+            x_0, k, n_fine, jnp.asarray(T / n_fine), x_0.dtype
+        )
+    )(keys)
+
+    def sra1_endpoints(d_w, d_z, n):
+        return jax.vmap(
+            lambda w, z: (
+                solve_sde(
+                    cubic_drift,
+                    ou_diffusion,
+                    SRA1(),
+                    0.0,
+                    T,
+                    x_0,
+                    noise=(w, z),
+                    n_steps=n,
+                ).xs
+            )
+        )(d_w, d_z)
+
+    reference = sra1_endpoints(d_w, d_z, n_fine)
+    level_w, level_z, n, dt = d_w, d_z, n_fine, T / n_fine
+    errors = {}
+    em_error = None
+    while n > 16:
+        level_w, level_z = coarsen_sra1_noise(level_w, level_z, dt)
+        n, dt = n // 2, 2.0 * dt
+        errors[n] = float(
+            jnp.mean(jnp.abs(sra1_endpoints(level_w, level_z, n) - reference))
+        )
+        if n == 16:
+            em = jax.vmap(
+                lambda w: (
+                    solve_sde(
+                        cubic_drift,
+                        ou_diffusion,
+                        EulerMaruyama(),
+                        0.0,
+                        T,
+                        x_0,
+                        noise=w,
+                        n_steps=16,
+                    ).xs
+                )
+            )(level_w)
+            em_error = float(jnp.mean(jnp.abs(em - reference)))
+    ns = sorted(errors)
+    slope = np.polyfit(np.log([T / n for n in ns]), np.log([errors[n] for n in ns]), 1)[
+        0
+    ]
+    assert 1.3 < slope < 2.5, (slope, errors)
+    assert errors[16] < 0.2 * em_error, (errors[16], em_error)
+
+
+def test_zero_diffusion_sra1_is_ralston_rk2():
+    # With diffusion = 0 the tableau collapses to the deterministic two-stage
+    # Ralston scheme, exact per step on dx = mu x: growth 1 + mu h + (mu h)^2/2.
+    n = 16
+    h = T / n
+    sol = solve_sde(
+        lambda x: MU * x,
+        lambda x: 0.0 * x,
+        SRA1(),
+        0.0,
+        T,
+        jnp.asarray(X_0),
+        key=jax.random.PRNGKey(0),
+        n_steps=n,
+    )
+    growth = 1.0 + MU * h + 0.5 * (MU * h) ** 2
+    assert jnp.abs(sol.xs - X_0 * growth**n) < 1e-12
+
+
+def test_zero_diffusion_milstein_matches_euler_ode():
+    n = 16
+    sde = solve_sde(
+        lambda x: MU * x,
+        lambda x: 0.0 * x,
+        Milstein(),
+        0.0,
+        T,
+        jnp.asarray(X_0),
+        key=jax.random.PRNGKey(0),
+        n_steps=n,
+        save_at=SaveAt(steps=True),
+    )
+    ode = solve_ode(
+        lambda x: MU * x,
+        Euler(),
+        0.0,
+        T,
+        jnp.asarray(X_0),
+        dt_0=T / n,
+        max_steps=n,
+        save_at=SaveAt(steps=True),
+    )
+    assert jnp.max(jnp.abs(sde.xs - ode.xs)) < 1e-14
diff --git a/tests/test_solvers_fixed.py b/tests/test_solvers_fixed.py
index 27e0885..45805e8 100644
--- a/tests/test_solvers_fixed.py
+++ b/tests/test_solvers_fixed.py
@@ -2,18 +2,23 @@
 import jax.numpy as jnp
 import jax.scipy.linalg as jsp_linalg
 import numpy as np
+import pytest
 from kernels_reference import rk4_grid
 
-from tinydiffeq import RK4, Euler, SaveAt, solve_ode
+from tinydiffeq import RK4, Euler, IController, SaveAt, Tsit5, solve_ode
 
 
 def logistic_exact(x_0, t):
     return x_0 * jnp.exp(t) / (1.0 + x_0 * (jnp.exp(t) - 1.0))
 
 
-def test_linear_system_vs_expm():
-    A = jnp.asarray([[0.0, 1.0], [-1.0, -0.3]])
-    x_0 = jnp.asarray([1.0, 0.5])
+@pytest.mark.parametrize(
+    ("dtype", "euler_tol", "rk4_tol"),
+    [(jnp.float32, 4e-3, 1e-4), (jnp.float64, 2e-3, 1e-12)],
+)
+def test_linear_system_vs_expm(dtype, euler_tol, rk4_tol):
+    A = jnp.asarray([[0.0, 1.0], [-1.0, -0.3]], dtype)
+    x_0 = jnp.asarray([1.0, 0.5], dtype)
     T = 2.0
     exact = jsp_linalg.expm(A * T) @ x_0
 
@@ -24,20 +29,26 @@ def f(x):
     euler = solve_ode(f, Euler(), 0.0, T, x_0, dt_0=T / n, max_steps=n)
     rk4 = solve_ode(f, RK4(), 0.0, T, x_0, dt_0=T / n, max_steps=n)
     assert bool(euler.ok) and bool(rk4.ok)
+    assert euler.xs.dtype == dtype and rk4.xs.dtype == dtype
     assert int(euler.num_steps) == int(euler.num_accepted) == n
     assert int(rk4.num_steps) == int(rk4.num_accepted) == n
-    assert jnp.max(jnp.abs(euler.xs - exact)) < 2e-3
-    assert jnp.max(jnp.abs(rk4.xs - exact)) < 1e-12
+    assert jnp.max(jnp.abs(euler.xs - exact)) < euler_tol
+    assert jnp.max(jnp.abs(rk4.xs - exact)) < rk4_tol
 
 
-def test_logistic_closed_form():
-    x_0 = jnp.asarray(0.1)
+@pytest.mark.parametrize(
+    ("dtype", "tol"),
+    [(jnp.float32, 2e-5), (jnp.float64, 1e-9)],
+)
+def test_logistic_closed_form(dtype, tol):
+    x_0 = jnp.asarray(0.1, dtype)
     T = 3.0
     n = 300
     sol = solve_ode(
         lambda x: x * (1.0 - x), RK4(), 0.0, T, x_0, dt_0=T / n, max_steps=n
     )
-    assert jnp.abs(sol.xs - logistic_exact(x_0, T)) < 1e-9
+    assert sol.xs.dtype == dtype
+    assert jnp.abs(sol.xs - logistic_exact(x_0, jnp.asarray(T, dtype))) < tol
 
 
 def test_convergence_slopes():
@@ -59,6 +70,64 @@ def f(x):
         assert abs(slope - expected) < 0.3, (type(solver).__name__, slope)
 
 
+def test_unroll_matches_rolled_and_requires_fixed_stepping():
+    x_0 = jnp.asarray(0.1)
+    n = 16
+
+    def solve(unroll):
+        return solve_ode(
+            lambda x: x * (1.0 - x),
+            RK4(),
+            0.0,
+            1.0,
+            x_0,
+            dt_0=1.0 / n,
+            max_steps=n,
+            save_at=SaveAt(steps=True),
+            unroll=unroll,
+        )
+
+    rolled, unrolled = solve(1), solve(4)
+    assert jnp.array_equal(rolled.xs, unrolled.xs)
+    assert jnp.array_equal(rolled.ts, unrolled.ts)
+    grad_rolled = jax.grad(
+        lambda x: (
+            solve_ode(
+                lambda v: v * (1.0 - v), RK4(), 0.0, 1.0, x, dt_0=1.0 / n, max_steps=n
+            ).xs
+        )
+    )(x_0)
+    grad_unrolled = jax.grad(
+        lambda x: (
+            solve_ode(
+                lambda v: v * (1.0 - v),
+                RK4(),
+                0.0,
+                1.0,
+                x,
+                dt_0=1.0 / n,
+                max_steps=n,
+                unroll=4,
+            ).xs
+        )
+    )(x_0)
+    assert jnp.allclose(grad_rolled, grad_unrolled, rtol=1e-12, atol=1e-12)
+    with pytest.raises(ValueError, match="fixed stepping"):
+        solve_ode(
+            lambda x: -x,
+            Tsit5(),
+            0.0,
+            1.0,
+            x_0,
+            dt_0=0.1,
+            controller=IController(),
+            max_steps=32,
+            unroll=4,
+        )
+    with pytest.raises(ValueError, match="at least 1"):
+        solve(0)
+
+
 def test_non_dividing_dt_0_lands_on_t_1():
     x_0 = jnp.asarray(0.1)
     sol = solve_ode(

From 2e04b4df3e4e80d7aa848f54e737fe01b660cf84 Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Tue, 4 Aug 2026 17:39:54 -0700
Subject: [PATCH 2/5] docs: fold guide pages into per-problem pages

Nav is now Home / ODEs / SDEs / DAEs / SDAEs / Markov / Exponential /
API. static_shapes, adaptive_ad, aux, and rodas5p fold into the new
ode.md and the problem pages; migration notes are gone; README gains
the SRA1 explicit-noise example; llms.txt regenerated.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01C8SWZyLFpzVjEX5WCRLJnD
---
 README.md     | 216 ++++++++++++------------------
 docs/api.md   |   6 +
 docs/dae.md   | 363 +++++++++++++++++++++-----------------------------
 docs/index.md | 162 +++++++---------------
 docs/llms.txt |  17 +--
 docs/ode.md   | 187 ++++++++++++++++++++++++++
 docs/sdae.md  | 123 ++++++++---------
 docs/sde.md   | 129 ++++++++++++------
 mkdocs.yml    |   7 +-
 9 files changed, 642 insertions(+), 568 deletions(-)
 create mode 100644 docs/ode.md

diff --git a/README.md b/README.md
index f8e5296..c716e9a 100644
--- a/README.md
+++ b/README.md
@@ -8,71 +8,27 @@
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 
 Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX: fixed-step Euler/RK4,
-adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs,
-and Euler–Maruyama for Itô SDEs and semi-explicit index-1 SDAEs. The package
-also includes primal, vmap-friendly finite-state DTMC and CTMC simulators with
-sequential and associative parallel-prefix execution. Deterministic probability
-forecasts are differentiable in the initial mass and include DTMC matrix powers,
-dense CTMC exponentials, and matrix-free Arnoldi/Krylov actions over probability
-pytrees. The same dense and matrix-free backends are available directly through
-`solve_linear_ode` for any fixed homogeneous linear array or pytree operator;
-`jvp_linear_ode` and `vjp_linear_ode` apply the exact initial-state tangent and
-adjoint exponential actions without differentiating Arnoldi orthogonalization.
-Fixed stepping and the default adaptive path use bounded `lax.scan` loops with
-exactly `max_steps` attempt slots. Shapes stay static as tolerances or curvature
-change, and these solves support forward mode, reverse mode, and
-reverse-over-forward. Adaptive ODE and DAE solves may instead select
-`adaptive_loop="forward"`: a dynamic `lax.while_loop` that executes only actual
-attempts and supports JVP and nested forward AD, but not reverse mode. A vmapped
-forward loop runs until its slowest lane finishes.
+adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs, and
+fixed-step Euler–Maruyama, Milstein, and SRA1 for Itô SDEs and semi-explicit
+index-1 SDAEs. Solves run in bounded `lax.scan` loops with static shapes and
+compose with `jit`, `vmap`, forward mode, reverse mode, and
+reverse-over-forward. Finite-state Markov simulation, probability forecasts,
+and general fixed homogeneous linear solves (dense or matrix-free Krylov
+exponential actions, after SciML's
+[`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/))
+round out the package.
 
 This is a deliberately small, jvp/vjp-friendly package. Rodas5P is a JAX
-adaptation of Steinebach's method and follows SciML's
+adaptation of Steinebach's method following SciML's
 [`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock)
-implementation. Use [diffrax](https://docs.kidger.site/diffrax/) or
+implementation, and DAE algebraic roots delegate both the primal solve and
+the implicit derivative to
+[`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/). Use
+[diffrax](https://docs.kidger.site/diffrax/) or
 [SciML](https://docs.sciml.ai/DiffEqDocs/stable/) if you need general mass
-matrices, fully implicit or higher-index DAEs, events, continuous solution
-objects, sparse/Krylov linear solvers for ODE/DAE stages, or specialized
-adjoints. Initial DAE consistency and explicit DAE stages use `nlls-gram`;
-the same nlls solve supplies the square root's implicit derivative, whose
-default is a direct nonsymmetric `LU()` solve. `LMRootSolver` requires
-residual-only stopping (`gtol=xtol=0`) and accepts only `CONVERGED` roots whose
-Euclidean residual norm is below the root `atol`. Its
-`max_steps_is_success` field remains for source compatibility but does not make
-`MAX_STEPS` a valid DAE root.
-
-The linear exponential-action API follows SciML
-[`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/).
-It includes fixed and residual-controlled adaptive matrix-free time slicing;
-the latter keeps the Krylov dimension static for predictable JAX compilation.
-SciML's
-[`ExponentialIntegrators.jl`](https://docs.sciml.ai/ExponentialIntegrators/stable/)
-is the reference for the broader nonlinear exponential-integrator family.
-
-## 2.4.0 migration note
-
-- `SaveAt(ts=..., exact=True)` now gathers realized knots for explicit
-  fixed-step ODEs. Every query must align with a knot; adaptive methods,
-  Rodas5P, DAEs, SDEs, and SDAEs continue to reject exact mode.
-- `Solution.num_steps` and `DAESolution.num_steps` count logical attempts,
-  including rejections. DAE results additionally expose `num_root_solves` and
-  `num_root_steps`; `num_accepted` retains its existing meaning.
-- Adaptive ODE and DAE solves may opt into `adaptive_loop="forward"` for an
-  actual-work loop. It supports primal, JVP, and nested forward AD but not
-  reverse mode; `adaptive_loop="bounded"` remains the reverse-mode-capable
-  default. Under `vmap`, the forward loop runs to the slowest lane.
-- `LMRootSolver(predictor="secant")` is an opt-in continuation warm start for
-  locally unique algebraic branches; `predictor="previous"` remains the
-  default.
-
-DAE root acceptance is stricter in 2.4.0. nlls-gram owns both the primal root
-solve and implicit derivative; square implicit AD defaults to direct `LU()`.
-Only `CONVERGED` roots whose residual norm is below `atol` are accepted, so
-`gtol` and `xtol` must both be zero. `max_steps_is_success` remains for source
-compatibility, now defaults to `False`, and never makes `MAX_STEPS` a valid
-root. Upgrading configurations should remove nonzero `gtol`/`xtol`; if they
-relied on budget exhaustion, increase the root budget or adjust the residual
-tolerance instead.
+matrices, fully implicit or higher-index DAEs, adaptive SDE stepping, events,
+continuous solution objects, sparse/Krylov ODE/DAE stages, or specialized
+adjoints.
 
 ## Install
 
@@ -91,12 +47,8 @@ uv add tinydiffeq "jax[cuda13]"
 
 The vector field may take `(x)`, `(x, t)`, `(x, t, args)`, or
 `(x, t, args, p)` — always in that order. `args` is pass-through data (not an
-AD target by convention); `p` holds differentiable parameters (any pytree).
-The state may also be any JAX pytree. It must contain at least one leaf, and
-every leaf must be a nonempty real floating array with the same dtype; vector
-fields and `project` preserve that structure. Output keeps the structure and
-adds the saved-time axis to each
-leaf.
+AD target by convention); `p` holds differentiable parameters, and the state
+may be any pytree of same-dtype real floating arrays.
 
 ```python
 import jax
@@ -122,29 +74,61 @@ print(sol.xs)   # states on the grid
 print(sol.ok)   # reached t_1 with every requested output valid?
 ```
 
-`IController()` and `PIController()` choose tolerances from `x_0.dtype`:
-`rtol=1e-4, atol=1e-6` for float32 and `rtol=1e-7, atol=1e-9` for
-float64. Pass explicit values when tolerances are part of your model's
-scientific specification. The default `dt_min` is
-`10 * finfo(dtype).eps * max(1, abs(t_1))`.
-
-`max_steps` is the total internal **attempt budget**: accepted steps plus
-rejections. It is not normally the number of returned times. Endpoint mode
-returns one time/state, `SaveAt(ts=...)` returns the requested grid, and
-`SaveAt(steps=True)` returns the initial state and accepted internal steps as
-a contiguous prefix of `max_steps + 1` rows. The remaining rows repeat the
-last accepted state by default; `sol.accepted` distinguishes data from
-padding. Rejected attempts never appear in the returned trajectory.
-`sol.num_steps` reports the number of attempts actually made, while
-`sol.num_accepted` excludes rejections.
-
-`SaveAt(ts=...)` also accepts a Python sequence. These are observation times:
-the adaptive controller still chooses its own internal mesh. Explicit methods
-use cubic Hermite interpolation; Rodas5P uses its published stiff-aware
-fourth-order continuous extension. For an explicit fixed-step ODE,
-`SaveAt(ts=..., exact=True)` instead requires every requested time to be an
-internal knot and gathers the stored state directly. Exact mode does not apply
-to adaptive ODEs, Rodas5P, DAEs, SDEs, or SDAEs.
+`max_steps` is the internal attempt budget (accepted plus rejected steps),
+not the number of returned times: `SaveAt` picks the endpoint, a fixed
+interpolation grid, or the padded accepted-step prefix, so output shapes
+never depend on how many steps the controller took. Omitted controller
+tolerances follow the state dtype (`1e-4`/`1e-6` in float32,
+`1e-7`/`1e-9` in float64).
+
+## SDEs with first-class noise
+
+`solve_sde` integrates diagonal-noise Itô SDEs with `EulerMaruyama` (strong
+order 0.5), `Milstein` (1.0, commutative diagonal noise), or `SRA1` (1.5,
+additive noise). An Ornstein–Uhlenbeck process under SRA1:
+
+```python
+from tinydiffeq import solve_sde, SRA1
+
+theta, sigma, n = 1.0, 0.5, 256
+
+
+def ou_drift(x):
+    return -theta * x
+
+
+def ou_diffusion(x):
+    return sigma * jnp.ones_like(x)
+
+
+sol = solve_sde(
+    ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, jnp.asarray(1.0),
+    key=jax.random.key(0), n_steps=n,
+)
+```
+
+The noise realization can also be passed explicitly — the same pytree
+`sample_noise` would draw, now inspectable, storable data that is
+differentiable like any other input:
+
+```python
+x_0 = jnp.asarray(1.0)
+noise = SRA1().sample_noise(x_0, jax.random.key(0), n, jnp.asarray(1.0 / n), x_0.dtype)
+same_sol = solve_sde(
+    ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n
+)  # bit-identical to the key= call
+d_endpoint_d_noise = jax.grad(
+    lambda noise: solve_sde(
+        ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n
+    ).xs
+)(noise)
+```
+
+A fixed key (or fixed noise) pins the whole path, so gradients with respect
+to `x_0`, `p`, and `noise` are pathwise derivatives under common random
+numbers — the setup simulation-based estimators want. `vmap` over
+trajectories with per-trajectory `x_0` and noise composes with `jit` and
+`grad`.
 
 ## Semi-explicit DAEs
 
@@ -152,7 +136,7 @@ For a square index-1 system `dy/dt = f(y, z, t, args, p)` and
 `0 = g(y, z, t, args, p)`:
 
 ```python
-from tinydiffeq import IController, Rodas5P, Tsit5, solve_semi_explicit_dae
+from tinydiffeq import solve_semi_explicit_dae
 
 
 def dae_f(y, z, t, args, p):
@@ -171,36 +155,17 @@ dae_sol = solve_semi_explicit_dae(
     controller=IController(), max_steps=128,
 )
 print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"])
-
-# One initial nonlinear consistency solve, then linear Rodas5P stages.
-stiff_dae_sol = solve_semi_explicit_dae(
-    dae_f, dae_g, Rodas5P(), 0.0, 1.0,
-    jnp.asarray(1.0), jnp.asarray(0.5),
-    p=jnp.asarray(2.0), dt_0=0.1,
-    controller=IController(), max_steps=128,
-)
 ```
 
 `z_0` is a guess and is made consistent automatically. RK4 and Tsit5 restore
-the algebraic root at every stage. Rodas5P performs no nonlinear solves after
-initialization: it advances the corresponding block mass-matrix system using
-one reused LU factorization per attempt. Differential fields may return a
-floating saved-aux pytree stored at accepted nodes and interpolated on requested
-deterministic grids. Algebraic equations may separately return internal context
-passed to the dynamics. On the default bounded path, JVP, VJP, and
-reverse-over-forward propagate through both implicit initialization and the
-time integrator. See the
-[DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
-for root controls, `SaveAt`, and scope limits.
-
-DAE solutions expose `num_steps`, `num_root_solves`, and `num_root_steps` as
-logical per-trajectory work counters. Explicit methods default to reusing the
-previous algebraic root as the next stage guess; `LMRootSolver(predictor="secant")`
-is an opt-in continuation predictor for locally unique root branches.
-
-Fixed-step semi-explicit Itô SDAEs use the corresponding
-`solve_semi_explicit_sdae` interface with `EulerMaruyama`, a PRNG key, and
-`n_steps`; see the [SDAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/).
+the algebraic root at every stage through `nlls-gram`, which also supplies
+the root's implicit derivative; `Rodas5P()` instead performs one initial
+consistency solve and then advances the block mass-matrix system with one
+reused LU factorization per attempt — the stiff path. Stochastic
+semi-explicit systems use `solve_semi_explicit_sdae` with `EulerMaruyama` or
+`SRA1`. See the
+[DAE](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) and
+[SDAE](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) docs.
 
 ## Gradients through the solve
 
@@ -212,21 +177,16 @@ def endpoint(p):
         max_steps=512,
     ).xs
 
-jax.grad(endpoint)(jnp.asarray(1.3))                      # reverse mode
+jax.grad(endpoint)(jnp.asarray(1.3))                         # reverse mode
 jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),))  # forward mode
-jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])(
-    jnp.asarray(1.3)
-)                                                          # reverse-over-forward
 ```
 
 The step-size controller is wrapped in `stop_gradient` (accept/reject is
-non-differentiable either way, and the error-ratio power blows up at exactly
-zero error); states differentiate through the solver stages on the realized,
-frozen mesh. In particular, adaptive `SaveAt(steps=True)` does not include mesh
-motion in its time or state derivatives. See the
+non-differentiable either way); states differentiate through the solver
+stages on the realized, frozen mesh. See the
 [docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
-contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key
-semantics, and the package API.
+contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE
+noise semantics, and the package API.
 
 ## License
 
diff --git a/docs/api.md b/docs/api.md
index 7a2fc07..8f4a596 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -36,6 +36,10 @@
 
 ::: tinydiffeq.EulerMaruyama
 
+::: tinydiffeq.Milstein
+
+::: tinydiffeq.SRA1
+
 ::: tinydiffeq.DenseExponential
 
 ::: tinydiffeq.KrylovExponential
@@ -78,6 +82,8 @@
 
 ## Utilities
 
+::: tinydiffeq.diagonal_brownian_increments
+
 ::: tinydiffeq.hermite_interpolate
 
 ::: tinydiffeq.cumulative_trapezoid
diff --git a/docs/dae.md b/docs/dae.md
index 8aff8ec..ea5c966 100644
--- a/docs/dae.md
+++ b/docs/dae.md
@@ -8,36 +8,22 @@ $$
 $$
 
 where the algebraic equation is square and $g_z$ is nonsingular along the
-solution. `y` and `z` may independently be array or pytree states. Leaves
-within each state share one real floating dtype; the `y` and `z` dtypes may
-differ. The residual `g` is a single array whose flattened size matches the
-total size of `z`. The implementation supports root-restored RK4/Tsit5 and
-linearly implicit Rodas5P with fixed or adaptive control.
+solution. `y` and `z` may independently be array or pytree states; their
+dtypes may differ. The residual `g` is a single array whose flattened size
+matches the total size of `z`. Root-restored RK4/Tsit5 and linearly implicit
+Rodas5P are supported, with fixed or adaptive control.
 
 The algebraic solve uses
 [`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/)'s
-general Levenberg–Marquardt solver for both the primal root and its implicit
-derivative. The default `Cholesky()` primal solver selects the dense normal form
-for the square algebraic system, while nlls's default `ad_solver` selects the
-direct nonsymmetric `LU()` rule for that square system. The implicit rule
-differentiates the defining equation rather than the LM iterations; no implicit
-ridge is added. Tinydiffeq applies the DAE validity policy around that solve but
-does not implement a separate square-system IFT.
-
-Rodas5P is a JAX adaptation of Steinebach's method following SciML's
-[`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock)
-implementation. See [Rodas5P](rodas5p.md) for direct links to SciML's tableau,
-step, and interpolation sources.
+Levenberg–Marquardt solver for both the primal root and its implicit
+derivative: the primal defaults to dense `Cholesky()` normal equations, and
+the square implicit rule defaults to a direct nonsymmetric `LU()` solve of
+the defining constraint rather than a differentiation of the LM iterations.
 
 ## Minimal examples
 
-Consider
-
-$$
-\dot y = pz, \qquad 0=z-y,
-$$
-
-whose reduced solution is $y(t)=z(t)=y_0e^{pt}$.
+Consider $\dot y = pz$, $0 = z - y$, whose reduced solution is
+$y(t) = z(t) = y_0 e^{pt}$.
 
 ```python
 import jax.numpy as jnp
@@ -74,19 +60,12 @@ adaptive = solve_semi_explicit_dae(
 )
 
 adaptive.aux["flow"]
-
-linearly_implicit = solve_semi_explicit_dae(
-    f, g, Rodas5P(), 0.0, 1.0,
-    jnp.asarray(1.0), jnp.asarray(0.5),
-    p=jnp.asarray(2.0), dt_0=0.1,
-    controller=IController(), max_steps=128,
-)
 ```
 
-`z_0` is a root-finding guess, not an assumed-consistent initial value. Both
+`z_0` is a root-finding guess, not an assumed-consistent initial value: both
 calls first solve `g(y_0, z, t_0, args, p) = 0`, so the `0.5` guess becomes
-the consistent value `1.0`. RK4 and Tsit5 then solve the algebraic equation at
-every stage. Rodas5P performs no further nonlinear solves.
+the consistent value `1.0`. RK4 and Tsit5 then re-solve the algebraic
+equation at every stage.
 
 ## Nonlinear-solve and AD contract
 
@@ -97,7 +76,6 @@ from tinydiffeq import LMRootSolver
 
 root_solver = LMRootSolver(
     max_steps=8,
-    max_steps_is_success=False,  # compatibility field; MAX_STEPS stays invalid
     atol=None,          # 1e-6 float32, 1e-10 float64
     gtol=0.0,           # required: residual stopping only
     xtol=0.0,           # required: residual stopping only
@@ -106,30 +84,16 @@ root_solver = LMRootSolver(
 )
 ```
 
-The outer `max_steps` counts attempted time steps, including adaptive
-rejections. `root_solver.max_steps` separately bounds one algebraic root. For
-Rodas5P it affects only initial consistency; the method's later stages reuse
-one dense LU factorization per attempted time step.
-
-Root tolerances are independent of the outer controller tolerances. Explicit
-`atol` must be positive; `None` selects the dtype default, and `atol=0` is
-invalid. `gtol` and `xtol` must both equal zero; `LMRootSolver` rejects nonzero
-values so `CONVERGED` can only come from the residual test. Every accepted
-algebraic root must report `CONVERGED` and have Euclidean residual norm
-`sqrt(sum(residual**2))` strictly below the root `atol`.
-`max_steps_is_success` remains in the configuration for source compatibility,
-but Tinydiffeq always asks nlls to treat `MAX_STEPS` as a failed implicit solve;
-setting the field to `True` does not broaden DAE root acceptance.
-
-The primal nonlinear solve uses nlls-gram's dense `Cholesky()` normal-equation
-default. For a successful square root, nlls's implicit rule forms `dg/dz`,
-forms the right-hand side with respect to `(y, t, p)`, and applies its direct
-`LU()` square solve. Transposing that nlls rule supplies the VJP. The primal LM
-factorization is not reused because it represents a damped normal equation,
-not the accepted root's generally nonsymmetric Jacobian. A failed nlls status
-has zero implicit tangent by the nlls solve contract. For the rare root solve
-that needs to depart from the defaults, `solver_options` forwards constructor
-arguments for either the primal `linear_solver` or the implicit `ad_solver`:
+The outer `max_steps` counts attempted time steps; `root_solver.max_steps`
+separately bounds one algebraic root, and root tolerances are independent of
+the controller tolerances. `gtol` and `xtol` must be zero so `CONVERGED` can
+only come from the residual test: every accepted root must report
+`CONVERGED` with Euclidean residual norm strictly below `atol`, and a
+`MAX_STEPS` iterate is never a differentiable root.
+
+`solver_options` forwards constructor arguments verbatim to nlls-gram's
+`LevenbergMarquardt` — either the primal `linear_solver` or the implicit
+`ad_solver`:
 
 ```python
 from nlls_gram import QR
@@ -138,100 +102,111 @@ root_solver = LMRootSolver(solver_options={"linear_solver": QR()})
 ```
 
 The names and semantics are nlls-gram's, so they track that package rather
-than being mirrored here. Pass a mapping or key/value pairs; it is normalized
-to a sorted tuple so equal configurations remain hashable and share one
-compiled solver. Algebraic roots fix `cache_jacobian=False` and
-`geodesic_acceleration=False` — each DAE stage changes the root problem and
-the intended path is the ordinary dense LM step — and `solver_options` rejects
-those two options rather than silently honoring an override. `ad_solver`
-remains nlls-owned and may be supplied explicitly; the square default is
-`LU()`.
-
-`predictor="previous"` is the default: each explicit RK stage starts from the
-most recent successful algebraic root. `predictor="secant"` extrapolates from
-the accepted-step root through the most recent successful stage at a later
-time. Duplicate RK4 stage times, non-forward targets, and failed stages fall
-back to the previous root. The predictor is stop-gradiented, so a successful
-root still uses the same implicit derivative. Its time-derived extrapolation
-scale is cast separately to each algebraic leaf's dtype, preserving the `z`
-dtype when the differential/time and algebraic dtypes differ. Secant prediction
-assumes the continued branch is locally unique; with multiple roots or a tight
-finite iteration budget, changing the guess can change the selected branch,
-value, or status.
-
-Every nonlinear root passes `(y, t, p)` to nlls-gram. Its implicit rule then
+than being mirrored here. Options are normalized to a sorted tuple so equal
+configurations stay hashable and share one compiled solver.
+`cache_jacobian` and `geodesic_acceleration` are fixed to `False` and
+rejected — each DAE stage changes the root problem, and the intended path is
+the ordinary dense LM step.
+
+`predictor="previous"` (default) starts each explicit RK stage from the most
+recent successful root. `predictor="secant"` extrapolates from the
+accepted-step root through the most recent successful stage at a strictly
+later time; duplicate stage times and failed stages fall back to the
+previous root. Predictor values are stop-gradiented, so successful roots keep
+the same implicit derivative — but with multiple algebraic roots a different
+warm start can select a different branch, so secant assumes the continued
+root is locally unique.
+
+Every nonlinear root passes `(y, t, p)` to nlls-gram, whose implicit rule
 differentiates the defining constraint,
 
 $$
-\dot z = -g_z^{-1}
-  (g_y\dot y + g_t\dot t + g_p\dot p),
+\dot z = -g_z^{-1}(g_y\dot y + g_t\dot t + g_p\dot p),
 $$
 
-rather than differentiating the LM iterations. The warm-start guess has zero
-derivative by design. Rodas5P differentiates through its exact JAX Jacobian,
-time derivative, LU factorization, and linear stage solves. `args` is fixed
-data; put every differentiated model quantity in `p`. On the default bounded
-integration path, JVP, VJP, `vmap`, and reverse-over-forward compose through
-the complete DAE solve.
-
-`sol.num_steps` counts logical time-step attempts, including adaptive
-rejections. `sol.num_root_solves` counts active nonlinear root calls, including
-the initial consistency solve and failed calls, and `sol.num_root_steps` sums
-their LM update counts. Rodas5P therefore reports one root solve regardless of
-its time-step count; later linear stages are not nonlinear roots. These
-counters have exact-zero tangents. Under `vmap`, a masked lane may still execute
-physically while remaining absent from its logical counters.
-
-The default `adaptive_loop="bounded"` uses a reverse-mode-capable bounded scan.
-`adaptive_loop="forward"` uses a dynamic actual-work loop for adaptive Tsit5
-and Rodas5P. It supports primal evaluation, JVP, and nested forward mode, but
-not ordinary reverse mode. A vmapped forward loop runs until the slowest lane
-finishes.
-
-The differential field may return `(dy, saved_aux)`. Saved aux is a nonempty
-pytree of nonempty real floating arrays; different leaves may use different
-floating dtypes. tinydiffeq evaluates it at required saved nodes. Ordinary JAX
-differentiation composes with either the root's implicit derivative or the
-Rodas5P stages, so aux tangents and cotangents include both direct dependence
-on `p` and indirect dependence through `z`.
-
-The algebraic function may instead or additionally return
-`(residual, algebraic_aux)`. In that case `f` takes
-`(y, z, t, args, p, algebraic_aux)`. This value is internal cached context:
-the nonlinear solver sees only the residual, and only differential-field
-saved aux appears in `sol.aux`. See [Auxiliary Outputs](aux.md) for the four
-supported combinations and flag behavior.
-
-Every saved aux leaf and every inexact algebraic-aux leaf must be finite.
-Invalid algebraic context at initialization sets `ok=False` before any
-time-step work. `SaveAt(steps=True)` and `SaveAt(ts=...)` check saved aux at
-the initial and accepted nodes, so an invalid value freezes the previous valid
-prefix. Endpoint mode evaluates saved aux only after integration; an invalid
-final value retains the endpoint state, returns zero aux, and sets `ok=False`.
-
-An adaptive stage-root failure rejects the time-step attempt and asks the
-controller for a smaller step; a fixed-step failure terminates. Rodas5P linear
-failures follow the same controller policy.
-In either case `sol.ok` is false if the endpoint is not reached with valid
-algebraic states. nlls supplies the primal LM iterate, diagnostics, and the
-implicit JVP/VJP. Tinydiffeq applies the DAE's residual-and-status acceptance
-check and returns the differentiation-inert warm-start guess when that check
-fails; aux at a failed initial root is a zero pytree of the declared shape.
-Callers that want to retain successful-lane JVPs/VJPs after another lane has
-already become inactive should pass
-`failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)`, choosing a point where
-the residual, context, and saved-aux maps are finite and differentiable.
-tinydiffeq substitutes this point into an already-inactive root call before
-entering nlls, and also uses it for inactive algebraic aux, saved aux, and
-differential-field evaluations. It is not an nlls solve argument and never
-changes an active root attempt.
-
-Every newly attempted root must therefore be JVP-safe at its actual
-`(y, z_guess, t, p)`: the residual and any model context evaluated there must
-have valid derivatives. The reference cannot rescue an intrinsically invalid
-active attempt after the fact. Without an explicit reference, an all-ones
-best-effort default is used for inactive work; gradients are not guaranteed if
-the model is undefined there. A failed lane itself is never a valid solution.
+rather than the LM iterations; the warm-start guess has zero derivative by
+design. `args` is fixed data; put every differentiated quantity in `p`. On
+the default bounded path, JVP, VJP, `vmap`, and reverse-over-forward compose
+through the complete solve. `adaptive_loop="forward"` runs an actual-work
+loop for adaptive Tsit5 and Rodas5P (primal, JVP, and nested forward mode
+only), as for [ODEs](ode.md#static-shapes-and-saveat).
+
+`sol.num_steps` counts logical time-step attempts including rejections;
+`sol.num_root_solves` counts active nonlinear root calls (including the
+initial consistency solve and failures) and `sol.num_root_steps` sums their
+LM updates. All have exact-zero tangents. An adaptive stage-root failure
+rejects the attempt and retries with a smaller step; a fixed-step failure
+terminates. Either way `sol.ok` is false if `t_1` is not reached with valid
+algebraic states.
+
+For batched differentiation where a lane may fail or leave the model domain,
+pass `failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)` at a point where
+the residual, context, and aux maps are finite and differentiable. It is
+substituted into already-inactive root calls and inactive aux/field
+evaluations so masked lanes cannot poison the JVP/VJP of successful lanes. A
+newly attempted root must still be JVP-safe at its actual
+`(y, z_guess, t, p)`; the reference never rescues an active failure, and a
+failed lane's outputs are not a valid solution. Without an explicit
+reference, an all-ones best-effort default is used.
+
+## Auxiliary outputs
+
+Saved output belongs to the differential field, and the algebraic function
+may separately expose internal context:
+
+```python
+def g(y, z, t, args, p):
+    return residual, algebraic_aux
+
+
+def f(y, z, t, args, p, algebraic_aux):
+    return dy, saved_aux
+```
+
+All four combinations are supported: neither aux, saved aux only, algebraic
+aux only, or both. When algebraic aux is present, `f` must take the full
+six-argument form. `has_aux` / `has_algebraic_aux` default to abstract
+auto-detection; explicit `False` skips those traces.
+
+Algebraic aux is internal cached context: it is passed to `f`, included in
+the implicit derivative path, ignored by the nonlinear solver's residual
+interface, and never stored or interpolated. It may be a nonempty pytree of
+bool, integer, real, or complex arrays; every inexact leaf must be finite.
+Only the differential field's real-floating `saved_aux` becomes `sol.aux`,
+stored at accepted nodes and interpolated on requested grids. Invalid
+algebraic context at initialization fails before any time-step work; invalid
+saved aux in a prefix mode terminates at the previous accepted node, while
+endpoint mode keeps the endpoint state with zero aux and `ok=False`.
+
+## Rodas5P for DAEs
+
+For the stiff path, tinydiffeq constructs the flattened mass-matrix system
+internally:
+
+$$
+M\dot u=F(u,t), \qquad
+u=(y,z), \quad M=\operatorname{diag}(I_y,0_z), \quad F=(f,g).
+$$
+
+```python
+linearly_implicit = solve_semi_explicit_dae(
+    f, g, Rodas5P(), 0.0, 1.0,
+    jnp.asarray(1.0), jnp.asarray(0.5),
+    p=jnp.asarray(2.0), dt_0=0.1,
+    controller=IController(), max_steps=128,
+)
+```
+
+`LMRootSolver` is used once, for initial consistency. Every later Rodas5P
+stage solves $\left(M/(\gamma h) - F_u\right)k_i = r_i$ with one reused
+pivoted LU factorization per attempted step — no nonlinear endpoint
+restoration. Returned internal `z` values therefore satisfy the constraint
+to the method's integration accuracy, not to `LMRootSolver.atol`, and the
+solver reports one root call regardless of its attempt count. This
+intentionally differs from RK4/Tsit5, which root-solve every stage; Rodas5P
+is the choice when those algebraic solves dominate runtime or the coupled
+dynamics are stiff. See [Stiff ODEs: Rodas5P](ode.md#stiff-odes-rodas5p) for
+the method, its AD boundaries, and the SciML/Steinebach credit and links.
 
 ## Saving output
 
@@ -240,79 +215,47 @@ All `SaveAt` modes are supported:
 - `SaveAt(t_1=True)` returns the endpoint.
 - `SaveAt(steps=True)` returns the initial point and accepted internal steps
   as a padded `max_steps + 1` buffer with the usual `accepted` mask.
-- `SaveAt(ts=grid)` uses cubic Hermite for root-restored methods and Rodas5P's
-  stiff-aware continuous extension for `(y, z)`. Aux uses cubic Hermite in
-  both cases. It performs no query-time nonlinear solves.
-
-`SaveAt(ts=..., exact=True)` is not a DAE mode; exact knot gathering is limited
-to explicit fixed-step ODEs.
+- `SaveAt(ts=grid)` uses cubic Hermite for root-restored methods and
+  Rodas5P's stiff-aware continuous extension for `(y, z)`; aux uses cubic
+  Hermite in both cases. No query-time nonlinear solves are performed.
+  `exact=True` is not a DAE mode.
 
 The result is a `DAESolution` with `ts`, `ys`, `zs`, `ok`, `num_accepted`,
-`accepted`, `aux`, `num_steps`, `num_root_solves`, and `num_root_steps` fields.
-For pytree states, saved rows are a leading axis on every state and aux leaf;
-the one `accepted` mask applies to the complete output.
+`accepted`, `aux`, `num_steps`, `num_root_solves`, and `num_root_steps`.
 
 ### Dense output for root-restored RK4 and Tsit5
 
 At a consistent knot, differentiating the constraint gives
-
-$$
-g_z\dot z = -(g_y\dot y + g_t).
-$$
-
-tinydiffeq solves this linear system once per accepted knot only when a query
-grid is requested. It then obtains `aux_dot` by a JVP of the aux map along
-$(\dot y,\dot z,1)$. Values and total derivatives feed the same normalized
-cubic Hermite basis used for ODE states. This is an order-3 continuous
-extension—uniform interpolation error $O(h^4)$—when `f` and `g` are $C^4$,
-$g_z$ stays uniformly nonsingular near the solution, and root error is no
-larger than the desired dense-output error. RK4 and Tsit5 knot errors meet the
-required order under their usual assumptions.
-
-The normalized coordinate stays in `[0, 1]` and Hermite basis coefficients
-are bounded by 3, which is favorable in float32. SciML's specialized Tsit5
-dense polynomial has one higher order for `y`, but requires all seven stages,
-does not directly supply `z`/aux output, and has much larger coefficients.
-Using one Hermite construction keeps `y`, `z`, and aux at the same dense order
-with substantially less storage.
-
-Interpolated `z` and aux are approximations: away from accepted knots they
-need not satisfy `g=0` exactly. The constraint defect is $O(h^4)$ under the
-conditions above. Use `SaveAt(steps=True)` when every returned row must be an
-actual converged root. Dense output also requires one `g_z` factorization per
-accepted knot, rather than one nonlinear solve per requested time; its cost
-therefore scales with internal steps rather than grid length.
+$g_z\dot z = -(g_y\dot y + g_t)$. tinydiffeq solves this linear system once
+per accepted knot (only when a query grid is requested), obtains `aux_dot`
+by a JVP of the aux map along $(\dot y, \dot z, 1)$, and feeds values and
+derivatives into the same normalized cubic Hermite basis used for ODE
+states — an order-3 continuous extension with uniform error $O(h^4)$ when
+`f` and `g` are $C^4$, $g_z$ stays uniformly nonsingular, and root error is
+below the dense-output error. Interpolated `z` and aux are approximations:
+away from knots the constraint defect is $O(h^4)$. Use `SaveAt(steps=True)`
+when every returned row must be an actual converged root.
 
 ### Dense output for Rodas5P
 
-Rodas5P stores the three coefficient pytrees defined by Steinebach's
-fourth-order stiff-aware continuous extension. tinydiffeq evaluates the same
-polynomial form used by
+Rodas5P stores the three coefficient pytrees of Steinebach's fourth-order
+stiff-aware continuous extension and evaluates the same polynomial form as
 [SciML's Rosenbrock interpolant](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_interpolants.jl)
-for the combined `(y, z)` state. No `g_z` factorization or nonlinear solve is
-performed for requested times.
-
-Aux remains a stored accepted-knot quantity. Its cubic-Hermite endpoint
-tangents come from the Rodas polynomial's endpoint derivatives and a JVP of
-the aux map. Aux is therefore interpolated rather than recalculated at every
-query. Rodas5P accepted knots are not root-restored: their constraint defect,
-and that of dense output, is controlled by integration accuracy rather than
-`LMRootSolver.atol`.
-
-Knot selection and adaptive step sizes remain differentiation-inert under the
-frozen-controller convention. Values, implicit slopes, and aux differentiate
-on that realized mesh, but adaptive `SaveAt(steps=True)` omits mesh motion.
-See [Frozen mesh versus a moving-mesh derivative](adaptive_ad.md#frozen-mesh-versus-a-moving-mesh-derivative).
-If `sol.ok` is false, neither outputs nor their derivatives should be treated
-as a valid solution.
+for the combined `(y, z)` state — no $g_z$ factorization or nonlinear solve
+at query times. Aux endpoint tangents come from the Rodas polynomial's
+endpoint derivatives and a JVP of the aux map.
+
+Knot selection and step sizes are differentiation-inert under the
+frozen-controller convention; see
+[AD through adaptive stepping](ode.md#ad-through-adaptive-stepping). If
+`sol.ok` is false, neither outputs nor derivatives are a valid solution.
 
 ## Deliberate limits
 
 Only the internally constructed constant block mass matrix
-`diag(I_y, 0_z)` is supported; there is no public general mass-matrix or fully
-implicit residual API. Rodas5P uses dense Jacobians and dense pivoted LU, not
-sparse or Krylov linear algebra. Higher-index constraints and automatic index
-reduction are unsupported. This is an initial-value solver: it does not
-determine unknown initial costates or solve boundary-value or saddle-path
-conditions. Initial branch selection and jumps between multiple roots are not
-differentiable.
+`diag(I_y, 0_z)` is supported; there is no general mass-matrix or fully
+implicit residual API. Rodas5P uses dense Jacobians and dense pivoted LU.
+Higher-index constraints and automatic index reduction are unsupported. This
+is an initial-value solver: it does not determine unknown initial costates
+or solve boundary-value problems, and jumps between multiple root branches
+are not differentiable.
diff --git a/docs/index.md b/docs/index.md
index 2b74c08..9a19dfc 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -1,64 +1,25 @@
 # tinydiffeq
 
-`tinydiffeq` is a deliberately tiny set of ODE/SDE/DAE/SDAE integrators and
-finite-state Markov simulators for JAX: fixed-step Euler and RK4, adaptive Tsit5 with integral or
-proportional-integral step-size control, linearly implicit Rodas5P for stiff
-ODEs and index-1 DAEs, and fixed-step Euler–Maruyama for Itô SDEs and SDAEs.
-Fixed stepping and the default adaptive path use bounded `lax.scan` loops with
-static shapes. These solves support forward mode, reverse mode, and
-reverse-over-forward. Adaptive ODE and DAE solves also offer
-`adaptive_loop="forward"`, a dynamic actual-work loop for primal, JVP, and
-nested forward AD; JAX cannot reverse-transpose this path.
-Finite-state DTMC/CTMC simulation is primal-only and offers chronological scan
-and associative parallel-prefix methods. Fixed-chain probability forecasts use
-binary matrix powers for DTMC endpoints, dense exponentials for small CTMCs, or
-fixed or adaptive matrix-free Arnoldi/Krylov actions over array or pytree
-probabilities; see
-[Markov Chains](markov_chains.md). The same backends solve general fixed
-homogeneous linear systems; see [Linear Exponential Solves](exponential.md).
-
-Rodas5P follows SciML's
-[`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock)
-implementation and Steinebach's published method. **Use SciML or
+`tinydiffeq` is a deliberately tiny set of differentiable ODE/SDE/DAE/SDAE
+integrators and finite-state Markov simulators for JAX: fixed-step Euler and
+RK4, adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1
+DAEs, and fixed-step Euler–Maruyama, Milstein, and SRA1 for Itô SDEs and
+SDAEs. Solves run in bounded `lax.scan` loops with static shapes and support
+forward mode, reverse mode, and reverse-over-forward; adaptive ODE/DAE solves
+can opt into a dynamic actual-work loop (`adaptive_loop="forward"`, no
+reverse mode). Probability forecasts and general fixed homogeneous linear
+solves use matrix powers, dense exponentials, or matrix-free Krylov actions;
+see [Markov Chains](markov_chains.md) and
+[Linear Exponential Solves](exponential.md).
+
+**Use [SciML](https://docs.sciml.ai/DiffEqDocs/stable/) or
 [diffrax](https://docs.kidger.site/diffrax/) instead if you need any of:**
 
 - general mass matrices, fully implicit solvers, or higher-index DAEs
-- sparse/Krylov linear solves and preconditioners inside general ODE/DAE stages
-- full derivative-term PID step-size control
+- sparse/Krylov linear solves and preconditioners inside ODE/DAE stages
+- adaptive SDE stepping (Brownian-bridge noise), full PID step-size control
 - events, root-finding, or backward-time integration
-- dense output / continuous interpolation objects
-- checkpointed or backsolve adjoints for long horizons
-
-Use the default bounded loop when reverse mode is required. Use the forward
-loop when the actual adaptive attempt count is much smaller than `max_steps`
-and the caller needs only primal or forward-mode execution. Both retain static
-public output shapes; under `vmap`, a dynamic loop still runs until the slowest
-lane finishes.
-
-## 2.4.0 migration note
-
-- `SaveAt(ts=..., exact=True)` now gathers realized knots for explicit
-  fixed-step ODEs. Every query must align with a knot; adaptive methods,
-  Rodas5P, DAEs, SDEs, and SDAEs continue to reject exact mode.
-- `Solution.num_steps` and `DAESolution.num_steps` count logical attempts,
-  including rejections. DAE results additionally expose `num_root_solves` and
-  `num_root_steps`; `num_accepted` retains its existing meaning.
-- Adaptive ODE and DAE solves may opt into `adaptive_loop="forward"` for an
-  actual-work loop. It supports primal, JVP, and nested forward AD but not
-  reverse mode; `adaptive_loop="bounded"` remains the reverse-mode-capable
-  default. Under `vmap`, the forward loop runs to the slowest lane.
-- `LMRootSolver(predictor="secant")` is an opt-in continuation warm start for
-  locally unique algebraic branches; `predictor="previous"` remains the
-  default.
-
-DAE root acceptance is stricter in 2.4.0. nlls-gram owns both the primal root
-solve and implicit derivative; square implicit AD defaults to direct `LU()`.
-Only `CONVERGED` roots whose residual norm is below `atol` are accepted, so
-`gtol` and `xtol` must both be zero. `max_steps_is_success` remains for source
-compatibility, now defaults to `False`, and never makes `MAX_STEPS` a valid
-root. Upgrading configurations should remove nonzero `gtol`/`xtol`; if they
-relied on budget exhaustion, increase the root budget or adjust the residual
-tolerance instead.
+- dense output objects or checkpointed/backsolve adjoints for long horizons
 
 ## Install
 
@@ -86,26 +47,17 @@ f(x, t, args, p)
 ```
 
 - `x` is an array or pytree state. Leaves must share one real floating dtype;
-  the vector field returns the same structure and dtype.
-- `t` is time.
-- `args` is pass-through data. By convention it is **not** an AD target —
-  nothing stops you differentiating with respect to it, but the library's
-  contracts and tests treat it as constants.
-- `p` holds differentiable parameters — any pytree, e.g. neural-network
-  weights. jvp/vjp with respect to `p` and `x_0` are first-class and tested.
-
-The arity is inspected once and the function is wrapped into the canonical
-four-argument form, so the compiled code is identical for all four. There is
-no special autonomous code path: an unused `t` is dead-code-eliminated.
-`drift` and `diffusion` in [`solve_sde`](sde.md) follow the same convention.
-Semi-explicit DAE fields use `(y, z)`, `(y, z, t)`,
-`(y, z, t, args)`, or `(y, z, t, args, p)`; see
-[Semi-Explicit DAEs](dae.md).
-Differential fields and stochastic drifts may return `(value, saved_aux)`;
-DAE/SDAE algebraic functions may separately return internal context consumed
-by the dynamics. Saved aux is differentiated and follows `SaveAt`. See
-[Auxiliary Outputs](aux.md) for all supported contracts and
-[Semi-Explicit SDAEs](sdae.md) for the stochastic form.
+  the field returns the same structure and dtype.
+- `args` is pass-through data — by convention **not** an AD target.
+- `p` holds differentiable parameters (any pytree, e.g. network weights).
+  JVP/VJP with respect to `p` and `x_0` are first-class and tested.
+
+The arity is inspected once and wrapped into the canonical four-argument
+form, so the compiled code is identical for all four. `drift` and `diffusion`
+in [`solve_sde`](sde.md) follow the same convention; semi-explicit DAE fields
+use `(y, z)` through `(y, z, t, args, p)` — see
+[Semi-Explicit DAEs](dae.md) and [SDAEs](sdae.md). Fields may return
+`(value, saved_aux)` to save extra quantities with the solution.
 
 ## Minimal example
 
@@ -133,12 +85,6 @@ sol.xs  # (21,) states on the grid, however many internal steps were taken
 sol.ok  # False if integration or a requested output failed
 ```
 
-With no arguments, `IController()` and `PIController()` use precision-aware
-tolerances: `rtol=1e-4, atol=1e-6` for float32 states and
-`rtol=1e-7, atol=1e-9` for float64 states. Explicit values override the
-policy and are cast to the state dtype. Automatic `dt_min` is
-`10 * eps * max(1, abs(t_1))` in the time dtype.
-
 Gradients go straight through the solve:
 
 ```python
@@ -151,45 +97,31 @@ def endpoint(p):
 
 jax.grad(endpoint)(jnp.asarray(1.3))          # reverse mode
 jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),))  # forward mode
-jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])(
-    jnp.asarray(1.3)
-)  # reverse-over-forward
 ```
 
 ## Design contracts at a glance
 
 - **`dt_0` is required.** There is no initial-step heuristic.
-- **`max_steps` counts attempted internal steps**, including rejections. It
-  controls the bounded scan and only becomes an output-row count in
-  `SaveAt(steps=True)`, which returns `max_steps + 1` padded rows including
-  the initial state. Accepted steps form a contiguous prefix; rejected
-  attempts are omitted and the tail repeats the last accepted state.
-  `sol.num_steps` reports actual attempts and `sol.num_accepted` reports
-  successful advances.
+- **`max_steps` counts attempted internal steps**, including rejections.
+  `sol.num_steps` reports attempts, `sol.num_accepted` successful advances.
+- **`SaveAt` is the shape contract**: endpoint, fixed interpolation grid, or
+  padded accepted-step prefix — output shapes never depend on how many steps
+  the controller took. See [ODEs](ode.md#static-shapes-and-saveat).
 - **Fixed-step times do not depend on the attempt budget.** They are formed
-  arithmetically from the accepted-step index, with only a small local endpoint
-  snap. Increasing a nonbinding `max_steps` therefore does not change the
-  numerical method.
-- **Exact fixed-step output is explicit-ODE-only.**
-  `SaveAt(ts=grid, exact=True)` gathers internal knots without interpolation;
-  every query must align with a realized knot.
-- **Adaptive loop choice is an AD choice.** `adaptive_loop="bounded"` is the
-  reverse-mode-capable default. `adaptive_loop="forward"` executes actual
-  attempts but supports only primal, JVP, and nested forward mode.
+  arithmetically from the accepted-step index.
+- **The controller is stop-gradiented.** States differentiate through solver
+  stages on the realized, frozen mesh; see
+  [AD through adaptive stepping](ode.md#ad-through-adaptive-stepping).
 - **Forward time only**: `t_1 > t_0`.
-- **Never poisons.** `sol.ok` reports whether `t_1` was reached and every
-  requested output was valid; callers that want diverging residuals can map
-  `jnp.where(sol.ok, x, jnp.inf)` over `sol.xs`.
-- **`project`** (an idempotent clamp, e.g. positivity) is applied at every
-  point where the vector field is evaluated and to every accepted state.
-- **Never sets `jax_enable_x64`.** The time dtype follows the common state
-  dtype; float32 problems stay float32 even when x64 is enabled.
+- **Never poisons.** `sol.ok` reports failure; callers map
+  `jnp.where(sol.ok, x, jnp.inf)` when they want loud divergence.
+- **Never sets `jax_enable_x64`.** The time dtype follows the state dtype;
+  float32 problems stay float32 even when x64 is enabled.
 - Solvers, controllers, `SaveAt`, and `Solution` are frozen dataclasses
-  registered as pytrees: numeric fields (tolerances, grids, `dt_0`, `x_0`) are
-  data leaves, so changing them never recompiles.
-
-Read next: [Static Shapes](static_shapes.md) for the loop and
-`SaveAt`, [Adaptive Stepping and AD](adaptive_ad.md) for what is and is not
-differentiated, [Auxiliary Outputs](aux.md), [Rodas5P](rodas5p.md) for the SciML-derived linearly implicit
-method, [DAEs](dae.md), [SDEs](sde.md), [SDAEs](sdae.md), and the [API
-Reference](api.md).
+  registered as pytrees: numeric fields (tolerances, grids, `dt_0`, `x_0`)
+  are data leaves, so changing them never recompiles.
+
+Read next: [ODEs](ode.md), [SDEs](sde.md), [Semi-Explicit DAEs](dae.md),
+[SDAEs](sdae.md), [Markov Chains](markov_chains.md),
+[Linear Exponential Solves](exponential.md), and the
+[API Reference](api.md).
diff --git a/docs/llms.txt b/docs/llms.txt
index b2cb761..72362da 100644
--- a/docs/llms.txt
+++ b/docs/llms.txt
@@ -1,17 +1,14 @@
 # tinydiffeq
 
-> Tiny ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Fixed stepping and the default adaptive path use bounded lax.scan loops with static output shapes and forward/reverse AD, including reverse-over-forward. Adaptive Tsit5 and Rodas5P ODE/DAE solves can instead select adaptive_loop="forward", an actual-work lax.while_loop for primal, JVP, and nested forward AD; reverse mode is unsupported, and vmap runs to the slowest lane. max_steps is an attempt budget, while Solution.num_steps records actual attempts. Fixed-step times are arithmetic and budget-invariant. SaveAt(ts=..., exact=True) gathers aligned internal knots for explicit constant-step ODEs only; other requested-grid deterministic output uses Hermite or Rodas5P interpolation. Adaptive SaveAt(steps=True) AD freezes the internal mesh and omits mesh motion, an asymptotically exact residual Jacobian only near a smooth residual root. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, an embedded estimator, and its published stiff-aware fourth-order dense output. RK4/Tsit5 DAEs delegate both their stage roots and implicit derivatives to nlls-gram; the square primal defaults to normal Cholesky and the square implicit rule defaults to direct nonsymmetric LU. LMRootSolver requires residual-only stopping with gtol=xtol=0: every accepted algebraic root must report CONVERGED and have Euclidean residual norm below root atol. max_steps_is_success remains for source compatibility but never makes MAX_STEPS a valid DAE root. LMRootSolver defaults to the previous-root predictor and offers an opt-in secant predictor for locally unique branches. DAESolution reports logical time attempts, nonlinear root calls, and LM update counts; Rodas5P uses LM only for initial consistency. ODE fields and stochastic drifts may return (value, saved_aux), while DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary JAX pytrees. IController remains the adaptive default; PIController is opt-in. Finite-state DTMC/CTMC sampling is primal-only and offers sequential and associative parallel-prefix methods. Deterministic distribution forecasts are differentiable in their initial mass through matrix powers, dense exponentials, or matrix-free Arnoldi/Krylov actions. Use SciML/diffrax for general mass matrices, fully implicit or higher-index DAEs, adaptive stochastic stepping, events, continuous solution objects, or specialized adjoints.
+> Tiny differentiable ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Fixed stepping and the default adaptive path use bounded lax.scan loops with static output shapes and forward/reverse AD, including reverse-over-forward; adaptive Tsit5 and Rodas5P ODE/DAE solves can instead select adaptive_loop="forward", an actual-work lax.while_loop for primal, JVP, and nested forward AD only. max_steps is an attempt budget, Solution.num_steps records actual attempts, and fixed-step times are arithmetic and budget-invariant. SaveAt picks the endpoint, a fixed interpolation grid (exact=True gathers aligned knots for explicit constant-step ODEs), or the padded accepted-step prefix. Adaptive AD freezes the internal mesh (stop-gradiented controllers) and omits mesh motion. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, and Steinebach's stiff-aware fourth-order dense output, following SciML's OrdinaryDiffEqRosenbrock. RK4/Tsit5 DAEs delegate stage roots and implicit derivatives to nlls-gram; LMRootSolver requires residual-only stopping (gtol=xtol=0) and accepts only CONVERGED roots with Euclidean residual norm below the root atol, with previous-root or opt-in secant stage predictors. solve_sde integrates diagonal-noise Ito SDEs with EulerMaruyama (strong order 0.5), Milstein (1.0, commutative diagonal), or SRA1 (1.5, additive); each solver declares its per-step noise via sample_noise, drawn from a key or passed explicitly through noise= as validated, differentiable data. solve_semi_explicit_sdae applies EulerMaruyama or SRA1 to the reduced index-1 stochastic system with root-restored consistency; SaveAt(ts=...) raises for SDEs/SDAEs because interpolating rough paths is wrong. Fields and stochastic drifts may return (value, saved_aux); DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary pytrees of same-dtype real floating arrays, and the library never sets jax_enable_x64. Finite-state DTMC/CTMC sampling is primal-only; deterministic distribution forecasts are differentiable in their initial mass through matrix powers, dense exponentials, or matrix-free Arnoldi/Krylov actions. Use SciML/diffrax for general mass matrices, fully implicit or higher-index DAEs, adaptive stochastic stepping, events, continuous solution objects, or specialized adjoints.
 
 ## Docs
 
-- [Home — positioning, vector-field signature convention f(x, t, args, p), minimal examples](https://highdimensionaleconlab.github.io/tinydiffeq/)
-- [Static shapes — bounded and actual-work loops, budget-invariant fixed times, SaveAt modes, work counters](https://highdimensionaleconlab.github.io/tinydiffeq/static_shapes/)
-- [Adaptive stepping and AD — frozen-mesh derivative contract, loop-mode AD boundary, controller and horizon logic](https://highdimensionaleconlab.github.io/tinydiffeq/adaptive_ad/)
-- [Auxiliary outputs — saved versus algebraic aux, interpolation, AD, and failure behavior](https://highdimensionaleconlab.github.io/tinydiffeq/aux/)
-- [Rodas5P — Steinebach method, direct SciML implementation links, stiff ODE/DAE formulation, dense output, and AD](https://highdimensionaleconlab.github.io/tinydiffeq/rodas5p/)
-- [Semi-explicit DAEs — index-1 contract, algebraic LM roots and predictors, counters, implicit AD, failure behavior, and SaveAt](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
-- [SDEs — Euler-Maruyama orders, fixed-noise key semantics, shared-path strong-convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/)
-- [Semi-explicit SDAEs — reduced-SDE Euler-Maruyama, algebraic roots, aux, convergence assumptions, and pathwise AD](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/)
+- [Home — positioning, install, vector-field signature convention f(x, t, args, p), minimal example, design contracts](https://highdimensionaleconlab.github.io/tinydiffeq/)
+- [ODEs — solver table, fixed vs adaptive control, static shapes and SaveAt, AD through adaptive stepping, aux outputs, stiff Rodas5P with SciML credit](https://highdimensionaleconlab.github.io/tinydiffeq/ode/)
+- [SDEs — EulerMaruyama/Milstein/SRA1 orders and noise contracts, explicit noise= as differentiable data, fixed-key semantics, shared-path convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/)
+- [Semi-explicit DAEs — index-1 contract, nlls-gram LM roots and predictors, implicit AD, aux contracts, Rodas5P mass-matrix path, dense output, failure behavior](https://highdimensionaleconlab.github.io/tinydiffeq/dae/)
+- [Semi-explicit SDAEs — reduced-SDE EulerMaruyama and SRA1, algebraic roots, aux, convergence assumptions, pathwise AD](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/)
 - [Finite-state Markov chains — sampling, deterministic PMF forecasts, matrix-free Krylov CTMC actions, pytrees, vmap, and AD scope](https://highdimensionaleconlab.github.io/tinydiffeq/markov_chains/)
 - [Linear exponential solves — dense expm and matrix-free Arnoldi/Krylov actions for fixed linear array or pytree operators, with traced and hand-coded initial-state JVP/VJP](https://highdimensionaleconlab.github.io/tinydiffeq/exponential/)
-- [API reference — ODE/SDE/DAE solves, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/)
+- [API reference — solve functions, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/)
diff --git a/docs/ode.md b/docs/ode.md
new file mode 100644
index 0000000..38e54c9
--- /dev/null
+++ b/docs/ode.md
@@ -0,0 +1,187 @@
+# ODEs
+
+`solve_ode` integrates `dx/dt = f(x, t, args, p)` from `t_0` to `t_1 > t_0`.
+The state is an array or pytree whose leaves share one real floating dtype;
+`args` is pass-through data (not an AD target by convention) and `p` holds
+differentiable parameters. `dt_0` is required — there is no initial-step
+heuristic.
+
+## Solvers
+
+| Solver | Order | Stepping |
+|---|---|---|
+| `Euler()` | 1 | fixed only |
+| `RK4()` | 4 | fixed only |
+| `Tsit5()` | 5(4) | fixed or adaptive (embedded estimate, FSAL) |
+| `Rodas5P()` | 5(4) | fixed or adaptive, linearly implicit (stiff) |
+
+Fixed stepping uses `ConstantStepSize()` (the default controller):
+`dt_0 = (t_1 - t_0)/n` with `max_steps = n` reproduces a uniform grid
+exactly. Times are formed arithmetically as `t_0 + i * dt_0` with a small
+endpoint snap that never scales with `max_steps`, so increasing a nonbinding
+attempt budget cannot change the numerical method. `unroll=` (a static int,
+fixed stepping only) unrolls that many steps per scan iteration — identical
+values, fewer GPU dispatches, more compile time; on small-batch
+neural-network fields it cut reverse-mode solve time 2–3× on an L40S (see
+`benchmarks/results/`).
+
+Adaptive stepping uses `IController()` or `PIController()` with an embedded
+error estimate. Omitted tolerances follow the state precision:
+
+| State dtype | `rtol` | `atol` | unit-scale `dt_min` |
+|---|---:|---:|---:|
+| float32 | `1e-4` | `1e-6` | `10 * eps` ≈ `1.19e-6` |
+| float64 | `1e-7` | `1e-9` | `10 * eps` ≈ `2.22e-15` |
+
+Explicit values override the policy and are cast to the state dtype; use them
+whenever tolerances are part of a reproducibility contract. Every attempt is
+clipped to the remaining horizon, and the controller's next-step proposal is
+computed from the clipped step — the clip doubles as the growth guard that
+keeps near-flat fields from growing steps into quarter-horizon leaps.
+`PIController` additionally damps step-size oscillations through the previous
+accepted error ratio; `p_coeff=0, i_coeff=1` reproduces `IController` exactly.
+
+## Static shapes and `SaveAt`
+
+`max_steps` bounds attempted internal steps, including rejections. The
+default `adaptive_loop="bounded"` runs a bounded `lax.scan` with exactly
+`max_steps` attempt slots grouped into chunks that are skipped once the solve
+completes; shapes are static, and changing tolerances, curvature, or initial
+conditions never recompiles (pinned by `tests/test_recompile.py`).
+`adaptive_loop="forward"` instead runs a dynamic `lax.while_loop` that
+executes only actual attempts; it supports primal evaluation, JVP, and nested
+forward mode, but JAX cannot transpose it, so reverse mode requires the
+bounded loop. Under `vmap`, both run until the slowest lane finishes.
+
+Exactly one `SaveAt` mode is set:
+
+- **`SaveAt(t_1=True)`** (default): the endpoint only.
+- **`SaveAt(ts=grid)`**: dense interpolation onto a fixed query grid. Output
+  shape is `(len(grid), ...)` regardless of how many internal steps the
+  controller takes; `ts` is a data leaf, so a new grid of the same length
+  retraces nothing. Explicit methods use cubic Hermite (4th-order accurate
+  between 5th-order knots); Rodas5P uses its stiff-aware fourth-order
+  continuous extension. Queries are observation times, not internal stops.
+  For an explicit constant-step solve, `exact=True` instead requires every
+  query to coincide with a realized knot and gathers those states directly,
+  with no interpolation work.
+- **`SaveAt(steps=True)`**: the initial state plus accepted steps as a
+  contiguous prefix of a `max_steps + 1` buffer. Rejected attempts are
+  omitted; `fill="last"` (default) repeats the last valid row through the
+  tail, `fill="inf"` fills it with `inf`, and `sol.accepted` masks the valid
+  prefix.
+
+If the budget runs out before `t_1`, `sol.ok` is `False` and the outputs hold
+the reached prefix. Nothing is poisoned; callers that want diverging
+residuals map `jnp.where(sol.ok, x, jnp.inf)`. `sol.num_steps` counts
+attempts, `sol.num_accepted` counts advances.
+
+`project` (an idempotent clamp, e.g. positivity) is applied at every point
+where the field is evaluated and to every accepted state.
+
+## AD through adaptive stepping
+
+The bounded loop supports `jax.grad`, `jax.jvp`, `vmap`, and
+reverse-over-forward through one adaptive solve, verified against closed
+forms in `tests/test_ad.py`.
+
+The step-size controllers compute their error norms, decisions, and
+next-step factors inside `stop_gradient`: accept/reject is a discrete branch,
+and `E**(-1/order)` is singular at the exact-zero error of a flat-start
+policy — without the stop-gradient and the machine-epsilon error floor the
+backward pass is NaN from iteration one.
+
+States therefore differentiate through the solver stages on a **frozen
+mesh**: the derivative holds realized step sizes and accept/reject patterns
+fixed and omits the mesh-motion term `(dq/dt) * dt_k/dtheta`. Near a smooth
+residual root that term is `O(‖r‖)`, so a frozen-mesh residual Jacobian is
+asymptotically exact there; far from a root it can be material. Endpoint and
+fixed-requested-grid outputs do not have a moving output-time axis; adaptive
+`SaveAt(steps=True)` times have exactly zero tangent.
+
+`SaveAt(ts=...)` brackets queries with `searchsorted` — integer indices, no
+gradient. Values differentiate through the bracketing states and slopes;
+zero-width brackets are handled with the double-`where` trick so neither the
+primal nor its transpose evaluates `0/0`.
+
+When several tangent or cotangent directions are needed at one primal point,
+cache the linearization (`jax.linearize` / `jax.vjp`) and `vmap` the
+pushforward or pullback instead of repeating the primal trajectory in a fused
+`vmap(jvp)`.
+
+On float32 GPUs, XLA serves `dot_general` from TF32 tensor cores by default
+(~1e-3 precision). When adaptive accept/reject decisions depend on
+neural-network outputs, that roundoff can move the mesh; set
+`jax_default_matmul_precision="highest"` in reproducibility entry points.
+tinydiffeq never changes application-global JAX settings, including
+`jax_enable_x64` — a float32 problem stays float32 even with x64 enabled.
+
+## Auxiliary outputs
+
+The field may return `(dx, aux)` to save quantities it already computes
+without adding them to the state:
+
+```python
+def f(x, t, args, p):
+    dx = -p * x
+    return dx, {"flow": dx, "moment": x**2}
+```
+
+Saved aux is a nonempty pytree of real floating arrays (leaves may differ in
+dtype). It follows `SaveAt` — stored at the endpoint or accepted nodes, or
+interpolated onto a requested grid with endpoint slopes obtained by JVP — and
+participates in JVP/VJP. `has_aux=None` (default) detects the contract with
+one `jax.eval_shape` trace; `has_aux=False` skips detection and selects the
+minimal value-only path. Every saved aux leaf must remain finite: an invalid
+value in a prefix-saving mode terminates at the previous accepted node, while
+endpoint mode keeps the endpoint state, zero-fills aux, and sets `ok=False`.
+For batched differentiation where inactive lanes may leave the model domain,
+`failure_ad_reference=(x, t, p)` supplies a finite point used only for safe
+linearization.
+
+## Stiff ODEs: Rodas5P
+
+`Rodas5P()` is an eight-stage, fifth-order Rosenbrock–Wanner method with an
+embedded estimator and a stiff-aware fourth-order continuous extension. It is
+A-stable and stiffly accurate.
+
+```python
+import jax.numpy as jnp
+
+from tinydiffeq import IController, Rodas5P, solve_ode
+
+
+def stiff_field(x, t):
+    # Exact solution x(t) = cos(t), with a fast transient eigenvalue -1000.
+    return -1000.0 * (x - jnp.cos(t)) - jnp.sin(t)
+
+
+sol = solve_ode(
+    stiff_field, Rodas5P(), 0.0, 1.0, jnp.asarray(1.0),
+    dt_0=0.01, controller=IController(), max_steps=512,
+)
+```
+
+At each attempted step it forms the exact JAX Jacobian $J = \partial f /
+\partial x$, factors $W = I/(\gamma h) - J$ once with pivoted LU, and reuses
+the factors for all eight stage solves — linear solves only, no Newton
+iteration. JVP/VJP propagate through Jacobian construction and the stage
+solves via the linear-solve rule $dx = A^{-1}(db - dA\,x)$; pivot selection
+is not differentiated. Dense Jacobians and dense LU only — use SciML or
+diffrax for sparse/Krylov stages.
+
+### Credit
+
+Rodas5P was constructed by Gerd Steinebach and introduced in
+[Steinebach (2023)](https://doi.org/10.1007/s10543-023-00967-x). tinydiffeq's
+JAX implementation deliberately follows SciML's authoritative MIT-licensed
+implementation — the
+[`OrdinaryDiffEqRosenbrock` package](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock),
+its
+[tableau and dense coefficients](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_tableaus.jl),
+[step implementation](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_perform_step.jl),
+and
+[dense interpolants](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_interpolants.jl).
+The tableau, stage equations, embedded estimate, and continuous extension are
+ported rather than redesigned, and regression tests compare fixed steps
+against SciML-produced values to prevent silent divergence.
diff --git a/docs/sdae.md b/docs/sdae.md
index fd65636..bd61f64 100644
--- a/docs/sdae.md
+++ b/docs/sdae.md
@@ -7,9 +7,9 @@ dy = f(y,z,t)\,dt + \sigma(y,z,t)\,dW,
 \qquad 0=g(y,z,t),
 $$
 
-with fixed-step Euler–Maruyama and diagonal noise in the differential state.
-The algebraic Jacobian $g_z$ must be square and nonsingular along the chosen
-root branch. Arrays and pytrees are supported for both `y` and `z`.
+with fixed steps and diagonal noise in the differential state. The algebraic
+Jacobian $g_z$ must be square and nonsingular along the chosen root branch.
+Arrays and pytrees are supported for both `y` and `z`.
 
 ```python
 import jax
@@ -54,70 +54,73 @@ sol = solve_semi_explicit_sdae(
 
 ## Algorithm and convergence
 
-At a consistent node, the update is
+`EulerMaruyama` and `SRA1` are supported. At a consistent node,
+Euler–Maruyama updates
 
 $$
 y_{n+1}=y_n+f(y_n,z_n,t_n)h+\sigma(y_n,z_n,t_n)\Delta W_n,
 $$
 
-followed by a root solve for
-$g(y_{n+1},z_{n+1},t_{n+1})=0$. Locally writing the unique root as
-$z=Z(y,t)$ shows that this is exactly Euler–Maruyama on the reduced SDE. No
-Itô correction is missing: `z` is reconstructed from the constraint rather
-than advanced by a separate SDE.
-
-The method has strong order 0.5 and weak order 1 when $g_z$ is uniformly
-nonsingular near the path, the reduced drift and diffusion are globally
-Lipschitz with linear growth, and the usual additional smoothness assumptions
-for weak convergence hold. Additive noise gives strong order 1. A per-node
-root error $\epsilon$ should be $O(\sqrt h)$ or smaller to preserve strong
-order 0.5, and $O(h)$ or smaller to preserve weak order 1; normal root
-tolerances are typically much tighter.
+followed by a root solve for $g(y_{n+1},z_{n+1},t_{n+1})=0$. Locally writing
+the unique root as $z=Z(y,t)$ shows this is exactly Euler–Maruyama on the
+reduced SDE — no Itô correction is missing, since `z` is reconstructed from
+the constraint rather than advanced by a separate SDE. It has strong order
+0.5 and weak order 1 under the usual regularity assumptions on the reduced
+system (strong order 1 for additive noise). A per-node root error
+$\epsilon$ should be $O(\sqrt h)$ or smaller to preserve strong order 0.5;
+normal root tolerances are far tighter.
+
+`SRA1` applies the Rößler additive-noise stochastic Runge–Kutta scheme to
+the same reduced SDE: each step evaluates the drift at the node and at one
+internal stage, with a root solve restoring consistency at the stage time
+$t_n + \tfrac34 h$ and at the endpoint — two root solves per step. Its
+strong order 1.5 requires the diffusion to be independent of the state; in
+the SDAE setting, `z` depends on `y` through the constraint, so the
+requirement is that `diffusion(y, z, t)` depends only on `t`. As in
+`solve_sde`, the contract is documented, not runtime-checked, and the
+endpoint diffusion evaluation reuses the current node's `(y, z, context)`
+with the endpoint time. `Milstein` is deliberately not wired in: its
+commutativity condition and derivative correction do not translate cleanly
+through the implicit reduction $z = Z(y, t)$.
 
 ## Randomness, aux, and AD
 
-The Brownian increments use the same sampling contract as `solve_sde`.
-Arrays use the draw shape `(n_steps,) + y_0.shape`; pytrees use one flat draw
-partitioned in deterministic leaf order. A fixed key defines a fixed path,
-so JVP/VJP with respect to `y_0` and `p` are pathwise derivatives under common
-random numbers. The key is not differentiable.
-
-`z_0` is a root guess, receives zero tangent, and selects a local root branch.
-Algebraic solves use the same `LMRootSolver` configuration and
-implicit-AD contract as deterministic DAEs; see
-[Nonlinear-solve and AD contract](dae.md#nonlinear-solve-and-ad-contract).
-nlls-gram supplies both the primal LM root and its implicit derivative, using
-direct `LU()` by default for the square derivative. `gtol` and `xtol` must
-remain zero, and every accepted iterate must report `CONVERGED` with Euclidean
-residual norm strictly below the root `atol`. `max_steps_is_success` is retained
-for source compatibility but does not make `MAX_STEPS` valid.
+Noise follows the `solve_sde` sampling contract via
+`solver.sample_noise(y_0, key, n_steps, dt, dtype)`: Brownian increments
+for `EulerMaruyama`, an independent $(\Delta W, \Delta Z)$ pair for `SRA1`.
+Arrays use the draw shape `(n_steps,) + y_0.shape`; pytrees use one flat
+draw partitioned in deterministic leaf order. A fixed key defines a fixed
+path, so JVP/VJP with respect to `y_0` and `p` are pathwise derivatives
+under common random numbers. The key is not differentiable.
+
+`z_0` is a root guess, receives zero tangent, and selects a local root
+branch. Algebraic solves use the same `LMRootSolver` configuration and
+implicit-AD contract as deterministic DAEs — see
+[Nonlinear-solve and AD contract](dae.md#nonlinear-solve-and-ad-contract):
+residual-only stopping, `CONVERGED` with residual norm strictly below the
+root `atol`, and nlls-gram's direct square `LU()` implicit derivative.
+
 The algebraic function may return `(residual, algebraic_aux)`; that internal
-context is passed to both drift and diffusion but is not stored. The drift may
-return `(drift_value, saved_aux)`, and only that saved aux becomes `sol.aux`.
-Steps mode stores it at every consistent node; endpoint mode evaluates it only
-at the final node. Its derivatives include both direct parameter dependence
-and dependence through the implicit root. See
-[Auxiliary Outputs](aux.md) for all contract variations and explicit flags.
-Invalid algebraic context at initialization sets `ok=False` before any
-stochastic step. In steps mode, invalid saved aux terminates at the previous
-consistent node. In endpoint mode, invalid final saved aux retains the
-endpoint state, returns zero aux, and sets `ok=False`.
-
-`sol.num_root_solves` counts logical active root calls, including the initial
-consistency solve and failures, while `sol.num_root_steps` sums their LM update
-counts. Both are path diagnostics with exact-zero tangents.
-
-Only `SaveAt(t_1=True)` and `SaveAt(steps=True)` are supported. Stochastic
+context is passed to both drift and diffusion but is not stored. The drift
+may return `(drift_value, saved_aux)`, and only that saved aux becomes
+`sol.aux`: steps mode stores it at every consistent node, endpoint mode
+evaluates it only at the final node. Its derivatives include both direct
+parameter dependence and dependence through the implicit root. Invalid
+algebraic context at initialization sets `ok=False` before any stochastic
+step; invalid saved aux terminates at the previous consistent node (steps
+mode) or zero-fills aux at the retained endpoint (endpoint mode).
+
+`sol.num_root_solves` counts logical active root calls — for `SRA1`, a
+stage and an endpoint root per step plus the initial consistency solve —
+and `sol.num_root_steps` sums their LM updates; both have exact-zero
+tangents.
+
+Only `SaveAt(t_1=True)` and `SaveAt(steps=True)` are supported: stochastic
 paths are rough, so deterministic dense interpolation between nodes would be
-mathematically wrong. A failed root freezes the last consistent prefix,
-sets `ok=False`, and pads the remaining static buffer. Failed roots have zero
-implicit tangents, and aux at a failed initial root is zero-filled, so masked
-lanes can preserve successful JVPs or VJPs under `vmap`. For that contract,
-pass `failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)` at a point where the
-residual, context, and saved-aux maps are finite and differentiable. The
-reference is substituted before nlls only for already-inactive lanes and is
-also used by inactive context, drift, diffusion, and saved-aux evaluation. A
-newly attempted root must be JVP-safe at its actual initial point; the
-reference cannot repair an invalid active attempt after it fails. Without one,
-an all-ones best-effort default is used for inactive work. Outputs and
-gradients from a failed lane are not a valid solution.
+mathematically wrong. A failed root freezes the last consistent prefix, sets
+`ok=False`, and pads the remaining static buffer with zero implicit
+tangents, so masked lanes preserve successful JVPs/VJPs under `vmap`. For
+that contract, pass `failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)` at a
+point where the residual, context, and saved-aux maps are finite and
+differentiable — the same semantics as the
+[DAE reference point](dae.md#nonlinear-solve-and-ad-contract).
diff --git a/docs/sde.md b/docs/sde.md
index 85c8303..b0b9c73 100644
--- a/docs/sde.md
+++ b/docs/sde.md
@@ -6,16 +6,18 @@ $$
 dX_t = \mu(X_t, t)\,dt + \sigma(X_t, t)\,dW_t
 $$
 
-with fixed-step Euler–Maruyama and **diagonal noise**: `diffusion` returns the
-same pytree structure as the state, multiplied leafwise and elementwise by
-independent Brownian increments.
+with fixed steps and **diagonal noise**: `diffusion` returns the same pytree
+structure as the state, multiplied leafwise and elementwise by independent
+Brownian increments. `n_steps` must be a static Python int; adaptive SDE
+stepping requires Brownian-bridge noise that can be re-evaluated on
+subdivided intervals, which is diffrax territory.
 
 ```python
-from tinydiffeq import solve_sde, EulerMaruyama, SaveAt
+from tinydiffeq import solve_sde, SRA1, SaveAt
 
 sol = solve_sde(
-    drift, diffusion, EulerMaruyama(), 0.0, 1.0, x_0,
-    key=jax.random.PRNGKey(0),
+    drift, diffusion, SRA1(), 0.0, 1.0, x_0,
+    key=jax.random.key(0),
     n_steps=256,
     p=(mu, sigma),
     save_at=SaveAt(steps=True),
@@ -24,17 +26,47 @@ sol = solve_sde(
 
 `drift` and `diffusion` follow the same `(x)`, `(x, t)`, `(x, t, args)`,
 `(x, t, args, p)` signature convention as `solve_ode`.
-The drift may return `(drift_value, saved_aux)`; diffusion remains value-only.
-Aux is stored at the same endpoint or fixed grid nodes and participates in
-pathwise JVP/VJP. See [Auxiliary Outputs](aux.md).
 
-## The static-shape contract, honestly
+## Solvers
 
-`n_steps` must be a static Python int. There is no adaptive SDE stepping in
-the current release — adaptivity for SDEs requires noise that can be *re-evaluated* on
-subdivided intervals (a Brownian-bridge / VirtualBrownianTree construction),
-so that a rejected step resamples consistently. That is diffrax territory
-for now; a roadmap issue sketches what it would take here.
+Three fixed-step schemes, each a frozen dataclass:
+
+| Solver | Strong order | Noise contract |
+|---|---|---|
+| `EulerMaruyama()` | 0.5 | any diagonal diffusion |
+| `Milstein()` | 1.0 | diagonal *commutative*: each diffusion component depends only on its own state component |
+| `SRA1()` | 1.5 | *additive*: diffusion independent of the state (may depend on time) |
+
+`Milstein` adds the correction $\tfrac12 g g' (\Delta W^2 - \Delta t)$, with
+$g g'$ computed as the forward-mode derivative of the diffusion field in the
+direction of its own value. `SRA1` is the Rößler additive-noise stochastic
+Runge–Kutta scheme: two drift stages per step plus the time-Wiener integral
+$I_{10}$. The contracts are documented, not runtime-checked: a
+state-dependent diffusion under `SRA1` silently degrades the order, exactly
+as in other SDE libraries.
+
+Each solver declares its own per-step randomness through
+`solver.sample_noise(x_0, key, n_steps, dt, dtype)`: Brownian increments
+$\Delta W$ for `EulerMaruyama`/`Milstein`, an independent pair
+$(\Delta W, \Delta Z)$ for `SRA1` (from which $I_{10}/\Delta t =
+\tfrac12(\Delta W + \Delta Z/\sqrt 3)$ is formed internally).
+
+## Explicit noise: the shocks as data
+
+Pass `noise=` instead of `key=` (exactly one is required) to hand `solve_sde`
+the realization directly — the same pytree `sample_noise` would produce,
+validated against the solver's spec:
+
+```python
+noise = SRA1().sample_noise(x_0, key, n_steps, dt, x_0.dtype)
+sol = solve_sde(drift, diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n_steps)
+```
+
+This is bit-identical to the `key=` call that would have drawn the same
+`noise`, and it makes the noise **first-class data**: inspectable, storable
+(e.g. as part of a training set alongside the initial condition), and
+differentiable — `jax.grad` with respect to the `noise` pytree works, since
+the steppers consume it as ordinary arrays.
 
 ## Key semantics: a fixed noise process
 
@@ -47,44 +79,61 @@ d_w = jnp.sqrt(dt) * jax.random.normal(key, (n_steps,) + x_0.shape)
 so a fixed `key` pins the entire noise path:
 
 - **Reproducible** — the same key gives the same path, every call.
-- **Differentiable with respect to `x_0` and `p`** (not `key`): with the path
-  held fixed, the solution map is smooth, and jvp/vjp against finite
-  differences are tested in `tests/test_sde.py`. This is exactly the "common
-  random numbers" setup simulation-based estimators want.
+- **Differentiable with respect to `x_0` and `p`** (not `key`), and with
+  respect to `noise` when passed explicitly: with the path held fixed, the
+  solution map is smooth, and jvp/vjp against finite differences are tested
+  in `tests/test_sde.py`. This is exactly the common-random-numbers setup
+  simulation-based estimators want.
 
 For a pytree, tinydiffeq draws one `(n_steps, total_state_size)` array and
-partitions it into leaves in JAX's deterministic leaf order. This makes the
-noise identical to that of one equivalent flat array without flattening the
-state during integration. Changing the tree structure changes the assignment
-of random components and requires a new compilation.
+partitions it into leaves in JAX's deterministic leaf order, so the noise is
+identical to that of the equivalent flat array. Changing the tree structure
+changes the assignment of random components and requires a new compilation.
 
 ## Orders of convergence
 
-Euler–Maruyama is strong order 0.5 (pathwise) and weak order 1.0 (in
-distribution) for multiplicative noise. The test suite verifies the strong
-rate on geometric Brownian motion, where the exact solution driven by the
-*same* increments is available in closed form:
+The test suite verifies the strong rates on shared noise paths: because the
+increments are exactly reproducible from the key, the exact solution (for
+geometric Brownian motion, $X_T = X_0 \exp((\mu - \sigma^2/2)T + \sigma
+W_T)$ with $W_T = \sum_k \Delta W_k$) is evaluated on the *same* path as the
+numerical endpoint. Comparing against an independently sampled exact solution
+would measure nothing.
 
-$$
-X_T = X_0 \exp\left((\mu - \tfrac{\sigma^2}{2})T + \sigma W_T\right),
-\qquad W_T = \textstyle\sum_k \Delta W_k .
-$$
+## Auxiliary output
+
+Only the drift owns saved aux:
 
-Because the increments `d_w` are exactly reproducible from the key, the test
-regenerates them, computes the exact endpoint on the same path, and checks
-the mean absolute error slope across `dt` levels lands in `[0.35, 0.65]`.
-That shared-path construction is the right way to measure strong error —
-comparing against an independently sampled exact solution measures nothing.
+```python
+def drift(x, t, args, p):
+    return drift_value, saved_aux
+
+
+def diffusion(x, t, args, p):
+    return diffusion_value
+```
+
+Aux is stored at the endpoint or at the fixed grid nodes and is
+differentiated pathwise under the fixed noise realization. `has_aux=None`
+auto-detects the form; `has_aux=False` selects the no-aux scan without the
+detection trace.
 
 ## SaveAt for SDEs
 
 `SaveAt(t_1=True)` (default) and `SaveAt(steps=True)` (`n_steps + 1` rows,
 or a padded accepted prefix after an aux failure) are supported.
-`SaveAt(ts=...)` **raises**: cubic Hermite
-interpolation assumes smooth trajectories and is simply wrong between the
-points of a rough path. Land your grid on the step boundaries instead by
-choosing `n_steps`.
+`SaveAt(ts=...)` **raises**: cubic Hermite interpolation assumes smooth
+trajectories and is simply wrong between the points of a rough path. Land
+your grid on the step boundaries instead by choosing `n_steps`.
 
 `project` is applied to every drift/diffusion evaluation point and every
 accepted state, as in `solve_ode`. With `diffusion ≡ 0`, `solve_sde`
 reproduces `solve_ode` with `Euler()` on the same grid exactly.
+
+## Performance: `unroll`
+
+`unroll=` (a static int, default 1) unrolls that many steps per iteration of
+the integration scan — identical values and gradients, fewer and larger GPU
+dispatches, more compile time. For small-batch neural-network-drift
+ensembles on an L40S, `unroll=8` cut reverse-mode solve time 2–3× and primal
+time 1.2–1.7× (see `benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md`).
+The same argument exists on `solve_ode` for fixed stepping.
diff --git a/mkdocs.yml b/mkdocs.yml
index ab6685e..11db43a 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -43,13 +43,10 @@ extra_javascript:
 
 nav:
   - Home: index.md
-  - Static Shapes: static_shapes.md
-  - Adaptive Stepping and AD: adaptive_ad.md
-  - Auxiliary Outputs: aux.md
-  - Rodas5P: rodas5p.md
+  - ODEs: ode.md
+  - SDEs: sde.md
   - Semi-Explicit DAEs: dae.md
   - Semi-Explicit SDAEs: sdae.md
-  - SDEs: sde.md
   - Markov Chains: markov_chains.md
   - Linear Exponential Solves: exponential.md
   - API Reference: api.md

From c0719a0dc5bac9d0ac178c8000f0ef6640f311a2 Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Tue, 4 Aug 2026 17:39:54 -0700
Subject: [PATCH 3/5] bench: GPU ensembles with trained kernels policies

gpu_trajectories.py benchmarks vmapped per-trajectory solves (SRA1/EM/
RK4, primal and jit(grad) w.r.t. x_0/p/noise) with remat, unroll,
matmul-precision, and noise-mode knobs. export_growth_policies.py
trains the kernels growth policies and exports their weights; the
neoclassical and stochastic-growth drifts evaluate those trained MLPs.
Vulcan L40S results and summary under results/.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01C8SWZyLFpzVjEX5WCRLJnD
---
 benchmarks/README.md                          |    5 +
 benchmarks/export_growth_policies.py          |  126 +
 benchmarks/gpu_trajectories.py                |  590 ++++
 benchmarks/policies/neoclassical_growth.json  |   21 +
 benchmarks/policies/neoclassical_growth.npz   |  Bin 0 -> 19734 bytes
 benchmarks/policies/stochastic_growth.json    |   53 +
 benchmarks/policies/stochastic_growth.npz     |  Bin 0 -> 19990 bytes
 .../2026-08-04_vulcan-l40s-sde-fixed.md       |   69 +
 .../results/vulcan-l40s-a-baseline.json       | 1296 +++++++
 benchmarks/results/vulcan-l40s-a-baseline.md  |   87 +
 .../results/vulcan-l40s-a-kernels-scale.json  | 3084 +++++++++++++++++
 .../results/vulcan-l40s-a-kernels-scale.md    |  198 ++
 .../results/vulcan-l40s-b-primal-scale.json   | 1164 +++++++
 .../results/vulcan-l40s-b-primal-scale.md     |   78 +
 benchmarks/results/vulcan-l40s-c-unroll4.json |  780 +++++
 benchmarks/results/vulcan-l40s-c-unroll4.md   |   54 +
 benchmarks/results/vulcan-l40s-c-unroll8.json |  780 +++++
 benchmarks/results/vulcan-l40s-c-unroll8.md   |   54 +
 .../results/vulcan-l40s-d-noisemode.json      |  268 ++
 benchmarks/results/vulcan-l40s-d-noisemode.md |   22 +
 .../results/vulcan-l40s-e-matmul-highest.json |  140 +
 .../results/vulcan-l40s-e-matmul-highest.md   |   14 +
 .../results/vulcan-l40s-f-cudagraph.json      |  396 +++
 benchmarks/results/vulcan-l40s-f-cudagraph.md |   30 +
 .../vulcan-l40s-g-unroll8-promoted.json       |  396 +++
 .../results/vulcan-l40s-g-unroll8-promoted.md |   30 +
 26 files changed, 9735 insertions(+)
 create mode 100644 benchmarks/export_growth_policies.py
 create mode 100644 benchmarks/gpu_trajectories.py
 create mode 100644 benchmarks/policies/neoclassical_growth.json
 create mode 100644 benchmarks/policies/neoclassical_growth.npz
 create mode 100644 benchmarks/policies/stochastic_growth.json
 create mode 100644 benchmarks/policies/stochastic_growth.npz
 create mode 100644 benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md
 create mode 100644 benchmarks/results/vulcan-l40s-a-baseline.json
 create mode 100644 benchmarks/results/vulcan-l40s-a-baseline.md
 create mode 100644 benchmarks/results/vulcan-l40s-a-kernels-scale.json
 create mode 100644 benchmarks/results/vulcan-l40s-a-kernels-scale.md
 create mode 100644 benchmarks/results/vulcan-l40s-b-primal-scale.json
 create mode 100644 benchmarks/results/vulcan-l40s-b-primal-scale.md
 create mode 100644 benchmarks/results/vulcan-l40s-c-unroll4.json
 create mode 100644 benchmarks/results/vulcan-l40s-c-unroll4.md
 create mode 100644 benchmarks/results/vulcan-l40s-c-unroll8.json
 create mode 100644 benchmarks/results/vulcan-l40s-c-unroll8.md
 create mode 100644 benchmarks/results/vulcan-l40s-d-noisemode.json
 create mode 100644 benchmarks/results/vulcan-l40s-d-noisemode.md
 create mode 100644 benchmarks/results/vulcan-l40s-e-matmul-highest.json
 create mode 100644 benchmarks/results/vulcan-l40s-e-matmul-highest.md
 create mode 100644 benchmarks/results/vulcan-l40s-f-cudagraph.json
 create mode 100644 benchmarks/results/vulcan-l40s-f-cudagraph.md
 create mode 100644 benchmarks/results/vulcan-l40s-g-unroll8-promoted.json
 create mode 100644 benchmarks/results/vulcan-l40s-g-unroll8-promoted.md

diff --git a/benchmarks/README.md b/benchmarks/README.md
index 670ecd4..1f887d8 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -1,5 +1,10 @@
 # Benchmarks
 
+`gpu_trajectories.py` measures vmapped trajectory ensembles (SRA1/EM SDEs and
+fixed RK4, primal and reverse-mode, explicit or key-drawn noise) and is meant
+for GPU nodes; run `python -m benchmarks.gpu_trajectories --help` for the
+sweep axes. Committed results live under `results/`.
+
 Run the opt-in CPU suite with:
 
 ```bash
diff --git a/benchmarks/export_growth_policies.py b/benchmarks/export_growth_policies.py
new file mode 100644
index 0000000..c894975
--- /dev/null
+++ b/benchmarks/export_growth_policies.py
@@ -0,0 +1,126 @@
+"""Train the kernels growth policies and export their weights for benchmarks.
+
+Runs the kernels package's neoclassical and stochastic growth solvers
+in-process (the package keeps trained nnx modules in memory only), extracts
+the MLP weights and structural parameters as plain arrays, verifies a pure-jnp
+reimplementation of the policy forward pass against the nnx module, and writes
+`benchmarks/policies/<model>.npz` consumed by `gpu_trajectories.py`.
+
+Run from the kernels venv (it owns flax and the model modules):
+
+    uv run --project /Users/jlperla/GitHub/kernels python \
+        benchmarks/export_growth_policies.py
+"""
+
+import dataclasses
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, "/Users/jlperla/GitHub/kernels")
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+OUTPUT_DIR = Path(__file__).parent / "policies"
+
+
+def unit_softplus(raw):
+    log_two = jnp.log(jnp.asarray(2.0, dtype=raw.dtype))
+    return jax.nn.softplus(2.0 * log_two * raw) / log_two
+
+
+def extract(policy, structural):
+    weights = {}
+    for index, layer in enumerate(policy.layers):
+        weights[f"layer{index}_kernel"] = np.asarray(layer.kernel.value)
+        weights[f"layer{index}_bias"] = np.asarray(layer.bias.value)
+    weights["head_kernel"] = np.asarray(policy.head.kernel.value)
+    weights["head_bias"] = np.asarray(policy.head.bias.value)
+    metadata = {
+        "num_layers": len(policy.layers),
+        "homothetic": bool(policy.homothetic),
+        "log_z_scale": float(getattr(policy, "log_z_scale", 0.0)),
+        "log_k_scale": float(getattr(policy, "log_k_scale", 0.0)),
+        "k_scale": float(getattr(policy, "k_scale", 1.0)),
+        "structural": {
+            key: (
+                np.asarray(value).item()
+                if np.asarray(value).ndim == 0
+                else np.asarray(value).tolist()
+            )
+            for key, value in dataclasses.asdict(structural).items()
+        },
+    }
+    return weights, metadata
+
+
+def pure_policy(weights, metadata, gelu_approximate):
+    num_layers = metadata["num_layers"]
+    delta = metadata["structural"]["delta"]
+
+    def forward(log_state):
+        log_z, log_k = log_state
+        if metadata["homothetic"]:
+            h = jnp.atleast_1d(log_k - log_z)
+            anchor = delta * jnp.exp(log_k)
+        else:
+            h = jnp.stack(
+                [log_z - metadata["log_z_scale"], log_k - metadata["log_k_scale"]]
+            )
+            anchor = delta * metadata["k_scale"]
+        for index in range(num_layers):
+            h = jax.nn.gelu(
+                h @ weights[f"layer{index}_kernel"] + weights[f"layer{index}_bias"],
+                approximate=gelu_approximate,
+            )
+        raw = h @ weights["head_kernel"] + weights["head_bias"]
+        return anchor * unit_softplus(raw)[0]
+
+    return forward
+
+
+def verify_and_save(name, policy, structural):
+    weights, metadata = extract(policy, structural)
+    states = jnp.stack(
+        [
+            0.3 * jax.random.normal(jax.random.key(0), (256,)),
+            0.3 * jax.random.normal(jax.random.key(1), (256,)),
+        ],
+        axis=1,
+    )
+    reference = jax.vmap(lambda s: policy(s, structural))(states)
+    gelu_approximate = None
+    for candidate in (True, False):
+        ours = jax.vmap(pure_policy(weights, metadata, candidate))(states)
+        gap = float(jnp.max(jnp.abs(ours - reference)))
+        if gap < 1e-10:
+            gelu_approximate = candidate
+            break
+    if gelu_approximate is None:
+        raise AssertionError(f"{name}: pure-jnp policy mismatch, max gap {gap}")
+    metadata["gelu_approximate"] = gelu_approximate
+    OUTPUT_DIR.mkdir(exist_ok=True)
+    np.savez(OUTPUT_DIR / f"{name}.npz", **weights)
+    (OUTPUT_DIR / f"{name}.json").write_text(json.dumps(metadata, indent=2) + "\n")
+    print(
+        f"{name}: exported {sum(w.size for w in weights.values())} params, "
+        f"gelu_approximate={gelu_approximate}, parity gap < 1e-10"
+    )
+
+
+def main():
+    from neoclassical_growth_nn_recursive import neoclassical_growth_nn_recursive
+
+    results, extras = neoclassical_growth_nn_recursive(use_float64=True)
+    verify_and_save("neoclassical_growth", extras["model"], extras["parameters"])
+
+    from stochastic_growth_nn_recursive import stochastic_growth_nn_recursive
+
+    results, extras = stochastic_growth_nn_recursive(use_float64=True)
+    verify_and_save("stochastic_growth", extras["model"], extras["parameters"])
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/gpu_trajectories.py b/benchmarks/gpu_trajectories.py
new file mode 100644
index 0000000..852f951
--- /dev/null
+++ b/benchmarks/gpu_trajectories.py
@@ -0,0 +1,590 @@
+"""Trajectory-ensemble benchmarks for the kernels use case.
+
+Measures vmapped, jitted solves over B trajectories with per-trajectory x_0 —
+SRA1/EulerMaruyama SDEs with explicit or key-drawn noise, and fixed-step RK4
+ODEs — in primal mode and as jit(grad) of a scalar residual with respect to
+(x_0, p, noise). n_steps defaults include 31 and 127 so both a scan XLA may
+unroll and one it will not are covered. Beyond the synthetic OU/MLP drifts,
+`--drifts neoclassical stochastic-growth` load trained kernels-package
+investment policies (see export_growth_policies.py) so the measured workload
+is the real recursive-NN growth model, with the policy weights as the
+differentiable parameters.
+
+`--remat` sweeps the reverse-mode memory/compute trade (the axis where
+diffrax's default RecursiveCheckpointAdjoint differs from a plain scan):
+`none` is JAX's native O(n)-memory scan rule, `solve` wraps the whole
+per-trajectory solve in jax.checkpoint, and `chunked` is a script-local
+sqrt(n)-nested-scan checkpoint over the raw solver steps (endpoint-only,
+explicit noise; verified against the library solve before timing). A chunked
+win motivates a real checkpointing option in the package.
+
+Run on a GPU node (or CPU for a smoke test):
+
+    python -m benchmarks.gpu_trajectories --batch 1000 10000 --n-steps 31 127 1024
+    python -m benchmarks.gpu_trajectories --drifts ou mlp --modes grad --remat chunked
+
+Writes <output>.json (raw records) and <output>.md (summary table).
+"""
+
+import argparse
+import itertools
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+
+HIDDEN = 32
+SIGMA = 0.1
+THETA = 0.7
+TARGET_SECONDS_PER_REPEAT = 0.05
+MAX_ITERS = 1000
+POLICY_DIR = Path(__file__).parent / "policies"
+GROWTH_DRIFTS = {
+    "neoclassical": "neoclassical_growth",
+    "stochastic-growth": "stochastic_growth",
+}
+GROWTH_LOG_Z_VARIANCE = 0.012564
+GROWTH_LOG_K_VARIANCE = 0.01
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--solvers", nargs="+", default=["sra1", "em", "rk4"])
+    parser.add_argument("--batch", nargs="+", type=int, default=[100, 1000, 10000])
+    parser.add_argument("--dim", nargs="+", type=int, default=[1, 8])
+    parser.add_argument(
+        "--n-steps", nargs="+", type=int, default=[31, 127, 1024], dest="n_steps"
+    )
+    parser.add_argument("--dtypes", nargs="+", default=["float32"])
+    parser.add_argument(
+        "--drifts",
+        nargs="+",
+        default=["ou"],
+        choices=["ou", "mlp", "neoclassical", "stochastic-growth"],
+        help="ou/mlp are synthetic; neoclassical (RK4 only) and "
+        "stochastic-growth (SDE solvers only) load trained kernels policies "
+        "from benchmarks/policies/ (see export_growth_policies.py)",
+    )
+    parser.add_argument(
+        "--noise-modes",
+        nargs="+",
+        default=["explicit"],
+        choices=["explicit", "key"],
+        dest="noise_modes",
+    )
+    parser.add_argument("--save", nargs="+", default=["t1"], choices=["t1", "steps"])
+    parser.add_argument("--modes", nargs="+", default=["primal", "grad"])
+    parser.add_argument(
+        "--remat",
+        default="none",
+        choices=["none", "solve", "chunked"],
+        help="reverse-mode memory policy: plain scan, jax.checkpoint around "
+        "the per-trajectory solve, or a sqrt(n)-chunked nested scan",
+    )
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument(
+        "--matmul-precision",
+        default=None,
+        choices=["default", "tensorfloat32", "highest"],
+        dest="matmul_precision",
+    )
+    parser.add_argument(
+        "--scan-unroll",
+        type=int,
+        default=1,
+        dest="scan_unroll",
+        help="forwarded to the solvers' unroll= argument",
+    )
+    parser.add_argument(
+        "--output",
+        default="benchmarks/results/gpu_trajectories",
+        help="prefix for the .json and .md outputs",
+    )
+    return parser.parse_args()
+
+
+def unit_softplus(raw):
+    log_two = jnp.log(jnp.asarray(2.0, dtype=raw.dtype))
+    return jax.nn.softplus(2.0 * log_two * raw) / log_two
+
+
+def load_growth_model(drift_kind, dtype):
+    """Build (drift, diffusion, weights) from an exported kernels policy.
+
+    The drift evaluates the trained investment-policy MLP at the (log z,
+    log k) state, exactly as the kernels package does; the weights pytree is
+    the differentiable `p` so grad mode measures the real training shape.
+    """
+    name = GROWTH_DRIFTS[drift_kind]
+    metadata = json.loads((POLICY_DIR / f"{name}.json").read_text())
+    raw = np.load(POLICY_DIR / f"{name}.npz")
+    weights = {key: jnp.asarray(raw[key], dtype) for key in raw.files}
+    structural = metadata["structural"]
+    num_layers = metadata["num_layers"]
+    homothetic = metadata["homothetic"]
+    log_z_scale = metadata["log_z_scale"]
+    log_k_scale = metadata["log_k_scale"]
+    k_scale = metadata["k_scale"]
+    approximate = metadata["gelu_approximate"]
+    delta = structural["delta"]
+
+    def policy(log_state, w):
+        log_z, log_k = log_state
+        if homothetic:
+            h = jnp.atleast_1d(log_k - log_z)
+            anchor = delta * jnp.exp(log_k)
+        else:
+            h = jnp.stack([log_z - log_z_scale, log_k - log_k_scale])
+            anchor = delta * k_scale
+        for index in range(num_layers):
+            h = jax.nn.gelu(
+                h @ w[f"layer{index}_kernel"] + w[f"layer{index}_bias"],
+                approximate=approximate,
+            )
+        raw_head = h @ w["head_kernel"] + w["head_bias"]
+        return anchor * unit_softplus(raw_head)[0]
+
+    if drift_kind == "neoclassical":
+        g = structural.get("g", 0.0)
+
+        def drift(x, t, args, w):
+            log_k = x[1]
+            investment = policy(x, w)
+            return jnp.stack(
+                [
+                    jnp.asarray(g, x.dtype),
+                    investment / jnp.exp(log_k) - delta,
+                ]
+            )
+
+        return drift, None, weights
+
+    eta = structural["eta"]
+    log_z_bar = structural["log_z_bar"]
+    sigma_z = structural["sigma_z"]
+
+    def drift(x, t, args, w):
+        log_z, log_k = x
+        investment = policy(x, w)
+        return jnp.stack(
+            [
+                eta * (log_z_bar - log_z),
+                investment / jnp.exp(log_k) - delta,
+            ]
+        )
+
+    def diffusion(x, t, args, w):
+        return jnp.stack([jnp.asarray(sigma_z, x.dtype), jnp.zeros((), x.dtype)])
+
+    return drift, diffusion, weights
+
+
+def growth_initial_states(batch, dtype):
+    key_z, key_k = jax.random.split(jax.random.key(7))
+    log_z = GROWTH_LOG_Z_VARIANCE**0.5 * jax.random.normal(key_z, (batch,), dtype)
+    log_k = GROWTH_LOG_K_VARIANCE**0.5 * jax.random.normal(key_k, (batch,), dtype)
+    return jnp.stack([log_z, log_k], axis=1)
+
+
+def ou_drift(x, t, args, p):
+    return -p * x
+
+
+def mlp_drift(x, t, args, p):
+    return jnp.tanh(x @ p["w1"] + p["b1"]) @ p["w2"] + p["b2"]
+
+
+def additive_diffusion(x, t, args, p):
+    return jnp.full_like(x, SIGMA)
+
+
+def make_parameters(drift_kind, dim, dtype, key):
+    if drift_kind == "ou":
+        return jnp.asarray(THETA, dtype)
+    keys = jax.random.split(key, 4)
+    scale_1 = 1.0 / dim**0.5
+    scale_2 = 1.0 / HIDDEN**0.5
+    return {
+        "w1": scale_1 * jax.random.normal(keys[0], (dim, HIDDEN), dtype),
+        "b1": jnp.zeros((HIDDEN,), dtype),
+        "w2": scale_2 * jax.random.normal(keys[1], (HIDDEN, dim), dtype),
+        "b2": jnp.zeros((dim,), dtype),
+    }
+
+
+def chunk_counts(n):
+    """Split n steps into (n_chunks, chunk) with n_chunks the largest divisor
+    of n at most sqrt(n); prime n degenerates to one whole-solve chunk."""
+    best = 1
+    divisor = 1
+    while divisor * divisor <= n:
+        if n % divisor == 0:
+            best = divisor
+        divisor += 1
+    return best, n // best
+
+
+def chunked_sde_endpoint(solver, drift, diffusion, n, dt, dtype, unroll):
+    n_chunks, chunk = chunk_counts(n)
+    times = (jnp.arange(n, dtype=dtype) * dt).reshape(n_chunks, chunk)
+
+    def one(x_0, p, noise):
+        def g_drift(x, t):
+            return drift(x, t, None, p)
+
+        def g_diffusion(x, t):
+            return diffusion(x, t, None, p)
+
+        def inner(x, inputs):
+            t, w = inputs
+            return solver.step(g_drift, g_diffusion, t, x, dt, w, identity), None
+
+        def chunk_body(x, inputs):
+            x, _ = jax.lax.scan(inner, x, inputs, unroll=unroll)
+            return x, None
+
+        noise_chunks = jax.tree.map(
+            lambda leaf: leaf.reshape((n_chunks, chunk) + leaf.shape[1:]), noise
+        )
+        x, _ = jax.lax.scan(jax.checkpoint(chunk_body), x_0, (times, noise_chunks))
+        return x
+
+    return one
+
+
+def chunked_rk4_endpoint(solver, drift, n, dt, dtype, unroll):
+    n_chunks, chunk = chunk_counts(n)
+    times = (jnp.arange(n, dtype=dtype) * dt).reshape(n_chunks, chunk)
+
+    def one(x_0, p):
+        def g(x, t):
+            return drift(x, t, None, p)
+
+        def inner(x, t):
+            x_1, _, _ = solver.step_fixed(g, t, x, dt, None, identity)
+            return x_1, None
+
+        def chunk_body(x, ts):
+            x, _ = jax.lax.scan(inner, x, ts, unroll=unroll)
+            return x, None
+
+        x, _ = jax.lax.scan(jax.checkpoint(chunk_body), x_0, times)
+        return x
+
+    return one
+
+
+def identity(x):
+    return x
+
+
+def build_case(case, dtype, remat, unroll=1):
+    """Return (fn, inputs) where jit(fn)(*inputs) runs the whole ensemble."""
+    save_at = SaveAt(steps=True) if case["save"] == "steps" else SaveAt(t_1=True)
+    n = case["n_steps"]
+    batch = case["batch"]
+    dim = case["dim"]
+    if case["drift"] in GROWTH_DRIFTS:
+        drift, diffusion, p = load_growth_model(case["drift"], dtype)
+        x_0s = growth_initial_states(batch, dtype)
+    else:
+        drift = ou_drift if case["drift"] == "ou" else mlp_drift
+        diffusion = additive_diffusion
+        p = make_parameters(case["drift"], dim, dtype, jax.random.key(0))
+        x_0s = jnp.linspace(0.5, 2.0, batch * dim, dtype=dtype).reshape(batch, dim)
+    solver = {"sra1": SRA1(), "em": EulerMaruyama(), "rk4": RK4()}[case["solver"]]
+    dt = jnp.asarray(1.0 / n, dtype)
+
+    if case["solver"] == "rk4":
+        if remat == "chunked":
+            one = chunked_rk4_endpoint(solver, drift, n, dt, dtype, unroll)
+        else:
+
+            def one(x_0, p):
+                return solve_ode(
+                    drift,
+                    solver,
+                    0.0,
+                    1.0,
+                    x_0,
+                    p=p,
+                    dt_0=1.0 / n,
+                    max_steps=n,
+                    save_at=save_at,
+                    has_aux=False,
+                    unroll=unroll,
+                ).xs
+
+            if remat == "solve":
+                one = jax.checkpoint(one)
+
+        if case["mode"] == "primal":
+
+            def fn(x_0s, p):
+                return jax.vmap(lambda x_0: one(x_0, p))(x_0s)
+
+            return fn, (x_0s, p)
+
+        def loss(x_0s, p):
+            return jnp.sum(jax.vmap(lambda x_0: one(x_0, p))(x_0s) ** 2)
+
+        return jax.grad(loss, argnums=(0, 1)), (x_0s, p)
+
+    keys = jax.random.split(jax.random.key(1), batch)
+    if case["noise"] == "explicit":
+        noise = jax.vmap(
+            lambda k: solver.sample_noise(jnp.zeros((dim,), dtype), k, n, dt, dtype)
+        )(keys)
+        if remat == "chunked":
+            one = chunked_sde_endpoint(solver, drift, diffusion, n, dt, dtype, unroll)
+        else:
+
+            def one(x_0, p, w):
+                return solve_sde(
+                    drift,
+                    diffusion,
+                    solver,
+                    0.0,
+                    1.0,
+                    x_0,
+                    p=p,
+                    noise=w,
+                    n_steps=n,
+                    save_at=save_at,
+                    has_aux=False,
+                    unroll=unroll,
+                ).xs
+
+            if remat == "solve":
+                one = jax.checkpoint(one)
+
+        if case["mode"] == "primal":
+
+            def fn(x_0s, p, noise):
+                return jax.vmap(lambda x_0, w: one(x_0, p, w))(x_0s, noise)
+
+            return fn, (x_0s, p, noise)
+
+        def loss(x_0s, p, noise):
+            return jnp.sum(jax.vmap(lambda x_0, w: one(x_0, p, w))(x_0s, noise) ** 2)
+
+        return jax.grad(loss, argnums=(0, 1, 2)), (x_0s, p, noise)
+
+    def one(x_0, p, key):
+        return solve_sde(
+            drift,
+            diffusion,
+            solver,
+            0.0,
+            1.0,
+            x_0,
+            p=p,
+            key=key,
+            n_steps=n,
+            save_at=save_at,
+            has_aux=False,
+            unroll=unroll,
+        ).xs
+
+    if remat == "solve":
+        one = jax.checkpoint(one)
+
+    if case["mode"] == "primal":
+
+        def fn(x_0s, p, keys):
+            return jax.vmap(lambda x_0, k: one(x_0, p, k))(x_0s, keys)
+
+        return fn, (x_0s, p, keys)
+
+    def loss(x_0s, p, keys):
+        return jnp.sum(jax.vmap(lambda x_0, k: one(x_0, p, k))(x_0s, keys) ** 2)
+
+    return jax.grad(loss, argnums=(0, 1)), (x_0s, p, keys)
+
+
+def verify_chunked(case, dtype):
+    """The chunked path bypasses the solve functions; require endpoint parity
+    before timing it so a silent mismatch cannot produce wrong conclusions."""
+    primal = dict(case, mode="primal")
+    plain_fn, plain_inputs = build_case(primal, dtype, remat="none")
+    chunk_fn, chunk_inputs = build_case(primal, dtype, remat="chunked")
+    plain = jax.jit(plain_fn)(*plain_inputs)
+    chunked = jax.jit(chunk_fn)(*chunk_inputs)
+    tolerance = 200 * float(jnp.finfo(dtype).eps)
+    if not jnp.allclose(plain, chunked, atol=tolerance, rtol=tolerance):
+        gap = float(jnp.max(jnp.abs(plain - chunked)))
+        raise AssertionError(f"chunked endpoint mismatch for {case}: max |diff|={gap}")
+
+
+def time_case(fn, inputs, repeats):
+    compiled = jax.jit(fn)
+
+    def call():
+        return jax.block_until_ready(compiled(*inputs))
+
+    start = time.perf_counter()
+    call()
+    compile_seconds = time.perf_counter() - start
+    single = time.perf_counter()
+    call()
+    single = time.perf_counter() - single
+    iters = max(1, min(MAX_ITERS, int(TARGET_SECONDS_PER_REPEAT / max(single, 1e-9))))
+    best = float("inf")
+    for _ in range(repeats):
+        start = time.perf_counter()
+        for _ in range(iters):
+            call()
+        best = min(best, (time.perf_counter() - start) / iters)
+    return best, compile_seconds
+
+
+def case_records(args, flush):
+    dtypes = {"float32": jnp.float32, "float64": jnp.float64}
+    records = []
+    cases = [
+        dict(
+            zip(
+                ("solver", "drift", "batch", "dim", "n_steps", "dtype", "save", "mode"),
+                values,
+                strict=True,
+            )
+        )
+        for values in itertools.product(
+            args.solvers,
+            args.drifts,
+            args.batch,
+            args.dim,
+            args.n_steps,
+            args.dtypes,
+            args.save,
+            args.modes,
+        )
+    ]
+    for case in cases:
+        if case["drift"] in GROWTH_DRIFTS:
+            # Growth drifts bind to their model class: the neoclassical model
+            # is the deterministic ODE, stochastic growth is the SDE, and the
+            # state is fixed at (log z, log k).
+            if (case["drift"] == "neoclassical") != (case["solver"] == "rk4"):
+                continue
+            if case["dim"] != args.dim[0]:
+                continue
+            case = dict(case, dim=2)
+        noise_modes = ["none"] if case["solver"] == "rk4" else args.noise_modes
+        for noise_mode in noise_modes:
+            case = dict(case, noise=noise_mode)
+            if args.remat == "chunked" and (
+                case["save"] != "t1" or noise_mode == "key"
+            ):
+                print(f"skip {case}: chunked remat is endpoint/explicit-noise only")
+                continue
+            dtype = dtypes[case["dtype"]]
+            try:
+                if args.remat == "chunked":
+                    verify_chunked(case, dtype)
+                fn, inputs = build_case(case, dtype, args.remat, args.scan_unroll)
+                seconds, compile_seconds = time_case(fn, inputs, args.repeats)
+            except Exception as error:
+                # OOM is a measurement (the memory wall), not a crash; anything
+                # else still fails the run loudly.
+                if "RESOURCE_EXHAUSTED" not in str(error):
+                    raise
+                record = dict(case, remat=args.remat, oom=True)
+                records.append(record)
+                flush(records)
+                print(
+                    f"{case['solver']:>4} {case['drift']:>3} "
+                    f"B={case['batch']:<6} d={case['dim']:<3} "
+                    f"n={case['n_steps']:<5} {case['dtype']} "
+                    f"noise={noise_mode:<8} save={case['save']:<5} "
+                    f"{case['mode']:<6}       OOM",
+                    flush=True,
+                )
+                continue
+            record = dict(
+                case,
+                remat=args.remat,
+                oom=False,
+                seconds_per_call=seconds,
+                trajectories_per_second=case["batch"] / seconds,
+                compile_seconds=compile_seconds,
+            )
+            records.append(record)
+            flush(records)
+            print(
+                f"{case['solver']:>4} {case['drift']:>3} B={case['batch']:<6} "
+                f"d={case['dim']:<3} n={case['n_steps']:<5} {case['dtype']} "
+                f"noise={noise_mode:<8} save={case['save']:<5} "
+                f"{case['mode']:<6} {seconds * 1e3:9.3f} ms/call "
+                f"({record['trajectories_per_second']:.3g} traj/s)",
+                flush=True,
+            )
+    return records
+
+
+def write_outputs(records, args, device):
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    header = {
+        "device": device,
+        "backend": jax.default_backend(),
+        "jax_version": jax.__version__,
+        "matmul_precision": args.matmul_precision,
+        "scan_unroll": args.scan_unroll,
+        "remat": args.remat,
+    }
+    output.with_suffix(".json").write_text(
+        json.dumps({"config": header, "records": records}, indent=2) + "\n"
+    )
+    lines = [
+        f"# gpu_trajectories — {device} ({jax.default_backend()}, "
+        f"jax {jax.__version__})",
+        "",
+        f"matmul_precision={args.matmul_precision}, "
+        f"scan_unroll={args.scan_unroll}, remat={args.remat}",
+        "",
+        "| solver | drift | B | d | n_steps | dtype | noise | save | mode "
+        "| ms/call | traj/s | compile s |",
+        "|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|",
+    ]
+    for r in records:
+        if r.get("oom"):
+            timing = "| OOM | — | — |"
+        else:
+            timing = (
+                f"| {r['seconds_per_call'] * 1e3:.3f} "
+                f"| {r['trajectories_per_second']:.3g} "
+                f"| {r['compile_seconds']:.2f} |"
+            )
+        lines.append(
+            f"| {r['solver']} | {r['drift']} | {r['batch']} | {r['dim']} "
+            f"| {r['n_steps']} | {r['dtype']} | {r['noise']} | {r['save']} "
+            f"| {r['mode']} {timing}"
+        )
+    output.with_suffix(".md").write_text("\n".join(lines) + "\n")
+
+
+def main():
+    args = parse_args()
+    global jax, jnp, SRA1, EulerMaruyama, RK4, SaveAt, solve_ode, solve_sde
+    import jax
+    import jax.numpy as jnp
+
+    if "float64" in args.dtypes:
+        jax.config.update("jax_enable_x64", True)
+    if args.matmul_precision is not None:
+        jax.config.update("jax_default_matmul_precision", args.matmul_precision)
+    from tinydiffeq import RK4, SRA1, EulerMaruyama, SaveAt, solve_ode, solve_sde
+
+    device = jax.devices()[0].device_kind
+    print(f"device: {device} ({jax.default_backend()}), jax {jax.__version__}")
+    records = case_records(
+        args, flush=lambda records: write_outputs(records, args, device)
+    )
+    write_outputs(records, args, device)
+    print(f"done: {Path(args.output).with_suffix('.json')}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/policies/neoclassical_growth.json b/benchmarks/policies/neoclassical_growth.json
new file mode 100644
index 0000000..97c88c3
--- /dev/null
+++ b/benchmarks/policies/neoclassical_growth.json
@@ -0,0 +1,21 @@
+{
+  "num_layers": 3,
+  "homothetic": true,
+  "log_z_scale": 0.0,
+  "log_k_scale": 0.0,
+  "k_scale": 1.0,
+  "structural": {
+    "a": 0.3333333333333333,
+    "delta": 0.1,
+    "rho_hat": 0.041817831615435494,
+    "sigma_crra": 2.0,
+    "g": 0.0,
+    "z_0": [
+      1.0
+    ],
+    "k_0": [
+      1.0
+    ]
+  },
+  "gelu_approximate": true
+}
diff --git a/benchmarks/policies/neoclassical_growth.npz b/benchmarks/policies/neoclassical_growth.npz
new file mode 100644
index 0000000000000000000000000000000000000000..8bc0d9cc84bb6763a84943a2b9f22f6bf0a2a4a3
GIT binary patch
literal 19734
zcmd42c{Ei4xHqg2$y6kijF}}urSvqRlrkhyiev~yNFgax8OuDEd7kI_!8zvX48{gS
zL{XtgA${kr_ul)h_x}04>s#yH|Lk@4*?*jUp8X8J;n6uy&mc$lzYbwKlaHYv|NZB&
zosOH%&dk-q@xV=63&(pFcKh!+xNf53q<gyQf8G9H=kSXL+B!y?=v?UBC2m_dnLA3T
z3P~KZJS-tABw=aq=<H~A@20)uZHxc7u3={9WI?;`WNqePLHjJJAS-lGQA$?GL+JnO
z!6Luwj^E#b1r)Zbx)-!Q4`C(@oQb@DfMe{QT*9><@I9mGE>rX*oQ_}4v=Lc{&jit|
zRi}B}L3nRyWfuh<&i`WCIsPE=*2x3WCkN0f2JKm2eum)#$1=HN$*{|rD*jk-6`PvZ
zyjA}!VnvqK-V+>uk?zEewe(ki;AX`im45ZNDCrnrA29m^Z&*zl`X89Yf#!D!p9ss?
z@VRQx=jt*@tP~B!$5tWLAndEQS2u8&&3iG3jl-U*fu9m~Q?RxC?O^3<H)OWuo{Bpj
zgK0DTeXb6xusP~L>%%)spkC755MtI1^DJ6wTgFDPAm-lNiy2>0=YI>~o8C7{{}n=x
z|2KrU?wC3K{}RB1in9Ml;S&;KX^Alj3Efcw(*-OHJ3snBLCB8#EOuw0P3YjgC5Bv_
zICxHtDD4St9*0LcQXb=&vrOOpr%AYN_KEMVhlV)bDA|`D`5JtFJ&)ofIwSiw4YvK$
zq3|{M+pPM065PJSwP&ZKH$({k3mRp7hU@oA4{l8;0K0v}pX^^=AtnCGm!kcF5awI9
z%kTSLG~S~Av#$OrN-9<!>F!I0%F*3{Qr6Kp@wYQjl`9CIk2)#e+82ksS!Pdy{-)wq
z*XV`f=P9_bT(WuFR4g_~`?8)BNWq`qYEnBN7{HvtN#Doo(a6`YlO*ny0-YB{oo65p
z_W!f~(X0{(|9b!n^U3Z16~O<_L52UBgF^qAga3K{KgvN`2>%b0kfpy<DALg~6IFLM
zt)JULLH8hD@gve%SU&qElRE1Uf#O%sJ~`b8eQEm#@1E}lt&Ed>nJR<GAT#`}YiBd^
z*P4&?DpO(PoYp#FiHag8a`=Q6$RPF3*7TLg2wVxPv`rBz0?mJ$(*AC#!3QR?51Qh-
zA&@LYSFhWM<AD-p_Vy&u=_iZY4mCjd6<Mhnha$-1{SljIPDTY@v-$ONf#4D5tAy@E
z;CvolUc}xGhf27`M=kR}Kl@$xgMcPTxo~$lZ?O@+d<ZB%f29V4%q5x4+lnxLU08wB
zAPqfJj|TCYm%=&H#`C#Ht!S$x<2HI?5O(m2{_`&C1n~v=C^Pm3oVY=+IQ6a#&&W;a
zl@>=}m)SqP6BVsEu$3UBu}Fc8#2(JXzx~j1@Qr!z^;lRm`K%s!v>Eujo?Sc4(T-0;
zvuer>8&Njx!8)fs5sYFYuT9Pb03G|+L-z_RaJuYXZTM3X{vN9tYMGBipG&C^rze~6
zwDrereF8(MeevfEDU|}|j&axj84TewXW1pTEqyru>zj4TOEOlhHW1QY#=xS5zzUso
z96si}LFaE$k2?;Q8U389g)QGrCXLJcP)6_Q&|SWGbfLdQR`}?M8tG{QK|}3W^;`0Z
zc~TV)`@89=JuAa*7KWj&lkLFs-p1p0D?;qvobh)T8gR+RkXhEG9Ii<o(~fW_!`djN
zcw$!)VpTRlK&c4H--5W#{Tu|JtNoe#G>M2J;a3E-y5Otqpv|e~CQ!NS8gk-e6^_cv
z4CzVKV$khd?l!d@c#l(%VyQ@hrgp0Q-kU>6s46W84{OIm7h<VPk8)s}#nqyIrkmJ$
zFK5U<*AELELX-I1${;qR>4<Ge5!%iU@GVg@(No9xQm}S2UQJDfpz2moSk4r<J4S}2
z&Bv<B772Lj24Sb*318gy&-n23;XbJUUUjPBUJVkcXTMo65#jqy_!?qO1d)TEq+MqR
zu)>_kD8A@{iE}^SO@(xT=cKfRVG|h+^>q;=(uXi2n7irLfkxa?x8q?+J{1IW#Im=g
zTVSQ#fdhJ{O7SaOb;B!lKh%DqH|X8g4(&0w8THO}fOKdsxh1m}Cb;8I?3t*-yHZCL
zO#8Y)audt;D;GM@+ga%u%WEpiF^ma2XHhYcnKiKc6BVv#W*<9s<}H*uN7Ok#EQQ<p
zdn9CY2rzj*bFp|QEl=Wqj#>Q3g4XSqWg{+>ppU`gneBf<;FrI)q=;2Keom7XOL<I&
z76<*iRlo8;O6S3m35O!wc(`L_{Z0d-w@|Xw%|hJcb?mE(7#UpXd9T0dYzD4<S6$pA
z+c2oy&`r#?5I5O4@?Tkb4Kj?8|5iPmVB-GMToW%^-gk}KK5=Y?38pLCNVkY^uE<Wp
zKBfj=ycXM}klly7|Mu42_?-*WbKAeQ`42<G-x|4C-D2b?ikjP{rom?JbLo3&`PADy
z_1JrR3%>lVJ!^QH3ZJ!fhNuFA@Y#|nevZElHfL^%nRCj9F4CutHybapNdK6<Y(N)i
zwb^x5e=LSM_u4M;*g`zHq_;mYod`V9Pb@t>D0t<FyS=(=DyGk?I&asE#`#A>G7_3q
z5DVp|c@}M0ey3OWf>Ja>wIp*&dn1Mz23_94*Z`lpPqzM5selJF$HS$29e_QT`~A!|
z3alNRI9_ZX4xtYhaznNTLKWdeM{;)#T!@Gv<;o_5KDovC@%Lg3<w{AMt8T&19&UBm
z*Z@8vGQw3o1F+hwbh7OAFsRBH?0Ivu3_gCY=$ay=!{;G`Z|7wifE=)W`Hx@<Omgpw
z@O1A&jk`v+yaqnFe4#1e%kU5|9ofthm(!1*6pZHs<4R!n23vKAVh=u@%GJs5Nx{t~
z2W9-zuHy~CWac5eW=zT0eXP(n3CSt9j)|O)!++9pm5;JpfMejLD3@ykjvBa7{`?-m
zg~N9o-`%9mHmpAk(^+b;L@=QJLSzRp%1eKtMx~(Sb>jQ@s}&fizsz->xeqUFw|HAM
z*oGZ<i*yMMrLajqH#X3QfC@DVCPQDS_^-BZ@9bJPJbJ?_S})iLR^QhtapTSSv?%W0
zvML3iT`hFFT-%6eNWIl27s+@k^3sLYul-p0=jAT;t3|kJB*1v@Bn4j_-g_ssvmY1s
zTf7Up7=(7Y&Yq9yb20kWA5-$vF0?jC4pm-sMZG_dO__;AG&pf~LAI<Md$(V^uf@=e
zLd9#UcAkOo*{ahrwYdk`$PZP{izdKL6GjwguY_YIb&q@uC_tC8^k#LY9X4AH@+(Fr
zLGijZ`@m=+<V;YchtHeCmzlw-JwNkMb>R=o*vT|h3+CA1IYb1%Jw9eDo2ihoEUzhf
zvI#u+hyRHc_F{XewztpoYE&&2+JUY4P{Je2CR0U0(Y^D$ass)~tzzW+^HwT6%?lN}
zl=K{|J$EfP_Y{Fs0{iEghmYXYK!NH_B_j5h{1bem8wD%h?n!)dDS@rl_=wM>1s%h9
zHhUKn;9Dchgxl>Fcyv4HWNM!mh_mX4GyE-q1`CSXEF})w#6s)l_LPBqaN@ns`<ua_
zJ?+D}7gSJ6FdlvXW(a#VUhY@2jm6_+6P=ZaN}%&Jb-tzFfNE!+{t+4`;F-f~<1d)`
z&}kv>?eyON$lrDCJ1az>|6Lt6A2WnRzwoi04PEf~P{FyW!5Z8)<<sB<!`M|6e55V7
z6>99t$3hmnKvYkYFX%uwesHS77V!eqiZ{sG#>okqjs}t+mdWsn&)`Mya{{m*mL>AN
zZH1m^->k&ttC3>D8GqZM3f;<w|GiZx!<+;L!(g*o{MGBv%X7F570$)=Drhy~2ZLdC
z-uJEGr#gGfMJox-jkV5EE!%*Y?iIl$U50Ot7Q{CckZ|ixrq*QB6nK?;m|iL%1cO7u
zTJAFj16R*A^R9#ne5O|3Q!#`XS$9g;`%){$>In$xZy$nN1vs=hsuM|$GVc#qw1GxQ
zJt0T09tz)Hl1gBw!oWAR<myLJC_8<d+oP%iGQM<Du3e?gO-#FD%|n{e=<;H)$LU(I
zQoJ9LM5JJ9dXkGBcL^}EH~1b-reb-7)7r)%S6JQ5^da;jEnj$2zHV|SV*H1fPcJQn
zpj0+nHdk69<O;Xne3D89(ZTVB#=#OC(UvlM__Z0z{hKaD8m0rKM8;+5RT}(?&)YBA
zT7wlXX>3`%!=RF?8}8DQjbo|SD)a3fz~nlnVxQLumsQ_3%biFD!G**?4ap>6^3m;&
zIBEdD;+(Eq*R-O^9$~t-w!To^c~dxUCkc7aYfRqKs)kQ}ndNu+YS19bMvYin2~7!s
z+^;pV;nq|Br9IyZabmdMo{eb`tqBH_qQCkeBJ)l4rnD5u(s*R|c|8*j4m~@o{<a7Y
zZ;~)Iy*~s}kxOo&>NO~6`77Y;Oegk*acs*yPJqJr<1bl?x-sa--4BD_C8*Z-c48U2
zK*H0BK7IEfD1+VF<MT;q{MxcHU#|>sf9)6D&NjTgrSV~6QZ-yN+F2}PoCQ(FHMxF&
zGtrPYB0e%B0d2Q=+qd|20dvGx+37c}AbBXLWL-Q3&tS_zs{?It<lsM%X&(~OpKTKP
zBv6V`ibtdJezzmTuRHM-Iz(tXDCI0}K*f%qU1?R6I&ivlU#2gi35TZ)cj>(GhUeiI
z53ABQ!R4Yk=XV``NK!hWue<mP?<wgeBp48I`s~T)Aui7_#LU_0vST-F%hO&dDagWz
z-J2$-P6cB-dwlfUqArMAPje^;BBNVGx{Uk#cC>9fx>23bhr9kVkk48^0Y<s}JSH0w
zsMIv@Fz-!4v;7VR-s3~q$2@z-QlcI2V&l~xdwS5r_tPTxkPlGXUsr8jeu)JSBOjc6
zNXDI3T3a3bQb6Cpjnmbu1+TQHJzFmC0$nBZz&C#e(NVjMbC50^Rm{#!uzq+Be=Fw{
zyDG@Iv_5%3F*6h0L~jS}ovX$#y}SQNN)KT_)4BZED`j|pGmPfD*WlOj>P~t6G8p^g
zA@WsZ2$fk!g@u0AK+5BK1BU~tkbLv-H}}KEF!R|w(2Y44Zh1Wa@zbpfd$tqym4r4T
zQ6+GAX|@jEP^AmM^Lrv8w<(S}bqH=}&fPRy8h~B1Z+OqY8^q@$Lqx-Ut-yKXb#|i^
z5$G;TX1?FygvPr_HmqNhU@~VPTj$e3@HcbQeOTOxYfg?zi|SRVOOkJjW$(d{e_d`<
zWQLH*Ks5PcUJEMHOF4CXj)U-fN|BLVL^#0Sv{BR4h9w%7A#RBU@Zs@TSIulDxIFjR
zs6EgR+?hR8KdW>cyq{rr^CA_(BZ+MdcAa22EHli1j)LAER|w4|jo9QlQTO2a0FIlh
z{(H<?j4v~r0vVa9IP<KUZH>Jg!bi9wg@&6z-Co64m5>i_^td^5s)=}A)sr>GCK;=)
zv9a6gbmA_HH}!`#sJQEyBE`$J8w<s{>dvOOq6Yh9TKw%IcuZ_{`A*Bn0~3)l=8Pop
zPgb*TttP;$fZ)h8l_|Iy@LO9zYyh9dy;8oS(~g%ITswB<z5?M9JugGCI*i^W)l;;J
z5b|VZ<<vtmI1C3bGt1S3cp-Pp`&lw<b2DLJcpQce{~k5^Ei|KL!14SvgLq&Ir)b)@
zb;3TTg6}z%6pXl|@0(%RiI;UKGJT~<Xfi2Sutlv4kB9irZPTvDFI;B53!*vD%|5x7
z*GYv-_6st<estmeVdfCoz!H$HDLeJy2nB+yVr64gyWpe1*e5b$EgCeq3i7ulfJ^#%
z;%izzP(C_&A@fldHpr)1eW)TpL7d<-(duM46#rMfYReD`>Fu9dc~}h5&6(f-Nw&ah
z`1oM8ZZRZHXB|E&-VfX3<UjEFIHO>&^-=c^CE)OO>DjlQY7~iOdft)c2S2a*%;>x<
z#aB&!@?1S+99P6D*&pRVTHLhrDQ7wU`eP;d&X)xFPGb8KiE(Hmsqk_7uNO9Kc{$W5
z+6Y_r+_?EgAqMO+KW}|R=!A)0VUkHU35=a86<e3{kVE*$w18CuT<!{vZrjreJC56k
z8*B^#*TRsq+8x^SDCFiGmZ(QN!oEP+PoZGjdqB2WCJzhzN#4QfRcLf(%1tMqh?I+F
z_uW%mK<&eO%3tX)WY=)Ids(s>-JZMYvUrq$PJL)XKdr6>s$a6-t=5X1d)+7HXnp(U
z9j}hnkG)_mVyLDNIDjUU9+M5HVIY><l6AGp@y5z(D0sDj<oN7w0%JPlYQ@ah)SZNC
z+s%PD-!;P#gZj#I?OwQbenu_sP&ukMv-%I{)!?sVzd5ovi1<2%JzBr35BO&Dderv!
zqkwMPH&gRIFy#ATz4Dt1Can84Nm7VrHl@vLAJb6&Y*kE|76q6|i?#-w9r!}zVd%9d
zZP;gcyKcHK7Fpj42seD|gegHgMUS2q%=Qdt;yK<8&2l|w+^T8#k%3oi%ad{(E&i~0
zU!@I&ww7LbR7-$Mm6anL@l9}rytY>>p&T9ew=ibvm7wu?7ADorEqFTX@yFK#%^>i?
zc`|si18>xYUl5ul;^HB<4e2d)Nc<jJ$)!exf1E3A4>Q^z*s$31lsJtq`I>jrO7J1X
zw8Z@`sH}i(&{MW6w-(G}!%fGs2^evv`f6+BbCjRg&Y$+n02CaOyT9y>Wn&NesfO8b
zeRIN&Q~gaC{3<PPPaG9JYGe5J)DR%!W7W-YogQ@Pygo?u&BL{Kns=v=HfOc{zP0n7
z9SjbqJ<Q~<L2}T_<b7d3*tWcpBxOnlV~0E~oy;;k*mjL4=~Nl&TpvCDaBmNk&64P3
ztq8bs*>?QjVjZNG$KO+*&cz{4UhQ$SGJFzn)?<i7#JMN+PoD>Mp<L~5-NK`lnEXO_
z-ydxf4%=uJ*C?eSTc8fTg>@rz^WIt$G^X(dcjbb)FX>^tOWu3Z`EK-WIdi;BivZMP
zPUlY1D8Q^7iG*5@0w{2sNxk{22KiRk=Zi*Z^K96tuE9zRe3;S*-W$?~a|!;tZjH2}
zKJV2R@&;{iWkb;~&#D`pC9@WK&$Wa8{e&Ar9zAfCY*Tfbfr`&q10+mDI)KHdXkNUf
z1FSoF38(3sQ8a6V{r#1Gyl_=@^+jkDblIkakthxD^XMy^rUw<UTWX%Z^ZE!nNf;7T
zFbk|UE7A$C41n{*?4zZi68s#$@bm+z4FlO<Ur<x;0^#2W4}Rln!_MEz3|B8x@uFs1
zendh8a$IKB%Kg@fl5Tzx*C7$U9`fz3lFGpwXFRS-vsXd9=#bg=zct{h<5%>4e<ph5
zKk8Rk>VW$AZLwUptI=L){(;;TGFXgkIem&ug1pW%(N<B-7(6c%@vWsEzB!(u&ygZy
z`H?}<%lG^7-@qwf)jwp+iy0EM-%|@seTsBb@A|;@O<dcPeO@T&#nQuZCJ@s`J-5p7
zXTj{x(g@zF3Y>Bue8>=!hU7+BdCLdQ5R;Xb7to%Kk=sqhD;Qca-QH>NWK}B+iQ1`c
zWp2PfCL4RiY4|^*Bd=QjRTDm2eW)NH(1SkAEu?>^5#QTQ`}Ciz!8gIxJ}-{6z;e(0
zXc&DBM!c8Uwau#!IZu0M9g0ZDUva(NPuU9Kl#X&m#7QE$bmbQ5mJCAhCK0VLu3GF#
z>`#C4xdW6QZ?N!6H~?PH+<Sbm3rl1d`8hp@P@|#Gs9}-_sBqlc`&kL_dxt9Uy3p|Z
z`sz#3w<!QSe^jNsOowG*)>pqSw_ytR2vJ{p5T_1*`ds8zjHFZU3mbFY5MDCbt}@e)
zi?M^3&$|r4jpNj!XOcbWXV{UY`y~y<4vrCCFP0)(26w~pZC$vZrR-jT90A2;VmCcB
zAmY(2IWMlg>4NrS4eTA<6(G)Q@qW`{HJ+Uc$mX)GfxW99{BJw@anCg|v$W0*oaoMW
zS-#l~!K=&bUjE&9OCWgS(VZGh&=Xt;5hjE0>mMAkJ1O|;E&u)ut4!Q7yk1x6(g>S*
zKJmP(X~wm0=bfUs`XD?v<Il5q8A!L=maT9v8G`ROS0@|B;-f~FtF@oPApI!oBYJHI
zxVlX>J$y?6&S|io<WY);nLCsBQ|`5(=H-~%a(9b?&m?S9#2X^Kuy$N?2p**2+?OZZ
zcl6`5p26tDP4VDP+Aco1T#NeVKV>hbHGy@ap;!xJC#s)*FQ(yAhZc`zltr(Uz^Okk
zKTjSnMw_2Q%x7B$V8O2Uvu95Y{M@~++L~4iz6a^u!!}d#V6#)|qkG*@UtAMw>iGse
zxmr#R@|VIw;*et9y-dtiPj!kLa)i?%LgtJH1@KAwlDKYBF%lUsX=?VSg5{<yf-TR<
zurZ;S@rO3gmFwn<5v^KqWxK>5{$Gyppr?{giH1AvW!NhBTehL=?cn2=YjQz<h#^d@
zu@__`MrN!lnn8hASKy9(E^r2(NDNhO2G6qbKwQkl-3`yz()gd@!y3U}C370zl=*US
z?`jnaT^=Wh6t+WzJ)PiN0~-Ir57U%uBy6T|%(!n5F~HvSy!CJglH}dia`jvx!{Fq3
zf1_&9SbE4ZV_6RMd!F)GM@Jy@s~TntjbV^H!`U}?q8uLNUan=*Bm+;y$IMJ~3K-19
zw}^yPL&q1}gwG8Gs3d-LJ^MKutp;9Quv#xfWv@fl4pCVccku>){4oTrQyniiT_eD_
z){T9$?;_wS(pYZJDp)<#+A6-2fMLc3>qhJi5W*EcKHt>{f}M@wa(p=$G@8AmQ@;Wk
zExT1`qq<>df{AKjVk)p`J>9#iN(6=>j^mFGmq3gjbK?mbuXE|VTv!@I3?%X7{|x^=
z1fesLemlw<(KMl?(^0M({+yw+mnrPSfgwflO>c@pGr>fcUXuW+DlXfno|S?s%l)38
zpJ{b$tCsKr4TpXWb5!XLCP4>}T5;~C5wPZ_JNtsZ1I3lv=~7hcpyH@-=;8Jr6!b3_
z<=5-NLxjd$(WMk*v9xpwy^;%0TpD&7M-uU3TEt_4>rE)l@S86okXDy(J2mQ9mf*8p
z{ap$%`6!&mbKLvP0Q|aID#^B^2uo%J4@l|^fFk*`+-Lz6lQ%oA9sJvat2LiKDSoL!
z@f)tOyjLHf-l_YatBQK?9K$!4ZiiNMZA{#4I1vVK6hA)qWgCFQX%+E-p{<Av1dAW;
zRIE=qUmUT~iu8H!gKQX^!7)IzD$BkL-gcBH7o-v~MeD|mnj-@+6%?~+-pCaTOh%5a
zgj3*&Xv59pqt#$9eV=EIaR90N_K-9#=i{R;J>@mSIyin=Y>}C>87P8hv&m|WsQY1k
z(*6w@UypyS$kk5=($5<0Z+4y7^KMbAIll=ita!@3_6@<~#UqMx?M;}I*-y8jOY4sw
zPQqs9LBKI=rXwU$iFaaqDb?~_xYX44Gq|P{IgV~*ACIR(nO(5YmHuYX)X!#(zR-lz
zO#6qkS^Ht;O&J?^ND?l{Nw|qMQ87c}B&+y0BILXg<zBdsXl3*D!}Z-|xT*TV`MvV3
zQ0GNCp&rwY)MpNtZU$E3zbr5ATQA!2g2_9j!)!x<zN5G6UX7s60R^vATAlFms`O{!
zMkK@+pA?Nv0>5QzIzPR1=z21)Wcj=wuf%QKno|#j*9R4|Og<KX>GQ-nH-9P&#Z3n2
zPIf}a{hsEs*kqg|-L5EOGzH4&g+mw3Nzm{}$x3NC9l<aA)F$pe?2hVZN^Glw_Af@;
za{bydeeKXg`MGjTlk+3ok1K*q2|;?_`b?aTv+Dj~(F#$_Wg=^23M}4aJ9%iS9d2mr
zZZq{oXxT*dIsCH`R)(E7N5s^_t*wEb1~mM#RG8swo>L1)@wZL$$!bh3%~Jf8or!hd
znLb#xmf%{83(&{+LsIlc=<YxQ78JgCH<MTdS`14e{<Qg>+O##QyRZm&_+2yC+=}5=
z;J8-Zy*zBUejG|@uYsXtD6JGH!e6fVzjsfT;9vtOvoe8-0^=j0rNe_5R~VDUt(FN7
zqrN*hs|KRGg48q4%2HS>`?&eodN{7T9~HU%ya~kUbIF`nI#H;<mEAS90iycuMcV8l
zp@L50Uov|=yp+@9T3TqrmsfnG9*x#vJlAy%U+qU=Z@0s!A*czDJLX;tqUkd#oCV3@
z!yO?1C%j$!LpmDndYTk>YycXP89&bE48wMzj)~F`GTKt7a$9v%;Xq@}sNs_v=<2f`
zrLxqHE@5{Po@Nc8k$GbI@A6>4vGHZE2j<wHnJ%llr3GJX5m<_$@m>Q_`eWXtRve!w
zzo)#50$bz=alDCT&`;bVK4+1OOX(_RwIaFL5X(@@xBCe&yWUk1Y|nt78T~tSX;QI<
z*0sLJ-cP|@_Def`Lmfm~*J$x27NAE-ql?s^Mxg(YlPS7NM%lJ;g>H*%tP^*daHQ13
zaBKU|+NlPZyLG*v{!}O?#tEN3I1>%q-x~d(BZuROy7}tClVYgH)b{e)-5z-A^~@*F
zl!&f2oB!BzX9JzOkWyw&KZvjelxVk5AxisQBlC0{_TFLhbN@61uYP*^^^*tS+PHYf
z3}YiuZkctcl~h0whhz~&xD}l=x>=R?4}p72v|`RgKmN74w%6Ui8aM?RoE9d$!9aD*
z4NHT7Xhzqiu;z@-pQ+Wg#|NNYDNE(XR~o;)DZlA;M+YX3rf)vGpN4M&b62;tKY*a=
zr&VvW>!IIX_43l)4xA2O{5AC<7pIo&+g{W9QHRl^)9HTMV4HSoDR6)QKKixa^PLN!
z)&0JIP)!PGRTL;|tv6!si<4gkIf=k=C0b_cF%dh<=qa1&Dad>GXLC4h4q~nd?sD>|
z1KGxn-q4_S3_Gx+Mq98MSbIK{b7;n4tiFPPHI)P$_WS*fwaVaOq(YEfMGqWi_WPKs
zHw68z8MZu<eej6mZg<w68Yqv{ah4hAL)YQWEPtKIFsc;&?`mut`V2_fO{KS^NBm(c
z-Ly9B`lLd5E>w?C!{U^Pe@Za<o|tU;o+13M4JvZHaX{$F-#y_Mj|mrxsvq2~fSnS5
z{X&lBqH+;I_rgjI{1rad_;zU!WKTUw?s08~6}uC=r<TdsSuRb!la~WYt^K|eN3+o)
z_bmG{QWFaEjBuT5Ys4Q6TXXh^HAAOl<nu+f9=O*M+ITg*95*(vyI#o)h3~g|H|yDF
z;DHlsl~#JCSiR{_{j)-M_^38DZBx_)@#XtxFOZwDuC<Gi>t8%Z&MtFb)h2`b-M6o2
zggil+F;>xVy%lo)WPjcxmPymuUnRS~N`^TxVdgc~*N|V4^n~$j3Wf;H(+`LB!(#Ep
zktfz9wDk(HBTo=eW20pB3ttg7wx?{O_h^D08a(oX*UNCXdrRZE3Kh@ybXmqdDnfzv
zol0+VTK<QH;ldpI@_$(v{@a71aP$9W56Zrwb8MT=gu_kKtF~p=lJTzmEjpVtcii~U
za*<A^0#1fdZ9Wh);nxq@YezSi!-7<#o{CN-o^*WFDIk9rSIR>_eB&|02fA&3%Avj>
z>6NFq*DL~yT<sJu(}YSj`6;C`?@*jQ^|vucG8Qk{rf(k>cn*=B4~94Em4bnRhs5fK
zI^Y&-P9UtlK(^^$GILJhutHE>FI3M*RMUJmZs&`4==^P$Pd$U5>AjT)<s*=Mg_M`t
zdk;d&e>~v-M0;VBuG#(8@`2BbKcY7}_F=;gqg!89tnkj}?`Kv;!eM@`=XA1vD3(`v
zgnjc)!>jI9-=>x`|3?4|)_a5gD}ev)K~em_+=Jrv;@OUxnmVv~)HOF#)rO4+`J_c;
znxRfZO2)9d2L{=mspN&Fg1s|7X#Pv9F=iIObd1~4tM6%0{~ijc^Te~3(-*<s+@Kv<
z)q%M0okdJ}azBnR_MaJMD#k0G`z~Ft?1gL6kGM*{MdO;H!iAW$X8b1-m39_8kVlru
zZEds#tXFjo@BbPH_YUna53lV5kA0!<B5O&|roGC*%wGq3XLJ_$BKlEBDwcjMvK4>s
z^EELr?FNC>;;W}-n*ndWn2F0OL|^5bZL6K7$Q;hN)5fh4cx~@>kMMLs#Ba{=aqcX%
z<(Ok%%dJ3eLk(pP#b^*x)?5rrevZc4F^S32Bs|9d;E8!91=yMk4-J0of}<`~RI%z3
zxLY3($8<3nNTIH6K1WJHOnvU23wr_T9CZJ>BdrhYrwS_BgWK@Q#QZY1Q9rV34<7t_
zCm7t5vj3*AHGvbgkTvCZGMJ?7i|)8k2(qEAN2fe1u$jy=zdqLjZhgEZ{qM;@SEJK$
zq@@v#urh0~%CtjG$Trq2?lw@GJkk-d)C_zYwYeRaN<nX&jX^TJ9!J?09u<t0gQ^lA
z%f)Yz0H^otz9vbNzvoUn>Fte%H=7DfMoYXP*Zl1*X^~8@xiO)h7)gTB>sN<)xtc&)
zqgu=V@c^hHYue%PaJV|M9+&UlghR|RQJt?yaC_105&2~<Fy!}r6|t;?sIA|WLb_@p
zlK;_~>ZUk6l=NcCnMj5DwwsS;Y5$+((E9zb`Z9Pa;rH3DwGq#)e!BEJvk(RQbepyG
z$-r5zqsZ#kgt0ElH#HnP(8oJl##O2u(l(4%7TNj$b^JaWt94+c=K1!L&=N@X`yyhX
z+KGqHJPbOUI|6Ei_QP4iS$NJg$&AWPLXVf1mP=^gYg@U0@7D1H<vs0E^=qAw)O2?6
zkRTCQcUIaP@}<JMvDtMw-ad$wzxVa3P!lHJbFpOaq-m0OIQ;1rND%cRQ=esW05csf
zd#&3xL-g!VmWN%fXkYg5Dp9c-+KB=ghWW2?Kg}PbE?9&l7vIg=G%a7>%JgbYR1a`)
zdN%dUl3@4N%A_GuBU<yPW!&fR1KDD+*crV+Yz|(08$eeH3PGJs+s?#5!M7!qEjGg#
z7tzW#e!m=So^=MldD@HYG^?EO??Rl^9=?22s}>8MWcE@OT5$fw^uC++IjGO2`}-PO
zJ$Mu<`<K&Ep>q0uQraaFo}H_7-!(wOFFFD$GDAJs?q#rIYFCUqv^PibeQSqqO6kOF
zwIUdB@Eff$9>)1G{xhp7Eojyfo;BIt2St1p{t4ems8-|oqR1fzY}?)^gp9>Pr;RiV
zw+0Eu7UUU=*Yr?(Flw$Ts~ta`m^;N=7z|3;`m8L*O?VdOLgT1~FeXjUWUWd74%gCG
zQKt#O$#JGd_GKq7*5)5xp{9ceuhzb@;yyIvzw6dSR|iJZN-CUt!f`$=K0S$z0@s#o
z5^YUsGuY-^zn^K>0z*lQS_Dn2(K-9L=EUn_Y(6=)rNq7i*0MkPzkD_T-lk3M+~bL0
zeAM%>qkbup6^jX;Lgjd|neQ}Vk%n16KAedv%Y@nZwNK|bGO@wsnax;OF_;OF63#c4
zqh#{@(O0rHc!W#Uu~>%$nJTF*(~GsRC*<BhZg~-i*(-g}8*hhM*;m#C**g5qRN<Ys
zy%=j1JQ#P<lR(<|yQs2B4rt})e6M~?gvCp;-ZnIkoRaLZN%n>U@Luwj{riuAtm^Qf
zHMA7G843s&%c*d=Z&!ZT-Yy6TTV`^6kb|GL-qR8}NCL$d8C!Za+ri06%YB-$89!aT
z%ST$xgJ(&3sXTXyz}r&KS8P)Xo~kvij<=iO$gXcgpWE8d*-5^Bp_d9Lk<TXb%Qa*_
zb7S*LOAk(2h17n`Zoz?yh;nDeUKE^p-gEI$DQb?132oP=c^Y1~yM4M{4W`!5J^!tz
z;pi)irM$`{WRuy{SK`@)Vrs>7dI40dCa`>eTi*#=TVLi??P~!i#cNLjFZSSu-tv+4
zmo!Yi9KRamtP4VV6Eokx_Cj!b(E>v;!U6U0H0_3L_}gcySP_s9AIC<d`;#+qyH)aC
z=Mz*Q_3o|FsjWhf_XV@W0|W5)w^jW?$zH6UC?HX`l;WE+)SdRfBJrUSm;S)@K4=!U
ztYz`(1<U4rhjvUCL7XG$k%$rnSPdS1dN$t=C$C&Ls*LG`z|Xs#Qo@_CR9#rYbPoX&
zCBAFdHnpO%m@TV_I~BC<OckAY--MxZb}1H-ZQ!x=hc)OyJuKdd`CCOUg(P3|_X_Gw
znBi*e48Cp9`H*4lx-k)*IouX@ri8<?`a(<f_EyM>pNRF%bOVbk*{c7<yJ5_D>E5k^
zKJZ$;0!z}xz#|f%HshZO{^P$0$`^`YBW*JIcxWS3&`GWSqU(W;mXjCclZt`eSE+sW
zM<-n9yXk)Id^QfJ2YMKbk|F&=u$hb&5kFrJSlX<Z2g_^&aleL080Py0Ij2eB><}r|
z??r(zAO9PtPYfV)P|66a@Gv;e?SFLMWDwX-56{>Z_T#E{d*;(8G%d?)E#riwD+>Hx
zF>h+2X)l-eoDVwRj`x+m6P;;wxQUV3CMAN1Tchs;onNBWc_p8}enMo74nF3Tx=umP
z@EP8#yDD%eiLU)@bQ4b8w)~ew;~Zn?Gmh>LX~M56&Mw8$6?mje@bi5dh6(j{I~yi3
zj723@DjC^Cux0vi;ncYnuv#-bKa$l9CFTcS9^<8AMJVgGwNteFTU#FtNYMI5j^_8$
zl`d4?u{e;<^&AXp%Pjm~()yON=i}t>k@%dO{&EseC5XSCpV^|;fYRB&yGZ)|czNpX
zta48j>S<k$&jl)mMW(o44!6fKQDLtiY6EbE)Sf-0LE{jDDmBmSEC)}G=_G%T4it<2
zSB5}=rz_K!jL)^f?5dH+<w*j(=9aAGu+GQej02+Ax+v%o``G7~B4QG=fU?Z(9uR)O
zedxrw92nDxAbxq62;T2`J^bnham#GR=zX(%oG$rNo4C6P3pASNE{G;!@iNm*4#x-x
z*f2|Bu4#kU#1X-^Itm<K_38O_zZz_7$|@cDy>KReb7UJw5~TR`WcI}*10CJ%ZMV`0
za9p+0=&CZJ;5z@?GWs&SWVwww&8Y~F4^8n2ESEq=_w4z&pepRXpv@HZIuBcth^uOZ
zBoKPI-II`z25DOznLH|M(fgPoPso-slrugZrNrJ0$$8IC`rf8undX_iONvRb(JAYG
zjHX$&#O|_xT2TtlK65LF69yrLy+R@9dnVQ|HecTrP>BJjKiyb(PQty%PQ>lzr=X4I
zQ4Ql4bx5jrKeN$KMaCN)mi=y22%<RL7%CiwLWArky(0t++AShhd!q#V6tpkjjizb+
zKKAp;_!0_!99(vJ-isNR#Ku`<3PI#&LZ*6l9_mqEeA)<X#xny+UISW#pv5Kau{hQZ
zL07J^UZ^d>?|~yS5)I|ZV;jjbx1%1KZMQ4^qRn}&XRZe*2o2)0KDX!EU;xnHVwjY^
zIS8?N_I6vCnqaKo?~uP{EIhjUWNzxsQ{+wl6R!Qe3r6Xt-)d12?|n{G_|Syl{k262
zL(|~syH9+keAT$p^V3R*iv(s@co+)8$k@_D-@rB30V-aebGuWzaEI5{fy)j}Ao&MU
zOZHQMGD6*|rQZ#6QvYl$cVERlmw291NnNOC@mOjjX9W9+QI#q*&X@P}ztQW}jUZH6
zaq7u&B76=X_Gio=h7<afd3PGezUdv0<S)e<WKI2<9ab9+?blXs&;U0!e*6~m_FOCW
z44Jgp31s5|#ptw+<8^5L^N;O|(jF*vtP&R>_CePwYA#>vhd0?-S1vplfy7Fd>o<xU
z;kN_%!0pmnj4s+3-z(aP3sa%QTXl^%cmLHcYEC-bPrLp)MYI$Md&i;<^|ylYmiyhI
z8>J|jXu+E~pNK}2w}uxEbz&E5Vk7U7C^QktoY`iY3?4Mrd0sRN^{4*uH1`icr$t5_
ze-0ThhBsMlkF5ia$XbW2@_y)0=Ztco4q*%D#+~54Uif&fa9qqF2YDI4{oY8)!rH@Z
zzxsE!!9F&T9KVY-@JQ$IiMJ)MvB)X)$&TA?pt-WB`<smns!SKO-<Ia#K_BkpY&DI*
zXdUq6Sxzr->#kVH<`JRv>f#xX^MUxs?~HdGeK~FuFW5E~PV?Luyzy=R8HYr{Ig=f{
z6ilx^%ZM@c_+!P3SB>WNy{ab=87hMq7EJfvaGi)X!j07ZVYN80OaC+Z+YmmvzRfz)
zr2`C?>Q0x~_rUaXvEiE3Uf_Lik-_$n1cApLLQc`NVk5nC)$d3Y^m-ELQ=!}i<(0Eq
zul>llnEdMcNR>PM`)l?_Ks^ik9C*-1y#b5a`k3Bnv}3tW<B2<4NWfm8K$4(&eyC;N
zuPm!4BX^InRK&JsjAiC9jnnmq1skdB`<ycI09jP`HHiww<)O>JTHCOC@{63MRtq}a
z{I+grS`IUXMbG7Tbc2WPUU_!oDpWHKIVWV)h}Q~vx?D;^uwK$z)!|GzxOcF(l&BEU
zPp-%8SUF7t^nU&E&&4kAz5ZtEpI0-~`W!Iw&C0`(XUcT>egmjo*zq!$m4I>I{)*nx
z4u$f#n7#+{132>C#y!fN0>gW**X|%z;p9Kny8<*WiI^=@<6}|>zft}BbY33X4Hv(1
z``im#Y_ngS<7h_3uiirUEz7{t+$JM7zZZ=NR!s+IDag0&gR{noHq6yL=6VJEz#`jU
z3)E?xtKUxwD^mx4Ro=!HEYk~8s?<sRLBQEa!n+J(GSm`dc0AD{V``0!r^hV<+`BFm
zYPHY}ZcEj@IqOw8HQnjr9^8b*4GY^BwR^z*hnM&wPZ2truSjO(cA<7NJ>5HpdbG?G
zE*TK0g|nua_m2*h;=?iR3#Q*{XkHpdue1gtDwCxuS0`w5>VR+3^&>;Dt*ZNyo@5V9
z8XcQ{U0i{e-o0t?FpY;-#T=9y^*M0z>Br%R6Ev@uu4dQ$7d7ya>HXK^@)SIj+J8=^
zB^5lq->CT?AVOO1Xjn`YZP&se<&a!ZD9rSp3|4QUV2_EffNPBd{CGJk^6GmJo<5Xo
z`$MD$9Y2*Uh|g7lz*aI(es~YQJs^GK?a(lOQFCzXf7}L}+!-3O@?S$X#rL3a_hYoG
zf>EDi9jJQ=DDJ1Tu#Nu%Q+1R!uv}G$UAi5K^uMAO6KNi{<uBJXSJzshvelSMFV};g
za(6Ww1q{Kw;!L0fP4hpuXOUTAdp(en%$&Il5lM|>N116_+U}{_cl(C=(erid-cj1`
zJ}yu?rP54>&Jcxz582CreJ;2xoHq9zcj>Be^{qrJ&8R)bY_$5T7V^FBZ39f3avxSX
zLWYtHW7OG%2oT#P%t`Jd<GE|+c>Gwa@cpKYQo+zz__28_zH^~zdb?u!U);z8uS;g;
z+ZO9^`Wi*qtUnW4q?H{TB2(Z+MWs)GL=o;5W|C2K?Z=l4qotF?i@+ySm396500d4N
zZGD~E3U)sV1ZjFcib`xhT9B84w?y83U16|75|8FNJ{2+^7YdFNh(X}JbR~b}B?Vef
zWGHrjk3gx@k`WH|FX8pTz|yIcooF`x$1MApA8r*Mix{E$;`1{<er=5yhHFg0{*uk@
zDC2&o+wX4^J~1ulGX70POP+Z3Va{w!Q9j{ra~pB-sn#l)t`Aj>cr_27?E~j@|FKI(
zCLq!`H9YgN1O;52C+ulnkQy899~05dz?|{+Td7qeDqgj0r_waMY>{pw9-~I^KK8zQ
zsfL2*VjjjHy&DZTJOt+D^je@RyMe_kpa$Hc?6wN-ZGz~}@5r7yBdGc7ptwS7CtB}o
z<6wDQjMaYL0k{5DV6<I{{4p9IME9kWAX?Y}e`9h#1{l(I1Qs86y*<%{?M9XpY4aq^
zHBgJbKyHG@&JEsK0h%X5{`YaWJH0T!x_ET*k|`+8wrI~LwjjoFJBq*V#WK&haQ)2`
zNPfdfmg7u7?t+VjyC)Ic4U5YbdTBcf9G+tN#zpuvV3WCbcP~;3{BEC=Cqq#c&)E~D
zeQ-1&(f)~ZA5F8qY{Gkoj3H;F`LFARgN~N@935RdZtwW<Y;Q#mN{G6zS3@H@Ja~Cf
zYg-1$#FjO&1rbrhRW-prHw!gKx+5<#lfck!zyHt0b~Nqm)yvk|IBs@Tm9MYO2j;pL
z<$?}XpyFiB#|=s+q<s|_*QXoA3Xj1aVpJQ=J6kPgzLE#KBwQxiev>hrcc<9Zq7wLg
zGI;BEizwKhxXaYilXjgcKq1c}4~q`w<t(m{QAbd<Z1h<Lq~i*ool^%8IGztYH_rxw
z_AR3$+~z1y%cZIt+6mbyRSA`CL?E6RvJA=Xfv<b-ZasMS9;PiMY`37{b}<EM2GM!M
z<3FjB?)hZ&@0{8hG~5YSvgBm2vJE#o<tWx{r+JDjFFEwa*Q0-3#&qHbnpae+_f=(O
z2ePF5z4LEqfV=xHy>YQN!U3s|HyPWCaaO+n_kvg}-r`HzrzTwuD%k;vJo7I>+4Pvw
zv%W?&?)D_8skgxWuQhH}H2mxq!s|FcHGszy`2SjLAY5GjSLw990#-ESf9&_|py97C
zPQOIS2*yw5nh1WVcdxYeD5VI_?0;|c+^!F!r5Sprq8p(>Wa@(3(PCK8)~7OkXoqa(
z7+v}M6|lA1@^0N$5-!b`EyYU=;SIW~gv+c0Sh|aII_PN|@bhq329Nf@WqAvox$EJ$
z#l$q(i@gRmJ{`W6{IU-VthoLK()g10>-3jCy-bE&CI_vVEhA9d5csn{lLBPzFiMzB
z6<&ILc92o31nTT&jipXM#=CABvJvuCn53@i`=@Rgy3TVEe3A#zEm_%=EvydtZ{a;k
zc`G*Uw#k*^Xa!sIKmG2?rP#oIj3wAB1N@8aPn1jbU|VCuodo@EbQ0Xo>F8Y!t9*^G
zGQB%NQ88=z{pM1beWa+OywnX31q@gYjL~+idg>a2A3cD%hdCdP3U`8T;-&Iq19_mQ
z?<Oal)``67F@f*J$h7(X@6F8A4t#QgQKwwG6j^_l|4#Dn#IuIOMvcdt(9QBwjCEcn
zfMAtgTor9sB!p|3p0foE3g6w(yU_sD2PQ(h7P>&u;_u4s|G-0`bR883hJY&fL@w$?
zA;gGW(}|+>-O;h^;jr>#Y@|kARro=Hy3!ok%*-&<3>0z+f+*-V%ktB{P>$`dJm#Cd
z>e1!S+QM1Jr+7Wba8JmAGPLX&EOq_YPkWE0am1t-S(8k*#l<(^J{x0xQ6?%LiMrPw
za)ag(`Quynd*uz@n$?UG5ADOpk-LL}%?Y^2#_MpvBQlmma~afKY{#SJT2C{aI#5~f
z!;5#q&2WhP{M*BEDjGY!mb!7f3u3o(HUEyNM!Cti4B}QvurQrzr{2*D@w}7k`!9E)
zu_4E{hSf^EhRQ>Q-RW?NZs3^MYchy<5F_u6P*8B3ZuP?&t?s0oiu6v_LbA_G*0i-o
zSg80T8B7^Mi@mOGf2lb*m@Tp;**X-DtJd$^L@0x|9M`6EzeZrrj#D>OFLuDO@%|O@
z_F`a)WZlsH)(j#nWxJ~-QekxO0aF>8Hg-|g!)W<V4}AJF^|-;H1Oh~V>B!b}K>n*Y
zLZnAgFp<?ZZP!JFk;&jA1tUdx$Nqkr=i?Z#+2xxlwkHV(cC?C+j<})&nnuOhwE<6%
z@1M3|1i_V{9ShVtIQ>DNIpa_h{!8_l7NQW4WWwvYvfc(q>mu%s*7rjA(z3cUUk2J6
zT`jfF?FN^vT!bM)05*mnx*H-yMGdF>Vh&%(=rwvSB<~+$NUt!z@fZp9jRcM!+gpk!
zjsAH3k}5;_j$cvytvZkrOKyJJpM_t=5-mB-m0@{VwD8Zlb`08MaH49k2e*ZpoT{Yp
zw}zdU=~$|IF;Y7HKyVfjJ&qkWq7$n?gDd6auMQoUtR~<&POE#0_%~_O>0W3ZE415i
zBZICAyPLI1A|7G7{Wm0*wj+}CT;SAXI2a6gaCTlx$4wbB8(Cg;7#L5Y(>`5+oiW10
z0|M<ReAU%u;#>w^m#|mNjw3_g>D=f%n>OUTb(k?zu?31V9MZ|f6x<&C#I0hW0O=A$
z%Q^aeKsL(cwG(4H?!G=@;yfOXy0*+6Syv)Jehcr|yYOh*t{mg!Btt(aT{=~@w2J_I
zF^4tJ{B6a(0mma7B>KUn&uP~ZP5ZS!&Dzd5ABEm7T7@k9UGUA&^vgsqZC53S?(Cp=
zB_89-E#TQ#1vT4^=)<1(fs|MB4dOxCPNK^h=G?(7=-tV5(EnL8nv9xVcTA|qMadUF
zkydZ;fb=oBh8z;itkF%l^X6hc@3j)s;%t~4$P>?eS&f6%<+T#-?bv8!rr2{m0dl(y
z9<G`;VGo~i<tJ$Z$R7L{_vvv9jE?3{cJn46W7yFvw|<%9=WAxeN~8iPjCN-C{Z<Zn
z@?$>(cGY7w)A_CC+X%>6$^0=@B^kQTM7k%46yepr6&RuI;CM0=s1{{3!J$KkFS53g
zQMfDiwM|wjKH8NUXq7k!^!*0$*<Tvr<amI_gxDa6wx%X?DW}87HS^0wHoh3xe`Vhm
zm10nR%C9f`Bn2JEd-d~ZxR-^C(0a_b1&n*s7@mKv1b!DziU(T<aMdXm9;ET)A~Tno
zChxa`&d><`TOfhC&B^5?-+mYkRP(gF*aEL#ak-57r^8&Y<}6XH1T9ulWj)RnVZI#Y
zfuVd8(yetU-yd$m&qugIT;=m{{)w23#*;c&)S*i@rR{$1d(bQAc`YCA6ug)`8&LuO
zN^VUt9wULO<%QTF=`K7-%_M!C?1G1jPJvQo6kxKL-6U%kfJNQiR<Su0;Iw|6cSSY`
zhiB`GR9Nf5%n71246?u>T4G%zsu9EtYPoup@<4<8LSfjEK^#um=j`(_8JR5g9u#(W
zz+02djumPhmS2d@tAY`D(rRJ7N4AKzn;iZ4IZHKGO#TVjz0n7MFVY{eI#LOin)?UU
zGP<C`h|l-u{io<@o_q4tsb+lTqyO}zMi0)+?^M`!ssb{<;8M6O39XOm-@j{H1npUZ
zJI5Lu(aN)ZT!LH;Lo$O?B1Zdhf<3m|YD+#?S#3ab>LBvQIDVr@P_c-do;=oGhZpbg
z*WH*Z0s;Ae<tLZ|M`bj#4Qi?Q;{|KBqCE-qc2y4-&qSk3u(X!67DB`C1ZCbsRMbt8
ztgif3iR#m(+1p)v@#k0+x&OKcY`Q2d;1pAiEY;QD?%b|H`A^a>kKN3Ik)o%1B<(t|
z-t;t%r<I5w?5RWQx@{0Ln$tOA(F>YBUNi5m1i-nj9n#E<G_T9|v}Vg^mB>)hJ^h-U
zrl)O}*HSyx0tZ(4nRIL0P?_UzRGM@wmRXpbDKW^yE8=(G_sS$dVcg3zO$kk4`%OdP
zhh{UvzmEZ5>d4@&_jVyd{uPL4EUK~`=zx)%zuThzpJvWHn#wee<Aj?fV+@^e$_OcI
zvW;r^Sq_?pWSL~jmYP!z%_U?xX`zK<Bun<8k}XuY67tKEa3xEw6eY{JwwY5Z*V4>9
z??hqrSM$f*d!F;2d+zzZ-}k-GbKmDVzvp{CT;Z^3P$%WcC<6_QyUov*HXv`(k&7qU
zWL>Z)JNBVT5=yaXTGJ8UV53sm>)Ce)e9P?avE(aJvM|ceRk0C9qhF>ewzMPr0(UFT
ziey|gs|~qLNXApYte+Fjgq!PZ9-i*G3TdvvYBGjZAgVTK_1cvQS4Q8)`%LBI>rVOc
z-M{2x&}E+2o{Tye3y>X4eOU*xQ{P2OZU_cH@r<1_%o<d-{AsYfxe==MP14oQ1*2U^
z-xcK*ER@}R`Bdt`Qcyh|dpzFi1Vo5vNkk@f;H`kp1RG{9=*Bw)6!6CZ&89A5%&Q6o
z`z|%{skCDG<`A>EpAnrRA8U%IAdK_~_$lZ#AS2iG0bMg1T=s{7t5Oyi2M$Iit0bUH
zd{}R<dL^(Q?NXocsv&jw0JC3HV=;eaRIwqM=aRpe9;Znbp_IVwj$~#e?&Y^G(ATen
z-1CK8ZtFro(9`Wytri2B-pak(Gpg``z#*ERL^LYhx0Lg}R|?(oT09r?3dlNm>zNxn
z?}Dk<VIdbIQdg1~-r&4D8-Lh3+-c&~fF0?c!469G=yzl$O#LsC722hm>*L)9{?Go*
zl_qs+uN4VmR$J;pXbpR7jbtMl_1{o;luE>InoR$j8L~gIK1RhkoCS{cosE`HA3#NL
zX5p5^bR1{1-_=<s;Wxu0cX_Itab=8Raf@3Ojt3_Sh$u8ccIXWYYwdby@XzNfA@g-h
zy92LjNhm&Xe)1-=`yP}{bfk0=MnSv!aZ&9ZMQ|(rKv!5#2`bB(tyfLUL2;~h$T~<^
z2<|7BO-R?ns$gERj)4NQ|IB*Y?0y|Q4wTram9~Jb-%-~L)kIu~(Rj-4IXo9(_>e)-
z_+MvGq@r>-LMd)A;$C2BSe6ZA?>5e;R-Oj5Gt}bMh=rT`zs@=qlnH5`>2zUE8tlv{
zH+U|37dNt|Oe0zu@L13)TskNUx2)&T9eVD7Ew6b@+)y@Di0-SdIFku#lTtepyt6U1
z?wX@$Eh$tnmj`Y4xCW{kJ90Km7vU>M&f1~D7?{ZOJ$kh$4(NL69&IH>Ft#rFOr&E7
z8qgV?b}sSIf4wSgn^`h;i7`k$`Zic|wQkgG%>dKKWh?TW2}Huru6~#IC449$rWk#m
zKtVE1bt=~6U~`3<kwyAN5Y{!;;4Z#^%<}v032LQy_JWdnVut6U3<^g{OK5%v9z5W(
ze{llkza$Y+6qHS$g|ZQ3B!4TRY?%aJ;{`VT@EkInjB?^!3&CISQPW6ZEjanfWsW;^
zfO+jBv9Q!e;MgC{8~&1w&q6cJe0c^SmdCe_E_(;P5;JZ6Cr5!6-YC6%GYjlWUp1Dz
zYXjwdt5PhwTTzFPetos{AgSa#>1t%X02Qa(ff1_*P<?8Tk0sf$HgMDSs8AskgG!NM
zE1f==lv}=c==WkYYSryJ6jcjr-G_>0b(%1{fSuB5KLCm|^*OYjHV_{YJ2f8o8y25?
zTiZZV)zW(Q(R^REqwM65sp%~&jC;{lYEQ2LooU;iW7^edUTpkowPip4A;PDS6iNo<
z0uoVq<<#h02MQAs(LWtHKLd8|J;R1;?)X^{n@yjKA(9c~?D5$~FB~Da!M#kCXCW1f
zf?ZKfboQF5!+QYMzZhAI<rI-=PA;-QUQJFoE1&ya7K<qw&<PphMcP8|5T04_d^XTx
zIhDj)+S^S2gZxw8R0?cqZ+?YY@%&8KVsGjRytFsd^n?6U-qgQ;X>s^AuJ=RfBb98i
mGd1$4K}9uT(&k<#xGlENia+}D(WtNa$$OM^hxqr@r@sNfkR;Ur

literal 0
HcmV?d00001

diff --git a/benchmarks/policies/stochastic_growth.json b/benchmarks/policies/stochastic_growth.json
new file mode 100644
index 0000000..6937465
--- /dev/null
+++ b/benchmarks/policies/stochastic_growth.json
@@ -0,0 +1,53 @@
+{
+  "num_layers": 3,
+  "homothetic": false,
+  "log_z_scale": 0.0,
+  "log_k_scale": 0.0,
+  "k_scale": 1.0,
+  "structural": {
+    "a": 0.3333333333333333,
+    "delta": 0.1,
+    "rho_hat": 0.041817831615435494,
+    "sigma_crra": 2.0,
+    "eta": 0.2051731775502023,
+    "sigma_z": 0.07180274153740325,
+    "log_z_bar": 0.0,
+    "z_0": [
+      1.2296024610060845,
+      1.0405861133091159,
+      0.929477223262021,
+      0.9996573515960677,
+      0.8302528952965452,
+      1.2036043826999496,
+      1.0892246228947389,
+      0.9565238483903568,
+      1.1067454672391177,
+      1.0770675044786562,
+      0.8230255876255208,
+      0.9266844557879376,
+      0.9456675135848382,
+      0.9276125419113377,
+      1.056822169707839,
+      0.9585549936963458
+    ],
+    "k_0": [
+      0.9220961859877157,
+      1.249409084056467,
+      0.9329877067724575,
+      0.9241332002131825,
+      1.0010051935754307,
+      1.077808522260105,
+      0.8799155356142421,
+      1.0492419562499204,
+      0.9319484982916157,
+      1.1321340981748897,
+      0.9658447901861608,
+      1.0576707064445756,
+      1.123999157704123,
+      0.8562897707666826,
+      1.0412593731933835,
+      0.9356658429406364
+    ]
+  },
+  "gelu_approximate": true
+}
diff --git a/benchmarks/policies/stochastic_growth.npz b/benchmarks/policies/stochastic_growth.npz
new file mode 100644
index 0000000000000000000000000000000000000000..4f26a5e714879caf235e81f7d165c8b652b2aca1
GIT binary patch
literal 19990
zcmd42c{Ei4)HZG?QItlcGE{_AREE$d3X$eeqyeQUQyDT9GL(79JZ7Hf*}m>IkN3LQ
zJY+6W5}`=neb@KC?{9tIpTGB?@BDMtIqR;o?mBy)XYc(y&%UC;yg{7de;)!2Z585=
z|NZB&onaS)>3usR%M%7qjVzxTnI3y)ZpX;L$>75DzkmK;=OC?{m#*kCGQ4DXC2VM9
z_0UrI48O3vv5fFheqrP1me!W{pBX&2G&K5;`xow;S{c#rTRpySZbZL6B6*bmq@>7E
zen<ZQzsKemfsW=YHeb<-aqHZUZxgUrmd8rmaRJhA$cSB5n}vpdXNGO_KEXhA^IdJ)
zBCK}ipHrWmz@Gu@rejHkC^z~ypB%Y>cZb4c-Pt<u34`h02)|KCopuXV{E!J797POI
zr7og~;BLmS=pSfsbw>q(Z3VrNV^()&7KUZENJ;Sb;`5oS%XY3`FnpldH15m*QVQ<}
zu08#Rxz=$efBt=k^!)vkd7hn^*0C#ybg~o9O8vZ1)jb5Q7O$TO2aUlyy=C*~yX#@A
z`igf)iVuA6De}mCI}aIzi!SB<Gq7YA8CiW7VY$Oe`o#+pYU!rm*{L;!ou)qN_x6k;
z%j!zCozpDl{!K>;69KEwS*TH8FW{3A@@)QO6BtpMg2p_?P*0a+o_(bcBts~sb1Ges
zAndlqlYbh%JFHl4V@bkY&kknnZ|KEOubS;X?Vf^}D%$Jat_>*0f2ucld=&gSBBum`
zsQ6OjgV3>?U2yyFKqw)n3EO8@^bV;|QJZ0)(`=v{&CeK~c2gX~(D>okkq28a)8do(
zX!;l6klpJ#6HWrr?as-iTHQc15urS1`UIVhmy#sbzQTKImE+h&SInnPS9((?;LgS`
zh3ELEF!ay$6&tM)I70MXXJ%c*`aVrjN<0y-{2wXI<*z0BUn#up|4ZQqCikuWJAoPh
zGlBo}{69+IlafdON9oJYzgcirfTdos80&IHuJpchgm#mwH5VV0pm-0L#~WHJq%Y_)
zXs$Kju=x2$an)3a1RdUgsx|OO!e+IpArtrgxwnV)+5}$tbo0L1%Tz2QKQ{QBP>e!D
zd4Z2MSD~%7xYleB0a0Qb&1d^3Jm-5L_|*;q5SY~QQc)LfxTJk7K{p;V<!_XMNIsZI
zvZXd1OM;gIxuPFGRKkHHkrOXPydfeipRwqAFAi$n-7ad71Z*>+(UM1uVef2Emt!*t
z#Qq$*$=q0iT~(744Uam(=Coj3+js$3F`Rc@KHLm~+3QQ%=V<ukT*Q-v%_T@K@6f%&
zSo41c@J{N{>;DzN|HdGR{~Lq&|1*UDdH$b<@P7zGn+-Q4m`f2OQSa%syQw}jRBX6v
zvHM{@-X$(4xi2Syk*m-oLoWd=i&l1b%-6xTqi^=hu=e2*F=>18ttB|+Hj~hA)d`m@
zp=*VMhW>>w*y_3|P~#xHMO3Q~$Qujq?;sb#wz0-V206s+*&icJ-0ct(VGtbjfQVY0
zd2Xtq?XV?d)+Sko2tIp@CAqjtz@6L2Y~uk6Ub${KE;;)g-S_KJ{|;7zn6_<4iCY6`
zoxeD(lbQfWPY8c7+DU@3m)n*)U8=!^KkcY@SP`yjhw%Jz$wa@2V<-HoBQWNguwz<C
zJ{a$l@G#(RM#?p9J3+l}xVU+ZVfIHWd{i~g8QxHfJJu`>HXm-pxYwNr?0W3rV^Hkv
zw=YVu_f}pA$4M%z>zs@6X>A9VdEMK($!_>4F4VGrhzNmp>@59_b-4M@_w7QzYO(G@
zYqob+Er?HYi~P)YLm}f=vt`6e{BATe6L^V&@}Hx_CNDW)iePcQoj@G$mCWYFo#{oB
z)?@z^Cn(@jB}g@k?nm=wmSjJfR^$mW$-P-l!MF?0FYi2&1RnoV4X>m`;2-gQKF2f(
zn86+R<Un#MM9e5HCEugquP-ttpGqTepiYFtBr^<3@5r6kTp%Hzzd^v>N7<;}pFP6#
zwFs?mnCuSGtAn4!$b(x?*F#?Y(mMy{YIGc2RWr~d!etk?T%OT>uu1;>BriMwE^w}1
z6jpl$vBBRwt2IYpIBr$#z>hkNUSSebUv2{(yEA)OzSY8T!d%-!<#cpCcgwm}CIhMa
z9Ry~@Vv*I_KqKg9J1hsbE3J0-VV^-R-$l1{m}V2N5m5I+2c^r?*$D(V(cR!2+3yYB
zn6Y$(zY0bL)9N4IO~Vp(gRkdp!a(!=V9?yYTGZeC_nlc)Eo_>(Dy4L#6Bb+~*O%_r
zU>}tt`%duzHfTzn&MTt8R);eqA2xl&%~yQS(9SeN?#zo9iwd=HWOI^m%U3FnZOqAa
zK3IxPi}GY9t!9|J-FLq!svV+b`|fKNQPEpjCtl=OBMK!RQLdP%gYV3J7kU>O@x-a8
zRuh6bNR(MF>?(f&qI%4=F+(H}{GDdpO{)j~J@-igSzch@Yi%`kuLZ*=kEu;vZUeS1
z_KEL3^~fZ6rT2v;8E0pM1I4y<;gJ!>9omvDAXu)Nq&C_ORqyRH^&eD#Nd3oYZ;1x@
z{^rpM=VMji!!_CZA}|Yt&cCeI+F69}1T-&>?Fd13v9#Q~vN?F>+%FIEaSBA5PmL~J
zsDuE=q~%SL1^7i`DUa(;E=sMxh+2XYw5k*O^jw(?8i7&A56sj9cc1RqT2mc*+vV1!
za%bbz?s)Z${-Mwqr2b+dL=9`>d7dk;7a(z063a`L7O)g*$v)p*3zVKBrbNdgWEQ*a
z?0TXDvt6G(Vd2OF*Yi$6wR-&!{LRJKyxJ3<DZ8{?wG0K5-)A=b)Fp$Pr$oo6x;ktN
zHQmI<+zBBsHGV}K4nlN4+1<0S6)u0hJ?9yh3YnqxYn&@)AmYL}KgUmo1kU5^JzaV5
zBSQ3Lgh(339#m&7tgM65qd77TW)##fzj5j<Mq+4guj|+G7_6NwVq5ko0<P!zQg2pk
z@woRZIThtt{Nq}}w0WL@3^c~Y_rx-AQX6NL_N@TvR-Npk)<htGR%a2h8-gQSJs)aD
zMWFNp*2CN0=iw~vsgb>XEO^6$>^_G?7~1v2?v!*c8YL{(JbqS-&y@R(IR91y)j2Y~
zMTZ6n&(ajs+CRZ0@#In-r2qy`evenvD1h4w+ys-hUSKJ9FrYakqY~$-zxs|86lrfU
zJaI1{xm`>S<VEzsiY)JDv*<1))}`Ox;!y?+L8V#w=UOpzN1Lzp-(;-n<GagxCIEkG
zO?iCfsmH@TO4gDW<B{vtNyTr(C@5p#U>2}yhrWruwr5_a;BU@ve$LU|XcSTDP{3S_
zL%LmQ--@d6jf#U8ZA(2oQQE(k)>{lS-z!Dm{3!yPo9?Z>E*1ECbne#AyA-_K64;wS
zs{uVFC*E%y6?mrV)x_a+8d41k6>9pM;C!NshCd$xE@mG6#=>8RSu#^1p3i!)m0`(l
z!nhQLBUk&BkCfu?j-ywD`|>cWW|CsY;*M?shEKv`dQsbe_4p5sejMhln~j~x0?!i-
z+Ec>S$fI1BYrBDrJGP~7zsHh-Wn*IQUM3Bw%z0FD{!bwWoa81--L6Id-i!GbhoiC0
zM?c_#Vhb8}aGTj@mciN7y*yiEiy+m9RVK}u2G>(w1_7l3cGAi@-!&$I0_jWTKbL%X
z(8;Ufo|^#bI-LBm;t3dW@T@bd`Z?VDZ%x<8p&rhi$2YBdG?)weqa*yK7WUSf9a4SO
zhMVWxG%K+J-EYmWxJ;zMTE6+{j$>4uRSf+6gD(>T_j(n+7Egs_hrZ60ffNXw-4%B4
zM;Y9u=4@x44ubs_ts~x^wRq&;TXtRVNZ9BVr#Y`(0#BT6NI&N5ain42Xq*8FOmf_9
zw@i{i`sV!JaFtLHP}N`?9x4QPm-QM+@i=(?cVx13H4hY4-_4t~5@66P@!o(04Zhx4
z)1lROqfM1NORBano*-HI*=Og&nNmT?$*LTTY)Xma5-r7ttgQkIQZx{hZaAp<dkS_a
zY^{&z>4EWY3n^Q&2QXD2`CHDWMp$LCb$WP}2=*$PSN9uq;m_Sp2V4BpAzfqlw;k_l
zAnPV)`hY4K@x=EC@!oP|HyXHqYfmjaR#$Z0_5L%^I0+2x+de}~e)(mihZNwR+p58+
zPDAdKf%8d+I`K#ypPR@A0t9cXe3Vv6#?231r6bQYqp{loHQR+Y+<x48+}FMlHY?rH
z;`thdx4o3-vY(cLUp)K7rm!aT;qy~+4d}oZvX6C7(@2m^O*-j8$HP4WeWycDW&jI$
z$gORt5T70EiJ4if#z*^aHa~pUfW6!9ff8RYBsM$zP&?I#9oBr!+4~zoG1jwrYfm}6
z{yWiZ@PG!HI#Ss#H6d6Lt?a*-rv%oA!ooYYIKnND_?ulTH5e>mTsW4DAUXHA_WK?x
z9<+#jxSO{CY{gGXy6Se~!MpDykLacVdAHi!!RIwl)ERY{mPo{6t~<<=Jdt>q705Yl
zx%Amk!#?wGClt`~+V1_!!`*2C2A_@+A@o%L-ABTS(9ZP5jPmIt_+A!PU+%5IdpTn=
zLUlc$?;b&G?oPqjt=rQT&6}a{+PryNT@!eXDeiYYkO1x5Rt#UR#RB`Y!PV2P0T{I&
zzW=0tBhub^l(XH8f_<)u>_@-X<8krdb0T^K_$(Q0Jo1cyI%lsj={pd?|A=_^rOH?!
zyb%k%=$eI#j3w0D&fRFwnq;`}i3~SxZ5<Hb5)TJO%Z~nwDTL6vJB;jKOVB<$;{8{)
z9-tjbY!~y)2VT|#Z>a@s81aI!Q0HqcM0`7>tLWQ~nfe<JHebub&t>*SUu;`pC~x~h
zZ+$l$RoUoVtQv-^+cXv2!g8TY;cOjaRRc1ANG<vkR1DTT<KK3ZGNG}&BmDP|cvRyW
zPs{u11@CW8r)~Y)3bt$GN0+4QL0#q>%V<U(+NHQ?&dayJewV#hsfw*wr@?k5i1G-+
zgJ?C)sr4v!-bM9BYX>xow|P(*+Awu=huOw$4Y2jW8Nq|~IT+VwcIJIj5_}vdM7C5`
zzz+rWMaqE$wBw$R6tsxN{hYOfw~Z=seABY5c<fu4<9$AuYDb2=Yprf4qXIA{xf>E=
zyufhiNK!p(2W)-a5-@1ihW~8XGv#B5Xr;MyV_y~p2VHNwZzR2g8TLJpQP~P1I&o2i
zGFy<~Va|K%){R%v6C}BL>d^HOQ={UR4m4hNQQ(n_N5#ATgO`N8fa8n}6Q@rrX6-u_
zS}c<bky{K;wC*P2v}1?3g+LQ{oH=>X#B2Z^HrWWxSyaQ2A%jKt*9JJN6YZhC+K6j+
zt}~r#t;7?-@-7!vJmDt;O91g_4+;)O{{x*e+?riH%<o=|<oO2N*jf&q$`?W^hiE9r
z{9GvNeI67&Vjlj`nGQt3(00CsCWtOl7%x3fgR$0r0YU{4f1ldh<G8g7-5yoW8b%Xw
z)Z~ZGZ`}aU+x+fT{y{2qsdz3(EOkNqsl8FUiv8$taiR884gnJGd^mSShX~tAZ{rT#
z3<Zav`iZZzk?=LpNW|}6AN=)^-q%9!Tij_!x6RYPFP;YZ59bS;QBQK5wIRM69d(O*
zpRF`ucWB1*g?@khp>Jorp`#au`#gNVup=0}^2zd3rNUq*^Z8w+g}ABxvf%}*A$a-o
zHtqNAc=&cMFx!Tc3LoYPHl!*d1Sjjai5;inUScHcYpBEh0y8I`UFb)T+bMH*76LGr
zUFWb*L;&ap2c$J8Pr{>jwbynmG=slX+Uy3aA~a)=R4Ttf#hV$;5GPU$Bby(eDD3XU
zdFuPpq1tYoa4t)FK%$`P9gi<^>@5gF3wmo@1u!{Kc=Gvs6U@>Zc%`}C0Mv-Htu~4E
zAY=Y^aKqgSsOL=mGZ{_8&bZj;{I*04OfE|(h)jS1PC4#+!$y3}^f-@we<4W9Y@(?D
zpyS(g#iEo;GuV?T8xJ#5L3H`%<alKjt{Dw;YzQUc>l^ljFKwPU;yi!;`FagR$E6)}
z+S!V>BAwGbQEjMniC7pCoQz>XYL|)UJCW>Mzd4I958dASahl#M2iZ}|TQ8wb@HnP*
zB~-c=uN%~dlqbjGr*}(y?Mj`X9yEJJCA$&q1mqWu)@zZ|FsXKH`fFTQU{*LPOTuEM
zJ(QhOiMU^f7IRXe91Lu9lhgex;YLo$Jd}5#%JGO_9F~<Zuv_`*J3=E|T06(^snZE6
z!&E${hr7Teyu9=UV-^}@L~OF)ErheT_{SGE=HnnqZr{riFDyN5)lb|>?{Bs;sU-7t
z<0+f5zMQ#MD0I_KS~p9=;}RX(4DDo47(PC+ThSJdm4`hsd-4&sdW5RzPgH{HwgA=*
z?g8-W<36mq(F1dF-8b$VwL%G-vH10E%^2J=Yu2jp9u2J5`JSC4!iyPCD+S(BIKRdH
zPrL^O$m#muCUtADNV@is@R@87c^MnZXApqV*<T}1k0RRFt8N%JLA>~CMD{d?2YCMO
z<A~ADg+JBBwl43h&}U`(oV0W*Xxm>>(Hf=TjbA%C_eT+sb5`Z;pG{Rr7JE4E!%2n$
zL(?b)`us1Ry|K8bxCczv8@vZta!|xGR@gbO38?QmpORQ&Kx3$|<u{`ZC>3SRuUu>d
zGjbE{h=C{myx)5Kwrf6C55KUF396t|K)wo3lZhCgbnQsXKqvIf>(J`Hb>i+A;ek}S
z4xmUJ>ne1lf#mi?^?gc(Xg9nfcmF{gz6ocuWq3@-caxnSek@J6Gaw^CioFp(ZDg`=
z-js;e>A#+eN3{XVxT>4kr+R!WxfC8f-vVX(T7QLWX2N~$rmxjqMJT6dG|?5+fGxmI
zX__HH=5eP`7ugz!;C%nBowF2$lzi*6nex#!<DJ^SLL%z?Y<g4YRSQo`{ff7`wc+<H
z0{cWm2|f*Gm7Dwg5fhvAMstpL!d`3BEWOZru-KTpUas7Xw-}F~T-fgdmpptFoixiJ
z1%~W?+f;#M!9%|OrW&MeeX8<BE*I`AY?esO&VmQi<x!dhU7WJd-^gZB3|_L$g%_K$
z@$PK8@WB8I7GAgg63<A6p_BYs-0tmoFXcFY@a+uTqEf{5P>qE0_KwT`hvOizV^6(^
zZw<=7%KEapB#ADGmPy{$6mX3Z61nyv4}-JzQb}KmQCRK$Hs|Vgur2Z5E%B9zot2**
z3p$G-q4`*vws$g?NJ%{7*pr7+##+9LZbW=b{qx4~Q!Az@n6B;V$wkIFp+aN34*Zd7
z{iOb1K3W9_e{LA9ffC_e1(CINAfWY8@44s;yt`LHsamQT->@d^n-d^`XK?nDHR~iy
z@KkX%KZU@O=_<uloR2NKbxEfl=EFX<Dde1}1%Ewf;>7+ol>F7sa5uFM$>I;EHoc&L
zYNMmh&VSA5xq3ibC8i$acL?&||Jw{-Q@@1}w^H!t-)Xkij}6e;vm{!~-3*UaN;)6v
z)8~G+f))u*8uB0L5qWSg1YAAds7$I>L-m~#x28l3K=qQ(kCkox7~qwvy!&Z7h}`?5
z@sG6!PN@$U&7CJAzyH#Vp=L9VFu8s24r+m?hsU_mzc%4<{nL@1Z+nsb)t((;ku9kE
zNTlW^J-@mSABs3!<plPdJl<82^O5(3eCa{;D){3v(mftq4DmKB%s!RL7;dL-V=YI7
z3Mtd~HLr8f+g(e{AifD6Y&1S{yR{k)S#p+_vmwq$vS%^nm%_x*&7G~DWc-++xOnDi
zH-6bKrZfp%NIRMLz0afou5Y_%km=U~%}*cA7}rLCeT&K%_3204_w}Z-rAaBMFN@u6
zcwd0u%YJ#XCr9I3d8gs5_8PdUN<H{{xfqr2S;UZuwI~)VYf<*D7KFkipWPx9VgHUn
zn~d~&eAUmd%FkAX(F8fO@-YJPa7!Mq8KCoEi%!u)n{rW}!@<3<lmdQC-&Rg<sz!Z+
z)CVT7HY|SpHvdCKJUX8CJ)o^y43~qrU+vjXz*(F72Dv%I5Y99E(qy0p&(L(NIsVnc
zzB;E2tgZ((#RSFg$`rKA3at&Dqu|T(Pd$R;^&q(3mcNi)44Sh?zwdcn31P~4g?(QN
zVQwh6Y`}<)vtpO*9?|>NvQy85vZuRnFhBByvSBu=FrBp+y-3e{e0VQ(Hx2K+wpnpE
z=?9Ok#n-^25tU<?k3H>9M$X^4Wrfw{7@QoeYAw)+n~%wBGH<Ivv)=8gA{OpgqEXSO
ze2fMcisCOA>b67wqsg~2^gJ&8N0!uYZH1i-+#55i+R%>cE1!IREwbJUcqG3-=fCbZ
z5AHeGfviLeQBl<@bo${hUC!1C-%mG@mJX9~`Gw%GarQh|`MX2scwH|r^?iuQR4c{W
zxAHI8Ovvz?YJcVY_B7;D_9iopbb*hw+=tB11<3l+o#Q5V5Gb<e9!lu_i0g}KM^i`R
zaF_L`34X;4a7wtEKohD*KHg37Rksr$S=Z`8Kmr2)%T+-Gp#iutaJO`xiGqe)Lbup7
za-hB6AeR4X6~<(YTu;A8gu`dzgESKw@m^M#Y`I<~8vNWDVz8P8$;@%n=R5N8+}NRU
z?)V<C-gb|T__hQN-7u|uE?)%_4h@%$!#fdWeGjylwgF+ZXll?h2G3@)#w32u2e!|r
zjfs7s=r+*2$Kj18jAdyk=UhvH7=C>XuFxvX?Km*~`B5~~h%icTuIPj{-wgvBm>OVB
zd+xorY#nZxS>PY{a)akI=cxkjHi%Vsd1-Z~0m+hGwPk;^!E4Sjh<v&O0z;((n5Ic!
z^N=?<*dQD3Zst&#Sg3*fde%G>k$GrStk<ileh!Xwq*z*=rNX5RqFt{#s_~$??edql
zW;o>eInAlR9y>?WmRT0tF}-r+%hUQMbTg&oFg<U<+T~HX0oDrG>#*Q>{$Mdo6<tut
zRIY|!n)1G0WEY@a-yMHtu^0X*`4&&z&xeWpm87lz$gt<6d4)fr8|0r{HqHB73R<&Q
zv=7q|Hg*j!+P|#F22N$8*zkNb>=ml;>30DKwlUMZj7HcgH)x?PSPO1$Z5JNN$3R!?
z3m(0qDoDI%F4ve|i(FN96t^}a$hU_xT+%1Pui=yF7nlR^hsXng>G1;8k3147Olg80
z`qJHe+w6fe><U%Dv=QD;?DHYj<iqbRx^*-6DOmb|#q@ak8+h<Z%Ri#7AHHun@Gp2X
z0n_2O4R>fF6jWMk>mH%tX?5lkO(}IiW7nD7ODKkuPJ4N{ovT4p@L9FBS|>=BA4&af
zoC9JbU5=%v+adGbQq!z`7d(TirS4lK+_Y$I9`>ga62*+0d8;Y-=ANXWqQxK@9`@_$
zBBi5VwP&B?#}qJ}GMTt&UVx^_x(_#tAPVa#gbuaTV|hYi-!fM|I(}?vxZaV6k@MTy
zc{+Q5z?-#KLnagLoBgXAJ?MN|;Hx%<=uX^qjnV0>0vWZOZD@!2i0C(zB>-!2u)1;l
zVbqp3l$nwroWIb5;uAZ%e+4^%=iU$VEd4#elt?QXWy-@%57;LCLppJV!z9qIHy(ww
z((QU=nqh3FQEj7hI}9{5u^cMTMPBKH)88h$VYaN~aoY7nYz(-+Ba_krl^`Yc=yx1&
z*6P%<ig)3!H(9qjqyq5%E|%leraBmu5iEJQxdO<}{wvbqMR+D_xX$Br7jj*-e|0sx
z5EeY^Ohh;;(MqRZGNGUp1l!5`sd-&c<-B3=jY0#ae;&Ck%h?VmyHA)M?4{$O&SwGH
zt~A{5@_v>q*iN6zQhv{666pQ*?mu65{Kx;3a<e(p1f8#6{+Q&T^KcT@u5H{jXpuh0
zw=*>k_v`<C_(_Y3`Qn)ytm{dj6Wl6ia*Kc_9v6r9&J<(n-x10(UF&5x;7r)^hOUkB
zEe=MUF2Pye4V+5nn(%_$wCl=r51cdq%C!2a4<#9Q4b6C!f%L}IUmt%FaOKd^PYygE
zAvO8b-g<@v*mo1M#PfTx^yN@8yY)LDK4ZDd*%%M;r%%{i38BK2s@1qIwHXvHJw0<A
zQqX|%JVo>KJv=e(pw@7i49B=kHQc3=Q1TRq(jTo36mPC^zjd=5c7D6obe2BPWGOau
ziANV<de<$!+<{DZ!7Q`?h-L~lbvb$O-`4;<kt+6qajoDzebkPN&a2HO{W!Q&uNkVl
zjMEpqGGIh-F73FsCvLqdC#SKc9es63+vexmfZc;*+)J_s$Nsv152EW9q6Du!0j~=Y
zj&L+@>+1#`(M=Z%`y0`mN6F^$Odbk_^;IA1ErG^F;)NrFg>d(v=97b#Raifr_&D|Y
z2&!7`+4A#r8uZ6BnYI@cfmZb%;W?Hf=sq>n$7@7I&og3471muS>b3QpAJqW{bNa+g
zCxTH@d2Y-gz8+p~EE`XGorwFY<*3w!cG&+sbAHErB0ltCG)?_p3Zg5v($PCRkw4^;
zvlvq&oR65I9yKMQBYRm+=C%s_<1p3cY!ZYUrt#TFbpmEKkzY<n&@iRmA=Kqp7o27?
zH8WpM#M42sgx&RI5WF>3nxNJL4ixFPOY#)l+ih~qSR@mkcX=JET@S;TB1eX$MH^9Y
zp5dyWaW~3YE4njvq=Aaq-t0|_9{6TPB9R<LM9yxDbSw2<5I3(oRZD0^G1<wgN@@pO
zBUaPK-6Ek!GbKv8J_^N*&T4Vw=Hiam)}`L=@$m88G3#FKbWqfwSV`3d((CF1nfOy0
zOE-NU<o;C%Z^C~)dhjO`Prue9e2c1s%*BHRqs(o{AG;VOnqUmy<Sa$TI4VJZNJK`%
zp&Zs~cBVd4@_@0Y_Z}`xCBqe=hJ5FVXb=}P`tEO{4YZ`FZH@#oTqNJ-WpF4(S$5@$
z_77!HLiP&WVM)iIoo`fMR&_zVU+)pcnI`-!z@Qpr&;#xZ9i)dJyFttUPpwEwJ)9kQ
zz$+430^GyDa~9s$V&vZQhZE@h{>_XNw0B2o_-2>CrPZ@KnBV+DZ<kyu+z#cmWY>rR
zQqM7ipeZl33@kqOJ*XSr>Mz!su$o}HhO&_Bu6kt6aXH9lM}qdH)XTdD0wI<|Wyd8`
zI{G-r%pIY&Vn7tr*b&2YR9min<jh5dN56H#<kzxbVr^5Hd2~Je4E-|8D{&o0zSJdh
zZl>!JPA{l|y~SA3meuvMuntb`5;|u!nvFTJPu7t{!)LtXQ-h=?`rNr;aO0K==($!R
zC-kZjMf!_waWE2Lx#iJ6ePSHi`(F4m#nK5HD!w{rMKi(l%$Ik{dnnMRlrgc-xC@MK
zGxL5PXaSy%wIx;0BDi_$@N~CzJMO+-?UTIH3_s4S6)3uuV}jr1{*A#Tyml|Z_vAbY
z<9U8tkjo2^UE|3XLoC2eRyU5i5V~;u;7hp$qqoqk{bh>~Ap!L_<$k!-ngW+D`5*RX
zOvOE?j+VyXsD@U%p{&<`yPz_yBf9BfE%XU<wJTJRF=MNZ(|||>-Azz)VYMv;`jvfb
zP8#H)LT7_PVPGZREV;=S_qYa{qhs3~`>D_$c|pX|pb(q8HrAxHM8a;p%n?d`8mcci
zrYr2s!u!8-pEa!K!37cbYajJVut}_xH|;bLSz2w+seXupd^7ot1Eqby%2}H>b0i*o
zzL`m{#*iUh?^&1j?;Je-=8((rC#f)$kp6IM@;yxEpDvDy^Tfp*<qYqMP9R0!I@&_7
z-{uyt<oaDmSo!tFx39-?F(gkhu6Gq-Av{&7d@&av2edv@yWEM}{&3utku5@s{15-5
zKK1Z_dQklAfcjq+hX3}UNErOT>_NGveKh5wRUr=SlK6bfD-VKt#hZ6E5kRPN(<S3S
zH8>*ZcQSmq85srE{WGWDVOWIqgJN^KFze$il&V6*d}2)WRpw$aVx&$iE~KER{pX+M
zVni5p+lmR<RVe5@VmW#)9}ir)+@@U}f_)$EV<FuTacxV$?7|Kb?v+SdvVHvx-!87&
zOdg0pLiH~viO5bY4{5G3zFr6?-2-2~>Wqd(xr84ZEAml!)0!VQUp_jF+GMq``QU!Z
z-e?imOwd#IeEmu%5lycr=i2ts2f)&IOqt@bcz46M*`(8%IGVb>O3|<eYd@tATWzKj
zoBX0*vo7ZSUjd9Wsnh<i0RFcJMe_e~56bR}F#*NLbZ)ZrX9J@yy=ET}7b1z0K+bV|
zrhTjpuHWL&iryCm?GAh6BpRC0+|Mw2*KH!+kLv0ld`kgM;}=JlT65uu_l~mzjxKQV
zm(cFe1R72pS94_v$j7L2H(Q6j=rv`B-{P8wK0rp<P9{wf#t=4nU1cQV-dw3u8m%O7
zc=uNM==n5wzk_(O-=_^OC}gd)e(wZ^T$XcJoXR1$Be<aCDh0paj&WP`ZNw-&<6&2h
zRtRgnyzkn29h`lrF~gyej*IQz?aXr$ajemm_)wu8h>C``H^iDjne%C@(xU{NXdR>7
zDvm?ZJ2HEdOJhOgYBrj$N8$L3hFpcS6g;=9Tl<A;7v%h8IAQ;x4cvI#>y8x1faQF#
z`KddOxc=#F@cqSRa8h{EZnL=%zp?Fhw_54|Vy%TX)t-d5DO`5;oL%VLJ<NMp-VV-Y
z-1zeIDG4_4Dt$1Sj|XFsp?DjQ6!6ZG3cB(3352$*GHaw&!)?i}*XE?UA@r}*J08vw
zsCyq3%Kfkt8uVr23ajeC@5cH!DVAmsNNQ|N6o`brBKB29kvh~{+QsHm)CeOBrl*(J
zia=S+JwG8U78~wdt+_s001CfqofXd%;5!Kuzr2A$_}(mEExorNBrIN(Ruz$eSMDHl
zD`ywHFV&rXCHWDUN-sa@)2qQWqWdV<7c!IvKl2P~OM#>O%d9^pDquOZ(ZNQp0IKAj
z8kjBfkTWs=rkW8IUZnEv_%PdzdzQCN#S7;F?{o2D5r!K4I;N-j+b0u?l(<Y%gz22>
z#e$|rx<>Z&r&8p_uN`=~&at2AKoJbZtS%hr>HrJz_dY))n{k$BzgDeX4kUjv`D>0%
z*nj>9?GHl_+_@onzel$c^?2T}2oXCm^VJTLvPBc!dtt0HM~wqu`<gaF&x_1|pM$az
zQsA!kjo==x=fL#!(MyeWdhd9h$@7?KElTlOE0p#m!ZwKrzgQUx{QXcQl95Wo3J$KZ
z?<+OX1S}g`4mDr|r`#I1Of}58JTAG$9F25SWWjJ=E|#C?Zgiq!{d!2#iMZQs5Fe;%
z*`(eHjy27J4^v9;oHHdYYcvAHyb?-%qq^~sMO}kEZwVy7dN{kJmjMl?Ki_3@cH`~n
ztnI;SWw2?oQ~bkjGIE^x=lbej8nTZ&e(2g)f(yl2r+wOMv8P+T@yW9Uymnx(SPC!F
zHAzzz#g$IzR`4!&*w}%YH5GxqQxwb(yCUjG=QzAmPu^@wPQxc$?Ds}5HbLn06=!zF
zJRqA2o#p2kz~4bWOwt1NnAB(cCDWG-3_qTX*Z5HINvG4VFU;>jx$xk4z~LC+zkE_%
z`WY2Ya!<-sFFyns;^jA4{LMK1d|i0^iE7B;c4EJ|wFYnJ7MNBs^}^t;<njfE76{s0
z<IAeu4Cd$D>(RFb?bW#hcsu<;bL3&3kZ~t+KmP11b372LhAuMN9rS@|^6zv{IyWCY
z(s{V3s{;J0V-%vd7s1mn`EgO4RJ_(Q+s>s>i@Gkg9hY=_pw#NxqP}z=G`bUsF};pB
zG+xj>saA~LH$`0&-V;!@$4Qlx-+@tH+`EgDy<tuya@E8=8jW?o|53eFffv0sBtIW0
zMusc<ZnoZ|zvs*Tzaq=YuyV<4VauHoSYWXKz&}#}Tcb6!{hY`!Ro!P@C0v2grn-|O
zl_l7b->=R@_l~?i75Veok7Q64IO-bAJq(u(st$Wtwn4{1<xJ5?0_ZzwRA&^_VX66n
zmkfj?uxmCbkvmNVrFSBy*A_{T;P`Q2o4G40v!)mi+c(4cdj$^#IyxYu!Ta8D2oaT|
zOoK&Ra^WCZ`DvGI6r^XbjYfr2fRn|V>Oa;5uFrWt>dG|Y$JTj25giJA`1ws;^-&z&
z^&vfzoA1DTPk32-bn8)d*T0+-Nv$~ZYBbF3OFcHz7KA-&>QG*(Ymtww82tV!Z>nu8
z!Uct~FRNE<k>4rh=l-3|xJLr7ZJDOwV&?X`>)s7O%T*s%Ea`wNg%^w{<1Ki!cl4dh
zbrPgLC?HBjH^6aAFB&tc1EgIHD#s%TkW;R3Qc^b;9*L7@@?^uIs7p#h>2ewB{Z{39
zFy9VEcN(^f9I8gGGsn)ZrgwwNM~mGmU&#2a_K>>1MI=g3K7W7LITmZfc8)XOp~pe{
zYDN?b1(N!`CfE|PVeHu?%Ng@7$l~p;@TJebOr}lzA@AE@`SkuTuB{~8r&^!e7}SdD
zXK31g*+{4+VIk;L(haqjoYtPyA@bCh1k1l^0P(=;>z4m&Kt)^V^5Fhbh~Mgd&&Ui>
zi(R5Ab}I!E=k4+(&1-St(UV0n_izy9iHtSm>xR)C`&rsuo$;HWxr~kaIB3s?x5r0z
z!sehi`>vVS!bV!NtjRzESc-WJ44EcEm5543_s$v^^ys_!M!X+>j!tHf=$uyj3B}I)
z&+DOANU5jq_B(k0=hJN~b1ED+NX*z8K*ZT`ZBw1&NqEfS*6@qz6cntGXc$SO0ek%(
zU)v}eXx;MDJ51*|wCc^B$bP->L`Akm_z@ZGUEa>>J?g>4><rcKKg#e*cJj|JwgkM|
ze(E_<s~*D{xprSCH{vyJymhX)7#@b`sxHKm@a&D`&Rg#*!S<*E=Vwt0Qhi-1R}HCn
zNZtL>8|gfhJt3lXq>_MTo-K#P^;_Xs(0XdQJ6%Iq{qyhqhZ1~cxbUvty9Sk-=6E;G
z_TcDdQV-<R!JE%vUS~hm!{6D08^;^zUWl{G>Gif${CF>^VLhP&&z6@L=bs>8$?I~{
zUk!~Yr7W?wrOyi<>$*p+3v|MZ43Svx+Ekn<JSP>nTm(|Fad$#WitzR0hejF?sMs9A
zAvS#G37p|kp|QtwqV+r0+9`HVoY*3;q3K2!>{8}CmLo)hyCSA%JR^$W{k}WjNb1cf
z%1~$>{gw)+Q@&jvf7A${x*osU(nbQRg(H5VYxtR}O`=iPsmOA`+Hm?t0m>fFw_|!k
zhA~pFq4I1h^rWd(Ii!?;MqG@<s$mamvz|Y1yeA87uDxSbd_h2B1ZOO5HUU{p+G3k*
z=y}k5=|^u^0}NmJ+BLk^2@{-R?oA9;@H5imsyL$+l;1EbK6@z*f|g$QEU|wC!IkTx
zVFm=)H`F&h%2bN-2YLf2NAvJRN6M1=a4e?&)KUE5S_~StT}#*R<)XA&W@t`FBAv7T
z68hnWFI4qC-g%la6*Be_0!D)}P|EBS+iQn>+<yKu+4TznS`)aG73kWI+{N0P{AUut
zh`H%lu2&7N>1nu2XH>!OJxm*}F?GY2u`<Qt;zZmoxT3PmpN9$2nKy1;r(n_^*Uf|#
zDym&NR+c1IfmVdc`X*)y@>zV-$c>}I%twU?QF0I5ZVP8~bSTFcC*3G_&J<$fa<S@y
zaXr+Fjh87M_yD2X&p#_Z*Mj<69{rZNQ~)Mc3%M7~QZd^mrc^ez4Ts6CBIWtruyYrI
zEt|gq)C-EKCl93JmY7D{M|6&g6yUw<JS!0nwSUiz+gXqE)y+Jw+xzgl;n|5VOTlpP
z?ujCb2o2tSd3aKuy#^*Mk3Yob5V)1Sn0euoJvQ|J6w%9V1EbKVTv_E+n4y$^$%Ul~
zGK$@I_@9ag>Xrei%a02%f;7@QuStgU!mcV_V-(z@;-~XhpaCp>w<oP$ZAE{FF6Quy
zMA-GsQ~LXPx_8Lbk*#T<8M?PF9sbhr79Q<9f0M17fZQ@Jphwrl*7ut*XeU<TC$Bxv
zY>p5>Y(_wapxp@m|16m(6}=GgNBPIl)ed+omSP%ETa6n=0zGDyUP2(lmeT1bt++mL
z{l%dI5>8C;!3k14o-E|qTT@YtH6y$IUpO?wh}G{u-)@kh{!`kOXu~deI8YHU64(dn
zQ<@CfQB@GVt@ecRPQ)6`vEZ}$MaVC`-)n5L3g>Mued8?hf*2zmQrINDUn(*+d-9eB
z6Ne-WrDn46<;Mq}_XrVakT&5E6IO*KY7OBsv4JRu8Rk1o;z9J6Px;xyG3dheo)*c|
z1O7X=o_t{1iE)~@Hn#ktfJF`$o796&$Qk|1*EUVX8wtra(>2X-#q&Ke`(QSNemi<B
ze<cGQaPP+SZZc#}eEr48TMSx5zck|aWnoH75PPN#6--UnnbbpQAS09L&O!=+%)h(J
z%vRgMu^?3c=EVlkZzdig{A&cZK);w=ay9Jt<9J)In2Y<xa=JLL)S>M4_TA5Srl7w9
z*XqRAc1&!sV)VE}$IGBJ4W)?7!1>@qxK=v7e$?D~7hj%(&)I4s8240T;FRRw#L*s%
znxbCiqjOM~5;EV&#ka%I<&R#63GHy|KvltXG6h=i<&OAQ#=+rO)497|RGd%|6yk92
z0o&=Tb34Onc#gWApRFYm?xn2g?#=Xv@`8<4kKIDhR*)t!=+T6P`q{x7<upJBBYCH<
zz1SPVc}+hx2OHxWZ|kV!;|beeY}e@Dx4uixBMD1?sOOvYKgH1xqE7d=h_kn2pHIN`
z_YZ4O;4yUE{Z01{zVQwEW7iJHty!hiQp@qgTKo0BAtIV*%z6Kiegy{~W_*wHDS?T;
z?N`N>8!@z#`+SW;0kB!#YWZ8)0RoymH?nN2;Ms)sAuBPuheKgO<M7Wu+$1a5z-38)
zj*C^ELw*nlNy%9m(>W_L+u0^g{Z24;jCd0Mq#T8rxGue>YaB#r@rsyV1z-?%`_fXe
zFN*DpJyK%R0*8x!^-dcXpy=HzHrZnN@H%7m+?$+c)IGj~sdQ&ICe#KlY_)8~y&(-R
zBZ^buYiVm>^}YtY@3(ohe~gF^#>BY>Y$?F8$LvSBK{IA<^8Wq3rV|vVVnwIu^W!K!
zduB3Q3r2U3>`<C&1HPflk|Eu6D&zPY`>ZMwUY**tPt>{%&R+Q}rY{$WZ@vGHe5CiW
zg(I8WJ+BgxO?1Spa-{{9jb8oUw1@6VIydKVov#)y*dz|!(yu~Q)!*K?>Gk*esHyL8
zPd6NTs}#0U*9wg$mkguQ^YNNKDe6<%M>uABjPIm*3GDLUG4EW|hj%kLg+}yhpzgpw
z)uVqY_;$lpk0(EqK&YTuWXD(qtf!cT6i2oJk!sWVhc^{$TS6X1sdVB)9s8LB2`*Sa
zxF$X5*N*a)kEAU|ThTD}BiFZErEqfO?@YN?8}1oeYz(L#!GT1wmF$)VSef}+Y-^bc
zAu36CHVLPqX!7&@R;*1JWFc|x4(CTq6TZYZMJ<7Q*X+0Mu}a3$ruCO*9~!V|io<Y;
z+zI~zmI4*MTF|-Zz0geRJm3S;LndPqs9SVgzw(6$KMqHPzjH;j%MX{XkthKf@fYh)
zmIkoLJyFJXv=4dDo|{)rq(U(>?*hw60^D-rKHz($0zDSbGS@OyqX%j6jc8RB)EAui
z!`NX1)qC=`&?JW7K}6$K###d8-C#8JI7G#eVsf&;gF-x@zmef}Z6pT$_Q<e2Z-b`7
z0`HS;<FMk+FT#Ph`Dp&_{w7*mA)F6M>dD=c0(Rx^t?vKJMD3P5X`XLXyvSp+RI#Rn
zITvSpudH^#(AVsraIF^1{j%AIRlD&IZIfP?YCittm`^`D6^w=PyCo7HQSgDwwkb{q
zS4cg+PfC4TC$MKU-q?8hJ@&pxFO%~2f{`te66#NyaVSQZpm^00O$ml(%!C+xwu5DM
z^+6wGHP6ppFla(<C5PHwc{Nz2%J$LKiH1w9GQWjxq~LA88@b8cwdmC{c|atV?#(=i
zd{^(bg5s^(CDCifsQJB{Ew?5gcgn{--SnUpBmAv;40o1;mdCGzi>x)sQf^z$6yAtQ
z8RrdiB8q|BJ5(n()rl43yKYX%y1?bjR>v=Fs{>}K&fUC@Bxs}g-eR*SLB_iy*Za;5
z;EnWA|D?KhT$@k{ZEX(#Ep^uCFG=}WyQH>el}PvQ1?nIBKGzOpMdOCTAI-qyuHaL(
zj+nn+#bS2>8EtrZF>^bKzUxBJA370_f*Hme&ORQ66;5jSt&(`W+8ErwL9+_0mG&J>
z5G14J`RcEB#(mJgDd^ASp(uPEELleqY=MgVflMjBG#KUF5u4*&f<-A`RD0s;(P0&f
zT0T;twnj~JpFR!5Rkoi^o1|Zhl0IFrqGH&1DeFV$G~iU-)qdWf0!6-VQ2D&O6`M_Y
zJw15nJVo$oZ5DmUrk;z@=IKr%`fQ8!lcIdYO=s_Nzkb~g@_CFroJB3zDm!O-HM$g<
z0ws)(gvUd{WwPg{Ec#By(=ZRG-d>Q9%|Gt6PQ?ojvyRze4Hzq7xgZjn2Z!y%GPivv
z<89H=9JZNANFI7))uWey(#0C(ud6$;wwOaDhAk1h1I|31_(<R7T9Mgvj;{tDuZEIO
zJ2k>*hEQ&4!yHHv*>t9`BNJ32I$1~Ww?bN4a@-$TGL*mX5zszDft1^lYnxf)fa9x*
zq|QbP{4#gnlR@7}z`MSsgY0DF?LE4G65By}t&Kw>yACC8ULTe*C4eEH+V0YtYGe+6
ztFdcyIoPs%emCGq_i*o8<J?|V1LWIraw+=&w(ZtlJu{h#0Yy=gPmdCChdqHo-l7b$
zY!b`AYZt@8<-MdjarzEZw|4cxk$R|5IRjiCjUZ=l-4w#vg9%%RmKpo$eJp9CQFTKF
zu*E%mr}C{CHjjVwpZrP2=QfWH5!2k@t3=t^Gsg;{>$gQCi);se7E1N-wyT4~&-o9$
zXexvD#yUrZpjKcPXHJ?rR0?J3r+xnHA|vsbfbeE|JtmwuBp<WTiFuPZ|6Uks1m9#E
z(R~8#up#m0jMaoEES6ZM2|Jg;O&v3#G8Qtjb0lf!R3)PtV`*aM(KzTmV`P18trseP
z_2j-8r2-Pz|K0K|!ee_r>&4PFf&4!IJW16gNZ`>eel=W%V{S}k4G%`ZxAjBQug)Hv
zdDm^T+E|X78(-Y9pJ~7^ZIWN*SsLNNM6gw`Mn1mE|NCI1F9#?$PVS9wZN^?Er=_UW
zHr)8krstyt!f=xLtxpV%;8Gx)XLF?pUJHKPOiFEny?WahTAJv)V~hFEK0Aj(OByF_
zT&4~tgZF*+{E`5BwXMQW)A=)NQ+DBg#t!`0-F&=sV=amv2ro8#UW7;Pz2C-okBlC{
z^Wo-%8hmsQa$82?VB7dL^`bi@oZgfF@-|-+Y@JZEeC<pC_7PVtwbdr5zU7(mE{Fgf
zcBzB6orVB>4qCl6%Y$lHGN(vVGZy~5Y8SL_isij`<k{YI!dCCWa{Y=BZ1~K3%$+|L
zswTbfKc@37(fYxhmAzHCE)=`Zdzb(pSlg#Qjg@2Hb8qj#p&oe9#`uY}g^WvAX%XWW
z6L76uZ<7O1ap_g8$Vm-)|0>|H%ylatrv&W$E*nvBQSzqt++q>Nvv>S^Pw#WgIK<4w
z>0UIaEj`O-a|G<2Tv_R|FM|_Ge)D|-o#^z|Lb$ED3zp*opElVQ<BHJds^|20_m3WY
zDQe#YOFp>?R+n0ksV2fus;dM8OkU;0Q*uEoK=&o{8U^%xmzx@WsK~k9{EcR719sWI
z7GC;P4r7^=$CG>MyIkh#riD&DX!<tL=07}mZ9NR4TepQF-xjgZR_97!hl8mVzdZ2q
zt;4_OeaH|mo|&jLSPaVKIZsh$A{f*?rP_+;fYz6w)y+{v;CR^?^YlqOTs^Co*?XZ3
z-iLCK1P#crcBSjWf%q)=tS>PaSxxs|+wonB`H_o%>U}3qi$}rkN;}nxoG|neVEt1|
zh=;FFOdU2y)`O|nHF95l3B(phCdvv`g11)4uFy9`v~?3s8dacs-L1KuxSg9IcWAbw
zna*oJ5punm_c#fSm&g{>!4^>A?9@1x<_F(>Nq_eE()-xeO5TI&ofv&_?Y5OiGd_NQ
zIZAkUC$g!9X$ck;;>n%YtCZUF@le*CDZK!7ApJY=C-+Y*-d-yiF<(hTImdnN`z|-*
zDP5)mN9<d$d-{n@3>^>5ZghuC=(VEVmUo<+O=>Z=@6o*9+akQEBUrwZryYML{nqJ;
zLHK$6)1>Sj3ixnO)_kPTnVILh6k7LNAkpvUV5kNKLR4jq-Lm}AVr=Qb3AqGxe0)X6
zo;43Mk1bv8ra#Z_^UHQ-wMLBZ;lDAVmW=+9o3bZUC{U0--|5iNfP8lrrhZ+l0~w)9
z`!CWxEe~%s?v^3sAmR7XQFg&__-qr=&50jTYEgDzs*vveaXcMlDd+|nnn~t5${|oY
zwwuB>Ooe%mnk<(zBJ5S`Uh}@&fL<%hq{%~_ApYI#_2y4m@G+vV!KgeIzw<e4BMY>G
zyyEe}w>;S}r|Hecxu*nKlRCniTnb>QqVsP_eFp@eS3Q6KE(shp1^mxm2nSU)=+d$v
zp~?8Unjpz)Jl-v;%)cW7TZm*Kxgs*CdEJwXAg7@elNHm}Ye|r}<6=)}aTzLJ8rk6<
z*oa(Sx|8hbiEtrB3Q~q@Fm?CL4j)+(bh6}6dk54*Io<d6$M^&G<Q*w2{1bwwr8ojU
zu4e$vs5BvKe-ZqUm5CCGZp9zcSKD*ONyr@fjXIqc3HxqSON<@pbJo^C14~{i3Wuh1
zll4dtW?$HuKhurvvc*X^X9{rPM~_FBbvF!dURp+n8jz^_5w<Ir22EkHj-!{7=^nhO
z;6r06I2_+6RDOiMbK>##!olSfJnq0(#QUTkYecscU&*Y8n<vZ(b(KYs=6NDqeF<U4
z<wSk#It5M~V~ae$Tn8P8vw2M%+rj?9>DsL?XmG7FO+{{d9h}&sSNmc$4MHh-w_ecl
zFp$lcIif8Czkf|QznV|Pd;LS=D#-+l5q&NrBaw|&nx`hYaw_1(X-Or)a5s?JZlt5^
zTmdJD1TBm3VmQBUDuF{C@W`T{$p5$ncZM<YaoV@SmR*sXtBIZ9eZl7o7ttTxl*$a&
z1*st2Zu9qosxKz8I+#}nRe}7VEYA;#bdZ=)K6~YI4J=2f8_KAJ!|s)-#Q^0l^c>NN
zk=>ApDdj3DcQ-bJyue)69j$8Is&>6|HM18q!`5$kgruSGp5-<Re!5PTV5)OtTL*L}
zWxJh8%!4&oCZ7Y>TcF*uoWUpF6XnYcFQq>vp=fn=Xzal@<lV1qoKjN+PToBB{f%U_
z?qU|p=g9*5w(WDsU4uV0mV~)Bl)#<WY(~E-??Cde|JBU3hC`W#VdrB96)CaNN$o}s
z+1=FcgTXj;U^ZKAY$9YG+C_9pCh6!jNkyrh${|XPL&9VpjbY3hM#$Q>5FHe1i7gTP
z&MqBn{c3;gH`jB$*L>f1&pq=!?>F<#``*uOSti_=4*B1`uvjE&#0mYO-v}<%qqmF2
z_25(2ut)pOIXiw1ILNp@mDEasfh-A*55FE8CZNwc^&4oE-LWprF$TjP3-pzivw^Dq
zbxk#+4(EKs98r2wh|`B#PWU^Euz88Y4C9@LaH(S_{pXges3H}iu0ms@dSpU7D>w~S
z1qJn&5SxOV*;d>3>?p-jN<?0`T?O!+{G>Oo6#?&wR<7CBB-o$p=fByK)GSX<?+U%0
zk7{w>@?8^)&~RE~X3TRg7$j|2U^BwU-d|(oeV!8yS*6!HYIr(MP@q(95G10Sn-4^5
z3t>vKccQ43i#-Fotn7P!0~&BXfA%a5BoqyHJl7@jH6}Xs8t1md-L^jK9)C8n-EF)b
zLIgO8amx3Lz+ccX*Kwcf-a`<R(J<8P&%}x>4@Uzs@8}%R=R45O!x!{nGo3*pSk+q(
zrRe{Ti?1h5TyL6#UL}ei*1QDBEj}^eVOR+=c}LuLtvw3An^W&Ki8v6iudRGJqYQlg
zj_%z)bO}@PFZA1RSi}z_#$b+RK5TR;Sen<(2RvDHKUFUnnEBt)!~QHo-Rmr9rB}m+
zcJ0+w6?Bv^l@B^$Kr}@03yXV?6~dz>cN#T+VM6($AIh7hn#lf@Vf*i11le+1W*hlG
z+bbBB*LUKg@ypDg_N0)S85=6!eXs-zHr|e#b+Z&$_9u;>MGH{1K<?9rpSi)xThkTP
z8L5z^^(9ay#G~DEuE#AhK2%rRd1Ts+a%grocXdiSia!iF8Kg|E03IdYE9PDSa#@Fr
zq-ZBmcTvD$&SoJPZT6AzWF-Nn)yP(6=Af_QV&|xUa&$=0uu9p@B)?xK?N!NaFqW&h
z`TeF$Jjgua>VCTvr1quC2F3EB<N~AMiu*a7o2SB(Hxc5?*}ADMI%K{3go$0kU@2}N
zIj!tkMQV&Zxwgy~M0$|kTeyOZza6|o1{8CJD7n^h#8x5`T}Q0-zNi(!x0L<M!i0^W
zY((q6Q6YrI67`{Ls+m|kdqz%+M+{!$RTNA;UW?SH4C!qah*c3CQ#3h3h<(^b_v@{K
zeb(2^*tBwpI_&2a8nPP=57ba&D|z5K^3%hF-3_pQ(EM`e;bM3h?r*iJn}cPQMRHzb
z9?1Uc<-n`10{n7cw5gA2C8TKT-br_1!zTg;Bs*V%rA-TGpPgR|qUY3NXP49Xt_EfQ
z4!v1p5{B0r6r=yO2BkVHk|w(_8Xm=~hVTkgp-*4$s?yg-aBgVi#rub%@POOG2Z5#+
zpl^ZNuB>5O)ReB)q18G;k<9HsEY_aIic?<&w?El}pUn&jo8Em7r!ex89u-AHYl}mi
z%c^8p(A}=uamx|o{M$4pMP*>ut3b8b#kt_C6Cw35t_1y;8GCNtu?J!TFMevqjf8m&
zsfIX{B*-v&aDD2rbEuIWZ16yC1KxbFX8y#c5O{i*TVNi*LWc3o-0<dDY)c=pbX)2R
zygIj?rMuF=rEhs=zeOUBNZC;~tvrtr5f&LQNs(-gVo8f*y(cc|(pX>8kPibM%t4hI
ztHvde()%>@?F4#naq}3jK>3f2Y_ttO%F0HGa>c!9!xo}Wef26(|Ev)0<gb{oD-b|1
zoT+UPvEdbW(F~{gt(YA-ZF!bo83tThu_{?~ACx4{CuALH$9cojPE#-EK%e=U!={gW
zQAcZTZ=0|VX?5&m%hoy~GAg%!+?fNZllnAIRgsmsI>$aY7O754+)7JR_yhX+9KUJP
z>u^obOvU3`wNPa49$_C#s^FDXoW1m$!D?^+YTB`EG%}X-t`t6osT?|`JfQ;vehf{&
zC*Ot*t0{HM1g%iBqEmIqkIWeH54L2B8ev;HduzB&Cs`r5b8%J^86h{*u4o8~p|xp~
zYEW1Q#AW=MtR7o9PDB*mGfR0affI=n(ccMtTLX6B9}gRevG$`(Yz#lvhDcj!CY@&{
zoYzG@hNP69^t-NLVzV!COc{O8;)ecr_xqPe*m$_ukx6{2I3~XXFC#BF3Lm>J<H2G#
zAi6dIlX#aqN@f)Nwl~mtxLAqzes1dg*YFQ<i&bFn=T^`i1;4Ec8_zBN27W(x74<dz
zL)_xK|NY<?TZw<V(i>N{@yz0$C+<|@m`Gx5qf4%|9|gbRm6a18Rv?cqk|0$F@w2}G
DZ9|>=

literal 0
HcmV?d00001

diff --git a/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md b/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md
new file mode 100644
index 0000000..3fa453d
--- /dev/null
+++ b/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md
@@ -0,0 +1,69 @@
+# Vulcan L40S: fixed-step SDE/ODE trajectory ensembles (2026-08-04)
+
+Environment: NVIDIA L40S (48 GB), jax 0.9.1 + CUDA-12 plugin (Alliance
+wheelhouse), Python 3.11, float32, TF32 matmul default. Workload:
+`benchmarks/gpu_trajectories.py` with the **trained kernels investment
+policies** (`benchmarks/policies/`, exported by `export_growth_policies.py`) —
+the neoclassical growth ODE under RK4 and the stochastic growth SDE
+(additive OU noise in log z) under SRA1/EulerMaruyama. Per-trajectory `x_0`
+and explicit presampled `noise=`, vmapped and jitted; grad mode is
+`jit(grad)` of a scalar loss with respect to `(x_0, policy weights, noise)`.
+Raw records: `vulcan-l40s-*.json` in this directory.
+
+## Headline numbers (kernels scale: B ≤ 32, `SaveAt(steps=True)`)
+
+Milliseconds per whole-ensemble call, `unroll=1` → `unroll=8`:
+
+| model / solver | B | n_steps | primal | grad |
+|---|---:|---:|---:|---:|
+| stochastic growth / SRA1 | 32 | 63 | 1.35 → 1.09 | 8.03 → 3.43 |
+| stochastic growth / SRA1 | 32 | 255 | 5.14 → 4.07 | 32.0 → 13.3 |
+| stochastic growth / EM | 32 | 255 | 3.07 → 2.19 | 18.7 → 7.1 |
+| neoclassical / RK4 | 32 | 63 | 2.11 → 1.62 | 14.3 → 6.0 |
+| neoclassical / RK4 | 32 | 255 | 8.25 → 6.13 | 57.1 → 25.1 |
+
+Cost is set almost entirely by `n_steps` (sequential scan latency), not by
+batch: times are flat in B from 1 to 32 — and in primal mode flat to
+**B = 1024** (SRA1 n=255: 4.49 ms at B=16 vs 5.22 ms at B=1024). Forward
+ensembles of hundreds of trajectories are effectively free; the grad-mode
+per-LM-iteration cost at the kernels training shape (B=16, n=63, steps
+saved, SRA1) is ~8 ms rolled and **~3.4 ms with `unroll=8`**.
+
+## Findings
+
+1. **`lax.scan` unroll is the one big lever: 2.0–3.1× faster grad, 1.2–1.7×
+   primal** at B ≤ 32 (per-iteration dispatch overhead dominates these small
+   kernels; the scan bodies are pure math — matmuls, elementwise, and their
+   transposes; no branches or bookkeeping to strip). Most of the gain is at
+   `unroll=4`; 8 adds a little more; compile time grows with the unroll
+   factor. Promoted to a real `unroll=` argument on `solve_sde` and
+   fixed-step `solve_ode`; `vulcan-l40s-g-unroll8-promoted` confirms the
+   argument reproduces the measured numbers.
+2. **Explicit `noise=` is free**: identical timing to in-solve `key=` draws
+   in every configuration, so first-class noise costs nothing.
+3. **`SaveAt(steps=True)` is free at this scale** (≤ 5% over endpoint-only),
+   so the collocation rollout shape carries no memory-traffic penalty.
+4. **`jax_default_matmul_precision="highest"` costs 15–35%** on these
+   width-32 policy matmuls versus the TF32 default. Use it for
+   reproducibility when needed, knowingly.
+5. **CUDA-graph command-buffer capture of the while loop
+   (`--xla_gpu_enable_command_buffer=...,WHILE`) had zero effect** — timings
+   identical to baseline. Unroll is the only dispatch-overhead remedy that
+   worked.
+6. **The reverse-mode memory wall is real but far from kernels scale**: the
+   plain-scan grad tape OOMs the 48 GB card at B=10,000 × n=1024 (SRA1,
+   ~45 GiB requested; `vulcan-l40s-a-baseline` records the failure). At
+   B ≤ 32 the tape is tens of MB. The benchmark's `--remat chunked` mode
+   (√n nested-scan checkpoint, endpoint-parity-verified) is the demonstrated
+   remedy if ensembles ever grow to B·n ≳ 10⁷ in grad mode; a package
+   `checkpoint_every=` option was considered and deliberately not added.
+
+## Reproduce
+
+```bash
+python -m benchmarks.gpu_trajectories \
+  --drifts neoclassical stochastic-growth --solvers sra1 em rk4 \
+  --batch 1 4 16 32 --dim 1 --n-steps 31 63 127 255 \
+  --modes primal grad --save t1 steps --output <prefix>
+# knobs: --scan-unroll 8 | --matmul-precision highest | --noise-modes key
+```
diff --git a/benchmarks/results/vulcan-l40s-a-baseline.json b/benchmarks/results/vulcan-l40s-a-baseline.json
new file mode 100644
index 0000000..12381d9
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-a-baseline.json
@@ -0,0 +1,1296 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 1,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006979090666087965,
+      "trajectories_per_second": 143285.14241247656,
+      "compile_seconds": 0.7613029390340671
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004477978091348301,
+      "trajectories_per_second": 22331.507202593395,
+      "compile_seconds": 2.764506147010252
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013014882979159419,
+      "trajectories_per_second": 76835.11266304033,
+      "compile_seconds": 0.44207365496549755
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008958171610720456,
+      "trajectories_per_second": 11162.992220457982,
+      "compile_seconds": 1.783474794938229
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002520166680608925,
+      "trajectories_per_second": 39679.915129993664,
+      "compile_seconds": 0.471507192007266
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01790090400027111,
+      "trajectories_per_second": 5586.310054424373,
+      "compile_seconds": 1.7965833969647065
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.019478568516205996,
+      "trajectories_per_second": 5133.847485599411,
+      "compile_seconds": 0.45416608604136854
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.15820773702580482,
+      "trajectories_per_second": 632.0803386732551,
+      "compile_seconds": 1.9938357799546793
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007358905841381504,
+      "trajectories_per_second": 1358897.6697822073,
+      "compile_seconds": 1.2876181559404358
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005326191778294742,
+      "trajectories_per_second": 187751.40693866726,
+      "compile_seconds": 4.128417236963287
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001382418572237449,
+      "trajectories_per_second": 723369.9113152804,
+      "compile_seconds": 0.4274732309859246
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01103723022970371,
+      "trajectories_per_second": 90602.44093747102,
+      "compile_seconds": 1.6469531509792432
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00268453383533698,
+      "trajectories_per_second": 372504.1520568034,
+      "compile_seconds": 0.45855261699762195
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.02228194155031815,
+      "trajectories_per_second": 44879.39247761475,
+      "compile_seconds": 1.7092685449169949
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.02078921248903498,
+      "trajectories_per_second": 48101.87016595448,
+      "compile_seconds": 0.4331063599092886
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.18087553698569536,
+      "trajectories_per_second": 5528.66361402474,
+      "compile_seconds": 1.86476408992894
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009766825313242723,
+      "trajectories_per_second": 10238741.535021741,
+      "compile_seconds": 0.5928605200024322
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.011672457243548706,
+      "trajectories_per_second": 856717.638055769,
+      "compile_seconds": 2.0495360889472067
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0018726597683360944,
+      "trajectories_per_second": 5339998.31100406,
+      "compile_seconds": 0.44451983796898276
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.023683579463977367,
+      "trajectories_per_second": 422233.4725715748,
+      "compile_seconds": 1.7391729099908844
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003682566538023261,
+      "trajectories_per_second": 2715497.438199129,
+      "compile_seconds": 0.4114024490118027
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.04794425202999264,
+      "trajectories_per_second": 208575.5763536423,
+      "compile_seconds": 1.8183223910164088
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.029412479023449123,
+      "trajectories_per_second": 339991.74268947175,
+      "compile_seconds": 0.4274300360120833
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": true
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002987449995998759,
+      "trajectories_per_second": 33473363.615770977,
+      "compile_seconds": 0.6725022579776123
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.09317554696463048,
+      "trajectories_per_second": 1073242.961889562,
+      "compile_seconds": 4.375617903075181
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00600924136233516,
+      "trajectories_per_second": 16641035.693254383,
+      "compile_seconds": 0.4181409140583128
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.19127922097686678,
+      "trajectories_per_second": 522795.9393043218,
+      "compile_seconds": 2.0444391909986734
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.012165952502982691,
+      "trajectories_per_second": 8219660.562991948,
+      "compile_seconds": 0.49205413402523845
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": true
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.09763287496753037,
+      "trajectories_per_second": 1024245.1636629247,
+      "compile_seconds": 0.571541424957104
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": true
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0004140019734305414,
+      "trajectories_per_second": 241544.74233871582,
+      "compile_seconds": 0.26524597394745797
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0026054690550903186,
+      "trajectories_per_second": 38380.80510095849,
+      "compile_seconds": 1.145113329985179
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007423295588153651,
+      "trajectories_per_second": 134711.0576595972,
+      "compile_seconds": 0.24592959508299828
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0051357642255930435,
+      "trajectories_per_second": 19471.298838383234,
+      "compile_seconds": 1.1722897390136495
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0014299512635368634,
+      "trajectories_per_second": 69932.45332897463,
+      "compile_seconds": 0.2810975379543379
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01015592849580571,
+      "trajectories_per_second": 9846.465543874096,
+      "compile_seconds": 1.1618733649374917
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01042707500164397,
+      "trajectories_per_second": 9590.417253566666,
+      "compile_seconds": 0.2653352909255773
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.08733745594508946,
+      "trajectories_per_second": 1144.9841184161776,
+      "compile_seconds": 1.2963226590072736
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00045979170171522545,
+      "trajectories_per_second": 2174897.885867796,
+      "compile_seconds": 0.27072585409041494
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0028724446413381133,
+      "trajectories_per_second": 348135.5169073529,
+      "compile_seconds": 1.2190662970533594
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008201496548612009,
+      "trajectories_per_second": 1219289.6675293196,
+      "compile_seconds": 0.2859465579967946
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00596189200587105,
+      "trajectories_per_second": 167731.98827070955,
+      "compile_seconds": 1.144728722050786
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0015455124666914344,
+      "trajectories_per_second": 647034.573678177,
+      "compile_seconds": 0.28703883197158575
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.012067058734828606,
+      "trajectories_per_second": 82870.23557064036,
+      "compile_seconds": 1.1576002539368346
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.011439821013482288,
+      "trajectories_per_second": 87413.95506288603,
+      "compile_seconds": 0.2860855530016124
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.09809419699013233,
+      "trajectories_per_second": 10194.282951320696,
+      "compile_seconds": 1.2522407149663195
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0005777520452474439,
+      "trajectories_per_second": 17308463.175958343,
+      "compile_seconds": 0.24972683808300644
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006051044125342742,
+      "trajectories_per_second": 1652607.35054276,
+      "compile_seconds": 1.1796769490465522
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001048030680976808,
+      "trajectories_per_second": 9541705.392326478,
+      "compile_seconds": 0.2818208369426429
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.012286034994758666,
+      "trajectories_per_second": 813932.2413021036,
+      "compile_seconds": 1.1883608979405835
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0020266218293134284,
+      "trajectories_per_second": 4934319.691694905,
+      "compile_seconds": 0.29116004006937146
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.024754316080361605,
+      "trajectories_per_second": 403969.956897065,
+      "compile_seconds": 1.2487655740696937
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.015777506011848647,
+      "trajectories_per_second": 633813.7340901765,
+      "compile_seconds": 0.2851269660750404
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.19926861603744328,
+      "trajectories_per_second": 50183.51709795066,
+      "compile_seconds": 1.4213094549486414
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0015345037681981921,
+      "trajectories_per_second": 65167647.07421969,
+      "compile_seconds": 0.2797583380015567
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.04465881793294102,
+      "trajectories_per_second": 2239199.4376151743,
+      "compile_seconds": 1.4061723200138658
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00315982960164547,
+      "trajectories_per_second": 31647276.15309552,
+      "compile_seconds": 0.3583576310193166
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.09023767802864313,
+      "trajectories_per_second": 1108184.5431379355,
+      "compile_seconds": 1.484362359973602
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006186905716146741,
+      "trajectories_per_second": 16163168.56728841,
+      "compile_seconds": 0.3430421269731596
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.18328238592948765,
+      "trajectories_per_second": 545606.1666420688,
+      "compile_seconds": 1.5297465489711612
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.049338942975737154,
+      "trajectories_per_second": 2026796.5620823263,
+      "compile_seconds": 0.3346744281006977
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 100000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": true
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001035444549329224,
+      "trajectories_per_second": 96576.87614926502,
+      "compile_seconds": 0.5566155769629404
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007564925181213766,
+      "trajectories_per_second": 13218.901390899855,
+      "compile_seconds": 2.647828501998447
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001985040695770927,
+      "trajectories_per_second": 50376.800945717216,
+      "compile_seconds": 0.4893043019110337
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01516536200263848,
+      "trajectories_per_second": 6593.973818930398,
+      "compile_seconds": 2.6690926519222558
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003880432087250731,
+      "trajectories_per_second": 25770.32602337065,
+      "compile_seconds": 0.47296555899083614
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.031245502992533147,
+      "trajectories_per_second": 3200.460559841119,
+      "compile_seconds": 2.734248523018323
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.030213839025236666,
+      "trajectories_per_second": 3309.741602729569,
+      "compile_seconds": 0.4995520890224725
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 100,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.26694091293029487,
+      "trajectories_per_second": 374.614737404875,
+      "compile_seconds": 3.0433065260294825
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011527200516026753,
+      "trajectories_per_second": 867513.3208706293,
+      "compile_seconds": 0.49437285808380693
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.010639087238814682,
+      "trajectories_per_second": 93993.02567533149,
+      "compile_seconds": 2.76271853398066
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021989913181063125,
+      "trajectories_per_second": 454753.95549135766,
+      "compile_seconds": 0.5211080589797348
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.021824156516231596,
+      "trajectories_per_second": 45820.785754366065,
+      "compile_seconds": 2.7108210891019553
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004315479097633876,
+      "trajectories_per_second": 231723.9818281793,
+      "compile_seconds": 0.49992516089696437
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.04420895501971245,
+      "trajectories_per_second": 22619.85155618601,
+      "compile_seconds": 2.7040379130048677
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03372074698563665,
+      "trajectories_per_second": 29655.33356737174,
+      "compile_seconds": 0.5715859179617837
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1000,
+      "dim": 2,
+      "n_steps": 1024,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.35841227206401527,
+      "trajectories_per_second": 2790.0830354977134,
+      "compile_seconds": 3.108547538984567
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 10000,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001657568448191059,
+      "trajectories_per_second": 6032933.367495756,
+      "compile_seconds": 0.48741554596927017
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-a-baseline.md b/benchmarks/results/vulcan-l40s-a-baseline.md
new file mode 100644
index 0000000..c827ecf
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-a-baseline.md
@@ -0,0 +1,87 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=1, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | primal | 0.698 | 1.43e+05 | 0.76 |
+| sra1 | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | grad | 4.478 | 2.23e+04 | 2.76 |
+| sra1 | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | primal | 1.301 | 7.68e+04 | 0.44 |
+| sra1 | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | grad | 8.958 | 1.12e+04 | 1.78 |
+| sra1 | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | primal | 2.520 | 3.97e+04 | 0.47 |
+| sra1 | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | grad | 17.901 | 5.59e+03 | 1.80 |
+| sra1 | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | primal | 19.479 | 5.13e+03 | 0.45 |
+| sra1 | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | grad | 158.208 | 632 | 1.99 |
+| sra1 | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | primal | 0.736 | 1.36e+06 | 1.29 |
+| sra1 | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | grad | 5.326 | 1.88e+05 | 4.13 |
+| sra1 | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | primal | 1.382 | 7.23e+05 | 0.43 |
+| sra1 | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | grad | 11.037 | 9.06e+04 | 1.65 |
+| sra1 | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | primal | 2.685 | 3.73e+05 | 0.46 |
+| sra1 | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | grad | 22.282 | 4.49e+04 | 1.71 |
+| sra1 | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | primal | 20.789 | 4.81e+04 | 0.43 |
+| sra1 | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | grad | 180.876 | 5.53e+03 | 1.86 |
+| sra1 | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | primal | 0.977 | 1.02e+07 | 0.59 |
+| sra1 | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | grad | 11.672 | 8.57e+05 | 2.05 |
+| sra1 | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | primal | 1.873 | 5.34e+06 | 0.44 |
+| sra1 | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | grad | 23.684 | 4.22e+05 | 1.74 |
+| sra1 | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | primal | 3.683 | 2.72e+06 | 0.41 |
+| sra1 | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | grad | 47.944 | 2.09e+05 | 1.82 |
+| sra1 | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | primal | 29.412 | 3.4e+05 | 0.43 |
+| sra1 | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | grad | OOM | — | — |
+| sra1 | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | primal | 2.987 | 3.35e+07 | 0.67 |
+| sra1 | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | grad | 93.176 | 1.07e+06 | 4.38 |
+| sra1 | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | primal | 6.009 | 1.66e+07 | 0.42 |
+| sra1 | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | grad | 191.279 | 5.23e+05 | 2.04 |
+| sra1 | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | primal | 12.166 | 8.22e+06 | 0.49 |
+| sra1 | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | grad | OOM | — | — |
+| sra1 | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | primal | 97.633 | 1.02e+06 | 0.57 |
+| sra1 | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | grad | OOM | — | — |
+| em | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | primal | 0.414 | 2.42e+05 | 0.27 |
+| em | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | grad | 2.605 | 3.84e+04 | 1.15 |
+| em | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | primal | 0.742 | 1.35e+05 | 0.25 |
+| em | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | grad | 5.136 | 1.95e+04 | 1.17 |
+| em | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | primal | 1.430 | 6.99e+04 | 0.28 |
+| em | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | grad | 10.156 | 9.85e+03 | 1.16 |
+| em | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | primal | 10.427 | 9.59e+03 | 0.27 |
+| em | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | grad | 87.337 | 1.14e+03 | 1.30 |
+| em | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | primal | 0.460 | 2.17e+06 | 0.27 |
+| em | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | grad | 2.872 | 3.48e+05 | 1.22 |
+| em | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | primal | 0.820 | 1.22e+06 | 0.29 |
+| em | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | grad | 5.962 | 1.68e+05 | 1.14 |
+| em | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | primal | 1.546 | 6.47e+05 | 0.29 |
+| em | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | grad | 12.067 | 8.29e+04 | 1.16 |
+| em | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | primal | 11.440 | 8.74e+04 | 0.29 |
+| em | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | grad | 98.094 | 1.02e+04 | 1.25 |
+| em | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | primal | 0.578 | 1.73e+07 | 0.25 |
+| em | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | grad | 6.051 | 1.65e+06 | 1.18 |
+| em | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | primal | 1.048 | 9.54e+06 | 0.28 |
+| em | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | grad | 12.286 | 8.14e+05 | 1.19 |
+| em | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | primal | 2.027 | 4.93e+06 | 0.29 |
+| em | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | grad | 24.754 | 4.04e+05 | 1.25 |
+| em | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | primal | 15.778 | 6.34e+05 | 0.29 |
+| em | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | grad | 199.269 | 5.02e+04 | 1.42 |
+| em | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | primal | 1.535 | 6.52e+07 | 0.28 |
+| em | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | grad | 44.659 | 2.24e+06 | 1.41 |
+| em | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | primal | 3.160 | 3.16e+07 | 0.36 |
+| em | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | grad | 90.238 | 1.11e+06 | 1.48 |
+| em | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | primal | 6.187 | 1.62e+07 | 0.34 |
+| em | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | grad | 183.282 | 5.46e+05 | 1.53 |
+| em | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | primal | 49.339 | 2.03e+06 | 0.33 |
+| em | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | grad | OOM | — | — |
+| rk4 | neoclassical | 100 | 2 | 31 | float32 | none | t1 | primal | 1.035 | 9.66e+04 | 0.56 |
+| rk4 | neoclassical | 100 | 2 | 31 | float32 | none | t1 | grad | 7.565 | 1.32e+04 | 2.65 |
+| rk4 | neoclassical | 100 | 2 | 63 | float32 | none | t1 | primal | 1.985 | 5.04e+04 | 0.49 |
+| rk4 | neoclassical | 100 | 2 | 63 | float32 | none | t1 | grad | 15.165 | 6.59e+03 | 2.67 |
+| rk4 | neoclassical | 100 | 2 | 127 | float32 | none | t1 | primal | 3.880 | 2.58e+04 | 0.47 |
+| rk4 | neoclassical | 100 | 2 | 127 | float32 | none | t1 | grad | 31.246 | 3.2e+03 | 2.73 |
+| rk4 | neoclassical | 100 | 2 | 1024 | float32 | none | t1 | primal | 30.214 | 3.31e+03 | 0.50 |
+| rk4 | neoclassical | 100 | 2 | 1024 | float32 | none | t1 | grad | 266.941 | 375 | 3.04 |
+| rk4 | neoclassical | 1000 | 2 | 31 | float32 | none | t1 | primal | 1.153 | 8.68e+05 | 0.49 |
+| rk4 | neoclassical | 1000 | 2 | 31 | float32 | none | t1 | grad | 10.639 | 9.4e+04 | 2.76 |
+| rk4 | neoclassical | 1000 | 2 | 63 | float32 | none | t1 | primal | 2.199 | 4.55e+05 | 0.52 |
+| rk4 | neoclassical | 1000 | 2 | 63 | float32 | none | t1 | grad | 21.824 | 4.58e+04 | 2.71 |
+| rk4 | neoclassical | 1000 | 2 | 127 | float32 | none | t1 | primal | 4.315 | 2.32e+05 | 0.50 |
+| rk4 | neoclassical | 1000 | 2 | 127 | float32 | none | t1 | grad | 44.209 | 2.26e+04 | 2.70 |
+| rk4 | neoclassical | 1000 | 2 | 1024 | float32 | none | t1 | primal | 33.721 | 2.97e+04 | 0.57 |
+| rk4 | neoclassical | 1000 | 2 | 1024 | float32 | none | t1 | grad | 358.412 | 2.79e+03 | 3.11 |
+| rk4 | neoclassical | 10000 | 2 | 31 | float32 | none | t1 | primal | 1.658 | 6.03e+06 | 0.49 |
diff --git a/benchmarks/results/vulcan-l40s-a-kernels-scale.json b/benchmarks/results/vulcan-l40s-a-kernels-scale.json
new file mode 100644
index 0000000..0c17191
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-a-kernels-scale.json
@@ -0,0 +1,3084 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 1,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00044518015159275196,
+      "trajectories_per_second": 2246.28163771954,
+      "compile_seconds": 0.46882447390817106
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002341831762654086,
+      "trajectories_per_second": 427.01615715838716,
+      "compile_seconds": 1.4274521339684725
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0005290067263508904,
+      "trajectories_per_second": 1890.3351322166357,
+      "compile_seconds": 0.33365265605971217
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0024705242766584787,
+      "trajectories_per_second": 404.7723835171357,
+      "compile_seconds": 1.284887518035248
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.000821885879649299,
+      "trajectories_per_second": 1216.713931655211,
+      "compile_seconds": 0.3873317010002211
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004553293704520911,
+      "trajectories_per_second": 219.62123791995054,
+      "compile_seconds": 1.376192542957142
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009717659146796601,
+      "trajectories_per_second": 1029.0544100115378,
+      "compile_seconds": 0.34592496894765645
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004817036399617791,
+      "trajectories_per_second": 207.59652139629776,
+      "compile_seconds": 1.2965855040820315
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0015334397734653566,
+      "trajectories_per_second": 652.128643918073,
+      "compile_seconds": 0.3606032660463825
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00902233279775828,
+      "trajectories_per_second": 110.83608002671586,
+      "compile_seconds": 1.3879702399717644
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0018500350834801792,
+      "trajectories_per_second": 540.5302899006962,
+      "compile_seconds": 0.34839765599463135
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.009488765196874738,
+      "trajectories_per_second": 105.38779063996276,
+      "compile_seconds": 1.3026674310676754
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003001927128934767,
+      "trajectories_per_second": 333.119345356944,
+      "compile_seconds": 0.3560068689985201
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01782353053567931,
+      "trajectories_per_second": 56.1056070231535,
+      "compile_seconds": 1.3728208880638704
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0035638541538411607,
+      "trajectories_per_second": 280.5950964413595,
+      "compile_seconds": 0.3729558620834723
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01892375253373757,
+      "trajectories_per_second": 52.84364177861573,
+      "compile_seconds": 1.2915878460044041
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006183858495205641,
+      "trajectories_per_second": 6468.453317780814,
+      "compile_seconds": 0.5545572639675811
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004215102452276783,
+      "trajectories_per_second": 948.9686301312568,
+      "compile_seconds": 1.6543556730030105
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006971375533699323,
+      "trajectories_per_second": 5737.7485700837315,
+      "compile_seconds": 0.37709420709870756
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004352877636186101,
+      "trajectories_per_second": 918.9323326590718,
+      "compile_seconds": 1.3398378629935905
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011719442831928887,
+      "trajectories_per_second": 3413.131543337752,
+      "compile_seconds": 0.3531762290513143
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008354840008541942,
+      "trajectories_per_second": 478.76440433454405,
+      "compile_seconds": 1.3873992669396102
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013069192043923099,
+      "trajectories_per_second": 3060.632965340743,
+      "compile_seconds": 0.3521754309767857
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008634187793359161,
+      "trajectories_per_second": 463.2746120111647,
+      "compile_seconds": 1.3849235730012879
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0022584000485949217,
+      "trajectories_per_second": 1771.1653887399737,
+      "compile_seconds": 0.3622201889520511
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.016870173043571413,
+      "trajectories_per_second": 237.10485895248414,
+      "compile_seconds": 1.391697199898772
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002554644420007734,
+      "trajectories_per_second": 1565.7756393306158,
+      "compile_seconds": 0.3469811760587618
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.017001587024424225,
+      "trajectories_per_second": 235.27215396148958,
+      "compile_seconds": 1.3942796170013025
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004364936189218001,
+      "trajectories_per_second": 916.3936943409518,
+      "compile_seconds": 0.34337302402127534
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03425101703032851,
+      "trajectories_per_second": 116.78485332152587,
+      "compile_seconds": 1.4128429669653997
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0050102312298905514,
+      "trajectories_per_second": 798.3663460752848,
+      "compile_seconds": 0.3538889540359378
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03452990192454308,
+      "trajectories_per_second": 115.84162644716027,
+      "compile_seconds": 1.4081455019768327
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006446458275119463,
+      "trajectories_per_second": 24819.83023415054,
+      "compile_seconds": 0.491259106900543
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003999261545356025,
+      "trajectories_per_second": 4000.7385909979635,
+      "compile_seconds": 1.5653145679971203
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007215623944913121,
+      "trajectories_per_second": 22174.104584926015,
+      "compile_seconds": 0.35081532900221646
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004105277084322138,
+      "trajectories_per_second": 3897.422676073012,
+      "compile_seconds": 1.3986358649563044
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001194223682055386,
+      "trajectories_per_second": 13397.825081196093,
+      "compile_seconds": 0.3312612830195576
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007826191159741333,
+      "trajectories_per_second": 2044.4172233238453,
+      "compile_seconds": 1.4205455440096557
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013475033638039322,
+      "trajectories_per_second": 11873.810804325438,
+      "compile_seconds": 0.35316630301531404
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008141389000229537,
+      "trajectories_per_second": 1965.2666147691625,
+      "compile_seconds": 1.3912285269470885
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002282388186791823,
+      "trajectories_per_second": 7010.201022153889,
+      "compile_seconds": 0.36521074501797557
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.015496470034122467,
+      "trajectories_per_second": 1032.4932042438559,
+      "compile_seconds": 1.3830204929690808
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002586590333117379,
+      "trajectories_per_second": 6185.749554208174,
+      "compile_seconds": 0.35045432997867465
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.016209229012019932,
+      "trajectories_per_second": 987.0919824832644,
+      "compile_seconds": 1.427961067063734
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004546631302218884,
+      "trajectories_per_second": 3519.0889554188284,
+      "compile_seconds": 0.343484070035629
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03083732898812741,
+      "trajectories_per_second": 518.8516815499848,
+      "compile_seconds": 1.4297372420551255
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00505759821842528,
+      "trajectories_per_second": 3163.5569511454223,
+      "compile_seconds": 0.3520263599930331
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03187646996229887,
+      "trajectories_per_second": 501.93763672463155,
+      "compile_seconds": 1.4496542459819466
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006498913063357274,
+      "trajectories_per_second": 49239.00302717562,
+      "compile_seconds": 0.6055959139484912
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003941296415481095,
+      "trajectories_per_second": 8119.155888480393,
+      "compile_seconds": 1.6412975949933752
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007238080008671834,
+      "trajectories_per_second": 44210.61933780959,
+      "compile_seconds": 0.407810817938298
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004049962997669354,
+      "trajectories_per_second": 7901.306757225967,
+      "compile_seconds": 1.5103212649701163
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0012194689998805926,
+      "trajectories_per_second": 26240.929456290698,
+      "compile_seconds": 0.43304388399701566
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007830712827853858,
+      "trajectories_per_second": 4086.47344162283,
+      "compile_seconds": 1.4790805460652336
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013507371259038337,
+      "trajectories_per_second": 23690.768089747653,
+      "compile_seconds": 0.4133943960769102
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008031683488904187,
+      "trajectories_per_second": 3984.220748266309,
+      "compile_seconds": 1.5542903270106763
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00238833566462355,
+      "trajectories_per_second": 13398.451680804192,
+      "compile_seconds": 0.4024871389847249
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.015544181650814911,
+      "trajectories_per_second": 2058.648098616525,
+      "compile_seconds": 1.527874932042323
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002630457443754292,
+      "trajectories_per_second": 12165.184453365779,
+      "compile_seconds": 0.4074179329909384
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0159813833112518,
+      "trajectories_per_second": 2002.3297969125229,
+      "compile_seconds": 1.5075764500070363
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00456721659284085,
+      "trajectories_per_second": 7006.455540155522,
+      "compile_seconds": 0.4330545160919428
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.031249569030478597,
+      "trajectories_per_second": 1024.014122204037,
+      "compile_seconds": 1.547785917064175
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00513628844378723,
+      "trajectories_per_second": 6230.1797008123,
+      "compile_seconds": 0.4154914850369096
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03203625499736518,
+      "trajectories_per_second": 998.8683134976869,
+      "compile_seconds": 1.5479462939547375
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00030696436657871997,
+      "trajectories_per_second": 3257.7071115632352,
+      "compile_seconds": 0.18483830604236573
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013746748898281818,
+      "trajectories_per_second": 727.4447270401428,
+      "compile_seconds": 0.7391808689571917
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0003924223752359727,
+      "trajectories_per_second": 2548.274673172438,
+      "compile_seconds": 0.22195479401852936
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0014773202728423657,
+      "trajectories_per_second": 676.901291062634,
+      "compile_seconds": 0.7877419149735942
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0005372288382008068,
+      "trajectories_per_second": 1861.4041706119608,
+      "compile_seconds": 0.19184806395787746
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0026184850010193056,
+      "trajectories_per_second": 381.90022078061435,
+      "compile_seconds": 0.7403571859467775
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006974619644661916,
+      "trajectories_per_second": 1433.7699415126651,
+      "compile_seconds": 0.22224036301486194
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0028671830043415815,
+      "trajectories_per_second": 348.7743888289543,
+      "compile_seconds": 0.7475595809519291
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00101398247749206,
+      "trajectories_per_second": 986.2103361720375,
+      "compile_seconds": 0.19123999297153205
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005112051782715652,
+      "trajectories_per_second": 195.6161718433874,
+      "compile_seconds": 0.74928293004632
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013029114411705557,
+      "trajectories_per_second": 767.5118725656324,
+      "compile_seconds": 0.2219563260441646
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005600566626526415,
+      "trajectories_per_second": 178.55336195156028,
+      "compile_seconds": 0.7629783579614013
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0019393865416835372,
+      "trajectories_per_second": 515.6269668304096,
+      "compile_seconds": 0.19201537396293133
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.010114373755641282,
+      "trajectories_per_second": 98.86919587505366,
+      "compile_seconds": 0.7507086748955771
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0025219939480949608,
+      "trajectories_per_second": 396.511657276327,
+      "compile_seconds": 0.2333108299644664
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.010971961746690795,
+      "trajectories_per_second": 91.1414041615307,
+      "compile_seconds": 0.7984069479862228
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0003989394178491209,
+      "trajectories_per_second": 10026.585042826735,
+      "compile_seconds": 0.2265496940817684
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002401275502052158,
+      "trajectories_per_second": 1665.7813718507323,
+      "compile_seconds": 0.9170253330375999
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00047566794711066055,
+      "trajectories_per_second": 8409.227538448014,
+      "compile_seconds": 0.24313144909683615
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0025512944428353673,
+      "trajectories_per_second": 1567.8315810364177,
+      "compile_seconds": 0.9244248060276732
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007076707728573996,
+      "trajectories_per_second": 5652.345911996605,
+      "compile_seconds": 0.21988887200132012
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0046742793987505134,
+      "trajectories_per_second": 855.7468774907303,
+      "compile_seconds": 0.9217888719867915
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.000854283126748421,
+      "trajectories_per_second": 4682.288429627342,
+      "compile_seconds": 0.24327497300691903
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005006728784388138,
+      "trajectories_per_second": 798.9248413999784,
+      "compile_seconds": 0.9286532551050186
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013328823331517704,
+      "trajectories_per_second": 3001.0150937641206,
+      "compile_seconds": 0.22099014394916594
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.009231478604488075,
+      "trajectories_per_second": 433.30003473715647,
+      "compile_seconds": 0.9346064699348062
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0016363451431971043,
+      "trajectories_per_second": 2444.4720703511043,
+      "compile_seconds": 0.2426328459987417
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00998435850488022,
+      "trajectories_per_second": 400.62663996338415,
+      "compile_seconds": 0.9612002939684317
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0025957618398885977,
+      "trajectories_per_second": 1540.9734200313492,
+      "compile_seconds": 0.22377602802589536
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01833497901679948,
+      "trajectories_per_second": 218.16223494638243,
+      "compile_seconds": 0.9413322400068864
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0031615220631162327,
+      "trajectories_per_second": 1265.213375122646,
+      "compile_seconds": 0.24617118400055915
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01979129499522969,
+      "trajectories_per_second": 202.10905860198247,
+      "compile_seconds": 0.9616675300057977
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0003946929003347407,
+      "trajectories_per_second": 40537.84597197044,
+      "compile_seconds": 0.21592206100467592
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002193548297509551,
+      "trajectories_per_second": 7294.117944959603,
+      "compile_seconds": 0.9024808730464429
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00048250114121354924,
+      "trajectories_per_second": 33160.5433300283,
+      "compile_seconds": 0.23947630892507732
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002308497050156196,
+      "trajectories_per_second": 6930.916372155389,
+      "compile_seconds": 0.9371393930632621
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007099968506329096,
+      "trajectories_per_second": 22535.311228123315,
+      "compile_seconds": 0.21416236797813326
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004313542645170607,
+      "trajectories_per_second": 3709.2481322546832,
+      "compile_seconds": 0.9356112200766802
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008749259607296209,
+      "trajectories_per_second": 18287.261686299982,
+      "compile_seconds": 0.2405970450490713
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004538466199301183,
+      "trajectories_per_second": 3525.4201083316702,
+      "compile_seconds": 0.9290328299393877
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013390144723577097,
+      "trajectories_per_second": 11949.086683004645,
+      "compile_seconds": 0.2172669160645455
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00847135940566659,
+      "trajectories_per_second": 1888.7169383106818,
+      "compile_seconds": 0.9089790439466015
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001627106498926878,
+      "trajectories_per_second": 9833.406731859559,
+      "compile_seconds": 0.2721616820199415
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008918875199742615,
+      "trajectories_per_second": 1793.9481875990075,
+      "compile_seconds": 0.920522270957008
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0025571096145237484,
+      "trajectories_per_second": 6257.064581480578,
+      "compile_seconds": 0.21738911396823823
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.016777171520516276,
+      "trajectories_per_second": 953.6768447788772,
+      "compile_seconds": 0.929095889092423
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003161054065761467,
+      "trajectories_per_second": 5061.602765135166,
+      "compile_seconds": 0.2428266960196197
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.017750939005054533,
+      "trajectories_per_second": 901.3607671934449,
+      "compile_seconds": 0.9496175439562649
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00038076739785877073,
+      "trajectories_per_second": 84040.80858800055,
+      "compile_seconds": 0.2364281170303002
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002288160718134826,
+      "trajectories_per_second": 13985.031622291164,
+      "compile_seconds": 0.9443110490683466
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00046480763727102916,
+      "trajectories_per_second": 68845.68461025698,
+      "compile_seconds": 0.2578090919414535
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0024021743171799338,
+      "trajectories_per_second": 13321.264727185515,
+      "compile_seconds": 0.9402767500141636
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006878971490853776,
+      "trajectories_per_second": 46518.582091155535,
+      "compile_seconds": 0.23492574703413993
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004524712641300125,
+      "trajectories_per_second": 7072.272326846631,
+      "compile_seconds": 0.924483462003991
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.000842625876040464,
+      "trajectories_per_second": 37976.52185851378,
+      "compile_seconds": 0.25491039804182947
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004737628798466176,
+      "trajectories_per_second": 6754.4337813802795,
+      "compile_seconds": 0.970670091919601
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0012932961076699398,
+      "trajectories_per_second": 24742.98021174179,
+      "compile_seconds": 0.23251260293181986
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00888781079556793,
+      "trajectories_per_second": 3600.436680758032,
+      "compile_seconds": 0.9200960400048643
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001609387603821233,
+      "trajectories_per_second": 19883.339429246957,
+      "compile_seconds": 0.25817957904655486
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.009383920207619667,
+      "trajectories_per_second": 3410.0886721112847,
+      "compile_seconds": 0.9416598289972171
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002494488738624281,
+      "trajectories_per_second": 12828.280001635969,
+      "compile_seconds": 0.2375009689712897
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.017782647511921823,
+      "trajectories_per_second": 1799.5070744413392,
+      "compile_seconds": 0.9460055059753358
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0030722367446287535,
+      "trajectories_per_second": 10415.863964893388,
+      "compile_seconds": 0.2579932389780879
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01872129802359268,
+      "trajectories_per_second": 1709.2831896417347,
+      "compile_seconds": 0.9482306479476392
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007537903852368656,
+      "trajectories_per_second": 1326.6287546049919,
+      "compile_seconds": 0.5394193780375645
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00450661373113028,
+      "trajectories_per_second": 221.8960975271327,
+      "compile_seconds": 2.2642879619961604
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009075410945235559,
+      "trajectories_per_second": 1101.8784780484057,
+      "compile_seconds": 0.7249034650158137
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004634865303523838,
+      "trajectories_per_second": 215.75600033936493,
+      "compile_seconds": 2.511369476094842
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0014247972079936195,
+      "trajectories_per_second": 701.85426697192,
+      "compile_seconds": 0.4596191640011966
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00899195959791541,
+      "trajectories_per_second": 111.21046409414787,
+      "compile_seconds": 2.1881700130179524
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001728332252241671,
+      "trajectories_per_second": 578.59245449073,
+      "compile_seconds": 0.7531253910856321
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.009198783803731203,
+      "trajectories_per_second": 108.71002312222848,
+      "compile_seconds": 2.458911601919681
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0027445124681381616,
+      "trajectories_per_second": 364.36343853755045,
+      "compile_seconds": 0.4927656380459666
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.017914881464093924,
+      "trajectories_per_second": 55.819515300966955,
+      "compile_seconds": 2.2372332849772647
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0033871910750998984,
+      "trajectories_per_second": 295.2298756781848,
+      "compile_seconds": 0.7282182819908485
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01836333196843043,
+      "trajectories_per_second": 54.456348211706,
+      "compile_seconds": 2.5651017769705504
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005403019772428606,
+      "trajectories_per_second": 185.0816843393689,
+      "compile_seconds": 0.45458097406663
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.035812951042316854,
+      "trajectories_per_second": 27.92285949343835,
+      "compile_seconds": 2.2190580539172515
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006687036000325212,
+      "trajectories_per_second": 149.5430860475952,
+      "compile_seconds": 0.7877713339403272
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03657345997635275,
+      "trajectories_per_second": 27.342231242178578,
+      "compile_seconds": 2.669947710004635
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0012014583247946576,
+      "trajectories_per_second": 3329.2873480931135,
+      "compile_seconds": 0.7411288149887696
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007314832667664935,
+      "trajectories_per_second": 546.8341084112445,
+      "compile_seconds": 2.7222610749304295
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013529325757796566,
+      "trajectories_per_second": 2956.5405339544827,
+      "compile_seconds": 0.7850496270693839
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00741405082711329,
+      "trajectories_per_second": 539.5161286690864,
+      "compile_seconds": 2.830660925945267
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0022959496681800202,
+      "trajectories_per_second": 1742.1984703918906,
+      "compile_seconds": 0.5999633140163496
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.014622149989008904,
+      "trajectories_per_second": 273.5575823669363,
+      "compile_seconds": 2.507592984009534
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0026088268884147205,
+      "trajectories_per_second": 1533.2561994677383,
+      "compile_seconds": 0.8342320459196344
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.014867656670200327,
+      "trajectories_per_second": 269.04037998249686,
+      "compile_seconds": 3.0378022770164534
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004471874361942438,
+      "trajectories_per_second": 894.4795126718474,
+      "compile_seconds": 0.6013494469225407
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03060528391506523,
+      "trajectories_per_second": 130.69638599336858,
+      "compile_seconds": 2.5861572610447183
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005181894327203433,
+      "trajectories_per_second": 771.9184814327779,
+      "compile_seconds": 0.9482657789485529
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.030956086004152894,
+      "trajectories_per_second": 129.215301942997,
+      "compile_seconds": 2.820216372027062
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008888749778270722,
+      "trajectories_per_second": 450.00704258526076,
+      "compile_seconds": 0.6464779659872875
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.06418487604241818,
+      "trajectories_per_second": 62.319977020077125,
+      "compile_seconds": 2.5713397080544382
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.010260737006319687,
+      "trajectories_per_second": 389.8355447114921,
+      "compile_seconds": 0.9137451130663976
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.06488121603615582,
+      "trajectories_per_second": 61.65112561655677,
+      "compile_seconds": 3.0693628640146926
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009877368994057178,
+      "trajectories_per_second": 16198.645620738242,
+      "compile_seconds": 0.46920299099292606
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007081081663879256,
+      "trajectories_per_second": 2259.541798764493,
+      "compile_seconds": 2.8289907249854878
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011576078804431571,
+      "trajectories_per_second": 13821.605977556801,
+      "compile_seconds": 0.657542129047215
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007209093833807856,
+      "trajectories_per_second": 2219.419023923118,
+      "compile_seconds": 3.0945253369864076
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0019127228297293186,
+      "trajectories_per_second": 8365.03844222127,
+      "compile_seconds": 0.5078434629831463
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01422940733997772,
+      "trajectories_per_second": 1124.4319329483087,
+      "compile_seconds": 2.853143715998158
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002240234276872467,
+      "trajectories_per_second": 7142.11016462849,
+      "compile_seconds": 0.6932006339775398
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.014579707329782346,
+      "trajectories_per_second": 1097.4157188544098,
+      "compile_seconds": 3.2591083210427314
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0037042140029370785,
+      "trajectories_per_second": 4319.4048689718165,
+      "compile_seconds": 0.4679751800140366
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.028543114895001054,
+      "trajectories_per_second": 560.5554985451917,
+      "compile_seconds": 2.8852242149878293
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004348183295223862,
+      "trajectories_per_second": 3679.697683760191,
+      "compile_seconds": 0.8160997150698677
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.02922388003207743,
+      "trajectories_per_second": 547.4974569577239,
+      "compile_seconds": 3.5154968770220876
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007285461489421626,
+      "trajectories_per_second": 2196.1546325145973,
+      "compile_seconds": 0.47133537696208805
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.05705096002202481,
+      "trajectories_per_second": 280.4510212242374,
+      "compile_seconds": 2.8587830069009215
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008682303386740387,
+      "trajectories_per_second": 1842.8289461106829,
+      "compile_seconds": 0.8617396949557588
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.05849828303325921,
+      "trajectories_per_second": 273.51230105169407,
+      "compile_seconds": 4.178306789021008
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009410176859876397,
+      "trajectories_per_second": 34005.73706158836,
+      "compile_seconds": 0.4545221939915791
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006921677135064134,
+      "trajectories_per_second": 4623.15698573876,
+      "compile_seconds": 2.5349392690695822
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0010999878868460655,
+      "trajectories_per_second": 29091.229442309435,
+      "compile_seconds": 0.6856201749760658
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0070723916675585015,
+      "trajectories_per_second": 4524.636290547366,
+      "compile_seconds": 2.8466071450384334
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001806339223144783,
+      "trajectories_per_second": 17715.38789059175,
+      "compile_seconds": 0.4752070080721751
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0138525446333612,
+      "trajectories_per_second": 2310.0448940575243,
+      "compile_seconds": 2.5599693639669567
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021056555636713038,
+      "trajectories_per_second": 15197.16735827705,
+      "compile_seconds": 0.6845111540751532
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01424610634179165,
+      "trajectories_per_second": 2246.227792511027,
+      "compile_seconds": 3.022973488084972
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0034967176974392855,
+      "trajectories_per_second": 9151.439369393252,
+      "compile_seconds": 0.46306284598540515
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0278529740171507,
+      "trajectories_per_second": 1148.8898808542217,
+      "compile_seconds": 2.5829724160721526
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004175346630456095,
+      "trajectories_per_second": 7664.034350246143,
+      "compile_seconds": 0.8565303880022839
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.028483098023571074,
+      "trajectories_per_second": 1123.4732954090362,
+      "compile_seconds": 3.2766328420257196
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006881416143317308,
+      "trajectories_per_second": 4650.20561662673,
+      "compile_seconds": 0.461951993056573
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.05579029710497707,
+      "trajectories_per_second": 573.5764399997302,
+      "compile_seconds": 2.682668666006066
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008254363667219877,
+      "trajectories_per_second": 3876.7373585779756,
+      "compile_seconds": 0.8560209930874407
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.05706574197392911,
+      "trajectories_per_second": 560.756749901183,
+      "compile_seconds": 3.9734320789575577
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-a-kernels-scale.md b/benchmarks/results/vulcan-l40s-a-kernels-scale.md
new file mode 100644
index 0000000..445cf37
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-a-kernels-scale.md
@@ -0,0 +1,198 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=1, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | primal | 0.445 | 2.25e+03 | 0.47 |
+| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | grad | 2.342 | 427 | 1.43 |
+| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | primal | 0.529 | 1.89e+03 | 0.33 |
+| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | grad | 2.471 | 405 | 1.28 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | primal | 0.822 | 1.22e+03 | 0.39 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | grad | 4.553 | 220 | 1.38 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.972 | 1.03e+03 | 0.35 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 4.817 | 208 | 1.30 |
+| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | primal | 1.533 | 652 | 0.36 |
+| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | grad | 9.022 | 111 | 1.39 |
+| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | primal | 1.850 | 541 | 0.35 |
+| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | grad | 9.489 | 105 | 1.30 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | primal | 3.002 | 333 | 0.36 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | grad | 17.824 | 56.1 | 1.37 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 3.564 | 281 | 0.37 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 18.924 | 52.8 | 1.29 |
+| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | primal | 0.618 | 6.47e+03 | 0.55 |
+| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | grad | 4.215 | 949 | 1.65 |
+| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | primal | 0.697 | 5.74e+03 | 0.38 |
+| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | grad | 4.353 | 919 | 1.34 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | primal | 1.172 | 3.41e+03 | 0.35 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | grad | 8.355 | 479 | 1.39 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 1.307 | 3.06e+03 | 0.35 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 8.634 | 463 | 1.38 |
+| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | primal | 2.258 | 1.77e+03 | 0.36 |
+| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | grad | 16.870 | 237 | 1.39 |
+| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | primal | 2.555 | 1.57e+03 | 0.35 |
+| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | grad | 17.002 | 235 | 1.39 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | primal | 4.365 | 916 | 0.34 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | grad | 34.251 | 117 | 1.41 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 5.010 | 798 | 0.35 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 34.530 | 116 | 1.41 |
+| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.645 | 2.48e+04 | 0.49 |
+| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | grad | 3.999 | 4e+03 | 1.57 |
+| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.722 | 2.22e+04 | 0.35 |
+| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | grad | 4.105 | 3.9e+03 | 1.40 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 1.194 | 1.34e+04 | 0.33 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | grad | 7.826 | 2.04e+03 | 1.42 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.348 | 1.19e+04 | 0.35 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 8.141 | 1.97e+03 | 1.39 |
+| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 2.282 | 7.01e+03 | 0.37 |
+| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | grad | 15.496 | 1.03e+03 | 1.38 |
+| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 2.587 | 6.19e+03 | 0.35 |
+| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | grad | 16.209 | 987 | 1.43 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 4.547 | 3.52e+03 | 0.34 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | grad | 30.837 | 519 | 1.43 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 5.058 | 3.16e+03 | 0.35 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 31.876 | 502 | 1.45 |
+| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | primal | 0.650 | 4.92e+04 | 0.61 |
+| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | grad | 3.941 | 8.12e+03 | 1.64 |
+| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | primal | 0.724 | 4.42e+04 | 0.41 |
+| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | grad | 4.050 | 7.9e+03 | 1.51 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | primal | 1.219 | 2.62e+04 | 0.43 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | grad | 7.831 | 4.09e+03 | 1.48 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.351 | 2.37e+04 | 0.41 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 8.032 | 3.98e+03 | 1.55 |
+| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | primal | 2.388 | 1.34e+04 | 0.40 |
+| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | grad | 15.544 | 2.06e+03 | 1.53 |
+| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | primal | 2.630 | 1.22e+04 | 0.41 |
+| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | grad | 15.981 | 2e+03 | 1.51 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | primal | 4.567 | 7.01e+03 | 0.43 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | grad | 31.250 | 1.02e+03 | 1.55 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 5.136 | 6.23e+03 | 0.42 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 32.036 | 999 | 1.55 |
+| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | primal | 0.307 | 3.26e+03 | 0.18 |
+| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | grad | 1.375 | 727 | 0.74 |
+| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | primal | 0.392 | 2.55e+03 | 0.22 |
+| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | grad | 1.477 | 677 | 0.79 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | primal | 0.537 | 1.86e+03 | 0.19 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | grad | 2.618 | 382 | 0.74 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.697 | 1.43e+03 | 0.22 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 2.867 | 349 | 0.75 |
+| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | primal | 1.014 | 986 | 0.19 |
+| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | grad | 5.112 | 196 | 0.75 |
+| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | primal | 1.303 | 768 | 0.22 |
+| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | grad | 5.601 | 179 | 0.76 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | primal | 1.939 | 516 | 0.19 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | grad | 10.114 | 98.9 | 0.75 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.522 | 397 | 0.23 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 10.972 | 91.1 | 0.80 |
+| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | primal | 0.399 | 1e+04 | 0.23 |
+| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | grad | 2.401 | 1.67e+03 | 0.92 |
+| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | primal | 0.476 | 8.41e+03 | 0.24 |
+| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | grad | 2.551 | 1.57e+03 | 0.92 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | primal | 0.708 | 5.65e+03 | 0.22 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | grad | 4.674 | 856 | 0.92 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 0.854 | 4.68e+03 | 0.24 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 5.007 | 799 | 0.93 |
+| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | primal | 1.333 | 3e+03 | 0.22 |
+| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | grad | 9.231 | 433 | 0.93 |
+| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | primal | 1.636 | 2.44e+03 | 0.24 |
+| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | grad | 9.984 | 401 | 0.96 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | primal | 2.596 | 1.54e+03 | 0.22 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | grad | 18.335 | 218 | 0.94 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 3.162 | 1.27e+03 | 0.25 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 19.791 | 202 | 0.96 |
+| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.395 | 4.05e+04 | 0.22 |
+| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | grad | 2.194 | 7.29e+03 | 0.90 |
+| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.483 | 3.32e+04 | 0.24 |
+| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | grad | 2.308 | 6.93e+03 | 0.94 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 0.710 | 2.25e+04 | 0.21 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | grad | 4.314 | 3.71e+03 | 0.94 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.875 | 1.83e+04 | 0.24 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 4.538 | 3.53e+03 | 0.93 |
+| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 1.339 | 1.19e+04 | 0.22 |
+| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | grad | 8.471 | 1.89e+03 | 0.91 |
+| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 1.627 | 9.83e+03 | 0.27 |
+| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | grad | 8.919 | 1.79e+03 | 0.92 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 2.557 | 6.26e+03 | 0.22 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | grad | 16.777 | 954 | 0.93 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.161 | 5.06e+03 | 0.24 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 17.751 | 901 | 0.95 |
+| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | primal | 0.381 | 8.4e+04 | 0.24 |
+| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | grad | 2.288 | 1.4e+04 | 0.94 |
+| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | primal | 0.465 | 6.88e+04 | 0.26 |
+| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | grad | 2.402 | 1.33e+04 | 0.94 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | primal | 0.688 | 4.65e+04 | 0.23 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | grad | 4.525 | 7.07e+03 | 0.92 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.843 | 3.8e+04 | 0.25 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 4.738 | 6.75e+03 | 0.97 |
+| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | primal | 1.293 | 2.47e+04 | 0.23 |
+| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | grad | 8.888 | 3.6e+03 | 0.92 |
+| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | primal | 1.609 | 1.99e+04 | 0.26 |
+| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | grad | 9.384 | 3.41e+03 | 0.94 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | primal | 2.494 | 1.28e+04 | 0.24 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | grad | 17.783 | 1.8e+03 | 0.95 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 3.072 | 1.04e+04 | 0.26 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 18.721 | 1.71e+03 | 0.95 |
+| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | t1 | primal | 0.754 | 1.33e+03 | 0.54 |
+| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | t1 | grad | 4.507 | 222 | 2.26 |
+| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | steps | primal | 0.908 | 1.1e+03 | 0.72 |
+| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | steps | grad | 4.635 | 216 | 2.51 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | t1 | primal | 1.425 | 702 | 0.46 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | t1 | grad | 8.992 | 111 | 2.19 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.728 | 579 | 0.75 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 9.199 | 109 | 2.46 |
+| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | t1 | primal | 2.745 | 364 | 0.49 |
+| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | t1 | grad | 17.915 | 55.8 | 2.24 |
+| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | steps | primal | 3.387 | 295 | 0.73 |
+| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | steps | grad | 18.363 | 54.5 | 2.57 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | t1 | primal | 5.403 | 185 | 0.45 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | t1 | grad | 35.813 | 27.9 | 2.22 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 6.687 | 150 | 0.79 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 36.573 | 27.3 | 2.67 |
+| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | t1 | primal | 1.201 | 3.33e+03 | 0.74 |
+| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | t1 | grad | 7.315 | 547 | 2.72 |
+| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | steps | primal | 1.353 | 2.96e+03 | 0.79 |
+| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | steps | grad | 7.414 | 540 | 2.83 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | t1 | primal | 2.296 | 1.74e+03 | 0.60 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | t1 | grad | 14.622 | 274 | 2.51 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | primal | 2.609 | 1.53e+03 | 0.83 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | grad | 14.868 | 269 | 3.04 |
+| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | t1 | primal | 4.472 | 894 | 0.60 |
+| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | t1 | grad | 30.605 | 131 | 2.59 |
+| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | steps | primal | 5.182 | 772 | 0.95 |
+| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | steps | grad | 30.956 | 129 | 2.82 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | t1 | primal | 8.889 | 450 | 0.65 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | t1 | grad | 64.185 | 62.3 | 2.57 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | primal | 10.261 | 390 | 0.91 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | grad | 64.881 | 61.7 | 3.07 |
+| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | t1 | primal | 0.988 | 1.62e+04 | 0.47 |
+| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | t1 | grad | 7.081 | 2.26e+03 | 2.83 |
+| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | steps | primal | 1.158 | 1.38e+04 | 0.66 |
+| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | steps | grad | 7.209 | 2.22e+03 | 3.09 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | t1 | primal | 1.913 | 8.37e+03 | 0.51 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | t1 | grad | 14.229 | 1.12e+03 | 2.85 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 2.240 | 7.14e+03 | 0.69 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | grad | 14.580 | 1.1e+03 | 3.26 |
+| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | t1 | primal | 3.704 | 4.32e+03 | 0.47 |
+| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | t1 | grad | 28.543 | 561 | 2.89 |
+| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | steps | primal | 4.348 | 3.68e+03 | 0.82 |
+| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | steps | grad | 29.224 | 547 | 3.52 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | t1 | primal | 7.285 | 2.2e+03 | 0.47 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | t1 | grad | 57.051 | 280 | 2.86 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 8.682 | 1.84e+03 | 0.86 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | grad | 58.498 | 274 | 4.18 |
+| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | t1 | primal | 0.941 | 3.4e+04 | 0.45 |
+| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | t1 | grad | 6.922 | 4.62e+03 | 2.53 |
+| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | steps | primal | 1.100 | 2.91e+04 | 0.69 |
+| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | steps | grad | 7.072 | 4.52e+03 | 2.85 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | t1 | primal | 1.806 | 1.77e+04 | 0.48 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | t1 | grad | 13.853 | 2.31e+03 | 2.56 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 2.106 | 1.52e+04 | 0.68 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 14.246 | 2.25e+03 | 3.02 |
+| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | t1 | primal | 3.497 | 9.15e+03 | 0.46 |
+| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | t1 | grad | 27.853 | 1.15e+03 | 2.58 |
+| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | steps | primal | 4.175 | 7.66e+03 | 0.86 |
+| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | steps | grad | 28.483 | 1.12e+03 | 3.28 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | t1 | primal | 6.881 | 4.65e+03 | 0.46 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | t1 | grad | 55.790 | 574 | 2.68 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 8.254 | 3.88e+03 | 0.86 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 57.066 | 561 | 3.97 |
diff --git a/benchmarks/results/vulcan-l40s-b-primal-scale.json b/benchmarks/results/vulcan-l40s-b-primal-scale.json
new file mode 100644
index 0000000..fb6d8c7
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-b-primal-scale.json
@@ -0,0 +1,1164 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 1,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006354805152637489,
+      "trajectories_per_second": 25177.7979272258,
+      "compile_seconds": 0.5559755880385637
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007129469992254268,
+      "trajectories_per_second": 22442.06093494049,
+      "compile_seconds": 0.39084272703621536
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011766815370675631,
+      "trajectories_per_second": 13597.561868671784,
+      "compile_seconds": 0.35566729691345245
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013372436967781847,
+      "trajectories_per_second": 11964.909641039048,
+      "compile_seconds": 0.36642879201099277
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0022774502518586814,
+      "trajectories_per_second": 7025.400439347476,
+      "compile_seconds": 0.3300260939868167
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0026012061086172857,
+      "trajectories_per_second": 6150.992782538507,
+      "compile_seconds": 0.4029288350138813
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004493608636866239,
+      "trajectories_per_second": 3560.6127041713426,
+      "compile_seconds": 0.3373294989578426
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005107863663902713,
+      "trajectories_per_second": 3132.42503183318,
+      "compile_seconds": 0.4059816070366651
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006648337464809181,
+      "trajectories_per_second": 385058.67272089154,
+      "compile_seconds": 1.4882204619934782
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007188769645596805,
+      "trajectories_per_second": 356111.0073360086,
+      "compile_seconds": 0.4338768139714375
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011983258737018332,
+      "trajectories_per_second": 213631.37158106442,
+      "compile_seconds": 0.410205323016271
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013713163461943623,
+      "trajectories_per_second": 186681.94301806716,
+      "compile_seconds": 0.46616124699357897
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0023338496685028076,
+      "trajectories_per_second": 109690.01279513734,
+      "compile_seconds": 0.391712898039259
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002635435725096613,
+      "trajectories_per_second": 97137.63745485207,
+      "compile_seconds": 0.4829016280127689
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004540593200363219,
+      "trajectories_per_second": 56380.29849921847,
+      "compile_seconds": 0.39958338206633925
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005143783119921055,
+      "trajectories_per_second": 49768.81684776962,
+      "compile_seconds": 0.48374091600999236
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007221257656055968,
+      "trajectories_per_second": 1418035.540029848,
+      "compile_seconds": 1.3409415270434693
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008058236046465781,
+      "trajectories_per_second": 1270749.5711162626,
+      "compile_seconds": 0.47849129198584706
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001375524909235537,
+      "trajectories_per_second": 744443.079965087,
+      "compile_seconds": 0.3930518280249089
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0015292710985147183,
+      "trajectories_per_second": 669600.0473654048,
+      "compile_seconds": 0.4506659930339083
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0026392538889518213,
+      "trajectories_per_second": 387988.4403264747,
+      "compile_seconds": 0.39985329599585384
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0029785154983983375,
+      "trajectories_per_second": 343795.42444907344,
+      "compile_seconds": 0.4705504380399361
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005217596215920316,
+      "trajectories_per_second": 196258.95865139877,
+      "compile_seconds": 0.4048258389811963
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005795780627522618,
+      "trajectories_per_second": 176680.25513893622,
+      "compile_seconds": 0.46957040403503925
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0003995470896673699,
+      "trajectories_per_second": 40045.34237333648,
+      "compile_seconds": 0.2224262390518561
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00047845955766541394,
+      "trajectories_per_second": 33440.65291133504,
+      "compile_seconds": 0.24956504988949746
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007118711664313168,
+      "trajectories_per_second": 22475.97705102967,
+      "compile_seconds": 0.2086654029553756
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008631101975749646,
+      "trajectories_per_second": 18537.609733906935,
+      "compile_seconds": 0.2661267350194976
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013300093808485305,
+      "trajectories_per_second": 12029.99033720513,
+      "compile_seconds": 0.21264974900987
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00162748118080864,
+      "trajectories_per_second": 9831.14286584263,
+      "compile_seconds": 0.2523548499448225
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0025666021133979987,
+      "trajectories_per_second": 6233.92302082115,
+      "compile_seconds": 0.21147427102550864
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003156211933431526,
+      "trajectories_per_second": 5069.368070795022,
+      "compile_seconds": 0.26409248902928084
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0004045593628563501,
+      "trajectories_per_second": 632787.2334792553,
+      "compile_seconds": 0.24785651196725667
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00048128006388791476,
+      "trajectories_per_second": 531914.8230075447,
+      "compile_seconds": 0.29310431494377553
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007103458220667359,
+      "trajectories_per_second": 360387.84497271135,
+      "compile_seconds": 0.25716804608237
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008685059104622765,
+      "trajectories_per_second": 294759.0763818059,
+      "compile_seconds": 0.29665725002996624
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001335371148717754,
+      "trajectories_per_second": 191707.00239091998,
+      "compile_seconds": 0.2603799499338493
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001629663723649393,
+      "trajectories_per_second": 157087.622608869,
+      "compile_seconds": 0.3043096459005028
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0025974840575240946,
+      "trajectories_per_second": 98556.90904375273,
+      "compile_seconds": 0.25925671693403274
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0031994502060115337,
+      "trajectories_per_second": 80013.74721162862,
+      "compile_seconds": 0.31164434493985027
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0004323138322280807,
+      "trajectories_per_second": 2368649.6328892773,
+      "compile_seconds": 0.2987338110106066
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.000522802716020394,
+      "trajectories_per_second": 1958673.8335920484,
+      "compile_seconds": 0.2993554960703477
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007981986331287771,
+      "trajectories_per_second": 1282888.6914853854,
+      "compile_seconds": 0.2597087600734085
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009497540397569537,
+      "trajectories_per_second": 1078173.8820105952,
+      "compile_seconds": 0.30768832308240235
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0015057500313560013,
+      "trajectories_per_second": 680059.7567166166,
+      "compile_seconds": 0.26700164598878473
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0018168388400226832,
+      "trajectories_per_second": 563616.308415784,
+      "compile_seconds": 0.31870555493514985
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002913577730456988,
+      "trajectories_per_second": 351457.93067253707,
+      "compile_seconds": 0.26981633191462606
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0035207956903972304,
+      "trajectories_per_second": 290843.34623360896,
+      "compile_seconds": 0.311832238920033
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009868918520320828,
+      "trajectories_per_second": 16212.516059439364,
+      "compile_seconds": 0.463355409912765
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011519278826502462,
+      "trajectories_per_second": 13889.758413685344,
+      "compile_seconds": 0.7675077740568668
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0019186041657424842,
+      "trajectories_per_second": 8339.396049319079,
+      "compile_seconds": 0.46876115596387535
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002235232497861778,
+      "trajectories_per_second": 7158.092062148161,
+      "compile_seconds": 0.7880562850041315
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0037642026929041515,
+      "trajectories_per_second": 4250.56812964971,
+      "compile_seconds": 0.47365841490682214
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004370011537420479,
+      "trajectories_per_second": 3661.3175647230546,
+      "compile_seconds": 0.9838717890670523
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007313157839234918,
+      "trajectories_per_second": 2187.8373681695175,
+      "compile_seconds": 0.46887913800310344
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008711993182078003,
+      "trajectories_per_second": 1836.5487283569757,
+      "compile_seconds": 1.005123715964146
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009848584997512045,
+      "trajectories_per_second": 259935.8182568063,
+      "compile_seconds": 0.4887423829641193
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011351239775459086,
+      "trajectories_per_second": 225526.02628786105,
+      "compile_seconds": 0.7136238029925153
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0018564381887419866,
+      "trajectories_per_second": 137898.47760752979,
+      "compile_seconds": 0.4788528170902282
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021981775228466305,
+      "trajectories_per_second": 116460.11176953584,
+      "compile_seconds": 0.7647949910024181
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0036615609961490217,
+      "trajectories_per_second": 69915.53609764884,
+      "compile_seconds": 0.5224804749013856
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004349283916367726,
+      "trajectories_per_second": 58860.26410844123,
+      "compile_seconds": 0.8992384019074962
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007238235344023754,
+      "trajectories_per_second": 35367.73644854837,
+      "compile_seconds": 0.4920805470319465
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 256,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008535689604468644,
+      "trajectories_per_second": 29991.718521017643,
+      "compile_seconds": 0.9124668670119718
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011423759502716935,
+      "trajectories_per_second": 896377.413894664,
+      "compile_seconds": 0.4886601109756157
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 31,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0012931094592990907,
+      "trajectories_per_second": 791889.6522147807,
+      "compile_seconds": 0.6905786630231887
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002228784681805833,
+      "trajectories_per_second": 459443.21511144016,
+      "compile_seconds": 0.4846900930861011
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0025102954500147868,
+      "trajectories_per_second": 407920.1115525936,
+      "compile_seconds": 0.766396596096456
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004280094639398158,
+      "trajectories_per_second": 239247.04621577924,
+      "compile_seconds": 0.5410391429904848
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 127,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004943655990064144,
+      "trajectories_per_second": 207134.1537635416,
+      "compile_seconds": 0.8720746849430725
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "t1",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008362751407548785,
+      "trajectories_per_second": 122447.73879989646,
+      "compile_seconds": 0.4981267689727247
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1024,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.009781642816960812,
+      "trajectories_per_second": 104685.8916402511,
+      "compile_seconds": 0.9231459660222754
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-b-primal-scale.md b/benchmarks/results/vulcan-l40s-b-primal-scale.md
new file mode 100644
index 0000000..2b9b872
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-b-primal-scale.md
@@ -0,0 +1,78 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=1, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.635 | 2.52e+04 | 0.56 |
+| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.713 | 2.24e+04 | 0.39 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 1.177 | 1.36e+04 | 0.36 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.337 | 1.2e+04 | 0.37 |
+| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 2.277 | 7.03e+03 | 0.33 |
+| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 2.601 | 6.15e+03 | 0.40 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 4.494 | 3.56e+03 | 0.34 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 5.108 | 3.13e+03 | 0.41 |
+| sra1 | stochastic-growth | 256 | 2 | 31 | float32 | explicit | t1 | primal | 0.665 | 3.85e+05 | 1.49 |
+| sra1 | stochastic-growth | 256 | 2 | 31 | float32 | explicit | steps | primal | 0.719 | 3.56e+05 | 0.43 |
+| sra1 | stochastic-growth | 256 | 2 | 63 | float32 | explicit | t1 | primal | 1.198 | 2.14e+05 | 0.41 |
+| sra1 | stochastic-growth | 256 | 2 | 63 | float32 | explicit | steps | primal | 1.371 | 1.87e+05 | 0.47 |
+| sra1 | stochastic-growth | 256 | 2 | 127 | float32 | explicit | t1 | primal | 2.334 | 1.1e+05 | 0.39 |
+| sra1 | stochastic-growth | 256 | 2 | 127 | float32 | explicit | steps | primal | 2.635 | 9.71e+04 | 0.48 |
+| sra1 | stochastic-growth | 256 | 2 | 255 | float32 | explicit | t1 | primal | 4.541 | 5.64e+04 | 0.40 |
+| sra1 | stochastic-growth | 256 | 2 | 255 | float32 | explicit | steps | primal | 5.144 | 4.98e+04 | 0.48 |
+| sra1 | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | t1 | primal | 0.722 | 1.42e+06 | 1.34 |
+| sra1 | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | steps | primal | 0.806 | 1.27e+06 | 0.48 |
+| sra1 | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | t1 | primal | 1.376 | 7.44e+05 | 0.39 |
+| sra1 | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | steps | primal | 1.529 | 6.7e+05 | 0.45 |
+| sra1 | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | t1 | primal | 2.639 | 3.88e+05 | 0.40 |
+| sra1 | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | steps | primal | 2.979 | 3.44e+05 | 0.47 |
+| sra1 | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | t1 | primal | 5.218 | 1.96e+05 | 0.40 |
+| sra1 | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | steps | primal | 5.796 | 1.77e+05 | 0.47 |
+| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.400 | 4e+04 | 0.22 |
+| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.478 | 3.34e+04 | 0.25 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 0.712 | 2.25e+04 | 0.21 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.863 | 1.85e+04 | 0.27 |
+| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 1.330 | 1.2e+04 | 0.21 |
+| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 1.627 | 9.83e+03 | 0.25 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 2.567 | 6.23e+03 | 0.21 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.156 | 5.07e+03 | 0.26 |
+| em | stochastic-growth | 256 | 2 | 31 | float32 | explicit | t1 | primal | 0.405 | 6.33e+05 | 0.25 |
+| em | stochastic-growth | 256 | 2 | 31 | float32 | explicit | steps | primal | 0.481 | 5.32e+05 | 0.29 |
+| em | stochastic-growth | 256 | 2 | 63 | float32 | explicit | t1 | primal | 0.710 | 3.6e+05 | 0.26 |
+| em | stochastic-growth | 256 | 2 | 63 | float32 | explicit | steps | primal | 0.869 | 2.95e+05 | 0.30 |
+| em | stochastic-growth | 256 | 2 | 127 | float32 | explicit | t1 | primal | 1.335 | 1.92e+05 | 0.26 |
+| em | stochastic-growth | 256 | 2 | 127 | float32 | explicit | steps | primal | 1.630 | 1.57e+05 | 0.30 |
+| em | stochastic-growth | 256 | 2 | 255 | float32 | explicit | t1 | primal | 2.597 | 9.86e+04 | 0.26 |
+| em | stochastic-growth | 256 | 2 | 255 | float32 | explicit | steps | primal | 3.199 | 8e+04 | 0.31 |
+| em | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | t1 | primal | 0.432 | 2.37e+06 | 0.30 |
+| em | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | steps | primal | 0.523 | 1.96e+06 | 0.30 |
+| em | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | t1 | primal | 0.798 | 1.28e+06 | 0.26 |
+| em | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | steps | primal | 0.950 | 1.08e+06 | 0.31 |
+| em | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | t1 | primal | 1.506 | 6.8e+05 | 0.27 |
+| em | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | steps | primal | 1.817 | 5.64e+05 | 0.32 |
+| em | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | t1 | primal | 2.914 | 3.51e+05 | 0.27 |
+| em | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | steps | primal | 3.521 | 2.91e+05 | 0.31 |
+| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | t1 | primal | 0.987 | 1.62e+04 | 0.46 |
+| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | steps | primal | 1.152 | 1.39e+04 | 0.77 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | t1 | primal | 1.919 | 8.34e+03 | 0.47 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 2.235 | 7.16e+03 | 0.79 |
+| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | t1 | primal | 3.764 | 4.25e+03 | 0.47 |
+| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | steps | primal | 4.370 | 3.66e+03 | 0.98 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | t1 | primal | 7.313 | 2.19e+03 | 0.47 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 8.712 | 1.84e+03 | 1.01 |
+| rk4 | neoclassical | 256 | 2 | 31 | float32 | none | t1 | primal | 0.985 | 2.6e+05 | 0.49 |
+| rk4 | neoclassical | 256 | 2 | 31 | float32 | none | steps | primal | 1.135 | 2.26e+05 | 0.71 |
+| rk4 | neoclassical | 256 | 2 | 63 | float32 | none | t1 | primal | 1.856 | 1.38e+05 | 0.48 |
+| rk4 | neoclassical | 256 | 2 | 63 | float32 | none | steps | primal | 2.198 | 1.16e+05 | 0.76 |
+| rk4 | neoclassical | 256 | 2 | 127 | float32 | none | t1 | primal | 3.662 | 6.99e+04 | 0.52 |
+| rk4 | neoclassical | 256 | 2 | 127 | float32 | none | steps | primal | 4.349 | 5.89e+04 | 0.90 |
+| rk4 | neoclassical | 256 | 2 | 255 | float32 | none | t1 | primal | 7.238 | 3.54e+04 | 0.49 |
+| rk4 | neoclassical | 256 | 2 | 255 | float32 | none | steps | primal | 8.536 | 3e+04 | 0.91 |
+| rk4 | neoclassical | 1024 | 2 | 31 | float32 | none | t1 | primal | 1.142 | 8.96e+05 | 0.49 |
+| rk4 | neoclassical | 1024 | 2 | 31 | float32 | none | steps | primal | 1.293 | 7.92e+05 | 0.69 |
+| rk4 | neoclassical | 1024 | 2 | 63 | float32 | none | t1 | primal | 2.229 | 4.59e+05 | 0.48 |
+| rk4 | neoclassical | 1024 | 2 | 63 | float32 | none | steps | primal | 2.510 | 4.08e+05 | 0.77 |
+| rk4 | neoclassical | 1024 | 2 | 127 | float32 | none | t1 | primal | 4.280 | 2.39e+05 | 0.54 |
+| rk4 | neoclassical | 1024 | 2 | 127 | float32 | none | steps | primal | 4.944 | 2.07e+05 | 0.87 |
+| rk4 | neoclassical | 1024 | 2 | 255 | float32 | none | t1 | primal | 8.363 | 1.22e+05 | 0.50 |
+| rk4 | neoclassical | 1024 | 2 | 255 | float32 | none | steps | primal | 9.782 | 1.05e+05 | 0.92 |
diff --git a/benchmarks/results/vulcan-l40s-c-unroll4.json b/benchmarks/results/vulcan-l40s-c-unroll4.json
new file mode 100644
index 0000000..03189f6
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-c-unroll4.json
@@ -0,0 +1,780 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 4,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006971989419338677,
+      "trajectories_per_second": 1434.3108399250184,
+      "compile_seconds": 1.2450868539744988
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002124298398848623,
+      "trajectories_per_second": 470.7436584907296,
+      "compile_seconds": 7.255458320956677
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002563808891481083,
+      "trajectories_per_second": 390.0446727221979,
+      "compile_seconds": 1.1820683469995856
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008215861841260145,
+      "trajectories_per_second": 121.71577605869533,
+      "compile_seconds": 7.17904240405187
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0010690772451642083,
+      "trajectories_per_second": 3741.5444188839765,
+      "compile_seconds": 1.2497892199316993
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004303974004208364,
+      "trajectories_per_second": 929.3736430770393,
+      "compile_seconds": 6.750535171944648
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003933795417348544,
+      "trajectories_per_second": 1016.8296964197693,
+      "compile_seconds": 0.9893559309421107
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01660529401851818,
+      "trajectories_per_second": 240.88703250536912,
+      "compile_seconds": 5.85350946104154
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001099110838177419,
+      "trajectories_per_second": 14557.221568783469,
+      "compile_seconds": 1.2533364739501849
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004105353242872904,
+      "trajectories_per_second": 3897.3503748494213,
+      "compile_seconds": 6.141521341982298
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004040575159403185,
+      "trajectories_per_second": 3959.8322933716413,
+      "compile_seconds": 1.0728455058997497
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01586910302285105,
+      "trajectories_per_second": 1008.2485429050691,
+      "compile_seconds": 5.67013100697659
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011212968648495998,
+      "trajectories_per_second": 28538.383547779005,
+      "compile_seconds": 1.765444446937181
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004042536909268661,
+      "trajectories_per_second": 7915.821356295088,
+      "compile_seconds": 6.6290037749568
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0041155470826197416,
+      "trajectories_per_second": 7775.393977422433,
+      "compile_seconds": 1.361965456046164
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.015573241321059564,
+      "trajectories_per_second": 2054.8066610081146,
+      "compile_seconds": 6.725575314951129
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00046295225038193164,
+      "trajectories_per_second": 2160.0499817745967,
+      "compile_seconds": 0.7804745350731537
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0011810132515771936,
+      "trajectories_per_second": 846.7305499447547,
+      "compile_seconds": 3.3045207059476525
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0015336008145823143,
+      "trajectories_per_second": 652.0601648691456,
+      "compile_seconds": 0.7775656069861725
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004280021551742472,
+      "trajectories_per_second": 233.64368331109978,
+      "compile_seconds": 3.313947707065381
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0005910745366355,
+      "trajectories_per_second": 6767.336016145616,
+      "compile_seconds": 0.6687185120536014
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0022954671003390104,
+      "trajectories_per_second": 1742.5647265470511,
+      "compile_seconds": 3.773560242028907
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002103084047451954,
+      "trajectories_per_second": 1901.9686849159943,
+      "compile_seconds": 0.6551523020025343
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008795526390895248,
+      "trajectories_per_second": 454.77664692594504,
+      "compile_seconds": 3.7498066320549697
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006124911045557574,
+      "trajectories_per_second": 26122.828365980713,
+      "compile_seconds": 0.6658988069975749
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021762345938689327,
+      "trajectories_per_second": 7352.1485436710345,
+      "compile_seconds": 3.589312301017344
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021402357378974557,
+      "trajectories_per_second": 7475.811994298453,
+      "compile_seconds": 0.6453863370697945
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00822971984356021,
+      "trajectories_per_second": 1944.1731072437499,
+      "compile_seconds": 3.514980165986344
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00060828894015346,
+      "trajectories_per_second": 52606.578695852986,
+      "compile_seconds": 1.0748092120047659
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021829829572445965,
+      "trajectories_per_second": 14658.840965204337,
+      "compile_seconds": 3.92439767497126
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002194125268777663,
+      "trajectories_per_second": 14584.399740233177,
+      "compile_seconds": 0.942150270100683
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008260089671239257,
+      "trajectories_per_second": 3874.0499526803633,
+      "compile_seconds": 4.008594830986112
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001262728130984071,
+      "trajectories_per_second": 791.9361068012943,
+      "compile_seconds": 2.060194884077646
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004306531363082203,
+      "trajectories_per_second": 232.20543766905152,
+      "compile_seconds": 14.47477564692963
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00470992939081043,
+      "trajectories_per_second": 212.31740797454538,
+      "compile_seconds": 2.1245737449498847
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.017145550984423608,
+      "trajectories_per_second": 58.3241682293255,
+      "compile_seconds": 14.554241691017523
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021403831418692356,
+      "trajectories_per_second": 1868.8242874622565,
+      "compile_seconds": 3.1521481269737706
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007487423504547526,
+      "trajectories_per_second": 534.2291640870293,
+      "compile_seconds": 17.594308632891625
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008161732674731562,
+      "trajectories_per_second": 490.0920134745236,
+      "compile_seconds": 3.1455108739901334
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.029878524015657604,
+      "trajectories_per_second": 133.87542162068755,
+      "compile_seconds": 17.492421875009313
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0017456265545829578,
+      "trajectories_per_second": 9165.76340912877,
+      "compile_seconds": 2.474717544973828
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007120985343741874,
+      "trajectories_per_second": 2246.8800633133246,
+      "compile_seconds": 18.362305072019808
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006705917418003082,
+      "trajectories_per_second": 2385.9524361343165,
+      "compile_seconds": 2.6125571640441194
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.02837010205257684,
+      "trajectories_per_second": 563.9740022911454,
+      "compile_seconds": 24.234453483950347
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0016487168613821268,
+      "trajectories_per_second": 19409.033018061244,
+      "compile_seconds": 2.3910328299971297
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007063096001123388,
+      "trajectories_per_second": 4530.591116828992,
+      "compile_seconds": 15.935979263973422
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006319132284261286,
+      "trajectories_per_second": 5063.986408339739,
+      "compile_seconds": 2.5525463869562373
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.029351581004448235,
+      "trajectories_per_second": 1090.2308804132356,
+      "compile_seconds": 23.500358068966307
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-c-unroll4.md b/benchmarks/results/vulcan-l40s-c-unroll4.md
new file mode 100644
index 0000000..cc74030
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-c-unroll4.md
@@ -0,0 +1,54 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=4, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.697 | 1.43e+03 | 1.25 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 2.124 | 471 | 7.26 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.564 | 390 | 1.18 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 8.216 | 122 | 7.18 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 1.069 | 3.74e+03 | 1.25 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 4.304 | 929 | 6.75 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 3.934 | 1.02e+03 | 0.99 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 16.605 | 241 | 5.85 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.099 | 1.46e+04 | 1.25 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 4.105 | 3.9e+03 | 6.14 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 4.041 | 3.96e+03 | 1.07 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 15.869 | 1.01e+03 | 5.67 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.121 | 2.85e+04 | 1.77 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 4.043 | 7.92e+03 | 6.63 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 4.116 | 7.78e+03 | 1.36 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 15.573 | 2.05e+03 | 6.73 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.463 | 2.16e+03 | 0.78 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.181 | 847 | 3.30 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 1.534 | 652 | 0.78 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 4.280 | 234 | 3.31 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 0.591 | 6.77e+03 | 0.67 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 2.295 | 1.74e+03 | 3.77 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 2.103 | 1.9e+03 | 0.66 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 8.796 | 455 | 3.75 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.612 | 2.61e+04 | 0.67 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 2.176 | 7.35e+03 | 3.59 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 2.140 | 7.48e+03 | 0.65 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 8.230 | 1.94e+03 | 3.51 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.608 | 5.26e+04 | 1.07 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 2.183 | 1.47e+04 | 3.92 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 2.194 | 1.46e+04 | 0.94 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 8.260 | 3.87e+03 | 4.01 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.263 | 792 | 2.06 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 4.307 | 232 | 14.47 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 4.710 | 212 | 2.12 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 17.146 | 58.3 | 14.55 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | primal | 2.140 | 1.87e+03 | 3.15 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | grad | 7.487 | 534 | 17.59 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | primal | 8.162 | 490 | 3.15 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | grad | 29.879 | 134 | 17.49 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 1.746 | 9.17e+03 | 2.47 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | grad | 7.121 | 2.25e+03 | 18.36 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 6.706 | 2.39e+03 | 2.61 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | grad | 28.370 | 564 | 24.23 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 1.649 | 1.94e+04 | 2.39 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 7.063 | 4.53e+03 | 15.94 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 6.319 | 5.06e+03 | 2.55 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 29.352 | 1.09e+03 | 23.50 |
diff --git a/benchmarks/results/vulcan-l40s-c-unroll8.json b/benchmarks/results/vulcan-l40s-c-unroll8.json
new file mode 100644
index 0000000..9fc5cf7
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-c-unroll8.json
@@ -0,0 +1,780 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 8,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006998448846085618,
+      "trajectories_per_second": 1428.8880607583799,
+      "compile_seconds": 2.1025211340747774
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0017523762295380808,
+      "trajectories_per_second": 570.6537118821773,
+      "compile_seconds": 15.505087749916129
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002500084890542846,
+      "trajectories_per_second": 399.98641797433885,
+      "compile_seconds": 2.134946339065209
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006634242716245353,
+      "trajectories_per_second": 150.73310440561474,
+      "compile_seconds": 15.41392176807858
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0010648247113244401,
+      "trajectories_per_second": 3756.486825915937,
+      "compile_seconds": 2.01313587394543
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003708232230005356,
+      "trajectories_per_second": 1078.681094359137,
+      "compile_seconds": 13.491816036985256
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0039193480818842845,
+      "trajectories_per_second": 1020.5778911264603,
+      "compile_seconds": 1.7366818400332704
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.014399717367875079,
+      "trajectories_per_second": 277.7832298933703,
+      "compile_seconds": 12.477905547944829
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001080075476784259,
+      "trajectories_per_second": 14813.779540331087,
+      "compile_seconds": 2.114171709981747
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0035481390077620745,
+      "trajectories_per_second": 4509.406188708405,
+      "compile_seconds": 12.483813125058077
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003984976414358243,
+      "trajectories_per_second": 4015.0802254062287,
+      "compile_seconds": 1.944047044031322
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.013775509665720165,
+      "trajectories_per_second": 1161.4815268734046,
+      "compile_seconds": 11.734136439976282
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001090569411065768,
+      "trajectories_per_second": 29342.469791746436,
+      "compile_seconds": 2.712940454017371
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0034320351745312414,
+      "trajectories_per_second": 9323.913763317028,
+      "compile_seconds": 13.425745286047459
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0040747205493971705,
+      "trajectories_per_second": 7853.299290606371,
+      "compile_seconds": 2.3499814469832927
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.013294546360460421,
+      "trajectories_per_second": 2407.0020241662287,
+      "compile_seconds": 13.817451639915816
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00043423763891480054,
+      "trajectories_per_second": 2302.8865081780823,
+      "compile_seconds": 1.2170277669792995
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0010025063681456706,
+      "trajectories_per_second": 997.4998980302672,
+      "compile_seconds": 7.303474535932764
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0014253398403525352,
+      "trajectories_per_second": 701.5870683532328,
+      "compile_seconds": 1.2481635819422081
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0035663329942438463,
+      "trajectories_per_second": 280.40006404730735,
+      "compile_seconds": 7.355043246061541
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0005990132657056556,
+      "trajectories_per_second": 6677.6484412041855,
+      "compile_seconds": 0.9456670730141923
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00201770125325614,
+      "trajectories_per_second": 1982.4540394891724,
+      "compile_seconds": 6.381001956993714
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002031796752514007,
+      "trajectories_per_second": 1968.7008530999335,
+      "compile_seconds": 0.9697260030079633
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0075386849930509925,
+      "trajectories_per_second": 530.5965169903131,
+      "compile_seconds": 6.386174759012647
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.000614990030457689,
+      "trajectories_per_second": 26016.681909611525,
+      "compile_seconds": 1.014958688057959
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0019072774425148964,
+      "trajectories_per_second": 8388.921109926583,
+      "compile_seconds": 6.508749642991461
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0020889290406004243,
+      "trajectories_per_second": 7659.427241913921,
+      "compile_seconds": 0.985628140042536
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007178531843237579,
+      "trajectories_per_second": 2228.868012206778,
+      "compile_seconds": 6.388680926989764
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.000623197469394654,
+      "trajectories_per_second": 51348.09040717601,
+      "compile_seconds": 1.2464099749922752
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001897456720471382,
+      "trajectories_per_second": 16864.67978676757,
+      "compile_seconds": 6.798548002028838
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021914461511187257,
+      "trajectories_per_second": 14602.229666316058,
+      "compile_seconds": 1.152224877034314
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007048038998618722,
+      "trajectories_per_second": 4540.269996558103,
+      "compile_seconds": 7.116658751037903
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0012470546781140216,
+      "trajectories_per_second": 801.88945805676,
+      "compile_seconds": 3.154672044911422
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0033622876425007625,
+      "trajectories_per_second": 297.4165527540148,
+      "compile_seconds": 32.872302406001836
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004600644600577653,
+      "trajectories_per_second": 217.3608454507529,
+      "compile_seconds": 3.0276695468928665
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.013169295697783431,
+      "trajectories_per_second": 75.93420505914476,
+      "compile_seconds": 32.81373755901586
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021212612378544043,
+      "trajectories_per_second": 1885.6706230326854,
+      "compile_seconds": 5.318178678979166
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006540245576096433,
+      "trajectories_per_second": 611.597829693027,
+      "compile_seconds": 40.33293925994076
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008166779787279666,
+      "trajectories_per_second": 489.78913405137746,
+      "compile_seconds": 5.438302522990853
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 4,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.025998475961387157,
+      "trajectories_per_second": 153.85517235474825,
+      "compile_seconds": 39.74205899902154
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.001730186678469181,
+      "trajectories_per_second": 9247.557040582658,
+      "compile_seconds": 4.074979388038628
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006082673120545223,
+      "trajectories_per_second": 2630.4224611310087,
+      "compile_seconds": 41.33263551397249
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006528985420508044,
+      "trajectories_per_second": 2450.6104654090313,
+      "compile_seconds": 4.321440764004365
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.024196182028390467,
+      "trajectories_per_second": 661.2613502918139,
+      "compile_seconds": 48.491429385961965
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0016249311718575913,
+      "trajectories_per_second": 19693.141810688627,
+      "compile_seconds": 4.038868790958077
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006009442877257243,
+      "trajectories_per_second": 5324.952853966565,
+      "compile_seconds": 35.891845369013026
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006131911504780874,
+      "trajectories_per_second": 5218.601079785729,
+      "compile_seconds": 4.401859620003961
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.025074374047107995,
+      "trajectories_per_second": 1276.2033436958634,
+      "compile_seconds": 45.259774502017535
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-c-unroll8.md b/benchmarks/results/vulcan-l40s-c-unroll8.md
new file mode 100644
index 0000000..2ca4806
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-c-unroll8.md
@@ -0,0 +1,54 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=8, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.700 | 1.43e+03 | 2.10 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.752 | 571 | 15.51 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.500 | 400 | 2.13 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 6.634 | 151 | 15.41 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 1.065 | 3.76e+03 | 2.01 |
+| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 3.708 | 1.08e+03 | 13.49 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 3.919 | 1.02e+03 | 1.74 |
+| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 14.400 | 278 | 12.48 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.080 | 1.48e+04 | 2.11 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 3.548 | 4.51e+03 | 12.48 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.985 | 4.02e+03 | 1.94 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 13.776 | 1.16e+03 | 11.73 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.091 | 2.93e+04 | 2.71 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 3.432 | 9.32e+03 | 13.43 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 4.075 | 7.85e+03 | 2.35 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 13.295 | 2.41e+03 | 13.82 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.434 | 2.3e+03 | 1.22 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.003 | 997 | 7.30 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 1.425 | 702 | 1.25 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 3.566 | 280 | 7.36 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 0.599 | 6.68e+03 | 0.95 |
+| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 2.018 | 1.98e+03 | 6.38 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 2.032 | 1.97e+03 | 0.97 |
+| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 7.539 | 531 | 6.39 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.615 | 2.6e+04 | 1.01 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 1.907 | 8.39e+03 | 6.51 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 2.089 | 7.66e+03 | 0.99 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 7.179 | 2.23e+03 | 6.39 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.623 | 5.13e+04 | 1.25 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 1.897 | 1.69e+04 | 6.80 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 2.191 | 1.46e+04 | 1.15 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 7.048 | 4.54e+03 | 7.12 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.247 | 802 | 3.15 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 3.362 | 297 | 32.87 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 4.601 | 217 | 3.03 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 13.169 | 75.9 | 32.81 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | primal | 2.121 | 1.89e+03 | 5.32 |
+| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | grad | 6.540 | 612 | 40.33 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | primal | 8.167 | 490 | 5.44 |
+| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | grad | 25.998 | 154 | 39.74 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 1.730 | 9.25e+03 | 4.07 |
+| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | grad | 6.083 | 2.63e+03 | 41.33 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 6.529 | 2.45e+03 | 4.32 |
+| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | grad | 24.196 | 661 | 48.49 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 1.625 | 1.97e+04 | 4.04 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 6.009 | 5.32e+03 | 35.89 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 6.132 | 5.22e+03 | 4.40 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 25.074 | 1.28e+03 | 45.26 |
diff --git a/benchmarks/results/vulcan-l40s-d-noisemode.json b/benchmarks/results/vulcan-l40s-d-noisemode.json
new file mode 100644
index 0000000..f46ba48
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-d-noisemode.json
@@ -0,0 +1,268 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 1,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013318749124129467,
+      "trajectories_per_second": 12013.14016119797,
+      "compile_seconds": 0.6122790339868516
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013370389707334754,
+      "trajectories_per_second": 11966.741695810624,
+      "compile_seconds": 0.7160951240221038
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008066055675347647,
+      "trajectories_per_second": 1983.621319265243,
+      "compile_seconds": 1.5805910240160301
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007812453336858501,
+      "trajectories_per_second": 2048.0122325356288,
+      "compile_seconds": 1.5591626500245184
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.00505731444961081,
+      "trajectories_per_second": 3163.7344601404593,
+      "compile_seconds": 0.40626175701618195
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0050121010022444856,
+      "trajectories_per_second": 3192.2740568945014,
+      "compile_seconds": 0.7446511259768158
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03158391290344298,
+      "trajectories_per_second": 506.58700994124865,
+      "compile_seconds": 1.4282209059456363
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.030595276970416307,
+      "trajectories_per_second": 522.9565339601595,
+      "compile_seconds": 1.6230945909628645
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008508051179495512,
+      "trajectories_per_second": 18805.716682288134,
+      "compile_seconds": 0.22564140998292714
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008544241986237467,
+      "trajectories_per_second": 18726.06139406141,
+      "compile_seconds": 0.3509363259654492
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004469206696376205,
+      "trajectories_per_second": 3580.0537068409435,
+      "compile_seconds": 0.9315736360149458
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004436203818344934,
+      "trajectories_per_second": 3606.6873063486305,
+      "compile_seconds": 1.0160666520241648
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0031061899305010834,
+      "trajectories_per_second": 5151.005044118122,
+      "compile_seconds": 0.22345678601413965
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003116562000165383,
+      "trajectories_per_second": 5133.862249219154,
+      "compile_seconds": 0.3596626119688153
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.017609555507078767,
+      "trajectories_per_second": 908.597607337008,
+      "compile_seconds": 0.9339326129993424
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 16,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "key",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.01742742903297767,
+      "trajectories_per_second": 918.0929653894119,
+      "compile_seconds": 1.0483848019503057
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-d-noisemode.md b/benchmarks/results/vulcan-l40s-d-noisemode.md
new file mode 100644
index 0000000..690f1b8
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-d-noisemode.md
@@ -0,0 +1,22 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=1, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.332 | 1.2e+04 | 0.61 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | primal | 1.337 | 1.2e+04 | 0.72 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 8.066 | 1.98e+03 | 1.58 |
+| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | grad | 7.812 | 2.05e+03 | 1.56 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 5.057 | 3.16e+03 | 0.41 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | primal | 5.012 | 3.19e+03 | 0.74 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 31.584 | 507 | 1.43 |
+| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | grad | 30.595 | 523 | 1.62 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.851 | 1.88e+04 | 0.23 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | primal | 0.854 | 1.87e+04 | 0.35 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 4.469 | 3.58e+03 | 0.93 |
+| em | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | grad | 4.436 | 3.61e+03 | 1.02 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.106 | 5.15e+03 | 0.22 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | primal | 3.117 | 5.13e+03 | 0.36 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 17.610 | 909 | 0.93 |
+| em | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | grad | 17.427 | 918 | 1.05 |
diff --git a/benchmarks/results/vulcan-l40s-e-matmul-highest.json b/benchmarks/results/vulcan-l40s-e-matmul-highest.json
new file mode 100644
index 0000000..b21f9b0
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-e-matmul-highest.json
@@ -0,0 +1,140 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": "highest",
+    "scan_unroll": 1,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0018264910009188147,
+      "trajectories_per_second": 17519.933021242607,
+      "compile_seconds": 0.4346738619497046
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.009604040393605828,
+      "trajectories_per_second": 3331.9310090891477,
+      "compile_seconds": 1.2720381129765883
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.007048667857556471,
+      "trajectories_per_second": 4539.864928618341,
+      "compile_seconds": 0.3383073400473222
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03859191294759512,
+      "trajectories_per_second": 829.1892667630539,
+      "compile_seconds": 1.31303374201525
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0028107422403991222,
+      "trajectories_per_second": 11384.893121845294,
+      "compile_seconds": 0.7347728699678555
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.016519141655104857,
+      "trajectories_per_second": 1937.1466549601955,
+      "compile_seconds": 2.9693069600034505
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.011021795507986099,
+      "trajectories_per_second": 2903.338206267178,
+      "compile_seconds": 0.8703918239334598
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.06651889497879893,
+      "trajectories_per_second": 481.0663197306438,
+      "compile_seconds": 3.8560943111078814
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-e-matmul-highest.md b/benchmarks/results/vulcan-l40s-e-matmul-highest.md
new file mode 100644
index 0000000..f57d609
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-e-matmul-highest.md
@@ -0,0 +1,14 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=highest, scan_unroll=1, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.826 | 1.75e+04 | 0.43 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 9.604 | 3.33e+03 | 1.27 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 7.049 | 4.54e+03 | 0.34 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 38.592 | 829 | 1.31 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 2.811 | 1.14e+04 | 0.73 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 16.519 | 1.94e+03 | 2.97 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 11.022 | 2.9e+03 | 0.87 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 66.519 | 481 | 3.86 |
diff --git a/benchmarks/results/vulcan-l40s-f-cudagraph.json b/benchmarks/results/vulcan-l40s-f-cudagraph.json
new file mode 100644
index 0000000..0e8940a
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-f-cudagraph.json
@@ -0,0 +1,396 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 1,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0009622347191907465,
+      "trajectories_per_second": 1039.247472634343,
+      "compile_seconds": 0.4700410889927298
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005052514995137851,
+      "trajectories_per_second": 197.92123347725294,
+      "compile_seconds": 1.5074586200062186
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003587769236988746,
+      "trajectories_per_second": 278.7247266881944,
+      "compile_seconds": 0.3678640240104869
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.019817276508547366,
+      "trajectories_per_second": 50.46102069417517,
+      "compile_seconds": 1.3597999310586601
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0013455422789168854,
+      "trajectories_per_second": 23782.23300850783,
+      "compile_seconds": 0.6324994110036641
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008021229334796468,
+      "trajectories_per_second": 3989.413425842658,
+      "compile_seconds": 1.570612299023196
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.005146991215749747,
+      "trajectories_per_second": 6217.224521790573,
+      "compile_seconds": 0.4534137600567192
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03185791301075369,
+      "trajectories_per_second": 1004.4600218852487,
+      "compile_seconds": 1.5433854449074715
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006954732269867126,
+      "trajectories_per_second": 1437.869872191508,
+      "compile_seconds": 0.2018867180449888
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0030720798749825917,
+      "trajectories_per_second": 325.51236969568265,
+      "compile_seconds": 0.7667697980068624
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002501146420591364,
+      "trajectories_per_second": 399.8166567807585,
+      "compile_seconds": 0.2249728940660134
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.011790249263867736,
+      "trajectories_per_second": 84.81584889511952,
+      "compile_seconds": 0.7430006139911711
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0008435127194644066,
+      "trajectories_per_second": 37936.59450721572,
+      "compile_seconds": 0.2570032199146226
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004686483694240451,
+      "trajectories_per_second": 6828.147090179157,
+      "compile_seconds": 0.9227630530949682
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003072791689191945,
+      "trajectories_per_second": 10413.982865338676,
+      "compile_seconds": 0.24970263394061476
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.018492621486075222,
+      "trajectories_per_second": 1730.4198879588657,
+      "compile_seconds": 0.9388698430266231
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0017479369237732429,
+      "trajectories_per_second": 572.1030240847115,
+      "compile_seconds": 0.7169918889412656
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.009450687794014812,
+      "trajectories_per_second": 105.81240453560504,
+      "compile_seconds": 2.523357450030744
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006671004501792292,
+      "trajectories_per_second": 149.90246217512384,
+      "compile_seconds": 0.8463023368967697
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.03746189596131444,
+      "trajectories_per_second": 26.69379043262157,
+      "compile_seconds": 2.656470412039198
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021086527101163354,
+      "trajectories_per_second": 15175.56677136964,
+      "compile_seconds": 0.6879883620422333
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.014226612984202802,
+      "trajectories_per_second": 2249.3055821180155,
+      "compile_seconds": 3.116745045990683
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.008247273838302741,
+      "trajectories_per_second": 3880.0700240342067,
+      "compile_seconds": 0.8314209800446406
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.05689899192657322,
+      "trajectories_per_second": 562.40012197923,
+      "compile_seconds": 3.9393029569182545
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-f-cudagraph.md b/benchmarks/results/vulcan-l40s-f-cudagraph.md
new file mode 100644
index 0000000..3f70297
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-f-cudagraph.md
@@ -0,0 +1,30 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=1, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.962 | 1.04e+03 | 0.47 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 5.053 | 198 | 1.51 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 3.588 | 279 | 0.37 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 19.817 | 50.5 | 1.36 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.346 | 2.38e+04 | 0.63 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 8.021 | 3.99e+03 | 1.57 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 5.147 | 6.22e+03 | 0.45 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 31.858 | 1e+03 | 1.54 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.695 | 1.44e+03 | 0.20 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 3.072 | 326 | 0.77 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.501 | 400 | 0.22 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 11.790 | 84.8 | 0.74 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.844 | 3.79e+04 | 0.26 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 4.686 | 6.83e+03 | 0.92 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 3.073 | 1.04e+04 | 0.25 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 18.493 | 1.73e+03 | 0.94 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.748 | 572 | 0.72 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 9.451 | 106 | 2.52 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 6.671 | 150 | 0.85 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 37.462 | 26.7 | 2.66 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 2.109 | 1.52e+04 | 0.69 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 14.227 | 2.25e+03 | 3.12 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 8.247 | 3.88e+03 | 0.83 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 56.899 | 562 | 3.94 |
diff --git a/benchmarks/results/vulcan-l40s-g-unroll8-promoted.json b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.json
new file mode 100644
index 0000000..f2cb095
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.json
@@ -0,0 +1,396 @@
+{
+  "config": {
+    "device": "NVIDIA L40S",
+    "backend": "gpu",
+    "jax_version": "0.9.1",
+    "matmul_precision": null,
+    "scan_unroll": 8,
+    "remat": "none"
+  },
+  "records": [
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0007148562263782052,
+      "trajectories_per_second": 1398.8826887141574,
+      "compile_seconds": 2.184549228986725
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0017632524977670982,
+      "trajectories_per_second": 567.1337492879517,
+      "compile_seconds": 15.08952729194425
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.002500937333227032,
+      "trajectories_per_second": 399.85008289258934,
+      "compile_seconds": 2.118349283002317
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006678907865924495,
+      "trajectories_per_second": 149.72507782327074,
+      "compile_seconds": 15.132703572977334
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0010891472568822114,
+      "trajectories_per_second": 29380.783725795787,
+      "compile_seconds": 2.608467624988407
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003410756928912763,
+      "trajectories_per_second": 9382.081651359584,
+      "compile_seconds": 13.017872113967314
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004007596903565255,
+      "trajectories_per_second": 7984.8349946403105,
+      "compile_seconds": 2.1188945450121537
+    },
+    {
+      "solver": "sra1",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.013283782677414516,
+      "trajectories_per_second": 2408.9523878170153,
+      "compile_seconds": 13.366022419999354
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.000445107000857769,
+      "trajectories_per_second": 2246.650800982444,
+      "compile_seconds": 1.2408874860266224
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0010245928041091778,
+      "trajectories_per_second": 975.9974850393763,
+      "compile_seconds": 7.035418029059656
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0014510225695890508,
+      "trajectories_per_second": 689.169156261445,
+      "compile_seconds": 1.2501748150680214
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003635078767099633,
+      "trajectories_per_second": 275.09720258355856,
+      "compile_seconds": 7.035499937948771
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0006264614267274737,
+      "trajectories_per_second": 51080.559208828665,
+      "compile_seconds": 1.2203648679424077
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0018924760855346297,
+      "trajectories_per_second": 16909.06439695375,
+      "compile_seconds": 6.612741588032804
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0021858520446006546,
+      "trajectories_per_second": 14639.600186592801,
+      "compile_seconds": 1.1961514059221372
+    },
+    {
+      "solver": "em",
+      "drift": "stochastic-growth",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "explicit",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0070658908419621485,
+      "trajectories_per_second": 4528.799087860494,
+      "compile_seconds": 6.868361442000605
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0012508814995128073,
+      "trajectories_per_second": 799.436237876633,
+      "compile_seconds": 2.7070158420829102
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.003360000147949904,
+      "trajectories_per_second": 297.6190345140751,
+      "compile_seconds": 32.03250827000011
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.004642306303139776,
+      "trajectories_per_second": 215.41017216456837,
+      "compile_seconds": 3.114756237016991
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 1,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.013239547687893113,
+      "trajectories_per_second": 75.53128124720217,
+      "compile_seconds": 32.787103537935764
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0016123186422711505,
+      "trajectories_per_second": 19847.193452358795,
+      "compile_seconds": 3.9508127509616315
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 63,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.0060252813855186105,
+      "trajectories_per_second": 5310.955282007246,
+      "compile_seconds": 34.89513242407702
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "primal",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.006146438128780574,
+      "trajectories_per_second": 5206.267325812756,
+      "compile_seconds": 4.017467572935857
+    },
+    {
+      "solver": "rk4",
+      "drift": "neoclassical",
+      "batch": 32,
+      "dim": 2,
+      "n_steps": 255,
+      "dtype": "float32",
+      "save": "steps",
+      "mode": "grad",
+      "noise": "none",
+      "remat": "none",
+      "oom": false,
+      "seconds_per_call": 0.025152009911835194,
+      "trajectories_per_second": 1272.264129672695,
+      "compile_seconds": 44.808164507034235
+    }
+  ]
+}
diff --git a/benchmarks/results/vulcan-l40s-g-unroll8-promoted.md b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.md
new file mode 100644
index 0000000..d947451
--- /dev/null
+++ b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.md
@@ -0,0 +1,30 @@
+# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1)
+
+matmul_precision=None, scan_unroll=8, remat=none
+
+| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s |
+|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.715 | 1.4e+03 | 2.18 |
+| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.763 | 567 | 15.09 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.501 | 400 | 2.12 |
+| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 6.679 | 150 | 15.13 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.089 | 2.94e+04 | 2.61 |
+| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 3.411 | 9.38e+03 | 13.02 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 4.008 | 7.98e+03 | 2.12 |
+| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 13.284 | 2.41e+03 | 13.37 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.445 | 2.25e+03 | 1.24 |
+| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.025 | 976 | 7.04 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 1.451 | 689 | 1.25 |
+| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 3.635 | 275 | 7.04 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.626 | 5.11e+04 | 1.22 |
+| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 1.892 | 1.69e+04 | 6.61 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 2.186 | 1.46e+04 | 1.20 |
+| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 7.066 | 4.53e+03 | 6.87 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.251 | 799 | 2.71 |
+| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 3.360 | 298 | 32.03 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 4.642 | 215 | 3.11 |
+| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 13.240 | 75.5 | 32.79 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 1.612 | 1.98e+04 | 3.95 |
+| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 6.025 | 5.31e+03 | 34.90 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 6.146 | 5.21e+03 | 4.02 |
+| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 25.152 | 1.27e+03 | 44.81 |

From f35bad1a62e1ad11f6ae12a866fa40c085605c61 Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Tue, 4 Aug 2026 17:53:46 -0700
Subject: [PATCH 4/5] perf: skip frozen tails in vmapped solver loops

A batched-predicate lax.cond lowers under vmap to a both-branches
select, so vmapped adaptive solves executed every max_steps attempt
slot regardless of actual work (8-attempt Tsit5 solves at B=32 cost
47.5 ms with max_steps=1024 vs 2.6 ms with 64 on CPU). The skip conds
in the bounded ODE loop (chunk and attempt level), the fixed clipped
scan, both DAE adaptive bodies, and the SDE/SDAE failure tails are now
gated on unvmap_all, a primitive whose batching rule reduces the
predicate over the batch axis, so the tail after the slowest lane
skips for real. Vmapped adaptive primal cost is now budget-invariant
(0.5 ms at every budget, 95x at max_steps=1024); reverse mode improves
3.9x but still stores per-slot scan residuals. Per-lane adaptivity,
values, counters, and AD are unchanged, pinned by
tests/test_vmap_adaptive.py, which also pins the documented
masked-residual pattern for collocation on adaptive output.

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01C8SWZyLFpzVjEX5WCRLJnD
---
 .../2026-08-04_vulcan-l40s-sde-fixed.md       |  24 +++
 docs/ode.md                                   |  38 +++-
 src/tinydiffeq/_unvmap.py                     |  43 +++++
 src/tinydiffeq/dae.py                         |  18 +-
 src/tinydiffeq/ode.py                         |  18 +-
 src/tinydiffeq/sdae.py                        |  13 +-
 src/tinydiffeq/sde.py                         |  15 +-
 tests/test_vmap_adaptive.py                   | 176 ++++++++++++++++++
 8 files changed, 334 insertions(+), 11 deletions(-)
 create mode 100644 src/tinydiffeq/_unvmap.py
 create mode 100644 tests/test_vmap_adaptive.py

diff --git a/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md b/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md
index 3fa453d..48d8623 100644
--- a/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md
+++ b/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md
@@ -58,6 +58,30 @@ saved, SRA1) is ~8 ms rolled and **~3.4 ms with `unroll=8`**.
    remedy if ensembles ever grow to B·n ≳ 10⁷ in grad mode; a package
    `checkpoint_every=` option was considered and deliberately not added.
 
+## Follow-up: vmapped adaptive solves (CPU measurement, same session)
+
+A batched-predicate `lax.cond` lowers under `vmap` to a both-branches
+select, so vmapped **adaptive** solves used to execute every `max_steps`
+attempt slot: a B=32 Tsit5 solve of the neoclassical policy model that
+needs 8 attempts took 2.6 ms at `max_steps=64` but 47.5 ms at
+`max_steps=1024` (CPU). Gating the skip conds on a scalar `unvmap_all`
+predicate (batching rule reduces over the batch axis; the diffrax
+`unvmap_any` trick) makes the vmapped primal budget-invariant: 0.5 ms at
+every budget, a 95× win at `max_steps=1024`. Reverse mode improved 3.9×
+(129 → 33 ms) but still scales with the budget through the scan's stacked
+per-slot residuals — keep `max_steps` realistic when differentiating
+adaptive solves.
+
+Validated on the L40S (same probe, pre-gate vs gated, ms per call):
+
+| B | max_steps | primal before → after | grad before → after |
+|---:|---:|---|---|
+| 32 | 64 | 2.05 → 0.87 (2.4×) | 9.0 → 5.7 (1.6×) |
+| 32 | 256 | 7.82 → 1.16 (6.7×) | 36.2 → 15.4 (2.3×) |
+| 32 | 1024 | 30.7 → 2.50 (12.3×) | 144.3 → 54.7 (2.6×) |
+| 256 | 64 | 2.71 → 1.00 (2.7×) | — |
+| 256 | 1024 | 40.9 → 2.59 (15.8×) | — |
+
 ## Reproduce
 
 ```bash
diff --git a/docs/ode.md b/docs/ode.md
index 38e54c9..3887021 100644
--- a/docs/ode.md
+++ b/docs/ode.md
@@ -51,7 +51,13 @@ conditions never recompiles (pinned by `tests/test_recompile.py`).
 `adaptive_loop="forward"` instead runs a dynamic `lax.while_loop` that
 executes only actual attempts; it supports primal evaluation, JVP, and nested
 forward mode, but JAX cannot transpose it, so reverse mode requires the
-bounded loop. Under `vmap`, both run until the slowest lane finishes.
+bounded loop. Under `vmap`, both run until the slowest lane finishes: the
+bounded loop's skip conds are gated on a scalar all-lanes predicate whose
+batching rule reduces over the batch axis, so the frozen tail after the
+slowest lane is genuinely skipped rather than lowered to a both-branches
+select — a vmapped adaptive solve costs actual attempts, not `max_steps`
+(reverse mode still stores per-slot scan residuals, so keep budgets
+realistic when differentiating).
 
 Exactly one `SaveAt` mode is set:
 
@@ -76,6 +82,36 @@ the reached prefix. Nothing is poisoned; callers that want diverging
 residuals map `jnp.where(sol.ok, x, jnp.inf)`. `sol.num_steps` counts
 attempts, `sol.num_accepted` counts advances.
 
+### Residuals on adaptive output
+
+Collocation-style residuals evaluated on an adaptive
+`SaveAt(steps=True)` rollout keep a static shape by evaluating the pointwise
+residual on **every** padded row and zeroing the tail with the `accepted`
+mask. With the default `fill="last"`, padded rows repeat the last accepted
+state, so the residual there is finite wherever the endpoint is — a single
+`where` is safe for both values and gradients (no double-`where` needed):
+
+```python
+sol = jax.vmap(solve_one)(x_0s)            # SaveAt(steps=True), (B, rows, ...)
+rows = pointwise_residual(sol.xs, p)       # evaluated on all rows in parallel
+masked = jnp.where(sol.accepted, rows, 0.0)
+count = jax.lax.stop_gradient(sol.accepted.sum())      # actual points, inert
+residual = jnp.where(
+    sol.ok[:, None], masked / jnp.sqrt(count.astype(masked.dtype)), jnp.inf
+).reshape(-1)                              # static length B * rows
+```
+
+Zero rows contribute nothing to a least-squares objective, so this is
+exactly the fixed-shape analogue of dropping rejected steps. Two properties
+to choose deliberately: normalizing by the accepted count makes the loss a
+mean over actual collocation points but changes discontinuously when the
+mesh changes, and lanes with more accepted steps contribute more rows — the
+adaptive mesh curvature-weights the collocation. When the residual
+definition needs fixed sample times and a smoother parameter dependence,
+use `SaveAt(ts=fixed_grid)` instead: no mask, no mesh-dependent weighting,
+and the output times carry none of the frozen-mesh ambiguity of adaptive
+internal knots.
+
 `project` (an idempotent clamp, e.g. positivity) is applied at every point
 where the field is evaluated and to every accepted state.
 
diff --git a/src/tinydiffeq/_unvmap.py b/src/tinydiffeq/_unvmap.py
new file mode 100644
index 0000000..c06c58e
--- /dev/null
+++ b/src/tinydiffeq/_unvmap.py
@@ -0,0 +1,43 @@
+import jax
+import numpy as np
+from jax.extend import core
+from jax.interpreters import ad, batching, mlir
+
+# A lax.cond whose predicate is batched is lowered under vmap to a select
+# that executes BOTH branches for every lane, so the adaptive loops'
+# skip-the-frozen-tail conds do no skipping in a vmapped solve: every attempt
+# slot up to max_steps runs for every lane. unvmap_all reduces its boolean
+# across any vmapped axes (its batching rule maps to an all-reduce and an
+# unbatched result), so a cond gated on it keeps a scalar predicate under
+# vmap and skips for real once every lane is finished. Its output is a
+# nondifferentiable bool, hence the zero-JVP rule.
+
+unvmap_all_p = core.Primitive("unvmap_all")
+
+
+def unvmap_all(x):
+    """All-reduce a boolean over any vmapped axes to a scalar predicate."""
+    return unvmap_all_p.bind(x)
+
+
+def unvmap_all_impl(x):
+    return jax.numpy.all(x)
+
+
+unvmap_all_p.def_impl(unvmap_all_impl)
+unvmap_all_p.def_abstract_eval(lambda x: jax.core.ShapedArray((), np.bool_))
+
+
+def unvmap_all_batch(args, dims):
+    (x,) = args
+    (dim,) = dims
+    # None is JAX's not_mapped sentinel: the reduced output has no batch axis.
+    return unvmap_all(jax.numpy.all(x, axis=dim)), None
+
+
+batching.primitive_batchers[unvmap_all_p] = unvmap_all_batch
+# The output is a nondifferentiable bool: no tangent contribution.
+ad.defjvp(unvmap_all_p, None)
+mlir.register_lowering(
+    unvmap_all_p, mlir.lower_fun(unvmap_all_impl, multiple_results=False)
+)
diff --git a/src/tinydiffeq/dae.py b/src/tinydiffeq/dae.py
index dbf15ca..732567d 100644
--- a/src/tinydiffeq/dae.py
+++ b/src/tinydiffeq/dae.py
@@ -31,6 +31,7 @@
     where,
     zeros_like,
 )
+from tinydiffeq._unvmap import unvmap_all
 from tinydiffeq.controllers import ConstantStepSize
 from tinydiffeq.interpolation import (
     hermite_interpolate,
@@ -613,7 +614,14 @@ def skip_step(carry):
         return carry, out
 
     def body(carry, _):
-        return jax.lax.cond(carry[5] | carry[6], skip_step, attempt_step, carry)
+        # Scalar-predicate outer cond under vmap (see _unvmap): the frozen
+        # tail skips for real; the inner cond keeps per-lane freezing.
+        def live_lanes(carry):
+            return jax.lax.cond(carry[5] | carry[6], skip_step, attempt_step, carry)
+
+        return jax.lax.cond(
+            unvmap_all(carry[5] | carry[6]), skip_step, live_lanes, carry
+        )
 
     carry_0 = (
         t_0,
@@ -1528,8 +1536,12 @@ def skip_step(carry):
         return carry, out
 
     def body(carry, _):
-        terminated = carry[8] | carry[9]
-        return jax.lax.cond(terminated, skip_step, attempt_step, carry)
+        def live_lanes(carry):
+            return jax.lax.cond(carry[8] | carry[9], skip_step, attempt_step, carry)
+
+        return jax.lax.cond(
+            unvmap_all(carry[8] | carry[9]), skip_step, live_lanes, carry
+        )
 
     carry_0 = (
         t_0,
diff --git a/src/tinydiffeq/ode.py b/src/tinydiffeq/ode.py
index abba2e9..72839ef 100644
--- a/src/tinydiffeq/ode.py
+++ b/src/tinydiffeq/ode.py
@@ -25,6 +25,7 @@
     zero_tangent,
     zeros_like,
 )
+from tinydiffeq._unvmap import unvmap_all
 from tinydiffeq.controllers import ConstantStepSize
 from tinydiffeq.interpolation import (
     hermite_interpolate,
@@ -427,7 +428,15 @@ def skip_step(carry):
         return carry, out
 
     def body(carry, _):
-        return jax.lax.cond(carry[6] | carry[7], skip_step, attempt_step, carry)
+        # The outer cond keeps a scalar predicate under vmap (see _unvmap),
+        # so the frozen tail is skipped for real once every lane finishes;
+        # the inner cond preserves per-lane freezing while any lane is live.
+        def live_lanes(carry):
+            return jax.lax.cond(carry[6] | carry[7], skip_step, attempt_step, carry)
+
+        return jax.lax.cond(
+            unvmap_all(carry[6] | carry[7]), skip_step, live_lanes, carry
+        )
 
     def fixed_attempt_step(carry):
         t, x, f_cur, done, num_accepted = carry
@@ -467,7 +476,10 @@ def fixed_skip_step(carry):
         return carry, output
 
     def fixed_body(carry, _):
-        return jax.lax.cond(carry[3], fixed_skip_step, fixed_attempt_step, carry)
+        def live_lanes(carry):
+            return jax.lax.cond(carry[3], fixed_skip_step, fixed_attempt_step, carry)
+
+        return jax.lax.cond(unvmap_all(carry[3]), fixed_skip_step, live_lanes, carry)
 
     def uniform_fixed_body(carry, step_index):
         t, x, f_cur, _, num_accepted = carry
@@ -532,7 +544,7 @@ def skip_chunk(chunk_carry, chunk_valid):
             return chunk_carry, repeat_output(output)
 
         def outer(chunk_carry, chunk_valid):
-            inactive = chunk_carry[6] | chunk_carry[7]
+            inactive = unvmap_all(chunk_carry[6] | chunk_carry[7])
             return jax.lax.cond(
                 inactive, skip_chunk, run_chunk, chunk_carry, chunk_valid
             )
diff --git a/src/tinydiffeq/sdae.py b/src/tinydiffeq/sdae.py
index da5d028..9e19e43 100644
--- a/src/tinydiffeq/sdae.py
+++ b/src/tinydiffeq/sdae.py
@@ -20,6 +20,7 @@
     take,
     where,
 )
+from tinydiffeq._unvmap import unvmap_all
 from tinydiffeq.dae import (
     LMRootSolver,
     _canonicalize_cached_dae_field,
@@ -384,10 +385,18 @@ def skip_step(carry, _):
         return carry, out
 
     def body(carry, inputs):
+        def live_lanes(pair):
+            return jax.lax.cond(
+                pair[0][5],
+                lambda pair: skip_step(*pair),
+                lambda pair: attempt_step(*pair),
+                pair,
+            )
+
         return jax.lax.cond(
-            carry[5],
+            unvmap_all(carry[5]),
             lambda pair: skip_step(*pair),
-            lambda pair: attempt_step(*pair),
+            live_lanes,
             (carry, inputs),
         )
 
diff --git a/src/tinydiffeq/sde.py b/src/tinydiffeq/sde.py
index a8224cd..1ba278d 100644
--- a/src/tinydiffeq/sde.py
+++ b/src/tinydiffeq/sde.py
@@ -15,6 +15,7 @@
     take,
     where,
 )
+from tinydiffeq._unvmap import unvmap_all
 from tinydiffeq.ode import canonicalize_field, identity_project
 from tinydiffeq.save_at import SaveAt
 from tinydiffeq.solution import Solution
@@ -239,10 +240,20 @@ def aux_skip(carry, inputs):
             return carry, (t, x, aux, jnp.asarray(False))
 
         def aux_body(carry, inputs):
+            # Scalar-predicate outer cond under vmap (see _unvmap): once every
+            # lane has failed, the frozen tail skips for real.
+            def live_lanes(pair):
+                return jax.lax.cond(
+                    pair[0][3],
+                    lambda pair: aux_skip(*pair),
+                    lambda pair: aux_attempt(*pair),
+                    pair,
+                )
+
             return jax.lax.cond(
-                carry[3],
+                unvmap_all(carry[3]),
                 lambda pair: aux_skip(*pair),
-                lambda pair: aux_attempt(*pair),
+                live_lanes,
                 (carry, inputs),
             )
 
diff --git a/tests/test_vmap_adaptive.py b/tests/test_vmap_adaptive.py
new file mode 100644
index 0000000..e5f4e59
--- /dev/null
+++ b/tests/test_vmap_adaptive.py
@@ -0,0 +1,176 @@
+import jax
+import jax.numpy as jnp
+import pytest
+
+from tinydiffeq import (
+    IController,
+    Rodas5P,
+    SaveAt,
+    Tsit5,
+    solve_ode,
+    solve_semi_explicit_dae,
+)
+
+# Under vmap, a batched-predicate lax.cond lowers to a select that executes
+# both branches, so without the scalar unvmap_all gates the adaptive loops
+# would run every max_steps attempt slot for every lane. These tests pin the
+# gates' semantics (per-lane adaptivity and results unchanged) and the work
+# skipping itself.
+
+X_0S = jnp.linspace(0.5, 2.0, 8)
+
+
+def adaptive_solve(x_0, max_steps=128, save_at=None):
+    return solve_ode(
+        lambda x, t, args, p: p * x * (1.0 - x),
+        Tsit5(),
+        0.0,
+        4.0,
+        x_0,
+        p=jnp.asarray(1.0),
+        dt_0=0.2,
+        controller=IController(),
+        max_steps=max_steps,
+        save_at=save_at,
+    )
+
+
+def test_vmapped_adaptive_matches_scalar_lanes():
+    batched = jax.jit(jax.vmap(lambda x: adaptive_solve(x).xs))(X_0S)
+    stacked = jnp.stack([adaptive_solve(X_0S[i]).xs for i in range(len(X_0S))])
+    assert jnp.allclose(batched, stacked, rtol=1e-12, atol=1e-12)
+    counts = jax.jit(jax.vmap(lambda x: adaptive_solve(x).num_steps))(X_0S)
+    scalar_counts = jnp.stack(
+        [adaptive_solve(X_0S[i]).num_steps for i in range(len(X_0S))]
+    )
+    # Per-lane adaptivity: heterogeneous attempt counts survive batching.
+    assert jnp.array_equal(counts, scalar_counts)
+    assert int(counts.min()) != int(counts.max())
+
+
+def test_vmapped_adaptive_is_budget_invariant():
+    small = jax.jit(jax.vmap(lambda x: adaptive_solve(x, max_steps=64).xs))(X_0S)
+    large = jax.jit(jax.vmap(lambda x: adaptive_solve(x, max_steps=512).xs))(X_0S)
+    assert jnp.array_equal(small, large)
+
+
+def test_vmapped_adaptive_skips_frozen_tail():
+    calls = []
+
+    def field(x, t, args, p):
+        jax.debug.callback(lambda: calls.append(1))
+        return p * x * (1.0 - x)
+
+    def solve(x_0):
+        return solve_ode(
+            field,
+            Tsit5(),
+            0.0,
+            4.0,
+            x_0,
+            p=jnp.asarray(1.0),
+            dt_0=0.2,
+            controller=IController(),
+            max_steps=256,
+        ).xs
+
+    result = jax.block_until_ready(jax.vmap(solve)(X_0S))
+    assert jnp.all(jnp.isfinite(result))
+    # ~20 attempts x 7 stages of actual work; without the scalar gates every
+    # one of the 256 slots would evaluate the field (>1700 calls).
+    assert len(calls) < 600, len(calls)
+
+
+def test_vmapped_adaptive_keeps_scalar_conds_in_jaxpr():
+    jaxpr = str(jax.make_jaxpr(jax.vmap(lambda x: adaptive_solve(x).xs))(X_0S))
+    assert "cond[" in jaxpr
+    assert "unvmap_all" in jaxpr
+
+
+def test_grad_through_vmapped_adaptive_matches_scalar():
+    def batched_loss(x_0s):
+        return jnp.sum(jax.vmap(lambda x: adaptive_solve(x).xs)(x_0s) ** 2)
+
+    def scalar_loss(x_0s):
+        return sum(adaptive_solve(x_0s[i]).xs ** 2 for i in range(len(X_0S)))
+
+    grad_batched = jax.jit(jax.grad(batched_loss))(X_0S)
+    grad_scalar = jax.grad(scalar_loss)(X_0S)
+    assert jnp.allclose(grad_batched, grad_scalar, rtol=1e-10, atol=1e-12)
+    tangent = jax.jvp(batched_loss, (X_0S,), (jnp.ones_like(X_0S),))[1]
+    assert jnp.isfinite(tangent)
+
+
+def test_masked_residuals_on_vmapped_adaptive_output():
+    # The collocation pattern from the docs: evaluate the pointwise residual
+    # on every padded row, zero the tail with `accepted`, normalize by the
+    # inert accepted count. fill="last" keeps padded rows finite, so the
+    # single where is safe for values and gradients.
+    def rollout(p):
+        return jax.vmap(
+            lambda x: solve_ode(
+                lambda x, t, args, p: p * x * (1.0 - x),
+                Tsit5(),
+                0.0,
+                4.0,
+                x,
+                p=p,
+                dt_0=0.2,
+                controller=IController(),
+                max_steps=64,
+                save_at=SaveAt(steps=True),
+            )
+        )(X_0S)
+
+    def residual(p):
+        sol = rollout(p)
+        rows = p * sol.xs * (1.0 - sol.xs) - (sol.xs - 1.0)
+        masked = jnp.where(sol.accepted, rows, 0.0)
+        count = jax.lax.stop_gradient(sol.accepted.sum())
+        scaled = masked / jnp.sqrt(count.astype(masked.dtype))
+        return jnp.where(sol.ok[:, None], scaled, jnp.inf).reshape(-1)
+
+    p = jnp.asarray(1.2)
+    value = jax.jit(residual)(p)
+    assert value.shape == (len(X_0S) * 65,)
+    assert bool(jnp.all(jnp.isfinite(value)))
+
+    # Padded rows contribute exactly zero: the flattened sum of squares
+    # equals the accepted-rows-only sum built lane by lane outside jit.
+    sol = rollout(p)
+    per_lane = 0.0
+    for i in range(len(X_0S)):
+        lane_rows = p * sol.xs[i] * (1.0 - sol.xs[i]) - (sol.xs[i] - 1.0)
+        per_lane += float(jnp.sum(lane_rows[sol.accepted[i]] ** 2))
+    count = float(sol.accepted.sum())
+    assert jnp.allclose(jnp.sum(value**2), per_lane / count, rtol=1e-12)
+
+    grad = jax.jit(jax.grad(lambda p: jnp.sum(residual(p) ** 2)))(p)
+    assert bool(jnp.isfinite(grad))
+
+
+@pytest.mark.parametrize("solver", [Tsit5(), Rodas5P()])
+def test_vmapped_adaptive_dae_matches_scalar_lanes(solver):
+    y_0s = jnp.linspace(0.8, 1.6, 4)
+
+    def solve(y_0):
+        return solve_semi_explicit_dae(
+            lambda y, z, t, args, p: p * z,
+            lambda y, z: z - y,
+            solver,
+            0.0,
+            1.0,
+            y_0,
+            jnp.asarray(0.5),
+            p=jnp.asarray(-0.4),
+            dt_0=0.1,
+            controller=IController(),
+            max_steps=64,
+            save_at=SaveAt(t_1=True),
+        )
+
+    batched = jax.jit(jax.vmap(lambda y: solve(y).ys))(y_0s)
+    stacked = jnp.stack([solve(y_0s[i]).ys for i in range(len(y_0s))])
+    assert jnp.allclose(batched, stacked, rtol=1e-12, atol=1e-12)
+    ok = jax.jit(jax.vmap(lambda y: solve(y).ok))(y_0s)
+    assert bool(jnp.all(ok))

From 6066b8237f5b5ced97c0f1ec55010202753bbd66 Mon Sep 17 00:00:00 2001
From: Jesse Perla <jesseperla@gmail.com>
Date: Tue, 4 Aug 2026 18:26:30 -0700
Subject: [PATCH 5/5] release: tinydiffeq 2.5.0

Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01C8SWZyLFpzVjEX5WCRLJnD
---
 pyproject.toml | 2 +-
 uv.lock        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index bdbf795..1be5027 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tinydiffeq"
-version = "2.4.0"
+version = "2.5.0"
 description = "Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX with static shapes and composable AD"
 readme = "README.md"
 license = "MIT"
diff --git a/uv.lock b/uv.lock
index df761a5..7373b95 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1227,7 +1227,7 @@ wheels = [
 
 [[package]]
 name = "tinydiffeq"
-version = "2.4.0"
+version = "2.5.0"
 source = { editable = "." }
 dependencies = [
     { name = "jax" },