XLB: A differentiable massively parallel lattice Boltzmann library in Python
Not on the map. Why?
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.
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
19 of 89 references.
References · 19
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2021Machine learning accelerated computational fluid dynamicscited by 471
- 2021Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulationcited by 407
- 2019DiffTaichi: Differentiable Programming for Physical Simulationcited by 258
- 2020Solver-in-the-Loop: Learning from Differentiable Physics to Interact with Iterative PDE-Solverscited by 205
- 2021Equinox: neural networks in JAX via callable PyTrees and filtered transformationscited by 148
- 2020Combining Differentiable PDE Solvers and Graph Neural Networks for Fluid Flow Predictioncited by 138
- 2021Enhancing Computational Fluid Dynamics with Machine Learningcited by 138
- 2017DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamicscited by 79
- 2020MLIR: A Compiler Infrastructure for the End of Moore's Lawcited by 65
- 2020Physics-Guided Machine Learning for Scientific Discovery: An Application in Simulating Lake Temperature Profilescited by 59
- 2022JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data sciencecited by 41
7 more
- 2022Learned Turbulence Modelling with Differentiable Fluid Solvers: Physics-based Loss-functions and Optimisation Horizonscited by 39
- 2022JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flowscited by 36
- 2021Applying Machine Learning to Study Fluid Mechanicscited by 13
- 2017Lat-Net: Compressing Lattice Boltzmann Flow Simulations using Deep Neural Networkscited by 11
- 2022WaveY-Net: Physics-augmented deep learning for high-speed electromagnetic simulation and optimizationcited by 8
- 2020ML-LBM: Machine Learning Aided Flow Simulation in Porous Mediacited by 3
- 2020Cross-platform programming model for many-core lattice Boltzmann simulationscited by 1
Built on by
Citing papers per year: 2024: 3; 2025: 6; 2026: 2.
11 papers.
- 2024JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flowscited by 7
- 2025Adaptive Physics-informed Neural Networks: A Surveycited by 5
- 2024Reduced-order modeling of unsteady fluid flow using neural network ensemblescited by 4
- 2024Differentiable programming across the PDE and Machine Learning barriercited by 4
- 2025Optimal Lattice Boltzmann Closures through Multi-Agent Reinforcement Learningcited by 2
- 2025Physics-informed Reduced Order Modeling of Time-dependent PDEs via Differentiable Solverscited by 2
- 2025mPOLICE: Provable Enforcement of Multi-Region Affine Constraints in Deep Neural Networkscited by 1
- 2025Beyond Blur: A Fluid Perspective on Generative Diffusion Modelscited by 1
- 2025GLUE: Coordinating Pre-Trained Generative Models for System-Level Designcited by 0
- 2026NewPINNs: Physics-Informing Neural Networks Using Conventional Solvers for Partial Differential Equationscited by 0
- 2026Mosaic: A Benchmark Suite for Differentiable Physics Solverscited by 0
- 2026Mosaic: A Benchmark Suite for Differentiable Physics Solvers26 Jun 2026
- 2026NewPINNs: Physics-Informing Neural Networks Using Conventional Solvers for Partial Differential Equations23 Jan 2026
- 2025GLUE: Coordinating Pre-Trained Generative Models for System-Level Design22 Dec 2025
- 2025Beyond Blur: A Fluid Perspective on Generative Diffusion Models20 Jun 2025
- 2025Physics-informed Reduced Order Modeling of Time-dependent PDEs via Differentiable Solvers20 May 2025
- 2025Optimal Lattice Boltzmann Closures through Multi-Agent Reinforcement Learning19 Apr 2025
- 2025Adaptive Physics-informed Neural Networks: A Survey23 Mar 2025
- 2025mPOLICE: Provable Enforcement of Multi-Region Affine Constraints in Deep Neural Networks4 Feb 2025
- 2024Differentiable programming across the PDE and Machine Learning barrier9 Sep 2024
- 2024Reduced-order modeling of unsteady fluid flow using neural network ensembles8 Feb 2024
- 2024JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows7 Feb 2024