闭包常量进入了输入序列
JAX 0.11.1 的私有 `trace_to_jaxpr_dynamic` 会把闭包捕获的数组常量提升进 `jaxpr.invars`，NumPyro 却只为显式关键字参数生成 provenance 集合。两组序列长度因此不一致，只要被追踪函数闭包中含有数组常量，变量与来源信息就会错位或直接触发映射失败。
NumPyro
—Stars建立在 JAX 之上的概率编程框架，覆盖 HMC、NUTS、SVI 等贝叶斯推断算法，并利用 JIT、自动微分和多设备并行提升推断效率。
从私有追踪假设回到公共接口
闭包常量参与计算，却不应占据显式参数的 provenance 位置。
PRIVATE TRACE · JAX 0.11.1
invarsclosed constdynamic x
provenancex
[2] ≠ [1]PUBLIC jax.make_jaxpr
constsclosed constinvarsdynamic x
ClosedJaxpr 对齐私有 API 隐藏了布局假设
闭包常量参与计算，但不属于调用方显式提供的动态输入，也不应继承任何输入 provenance。根因是旧实现依赖了私有 tracing API 的输入布局；要恢复一一对应关系，必须在表示层明确区分 `constvars` 与 `invars`，而不是按位置补齐错误序列。
ClosedJaxpr 显式拆分常量
改用公共 `jax.make_jaxpr` 获取 `ClosedJaxpr`，让闭包常量与显式输入分别进入常量和动态参数通道；同步适配 JAX 0.11.1 的 `call_jaxpr` 绑定方式，并加入捕获数组常量的回归用例，直接验证关键字参数与 provenance 的位置对应。
Provenance 再次准确对齐
闭包常量现在保持空 provenance，显式参数则重新与 `jaxpr.invars` 精确对齐；JAX 0.11.1 下相关 415 项测试通过，同时在 JAX 0.7.0 上保持兼容。实现也不再直接耦合容易变化的私有 partial-evaluation API。
PR #2245 已随 NumPyro 0.22.0 发布
查看正式版本 - JAXPR provenance 对齐修复
- Release Notes 收录 PR #2245
- 0.22.0 已发布