tinydiffeq 2.5.0: SRA1 SDAEs, unroll, vmapped-loop skipping - #10
Merged
Conversation
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
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
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
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
Co-Authored-By: Mecha Perla (Claude) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01C8SWZyLFpzVjEX5WCRLJnD
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
solve_semi_explicit_sdaeaccepts SRA1 (stage root at t + 3/4 dt plus the endpoint root; noise viasolver.sample_noise, matchingsolve_sde); Milstein stays excluded since its commutativity contract does not survive the implicit reduction.unroll=onsolve_sdeand fixed-stepsolve_ode: identical values, fewer GPU dispatches — 2–3× faster reverse mode on L40S NN-policy ensembles.max_stepsslot. Skip conds are now gated on anunvmap_allprimitive (scalar predicate under vmap). Measured on L40S: vmapped adaptive Tsit5 primal 12–16× faster atmax_steps=1024, grad 1.6–2.6×; per-lane adaptivity, values, and AD unchanged.max_steps_is_successcompat field and the private brownian-increment alias.benchmarks/gpu_trajectories.py) driving trained kernels investment policies as the drift, plus exported policy weights and the Vulcan L40S results memo (benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md).Test plan
uv run pytest -q(340 passed), ruff clean,mkdocs build --strictcleantests/test_gpu.pygreen on Vulcan L40S (24/24)benchmarks/results/🤖 Generated with Claude Code
https://claude.ai/code/session_01C8SWZyLFpzVjEX5WCRLJnD