solve_sde is fixed-step only, and this is an honest limitation: adaptive stepping for SDEs is not just wiring IController into the EM loop. A rejected step must retry over a sub-interval of the same Brownian path — with v1's presampled dW = sqrt(dt) * normal(key, (n_steps,) + shape), the increments are tied to a fixed grid and cannot be subdivided consistently.
What it takes
Re-evaluable noise. A Brownian-bridge data structure that can answer W(s, t) (the increment over any interval) consistently: querying [t0, tm] and [tm, t1] must sum to the previously returned [t0, t1] sample. The standard construction is diffrax's VirtualBrownianTree: a virtual binary tree over [t0, t1], descended to a fixed tolerance/depth, deriving one PRNG key per node by splitting along the path, and sampling midpoints from the bridge conditional
W(mid) | W(l), W(r) ~ N( (W(l)+W(r))/2 + bias, (r-mid)(mid-l)/(r-l) )
Static depth = static shapes, which fits the bounded-scan design: each query costs depth iterations of a lax.scan descent, no dynamic allocation.
Controller interaction. With re-evaluable noise, the bounded scan in ode.py generalizes: the EM (or Milstein) step consumes W(t, t+h) from the tree instead of a scanned-in increment, and a rejection shrinks h and re-queries — the bridge guarantees the retried sub-increment is consistent with the rejected one. The error estimate for accept/reject can be the standard step-doubling comparison (two half-steps vs one full step, sharing the bridge increments), since EM has no embedded pair.
Reference implementation. diffrax's VirtualBrownianTree (Li et al. 2020, "Scalable Gradients for SDEs") is the model — including the Lévy-area caveats that decide which solvers may legally use adaptive steps (EM with diagonal noise is fine).
Suggested scope
VirtualBrownianTree(t0, t1, tol, shape, key) with evaluate(s, t) — standalone, tested for consistency (nested queries sum, same key reproducible, increments N(0, t-s))solve_sde(..., controller=...) accepting it, sharing the ODE loop's carry structure- strong-convergence test reusing the shared-path GBM construction from
tests/test_sde.py, plus a consistency test that halving the tolerance leaves the path (not just the law) fixed
Until then, SaveAt(ts=...)-style output for SDEs also stays out (Hermite interpolation is wrong for rough paths; the tree would enable Brownian-bridge-correct interpolation later).
solve_sdeis fixed-step only, and this is an honest limitation: adaptive stepping for SDEs is not just wiringIControllerinto the EM loop. A rejected step must retry over a sub-interval of the same Brownian path — with v1's presampleddW = sqrt(dt) * normal(key, (n_steps,) + shape), the increments are tied to a fixed grid and cannot be subdivided consistently.What it takes
Re-evaluable noise. A Brownian-bridge data structure that can answer
W(s, t)(the increment over any interval) consistently: querying[t0, tm]and[tm, t1]must sum to the previously returned[t0, t1]sample. The standard construction is diffrax'sVirtualBrownianTree: a virtual binary tree over[t0, t1], descended to a fixed tolerance/depth, deriving one PRNG key per node by splitting along the path, and sampling midpoints from the bridge conditionalStatic depth = static shapes, which fits the bounded-scan design: each query costs
depthiterations of alax.scandescent, no dynamic allocation.Controller interaction. With re-evaluable noise, the bounded scan in
ode.pygeneralizes: the EM (or Milstein) step consumesW(t, t+h)from the tree instead of a scanned-in increment, and a rejection shrinkshand re-queries — the bridge guarantees the retried sub-increment is consistent with the rejected one. The error estimate for accept/reject can be the standard step-doubling comparison (two half-steps vs one full step, sharing the bridge increments), since EM has no embedded pair.Reference implementation. diffrax's
VirtualBrownianTree(Li et al. 2020, "Scalable Gradients for SDEs") is the model — including the Lévy-area caveats that decide which solvers may legally use adaptive steps (EM with diagonal noise is fine).Suggested scope
VirtualBrownianTree(t0, t1, tol, shape, key)withevaluate(s, t)— standalone, tested for consistency (nested queries sum, same key reproducible, increments N(0, t-s))solve_sde(..., controller=...)accepting it, sharing the ODE loop's carry structuretests/test_sde.py, plus a consistency test that halving the tolerance leaves the path (not just the law) fixedUntil then,
SaveAt(ts=...)-style output for SDEs also stays out (Hermite interpolation is wrong for rough paths; the tree would enable Brownian-bridge-correct interpolation later).