AeroJAX: JAX-native CFD, differentiable end-to-end. ~560 FPS at 128x128 on CPU [P]

About this article

I have been building a JAX based CFD framework for differentiable Navier Stokes simulation inside ML loops such as inverse design and learned closures. The goal is to keep the full solver stack differentiable so it can sit inside optimisation and learning pipelines. Design choices: Fully JAX native with no external dependencies CPU first vectorized implementation End to end differentiability through velocity, pressure, and vorticity fields Navier Stokes (projection method) and LBM (D2Q9) supp...