Add cat constraint and CatTransform

[aswanth-07]

Closes #1872.

Summary

• add constraints.cat for applying heterogeneous constraints to consecutive slices, including shape validation, feasible-value construction, equality, and JAX pytree support
• add CatTransform with forward/inverse operations, per-component Jacobians, intermediate propagation, and biject_to registration
• cover mixed interval/positive domains, JIT execution, pytree round-trips, invalid layouts, and transform round-trips

The public API follows the corresponding PyTorch constraints.cat / CatTransform naming and slice semantics. Malformed layouts are rejected explicitly so a constraint cannot silently leave trailing values unchecked.

Tests

• pytest -q test/test_transforms.py — 239 passed
• focused constraint tests — 42 passed
• ruff check and ruff format on changed files — passed
• ty check numpyro/distributions/constraints.py numpyro/distributions/transforms.py — passed
• scripts/update_headers.py --check — passed

No new dependencies.

I used OpenAI Codex to assist with source comparison, implementation, and test execution; I reviewed the final diff and validation results.

[github-actions]

Benchmark report

this PR agent/add-cat-constraint-transform at 956e3073 vs baseline master at 999d8d1f

+ run time:     1 faster
  compile time: unchanged across 32 benchmarks

Significant changes (1)

                             ─────── run time ───────     ───── compile time ─────
  benchmark                  baseline  this PR      Δ     baseline  this PR      Δ
──────────────────────────────────────────────────────────────────────────────────
+ log_density_hierarchical     4.0 ms   3.7 ms  -5.6%       1.27 s   1.28 s  +1.0%

_{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.0 ms   4.1 ms    +2.4%     361.0 ms  361.6 ms   +0.2%
  categorical_log_prob             2.5 ms   2.4 ms    -2.6%      77.1 ms   77.6 ms   +0.7%
  dirichlet_log_prob               650 µs   644 µs    -0.9%     603.1 ms  609.2 ms   +1.0%
  dirichlet_sample                46.4 ms  46.7 ms    +0.5%     919.3 ms  908.3 ms   -1.2%
  gamma_log_prob                   2.2 ms   2.1 ms    -0.6%       2.65 s    2.81 s   +6.0%
  gamma_sample                    19.7 ms  19.4 ms    -1.3%     879.9 ms  840.7 ms   -4.4%
  lkj_cholesky_sample              5.5 ms   5.3 ms    -3.2%       1.27 s    1.27 s   -0.2%
  mixture_same_family_log_prob     2.1 ms   2.1 ms    +1.1%     114.6 ms  106.7 ms   -6.9%
  multivariate_normal_log_prob     266 µs   262 µs    -1.5%     175.0 ms  172.0 ms   -1.7%
  normal_log_prob                  822 µs   757 µs  (-7.9%)      69.5 ms   68.0 ms   -2.2%
  normal_sample                   31.6 ms  30.5 ms    -3.7%     216.8 ms  251.1 ms  +15.9%
  stick_breaking_transform         6.4 ms   6.4 ms    -0.6%     224.0 ms  227.3 ms   +1.5%
  student_t_log_prob               3.1 ms   3.0 ms    -3.0%      91.6 ms   91.0 ms   -0.6%
  truncated_normal_log_prob        880 µs   876 µs    -0.4%      56.4 ms   55.2 ms   -2.0%

handlers

                                  ───────── run time ─────────     ────── compile time ──────
  benchmark                       baseline   this PR         Δ     baseline   this PR       Δ
─────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    49.5 ms   48.8 ms     -1.5%       4.07 s    4.14 s   +1.7%
+ log_density_hierarchical          4.0 ms    3.7 ms     -5.6%       1.27 s    1.28 s   +1.0%
  nested_handler_stack              1.4 ms    1.4 ms     -0.6%       657 µs    716 µs   +9.0%
  potential_energy_and_grad          26 µs     22 µs  (-15.5%)     111.5 ms  105.4 ms   -5.5%
  predictive_forward_sampling     786.2 ms  771.3 ms     -1.9%     185.9 ms  217.6 ms  +17.0%
  trace_seeded_model                917 µs    928 µs     +1.2%     575.7 ms  587.2 ms   +2.0%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    802.6 ms  800.2 ms  -0.3%       3.63 s   3.55 s  -2.1%
  nuts_dense_mass_funnel       1.29 s    1.27 s  -1.5%       2.90 s   2.84 s  -2.0%
  nuts_eight_schools           1.27 s    1.27 s  +0.4%       2.73 s   2.68 s  -1.8%
  nuts_hierarchical_glm        5.04 s    5.05 s  +0.2%       5.21 s   5.10 s  -2.2%
  nuts_logistic_regression     1.18 s    1.20 s  +1.9%       3.60 s   3.62 s  +0.3%
  nuts_vectorized_chains       2.78 s    2.77 s  -0.4%       3.10 s   3.14 s  +1.4%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 344.0 ms  341.3 ms  -0.8%       3.43 s   3.35 s  -2.2%
  svi_autodiagonalnormal_hierarchical          1.14 s    1.11 s  -2.1%       5.38 s   5.26 s  -2.2%
  svi_automultivariatenormal_eight_schools   813.3 ms  819.9 ms  +0.8%       4.32 s   4.31 s  -0.1%
  svi_autonormal_logistic                    858.6 ms  838.1 ms  -2.4%       3.74 s   3.66 s  -2.0%
  svi_multi_particle_elbo                      1.65 s    1.62 s  -1.9%       3.73 s   3.74 s  +0.3%
  svi_trace_mean_field_elbo                    1.43 s    1.44 s  +0.8%       5.68 s   5.70 s  +0.3%

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	agent/add-cat-constraint-transform
commit	999d8d1f	956e3073
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.}