Diffrax and JAX: a practical guide to ODEs, stochastic simulations, and neural ODE

A comprehensive practical guide to Diffrax and JAX has been released, taking readers from their first launch of a scientific Python stack through building and training neural ODE. This is not a theoretical overview, but a sequential code walkthrough where ordinary and stochastic differential equations are assembled into a single working pipeline.

Where it starts

The guide begins with the most practical but essential layer: a clean computational environment. The author reinstalls `numpy`, `jax`, `jaxlib`, `diffrax`, `equinox`, `optax`, and `matplotlib` to eliminate dependency conflicts and obtain a reproducible notebook. After this, using the example of logistic growth, the main workflow with Diffrax is demonstrated: specifying the term, choosing the adaptive solver `Tsit5`, configuring the step `dt0`, and saving results via `SaveAt`. All of this is immediately accompanied by verifiable and executable code.

The focus then shifts to numerical accuracy, not just the fact of solving the equation. The example uses `PIDController` to manage precision through `rtol` and `atol`, and dense interpolation allows querying values at arbitrary time points without recalculating the entire trajectory. For a researcher, this is an important detail: the obtained solution can immediately be used both for plots and analysis, and as a foundation for training a subsequent model.

What scenarios are covered

After the basic ODE example, the material expands to tasks that are closer to real research or ML pipelines. Here, Diffrax is presented not as an educational toy, but as a flexible interface over JAX that equally comfortably works with classical dynamical systems, structured states, and batch computations. Thanks to this, the library is shown to be suitable not only for a single textbook equation, but for serial simulations with different input data structures.

• Lotka–Volterra system for modeling predator–prey dynamics
• PyTree-state for a spring–mass–damper system
• Batch runs through `jax.vmap` for multiple trajectories at once
• Ornstein–Uhlenbeck stochastic equation with `VirtualBrownianTree`
• Plots of trajectories and metrics to verify the result

It is particularly useful that all these scenarios are arranged in increasing order rather than dumped into one overloaded notebook. First the reader sees ordinary ODE, then work with PyTree-states, after this batched solves, and only then SDE with a Brownian process. This order lowers the entry barrier and provides a clear mental model: the same API expands for new types of tasks without changing tools and without switching to another stack of numerical libraries.

How neural ODE is assembled

The final section is devoted to neural ODE and is constructed as practically as possible. First, a synthetic dataset is created from a second-order physical system: the basic dynamics are solved by an ordinary solver, noise is added to the trajectory, and the resulting series become the training target. Then a compact model is built on Equinox, where an MLP takes the current state and time, predicts derivatives, and Diffrax again integrates them into a continuous trajectory. This approach clearly shows the connection between simulation and training.

On top of this scheme, the author adds a complete training loop with a loss function that calculates the mean squared error between the target and predicted trajectory, and the `optax.adam` optimizer. Thanks to `eqx.filter_jit`, both training and solver are compiled into JAX, and at the end latency is also measured for the already compiled solution. As a result, the guide answers two applied questions at once: how to train neural ODE on data and what computational cost this might have in practice.

What this means

Diffrax is increasingly transforming from a narrow library for numerical analysis into a convenient entry point to continuous-time ML within the JAX ecosystem. The value of this guide is that it links ODE, SDE, batching, JIT compilation, and neural ODE into one sequential scenario and helps transition faster from a formula on paper to code that can be executed, measured, and further trained in real experiments without unnecessary tool integration.