XLB: A differentiable massively parallel lattice Boltzmann library in Python

TL;DR

XLB presents a Python-based, differentiable lattice Boltzmann library built on JAX that scales from single GPUs to distributed multi-GPU systems, enabling large-scale fluid simulations and physics-based machine learning. It combines a modular LBM implementation (BGK, MRT, KBC) with multi-device streaming, boundary handling, and mixed-precision support, while delivering differentiable steps suitable for gradient-based optimization and ML integration. Through benchmarks (Taylor–Green, lid-driven cavity, flow over cylinder, turbulent channel), XLB demonstrates accuracy and strong scalability, and its ML demos show feasibility for deep-learning correctors and differentiable flow control. The work emphasizes accessibility, extensibility, and Apache-2.0 licensing, outlining a roadmap toward adjoint methods, data assimilation, and alternative backends to push toward state-of-the-art performance in physics-informed ML and large-scale CFD.

Abstract

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.

Figures (15)

• Figure 1: D2Q9, D3Q19, and D3Q27 lattices containing $q=9, 19, 27$ vectors respectively.
• Figure 2: Illustration of the array sharding strategy employed by XLB for efficient computational load distribution across multiple devices. The strategy requires that the number of grid points along the x-axis to be divisible by the number of available devices. This requirement is automatically handled by the library based on the user's choice for nx.
• Figure 3: The streaming operation in distributed computing, showcasing how population arrays are communicated between devices in XLB library.
• Figure 4: Visual representation of utilizing a mask array for identifying normal vector and distinguishing known and unknown boundary populations post-streaming.
• Figure 5: Illustrating the variation of $\epsilon_{u}$ and $\epsilon_{\rho}$\ref{['eqn:error_tg']} associated with the 2D Taylor-Green Vortex problem with increasing resolution for both BGK (a, b) and KBC (c, d) collision models. The results are obtained with various 'computation/storage' pairs where for instance f64/f32 represents double precision accuracy for computations with single precision accuracy for storage.