diffrax: one solver machinery for three equation types, and a version floor the readme gets wrong

Three equation types, one formulation, and everything else follows

The architectural claim is stated once, in a sentence marked as a technical aside, and it is the sentence the rest of the page depends on. Ordinary, stochastic and controlled differential equations are solved in a unified way rather than being treated separately, which the page credits for producing a small, tightly written library.

That decision is what makes the rest of the feature list coherent. If the three equation types share a formulation, then the things that vary between problem classes can be uniform too. So the integration region itself can be transformed, which is not a detail: it means an integral over an arbitrary domain, or a domain that a transformation code can construct, is not a separate code path. The state can be a pytree rather than a flat vector. Dense solutions are available. And several adjoint methods exist for backpropagation, rather than one, because the differentiation strategy is separable from the equation type.

The feature list is short enough to print in full: solvers for the three equation types, a selection of solvers including a high order explicit Runge-Kutta pair, a symplectic family and implicit ones, transformable everything including the region of integration, pytree state, dense solutions, multiple adjoint methods for backpropagation, and support for neural differential equations.

Read in that order, the list is not a feature list at all. It is one design decision and its consequences, which is a good sign for a library that describes itself as tightly written.

The wheel ships one directory, which is how small the claim is

The claim that the library is small is checkable, because the build configuration is one line. The packaging includes a single pattern, the package directory and nothing else, so the distributed artifact is exactly one package.

That is worth noting for a library that supports three equation types, several solver families, dense output, multiple adjoints and automatic differentiation. There is no plugin directory, no optional backend package, no per-solver subpackage visible in the build configuration. If the solvers for ordinary, stochastic and controlled equations share a formulation, they share a directory too, and the wheel reflects the design rather than the marketing.

The rest of the repository layout supports the same picture. The top level holds the package directory, a documentation directory, a configuration file for the documentation build, the readme, a contributing guide, a license, a benchmarks directory, an images directory, a pre-commit configuration and a test directory.

Two of those deserve a note. The benchmarks directory is real and present, and the readme never mentions it, which is the subject of a later section. And the test directory is a single top level entry in a project that puts its runtime type assertions in the test dependency group, which says something about how this codebase checks itself.

The readme promises Python 3.10 and the manifest requires 3.11

There is a version floor conflict, and it is the kind that produces a confusing error rather than an obvious one.

The installation section says the library requires Python 3.10 and newer. The project manifest says the same thing as 3.11 and newer. Both statements are in the repository, four sections apart, and they disagree.

The practical consequence is that the manifest wins and the readme misleads. Package resolution reads the manifest, so somebody on 3.10 following the readme will not get an import error deep inside the library, they will get the installer refusing to install it, with no mention of a readme to consult. The classifiers do not help, because the list names Python 3 generically and no minor version at all.

Two other metadata details are worth a look. The development status is classified as alpha while the version sits at 0.7.2, which is a defensible combination for a numerically sensitive library but does mean the classifier and the version number are telling you different things about maturity.

And the intended audience list includes financial services and insurance alongside developers and science. For a library of differential equation solvers that is an odd entry, and it appears to be one of those classifier sets copied from another project rather than chosen, which is a small reminder that the metadata is written by hand.

The version number itself is static rather than derived, and it currently matches the newest release exactly. The branch has been pushed considerably more recently than that release, so the tree carries the release's version number while containing later work.

The linting tools are pinned exactly and the test tools are not pinned at all

The manifest declares four dependency groups, and comparing them side by side shows where the reproducibility effort went.

The development group has four entries, every one pinned to an exact version: a pre-commit runner, a type checker, a linter and formatter, and a toml sorting tool. The documentation group has eleven entries, and every one of those is pinned to an exact version as well, down to the patch level of a plugin that includes or excludes files from a documentation build.

The test group has six entries and not one of them carries a version constraint. A runtime type assertion library, the JAX binary package, an optimiser library, the test runner, a scientific library and a progress bar. The runtime dependencies are in the same style, seven of them with version floors.

So the two groups that only ever run on one maintainer's machine are pinned to the patch, and the two groups that run in continuous integration on every commit are allowed to float. That is a strange allocation of effort if reproducibility is the goal, and it is a specific one: a documentation build is reproducible, a test run is not.

The type checking choice is worth naming as well, since it is unusual. The project uses pyright rather than the more common mypy, with a configuration section of its own, and pairs it with a runtime assertion library in the test group. So shapes and dtypes are annotated statically through one library and checked at runtime through another, which is a heavier setup than most projects adopt and a reasonable one for a library where a silently wrong shape should be an error rather than a wrong answer.

Three of the seven runtime dependencies are by the same author

