JAX-FVM: A differentiable, entropy-stable finite volume solver on unstructured meshes for compressible flows

Abstract

We present JAX-FVM, an open-source, fully differentiable finite volume method (FVM) for the two-dimensional compressible Euler and Navier-Stokes equations on unstructured triangular meshes. The solver is written entirely in JAX, so that every operation : mesh connectivity, flux evaluation, slope limiting, and time integration is just-in-time compiled, vectorised, and end-to-end differentiable through automatic differentiation (AD), and runs transparently on CPU or GPU. On the numerical side, JAX-FVM is built around an entropy-conservative Tadmor/Ismail-Roe two-point flux supplemented with entropy-variable Rusanov or Roe dissipation, second-order MUSCL reconstruction of primitive variables with least-squares gradients and Venkatakrishnan limiting, and a family of explicit (RK2-4) and matrix-free implicit (Newton, SDIRK2) time integrators whose Jacobian actions are obtained by AD. The combination of an unstructured-mesh compressible FVM with end-to-end differentiability fills a gap left by existing differentiable CFD frameworks, which are almost exclusively restricted to structured grids or spectral discretisations. We describe the governing equations, the discretisation, the software architecture, and a set of standard verification cases. The code is openly available at https://github.com/guigzair/jax_fvm.

Disclosure

“t on the existing flax/optax dependencies, and completion of the VTK mesh-input path. Acknowledgements This work was carried out at INRIA. A deep thank you to Julien, Anass and Mathias for their support and feedback. Declaration of Generative AI and AI-assisted technologies in the writing process During the preparation of this work the authors used the Claude large language model inside the VSCode IDE to assist in writing code. After using this tool, the authors reviewed and e”

PDF page 8

Classification

Code generation, completion, or debugging

Multiplier

2

Verified

2026-08-24T00:34:15.228Z

Structural counts