Make Gibbs kernels composable

[juanitorduz]

Implements the design in #2243 for #898.

Summary

• New Gibbs([(kernel, sites), ...]) composite kernel plus CustomGibbs and DiscreteGibbs block kernels (numpyro/infer/gibbs.py, gibbs_util.py). Blocks can be HMC/NUTS, DiscreteGibbs, CustomGibbs, MixedHMC or a nested Gibbs.
• Opt-in MCMCKernel.refresh / wrap_model / get_constrain_fn hooks and public HMC.get_potential_fn / get_constrain_fn / refresh / wrap_model, so no kernel reads another kernel's private attributes.
• HMCGibbs and DiscreteHMCGibbs are thin facades over Gibbs (constructors, behavior and extra_fields=["hmc_state.*"] unchanged); HMCECS and MixedHMC are standalone wrappers on the public accessors.
• Fixes: HMCECS no longer fails on warmup() followed by run() (model is wrapped once at construction); unused/reused RNG keys; the dead stay branch of the modified random-walk proposal; init_params is no longer mutated.

Deviations from the design doc (found during implementation)

1. A third hook, MCMCKernel.get_constrain_fn, is needed so the composite can condition siblings on any block's constrained values (the design only had it on HMC).
2. The changed-gated refresh is done at two points per sweep (before a block, on earlier siblings; after the sweep, on later siblings), which keeps it state-free and sound.
3. _prepare_model_for_potential(..., enum=) is explicit and DiscreteGibbs passes enum= to potential_energy, otherwise enumerate="parallel" sites give wrong energies and every discrete block would import funsor.
4. v1 raises for an HMC block with value-dependent supports inside Gibbs (its constrain function would need the siblings' values, which is circular).
5. HMCECS was decoupled before the facades, since it subclassed HMCGibbs.

Notes for reviewers

• The DiscreteHMCGibbs and HMCECS docstring assertions were loosened from 0.1 to 0.2: the refactor changes the RNG streams, and across 8 to 10 seeds the seed-to-seed sd of the asserted means is 0.07 (GMM) and 0.1 to 0.14 (ECS) on master and this branch alike, with no bias.
• Minor break: HMCGibbsState is now a GibbsState subclass (z, block_states, rng_key) with an hmc_state property; positional construction of the old 3-tuple no longer works.
• Tests: new test/infer/test_gibbs.py and test/infer/test_kernel_hooks.py; existing Gibbs/ECS/pickle tests pass.

[github-actions]

Benchmark report

this PR composable-gibbs at a231cc86 vs baseline master at 999d8d1f

- run time:     1 slower, 0 faster
  compile time: unchanged across 32 benchmarks

Significant changes (1)

                         ──────── run time ───────     ───── compile time ─────
  benchmark              baseline  this PR       Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────
- categorical_log_prob     2.0 ms   2.3 ms  +10.5%      71.2 ms  73.8 ms  +3.7%

_{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     +0.1%     357.9 ms  383.2 ms   +7.0%
- categorical_log_prob             2.0 ms   2.3 ms    +10.5%      71.2 ms   73.8 ms   +3.7%
  dirichlet_log_prob               671 µs   634 µs   (-5.5%)     462.4 ms  541.5 ms  +17.1%
  dirichlet_sample                44.5 ms  44.6 ms     +0.1%     837.1 ms  864.5 ms   +3.3%
  gamma_log_prob                   2.2 ms   2.2 ms     +0.6%       2.29 s    2.35 s   +2.8%
  gamma_sample                    19.7 ms  19.5 ms     -1.2%     838.6 ms  821.1 ms   -2.1%
  lkj_cholesky_sample              5.4 ms   5.5 ms     +2.2%       1.17 s    1.21 s   +3.2%
  mixture_same_family_log_prob     2.1 ms   2.1 ms     +1.1%     103.0 ms  106.6 ms   +3.5%
  multivariate_normal_log_prob     264 µs   317 µs  (+20.4%)     158.6 ms  158.3 ms   -0.2%
  normal_log_prob                  674 µs   647 µs     -4.0%      58.5 ms   64.0 ms   +9.5%
  normal_sample                   22.3 ms  22.3 ms     -0.0%     212.0 ms  204.7 ms   -3.4%
  stick_breaking_transform         6.4 ms   6.2 ms     -2.5%     221.5 ms  219.9 ms   -0.7%
  student_t_log_prob               3.0 ms   3.1 ms     +0.9%      78.6 ms   80.5 ms   +2.4%
  truncated_normal_log_prob        723 µs   722 µs     -0.2%      54.4 ms   56.9 ms   +4.7%

handlers

                                  ───────── run time ─────────     ────── compile time ──────
  benchmark                       baseline   this PR         Δ     baseline   this PR       Δ
─────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    44.6 ms   46.1 ms     +3.5%       3.50 s    3.73 s   +6.4%
  log_density_hierarchical          3.8 ms    3.9 ms     +3.4%       1.15 s    1.15 s   -0.2%
  nested_handler_stack              1.4 ms    1.4 ms     +1.5%       653 µs    796 µs  +21.9%
  potential_energy_and_grad          27 µs     22 µs  (-19.0%)      97.4 ms  102.4 ms   +5.1%
  predictive_forward_sampling     727.0 ms  731.8 ms     +0.7%     191.7 ms  204.0 ms   +6.4%
  trace_seeded_model                837 µs    833 µs     -0.5%     551.9 ms  578.0 ms   +4.7%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    743.3 ms  742.6 ms  -0.1%       3.08 s   3.20 s  +3.8%
  nuts_dense_mass_funnel       1.17 s    1.16 s  -0.3%       2.38 s   2.40 s  +0.7%
  nuts_eight_schools           1.15 s    1.15 s  +0.4%       2.29 s   2.35 s  +3.0%
  nuts_hierarchical_glm        4.85 s    4.82 s  -0.6%       4.57 s   4.71 s  +3.0%
  nuts_logistic_regression     1.09 s    1.10 s  +1.1%       3.06 s   3.14 s  +2.6%
  nuts_vectorized_chains       2.55 s    2.59 s  +1.5%       2.81 s   2.78 s  -0.8%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 311.2 ms  319.3 ms  +2.6%       3.05 s   3.10 s  +1.6%
  svi_autodiagonalnormal_hierarchical          1.04 s    1.08 s  +4.0%       4.52 s   4.87 s  +7.8%
  svi_automultivariatenormal_eight_schools   758.5 ms  773.2 ms  +1.9%       3.56 s   3.76 s  +5.5%
  svi_autonormal_logistic                    777.5 ms  783.5 ms  +0.8%       3.19 s   3.31 s  +3.7%
  svi_multi_particle_elbo                      1.52 s    1.53 s  +0.4%       3.46 s   3.25 s  -6.1%
  svi_trace_mean_field_elbo                    1.33 s    1.33 s  -0.3%       4.89 s   4.91 s  +0.4%

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	composable-gibbs
commit	999d8d1f	a231cc86
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.}

[juanitorduz]

Pushed a231cc8 addressing the CI failures plus two structural follow-ups:

• CI fix: the three test_discrete_gibbs_multiple_sites_chain[2] failures were a tracer leak in DiscreteGibbs.init/MixedHMC.init: np.asarray(ravel_pytree(support_sizes)[0]) raises under pmap's staging trace because ravel_pytree stages a jnp.concatenate even on constant inputs (CI's upgraded JAX exposes multiple CPU devices by default, so MCMC.run takes the pmap branch; a single-device machine falls back to the eager loop, which is why it passed locally). The flat support sizes are now built with pure numpy in ravel_pytree leaf order. Added a device-independent regression test (single-chain pmap over init/sample for DiscreteGibbs, MixedHMC, DiscreteHMCGibbs) and added test_gibbs.py -k chain to the forced multi-device CI step. Side effect: the stored dtype changes int32 -> int64 (benign, consumers re-wrap with jnp.asarray).
• Module layout: numpyro/infer/gibbs_util.py is merged into numpyro/infer/gibbs.py (helpers at the top, discrete proposals next to DiscreteGibbs), so the feature lives in one module.
• Type checking: every module this PR touches is now in the ty include list (make lint covers them) and all diagnostics are fixed. Two are worth reviewer attention: numpyro.util.cond's pred annotation widened to bool | jax.Array (it was already called with traced booleans everywhere), and mixed_hmc.py's one deprecated 5-arg lax.cond call now uses numpyro's own cond wrapper like the surrounding code (same argument order and behavior; additionally honors NUMPYRO_DISABLE_CONTROL_FLOW_PRIM).

[juanitorduz]

@fehiepsi @Qazalbash I think this one si ready for review 🤗