fix(distributions): keep Cauchy and StudentT log_prob finite for very large values

[yupengtang]

Cauchy.log_prob and StudentT.log_prob square the standardized value inside log1p. In float32 the square overflows once |z| passes about 1.8e19, so the log density becomes -inf and its gradient 0, although the true values are small finite numbers:

import numpyro.distributions as dist
import scipy.stats

print(dist.Cauchy(0.5, 2.0).log_prob(1e20), scipy.stats.cauchy(0.5, 2.0).logpdf(1e20))
print(dist.StudentT(3.0, 0.5, 2.0).log_prob(1e20), scipy.stats.t(3.0, 0.5, 2.0).logpdf(1e20))

-inf -92.55498642505128
-inf -180.93103017013112

HalfCauchy and the truncated Cauchy variants go through Cauchy.log_prob, so they have the same problem. With x64 it starts around 1e154.

This writes log(1 + z**2) as 2 * log(hypot(1, z)), which is the same quantity without the intermediate square. After the change both lines above print the scipy values, and the gradient at 1e20 is -4e-20 for the StudentT example instead of 0. For ordinary values the result differs from the old expression only at rounding level (max 9.5e-7 over 200k Cauchy draws in float32).

pytest test/test_distributions.py -k "Cauchy or StudentT" gives 872 passed in float32. With JAX_ENABLE_X64=1 it gives 870 passed and 2 failed, test_left_truncated_cdf[0.0-StudentT-base_params5] and test_right_truncated_cdf[0.0-StudentT-base_params5], and the same two fail on master without this change. ruff and ty are clean. I left out a regression test, following the reviews on #2331 and #2332, but can add one if wanted.

#2328 fixed the Cauchy cdf / icdf tails, not log_prob. I didn't find another issue or PR for this (searched "Cauchy log_prob", "StudentT log_prob overflow", "HalfCauchy log_prob").

AI assistance: I used Claude Code to find this, with the same scipy comparison sweep as #2332, and to draft the change. I checked the numbers above and ran the tests myself.

[juanitorduz]

I don't think these AI generated PRs are very useful as they introduce us hard to read at the expense of these unreasonable outliers. If you have not suffer from these overflows in a model, my suggestion is to close this (and all future PRs)

[juanitorduz]

Im also not fond of "AI assistance: I used Claude Code to find this, with the same scipy comparison sweep" as this increases maintenance burden. IMO these AI scans PRs should be blocked.

[github-actions]

Benchmark report

this PR fix-cauchy-studentt-log-prob-tails at e4474b2b vs baseline master at 31dcd985

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

Significant changes (3)

                           ──────── run time ───────     ───── compile time ─────
  benchmark                baseline   this PR      Δ     baseline  this PR      Δ
─────────────────────────────────────────────────────────────────────────────────
+ categorical_log_prob       1.2 ms    1.2 ms  -6.0%      47.6 ms  52.1 ms  +9.5%
- nuts_eight_schools       802.9 ms  847.0 ms  +5.5%       1.72 s   1.85 s  +7.5%
- nuts_vectorized_chains     1.75 s    1.86 s  +6.5%       1.96 s   1.93 s  -1.9%

_{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            2.2 ms   2.1 ms    -3.0%     200.5 ms  196.0 ms   -2.3%
+ categorical_log_prob             1.2 ms   1.2 ms    -6.0%      47.6 ms   52.1 ms   +9.5%
  dirichlet_log_prob               394 µs   400 µs    +1.4%     319.0 ms  312.4 ms   -2.1%
  dirichlet_sample                41.9 ms  41.9 ms    +0.2%     536.5 ms  558.8 ms   +4.2%
  gamma_log_prob                   1.3 ms   1.3 ms    +0.1%       1.88 s    1.85 s   -1.7%
  gamma_sample                    19.3 ms  19.5 ms    +1.0%     504.9 ms  517.2 ms   +2.4%
  lkj_cholesky_sample              4.3 ms   4.3 ms    -0.1%     783.4 ms  790.6 ms   +0.9%
  mixture_same_family_log_prob     570 µs   573 µs    +0.4%      59.7 ms   60.1 ms   +0.7%
  multivariate_normal_log_prob     461 µs   421 µs  (-8.5%)     395.0 ms  406.3 ms   +2.8%
  normal_log_prob                  426 µs   431 µs    +1.2%      46.1 ms   46.0 ms   -0.2%
  normal_sample                   18.4 ms  18.5 ms    +0.4%     148.7 ms  152.2 ms   +2.3%
  stick_breaking_transform         3.6 ms   3.6 ms    +0.4%     166.9 ms  147.8 ms  -11.4%
  student_t_log_prob               2.4 ms   2.4 ms    -2.0%      63.8 ms   67.2 ms   +5.4%
  truncated_normal_log_prob        472 µs   446 µs  (-5.5%)      40.2 ms   40.5 ms   +0.7%

handlers

                                  ──────── run time ───────     ────── compile time ──────
  benchmark                       baseline   this PR      Δ     baseline   this PR       Δ
──────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    29.3 ms   29.8 ms  +1.8%       2.69 s    2.68 s   -0.1%
  log_density_hierarchical          2.4 ms    2.5 ms  +2.7%     923.5 ms  948.5 ms   +2.7%
  nested_handler_stack              990 µs    1.0 ms  +2.7%       662 µs    562 µs  -15.2%
  potential_energy_and_grad          20 µs     19 µs  -2.6%      76.6 ms   75.3 ms   -1.6%
  predictive_forward_sampling     462.4 ms  459.1 ms  -0.7%     118.6 ms  131.3 ms  +10.7%
  trace_seeded_model                552 µs    573 µs  +3.8%     395.9 ms  403.6 ms   +2.0%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    519.7 ms  523.0 ms  +0.6%       2.19 s   2.21 s  +1.0%
  nuts_dense_mass_funnel     831.3 ms  838.1 ms  +0.8%       1.70 s   1.74 s  +2.5%
- nuts_eight_schools         802.9 ms  847.0 ms  +5.5%       1.72 s   1.85 s  +7.5%
  nuts_hierarchical_glm        2.22 s    2.21 s  -0.7%       3.15 s   3.21 s  +1.8%
  nuts_logistic_regression   780.1 ms  783.1 ms  +0.4%       2.25 s   2.24 s  -0.4%
- nuts_vectorized_chains       1.75 s    1.86 s  +6.5%       1.96 s   1.93 s  -1.9%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 221.0 ms  226.3 ms  +2.4%       2.30 s   2.25 s  -2.2%
  svi_autodiagonalnormal_hierarchical        743.9 ms  751.3 ms  +1.0%       3.48 s   3.55 s  +2.2%
  svi_automultivariatenormal_eight_schools   523.0 ms  546.7 ms  +4.5%       3.19 s   3.29 s  +3.0%
  svi_autonormal_logistic                    527.5 ms  536.1 ms  +1.6%       2.49 s   2.49 s  -0.2%
  svi_multi_particle_elbo                      1.03 s    1.02 s  -0.6%       2.62 s   2.58 s  -1.7%
  svi_trace_mean_field_elbo                  867.4 ms  845.1 ms  -2.6%       3.81 s   3.92 s  +2.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-cauchy-studentt-log-prob-tails
commit	31dcd985	e4474b2b
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.}