# JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows

arXiv: https://arxiv.org/abs/2203.13760 · HTML: https://mlsanity.com/paper/2203.13760/

## Key facts

- Cited by 36 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 98 references that are in the map.
- First posted on arXiv 25 Mar 2022.

## Builds on (most cited first, 12 of 12)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Machine Learning for Fluid Mechanics](https://mlsanity.com/paper/1905.11075/) (2019)
- [Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers](https://mlsanity.com/paper/2007.00016/) (2020)
- [DPM: A deep learning PDE augmentation method (with application to large-eddy simulation)](https://mlsanity.com/paper/1911.09145/) (2019)
- [Deep Neural Networks for Data-Driven Turbulence Models](https://mlsanity.com/paper/1806.04482/) (2018)
- [ChemNODE: A Neural Ordinary Differential Equations Approach for Chemical Kinetics Solvers](https://mlsanity.com/paper/2101.04749/) (2020)
- [Enhancement of shock-capturing methods via machine learning](https://mlsanity.com/paper/2002.02521/) (2020)
- [Inferring incompressible two-phase flow fields from the interface motion using physics-informed neural networks](https://mlsanity.com/paper/2101.09833/) (2021)
- [Consistent and symmetry preserving data-driven interface reconstruction for the level-set method](https://mlsanity.com/paper/2104.11578/) (2021)

## Built on by (most cited first, 30 of 36)

- [PDEBENCH: An Extensive Benchmark for Scientific Machine Learning](https://mlsanity.com/paper/2210.07182/) (2022)
- [FluidLab: A Differentiable Environment for Benchmarking Complex Fluid Manipulation](https://mlsanity.com/paper/2303.02346/) (2023)
- [LagrangeBench: A Lagrangian Fluid Mechanics Benchmarking Suite](https://mlsanity.com/paper/2309.16342/) (2023)
- [Differentiable Programming for Earth System Modeling](https://mlsanity.com/paper/2208.13825/) (2022)
- [XLB: A differentiable massively parallel lattice Boltzmann library in Python](https://mlsanity.com/paper/2311.16080/) (2023)
- [JAX-DIPS: Neural bootstrapping of finite discretization methods and application to elliptic problems with discontinuities](https://mlsanity.com/paper/2210.14312/) (2022)
- [JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows](https://mlsanity.com/paper/2402.05193/) (2024)
- [A Comparison of Mesh-Free Differentiable Programming and Data-Driven Strategies for Optimal Control under PDE Constraints](https://mlsanity.com/paper/2310.02286/) (2023)
- [PICT -- A Differentiable, GPU-Accelerated Multi-Block PISO Solver for Simulation-Coupled Learning Tasks in Fluid Dynamics](https://mlsanity.com/paper/2505.16992/) (2025)
- [What You See is Not What You Get: Neural Partial Differential Equations and The Illusion of Learning](https://mlsanity.com/paper/2411.15101/) (2024)
- [Neural-Integrated Meshfree (NIM) Method: A differentiable programming-based hybrid solver for computational mechanics](https://mlsanity.com/paper/2311.12915/) (2023)
- [Neural Physics: Using AI Libraries to Develop Physics-Based Solvers for Incompressible Computational Fluid Dynamics](https://mlsanity.com/paper/2402.17913/) (2024)
- [JAX-SPH: A Differentiable Smoothed Particle Hydrodynamics Framework](https://mlsanity.com/paper/2403.04750/) (2024)
- [Scientific Machine Learning of Flow Resistance Using Universal Shallow Water Equations with Differentiable Programming](https://mlsanity.com/paper/2502.12396/) (2025)
- [Physics-Informed DeepONet Coupled with FEM for Convective Transport in Porous Media with Sharp Gaussian Sources](https://mlsanity.com/paper/2508.19847/) (2025)
- [Deep Learning of Solver-Aware Turbulence Closures from Nudged LES Dynamics](https://mlsanity.com/paper/2604.23874/) (2026)
- [Differentiable Neural-Integrated Meshfree Method for Forward and Inverse Modeling of Finite Strain Hyperelasticity](https://mlsanity.com/paper/2407.11183/) (2024)
- [Differentiable Inverse Modeling with Physics-Constrained Latent Diffusion for Heterogeneous Subsurface Parameter Fields](https://mlsanity.com/paper/2512.22421/) (2025)
- [Systematic LLM Translation of Legacy Scientific Code to Differentiable Frameworks: Application to a Land Surface Model](https://mlsanity.com/paper/2606.07681/) (2026)
- [Neuro-symbolic partial differential equation solver](https://mlsanity.com/paper/2210.14907/) (2022)
- [Differentiable Turbulence: Closure as a partial differential equation constrained optimization](https://mlsanity.com/paper/2307.03683/) (2023)
- [Solving Euler equations with Multiple Discontinuities via Separation-Transfer Physics-Informed Neural Networks](https://mlsanity.com/paper/2505.20361/) (2025)
- [Learning Particle Dynamics Subject to Rigid Body Manipulations Using Graph Neural Networks](https://mlsanity.com/paper/2509.03446/) (2025)
- [ADDA: a Modular Framework for Representing, Simulating and Assimilating Dynamics with End-to-end Differentiability](https://mlsanity.com/paper/2608.23297/) (2026)
- [A differentiable programming framework for spin models](https://mlsanity.com/paper/2304.01772/) (2023)
- [SBMLtoODEjax: Efficient Simulation and Optimization of Biological Network Models in JAX](https://mlsanity.com/paper/2307.08452/) (2023)
- [Rational-WENO: A lightweight, physically-consistent three-point weighted essentially non-oscillatory scheme](https://mlsanity.com/paper/2409.09217/) (2024)
- [Learning second-order TVD flux limiters using differentiable solvers](https://mlsanity.com/paper/2503.09625/) (2025)
- [Neural network-based Godunov corrections for approximate Riemann solvers using bi-fidelity learning](https://mlsanity.com/paper/2503.13248/) (2025)
- [Parameter estimation for land-surface models using Neural Physics](https://mlsanity.com/paper/2505.02979/) (2025)

## Abstract (arXiv)

Physical systems are governed by partial differential equations (PDEs). The Navier-Stokes equations describe fluid flows and are representative of nonlinear physical systems with complex spatio-temporal interactions. Fluid flows are omnipresent in nature and engineering applications, and their accurate simulation is essential for providing insights into these processes. While PDEs are typically solved with numerical methods, the recent success of machine learning (ML) has shown that ML methods can provide novel avenues of finding solutions to PDEs. ML is becoming more and more present in computational fluid dynamics (CFD). However, up to this date, there does not exist a general-purpose ML-CFD package which provides 1) powerful state-of-the-art numerical methods, 2) seamless hybridization of ML with CFD, and 3) automatic differentiation (AD) capabilities. AD in particular is essential to ML-CFD research as it provides gradient information and enables optimization of preexisting and novel CFD models. In this work, we propose JAX-FLUIDS: a comprehensive fully-differentiable CFD Python solver for compressible two-phase flows. JAX-FLUIDS allows the simulation of complex fluid dynamics with phenomena like three-dimensional turbulence, compressibility effects, and two-phase flows. Written entirely in JAX, it is straightforward to include existing ML models into the proposed framework. Furthermore, JAX-FLUIDS enables end-to-end optimization. I.e., ML models can be optimized with gradients that are backpropagated through the entire CFD algorithm, and therefore contain not only information of the underlying PDE but also of the applied numerical methods. We believe that a Python package like JAX-FLUIDS is crucial to facilitate research at the intersection of ML and CFD and may pave the way for an era of differentiable fluid dynamics.