docs: add math for FoldedDistribution density

[cnYui]

Summary

Adds the mathematical description of `FoldedDistribution`'s density to its docstring, as tracked in #2187 (documentation with math for distributions).

• For `Y = |X|`, the density is `p_Y(y) = p_X(y) + p_X(-y)` for `y >= 0`.
• `log_prob` computes this as `logsumexp{log p_X(y), log p_X(-y)}`, which the docstring now states.

Docstring-only change: no code path is modified.

Verification

• `ruff check` and `ruff format --check` on `numpyro/distributions/distribution.py`: pass.
• `git diff --check`: clean.
• Numerical check with `Normal(0, 1)` base: `FoldedDistribution.log_prob(y)` matches `log(2 * phi(y))` at y = 0.3, 1.0, 2.5 to within float32 precision (~1e-7).
• `pytest test/test_distributions.py -k "fold or Fold"`: 52 passed, 52 skipped.

Not run: full test suite, `make doctest`, Sphinx build.

Part of #2187.

🤖 Generated with Claude Code

[github-actions]

Benchmark report

this PR doc/folded-distribution-math at e1d0aa40 vs baseline master at 31dcd985

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

Significant changes (5)

                               ──────── run time ────────     ────── compile time ──────
  benchmark                    baseline   this PR       Δ     baseline   this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────
- gamma_sample                  10.0 ms   10.9 ms   +8.8%     402.3 ms  457.0 ms  +13.6%
+ lkj_cholesky_sample            2.3 ms    2.1 ms   -7.3%     693.4 ms  691.2 ms   -0.3%
+ nuts_dense_mass_funnel       716.4 ms  677.0 ms   -5.5%       2.04 s    1.69 s  -16.9%
+ svi_autodelta_map_logistic   199.0 ms  178.5 ms  -10.3%       2.36 s    2.30 s   -2.8%
+ svi_trace_mean_field_elbo    748.7 ms  691.7 ms   -7.6%       4.00 s    4.09 s   +2.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            1.7 ms   1.7 ms     +2.1%     161.9 ms  156.6 ms   -3.3%
  categorical_log_prob             997 µs   823 µs  (-17.4%)      45.3 ms   45.5 ms   +0.4%
  dirichlet_log_prob               268 µs   259 µs     -3.4%     327.7 ms  294.0 ms  -10.3%
  dirichlet_sample                25.7 ms  24.7 ms     -4.1%     494.5 ms  467.0 ms   -5.6%
  gamma_log_prob                   1.1 ms   1.1 ms     +3.9%       1.81 s    1.81 s   +0.0%
- gamma_sample                    10.0 ms  10.9 ms     +8.8%     402.3 ms  457.0 ms  +13.6%
+ lkj_cholesky_sample              2.3 ms   2.1 ms     -7.3%     693.4 ms  691.2 ms   -0.3%
  mixture_same_family_log_prob     345 µs   377 µs   (+9.3%)      50.7 ms   54.9 ms   +8.2%
  multivariate_normal_log_prob     323 µs   349 µs   (+8.1%)     437.5 ms  434.4 ms   -0.7%
  normal_log_prob                  368 µs   363 µs     -1.5%      44.1 ms   43.4 ms   -1.6%
  normal_sample                   16.2 ms  16.5 ms     +1.5%     121.9 ms  118.3 ms   -3.0%
  stick_breaking_transform         2.9 ms   2.9 ms     -0.1%     136.8 ms  131.6 ms   -3.9%
  student_t_log_prob               1.8 ms   1.8 ms     -2.5%      61.1 ms   59.1 ms   -3.4%
  truncated_normal_log_prob        415 µs   377 µs   (-9.2%)      39.3 ms   36.5 ms   -7.2%

handlers

                                  ───────── run time ─────────     ────── compile time ──────
  benchmark                       baseline   this PR         Δ     baseline   this PR       Δ
─────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    21.7 ms   21.0 ms     -3.5%       2.97 s    2.54 s  -14.7%
  log_density_hierarchical          1.6 ms    1.6 ms     +4.8%       1.05 s  969.1 ms   -7.9%
  nested_handler_stack              625 µs    703 µs  (+12.5%)       715 µs    876 µs  +22.6%
  potential_energy_and_grad          15 µs     15 µs     +1.5%      72.7 ms   74.9 ms   +3.0%
  predictive_forward_sampling     397.8 ms  390.0 ms     -1.9%      94.0 ms   98.4 ms   +4.6%
  trace_seeded_model                332 µs    338 µs     +1.7%     402.5 ms  389.2 ms   -3.3%

