torchdiffeq

torchdiffeq provides a family of ordinary differential equation (ODE) solvers implemented natively in PyTorch, letting neural ODE models run and backpropagate entirely on GPU. It exposes a single odeint interface shared across more than a dozen fixed-step and adaptive-step Runge-Kutta and Adams-Bashforth methods, so switching between solvers such as dopri5, rk4, or tsit5 is a one-line change with no code rewrite.

For memory-constrained training, odeint_adjoint implements the adjoint sensitivity method from the original Neural Ordinary Differential Equations paper, solving a second, backward-time ODE instead of storing the full forward trajectory for backpropagation — giving constant memory cost regardless of how many solver steps were taken. The library also supports differentiable event handling via odeint_event, letting gradients flow through the time at which a trajectory crosses a user-defined condition, useful for physical simulations like a bouncing ball or other early-stopping dynamics.

What You Get

• A single odeint(func, y0, t) interface covering 20+ interchangeable solvers, from adaptive Runge-Kutta (dopri5, dopri8, tsit5, bosh3) to fixed-step and implicit methods (euler, rk4, implicit_adams, radauIIA5).
• odeint_adjoint, a drop-in constant-memory alternative that backpropagates by solving the adjoint ODE rather than storing every intermediate state.
• odeint_event and differentiable event handling for terminating a solve when a user-defined condition is met, with gradients flowing through the event time itself.
• A SciPy solver wrapper (scipy_solver) for accessing SciPy’s ODE integrators through the same PyTorch-native interface.
• Full GPU support end-to-end, since every solver step is expressed as PyTorch tensor operations rather than calling out to a CPU-only numerical library.

Common Use Cases

• Training Neural ODE models where a network layer is replaced by a continuous-time ODE solve, using odeint_adjoint to keep memory usage flat as solver precision increases.
• Building continuous normalizing flows (CNFs) for density estimation and generative modeling, following the pattern in the bundled examples/cnf.py.
• Simulating physical systems with discontinuities — e.g. a bouncing ball — where odeint_event differentiates through the moment of impact.
• Research comparing solver accuracy/speed tradeoffs, by swapping the method argument across adaptive and fixed-step integrators without touching model code.

Under The Hood

Architecture torchdiffeq is organized as a thin public API (torchdiffeq/__init__.py exporting odeint, odeint_adjoint, odeint_event, odeint_dense) sitting on top of an internal _impl package. _impl/odeint.py holds a SOLVERS registry mapping method names to solver classes and dispatches to whichever one is selected after _check_inputs validates and flattens tensor/tuple state. Each solver (dopri5.py, rk4 in fixed_grid.py, tsit5.py, etc.) subclasses one of two abstract base classes in solvers.py — AdaptiveStepsizeODESolver or FixedGridODESolver — giving a uniform integrate(t) contract regardless of stepping strategy. adjoint.py implements the memory-efficient path as a custom torch.autograd.Function (OdeintAdjointMethod) whose backward constructs an augmented dynamics function and re-solves the ODE in reverse time, coupling gradient computation to the choice of solver via the same registry. Event handling (event_handling.py) and dense output (interp.py) are separate concerns layered on top rather than baked into each solver, so most solvers only implement the numerical stepping.

Tech Stack The library targets PyTorch >=1.5.0 for tensor operations and autograd, with SciPy >=1.4.0 pulled in solely to back the scipy_solver wrapper; NumPy is used internally for a handful of coefficient tables. There is no build system beyond a standard setuptools setup.py (no compiled extensions), no web framework, no database, and no CI configuration checked into the repo. Distribution is a pure-Python package published to PyPI.

Code Quality Tests live under tests/ as unittest.TestCase classes and are notably thorough for a research library — odeint_tests.py cross-checks every solver against reference solutions across multiple problems (construct_problem), dtypes (including complex64), devices, and both forward and reverse time, with per-method error tolerances tuned individually (e.g. relaxed rtol/atol for dopri8). gradient_tests.py and event_tests.py extend this to backprop correctness and event-time differentiation. There is no type-hinting and no linter/formatter configuration, and no CI workflow runs these tests automatically on push — they must be run manually via tests/run_all.py.

What Makes It Unique Unlike general-purpose SciPy or MATLAB ODE suites, every solver step here is a native PyTorch tensor op, so integration participates directly in autograd and runs on GPU without any wrapper overhead. Its adjoint implementation is the reference companion to the original Neural ODEs paper and the seminorm-based adjoint speedup paper, making it the de facto standard other continuous-depth-model libraries (e.g. torchdyn, FFJORD) build on or benchmark against, rather than a novel solver algorithm in its own right.