fix(distributions): handle zero-mean negative binomial

[aswanth-07]

Summary

• Handle zero-mean Negative Binomial distributions without scalar division errors or NaN log probabilities.
• Evaluate infinite-rate GammaPoisson boundaries with stable log-probability, variance, CDF, and sampling behavior.
• Preserve the strict-positive mean requirement for HurdleNegativeBinomial2.

Fixes #2193

Tests

• PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .\.venv\Scripts\python.exe -m pytest -q test/test_distributions.py -k 'zero_mean_negative_binomial or negative_binomial_log_prob_gradient or hurdle_negative_binomial_requires_positive_mean or gamma_poisson_mixed_finite_and_infinite_rates or gamma_poisson_log_prob or test_log_prob_gradient and (GammaPoisson or NegativeBinomial)' (27 passed)
• PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .\.venv\Scripts\python.exe -m pytest -q test/test_distributions.py -k 'HurdleNegativeBinomial2' (5 passed)
• py -3.13 -m ruff check numpyro/distributions/conjugate.py test/test_distributions.py
• py -3.13 -m ruff format --check numpyro/distributions/conjugate.py test/test_distributions.py

AI assistance disclosure: AI-assisted development tools were used during investigation and implementation. The reported tests were run against the final diff.

[github-actions]

Benchmark report

this PR fix/negative-binomial-zero-mean at 9236dc47 vs baseline master at 0059d3af

+ run time:     1 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.2 ms   2.0 ms  -6.2%      73.0 ms  73.1 ms  +0.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     +4.0%     375.6 ms  375.6 ms  -0.0%
+ categorical_log_prob             2.2 ms   2.0 ms     -6.2%      73.0 ms   73.1 ms  +0.0%
  dirichlet_log_prob               633 µs   635 µs     +0.2%     451.9 ms  430.6 ms  -4.7%
  dirichlet_sample                45.0 ms  45.2 ms     +0.6%     918.9 ms  846.2 ms  -7.9%
  gamma_log_prob                   2.1 ms   2.2 ms     +3.1%       2.44 s    2.32 s  -4.8%
  gamma_sample                    20.1 ms  20.0 ms     -0.2%     856.3 ms  821.3 ms  -4.1%
  lkj_cholesky_sample              5.5 ms   5.5 ms     +1.3%       1.20 s    1.18 s  -2.1%
  mixture_same_family_log_prob     981 µs   975 µs     -0.5%      82.3 ms   81.2 ms  -1.3%
  multivariate_normal_log_prob     351 µs   269 µs  (-23.2%)     158.5 ms  167.4 ms  +5.6%
  normal_log_prob                  816 µs   742 µs   (-9.0%)      58.7 ms   62.1 ms  +5.8%
  normal_sample                   22.6 ms  22.6 ms     -0.0%     211.3 ms  224.7 ms  +6.4%
  stick_breaking_transform         6.5 ms   6.3 ms     -2.4%     207.5 ms  212.7 ms  +2.5%
  student_t_log_prob               3.1 ms   3.0 ms     -1.2%      81.7 ms   79.3 ms  -2.8%
  truncated_normal_log_prob        900 µs   885 µs     -1.6%      55.1 ms   55.0 ms  -0.1%

handlers

                                  ──────── run time ───────     ─────── compile time ───────
  benchmark                       baseline   this PR      Δ     baseline   this PR         Δ
────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    45.1 ms   45.1 ms  +0.1%       3.64 s    3.64 s     -0.1%
  log_density_hierarchical          3.8 ms    3.8 ms  -1.3%       1.11 s    1.18 s     +5.8%
  nested_handler_stack              1.4 ms    1.4 ms  +0.2%       434 µs    793 µs  (+82.6%)
  potential_energy_and_grad          22 µs     22 µs  +1.2%     103.3 ms  100.8 ms     -2.4%
  predictive_forward_sampling     735.0 ms  756.4 ms  +2.9%     217.5 ms  197.9 ms     -9.0%
  trace_seeded_model                853 µs    893 µs  +4.7%     581.2 ms  600.9 ms     +3.4%

mcmc

                             ──────── run time ───────     ────── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR       Δ
────────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    739.9 ms  737.5 ms  -0.3%       3.11 s   3.02 s   -2.9%
  nuts_dense_mass_funnel       1.15 s    1.16 s  +0.8%       2.47 s   2.37 s   -4.0%
  nuts_eight_schools           1.15 s    1.15 s  +0.1%       2.29 s   2.41 s   +5.6%
  nuts_hierarchical_glm        4.87 s    4.83 s  -1.0%       4.39 s   4.93 s  +12.5%
  nuts_logistic_regression     1.10 s    1.09 s  -0.1%       3.09 s   3.07 s   -0.8%
  nuts_vectorized_chains       2.53 s    2.56 s  +1.2%       2.82 s   2.80 s   -0.8%

svi

                                             ──────── run time ───────     ────── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 315.0 ms  311.3 ms  -1.2%       3.71 s   3.01 s  -18.9%
  svi_autodiagonalnormal_hierarchical          1.07 s    1.06 s  -1.5%       4.71 s   4.62 s   -1.9%
  svi_automultivariatenormal_eight_schools   779.2 ms  779.8 ms  +0.1%       3.74 s   3.65 s   -2.2%
  svi_autonormal_logistic                    780.8 ms  766.7 ms  -1.8%       3.35 s   3.24 s   -3.5%
  svi_multi_particle_elbo                      1.51 s    1.54 s  +1.8%       3.44 s   3.34 s   -3.0%
  svi_trace_mean_field_elbo                    1.35 s    1.37 s  +1.3%       5.02 s   5.21 s   +3.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	fix/negative-binomial-zero-mean
commit	0059d3af	9236dc47
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.}