blackjax-devs/blackjax

解决的问题

BlackJAX 提供了一个模块化的贝叶斯推断采样器库，这些采样器与概率编程语言（PPL）解耦。它解决了采样算法过于庞大或被锁定在特定建模框架中的问题，使用户无需使用完整的 PPL 即可在 CPU、GPU 或 TPU 上执行采样。

工作原理

该库基于一种通用的转移核模式，其中无状态函数接收当前状态并返回新状态以及关于转移的信息。这种一致的 API 使得用户可以轻松地组合、替换和构建复杂的采样算法，这些算法由积分器、提议函数和动量生成器等基本构建块组成。

适用人群

• 正在开发新型采样算法的研究人员和开发者。
• 需要稳健、高性能采样器实现的 PPL 开发者。
• 拥有对数概率密度函数（logpdf）但仅需采样器生成样本的用户。
• 希望利用硬件加速（GPU/TPU）进行采样的用户。

主要亮点

• 硬件加速：通过 JAX 在 CPU、GPU 和 TPU 上运行。
• 可组合设计：采用转移核的模块化架构，便于自定义并创建新算法。
• 与 PPL 无关：可与任何能提供 JAX 兼容对数概率密度函数的 PPL 集成。
• 研究就绪：公开低层构建块，加速新推断方案的开发与测试。