# CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators

arXiv: https://arxiv.org/abs/2305.06347 · HTML: https://mlsanity.com/paper/2305.06347/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 0 of its 13 references that are in the map.
- First posted on arXiv 10 May 2023.

## Built on by (most cited first, 5 of 5)

- [Differentiable and accelerated spherical harmonic and Wigner transforms](https://mlsanity.com/paper/2311.14670/) (2023)
- [Automatic Functional Differentiation in JAX](https://mlsanity.com/paper/2311.18727/) (2023)
- [A representation learning approach to probe for dynamical dark energy in matter power spectra](https://mlsanity.com/paper/2310.10717/) (2023)
- [Anchors no more: Using peculiar velocities to constrain H0 and the primordial Universe without calibrators](https://mlsanity.com/paper/2504.10453/) (2025)
- [Symbolic Emulators for Cosmology: Accelerating Cosmological Analyses Without Sacrificing Precision](https://mlsanity.com/paper/2510.18749/) (2025)

## Abstract (arXiv)

We present CosmoPower-JAX, a JAX-based implementation of the CosmoPower framework, which accelerates cosmological inference by building neural emulators of cosmological power spectra. We show how, using the automatic differentiation, batch evaluation and just-in-time compilation features of JAX, and running the inference pipeline on graphics processing units (GPUs), parameter estimation can be accelerated by orders of magnitude with advanced gradient-based sampling techniques. These can be used to efficiently explore high-dimensional parameter spaces, such as those needed for the analysis of next-generation cosmological surveys. We showcase the accuracy and computational efficiency of CosmoPower-JAX on two simulated Stage IV configurations. We first consider a single survey performing a cosmic shear analysis totalling 37 model parameters. We validate the contours derived with CosmoPower-JAX and a Hamiltonian Monte Carlo sampler against those derived with a nested sampler and without emulators, obtaining a speed-up factor of O(103). We then consider a combination of three Stage IV surveys, each performing a joint cosmic shear and galaxy clustering (3x2pt) analysis, for a total of 157 model parameters. Even with such a high-dimensional parameter space, CosmoPower-JAX provides converged posterior contours in 3 days, as opposed to the estimated 6 years required by standard methods. CosmoPower-JAX is fully written in Python, and we make it publicly available to help the cosmological community meet the accuracy requirements set by next-generation surveys.