NumPyro 0.22.0 发布：Gibbs 采样改为可组合内核
2026 年 9 月 18 日，基于 JAX 的概率编程库 NumPyro 发布 0.22.0 版。核心改动是把 Gibbs 采样重构为可组合内核，同一模型的不同参数块可分别交给 NUTS、离散 Gibbs 或手写的条件抽样器。新版还加入 hurdle 分布，后验诊断计算提速约 1.6–2.1 倍，并修复了集成采样器的一处严重错误。
发布信息与改动概况
NumPyro 是 Pyro 项目旗下以 JAX 为后端的概率编程库，常用于以 NUTS 等采样器做贝叶斯推断。按 GitHub 发布页记录，0.22.0 由维护者 Juan Orduz 于 2026 年 9 月 18 日 08:51（UTC）发布；他在同日的发布议题中确认新版已上传 PyPI。上一版 0.21.0 发布于 2026 年 5 月 2 日，两版相隔约四个半月。发布说明列出 29 位贡献者，其中 20 位首次参与。合并内容中有相当一部分是为各分布补写数学说明和持续集成维护；对建模者影响较大的是四类改动：可组合 Gibbs 内核，hurdle 等新分布与变换，后验诊断提速，以及一组正确性修复。
可组合 Gibbs 内核
在 0.21 及更早版本中，NumPyro 的混合采样只有几种固定搭配：HMCGibbs 让 HMC/NUTS 处理一部分参数，其余参数交给用户自写的 gibbs_fn；DiscreteHMCGibbs 让离散潜变量走 Gibbs 更新、连续参数走 HMC。想把模型拆成三块以上，或在一个块里再嵌套另一种更新方式，就没有现成接口，对应的需求记录在编号较早的议题 #898 中。
0.22.0 新增 Gibbs([(kernel, sites), ...]) 组合内核（合并请求 #2249）。可以把它想象成一条流水线：每个工位只负责一类参数，一轮扫描中各工位依次在“其余参数固定”的条件下更新自己那一块。工位可以是 HMC/NUTS、新增的 DiscreteGibbs（更新离散变量）、CustomGibbs（按用户写出的精确条件分布抽样）、MixedHMC，也可以是另一个嵌套的 Gibbs。原有的 HMCGibbs 和 DiscreteHMCGibbs 改成建在新内核之上的薄封装，开发者称其构造方式、行为和 extra_fields 用法不变。同一合并请求还修复了 HMCECS 先调用 warmup() 再调用 run() 时报错、随机数密钥被重复使用，以及 init_params 被意外改写等问题。
目前能看到的证据都来自开发过程，没有独立评测。第一，持续集成在 32 项基准任务上未检测到运行时间变化，说明重构没有拖慢 NUTS、SVI 等原有路径。第二，重构改变了随机数流，开发者把文档示例中断言的容差从 0.1 放宽到 0.2，并报告：在 8 到 10 个随机种子下，被断言均值的种子间标准差在高斯混合示例中约为 0.07，在 HMCECS 示例中约为 0.10–0.14，新旧代码水平相当，未见系统偏差。第三，维护者在评审中给出一个三块示例（离散类别 c、按共轭条件分布手写抽样的 y、由 NUTS 处理的 x），并以把 c 解析积分掉后直接用 NUTS 拟合的模型作为参照，但发布材料没有公布两者的数值比较。换言之，现有证据只能说明新内核与旧实现结果一致、速度没有变慢，还不能说明组合 Gibbs 的混合效率优于直接使用 NUTS。
新分布、诊断提速与正确性修复
hurdle 分布（#2195）适合零值很多、但零值来源单一的数据。以每周就诊次数为例：hurdle 模型假定先有一道“是否就诊”的门槛，没跨过就记为 0，跨过之后的次数来自一个去掉 0 的计数分布；零膨胀模型则假定 0 有两个来源，一是根本不会就诊的人，二是会就诊但本周恰好没去的人。NumPyro 的实现仿照库中已有的 ZeroInflated 分布设计，PyMC 此前已提供 HurdlePoisson、HurdleNegativeBinomial、HurdleGamma 和 HurdleLogNormal。同一版本还加入了 TruncatedGamma 分布、HaarTransform 与 DiscreteCosineTransform 两种变换，以及 CatTransform。
后验诊断提速（#2251）针对的是 summary()、effective_sample_size 等采样之后的计算：自协方差改由 FFT 一步得到，并借助 scipy.fft 多线程计算；summary() 对每个参数只排序一次，同时用于中位数和 HPDI。贡献者在 Apple M2 上报告，summary() 在两组测试数据上分别从 0.317 秒降到 0.199 秒、从 3.50 秒降到 2.08 秒（约 1.6–1.7 倍），有效样本量计算从 0.107 秒降到 0.052 秒（约 2.1 倍）；评审者在 Intel i5-13420H 上复测，较大数据的 summary() 从 5.281 秒降到 2.539 秒。float64 结果与旧版在相对误差 1e-9 以内一致。这项改动只缩短 mcmc.print_summary() 这类后处理的时间，采样本身的速度不变。
正确性修复中最值得注意的是 EnsembleSampler（#2189），它是 AIES、ESS 两种集成采样器的共同基类。集成采样器把多条“行者”链分成两半，轮流以另一半为参照提出新位置；旧代码在第二个子步骤里却让后一半链以自己为参照更新。提交者称之为严重错误，并补充了在旧逻辑下会失败的回归测试。发布材料没有量化这一错误对后验估计的影响，此前用 NumPyro 集成采样器得到的结果，稳妥做法是在 0.22.0 下重跑核对。此外，#2203 为接入外部采样器提供了更简洁的接口，并以 BlackJAX 的 MCLMC 采样器写了示例笔记本。
适用边界与升级注意
升级前有几处接口变化需要留意。其一，组合 Gibbs 的首个版本不支持取值范围依赖其他参数的 HMC 块，遇到会直接报错，因为这类块的约束变换需要兄弟块的取值，会形成循环依赖。其二，HMCGibbsState 现在是 GibbsState 的子类，按旧的三元组位置参数构造会失败。其三，Predictive 的 exclude_deterministic 参数改名为 condition_deterministic（#2259），旧名仍可使用但会给出 FutureWarning，两者同时传入会报错，默认行为不变。其四，发布说明列出截断分布的 var 改名为 variance，本文未核实旧名是否保留兼容。
在复现性上，Gibbs 重构改变了随机数流，同一随机种子在 0.21 和 0.22 下的抽样不会逐值相同，论文或报告应写明所用的 NumPyro 版本。在证据性质上，本文引用的数字都来自 GitHub 合并请求中的开发者自测和持续集成基准，没有配套论文或第三方评测。组合 Gibbs 适合哪些模型、何时不如先把离散变量积分掉再直接用 NUTS，仍需使用者在自己的模型上比较有效样本量和 R-hat 等诊断指标。