# XLB: A differentiable massively parallel lattice Boltzmann library in Python

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 89 references that are in the map.
- First posted on arXiv 27 Nov 2023.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [DiffTaichi: Differentiable Programming for Physical Simulation](https://mlsanity.com/paper/1910.00935/) (2019)
- [Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solvers](https://mlsanity.com/paper/2007.00016/) (2020)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Combining Differentiable PDE Solvers and Graph Neural Networks for Fluid Flow Prediction](https://mlsanity.com/paper/2007.04439/) (2020)
- [Enhancing Computational Fluid Dynamics with Machine Learning](https://mlsanity.com/paper/2110.02085/) (2021)
- [DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics](https://mlsanity.com/paper/1712.03641/) (2017)
- [MLIR: A Compiler Infrastructure for the End of Moore's Law](https://mlsanity.com/paper/2002.11054/) (2020)
- [Physics-Guided Machine Learning for Scientific Discovery: An Application in Simulating Lake Temperature Profiles](https://mlsanity.com/paper/2001.11086/) (2020)
- [JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science](https://mlsanity.com/paper/2212.00964/) (2022)
- [Learned Turbulence Modelling with Differentiable Fluid Solvers: Physics-based Loss-functions and Optimisation Horizons](https://mlsanity.com/paper/2202.06988/) (2022)
- [JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows](https://mlsanity.com/paper/2203.13760/) (2022)
- [Applying Machine Learning to Study Fluid Mechanics](https://mlsanity.com/paper/2110.02083/) (2021)
- [Lat-Net: Compressing Lattice Boltzmann Flow Simulations using Deep Neural Networks](https://mlsanity.com/paper/1705.09036/) (2017)
- [WaveY-Net: Physics-augmented deep learning for high-speed electromagnetic simulation and optimization](https://mlsanity.com/paper/2203.01248/) (2022)
- [ML-LBM: Machine Learning Aided Flow Simulation in Porous Media](https://mlsanity.com/paper/2004.11675/) (2020)
- [Cross-platform programming model for many-core lattice Boltzmann simulations](https://mlsanity.com/paper/2010.11751/) (2020)

## Built on by (most cited first, 11 of 11)

- [JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows](https://mlsanity.com/paper/2402.05193/) (2024)
- [Adaptive Physics-informed Neural Networks: A Survey](https://mlsanity.com/paper/2503.18181/) (2025)
- [Reduced-order modeling of unsteady fluid flow using neural network ensembles](https://mlsanity.com/paper/2402.05372/) (2024)
- [Differentiable programming across the PDE and Machine Learning barrier](https://mlsanity.com/paper/2409.06085/) (2024)
- [Optimal Lattice Boltzmann Closures through Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2504.14422/) (2025)
- [Physics-informed Reduced Order Modeling of Time-dependent PDEs via Differentiable Solvers](https://mlsanity.com/paper/2505.14595/) (2025)
- [mPOLICE: Provable Enforcement of Multi-Region Affine Constraints in Deep Neural Networks](https://mlsanity.com/paper/2502.02434/) (2025)
- [Beyond Blur: A Fluid Perspective on Generative Diffusion Models](https://mlsanity.com/paper/2506.16827/) (2025)
- [GLUE: Coordinating Pre-Trained Generative Models for System-Level Design](https://mlsanity.com/paper/2512.19469/) (2025)
- [NewPINNs: Physics-Informing Neural Networks Using Conventional Solvers for Partial Differential Equations](https://mlsanity.com/paper/2601.17207/) (2026)
- [Mosaic: A Benchmark Suite for Differentiable Physics Solvers](https://mlsanity.com/paper/2606.27895/) (2026)

## Abstract (arXiv)

The lattice Boltzmann method (LBM) has emerged as a prominent technique for solving fluid dynamics problems due to its algorithmic potential for computational scalability. We introduce XLB library, a Python-based differentiable LBM library based on the JAX platform. The architecture of XLB is predicated upon ensuring accessibility, extensibility, and computational performance, enabling scaling effectively across CPU, TPU, multi-GPU, and distributed multi-GPU or TPU systems. The library can be readily augmented with novel boundary conditions, collision models, or multi-physics simulation capabilities. XLB's differentiability and data structure is compatible with the extensive JAX-based machine learning ecosystem, enabling it to address physics-based machine learning, optimization, and inverse problems. XLB has been successfully scaled to handle simulations with billions of cells, achieving giga-scale lattice updates per second. XLB is released under the permissive Apache-2.0 license and is available on GitHub at https://github.com/Autodesk/XLB.