numpyro

什么是 NumPyro？

NumPyro 是一个轻量级的概率编程库，为 Pyro 提供 NumPy 后端。我们依赖 [JAX][JAX] 实现自动微分以及 GPU/CPU 的 JIT 编译。NumPyro 正在积极开发中，因此请注意其脆弱性、错误以及随着设计演进 API 可能发生的变化。

NumPyro 旨在做到轻量级，并专注于提供灵活的基础，供用户在其上构建。

• Pyro 原语： NumPyro 程序除了可以包含 Pyro 原语（如 sample 和 param）外，还可以包含常规的 Python 和 NumPy 代码。模型代码应与 Pyro 非常相似，只是 PyTorch 和 Numpy 的 API 之间存在一些细微差异。请参见下面的示例。
• 推断算法： NumPyro 支持多种推断算法，特别关注 MCMC 算法，如 Hamiltonian Monte Carlo，包括 No U-Turn Sampler 的实现。其他 MCMC 算法包括 MixedHMC（可处理离散潜变量）以及 HMCECS（每次迭代仅计算数据子集的似然）。NumPyro 的动机之一是通过 JIT 编译包含多次梯度计算的 verlet 积分器来加速 Hamiltonian Monte Carlo。借助 JAX，我们可以组合 jit 和 grad 将整个积分步骤编译为 XLA 优化内核。我们还通过 JIT 编译 NUTS 中的整个树构建阶段来消除 Python 开销（这可以通过 Iterative NUTS 实现）。此外，还有一个基本的变分推断实现，以及许多灵活的（自动）指南用于自动微分变分推断（ADVI）。变分推断实现支持多种功能，包括支持具有离散潜变量的模型（参见 TraceGraph_ELBO）。