Laplax is a new open-source Python package that performs Laplace approximation using jax. It is designed with a modular, purely functional architecture and minimal external dependencies, providing a flexible, researcher-friendly framework for rapid prototyping and experimentation. Laplax approximation provides a scalable and efficient way to quantify weight space uncertainty in deep neural networks, and enables the application of Bayesian tools such as prediction uncertainty and model selection via Occam's razor. The goal of Laplax is to facilitate Bayesian neural network research, uncertainty quantification for deep learning, and the development of improved Laplace approximation techniques.
