  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.
