Differentiable centisecond halo-model predictions in ΛCDM and beyond

Abstract: The cost of a halo-model Cℓyy evaluation has dropped from ~30 s per evaluation a decade ago to ~5 ms today. We present classy_szlite, a pure-JAX cosmology and halo-model code that combines neural-network emulators for cosmological distances and the linear and non-linear matter power spectrum with FFTLog for profile Fourier transforms and a fully JIT-friendly Tinker-class halo-model integrator. The result is a fully differentiable pipeline: gradient-based optimisation reaches the MAP in fewer than ~40 forward-and-gradient evaluations (~0.4 s wall), and NUTS on a real Cℓyy bandpower dataset reaches a publication-grade posterior (R-hat ≤ 1.05, |μ̂ − μgold| < 0.1 σ) in ~10 s wall. We compare random-walk Metropolis and NUTS quantitatively: at matched accuracy, NUTS is ~100× faster wall-for-wall on this 2D problem; the gap widens with parameter-space dimension as predicted by the well-known scaling arguments. The architecture generalises to any halo-model tracer; we outline the recipe for kSZ², CIB, galaxy–lensing, and cluster counts.

Subjects:	astro-ph.CO; astro-ph.IM; cs.LG
Cite as:	PX:2605.00009