Fix provenance tracking for closed-over constants on JAX 0.11.1

[tandede]

Fixes #2244.

Changes made

• Replace the private trace_to_jaxpr_dynamic call with the public jax.make_jaxpr tracing API.
• Keep closed-over constants separate from the explicit inputs whose provenance is tracked, including on JAX 0.11.1.
• Adapt the call-primitive regression helper to JAX 0.11.1's call_jaxpr binding protocol.
• Add regression coverage for an array constant captured by a closure.

Why

JAX 0.11.1 hoists constants returned by trace_to_jaxpr_dynamic into jaxpr.invars. NumPyro then paired those invars with provenance values built only from explicit keyword arguments, causing a length mismatch whenever a traced function closed over an array.

jax.make_jaxpr preserves the distinction between constants and explicit inputs through its public API. Constants therefore retain empty provenance while tracked keyword arguments remain aligned with jaxpr.invars. This also removes a direct dependency on JAX's private partial-evaluation API.

Tests

• JAX/JAXLIB 0.11.1, Python 3.13: pytest -q test/ops test/infer/test_inspect.py test/infer/test_svi.py — 415 passed, 4 skipped.
• JAX/JAXLIB 0.7.0, Python 3.11: pytest -q test/ops/test_provenance.py — 11 passed.
• ruff check .
• ruff format . --check
• python scripts/update_headers.py --check
• ty check numpyro/ops/provenance.py test/ops/test_provenance.py

Dependencies

No new dependencies.

[github-actions]

Benchmark report

this PR fix/jax-hoisted-provenance-constants at d34087fd vs baseline master at 26cc211a

  run time:     unchanged across 32 benchmarks
  compile time: unchanged across 32 benchmarks

No significant changes

Every benchmark stayed within ±5% run time and ±25% compile time.

_{Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below.}

Full results

distributions

                                 ───────── run time ────────     ────── compile time ─────
  benchmark                      baseline  this PR         Δ     baseline   this PR      Δ
──────────────────────────────────────────────────────────────────────────────────────────
  biject_to_constraints            4.1 ms   4.1 ms     -1.5%     378.0 ms  357.2 ms  -5.5%
  categorical_log_prob             2.1 ms   2.2 ms     +2.5%      73.2 ms   73.5 ms  +0.4%
  dirichlet_log_prob               612 µs   654 µs   (+6.9%)     490.7 ms  475.4 ms  -3.1%
  dirichlet_sample                46.1 ms  45.5 ms     -1.3%     907.2 ms  910.6 ms  +0.4%
  gamma_log_prob                   2.2 ms   2.1 ms     -4.3%       2.51 s    2.47 s  -1.9%
  gamma_sample                    20.5 ms  20.3 ms     -0.7%     847.9 ms  857.4 ms  +1.1%
  lkj_cholesky_sample              5.5 ms   5.6 ms     +0.8%       1.20 s    1.26 s  +5.0%
  mixture_same_family_log_prob     2.1 ms   2.1 ms     -0.1%     108.2 ms  107.2 ms  -0.9%
  multivariate_normal_log_prob     272 µs   273 µs     +0.3%     173.9 ms  172.1 ms  -1.0%
  normal_log_prob                  675 µs   612 µs   (-9.3%)      59.8 ms   60.9 ms  +1.9%
  normal_sample                   24.8 ms  24.8 ms     +0.2%     221.6 ms  241.4 ms  +9.0%
  stick_breaking_transform         6.3 ms   6.5 ms     +2.7%     215.7 ms  226.1 ms  +4.8%
  student_t_log_prob               3.1 ms   3.1 ms     +0.0%      85.4 ms   82.7 ms  -3.1%
  truncated_normal_log_prob        781 µs   684 µs  (-12.5%)      55.1 ms   57.8 ms  +5.0%

handlers

                                  ──────── run time ───────     ────── compile time ──────
  benchmark                       baseline   this PR      Δ     baseline   this PR       Δ
