Skip to content

Developer's Guide

Setup dev environment#

For development purposes, e.g. if you would like to make contributions, follow the following steps:

With uv

  1. Install uv, e.g. pip install --upgrade uv
  2. Then clone this repository and install the development dependencies:
git clone [email protected]:aleximmer/Laplace.git
uv sync --all-extras
  1. laplace-torch is now available in editable mode, e.g. you can run:
uv run python examples/regression_example.py

# Or, equivalently:
uv venv
source .venv/bin/activate
python examples/regression_example.py

With pip

git clone [email protected]:aleximmer/Laplace.git

# Recommended to create a virtualenv before the following step
pip install -e ".[dev]"

# Run as usual, e.g.
python examples/regression_examples.py

Contributing#

Pull requests are very welcome. Please follow these guidelines:

  1. Follow the development setup.
  2. Use ruff as autoformatter. Please refer to the following makefile and run it via make ruff. Please note that the order of ruff check --fix and ruff format is important!
  3. Also use ruff as linter. Please manually fix all linting errors/warnings before opening a pull request.
  4. Fully document your changes in the form of Python docstrings, typehinting, and (if applicable) code/markdown examples in the ./examples subdirectory.
  5. See docs/api_reference/*.md on how to include a newly added class in the docs.
  6. Provide as many test cases as possible. Make sure all test cases pass.

Issues, bug reports, and ideas are also very welcome!

Documentation#

The documentation is available here or can be generated and/or viewed locally:

With uv

# assuming the repository was cloned
uv sync --all-extras
# serve the docs locally
uv run mkdocs serve

With pip

# assuming the repository was cloned
pip install -e ".[dev]"
# serve the docs locally
mkdocs serve

Publishing the laplace-torch package to PyPI#

  1. Update the package version in pyproject.toml
  2. Have your PyPI token ready
  3. Build the wheel: uv build
  4. Run uv publish
  5. Create a new release on Github

Tip

It's a good idea to test on TestPyPI first. To do so, simply run the above but with

uv publish --publish-url https://test.pypi.org/legacy/

Structure#

The laplace package consists of two main components:

  1. The subclasses of laplace.BaseLaplace that implement different sparsity structures: different subsets of weights ('all', 'subnetwork' and 'last_layer') and different structures of the Hessian approximation ('full', 'kron', 'lowrank', 'diag' and 'gp'). The weight-space and subset-of-data GP estimators include laplace.FullLaplace, laplace.KronLaplace, laplace.DiagLaplace, laplace.FunctionalLaplace, their last-layer variations, the subnetwork variants, and laplace.LowRankLaplace. The function-space ELLA and VaLLA variants are selected through the same laplace.Laplace function.
  2. The backends in laplace.curvature which provide access to Hessian approximations of the corresponding sparsity structures, for example, the diagonal GGN.

For hessian_structure="gp", the factory also accepts functional_approximation="sod" | "nystrom" | "variational". The latter two select ELLA and VaLLA, respectively. These and FunctionalLaplace share BaseFunctionalLaplace prediction and sampling methods, while each estimator implements its own fitting and posterior covariance. New function-space methods should implement _glm_predictive_distribution, fit, and method-specific checkpoint state; they must not inherit another method's marginal-likelihood formula by default.

The predictive-moment method returns a mean shaped [batch, outputs] and covariance shaped [batch, outputs, outputs]. With joint=True, flatten the mean to [batch * outputs] and return covariance shaped [batch * outputs, batch * outputs], ordered by input and then output. These moments exclude observation noise. Implement _refresh_posterior if changing hyperparameters invalidates cached posterior factors.

Reusable function-space helpers live in laplace/utils/functional_utils.py and are exported through laplace.utils. See the utilities API reference for sampling, batch handling, target shaping, and device conversion. Helpers prefixed with an underscore support specific estimator internals and are not part of the public API.

Additionally, the package provides utilities for decomposing a neural network into feature extractor and last layer for LLLaplace subclasses (laplace.utils.feature_extractor) and effectively dealing with Kronecker factors (laplace.utils.matrix).

Finally, the package implements several options to select/specify a subnetwork for SubnetLaplace (as subclasses of laplace.utils.subnetmask.SubnetMask). Automatic subnetwork selection strategies include: uniformly at random (laplace.utils.subnetmask.RandomSubnetMask), by largest parameter magnitudes (LargestMagnitudeSubnetMask), and by largest marginal parameter variances (LargestVarianceDiagLaplaceSubnetMask and LargestVarianceSWAGSubnetMask). In addition to that, subnetworks can also be specified manually, by listing the names of either the model parameters (ParamNameSubnetMask) or modules (ModuleNameSubnetMask) to perform Laplace inference over.

Extendability#

To extend the laplace package, new BaseLaplace subclasses can be designed, for example, Laplace with a block-diagonal Hessian structure. One can also implement custom subnetwork selection strategies as new subclasses of SubnetMask.

Alternatively, extending or integrating backends (subclasses of curvature.curvature) allows to provide different Hessian approximations to the Laplace approximations. For example, currently the curvature.CurvlinopsInterface based on Curvlinops and the native torch.func (previously known as functorch), curvature.BackPackInterface based on BackPACK and curvature.AsdlInterface based on ASDL are available.