JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library

Understand

We present jax-cosmo, a library for automatically differentiable cosmological theory calculations.

• It uses the JAX library, which has created a new coding ecosystem, especially in probabilistic programming.
• As well as batch acceleration, just-in-time compilation, and automatic optimization of code for different hardware modalities (CPU, GPU, TPU), JAX exposes an automatic differentiation (autodiff) mechanism.
• Thanks to autodiff, jax-cosmo gives access to the derivatives of cosmological likelihoods with respect to any of their parameters, and thus enables a range of powerful Bayesian inference algorithms, otherwise impractical in cosmology, such as Hamiltonian Monte Carlo and Variational Inference.