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

JAX-bandflux is a JAX (Bradbury et al., 2018) implementation of critical supernova modelling functionality for cosmological analysis.The codebase implements key components of the established library SNCosmo (Barbary et al., 2016) 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 (Kenworthy et al., 2021;Pierel et al., 2022) parameters necessary for cosmological analysis.When parameters are evaluated in batches on GPU (as is standard in JAX-based samplers), the bandflux kernels deliver order-of-magnitude speedups (≈ 100× per parameter set compared with SNCosmo) while matching fluxes to within 0.001% (Leeney et al., 2025), enabling fast evaluation of large, fused likelihoods.