tumaer/JAXFLUIDS

Differentiable Fluid Dynamics Package

What it solves

JAX-Fluids is designed to bridge the gap between machine learning and computational fluid dynamics (CFD). It provides a fully-differentiable solver for 3D, compressible single-phase and two-phase flows, allowing researchers to perform end-to-end optimization of numerical models using automatic differentiation.

How it works

Built entirely in JAX, the package solves the Navier-Stokes equations using the finite-volume method on a Cartesian grid. It supports various high-order adaptive spatial reconstructions (such as WENO and TENO), multiple Riemann solvers, and explicit time stepping. For complex flows, it includes a level-set method and diffuse-interface method for two-phase simulations and immersed solid boundaries.

Who it’s for

It is intended for researchers working at the intersection of ML and CFD who need a high-performance, differentiable solver that can scale across CPU, GPU, and TPU hardware, including large-scale HPC clusters.

Highlights

  • Fully Differentiable: Enables automatic differentiation for model optimization.
  • High Performance: Scales efficiently on HPC clusters, tested on up to 512 NVIDIA A100 GPUs and 2048 TPU-v3 cores.
  • Versatile Flow Models: Supports 3D compressible flows, two-phase simulations, and immersed solid boundaries.
  • Hardware Agnostic: Runs on CPU, GPU, and TPU.