doc(gh-2187): mathematical details for count-data distributions

[Yi-111-a]

Adds math-flavored docstrings to four more distributions on the #2187 checklist, following the template from #2188, #2192, #2199 and the recent #2268–#2270 series.

• MultinomialProbs / MultinomialLogits — class PMF, __init__ parameters, sample, log_prob, the probs/logits conversions, mean, variance, and support.
• BetaBinomial — class PMF plus sample, log_prob, mean, variance, and support.
• DirichletMultinomial — class PMF plus sample, log_prob, mean, variance, and support.

Docstrings only, no logic changes. ruff check and ruff format --check are clean on both files.

[github-actions]

Benchmark report

this PR doc-2187-more at e9bc2e79 vs baseline master at 0771b92c

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

Significant changes (2)

                                ──────── run time ───────     ────── compile time ──────
  benchmark                     baseline   this PR      Δ     baseline   this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────
- predictive_forward_sampling   671.5 ms  663.9 ms  -1.1%     176.4 ms  238.9 ms  +35.4%
- svi_autonormal_logistic       678.4 ms  726.6 ms  +7.1%       3.27 s    3.18 s   -2.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.9 ms   3.9 ms     +1.5%     385.3 ms  385.4 ms   +0.0%
  categorical_log_prob             2.1 ms   2.1 ms     +1.5%      67.0 ms   67.6 ms   +0.9%
  dirichlet_log_prob               638 µs   641 µs     +0.5%     430.8 ms  425.1 ms   -1.3%
  dirichlet_sample                45.3 ms  47.2 ms     +4.1%     823.4 ms  817.5 ms   -0.7%
  gamma_log_prob                   2.2 ms   2.1 ms     -3.8%       2.20 s    2.20 s   -0.1%
  gamma_sample                    21.9 ms  22.2 ms     +1.5%     794.9 ms  790.5 ms   -0.5%
  lkj_cholesky_sample              5.5 ms   5.5 ms     +0.4%       1.15 s    1.15 s   +0.2%
  mixture_same_family_log_prob     958 µs   969 µs     +1.2%      78.9 ms   78.7 ms   -0.2%
  multivariate_normal_log_prob     272 µs   303 µs  (+11.5%)     538.8 ms  542.3 ms   +0.6%
  normal_log_prob                  725 µs   760 µs     +4.8%      50.9 ms   56.2 ms  +10.6%
  normal_sample                   22.1 ms  22.2 ms     +0.5%     203.7 ms  196.6 ms   -3.5%
  stick_breaking_transform         6.2 ms   6.1 ms     -1.5%     214.9 ms  209.7 ms   -2.4%
  student_t_log_prob               3.0 ms   3.0 ms     -0.4%      79.7 ms   78.9 ms   -1.0%
  truncated_normal_log_prob        734 µs   745 µs     +1.5%      52.1 ms   51.8 ms   -0.6%

handlers

                                  ───────── run time ─────────     ────── compile time ──────
  benchmark                       baseline   this PR         Δ     baseline   this PR       Δ
─────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    45.0 ms   45.7 ms     +1.4%       3.51 s    3.46 s   -1.5%
  log_density_hierarchical          4.1 ms    4.1 ms     +0.0%       1.23 s    1.26 s   +2.2%
  nested_handler_stack              1.5 ms    1.5 ms     +0.7%       499 µs    457 µs   -8.5%
  potential_energy_and_grad          28 µs     23 µs  (-16.4%)     107.9 ms  103.5 ms   -4.0%
- predictive_forward_sampling     671.5 ms  663.9 ms     -1.1%     176.4 ms  238.9 ms  +35.4%
  trace_seeded_model                922 µs    901 µs     -2.3%     514.8 ms  515.6 ms   +0.2%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    687.4 ms  689.4 ms  +0.3%       3.01 s   3.07 s  +2.0%
  nuts_dense_mass_funnel       1.11 s    1.11 s  +0.3%       2.35 s   2.35 s  -0.1%
  nuts_eight_schools           1.09 s    1.08 s  -1.0%       2.48 s   2.42 s  -2.3%
  nuts_hierarchical_glm        3.20 s    3.18 s  -0.5%       4.30 s   4.39 s  +2.1%
  nuts_logistic_regression     1.03 s    1.03 s  +0.2%       2.99 s   2.98 s  -0.4%
  nuts_vectorized_chains       2.41 s    2.40 s  -0.1%       2.72 s   2.71 s  -0.2%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 304.4 ms  305.4 ms  +0.3%       2.97 s   2.94 s  -1.2%
  svi_autodiagonalnormal_hierarchical          1.01 s  995.4 ms  -1.3%       4.72 s   4.60 s  -2.6%
  svi_automultivariatenormal_eight_schools   723.5 ms  726.4 ms  +0.4%       4.07 s   4.07 s  +0.1%
- svi_autonormal_logistic                    678.4 ms  726.6 ms  +7.1%       3.27 s   3.18 s  -2.8%
  svi_multi_particle_elbo                      1.45 s    1.45 s  +0.4%       3.37 s   3.41 s  +1.3%
  svi_trace_mean_field_elbo                    1.25 s    1.25 s  +0.0%       5.01 s   5.03 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	doc-2187-more
commit	0771b92c	e9bc2e79
numpyro	0.21.0	0.21.0
jax	0.11.2	0.11.2
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.}