Poster
in
Workshop: CODEML: Championing Open-source DEvelopment in Machine Learning
laplax - Laplace Approximations with JAX
Tobias Weber · Bálint Mucsányi · Lenard Rommel · Thomas Christie · Lars Kasüschke · Marvin Pförtner · Philipp Hennig
The Laplace approximation provides a scalable and efficient means of quantifying weight-space uncertainty in deep neural networks, enabling the application of Bayesian tools such as predictive uncertainty and model selection via Occam's razor. In this work, we introduce laplax
, a new open-source Python package for performing Laplace approximations in jax
. Designed with a modular and purely functional architecture and minimal external dependencies, laplax
offers a flexible and researcher-friendly framework for rapid prototyping and experimentation. Its goal is to facilitate research on Bayesian neural networks, uncertainty quantification for deep learning, and the development of improved Laplace approximation techniques.