# tinydiffeq

> Tiny ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Euler, RK4, adaptive Tsit5, and the SciML-derived linearly implicit Rodas5P use bounded lax.scan loops with static shapes and forward/reverse AD, including reverse-over-forward. 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. See https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock and https://doi.org/10.1007/s10543-023-00967-x. RK4/Tsit5 DAEs use nlls-gram augmented-QR LM roots at every stage; Rodas5P uses LM only for initial consistency. 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: DTMC endpoints use binary matrix powers, dense CTMCs use matrix exponentials, and matrix-free CTMCs use fixed or residual-controlled adaptive Arnoldi time slicing over array or pytree probabilities. The adaptive method keeps the Krylov dimension static and bounds accepted plus rejected attempts with max_steps. The Krylov design follows SciML ExponentialUtilities.expv, stores basis vectors as contiguous rows, defaults to stable two-pass reorthogonalization, and offers an explicitly selected one-pass performance mode. ODE fields and stochastic drifts may return (value, saved_aux), a real-floating pytree saved with the solution and differentiated through JVP/VJP. DAE/SDAE algebraic functions may separately return (residual, algebraic_aux); that internal array pytree is passed to the differential fields but is not itself stored or interpolated. Requested-grid deterministic saved aux uses cubic Hermite interpolation and no query-time nonlinear solves. Fixed-step solve_semi_explicit_sdae applies Euler-Maruyama to the reduced Ito SDE. States may be arrays or arbitrary JAX pytrees; state leaves share one real floating dtype and saved-aux leaves may use different floating dtypes. Adaptive ODE/DAE defaults are rtol=1e-4, atol=1e-6 for float32 and rtol=1e-7, atol=1e-9 for float64; adaptive Krylov defaults are rtol=1e-5, atol=1e-7 and rtol=1e-10, atol=1e-12 respectively. 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 — the bounded-scan design, SaveAt observation grids, compact accepted steps and padding, why nothing recompiles](https://highdimensionaleconlab.github.io/tinydiffeq/static_shapes/)
- [Adaptive stepping and AD — stop-gradiented controller rationale, horizon-clip growth guard, non-differentiable interpolation knots, double-where NaN safety](https://highdimensionaleconlab.github.io/tinydiffeq/adaptive_ad/)
- [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, 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/)
- [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/)
