  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. In its initial release, jax-cosmo implements
background evolution, linear and non-linear power spectra (using halofit or the
Eisenstein and Hu transfer function), as well as angular power spectra with the
Limber approximation for galaxy and weak lensing probes, all differentiable
with respect to the cosmological parameters and their other inputs. We
illustrate how autodiff can be a game-changer for common tasks involving Fisher
matrix computations, or full posterior inference with gradient-based
techniques. In particular, we show how Fisher matrices are now fast, exact, no
longer require any fine tuning, and are themselves differentiable. Finally,
using a Dark Energy Survey Year 1 3x2pt analysis as a benchmark, we demonstrate
how jax-cosmo can be combined with Probabilistic Programming Languages to
perform posterior inference with state-of-the-art algorithms including a No
U-Turn Sampler, Automatic Differentiation Variational Inference,and Neural
Transport HMC. We further demonstrate that Normalizing Flows using Neural
Transport are a promising methodology for model validation in the early stages
of analysis.
