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

Understand

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.