Differential Equations with Diffrax

Diffrax, created by Patrick Kidger, is the premier library for solving differential equations in JAX. It supports ordinary differential equations (ODEs), stochastic differential equations (SDEs), and controlled differential equations (CDEs) — all differentiable and JIT-compilable.

import diffrax
import jax.numpy as jnp

# Solve a simple ODE: dy/dt = -y (exponential decay)
def vector_field(t, y, args):
    return -y

term = diffrax.ODETerm(vector_field)
solver = diffrax.Tsit5()  # 5th-order Runge-Kutta (recommended default)
saveat = diffrax.SaveAt(ts=jnp.linspace(0, 5, 100))
stepsize_controller = diffrax.PIDController(rtol=1e-5, atol=1e-5)

sol = diffrax.diffeqsolve(
    term,
    solver,
    t0=0,
    t1=5,
    dt0=0.1,
    y0=1.0,
    saveat=saveat,
    stepsize_controller=stepsize_controller,
)

print(sol.ts.shape)  # (100,) — time points
print(sol.ys.shape)  # (100,) — solution values
# sol.ys ≈ exp(-ts) — the exact solution

Neural ODEs: Combining Diffrax + Equinox

A Neural ODE uses a neural network as the vector field. Diffrax makes this straightforward:

import diffrax
import equinox as eqx
import jax
import jax.numpy as jnp

class NeuralODE(eqx.Module):
    """A Neural ODE: dy/dt = neural_network(t, y)"""
    net: eqx.nn.MLP

    def __init__(self, data_dim, hidden_dim, *, key):
        self.net = eqx.nn.MLP(
            in_size=data_dim + 1,  # +1 for time
            out_size=data_dim,
            width_size=hidden_dim,
            depth=2,
            key=key,
        )

    def __call__(self, t, y, args=None):
        t_expanded = jnp.broadcast_to(t, (1,))
        inp = jnp.concatenate([y, t_expanded])
        return self.net(inp)

# Create the Neural ODE model
key = jax.random.key(42)
model = NeuralODE(data_dim=2, hidden_dim=64, key=key)

# Solve the ODE (forward pass)
term = diffrax.ODETerm(model)
solver = diffrax.Tsit5()
controller = diffrax.PIDController(rtol=1e-3, atol=1e-3)

y0 = jnp.array([1.0, 0.0])
sol = diffrax.diffeqsolve(
    term, solver, t0=0, t1=1, dt0=0.1, y0=y0,
    stepsize_controller=controller,
)

# The whole thing is differentiable!
@eqx.filter_jit
@eqx.filter_value_and_grad
def loss_fn(model):
    term = diffrax.ODETerm(model)
    sol = diffrax.diffeqsolve(
        term, diffrax.Tsit5(), t0=0, t1=1, dt0=0.1, y0=y0,
        stepsize_controller=diffrax.PIDController(rtol=1e-3, atol=1e-3),
    )
    return jnp.mean(sol.ys[-1] ** 2)  # minimize final state

loss, grads = loss_fn(model)
print(f"Loss: {loss:.6f}")