# 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.

## 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/)
- [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/)
