JAX-bandflux: differentiable supernovae SALT modelling for cosmological analysis on GPUs

Abstract

JAX-bandflux is a JAX implementation of critical supernova modelling functionality for cosmological analysis. The codebase implements key components of the established library SNCosmo in a differentiable framework, offering efficient parallelisation and gradient-based optimisation capabilities through GPU acceleration. The package facilitates differentiable computation of supernova light curve measurements, supporting the inference of SALT parameters necessary for cosmological analysis.