tinydiffeq

Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX: fixed-step Euler/RK4, adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs, and Euler–Maruyama for Itô SDEs and semi-explicit index-1 SDAEs. The package also includes primal, vmap-friendly finite-state DTMC and CTMC simulators with sequential and associative parallel-prefix execution. Deterministic probability forecasts are differentiable in the initial mass and include DTMC matrix powers, dense CTMC exponentials, and matrix-free Arnoldi/Krylov actions over probability pytrees. The same dense and matrix-free backends are available directly through solve_linear_ode for any fixed homogeneous linear array or pytree operator; jvp_linear_ode and vjp_linear_ode apply the exact initial-state tangent and adjoint exponential actions without differentiating Arnoldi orthogonalization. Bounded lax.scan loops provide exactly max_steps attempt slots for fixed and adaptive stepping, so shapes are static, nothing recompiles as tolerances or curvature change, and every solve is differentiable in both forward and reverse mode — including reverse-over-forward, the pattern a Levenberg–Marquardt optimizer with geodesic acceleration needs when it differentiates through a rollout. Adaptive attempts are grouped into static chunks, allowing one lax.cond to skip solver and controller work for an entire padded chunk.

This is a deliberately small, jvp/vjp-friendly package. Rodas5P is a JAX adaptation of Steinebach's method and follows SciML's OrdinaryDiffEqRosenbrock implementation. Use diffrax or SciML if you need general mass matrices, fully implicit or higher-index DAEs, events, continuous solution objects, sparse/Krylov linear solvers for ODE/DAE stages, or specialized adjoints. Initial DAE consistency and explicit DAE stages use nlls-gram.

The linear exponential-action API follows SciML ExponentialUtilities.expv. It includes fixed and residual-controlled adaptive matrix-free time slicing; the latter keeps the Krylov dimension static for predictable JAX compilation. SciML's ExponentialIntegrators.jl is the reference for the broader nonlinear exponential-integrator family.

Install

uv add tinydiffeq

For GPU use, install the JAX accelerator build that matches your hardware, for example:

uv add tinydiffeq "jax[cuda13]"

Minimal example

The vector field may take (x), (x, t), (x, t, args), or (x, t, args, p) — always in that order. args is pass-through data (not an AD target by convention); p holds differentiable parameters (any pytree). The state may also be any JAX pytree. It must contain at least one leaf, and every leaf must be a nonempty real floating array with the same dtype; vector fields and project preserve that structure. Output keeps the structure and adds the saved-time axis to each leaf.

import jax
import jax.numpy as jnp
from tinydiffeq import solve_ode, Tsit5, IController, SaveAt

jax.config.update("jax_enable_x64", True)  # your call — the library never sets it


def f(x, t, args, p):
    return -p * x


sol = solve_ode(
    f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
    p=jnp.asarray(1.3),
    dt_0=0.1,
    controller=IController(rtol=1e-8, atol=1e-10),
    max_steps=512,
    save_at=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),  # fixed output shape,
)                                                  # however many steps adapt
print(sol.xs)   # states on the grid
print(sol.ok)   # reached t_1 within the max_steps budget?

IController() and PIController() choose tolerances from x_0.dtype: rtol=1e-4, atol=1e-6 for float32 and rtol=1e-7, atol=1e-9 for float64. Pass explicit values when tolerances are part of your model's scientific specification. The default dt_min is 10 * finfo(dtype).eps * max(1, abs(t_1)).

max_steps is the total internal attempt budget: accepted steps plus rejections. It is not normally the number of returned times. Endpoint mode returns one time/state, SaveAt(ts=...) returns the requested grid, and SaveAt(steps=True) returns the initial state and accepted internal steps as a contiguous prefix of max_steps + 1 rows. The remaining rows repeat the last accepted state by default; sol.accepted distinguishes data from padding. Rejected attempts never appear in the returned trajectory.

SaveAt(ts=...) also accepts a Python sequence. These are observation times: the adaptive controller still chooses its own internal mesh. Explicit methods use cubic Hermite interpolation; Rodas5P uses its published stiff-aware fourth-order continuous extension.

Semi-explicit DAEs

For a square index-1 system dy/dt = f(y, z, t, args, p) and 0 = g(y, z, t, args, p):

from tinydiffeq import IController, Rodas5P, Tsit5, solve_semi_explicit_dae


def dae_f(y, z, t, args, p):
    return p * z


def dae_g(y, z, t, args, p):
    return z - y, {"flow": p * z}


dae_sol = solve_semi_explicit_dae(
    dae_f, dae_g, Tsit5(), 0.0, 1.0,
    jnp.asarray(1.0), jnp.asarray(0.5),
    p=jnp.asarray(2.0), dt_0=0.1,
    controller=IController(), max_steps=128, has_aux=True,
)
print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"])

# One initial nonlinear consistency solve, then linear Rodas5P stages.
stiff_dae_sol = solve_semi_explicit_dae(
    dae_f, dae_g, Rodas5P(), 0.0, 1.0,
    jnp.asarray(1.0), jnp.asarray(0.5),
    p=jnp.asarray(2.0), dt_0=0.1,
    controller=IController(), max_steps=128, has_aux=True,
)

z_0 is a guess and is made consistent automatically. RK4 and Tsit5 restore the algebraic root at every stage. Rodas5P performs no nonlinear solves after initialization: it advances the corresponding block mass-matrix system using one reused LU factorization per attempt. Algebraic equations may return a floating aux pytree stored at every accepted node and interpolated on requested deterministic grids. JVP, VJP, and reverse-over-forward propagate through both implicit initialization and the time integrator. See the DAE documentation for root controls, SaveAt, and scope limits.

Fixed-step semi-explicit Itô SDAEs use the corresponding solve_semi_explicit_sdae interface with EulerMaruyama, a PRNG key, and n_steps; see the SDAE documentation.

Gradients through the solve

def endpoint(p):
    return solve_ode(
        f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
        dt_0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
        max_steps=512,
    ).xs

jax.grad(endpoint)(jnp.asarray(1.3))                      # reverse mode
jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),))  # forward mode
jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])(
    jnp.asarray(1.3)
)                                                          # reverse-over-forward

The step-size controller is wrapped in stop_gradient (accept/reject is non-differentiable either way, and the error-ratio power blows up at exactly zero error); the states differentiate fully through the solver stages. See the docs for the design contracts: static shapes and SaveAt, AD through adaptive stepping, SDE key semantics, and the package API.

License

MIT