Differentiable Navier-Stokes on 2,048 TPU cores

What it does

JAX-Fluids is a 3D compressible flow solver built entirely in JAX. It simulates single-phase and two-phase fluids—think shock-bubble interactions or a space shuttle at Mach 2—using finite-volume methods on a Cartesian grid. Because it is written in JAX, the same code runs on CPU, GPU, or TPU without rewrites, and every operation supports automatic differentiation.

The interesting bit

The project treats numerical methods as differentiable programs. That means you can backpropagate through the entire simulation stack, from WENO reconstruction stencils and Riemann solvers to the time-stepping loop, because the whole solver is standard JAX. It has been scaled to 512 NVIDIA A100 GPUs and 2,048 TPU-v3 cores using JAX’s native parallelism.

Key highlights

• Supports high-order adaptive spatial reconstruction (WENO-3/5/7, WENO-CU6, TENO) and multiple Riemann solvers (HLL, HLLC, Roe, etc.)
• Handles two-phase flows via both level-set and diffuse-interface methods, plus immersed solid boundaries
• Includes an implicit turbulence sub-grid scale model (ALDM) and positivity-preserving techniques
• Tested at massive scale: simulations up to approximately 2 billion cells on TPUv3-256
• Distributed entirely under the MIT License

Verdict

Researchers working at the intersection of deep learning and CFD should look here; if you just need a quick black-box flow simulation without touching gradients, it is probably overkill.