tinydiffeq

tinydiffeq is a deliberately tiny set of differentiable ODE/SDE integrators for JAX: fixed-step Euler and RK4, adaptive Tsit5 with an integral step-size controller, and fixed-step Euler–Maruyama for Itô SDEs. Everything runs inside one bounded lax.scan with static shapes, 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.

It is a jvp/vjp-friendly subset of diffrax. Use diffrax instead if you need any of:

• pytree states (tinydiffeq states are single arrays, scalar or vector)
• stiff or implicit solvers
• PID step-size control
• events, root-finding, or backward-time integration
• dense output / continuous interpolation objects
• checkpointed or backsolve adjoints for long horizons

tinydiffeq ships only the O(max_steps)-memory bounded-scan approach, because that is the one that composes cleanly with jax.jvp, jax.vjp, jax.vmap, and reverse-over-forward without custom adjoint machinery.

Install

uv add tinydiffeq

For accelerator use, install the JAX build matching your hardware alongside it, for example:

uv add tinydiffeq "jax[cuda13]"

Vector-field interface

The vector field may take one to four positional arguments — always in this order:

f(x)                # autonomous, closes over everything
f(x, t)
f(x, t, args)
f(x, t, args, p)

• x is the array state (scalar or vector).
• t is time.
• args is pass-through data. By convention it is not an AD target — nothing stops you differentiating with respect to it, but the library's contracts and tests treat it as constants.
• p holds differentiable parameters — any pytree, e.g. neural-network weights. jvp/vjp with respect to p and x0 are first-class and tested.

The arity is inspected once and the function is wrapped into the canonical four-argument form, so the compiled code is identical for all four. There is no special autonomous code path: an unused t is dead-code-eliminated. drift and diffusion in solve_sde follow the same convention.

Minimal example

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

jax.config.update("jax_enable_x64", True)  # your call, not the library's


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),
    dt0=0.1,
    controller=IController(rtol=1e-8, atol=1e-10),
    max_steps=512,
    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),
)
sol.xs  # (21,) states on the grid, however many internal steps were taken
sol.ok  # False if max_steps ran out before t1

Gradients go straight through the solve:

def endpoint(p):
    return solve_ode(
        f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
        dt0=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

Design contracts at a glance

• dt0 is required. There is no initial-step heuristic.
• Forward time only: t1 > t0.
• Never poisons. sol.ok reports whether t1 was reached; callers that want diverging residuals do jnp.where(sol.ok, sol.xs, jnp.inf).
• project (an idempotent clamp, e.g. positivity) is applied at every point where the vector field is evaluated and to every accepted state.
• Never sets jax_enable_x64. The time dtype follows jnp.result_type(x0, float); float32 problems stay float32 even under x64.
• Solvers, controllers, SaveAt, and Solution are frozen dataclasses registered as pytrees: numeric fields (tolerances, grids, dt0, x0) are data leaves, so changing them never recompiles.

Read next: Static Shapes for the bounded-scan design and SaveAt, Adaptive Stepping and AD for what is and is not differentiated, SDEs, and Migration if you are replacing hand-rolled RK4/Tsit5 loops.