The runtime dependency list has seven entries. Three of them are libraries by the project's own author: the neural network library that the readme describes as handling neural networks and everything not already in core JAX, the linear solver library, and the root finding and optimisation library. The remaining four are JAX itself, a typing annotation library, a typing extensions backport, and one more whose role is not numerical at all.

That last one is a pretty-printing library for structured data, and it is there so that expression trees print in a readable form when something goes wrong. It is the only dependency in the list with no obvious numerical role, and its presence tells you something about how much the project invests in error messages.

That coupling is not hidden, because the readme has a section pointing at the wider JAX ecosystem, and it is candid about the proportions. Three groups are given: libraries that are always useful, which is where the author's own neural network and typing libraries sit; deep learning libraries, including an optimiser library, a checkpointing library, a training library, and a parameterisation library; and scientific computing, which holds the author's own linear solver, root finding, probabilistic sampling, and symbolic expression conversion, alongside a symbolic regression package.

That last entry carries a parenthetical acknowledgement that it is not a JAX library, and a friendly label for it. Which is the tone of the whole section: a maintainer listing his own ecosystem and being upfront that the best tool in one category comes from elsewhere. It also means a reader evaluating dependencies is looking at a chain, since three of the seven runtime requirements are released by the same person and a failure in any of them is a failure here.

The four line example uses a solver the feature list does not name

The quick example is short enough to read in full, which is the point of including it:

from diffrax import diffeqsolve, ODETerm, Dopri5
import jax.numpy as jnp

def f(t, y, args):
    return -y

term = ODETerm(f)
solver = Dopri5()
y0 = jnp.array([2., 3.])
solution = diffeqsolve(term, solver, t0=0, t1=1, dt0=0.1, y0=y0)

Four things in nine lines are worth pointing out. The equation is the simplest possible decay, which keeps the example about the calling convention rather than the mathematics. The term and the solver are separate objects, so the equation and the integration method are independent choices. The initial state is a two element vector rather than a scalar, so the pytree state feature shows up even here. And the term function takes an extra argument that the example never uses, which is a hook for parameters without which a batched problem would need a different signature.

One small mismatch. The feature list names a fifth order explicit pair, an eighth order explicit pair, symplectic solvers and implicit solvers. The example uses a fifth order explicit pair, and the page explains that this refers to the standard fifth order method with a fourth order embedded estimate. So the example demonstrates a solver that is not among the four families advertised, while explaining the order pair rather than the family.

The step size argument is set explicitly rather than left to a default, which tells you the solver does not pick one for you, and that the cost of an adaptive method is on the caller to think about.

A benchmark directory in the tree, and no performance claim anywhere

The repository contains a benchmarks directory. The readme never mentions it, and makes no claim about speed, accuracy or memory at any point.

For most libraries that would be unremarkable. For a numerical solver library it is the absence that stands out, because in this field the comparison against other solvers on accuracy and cost is the first question a reader asks and usually the first thing the documentation answers. This page answers it in the ecosystem section instead, by pointing at other JAX projects and, at the end, at a longer list of projects to read.

So a reader is given a good picture of the architecture and a good picture of the neighbourhood, and no numbers. The information is not obviously absent from the project, given the directory, but it is absent from the place a reader would look.

The citation section has its own quirk. What the project asks you to cite, if you used it in academic research, is the author's 2021 doctoral thesis on neural differential equations from the University of Oxford, with a separate link to a preprint. So the requested citation predates the library by several years and describes a research direction rather than a description of the solvers. The section ends with a note that you might consider starring the project, which is a different kind of request attached to the same paragraph.

It is a small library with a clear idea behind it, a real benchmark suite, and a readme that spends its limited length on architecture and neighbours rather than on results.

Editorial conclusion

If you are coming from a solver that treats each equation type as its own problem, the reason to look at this library is that it does not, and every other feature on the page is downstream of that. Solving ordinary, stochastic and controlled equations through one formulation is what lets the integration region, the state representation and the adjoint method be swapped the same way in all three cases, and it is why the wheel is one package directory rather than three. Two things to check before you build on it. First, the readme says Python 3.10 and the manifest requires 3.11, and the installer will enforce the manifest, so on 3.10 you will get a resolution failure rather than an import error and the readme will not have warned you. Second, the package is classified as alpha while sitting at version 0.7.2, and that version number has not moved since the release in February even though the branch has been pushed since, so the version string does not tell you whether you have the release or the branch. On accuracy and speed, this page says nothing at all, which is a real gap in a numerical library even though the repository does contain a benchmark directory. And the citation it asks for is the author's 2021 doctoral thesis, not a paper about the library, so cite it if you are writing up work that used the library, but do not expect it to point at a description of the solvers.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. patrick-kidger/diffrax on GitHub
4. README
5. Releases