BlackJAX

BlackJAX 是什么？

BlackJAX 是一个适用于 JAX 的采样器库，可在 CPU 和 GPU 上运行。

它不是一个概率编程库。但是，只要概率编程库能提供与 JAX 兼容的（可能未归一化的）对数概率密度函数，BlackJAX 就能很好地与之集成。

谁应该使用 BlackJAX？

BlackJAX 适合以下人群：

• 已有对数概率密度函数，只需要一个采样器；
• 需要比通用采样器更高级的功能；
• 想在 GPU 上进行采样；
• 希望基于稳健的基本模块进行研究；
• 正在构建概率编程语言；
• 想了解采样算法的工作原理。

快速入门

安装

你可以使用 pip 安装 BlackJAX：

pip install blackjax

或通过 conda-forge 安装：

conda install -c conda-forge blackjax

BlackJAX 使用纯 Python 编写，但依赖于 JAX 提供的 XLA。默认情况下，随 BlackJAX 一起安装的 JAX 版本只能在 CPU 上运行代码。如果你想在 GPU/TPU 上使用 BlackJAX，我们建议你按照这些说明安装带有相关硬件加速支持的 JAX。

示例

让我们看一个简单的自包含示例，使用 NUTS 进行采样：

import jax
import jax.numpy as jnp
import jax.scipy.stats as stats
import numpy as np

import blackjax

observed = np.random.normal(10, 20, size=1_000)
def logdensity_fn(x):
    logpdf = stats.norm.logpdf(observed, x["loc"], x["scale"])
    return jnp.sum(logpdf)

# 构建核函数
step_size = 1e-3
inverse_mass_matrix = jnp.array([1., 1.])
nuts = blackjax.nuts(logdensity_fn, step_size, inverse_mass_matrix)

# 初始化状态
initial_position = {"loc": 1., "scale": 2.}
state = nuts.init(initial_position)

# 迭代
rng_key = jax.random.key(0)
step = jax.jit(nuts.step)
for i in range(100):
    nuts_key = jax.random.fold_in(rng_key, i)
    state, _ = step(nuts_key, state)

查看文档以获取更多关于如何使用该库的示例：如何为一个或多个链编写推断循环，如何使用 Stan 预热等。

理念

BlackJAX 是什么？

BlackJAX 弥合了"一行代码"框架和模块化、可定制库之间的差距。

用户可以导入该库，并通过几行代码与稳健、经过良好测试和高性能的采样器进行交互。这些采样器面向概率编程语言开发者，或那些有对数概率密度函数并且只需要一个可用采样器的人。

但 BlackJAX 的真正优势在于其内部结构以及如何利用它们快速实验现有或新的采样方案。这个较低层次暴露了推断算法的构建块：积分器、提议分布、动量生成器等，并使它们易于组合以构建新算法。它为采样算法研究提供了加速的机会，提供了稳健、高性能和可重用的代码。

为什么选择 BlackJAX？

采样算法经常集成到概率编程语言中，与框架的其他部分耦合，使得那些不需要建模语言来构建对数概率密度函数的人难以使用。它们的实现大多是单一的，不可能重用算法的部分来构建自定义核函数。BlackJAX 解决了这两个问题。

它是如何工作的？

BlackJAX 允许构建任意复杂的算法，因为它建立在一个非常通用的模式之上。任何接受状态并返回状态的东西都是一个转换核函数，实现如下：

new_state, info =  kernel(rng_key, state)

核函数是无状态的函数，都遵循相同的 API；状态和与转换相关的信息是分开返回的。因此，它们可以轻松组合和交换。我们通过闭包而不是传递参数来专门化这些核函数。

贡献

请遵循我们的简短指南。

引用 Blackjax

要引用此存储库：

@misc{cabezas2024blackjax,
      title={BlackJAX: Composable {B}ayesian inference in {JAX}},
      author={Alberto Cabezas and Adrien Corenflos and Junpeng Lao and Rémi Louf},
      year={2024},
      eprint={2402.10797},
      archivePrefix={arXiv},
      primaryClass={cs.MS}
}

在上述 bibtex 条目中，名字按字母顺序排列，版本号应为 main 分支上的最后一个标签。

致谢

NUTS 实现的一些细节很大程度上受到了 Numpyro 的启发。