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.