Surjectors: surjection layers for density estimation with normalizing flows
Simon Dirmeier · The Journal of Open Source Software · 2024
Normalizing flows (NFs, Papamakarios et al., 2021) are tractable neural density estimators which have in the recent past been applied successfully for, e.g., generative modelling (Kingma & Dhariwal (2018), Ping et al. ( 2020)), Bayesian inference (Rezende & Mohamed (2015), Hoffman et al. (2019)) or simulation-based inference (Papamakarios et al. (2019), Dirmeier, Albert, et al. (2023)).Surjectors is a Python library in particular for surjective, i.e., dimensionality-reducing normalizing flows (SNFs, Klein et al. (2021)).Surjectors is based on the libraries JAX, Haiku and Distrax (Bradbury et al. (2018), Babuschkin et al. (2020)) and is fully compatible with them.By virtue of being entirely written in JAX (Bradbury et al., 2018), Surjectors naturally supports usage on either CPU, GPU or TPU.