fix(distributions): keep Categorical entropy finite when a category has zero probability

[Dev-next-gen]

While reading the entropy methods in numpyro/distributions/discrete.py I noticed that CategoricalProbs.entropy is the only one in that file that multiplies by a bare logarithm:

probs = jnp.asarray(self.probs)
return -(probs * jnp.log(probs)).sum(axis=-1)

probs is constrained to the simplex, and a simplex point may contain an exact zero. When it does, that term is 0 * -inf = nan and the whole sum becomes NaN, even though a category that never occurs contributes nothing to the entropy under the usual 0 log 0 = 0 convention:

>>> dist.CategoricalProbs(jnp.array([0.25, 0.0, 0.75])).entropy()
Array(nan, dtype=float32)
>>> scipy.stats.entropy([0.25, 0.0, 0.75])
0.5623351446188083

BernoulliProbs.entropy, a few hundred lines above, already handles the same edge case with xlogy, and Poisson.log_prob was moved to xlogy for the zero-rate case in #2176. I applied the same treatment here; the change is one line and leaves every strictly positive probs vector bit-for-bit identical, since xlogy(x, x) == x * log(x) away from zero.

I added the case to the existing categorical entropy test, checked against scipy.stats.entropy because scipy.stats.multinomial cannot be used with a zero-probability category. On my machine (CPU jax 0.11.1) the new test fails on master with ACTUAL: nan / DESIRED: 0.562335 and passes with the fix; pytest test/test_distributions.py -k entropy gives 327 passed, 826 skipped afterwards.

I deliberately left CategoricalLogits.entropy alone: it hits the same 0 * -inf shape only when a logit is -inf, which is outside its declared real_vector constraint, so that felt like a separate decision rather than part of this fix.

Found by a defect-hunting pipeline I build and run (Dev-next-gen), using Claude Code with Anthropic's Claude Opus 5.

[fehiepsi]

Thanks @Dev-next-gen !

[github-actions]

Benchmark report

this PR fix-categorical-entropy-zero-probs at 73599fe1 vs baseline master at 0adc509b

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

Significant changes (4)

                              ──────── run time ───────     ─────── compile time ───────
  benchmark                   baseline  this PR       Δ     baseline   this PR         Δ
────────────────────────────────────────────────────────────────────────────────────────
+ categorical_log_prob          1.8 ms   1.7 ms   -6.5%      67.6 ms   69.3 ms     +2.4%
- stick_breaking_transform      6.3 ms   7.2 ms  +14.3%     210.1 ms  210.0 ms     -0.0%
- truncated_normal_log_prob     1.0 ms   1.1 ms   +7.8%      51.1 ms   52.6 ms     +3.1%
- nested_handler_stack          1.3 ms   1.4 ms   +6.2%       482 µs    604 µs  (+25.2%)

_{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.3 ms   4.1 ms    -4.6%     234.2 ms  225.3 ms   -3.8%
+ categorical_log_prob             1.8 ms   1.7 ms    -6.5%      67.6 ms   69.3 ms   +2.4%
  dirichlet_log_prob               574 µs   576 µs    +0.4%     601.4 ms  615.9 ms   +2.4%
  dirichlet_sample                81.6 ms  81.0 ms    -0.8%     718.9 ms  722.1 ms   +0.4%
  gamma_log_prob                   1.9 ms   1.8 ms    -2.5%       2.88 s    2.89 s   +0.2%
  gamma_sample                    31.4 ms  31.6 ms    +0.8%     660.0 ms  658.2 ms   -0.3%
  lkj_cholesky_sample              5.2 ms   5.2 ms    +0.2%       1.03 s    1.03 s   +0.2%
  mixture_same_family_log_prob     765 µs   765 µs    +0.1%      79.6 ms   81.0 ms   +1.7%
  multivariate_normal_log_prob     740 µs   690 µs  (-6.9%)     505.3 ms  567.8 ms  +12.4%
  normal_log_prob                  936 µs   961 µs    +2.6%      59.9 ms   57.4 ms   -4.1%
  normal_sample                   29.6 ms  30.6 ms    +3.2%     186.3 ms  198.7 ms   +6.7%
- stick_breaking_transform         6.3 ms   7.2 ms   +14.3%     210.1 ms  210.0 ms   -0.0%
  student_t_log_prob               3.0 ms   3.0 ms    -2.1%      86.5 ms   86.9 ms   +0.5%
- truncated_normal_log_prob        1.0 ms   1.1 ms    +7.8%      51.1 ms   52.6 ms   +3.1%

handlers

                                  ───────── run time ────────     ─────── compile time ───────
  benchmark                       baseline   this PR        Δ     baseline   this PR         Δ
──────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    39.0 ms   39.6 ms    +1.3%       3.61 s    3.58 s     -0.8%
  log_density_hierarchical          3.2 ms    3.4 ms    +4.8%       1.23 s    1.18 s     -3.8%
- nested_handler_stack              1.3 ms    1.4 ms    +6.2%       482 µs    604 µs  (+25.2%)
  potential_energy_and_grad          27 µs     25 µs  (-8.8%)      92.7 ms   93.5 ms     +0.9%
  predictive_forward_sampling     649.2 ms  662.8 ms    +2.1%     168.3 ms  165.7 ms     -1.5%
  trace_seeded_model                734 µs    763 µs    +4.0%     517.5 ms  520.2 ms     +0.5%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    726.4 ms  718.7 ms  -1.1%       2.91 s   3.02 s  +3.9%
  nuts_dense_mass_funnel       1.14 s    1.13 s  -0.7%       2.35 s   2.46 s  +4.7%
  nuts_eight_schools           1.15 s    1.16 s  +0.5%       2.47 s   2.43 s  -1.9%
  nuts_hierarchical_glm        4.05 s    4.01 s  -1.0%       4.59 s   4.37 s  -4.8%
  nuts_logistic_regression     1.07 s    1.07 s  -0.4%       3.01 s   3.05 s  +1.6%
  nuts_vectorized_chains       2.45 s    2.45 s  +0.1%       2.69 s   2.80 s  +4.1%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 297.0 ms  294.9 ms  -0.7%       2.88 s   2.95 s  +2.4%
  svi_autodiagonalnormal_hierarchical        983.6 ms  997.6 ms  +1.4%       4.62 s   4.87 s  +5.4%
  svi_automultivariatenormal_eight_schools   730.5 ms  746.8 ms  +2.2%       4.13 s   4.20 s  +1.9%
  svi_autonormal_logistic                    739.0 ms  752.2 ms  +1.8%       3.22 s   3.28 s  +2.0%
  svi_multi_particle_elbo                      1.43 s    1.43 s  +0.2%       3.35 s   3.38 s  +0.9%
  svi_trace_mean_field_elbo                    1.23 s    1.26 s  +2.7%       5.18 s   5.13 s  -0.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-categorical-entropy-zero-probs
commit	0adc509b	73599fe1
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.}