mcmc

                             ──────── run time ───────     ────── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR       Δ
────────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    441.8 ms  425.7 ms  -3.7%       2.44 s   2.31 s   -5.3%
+ nuts_dense_mass_funnel     716.4 ms  677.0 ms  -5.5%       2.04 s   1.69 s  -16.9%
  nuts_eight_schools         705.5 ms  713.9 ms  +1.2%       1.94 s   1.91 s   -1.5%
  nuts_hierarchical_glm        1.56 s    1.62 s  +3.6%       3.41 s   3.47 s   +1.7%
  nuts_logistic_regression   667.3 ms  648.5 ms  -2.8%       2.47 s   2.18 s  -11.7%
  nuts_vectorized_chains       1.51 s    1.51 s  -0.2%       1.55 s   1.80 s  +16.1%

svi

                                             ──────── run time ────────     ───── compile time ─────
  benchmark                                  baseline   this PR       Δ     baseline  this PR      Δ
────────────────────────────────────────────────────────────────────────────────────────────────────
+ svi_autodelta_map_logistic                 199.0 ms  178.5 ms  -10.3%       2.36 s   2.30 s  -2.8%
  svi_autodiagonalnormal_hierarchical        628.4 ms  615.8 ms   -2.0%       4.01 s   3.84 s  -4.1%
  svi_automultivariatenormal_eight_schools   471.0 ms  457.3 ms   -2.9%       3.66 s   3.59 s  -1.9%
  svi_autonormal_logistic                    457.4 ms  468.7 ms   +2.5%       2.79 s   2.79 s  +0.0%
  svi_multi_particle_elbo                    891.3 ms  890.3 ms   -0.1%       2.77 s   2.89 s  +4.5%
+ svi_trace_mean_field_elbo                  748.7 ms  691.7 ms   -7.6%       4.00 s   4.09 s  +2.2%

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/folded-distribution-math
commit	31dcd985	e1d0aa40
numpyro	0.22.0	0.22.0
jax	0.11.2	0.11.2
backend	cpu	cpu
python	3.14.8	3.14.8

Runner: Linux-6.17.0-1022-azure-x86_64-with-glibc2.39, 4 CPUs.

_{Produced by this benchmark run.}

[juanitorduz]

Thanks for the PR! The math and the reST both look correct. One small wording suggestion and two optional nits below.

Verified

• Math: |x| has derivative ±1, so $p_Y(y) = p_X(y) + p_X(-y)$. This matches log_prob, which takes logsumexp over base_dist.log_prob(±value) along axis 0.
• Numerics: with a Normal(0.5, 1) base, log_prob matches scipy.stats.foldnorm(0.5).logpdf at y = 0.3 (−0.38458) and y = 2.5 (−2.84005).
• reST: the docstring parses cleanly with docutils (math and math_block nodes, no warnings). The r""" prefix is needed for \qquad / \ge. The class is already registered via autoclass in docs/source/distributions.rst.
• Style: consistent with the already-merged Delta and HalfNormal docstrings.

Suggestion

1. Say that the base distribution must be continuous. The formula (and the code) double-counts an atom at 0 when the base is discrete. With validation off, FoldedDistribution(DiscreteUniform(-2, 2)) gives P(Y=0) = P(Y=1) = P(Y=2) = 0.4, which sums to 1.2. Since this PR documents the math, stating the restriction is cheap:

   :param Distribution base_dist: A univariate continuous distribution to reflect.
   

   Fixing the code for discrete bases is out of scope here.

Optional nits

2. support = constraints.positive is strict, so with validation on, log_prob(0.) returns -inf with an out-of-support warning. HalfNormal already documents x \ge 0 the same way, so I'm fine leaving it as is for consistency.
3. In the PR description, "matches to about 1e-9 (float32 precision)" doesn't quite add up, since float32 precision is around 1e-7. This affects only the description, not the code.

When this merges, FoldedDistribution can be ticked off in #2187.

[cnYui]

Thanks for the thorough check! Pushed e1d0aa4: the base_dist param now says "A univariate continuous distribution to reflect." Docstring-only change; ruff check/ruff format --check pass on the file. I left support as is for consistency with HalfNormal, and will fix the 1e-9/float32 wording in the PR description.