Developer's Guide
Setup dev environment#
For development purposes, e.g. if you would like to make contributions, follow the following steps:
With uv
- Install
uv, e.g.pip install --upgrade uv - Then clone this repository and install the development dependencies:
git clone [email protected]:aleximmer/Laplace.git
uv sync --all-extras
laplace-torchis 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:
- Follow the development setup.
- Use ruff as autoformatter. Please refer to the following makefile and run it via
make ruff. Please note that the order ofruff check --fixandruff formatis important! - Also use ruff as linter. Please manually fix all linting errors/warnings before opening a pull request.
- Fully document your changes in the form of Python docstrings, typehinting, and (if applicable) code/markdown examples in the
./examplessubdirectory. - See
docs/api_reference/*.mdon how to include a newly added class in the docs. - 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
Publishing the laplace-torch package to PyPI#
- Update the package version in
pyproject.toml - Have your PyPI token ready
- Build the wheel:
uv build - Run
uv publish - 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
Structure#
The laplace package consists of two main components:
- The subclasses of
laplace.BaseLaplacethat 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 includelaplace.FullLaplace,laplace.KronLaplace,laplace.DiagLaplace,laplace.FunctionalLaplace, their last-layer variations, the subnetwork variants, andlaplace.LowRankLaplace. The function-space ELLA and VaLLA variants are selected through the samelaplace.Laplacefunction. - The backends in
laplace.curvaturewhich 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.