Finite volume Navier-Stokes that backpropagates through itself

Finite volume on a Cartesian grid, written entirely in JAX

The solver integrates the Navier-Stokes equations for compressible flows with the finite volume method on a Cartesian grid, and the whole thing is expressed in JAX primitives rather than in a lower-level language with a Python wrapper. That single choice is what the package is for.

Because the numerical scheme is written in JAX, automatic differentiation can be applied through it, and the stated purpose is end-to-end optimisation of numerical models. A discrete scheme normally resists gradients, since the update is a branchy sequence of reconstructions and flux solves, so the usual approach is to hand-differentiate it or to treat the solver as a black box. Here the solver is the differentiable function.

The three hardware targets follow from the same decision. JAX-Fluids runs on CPU, GPU and TPU, is parallelised using JAX primitives, and the project reports being tested on up to 512 NVIDIA A100 GPUs and on up to 2048 TPU-v3 cores. The largest example cases sit in the same range: a shock-bubble interaction of roughly 800 million cells on a TPUv3-64 and another of roughly 2 billion cells on a TPUv3-256.

Five reconstructions, five Riemann solvers, three time steppers

The scheme choices are exposed as separate menu entries rather than fixed to one method, which is what a researcher changing a discretisation needs. Time stepping is explicit, with Euler, RK2 and RK3 available.

Spatial reconstruction is high-order and adaptive, and the list covers WENO at three, five and seven, WENO-CU6, WENO-3NN and TENO. Adapting the stencil order rather than fixing it is the mechanism the example notebooks demonstrate changing on a running case.

Fluxes come from five Riemann solvers: Lax-Friedrichs, Rusanov, HLL, HLLC and Roe. Together with the reconstruction family and the time stepper, that is a wide enough grid that a study of numerical choices is possible without leaving the package.

Turbulence is handled by one implicit sub-grid-scale model, ALDM, rather than by a menu of closures. Positivity-preserving techniques are available for the same reason the rest of the list is: the states these schemes produce can otherwise go non-physical under aggressive settings.

Two-phase arrives as level-set or diffuse interface

Two-phase flow is supported two ways. The level-set method tracks an interface through a scalar field, and the diffuse-interface method spreads the phase across a finite thickness. Both are named in the feature list, and the two example cases use one each, so the choice is a per-case decision rather than a global mode.

The same level-set machinery also handles immersed solid boundaries, which is how the space shuttle case at Mach 2 is set up: a solid body is represented in the fluid grid instead of being cut out of it. That avoids staircase geometry on curved surfaces, at the cost of an extra field to advect.

Forcings are available for temperature, mass flow rate and kinetic energy spectrum, and boundary conditions cover symmetry, periodic, wall, Dirichlet and Neumann. Those three forcings matter more than they sound: they are what lets a case hold a shock tube steady or drive turbulence to a target spectrum, rather than only playing out and decaying.

A video of shock-induced collapse of air bubbles in water is linked as a demonstration, which is the regime the level-set path is built for.

The GPU install and the CPU install are not the same lines

There are two documented install paths, and only the GPU one is worth reading closely. For a CPU-only machine, upgrade JAX to the cpu extra, clone the repository, change into the directory and install the package.

pip install --upgrade "jax[cpu]"
git clone https://github.com/tumaer/JAXFLUIDS.git
cd JAXFLUIDS
pip install .

For CPU and GPU together, the instruction is to install JAX with GPU support first, by either pip CUDA and cuDNN wheels or a manual CUDA install, with the pip wheels recommended. Then the same clone and an editable install:

pip install --upgrade "jax[cuda12]"
git clone https://github.com/tumaer/JAXFLUIDS.git
cd JAXFLUIDS
pip install -e .

So the GPU decision is made before the package is involved, and the solver inherits whatever JAX build you have. Two details are easy to miss. Editable installs use pip install -e . and matter for code development on your own machine. And on an M1 Mac the project points at a jaxlib discussion upstream rather than offering a fix of its own.

Version and metadata live in setup.py, not in pyproject.toml

The build configuration is split oddly. pyproject.toml contains only a build-system table naming setuptools and wheel, with setuptools.build_meta as the backend. Every piece of package metadata is in setup.py instead, which is the older path, kept working as a fallback.

