# laplax -- Laplace Approximations with JAX

arXiv: https://arxiv.org/abs/2507.17013 · HTML: https://mlsanity.com/paper/2507.17013/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 32 references that are in the map.
- First posted on arXiv 22 Jul 2025.

## Builds on (most cited first, 22 of 22)

- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks](https://mlsanity.com/paper/2002.10118/) (2020)
- [Online Structured Laplace Approximations For Overcoming Catastrophic Forgetting](https://mlsanity.com/paper/1805.07810/) (2018)
- [Functional Variational Bayesian Neural Networks](https://mlsanity.com/paper/1903.05779/) (2019)
- [Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning](https://mlsanity.com/paper/2104.04975/) (2021)
- [Priors in Bayesian Deep Learning: A Review](https://mlsanity.com/paper/2105.06868/) (2021)
- [Bayesian Deep Learning via Subnetwork Inference](https://mlsanity.com/paper/2010.14689/) (2020)
- [Benchmarking Uncertainty Disentanglement: Specialized Uncertainties for Specialized Tasks](https://mlsanity.com/paper/2402.19460/) (2024)
- [Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI](https://mlsanity.com/paper/2402.00809/) (2024)
- [All You Need is a Good Functional Prior for Bayesian Deep Learning](https://mlsanity.com/paper/2011.12829/) (2020)
- [The LLM Surgeon](https://mlsanity.com/paper/2312.17244/) (2023)
- [Fast Predictive Uncertainty for Classification with Bayesian Deep Networks](https://mlsanity.com/paper/2003.01227/) (2020)
- [Learnable Uncertainty under Laplace Approximations](https://mlsanity.com/paper/2010.02720/) (2020)
- [Mixtures of Laplace Approximations for Improved Post-Hoc Uncertainty in Deep Learning](https://mlsanity.com/paper/2111.03577/) (2021)
- [Learning Layer-wise Equivariances Automatically using Gradients](https://mlsanity.com/paper/2310.06131/) (2023)
- [Reparameterization invariance in approximate Bayesian inference](https://mlsanity.com/paper/2406.03334/) (2024)
- [Mean-Field Approximation to Gaussian-Softmax Integral with Application to Uncertainty Estimation](https://mlsanity.com/paper/2006.07584/) (2020)
- [Linearization Turns Neural Operators into Function-Valued Gaussian Processes](https://mlsanity.com/paper/2406.05072/) (2024)
- [FSP-Laplace: Function-Space Priors for the Laplace Approximation in Bayesian Deep Learning](https://mlsanity.com/paper/2407.13711/) (2024)
- [Promises and Pitfalls of the Linearized Laplace in Bayesian Optimization](https://mlsanity.com/paper/2304.08309/) (2023)
- [Position: Curvature Matrices Should Be Democratized via Linear Operators](https://mlsanity.com/paper/2501.19183/) (2025)
- [Shaving Weights with Occam's Razor: Bayesian Sparsification for Neural Networks Using the Marginal Likelihood](https://mlsanity.com/paper/2402.15978/) (2024)

## Built on by (most cited first, 5 of 5)

- [Richer Bayesian Last Layers with Subsampled NTK Features](https://mlsanity.com/paper/2602.01279/) (2026)
- [Efficient Credal Prediction through Decalibration](https://mlsanity.com/paper/2603.08495/) (2026)
- [A Ground-Truth Framework for Uncertainty Disentanglement with Posterior Risk](https://mlsanity.com/paper/2608.05995/) (2026)
- [VOLTA: The Surprising Ineffectiveness of Auxiliary Losses for Calibrated Deep Learning](https://mlsanity.com/paper/2604.08639/) (2026)
- [Subspace Inference Enables Efficient Active Reward Learning from Preferences](https://mlsanity.com/paper/2609.04066/) (2026)

## Abstract (arXiv)

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 with 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.
