[ColmTalbot]

With jax, the gradient of jnp.logaddexp(-inf, -inf) is nan. Every log-space mixture combines its components as logaddexp(log1p(-lam) + ln_a, log(lam) + ln_b), so the gradient with respect to the mixing fraction was nan as soon as any sample was outside the support of every component, e.g., mass_1 < mmin, which is common for PE samples. For ln_double_power_law_peak_primary_mass the nan also reached every power-law parameter. This breaks NUTS and other gradient-based samplers using ln_prob; ln_two_component_single and ln_two_component_primary_mass_ratio have had it since #144.

Changes

• utils.logaddexp: with jax, both terms are replaced by zero where both are -inf before calling xp.logaddexp, and its result by -inf after, so the gradient there is zero (numpy and cupy call xp.logaddexp directly). Used by every mixture: _ln_two_component_single, ln_three_component_single, ln_double_power_law_peak_primary_mass and ln_independent_spin_orientation_gaussian_isotropic. The logaddexp(0.0, x) calls in the smoothing window and Madau-Dickinson are already safe.
• BaseSmoothedMassDistribution.ln_p_q: the log of one, rather than of zero, where the mass ratio normalization vanishes (m1 < mmin), which gave nan gradients with respect to beta, mmin and delta_m for every smoothed and spline mass model.
• ln_beta_dist: one in place of scale - xx above the support, which gave nan gradients with respect to beta for spin magnitudes above amax (also for beta_dist, which is its exponential).

Gradients

jax.grad of logsumexp(ln_prob) with respect to every hyperparameter, for every model in test/ln_model_test.py (10 draws, its dataset with samples outside the support):

	before	after
power law + peak(s) mixtures	nan: lam, lam_1 (all parameters for the broken power law + peak)	finite
smoothed and spline mass models	nan: beta, mmin, delta_m (and lam, lam_1 with peaks)	finite
Beta spin magnitudes	nan: beta_chi	finite
all other models	finite	finite

In double precision, the gradients match those computed from the linear density to ~1e-11 (relative) wherever the latter are finite. In single precision, the log-space gradients are finite for every model, while the linear ones overflow for steep power laws.

Values

Unchanged: bit-identical for every model with numpy and with jit-compiled jax in double precision. In single precision under jit, XLA fuses ln_two_component_single differently, which changes values at the level of float32 rounding (at most 4e-6 on the test draws) with the same accuracy against double precision. The guard in ln_beta_dist is on scale - xx only: guarding xx too moves log(xx) from compile time (XLA constant-folds the logs of closed-over data) to run time and changes single precision rounding.

No measurable change in the speed of the jitted log density or its gradient.

Tests

test_ln_prob_gradient checks, for every model in single and double precision, that this gradient is finite and, in double precision, that it matches the gradient computed from the density wherever that is finite. It uses 3 draws per model; its run time is dominated by compiling the two gradients and the first (eager) call per model, not the draws. Without the source changes, 26 of its 60 cases fail. The full suite passes on CPU with jax 0.12 (460 passed, 2 skipped), where the new test takes ~40 s.

Still not finite

• The derivative with respect to a mixing fraction of exactly 0 or 1 in log space (log(0)); the linear density's is finite. This would need a custom derivative.
• The gradients of the linear smoothed models' p_q with respect to beta, mmin and delta_m. Dividing by where(norm > 0, norm, 1) fixes most, but not all, of the test draws.

🤖 Generated with [Cl

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 95.64%. Comparing base (29c1f31) to head (6f76fe2).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #151      +/-   ##
==========================================
+ Coverage   95.61%   95.64%   +0.03%     
==========================================
  Files          14       14              
  Lines        1252     1261       +9     
==========================================
+ Hits         1197     1206       +9     
  Misses         55       55              

Flag	Coverage Δ
python3.11	95.64% <100.00%> (+0.03%)	⬆️
python3.12	95.64% <100.00%> (+0.03%)	⬆️
python3.13	95.64% <100.00%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.