[WIP] Add GammaGaussianHMM, LinearHMM, GaussianMRF and LinearHMMReparam

[juanitorduz]

Stacked on #2273: this PR targets its branch, so the diff shows only the additions on top of it.
Completes the Pyro HMM family: numpyro.ops.gamma_gaussian and GammaGaussianHMM (Student-t marginal with filter), sampling-only LinearHMM, pairwise-factor GaussianMRF, plus StudentTReparam and LinearHMMReparam to marginalize heavy-tailed noise exactly through GaussianHMM.
Verified against fixed-scale Gaussian identities, unrolled Student-t and MRF oracles, closed-form Gamma posteriors, and a NUTS smoke test with Student-t transition noise.

[github-actions]

Benchmark report

this PR feat/gaussian-hmm-family at 911baab5 vs baseline master at 07501df3

- run time:     2 slower, 1 faster
- compile time: 1 slower, 0 faster

Significant changes (4)

                                  ──────── run time ────────     ────── compile time ──────
  benchmark                       baseline   this PR       Δ     baseline   this PR       Δ
───────────────────────────────────────────────────────────────────────────────────────────
+ gamma_log_prob                    1.5 ms    1.4 ms   -6.4%       2.16 s    2.00 s   -7.5%
- lkj_cholesky_sample               4.2 ms    4.4 ms   +6.1%     845.0 ms  841.5 ms   -0.4%
- initialize_model_hierarchical    28.4 ms   31.3 ms  +10.0%       3.26 s    3.19 s   -2.0%
- predictive_forward_sampling     531.7 ms  541.2 ms   +1.8%     116.9 ms  147.2 ms  +25.8%

_{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            3.0 ms   3.0 ms    -0.7%     208.6 ms  198.9 ms   -4.6%
  categorical_log_prob             1.4 ms   1.4 ms    -1.1%      55.7 ms   52.1 ms   -6.4%
  dirichlet_log_prob               400 µs   393 µs    -1.7%     380.3 ms  351.1 ms   -7.7%
  dirichlet_sample                42.8 ms  42.6 ms    -0.4%     572.1 ms  558.6 ms   -2.3%
+ gamma_log_prob                   1.5 ms   1.4 ms    -6.4%       2.16 s    2.00 s   -7.5%
  gamma_sample                    19.6 ms  19.4 ms    -0.7%     520.6 ms  539.4 ms   +3.6%
- lkj_cholesky_sample              4.2 ms   4.4 ms    +6.1%     845.0 ms  841.5 ms   -0.4%
  mixture_same_family_log_prob     564 µs   565 µs    +0.2%      63.0 ms   62.6 ms   -0.7%
  multivariate_normal_log_prob     477 µs   438 µs  (-8.1%)     520.3 ms  459.4 ms  -11.7%
  normal_log_prob                  435 µs   470 µs  (+8.0%)      49.7 ms   52.9 ms   +6.4%
  normal_sample                   19.8 ms  20.6 ms    +4.2%     158.7 ms  158.6 ms   -0.1%
  stick_breaking_transform         4.6 ms   4.6 ms    -0.1%     159.5 ms  151.8 ms   -4.8%
  student_t_log_prob               2.4 ms   2.5 ms    +3.2%      72.0 ms   70.9 ms   -1.5%
  truncated_normal_log_prob        505 µs   511 µs    +1.3%      43.9 ms   42.7 ms   -2.7%

handlers

                                  ──────── run time ────────     ────── compile time ──────
  benchmark                       baseline   this PR       Δ     baseline   this PR       Δ
───────────────────────────────────────────────────────────────────────────────────────────
- initialize_model_hierarchical    28.4 ms   31.3 ms  +10.0%       3.26 s    3.19 s   -2.0%
  log_density_hierarchical          2.3 ms    2.3 ms   +2.1%       1.15 s    1.09 s   -4.8%
  nested_handler_stack              951 µs    956 µs   +0.6%       652 µs    539 µs  -17.3%
  potential_energy_and_grad          19 µs     19 µs   +1.1%      76.2 ms   83.1 ms   +9.0%
- predictive_forward_sampling     531.7 ms  541.2 ms   +1.8%     116.9 ms  147.2 ms  +25.8%
  trace_seeded_model                520 µs    537 µs   +3.4%     443.1 ms  473.5 ms   +6.9%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    592.8 ms  569.2 ms  -4.0%       2.87 s   2.79 s  -2.8%
  nuts_dense_mass_funnel     935.2 ms  916.0 ms  -2.1%       2.30 s   2.22 s  -3.3%
  nuts_eight_schools         936.6 ms  918.8 ms  -1.9%       2.31 s   2.30 s  -0.3%
  nuts_hierarchical_glm        3.30 s    3.21 s  -2.8%       4.24 s   4.12 s  -2.8%
  nuts_logistic_regression   868.5 ms  865.2 ms  -0.4%       2.85 s   2.82 s  -1.3%
  nuts_vectorized_chains       2.03 s    1.94 s  -4.6%       2.34 s   2.26 s  -3.6%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 237.8 ms  245.0 ms  +3.0%       2.71 s   2.75 s  +1.5%
  svi_autodiagonalnormal_hierarchical        821.7 ms  804.2 ms  -2.1%       4.38 s   4.39 s  +0.2%
  svi_automultivariatenormal_eight_schools   600.0 ms  599.2 ms  -0.1%       3.96 s   3.91 s  -1.4%
  svi_autonormal_logistic                    600.0 ms  601.7 ms  +0.3%       3.07 s   3.04 s  -0.8%
  svi_multi_particle_elbo                      1.17 s    1.14 s  -1.7%       3.09 s   3.07 s  -0.8%
  svi_trace_mean_field_elbo                    1.00 s  991.6 ms  -1.0%       4.74 s   4.65 s  -1.8%

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	feat/gaussian-hmm-family
commit	07501df3	911baab5
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.}