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 framework, a Python-based differentiable LBM library which harnesses the capabilities of the JAX framework. The architecture of XLB is predicated upon ensuring accessibility, extensibility, and computational performance, enabling scaling effectively across CPU, multi-GPU, and distributed multi-GPU systems. The framework can be readily augmented with novel boundary conditions, collision models, or simulation capabilities. XLB offers the unique advantage of integration with JAX's extensive machine learning echosystem, and the ability to utilize automatic differentiation for tackling 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.

翻译：格子玻尔兹曼方法（LBM）因其算法在计算可扩展性方面的潜力，已成为求解流体动力学问题的重要技术。我们提出了XLB框架，这是一个基于Python的可微分LBM库，充分利用了JAX框架的能力。XLB的架构设计以确保可访问性、可扩展性和计算性能为前提，能够在CPU、多GPU及分布式多GPU系统上高效扩展。该框架可便捷地扩展新的边界条件、碰撞模型或模拟功能。XLB的独特优势在于可与JAX广泛的机器学习生态系统集成，并利用自动微分技术解决基于物理的机器学习、优化及反问题。XLB已成功扩展至处理数十亿单元的模拟，每秒实现十亿级格点更新。XLB采用宽松的Apache-2.0许可证发布，可在GitHub上获取：https://github.com/Autodesk/XLB。