# tinydiffeq

> 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, 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 — solve functions, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/)
