DiffraxDDE

Delay Differential Equation (DDE) solvers in JAX. Autodifferentiable and GPU-capable.

Built on top of Diffrax and the JAX ecosystem.

Installation

pip install diffraxdde

For development:

git clone https://github.com/YOUR_USERNAME/DiffraxDDE.git
cd DiffraxDDE
pip install -e .

Quick Example

import jax.numpy as jnp
import diffrax as dfx
from diffraxdde import ddeint, DenseHistory, FunctionHistory

# Mackey-Glass equation: dy/dt = β*y(t-τ)/(1 + y(t-τ)^n) - γ*y(t)
def vector_field(t, y, args, *, history):
    y_delayed = history(t - args["tau"])
    beta, n, gamma = args["beta"], args["n"], args["gamma"]
    return beta * y_delayed / (1 + y_delayed**n) - gamma * y

# Parameters
args = {"tau": 2.0, "beta": 2.0, "gamma": 1.0, "n": 9.7}

# History function for t < 0
def history_func(t):
    return jnp.array([0.5])

# Solve
sol = ddeint(
    dfx.Dopri5(),
    vector_field,
    t0=0.0,
    t1=100.0,
    dt0=0.1,
    y0=jnp.array([0.5]),
    args=args,
    history=FunctionHistory(history_func),
    saveat=dfx.SaveAt(ts=jnp.linspace(0, 100, 1000)),
)

print(sol.ys)  # Solution values

Features

• JAX-native: Full support for JIT compilation, automatic differentiation, and GPU/TPU acceleration
• Diffrax integration: Works seamlessly with any Diffrax solver (explicit, implicit, stochastic)
• Multiple history types:
  • FunctionHistory: For constant delays with analytical history
  • DenseHistory: For time-varying/state-dependent delays with dense output
  • CompositeHistory: Combine function and dense histories
• Advanced capabilities:
  • State-dependent delays
  • Neutral DDEs
  • Stochastic DDEs (SDDEs)
  • Vectorization with jax.vmap
  • Gradient-based optimization

Documentation

See the examples/ folder for comprehensive examples:

• 01_mackey_glass.py - Classic Mackey-Glass equation
• 02_hutchinson.py - Population dynamics
• 03_ikeda.py - Optical bistability
• 04_coupled_neurons.py - Neural network dynamics
• 09_stiff_dde_working.py - Stiff DDEs with implicit solvers

For API documentation, see API_DOCS.md.

Citation

If you use DiffraxDDE in your research, please cite:

@software{diffraxdde2026,
  author = {Merugu, Praneeth},
  title = {DiffraxDDE: Delay Differential Equation Solvers in JAX},
  year = {2026},
  url = {https://github.com/YOUR_USERNAME/DiffraxDDE}
}

See also: other libraries in the JAX ecosystem

• Diffrax: ODE/SDE/CDE solvers
• Equinox: Neural networks and utilities
• Lineax: Linear solvers
• Optimistix: Nonlinear optimization

License

Apache 2.0. See LICENSE.