──────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    45.0 ms   45.2 ms  +0.3%       3.84 s    3.84 s   -0.0%
  log_density_hierarchical          3.8 ms    3.8 ms  -0.8%       1.25 s    1.27 s   +1.9%
  nested_handler_stack              1.4 ms    1.4 ms  +1.0%       579 µs    643 µs  +11.2%
  potential_energy_and_grad          21 µs     22 µs  +3.4%     100.8 ms   99.8 ms   -1.0%
  predictive_forward_sampling     756.3 ms  750.6 ms  -0.8%     188.2 ms  169.4 ms  -10.0%
  trace_seeded_model                836 µs    837 µs  +0.1%     562.2 ms  565.0 ms   +0.5%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    758.3 ms  756.7 ms  -0.2%       3.37 s   3.28 s  -2.6%
  nuts_dense_mass_funnel       1.20 s    1.20 s  -0.4%       2.70 s   2.64 s  -2.1%
  nuts_eight_schools           1.20 s    1.19 s  -0.3%       2.63 s   2.53 s  -3.8%
  nuts_hierarchical_glm        4.95 s    4.93 s  -0.3%       5.01 s   4.83 s  -3.5%
  nuts_logistic_regression     1.13 s    1.14 s  +1.6%       3.39 s   3.47 s  +2.5%
  nuts_vectorized_chains       2.72 s    2.62 s  -3.4%       3.16 s   2.99 s  -5.3%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 323.4 ms  331.6 ms  +2.5%       3.28 s   3.31 s  +1.0%
  svi_autodiagonalnormal_hierarchical          1.10 s    1.07 s  -2.0%       5.12 s   5.07 s  -1.1%
  svi_automultivariatenormal_eight_schools   802.8 ms  805.3 ms  +0.3%       4.06 s   4.03 s  -0.8%
  svi_autonormal_logistic                    799.4 ms  795.6 ms  -0.5%       3.63 s   3.51 s  -3.1%
  svi_multi_particle_elbo                      1.57 s    1.57 s  -0.2%       3.76 s   3.68 s  -2.0%
  svi_trace_mean_field_elbo                    1.42 s    1.39 s  -2.0%       5.55 s   5.37 s  -3.2%

Methodology and environment

Each benchmark is set up untimed, then called once with the JAX caches cleared and several more times warm. Run is the fastest warm call; compile is the first call minus that, i.e. the tracing, lowering and XLA compilation the warm calls did not have to pay for.

Both refs were measured on the same runner over 2 interleaved round(s), taking the best observation per benchmark. A result is called neutral when it moves less than ±5% (run) or ±25% (compile), or when the measurement itself is under 1 ms (run) / 50 ms (compile) — a shared CI runner cannot resolve changes below that. Compile time gets the looser band because it is measured once per round rather than best-of-N, and swings by roughly 20% even between two runs of identical code. A delta shown in parentheses did clear its threshold, but on a measurement below the resolution floor, so it is reported without being called a change.

	baseline	this PR
ref	master	fix/jax-hoisted-provenance-constants
commit	26cc211a	d34087fd
numpyro	0.21.0	0.21.0
jax	0.11.1	0.11.1
backend	cpu	cpu
python	3.14.7	3.14.7

Runner: Linux-6.17.0-1022-azure-x86_64-with-glibc2.39, 4 CPUs.

_{Produced by this benchmark run.}

[fehiepsi]

Thanks @tandede! Would this change work for older jax releases?

[tandede]

Thanks @tandede! Would this change work for older jax releases?

Yes. NumPyro currently declares jax>=0.7.0, so I tested the current PR at that supported lower bound as well as on JAX 0.11.1.

• JAX/JAXLIB 0.7.0, Python 3.11: python -m pytest -q test/ops — 337 passed
• JAX/JAXLIB 0.11.1, Python 3.13: python -m pytest -q test/ops — 337 passed

Both runs include the new closed-over-constant regression. The production change uses the public jax.make_jaxpr API, which is available at the 0.7.0 lower bound, so no version-specific production path is needed.

[Qazalbash]

Awesome! Thank you @tandede