The version is not declared anywhere as a value. setup.py opens src/jaxfluids/__init__.py, scans for the line starting with __version__ and parses the string after the equals sign, so the version is whatever that file says and editing it is a source change rather than a packaging change.

The declared floor is Python 3.11 or newer. The dependencies are flax, gitpython, h5py, jax, jaxlib, matplotlib, numpy, optax and gymnasium, which says a lot about the design: optax for optimisation, flax alongside it for the JAX model ecosystem, h5py for field storage, gymnasium because the package is aimed at machine-learning workflows, and gitpython for data provenance. The only extra is cuda, and it pulls in jaxlib alone, so the CUDA setup is still yours to do. Packages are discovered under src/, which is why the import name and the repository name differ.

The Makefile still lints a package named qlbm

The Makefile is a help-target dispatcher: running make with no target parses the comments and prints the available goals, which is the first thing you see and the friendliest part of the build.

The targets themselves are a generic Python template rather than something written for this solver. The style targets run flake8 and black against qlbm and tests. The coverage target runs coverage with --source qlbm. Nothing in this repository is called qlbm: the package lives in src/ under the name jaxfluids.

What that means in practice is that the lint and coverage targets point at directories that do not exist, so they either fail or report nothing, and anyone trusting them is reading a green result that covers no code. The other targets are generic in the same way, with clean, test running pytest, test-all running tox across Python versions, and docs building Sphinx HTML.

Useful, but not specific to CFD. Anyone running quality checks here should invoke flake8, black and coverage against src/ themselves.

Three tagged releases, the newest from March 2025

The release history is short. Version 0.1.0 is dated 22 May 2024. Versions 0.2.0 and 0.2.1 were both published on 21 March 2025, an hour apart, which reads as a packaging correction rather than a feature release. Nothing has been tagged since.

The branch has not stopped moving. The last push to the main branch is dated 28 August 2026, and the repository is not archived, so the gap between the newest tag and the newest commit is a release-discipline gap rather than an abandonment one. Anyone installing should therefore pin a commit rather than a tag.

The licence is GPL-3.0-only, stated in the setup metadata, in the classifier and in the root licence file. For a solver meant to be imported into an optimisation pipeline or an RL environment, that is the fact to read twice, because it governs what you can do with a derivative work.

There are also two papers to cite if the work goes into a publication: the 2023 JAX-Fluids paper in Computer Physics Communications and the 2025 JAX-Fluids 2.0 paper in the same journal.

Three papers use it as infrastructure

The project keeps a list of work built on it, and it is short enough to read in full. Two entries are from the team itself. One is a 2025 Physical Review Fluids paper on data-driven shape inference in three-dimensional steady-state supersonic flows, which optimises a discrete loss with the solver. The other is a 2025 Journal of Computational Physics paper, ML-ILES, on end-to-end optimisation of high-order Godunov-type schemes for compressible turbulence.

The third is external: a 2026 Nature paper on HydroGym, a reinforcement learning platform for fluid dynamics. That is the clearest evidence for the differentiable claim, since a reinforcement learning environment needs to produce large numbers of gradient-friendly state transitions rather than one converged field.

What the list implies about the intended use is narrower than differentiable CFD in general. The work that cites the solver optimises things inside the scheme, or puts it inside a loop that has to run many times. It is not being used as a general-purpose incompressible flow package.

Editorial conclusion

It fits a researcher optimising a numerical model end to end, or a reinforcement learning environment where the state is a fluid and the simulator must be fast enough to sample. Three things to settle first. The licence is GPL-3.0-only, which matters if you intend to link the solver into a codebase you do not want to carry that licence. Python 3.11 or newer is required, and the CUDA route depends on your own JAX install rather than on the package, so getting the GPU path right is a JAX question before it is a JAX-Fluids question. And the Makefile still lints a package name from a project template, so do not take its lint and coverage targets at face value; run the tools against src/ yourself.

Frequently asked questions

Official sources

1. Issues
2. README
3. Releases
4. tumaer/JAXFLUIDS on GitHub