[juanitorduz]

Problem

In float32 the HSGP spectral densities under/overflow well before m ~ 100, and jnp.sqrt of an underflowed density has an infinite derivative, so the gradient w.r.t. the kernel hyperparameters is inf * 0 = nan and initialize_model fails with "Cannot find valid initial parameters" (centered and non-centered alike). (related to #2143)

Kernel	Defect	float32 symptom
Squared exponential	exp(-w^2 l^2 / 2) underflows to 0	92/100 zeros at length=1, ell=1; NaN gradients
Rational quadratic	log(exp(-z) * kve) -> log(0); jnp.where at w = 0 leaks NaN from the unselected branch	NaN dS/dlength at w = 0 and w >= 100
Matérn	gamma(nu + D/2), (2 nu)^nu overflow	NaN for nu >= 30.5
Periodic	exp(a) * ive(j, a) / exp(a) -> inf / inf	NaN for length <= 0.1

Fix

No API changes. All densities are evaluated in log space and the public spectral_density_* / diag_spectral_density_* return max(exp(log S), finfo.tiny), i.e. they are strictly positive as they are mathematically. This makes the usual pattern jnp.sqrt(diag_spectral_density_*(...)) safe to differentiate (the finite sqrt cotangent is routed into the constant floor instead of multiplying inf * 0), so hsgp_* and existing user code are fixed without changes on the user's side.

• SE / Matérn: log-space evaluation, gammaln instead of gamma.
• RQ: log_bessel_kve(nu, z) - z; the regular branch is evaluated at a safe argument near w = 0 so the jnp.where cannot leak NaN gradients.
• Periodic: bessel_ive(j, a) directly, no exp(a) round trip.
• TFP arguments cast to a common float dtype (TFP rejects mixed float32/float64 and integer orders).
• Removed the unused, undocumented helpers modified_bessel_first_kind / modified_bessel_second_kind.

Values agree with master to float32 noise wherever master did not under/overflow.

Tests

• diag_spectral_density_* at m=100 match float64 NumPy/SciPy references and the tiny floor (SE, RQ); Matérn correct at nu in {30.5, 60.5}; periodic correct at length in {0.1, 0.05}.
• check_grads of sum(sqrt(diag_spectral_density_*)) w.r.t. length at m=100 for all three kernels; RQ gradient at w = 0 exact.
• End-to-end initialize_model + grad(potential_fn) finite for hsgp_{squared_exponential,matern,rational_quadratic} with m=100.

pytest test/contrib/hsgp: 141 passed (also with JAX_ENABLE_X64=1). NUTS on hsgp_squared_exponential with m=100: 0 divergences. make lint and pre-commit pass.

Notes

• Not addressed here: contrib.hsgp: hsgp_periodic_non_centered scales beta by the variance coefficient instead of its square root #2366 (hsgp_periodic_non_centered uses q2 instead of sqrt(q2)).
• Pre-existing, now documented: gradients w.r.t. scale_mixture through the TFP Bessel function are silently dropped.
• Intermediate commits introduced and then removed a log_diag_spectral_density_* API; happy to squash.

[github-actions]

Benchmark report

this PR fix/hsgp-log-space-spectral-densities at 5b97a054 vs baseline master at 02556b0e

+ run time:     3 faster
  compile time: unchanged across 32 benchmarks

Significant changes (3)

                            ─────── run time ───────     ────── compile time ─────
  benchmark                 baseline  this PR      Δ     baseline   this PR      Δ
──────────────────────────────────────────────────────────────────────────────────
+ biject_to_constraints       2.3 ms   2.1 ms  -6.1%     188.8 ms  195.3 ms  +3.4%
+ categorical_log_prob        1.3 ms   1.3 ms  -6.0%      53.7 ms   48.6 ms  -9.6%
+ svi_multi_particle_elbo     1.10 s   1.02 s  -7.1%       2.83 s    2.73 s  -3.5%

_{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.3 ms   2.1 ms     -6.1%     188.8 ms  195.3 ms  +3.4%
+ categorical_log_prob             1.3 ms   1.3 ms     -6.0%      53.7 ms   48.6 ms  -9.6%
  dirichlet_log_prob               389 µs   399 µs     +2.6%     361.8 ms  364.6 ms  +0.8%
  dirichlet_sample                42.5 ms  42.2 ms     -0.8%     542.0 ms  553.9 ms  +2.2%
  gamma_log_prob                   1.4 ms   1.4 ms     +0.6%       1.99 s    1.94 s  -2.4%
  gamma_sample                    19.4 ms  19.3 ms     -0.5%     512.9 ms  504.6 ms  -1.6%
  lkj_cholesky_sample              4.2 ms   4.0 ms     -3.1%     784.2 ms  785.2 ms  +0.1%
  mixture_same_family_log_prob     574 µs   567 µs     -1.1%      60.9 ms   59.3 ms  -2.5%
  multivariate_normal_log_prob     401 µs   414 µs     +3.2%     444.1 ms  425.1 ms  -4.3%
  normal_log_prob                  429 µs   420 µs     -2.2%      48.5 ms   49.4 ms  +1.9%
  normal_sample                   19.5 ms  19.2 ms     -1.7%     151.0 ms  141.3 ms  -6.4%
  stick_breaking_transform         3.8 ms   3.7 ms     -3.3%     161.4 ms  154.3 ms  -4.4%
  student_t_log_prob               2.5 ms   2.4 ms     -1.5%      67.1 ms   67.3 ms  +0.3%
  truncated_normal_log_prob        463 µs   542 µs  (+17.0%)      45.0 ms   42.2 ms  -6.3%

handlers

                                  ──────── run time ───────     ────── compile time ─────
  benchmark                       baseline   this PR      Δ     baseline   this PR      Δ
─────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    30.2 ms   29.8 ms  -1.4%       2.96 s    2.96 s  -0.0%
  log_density_hierarchical          2.4 ms    2.4 ms  -0.3%       1.05 s  998.3 ms  -4.7%
  nested_handler_stack              991 µs    998 µs  +0.7%       605 µs    637 µs  +5.3%
  potential_energy_and_grad          20 µs     19 µs  -2.6%      77.2 ms   72.1 ms  -6.6%
  predictive_forward_sampling     475.7 ms  480.0 ms  +0.9%     135.7 ms  126.4 ms  -6.8%
  trace_seeded_model                552 µs    534 µs  -3.3%     377.1 ms  370.9 ms  -1.7%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    519.2 ms  531.8 ms  +2.4%       2.28 s   2.39 s  +5.1%
  nuts_dense_mass_funnel     845.2 ms  843.6 ms  -0.2%       1.97 s   1.88 s  -4.5%
  nuts_eight_schools         821.6 ms  832.7 ms  +1.4%       1.87 s   1.80 s  -3.9%
  nuts_hierarchical_glm        2.21 s    2.24 s  +1.4%       3.51 s   3.64 s  +3.7%
  nuts_logistic_regression   784.6 ms  799.7 ms  +1.9%       2.37 s   2.45 s  +3.2%
  nuts_vectorized_chains       1.80 s    1.82 s  +0.8%       2.12 s   2.11 s  -0.3%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 228.5 ms  230.5 ms  +0.9%       2.39 s   2.42 s  +1.0%
  svi_autodiagonalnormal_hierarchical        742.6 ms  763.6 ms  +2.8%       3.92 s   4.11 s  +4.9%
  svi_automultivariatenormal_eight_schools   539.6 ms  553.3 ms  +2.5%       3.59 s   3.59 s  +0.0%
  svi_autonormal_logistic                    545.6 ms  540.7 ms  -0.9%       2.76 s   2.69 s  -2.3%
+ svi_multi_particle_elbo                      1.10 s    1.02 s  -7.1%       2.83 s   2.73 s  -3.5%
  svi_trace_mean_field_elbo                  886.5 ms  866.2 ms  -2.3%       4.20 s   4.32 s  +3.1%

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/hsgp-log-space-spectral-densities
commit	02556b0e	5b97a054
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.}

[fehiepsi]

Thanks for the fix, @juanitorduz!