JAX-Fluids: A Differentiable Compressible CFD Solver for ML-CFD Research

The gap JAX-Fluids fills: gradients through a compressible flow solver

Most production CFD codes are written in Fortran or C++ and treat the solver as a black box. If you want to tune a turbulence model, infer an inflow profile, or optimize a geometry against a flow objective, you either wrap the solver in a finite-difference loop or you hand-derive an adjoint. JAX-Fluids takes a different route: the README states it is a "fully-differentiable CFD solver for 3D, compressible single-phase and two-phase flows," written entirely in JAX, which means the same code path that advances the solution in time can also be differentiated with respect to its inputs. The stated intent is to "facilitate research at the intersection of ML and CFD." That audience is narrow but real: PhD students and research engineers who already work in the JAX ecosystem, who need compressible flow physics, and who want to backpropagate a loss through a simulation rather than treat it as a numerical oracle. The two cited framework papers (Bezgin et al., Computer Physics Communications 2023 and 2025) are the primary references for numerical details, and the README points readers there rather than reproducing the schemes in prose.

Finite volumes on a Cartesian grid, with a wide menu of reconstruction and Riemann options

The solver uses the finite-volume method on a Cartesian grid and explicit time stepping. The README lists Euler, RK2 and RK3 for temporal integration. Spatial reconstruction is where the options multiply: WENO-3, WENO-5, WENO-7, WENO-CU6, WENO-3NN and TENO are all listed, which gives you a knob for trading stencil width against shock-capturing behaviour. Riemann solvers include Lax-Friedrichs, Rusanov, HLL, HLLC and Roe. Two-phase flows are handled by two distinct routes, a level-set method and a diffuse-interface method, and immersed solid boundaries are supported through the level-set approach. There is an implicit turbulence sub-grid scale model called ALDM, plus positivity-preserving techniques, and forcings for temperature, mass flow rate and kinetic energy spectrum. Boundary conditions listed are symmetry, periodic, wall, Dirichlet and Neumann. The architecture is a single JAX program: parallelization uses JAX primitives, so the same code runs on CPU, GPU or TPU without a separate MPI build. The README states the solver was tested on up to 512 NVIDIA A100 GPUs and up to 2048 TPU-v3 cores, and the example images cite roughly 800M cells on TPUv3-64 and roughly 2B cells on TPUv3-256. Those figures come from the project's own documentation, not from independent measurement.

Installation: JAX first, then the package, and the backend choice is yours

The install path is deliberately thin because JAX itself carries the hardware support. For CPU only, the README gives: pip install --upgrade "jax[cpu]", then git clone https://github.com/tumaer/JAXFLUIDS.git, cd JAXFLUIDS, pip install . For editable development, pip install -e . is the documented variant. For GPU, the README recommends installing CUDA and cuDNN through pip wheels rather than a manual CUDA install: pip install --upgrade "jax[cuda12]", followed by the same clone and pip install -e . sequence. The README notes that Mac M1 users should consult a linked JAX issue about jaxlib. There is no conda recipe and no container mentioned in the supplied material, so a reproducible environment is something you assemble yourself. The practical consequence is that your JAX version and your CUDA version must agree before JAX-Fluids is even in the picture; if jax[cuda12] does not import cleanly on your machine, the solver will not either. Five Jupyter notebooks in the notebooks directory are the documented starting point, covering cases such as a 1D Sod shock tube and a 2D air-helium shock-bubble interaction, and the README says they show how to switch the numerical setup, for example by changing the reconstruction order or the resolution.

Where the Cartesian, explicit design becomes the wrong tool

The Cartesian grid is the first constraint to take seriously. Immersed boundaries via level set let you place a solid inside the domain, which is how the README's Mach 2 space shuttle example is built, but you are still resolving that boundary against a uniform grid. A boundary layer on a swept wing, or any case where you would normally cluster cells near a wall, pays a resolution penalty that a body-fitted mesh would not. Explicit time stepping is the second constraint. It is a good fit for compressible flows where acoustic and convective timescales are comparable, which is exactly the regime the solver targets. It is a poor fit for low-Mach or stiff problems where an implicit solver would take far larger steps. Third, nothing in the supplied material claims incompressible or variable-density incompressible capability, so a channel-flow or bluff-body wake study at low Mach number is out of scope. Finally, differentiability has a memory cost: reverse-mode gradients through a long unsteady simulation require storing intermediate states, and the README does not discuss checkpointing or memory budgets. If your optimization horizon is thousands of time steps, that is an open question you should answer with a small experiment before building a workflow around it.

How it differs from SU2 and from differentiable Python CFD stacks

SU2 is the obvious comparison for compressible CFD, and the difference is architectural rather than numerical. SU2 is a C++ solver with an adjoint capability built for shape optimization; it supports unstructured, body-fitted meshes and offers implicit time integration. Its gradient path is a hand-derived adjoint, and its interface to Python is a wrapper around a compiled binary. JAX-Fluids inverts that: the solver is Python all the way down, the gradient path is automatic differentiation through the same code, and the mesh is a Cartesian grid. If your geometry is complex and your objective is aerodynamic shape optimization with an established adjoint workflow, SU2 is the more direct fit. If your objective is to train or calibrate a model that sits inside the solver, or to differentiate through a two-phase simulation where an adjoint would be painful to derive, JAX-Fluids is the more natural platform. The README's own publication list illustrates the second use case: one cited paper optimizes high-order Godunov-type finite-volume schemes end to end for compressible homogeneous isotropic turbulence, and another performs data-driven shape inference in three-dimensional steady-state supersonic flows.

Licence status and what the repository metadata does not settle

The repository metadata reports the licence as NOASSERTION, which means the automated classifier could not map the licence file to a known identifier. The supplied material does not include the licence text, so this review cannot tell you what terms apply. That matters more than usual here: JAX-Fluids is a research code from a university group, and research licences sometimes restrict commercial use or require specific citation. Before you build anything on top of it, read the LICENSE file in the repository and, if your use is commercial, get your own advice. The README does ask that users cite the relevant framework publications, and it lists both the 2023 and 2025 Computer Physics Communications articles with BibTeX entries, so citation is at minimum an expectation of the project.

Maintenance cost and upgrade exposure

The release history in the supplied material is short: v0.1.0 in May 2024, then v0.2.0 and v0.2.1 both in March 2025, with the last push to the default branch in August 2026. Three releases in roughly two years, and the version numbers never leave the 0.x range, tells you this is pre-1.0 software where the API can move. The heavier maintenance burden is not the package itself but the JAX dependency underneath it. A JAX upgrade can change how arrays are traced, how compilation caches behave, or which CUDA version is required, and any of those can break a solver that relies on JAX primitives for parallelization. Budget for pinning your JAX version in the environment and for re-validating a case after any bump. The upside of the design is that there is no separate MPI or CUDA build to maintain inside the project; the same install works across CPU, GPU and TPU as long as the JAX backend is correct. For a research group that already tracks JAX, that is a smaller ongoing cost than maintaining a bespoke Fortran build chain. For a team that does not, adopting JAX-Fluids means adopting the JAX release cadence as well.

Editorial conclusion

Adopt JAX-Fluids if your problem is compressible, fits a Cartesian grid, and you need gradients through the solver, for example to optimize a closure model or infer a shape. Do not adopt it if you need body-fitted meshes around complex geometry, implicit time stepping for stiff low-Mach flows, or incompressible solvers, since the README lists none of these. Before committing, verify the licence terms yourself, because the repository metadata reports NOASSERTION, and check that the JAX version you install matches the CUDA or TPU backend you intend to run on.

Official sources

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