jax-cfd: differentiable fluid dynamics in JAX, and why it stopped

The first line tells you what has happened

The readme opens with a banner, and it is the most important sentence in the document.

JAX-CFD is no longer maintained. For alternatives, the readme suggests JAX-Fluids, Phi Flow and Exponax, with links to each.

That is a good way to end a research project, and it is worth pausing on what the alternative list tells you. They are more or less contemporaneous with the wave of machine learning for scientific computing that JAX-CFD was part of.

So the project is best understood as a completed piece of research rather than a library with a maintenance burden. It is still Apache 2.0 licensed, still installable, and still the clearest description of what a differentiable CFD solver looks like when you write it from scratch in a framework that composes functions.

The last push was on 2026-10-15, so somebody has been touching the tree, but the tagged releases stop in 2024, and the banner is the statement of intent.

What the project was for

The description is precise about the intent, and the intent is exploration rather than production.

It is an experimental research project for exploring the potential of machine learning, automatic differentiation and hardware accelerators for computational fluid dynamics, implemented in JAX.

Three technologies, and the reason each is there is worth spelling out. Automatic differentiation is what makes a solver differentiable end to end, which is the precondition for learning a correction to it. Hardware accelerators are what make it worth doing, since a differentiable solver that runs on a CPU is a curiosity. And machine learning is the point of the exercise, and the paper the project is built around is the 2021 PNAS study on machine learning accelerated computational fluid dynamics.

The authors are listed, and they are the people behind the work rather than a company team: Dmitrii Kochkov, Jamie A. Smith, Peter Norgaard, Gideon Dresdner, Ayya Alieva and Stephan Hoyer.

The general approach is documented in a paper in the Proceedings of the National Academy of Sciences from 2021, on machine learning accelerated computational fluid dynamics, and the readme asks you to read that to understand the approach. The code is the artefact that follows from the paper, and the paper is where the argument is.

The project also lists a downstream use, a project on variational data assimilation with a learned inverse observation operator, which is a nice illustration of the kind of work the library enables.

The notebooks are the interface

The getting started section is short and points at a notebooks directory, and the list of what is in it is a good map of the project's capabilities.

There are demos of different simulation setups. A two-dimensional simulation with a finite volume method on a staggered grid. A two-dimensional simulation with a pseudo-spectral solver, which is the forced turbulence case. A two-dimensional simulation of channel flow. And a finite volume simulation on a collocated grid, which is marked as experimental.

Then there are two notebooks for reproducing the results from the paper: one for data analysis and evaluation, and one for running the pre-trained models.

That last pair is the part that makes this a research artefact rather than a library. Shipping the trained models and the analysis code alongside the solver means the paper's claims can be checked, and it means the interesting part of the project, the machine learning, is usable without retraining anything.

The notebooks are hosted so you can open them in Colab, which means the demos run on a hosted machine with a GPU and no local environment at all. For a project whose entire value proposition is running on accelerators, having the examples reachable in a hosted notebook is the difference between a demo people try and a demo people read about.

The primary language of the repository is listed as a Jupyter notebook, which is a direct consequence of that choice.

Four submodules and three install shapes

The organisation section describes the package as a set of submodules, and the split is a sensible one.

The base module holds the core finite volume and difference methods, written in JAX. The spectral module holds the core pseudo-spectral methods, also in JAX. The machine learning module holds the learning augmented models, written in JAX and in Haiku, the deep learning library. And the data module holds utilities for preparing, evaluating and post-processing the data the other modules produce, written in Xarray and Pillow.

The dependency structure follows from that. A base install needs only NumPy, SciPy and JAX, and the packaging file shows the actual base list also includes a tree library. The data extra adds the plotting and array libraries, the machine learning extra adds the Haiku ecosystem, and the test extra adds the test and image libraries. There is a complete extra that is all of them.

So the install is three shapes. If you want to run the solver, the base install is enough.

pip install jax-cfd

If you want the learning half, you add one extra, and if you want the data pipeline and the tests, you add the rest.

That granularity is unusual for a research code base and it is the right choice, because the base solver is the part that other people want to build on and the part that has the fewest moving pieces.

The numerics section, which is the honest part

The numerics section is the part of the readme most worth reading, because it says exactly what the solver does and, more importantly, what it does not.

The scope is unsteady turbulent flows. The spatial discretisation is finite volume or difference methods on a staggered grid, the Arakawa C or MAC grid, with pressure at the centre of each cell and the velocity components on the corresponding faces. Alongside that there are pseudo-spectral methods for vorticity, which use anti-aliasing filtering for the non-linear terms to maintain stability.

Then the limitations, listed plainly. Temporal discretisation is currently only first order, with explicit time stepping for advection and either implicit or explicit for diffusion. Pressure solves are either conjugate gradient or fast diagonalisation with real-valued fast Fourier transforms, which suits periodic boundary conditions. Boundary conditions are currently only periodic. Advection uses second order accurate Van Leer schemes. And the closures are Smagorinsky eddy viscosity models.

First-order time stepping is the one to notice, since a first order scheme in time will limit accuracy for anything unsteady.

The list of future work is honest in the same way. Colocated grids, alternative boundary conditions including non-periodic and immersed boundary methods, higher order time stepping, geometric multigrid, steady state simulation such as RANS, and distributed simulations across multiple accelerators. The readme asks people to reach out before starting significant work, which is the polite version of a note about how a research code base handles contributions.

Two papers to cite, and a development setup

The citation section asks for two different references depending on what you used, which is a small sign of how the project has grown.

If you use the finite volume method or the machine learning models, the request is the 2021 PNAS paper. If you use the spectral code, the request is a separate 2022 preprint on learning to correct spectral methods for simulating turbulent flows.

That split is a good indication of the project's structure. The spectral work came later, has its own authors and its own paper, and is a distinct contribution, and the citation instructions are precise enough that you get the right one.

The local development section is short. You clone the repository, install a matching JAX build, and install the package in editable mode with all the extras, and then the test suite is run with a command that spreads the tests over the available cores, distributes them by file rather than by test, and skips one specific file.

The skipped file is the base validation test, and the reason is not given. That is the kind of small gap you will find in a project whose maintainers have moved on, and it is a reminder that running the suite is a way to find out what state the tree is in rather than a formality.

The tree itself is small: a package directory, the notebooks, a setup file, a contributing guide and a licence. There is no separate documentation site, because the readme and the notebooks are the documentation.

Editorial conclusion

JAX-CFD is worth reading rather than adopting. The code is a clean, well documented example of how to write a CFD solver so that it can be differentiated and run on accelerators, and the numerics section is unusually honest about what the solver does and does not do. It is a poor fit as a dependency for new work, since the project is unmaintained and the readme itself points elsewhere, and it is a poor fit for anything beyond unsteady turbulent flows in periodic domains. If you want the ideas rather than the code, follow the three alternatives the readme names, and if you want to read the original, the two papers it asks you to cite are the better starting point than the repository.

Frequently asked questions

Official sources

1. google/jax-cfd on GitHub
2. Issues
3. License: Apache-2.0
4. README
5. Releases