Repository: aleximmer/Laplace (language: python)

Issue:
Project Documentation

Temperature scaling naming issue
We call it temperature but its actually `1/temperature`


Discussion / hints:

Fix patch:
diff --git a/.travis.yml b/.travis.yml
new file mode 100644
index 00000000..309bcab3
--- /dev/null
+++ b/.travis.yml
@@ -0,0 +1,11 @@
+language: python
+python:
+- '3.8'
+install:
+- pip install -r requirements.txt
+- pip install -r tests/requirements.txt
+- python setup.py install
+script:
+- pytest -vx --cov=laplace/ tests/
+after_success:
+- coveralls
diff --git a/README.md b/README.md
index 3d0f0e94..9d84c2e1 100644
--- a/README.md
+++ b/README.md
@@ -1,1 +1,104 @@
-# Laplace
\ No newline at end of file
+# Laplace
+
+Main: [![Main](https://travis-ci.com/AlexImmer/Laplace.svg?token=rpuRxEjQS6cCZi7ptL9y&branch=main)](https://travis-ci.com/AlexImmer/Laplace)
+Dev: [![Dev](https://travis-ci.com/AlexImmer/Laplace.svg?token=rpuRxEjQS6cCZi7ptL9y&branch=dev)](https://travis-ci.com/AlexImmer/Laplace)
+
+The documentation is available at [https://aleximmer.github.io/Laplace/index.html](https://aleximmer.github.io/Laplace/index.html).
+
+## Setup
+
+We assume `python3.8` since the package was developed with that version.
+
+```bash
+pip install -r requirements.txt
+# for development
+pip install -e .
+# for "production"
+pip install .
+
+
+# run tests
+pip install -r tests/requirements.txt
+pytest tests/
+```
+
+## Structure 
+The laplace package consists of two main components: 
+
+1. The subclasses of [`laplace.BaseLaplace`](laplace/baselaplace.py) that implement different sparsity structures: different subsets of weights (`'all'` and `'last_layer'`) and different structures of the Hessian approximation (`'full'`, `'kron'`, and `'diag'`). This results in six currently available options: `laplace.FullLaplace`, `laplace.KronLaplace`, `laplace.DiagLaplace`, and the corresponding last-layer variations `laplace.FullLLLaplace`, `laplace.KronLLLaplace`,  and `laplace.DiagLLLaplace`, which are all subclasses of [`laplace.LLLaplace`](laplace/lllaplace.py). All of these can be conveniently accessed via the [`laplace.Laplace`](laplace/laplace.py) function.
+2. The backends in [`laplace.curvature`](laplace/curvature/) which provide access to Hessian approximations of
+the corresponding sparsity structures, for example, the diagonal GGN.
+
+Additionally, the package provides utilities for
+decomposing a neural network into feature extractor and last layer for `LLLaplace` subclasses ([`laplace.feature_extractor`](laplace/feature_extractor.py))
+and
+effectively dealing with Kronecker factors ([`laplace.matrix`](laplace/matrix.py)).
+
+## Extendability
+To extend the laplace package, new `BaseLaplace` subclasses can be designed, for example,
+a block-diagonal structure or subset-of-weights Laplace.
+Alternatively, extending or integrating backends (subclasses of [`curvature.curvature`](laplace/curvature/curvature.py)) allows to provide different Hessian
+approximations to the Laplace approximations.
+For example, currently the [`curvature.BackPackInterface`](laplace/curvature/backpack.py) based on [BackPACK](https://github.com/f-dangel/backpack/) and [`curvature.AsdlInterface`](laplace/curvature/asdl.py) based on [ASDL](https://github.com/kazukiosawa/asdfghjkl) are available.
+The `curvature.AsdlInterface` provides a Kronecker factored empirical Fisher while the `curvature.BackPackInterface`
+does not, and only the `curvature.BackPackInterface` provides access to Hessian approximations
+for a regression (MSELoss) loss function.
+
+## Example usage
+
+### *Post-hoc* prior precision tuning of last-layer LA 
+
+In the following example, a pre-trained model is loaded,
+then the Laplace approximation is fit to the training data,
+and the prior precision is optimized with cross-validation `'CV'`.
+After that, the resulting LA is used for prediction with 
+the `'probit'` predictive for classification. 
+
+```python
+from laplace import Laplace
+
+# pre-trained model
+model = load_map_model()  
+
+# User-specified LA flavor
+la = Laplace(model, 'classification',
+             subset_of_weights='all', 
+             hessian_structure='diag')
+la.fit(train_loader)
+la.optimize_prior_precision(method='CV', val_loader=val_loader)
+
+# User-specified predictive approx.
+pred = la(x, link_approx='probit')
+```
+
+### Differentiating the log marginal likelihood w.r.t. hyperparameters
+
+The marginal likelihood can be used for model selection and is differentiable
+for continuous hyperparameters like the prior precision or observation noise.
+Here, we fit the library default, KFAC last-layer LA and differentiate
+the log marginal likelihood.
+
+```python
+from laplace import Laplace
+    
+# Un- or pre-trained model
+model = load_model()  
+    
+# Default to recommended last-layer KFAC LA:
+la = Laplace(model, likelihood='regression')
+la.fit(train_loader)
+    
+# ML w.r.t. prior precision and observation noise
+ml = la.log_marginal_likelihood(prior_prec, obs_noise)
+ml.backward()
+```
+
+## [Documentation](https://aleximmer.github.io/Laplace/index.html)
+
+```bash
+pip install pdoc3 matplotlib
+# create docs and write to html
+bash update_docs.sh
+# .. or serve the docs directly
+pdoc --http 0.0.0.0:8080 laplace --template-dir template
+```
diff --git a/docs/baselaplace.html b/docs/baselaplace.html
new file mode 100644
index 00000000..34f0e53a
--- /dev/null
+++ b/docs/baselaplace.html
@@ -0,0 +1,577 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.baselaplace API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.baselaplace</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.baselaplace.BaseLaplace"><code class="flex name class">
+<span>class <span class="ident">BaseLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Baseclass for all Laplace approximations in this library.
+Subclasses need to specify how the Hessian approximation is initialized,
+how to add up curvature over training data, how to sample from the
+Laplace approximation, and how to compute the functional variance.</p>
+<p>A Laplace approximation is represented by a MAP which is given by the
+<code>model</code> parameter and a posterior precision or covariance specifying
+a Gaussian distribution <span><span class="MathJax_Preview">\mathcal{N}(\theta_{MAP}, P^{-1})</span><script type="math/tex">\mathcal{N}(\theta_{MAP}, P^{-1})</script></span>.
+The goal of this class is to compute the posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>
+which sums as
+<span><span class="MathJax_Preview">
+P = \sum_{n=1}^N \nabla^2_\theta \log p(\mathcal{D}_n \mid \theta)
+\vert_{\theta_{MAP}} + \nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}}.
+</span><script type="math/tex; mode=display">
+P = \sum_{n=1}^N \nabla^2_\theta \log p(\mathcal{D}_n \mid \theta)
+\vert_{\theta_{MAP}} + \nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}}.
+</script></span>
+Every subclass implements different approximations to the log likelihood Hessians,
+for example, a diagonal one. The prior is assumed to be Gaussian and therefore we have
+a simple form for <span><span class="MathJax_Preview">\nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}} = P_0 </span><script type="math/tex">\nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}} = P_0 </script></span>.
+In particular, we assume a scalar, layer-wise, or diagonal prior precision so that in
+all cases <span><span class="MathJax_Preview">P_0 = \textrm{diag}(p_0)</span><script type="math/tex">P_0 = \textrm{diag}(p_0)</script></span> and the structure of <span><span class="MathJax_Preview">p_0</span><script type="math/tex">p_0</script></span> can be varied.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>determines the log likelihood Hessian approximation</dd>
+<dt><strong><code>sigma_noise</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>1</code></dt>
+<dd>observation noise for the regression setting; must be 1 for classification</dd>
+<dt><strong><code>prior_precision</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>1</code></dt>
+<dd>prior precision of a Gaussian prior (= weight decay);
+can be scalar, per-layer, or diagonal in the most general case</dd>
+<dt><strong><code>prior_mean</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>0</code></dt>
+<dd>prior mean of a Gaussian prior, useful for continual learning</dd>
+<dt><strong><code>temperature</code></strong> :&ensp;<code>float</code>, default=<code>1</code></dt>
+<dd>temperature of the likelihood; lower temperature leads to more
+concentrated posterior and vice versa.</dd>
+<dt><strong><code>backend</code></strong> :&ensp;<code>subclasses</code> of <code><a title="laplace.curvature.CurvatureInterface" href="curvature/index.html#laplace.curvature.CurvatureInterface">CurvatureInterface</a></code></dt>
+<dd>backend for access to curvature/Hessian approximations</dd>
+<dt><strong><code>backend_kwargs</code></strong> :&ensp;<code>dict</code>, default=<code>None</code></dt>
+<dd>arguments passed to the backend on initialization, for example to
+set the number of MC samples for stochastic approximations.</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.DiagLaplace" href="#laplace.baselaplace.DiagLaplace">DiagLaplace</a></li>
+<li><a title="laplace.baselaplace.FullLaplace" href="#laplace.baselaplace.FullLaplace">FullLaplace</a></li>
+<li><a title="laplace.baselaplace.KronLaplace" href="#laplace.baselaplace.KronLaplace">KronLaplace</a></li>
+<li>laplace.lllaplace.LLLaplace</li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.baselaplace.BaseLaplace.backend"><code class="name">var <span class="ident">backend</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.log_likelihood"><code class="name">var <span class="ident">log_likelihood</span></code></dt>
+<dd>
+<div class="desc"><p>Compute log likelihood on the training data after <code>.fit()</code> has been called.
+The log likelihood is computed on-demand based on the loss and, for example,
+the observation noise which makes it differentiable in the latter for
+iterative updates.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_likelihood</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.scatter"><code class="name">var <span class="ident">scatter</span></code></dt>
+<dd>
+<div class="desc"><p>Computes the <em>scatter</em>, a term of the log marginal likelihood that
+corresponds to L-2 regularization:
+<code>scatter</code> = <span><span class="MathJax_Preview">(\theta_{MAP} - \mu_0)^{T} P_0 (\theta_{MAP} - \mu_0) </span><script type="math/tex">(\theta_{MAP} - \mu_0)^{T} P_0 (\theta_{MAP} - \mu_0) </script></span>.</p>
+<h2 id="returns">Returns</h2>
+<p>[type]
+[description]</p></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.log_det_prior_precision"><code class="name">var <span class="ident">log_det_prior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Compute log determinant of the prior precision
+<span><span class="MathJax_Preview">\log \det P_0</span><script type="math/tex">\log \det P_0</script></span></p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_det</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.log_det_posterior_precision"><code class="name">var <span class="ident">log_det_posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Compute log determinant of the posterior precision
+<span><span class="MathJax_Preview">\log \det P</span><script type="math/tex">\log \det P</script></span> which depends on the subclasses structure
+used for the Hessian approximation.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_det</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.log_det_ratio"><code class="name">var <span class="ident">log_det_ratio</span></code></dt>
+<dd>
+<div class="desc"><p>Compute the log determinant ratio, a part of the log marginal likelihood.
+<span><span class="MathJax_Preview">
+\log \frac{\det P}{\det P_0} = \log \det P - \log \det P_0
+</span><script type="math/tex; mode=display">
+\log \frac{\det P}{\det P_0} = \log \det P - \log \det P_0
+</script></span></p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_det_ratio</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.prior_precision_diag"><code class="name">var <span class="ident">prior_precision_diag</span></code></dt>
+<dd>
+<div class="desc"><p>Obtain the diagonal prior precision <span><span class="MathJax_Preview">p_0</span><script type="math/tex">p_0</script></span> constructed from either
+a scalar, layer-wise, or diagonal prior precision.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>prior_precision_diag</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.prior_mean"><code class="name">var <span class="ident">prior_mean</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.prior_precision"><code class="name">var <span class="ident">prior_precision</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.sigma_noise"><code class="name">var <span class="ident">sigma_noise</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Compute or return the posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>posterior_prec</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.baselaplace.BaseLaplace.fit"><code class="name flex">
+<span>def <span class="ident">fit</span></span>(<span>self, train_loader)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Fit the local Laplace approximation at the parameters of the model.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>train_loader</code></strong> :&ensp;<code>torch.data.utils.DataLoader</code></dt>
+<dd>each iterate is a training batch (X, y);
+<code>train_loader.dataset</code> needs to be set to access <span><span class="MathJax_Preview">N</span><script type="math/tex">N</script></span>, size of the data set</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.log_marginal_likelihood"><code class="name flex">
+<span>def <span class="ident">log_marginal_likelihood</span></span>(<span>self, prior_precision=None, sigma_noise=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute the Laplace approximation to the log marginal likelihood subject
+to specific Hessian approximations that subclasses implement.
+Requires that the Laplace approximation has been fit before.
+The resulting torch.Tensor is differentiable in <code>prior_precision</code> and
+<code>sigma_noise</code> if these have gradients enabled.
+By passing <code>prior_precision</code> or <code>sigma_noise</code>, the current value is
+overwritten. This is useful for iterating on the log marginal likelihood.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>prior_precision</code></strong> :&ensp;<code>torch.Tensor</code>, optional</dt>
+<dd>prior precision if should be changed from current <code>prior_precision</code> value</dd>
+<dt><strong><code>sigma_noise</code></strong> :&ensp;<code>[type]</code>, optional</dt>
+<dd>observation noise standard deviation if should be changed</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_marglik</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.predictive"><code class="name flex">
+<span>def <span class="ident">predictive</span></span>(<span>self, x, pred_type='glm', link_approx='mc', n_samples=100)</span>
+</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.predictive_samples"><code class="name flex">
+<span>def <span class="ident">predictive_samples</span></span>(<span>self, x, pred_type='glm', n_samples=100)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Sample from the posterior predictive on input data <code>x</code>.
+Can be used, for example, for Thompson sampling.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch_size, input_shape)</code></dd>
+<dt><strong><code>pred_type</code></strong> :&ensp;<code>{'glm', 'nn'}</code>, default=<code>'glm'</code></dt>
+<dd>type of posterior predictive, linearized GLM predictive or neural
+network sampling predictive. The GLM predictive is consistent with
+the curvature approximations used here.</dd>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</code></dt>
+<dd>number of samples</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>samples</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>samples <code>(n_samples, batch_size, output_shape)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.functional_variance"><code class="name flex">
+<span>def <span class="ident">functional_variance</span></span>(<span>self, Jacs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute functional variance for the <code>'glm'</code> predictive:
+<code>f_var[i] = Jacs[i] @ P.inv() @ Jacs[i].T</code>, which is a output x output
+predictive covariance matrix.
+Mathematically, we have for a single Jacobian
+<span><span class="MathJax_Preview">\mathcal{J} = \nabla_\theta f(x;\theta)\vert_{\theta_{MAP}}</span><script type="math/tex">\mathcal{J} = \nabla_\theta f(x;\theta)\vert_{\theta_{MAP}}</script></span>
+the output covariance matrix
+<span><span class="MathJax_Preview"> \mathcal{J} P^{-1} \mathcal{J}^T </span><script type="math/tex"> \mathcal{J} P^{-1} \mathcal{J}^T </script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>Jacs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians of model output wrt parameters
+<code>(batch, outputs, parameters)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>f_var</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output covariance <code>(batch, outputs, outputs)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.sample"><code class="name flex">
+<span>def <span class="ident">sample</span></span>(<span>self, n_samples=100)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Sample from the Laplace posterior approximation, i.e.,
+<span><span class="MathJax_Preview"> \theta \sim \mathcal{N}(\theta_{MAP}, P^{-1})</span><script type="math/tex"> \theta \sim \mathcal{N}(\theta_{MAP}, P^{-1})</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>number of samples</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.BaseLaplace.optimize_prior_precision"><code class="name flex">
+<span>def <span class="ident">optimize_prior_precision</span></span>(<span>self, method='marglik', n_steps=100, lr=0.1, init_prior_prec=1.0, val_loader=None, loss=&lt;function get_nll&gt;, log_prior_prec_min=-4, log_prior_prec_max=4, grid_size=100, pred_type='glm', link_approx='probit', n_samples=100, verbose=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Optimize the prior precision post-hoc using the <code>method</code>
+specified by the user.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>method</code></strong> :&ensp;<code>{'marglik', 'CV'}</code>, default=<code>'marglik'</code></dt>
+<dd>specifies how the prior precision should be optimized.</dd>
+<dt><strong><code>n_steps</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>the number of gradient descent steps to take.</dd>
+<dt><strong><code>lr</code></strong> :&ensp;<code>float</code>, default=<code>1e-1</code></dt>
+<dd>the learning rate to use for gradient descent.</dd>
+<dt><strong><code>init_prior_prec</code></strong> :&ensp;<code>float</code>, default=<code>1.0</code></dt>
+<dd>initial prior precision before the first optimization step.</dd>
+<dt><strong><code>val_loader</code></strong> :&ensp;<code>torch.data.utils.DataLoader</code>, default=<code>None</code></dt>
+<dd>DataLoader for the validation set; each iterate is a training batch (X, y).</dd>
+<dt><strong><code>loss</code></strong> :&ensp;<code>callable</code>, default=<code>get_nll</code></dt>
+<dd>loss function to use for CV.</dd>
+<dt><strong><code>log_prior_prec_min</code></strong> :&ensp;<code>float</code>, default=<code>-4</code></dt>
+<dd>lower bound of gridsearch interval for CV.</dd>
+<dt><strong><code>log_prior_prec_max</code></strong> :&ensp;<code>float</code>, default=<code>4</code></dt>
+<dd>upper bound of gridsearch interval for CV.</dd>
+<dt><strong><code>grid_size</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>number of values to consider inside the gridsearch interval for CV.</dd>
+<dt><strong><code>pred_type</code></strong> :&ensp;<code>{'glm', 'nn'}</code>, default=<code>'glm'</code></dt>
+<dd>type of posterior predictive, linearized GLM predictive or neural
+network sampling predictive. The GLM predictive is consistent with
+the curvature approximations used here.</dd>
+<dt><strong><code>link_approx</code></strong> :&ensp;<code>{'mc', 'probit', 'bridge'}</code>, default=<code>'probit'</code></dt>
+<dd>how to approximate the classification link function for the <code>'glm'</code>.
+For <code>pred_type='nn'</code>, only <code>'mc'</code> is possible.</dd>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>number of samples for <code>link_approx='mc'</code>.</dd>
+<dt><strong><code>verbose</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>if true, the optimized prior precision will be printed
+(can be a large tensor if the prior has a diagonal covariance).</dd>
+</dl></div>
+</dd>
+</dl>
+</dd>
+<dt id="laplace.baselaplace.FullLaplace"><code class="flex name class">
+<span>class <span class="ident">FullLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Laplace approximation with full, i.e., dense, log likelihood Hessian approximation
+and hence posterior precision. Based on the chosen <code>backend</code> parameter, the full
+approximation can be, for example, a generalized Gauss-Newton matrix.
+Mathematically, we have <span><span class="MathJax_Preview">P \in \mathbb{R}^{P \times P}</span><script type="math/tex">P \in \mathbb{R}^{P \times P}</script></span>.
+See <code><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.lllaplace.FullLLLaplace" href="lllaplace.html#laplace.lllaplace.FullLLLaplace">FullLLLaplace</a></li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.baselaplace.FullLaplace.posterior_scale"><code class="name">var <span class="ident">posterior_scale</span></code></dt>
+<dd>
+<div class="desc"><p>Posterior scale (square root of the covariance), i.e.,
+<span><span class="MathJax_Preview">P^{-\frac{1}{2}}</span><script type="math/tex">P^{-\frac{1}{2}}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>scale</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.FullLaplace.posterior_covariance"><code class="name">var <span class="ident">posterior_covariance</span></code></dt>
+<dd>
+<div class="desc"><p>Posterior covariance, i.e., <span><span class="MathJax_Preview">P^{-1}</span><script type="math/tex">P^{-1}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>covariance</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.FullLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_posterior_precision" href="#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_prior_precision" href="#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_ratio" href="#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_likelihood" href="#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.prior_precision_diag" href="#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.scatter" href="#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.baselaplace.KronLaplace"><code class="flex name class">
+<span>class <span class="ident">KronLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, damping=False, **backend_kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Laplace approximation with Kronecker factored log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have for each parameter group, e.g., torch.nn.Module,
+that \P\approx Q \otimes H.
+See <code><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></code> for the full interface and see
+<code><a title="laplace.matrix.Kron" href="matrix.html#laplace.matrix.Kron">Kron</a></code> and <code><a title="laplace.matrix.KronDecomposed" href="matrix.html#laplace.matrix.KronDecomposed">KronDecomposed</a></code> for the structure of
+the Kronecker factors. <code>Kron</code> is used to aggregate factors by summing up and
+<code>KronDecomposed</code> is used to add the prior, a Hessian factor (e.g. temperature),
+and computing posterior covariances, marginal likelihood, etc.
+Damping can be enabled by setting <code>damping=True</code>.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.lllaplace.KronLLLaplace" href="lllaplace.html#laplace.lllaplace.KronLLLaplace">KronLLLaplace</a></li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.baselaplace.KronLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Kronecker factored Posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code><a title="laplace.matrix.KronDecomposed" href="matrix.html#laplace.matrix.KronDecomposed">KronDecomposed</a></code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.KronLaplace.prior_precision"><code class="name">var <span class="ident">prior_precision</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_posterior_precision" href="#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_prior_precision" href="#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_ratio" href="#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_likelihood" href="#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.prior_precision_diag" href="#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.scatter" href="#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.baselaplace.DiagLaplace"><code class="flex name class">
+<span>class <span class="ident">DiagLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Laplace approximation with diagonal log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have <span><span class="MathJax_Preview">P \approx \textrm{diag}(P)</span><script type="math/tex">P \approx \textrm{diag}(P)</script></span>.
+See <code><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.lllaplace.DiagLLLaplace" href="lllaplace.html#laplace.lllaplace.DiagLLLaplace">DiagLLLaplace</a></li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.baselaplace.DiagLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Diagonal posterior precision <span><span class="MathJax_Preview">p</span><script type="math/tex">p</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.DiagLaplace.posterior_scale"><code class="name">var <span class="ident">posterior_scale</span></code></dt>
+<dd>
+<div class="desc"><p>Diagonal posterior scale <span><span class="MathJax_Preview">\sqrt{p^{-1}}</span><script type="math/tex">\sqrt{p^{-1}}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.baselaplace.DiagLaplace.posterior_variance"><code class="name">var <span class="ident">posterior_variance</span></code></dt>
+<dd>
+<div class="desc"><p>Diagonal posterior variance <span><span class="MathJax_Preview">p^{-1}</span><script type="math/tex">p^{-1}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_posterior_precision" href="#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_prior_precision" href="#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_ratio" href="#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_likelihood" href="#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.prior_precision_diag" href="#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.scatter" href="#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace" href="index.html">laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.baselaplace.BaseLaplace" href="#laplace.baselaplace.BaseLaplace">BaseLaplace</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive" href="#laplace.baselaplace.BaseLaplace.predictive">predictive</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.baselaplace.FullLaplace" href="#laplace.baselaplace.FullLaplace">FullLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.baselaplace.KronLaplace" href="#laplace.baselaplace.KronLaplace">KronLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.baselaplace.DiagLaplace" href="#laplace.baselaplace.DiagLaplace">DiagLaplace</a></code></h4>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/curvature/asdl.html b/docs/curvature/asdl.html
new file mode 100644
index 00000000..2e4a037a
--- /dev/null
+++ b/docs/curvature/asdl.html
@@ -0,0 +1,204 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.curvature.asdl API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.curvature.asdl</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.curvature.asdl.AsdlInterface"><code class="flex name class">
+<span>class <span class="ident">AsdlInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface for asdfghjkl backend.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlEF" href="#laplace.curvature.asdl.AsdlEF">AsdlEF</a></li>
+<li><a title="laplace.curvature.asdl.AsdlGGN" href="#laplace.curvature.asdl.AsdlGGN">AsdlGGN</a></li>
+</ul>
+<h3>Static methods</h3>
+<dl>
+<dt id="laplace.curvature.asdl.AsdlInterface.jacobians"><code class="name flex">
+<span>def <span class="ident">jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_\theta f(x;\theta)</span><script type="math/tex">\nabla_\theta f(x;\theta)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>
+using asdfghjkl's gradient per output dimension.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.asdl.AsdlInterface.gradients"><code class="name flex">
+<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute gradients <span><span class="MathJax_Preview">\nabla_\theta \ell(f(x;\theta, y)</span><script type="math/tex">\nabla_\theta \ell(f(x;\theta, y)</script></span> at current parameter
+<span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span> using asdfghjkl's backend.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>Gs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>gradients <code>(batch, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.asdl.AsdlGGN"><code class="flex name class">
+<span>class <span class="ident">AsdlGGN</span></span>
+<span>(</span><span>model, likelihood, last_layer=False, stochastic=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of the <code>GGNInterface</code> using asdfghjkl.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlInterface" href="#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></li>
+<li><a title="laplace.curvature.curvature.GGNInterface" href="curvature.html#laplace.curvature.curvature.GGNInterface">GGNInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.asdl.AsdlInterface" href="#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.gradients" href="#laplace.curvature.asdl.AsdlInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.jacobians" href="#laplace.curvature.asdl.AsdlInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.asdl.AsdlEF"><code class="flex name class">
+<span>class <span class="ident">AsdlEF</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of the <code>EFInterface</code> using asdfghjkl.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlInterface" href="#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></li>
+<li><a title="laplace.curvature.curvature.EFInterface" href="curvature.html#laplace.curvature.curvature.EFInterface">EFInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.asdl.AsdlInterface" href="#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.gradients" href="#laplace.curvature.asdl.AsdlInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.jacobians" href="#laplace.curvature.asdl.AsdlInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace.curvature" href="index.html">laplace.curvature</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.curvature.asdl.AsdlInterface" href="#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.jacobians" href="#laplace.curvature.asdl.AsdlInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.gradients" href="#laplace.curvature.asdl.AsdlInterface.gradients">gradients</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.asdl.AsdlGGN" href="#laplace.curvature.asdl.AsdlGGN">AsdlGGN</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.asdl.AsdlEF" href="#laplace.curvature.asdl.AsdlEF">AsdlEF</a></code></h4>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/curvature/backpack.html b/docs/curvature/backpack.html
new file mode 100644
index 00000000..02120458
--- /dev/null
+++ b/docs/curvature/backpack.html
@@ -0,0 +1,204 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.curvature.backpack API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.curvature.backpack</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.curvature.backpack.BackPackInterface"><code class="flex name class">
+<span>class <span class="ident">BackPackInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface for Backpack backend.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.backpack.BackPackEF" href="#laplace.curvature.backpack.BackPackEF">BackPackEF</a></li>
+<li><a title="laplace.curvature.backpack.BackPackGGN" href="#laplace.curvature.backpack.BackPackGGN">BackPackGGN</a></li>
+</ul>
+<h3>Static methods</h3>
+<dl>
+<dt id="laplace.curvature.backpack.BackPackInterface.jacobians"><code class="name flex">
+<span>def <span class="ident">jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_{\theta} f(x;\theta)</span><script type="math/tex">\nabla_{\theta} f(x;\theta)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>
+using backpack's BatchGrad per output dimension.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.backpack.BackPackInterface.gradients"><code class="name flex">
+<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute gradients <span><span class="MathJax_Preview">\nabla_\theta \ell(f(x;\theta, y)</span><script type="math/tex">\nabla_\theta \ell(f(x;\theta, y)</script></span> at current parameter
+<span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span> using Backpack's BatchGrad.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>Gs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>gradients <code>(batch, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.backpack.BackPackGGN"><code class="flex name class">
+<span>class <span class="ident">BackPackGGN</span></span>
+<span>(</span><span>model, likelihood, last_layer=False, stochastic=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of the <code>GGNInterface</code> using Backpack.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.backpack.BackPackInterface" href="#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></li>
+<li><a title="laplace.curvature.curvature.GGNInterface" href="curvature.html#laplace.curvature.curvature.GGNInterface">GGNInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.backpack.BackPackInterface" href="#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.gradients" href="#laplace.curvature.backpack.BackPackInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.jacobians" href="#laplace.curvature.backpack.BackPackInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.backpack.BackPackEF"><code class="flex name class">
+<span>class <span class="ident">BackPackEF</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of <code>EFInterface</code> using Backpack.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.backpack.BackPackInterface" href="#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></li>
+<li><a title="laplace.curvature.curvature.EFInterface" href="curvature.html#laplace.curvature.curvature.EFInterface">EFInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.backpack.BackPackInterface" href="#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.gradients" href="#laplace.curvature.backpack.BackPackInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.jacobians" href="#laplace.curvature.backpack.BackPackInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace.curvature" href="index.html">laplace.curvature</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.curvature.backpack.BackPackInterface" href="#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.jacobians" href="#laplace.curvature.backpack.BackPackInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.gradients" href="#laplace.curvature.backpack.BackPackInterface.gradients">gradients</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.backpack.BackPackGGN" href="#laplace.curvature.backpack.BackPackGGN">BackPackGGN</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.backpack.BackPackEF" href="#laplace.curvature.backpack.BackPackEF">BackPackEF</a></code></h4>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/curvature/curvature.html b/docs/curvature/curvature.html
new file mode 100644
index 00000000..8f47d2f5
--- /dev/null
+++ b/docs/curvature/curvature.html
@@ -0,0 +1,396 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.curvature.curvature API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.curvature.curvature</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.curvature.curvature.CurvatureInterface"><code class="flex name class">
+<span>class <span class="ident">CurvatureInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface to access curvature for a model and corresponding likelihood.
+A <code><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></code> must inherit from this baseclass and implement the
+necessary functions <code>jacobians</code>, <code>full</code>, <code>kron</code>, and <code>diag</code>.
+The interface might be extended in the future to account for other curvature
+structures, for example, a block-diagonal one.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code> or <code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>torch model (neural network)</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>last_layer</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>only consider curvature of last layer</dd>
+</dl>
+<h2 id="attributes">Attributes</h2>
+<dl>
+<dt><strong><code>lossfunc</code></strong> :&ensp;<code>torch.nn.MSELoss</code> or <code>torch.nn.CrossEntropyLoss</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>factor</code></strong> :&ensp;<code>float</code></dt>
+<dd>conversion factor between torch losses and base likelihoods
+For example, <span><span class="MathJax_Preview">\frac{1}{2}</span><script type="math/tex">\frac{1}{2}</script></span> to get to <span><span class="MathJax_Preview">\mathcal{N}(f, 1)</span><script type="math/tex">\mathcal{N}(f, 1)</script></span> from MSELoss.</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlInterface" href="asdl.html#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></li>
+<li><a title="laplace.curvature.backpack.BackPackInterface" href="backpack.html#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></li>
+<li><a title="laplace.curvature.curvature.EFInterface" href="#laplace.curvature.curvature.EFInterface">EFInterface</a></li>
+<li><a title="laplace.curvature.curvature.GGNInterface" href="#laplace.curvature.curvature.GGNInterface">GGNInterface</a></li>
+</ul>
+<h3>Static methods</h3>
+<dl>
+<dt id="laplace.curvature.curvature.CurvatureInterface.jacobians"><code class="name flex">
+<span>def <span class="ident">jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_\theta f(x;\theta)</span><script type="math/tex">\nabla_\theta f(x;\theta)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians"><code class="name flex">
+<span>def <span class="ident">last_layer_jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})</span><script type="math/tex">\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})</script></span>
+only at current last-layer parameter <span><span class="MathJax_Preview">\theta_{\textrm{last}}</span><script type="math/tex">\theta_{\textrm{last}}</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, last-layer-parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.curvature.CurvatureInterface.gradients"><code class="name flex">
+<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute gradients <span><span class="MathJax_Preview">\nabla_\theta \ell(f(x;\theta, y)</span><script type="math/tex">\nabla_\theta \ell(f(x;\theta, y)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>Gs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>gradients <code>(batch, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.curvature.CurvatureInterface.full"><code class="name flex">
+<span>def <span class="ident">full</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute a dense curvature (approximation) in the form of a <span><span class="MathJax_Preview">P \times P</span><script type="math/tex">P \times P</script></span> matrix
+<span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> with respect to parameters <span><span class="MathJax_Preview">\theta \in \mathbb{R}^P</span><script type="math/tex">\theta \in \mathbb{R}^P</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Hessian approximation <code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.curvature.CurvatureInterface.kron"><code class="name flex">
+<span>def <span class="ident">kron</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute a Kronecker factored curvature approximation (such as KFAC).
+The approximation to <span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> takes the form of two Kronecker factors <span><span class="MathJax_Preview">Q, H</span><script type="math/tex">Q, H</script></span>,
+i.e., <span><span class="MathJax_Preview">H \approx Q \otimes H</span><script type="math/tex">H \approx Q \otimes H</script></span> for each Module in the neural network permitting
+such curvature.
+<span><span class="MathJax_Preview">Q</span><script type="math/tex">Q</script></span> is quadratic in the input-dimension of a module <span><span class="MathJax_Preview">p_{in} \times p_{in}</span><script type="math/tex">p_{in} \times p_{in}</script></span>
+and <span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> in the output-dimension <span><span class="MathJax_Preview">p_{out} \times p_{out}</span><script type="math/tex">p_{out} \times p_{out}</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H</code></strong> :&ensp;<code><a title="laplace.matrix.Kron" href="../matrix.html#laplace.matrix.Kron">Kron</a></code></dt>
+<dd>Kronecker factored Hessian approximation.</dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.curvature.CurvatureInterface.diag"><code class="name flex">
+<span>def <span class="ident">diag</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute a diagonal Hessian approximation to <span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> and is represented as a
+vector of the dimensionality of parameters <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>vector representing the diagonal of H</dd>
+</dl></div>
+</dd>
+</dl>
+</dd>
+<dt id="laplace.curvature.curvature.GGNInterface"><code class="flex name class">
+<span>class <span class="ident">GGNInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False, stochastic=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Generalized Gauss-Newton or Fisher Curvature Interface.
+The GGN is equal to the Fisher information for the available likelihoods.
+In addition to <code><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></code>, methods for Jacobians are required by subclasses.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code> or <code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>torch model (neural network)</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>last_layer</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>only consider curvature of last layer</dd>
+<dt><strong><code>stochastic</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>Fisher if stochastic else GGN</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlGGN" href="asdl.html#laplace.curvature.asdl.AsdlGGN">AsdlGGN</a></li>
+<li><a title="laplace.curvature.backpack.BackPackGGN" href="backpack.html#laplace.curvature.backpack.BackPackGGN">BackPackGGN</a></li>
+</ul>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.curvature.GGNInterface.full"><code class="name flex">
+<span>def <span class="ident">full</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute the full GGN <span><span class="MathJax_Preview">P \times P</span><script type="math/tex">P \times P</script></span> matrix as Hessian approximation
+<span><span class="MathJax_Preview">H_{ggn}</span><script type="math/tex">H_{ggn}</script></span> with respect to parameters <span><span class="MathJax_Preview">\theta \in \mathbb{R}^P</span><script type="math/tex">\theta \in \mathbb{R}^P</script></span>.
+For last-layer, reduced to <span><span class="MathJax_Preview">\theta_{last}</span><script type="math/tex">\theta_{last}</script></span></p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H_ggn</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>GGN <code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.gradients" href="#laplace.curvature.curvature.CurvatureInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.jacobians" href="#laplace.curvature.curvature.CurvatureInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.curvature.EFInterface"><code class="flex name class">
+<span>class <span class="ident">EFInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface for Empirical Fisher as Hessian approximation.
+In addition to <code><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></code>, methods for gradients are required by subclasses.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code> or <code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>torch model (neural network)</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>last_layer</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>only consider curvature of last layer</dd>
+</dl>
+<h2 id="attributes">Attributes</h2>
+<dl>
+<dt><strong><code>lossfunc</code></strong> :&ensp;<code>torch.nn.MSELoss</code> or <code>torch.nn.CrossEntropyLoss</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>factor</code></strong> :&ensp;<code>float</code></dt>
+<dd>conversion factor between torch losses and base likelihoods
+For example, <span><span class="MathJax_Preview">\frac{1}{2}</span><script type="math/tex">\frac{1}{2}</script></span> to get to <span><span class="MathJax_Preview">\mathcal{N}(f, 1)</span><script type="math/tex">\mathcal{N}(f, 1)</script></span> from MSELoss.</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlEF" href="asdl.html#laplace.curvature.asdl.AsdlEF">AsdlEF</a></li>
+<li><a title="laplace.curvature.backpack.BackPackEF" href="backpack.html#laplace.curvature.backpack.BackPackEF">BackPackEF</a></li>
+</ul>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.curvature.EFInterface.full"><code class="name flex">
+<span>def <span class="ident">full</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute the full EF <span><span class="MathJax_Preview">P \times P</span><script type="math/tex">P \times P</script></span> matrix as Hessian approximation
+<span><span class="MathJax_Preview">H_{ef}</span><script type="math/tex">H_{ef}</script></span> with respect to parameters <span><span class="MathJax_Preview">\theta \in \mathbb{R}^P</span><script type="math/tex">\theta \in \mathbb{R}^P</script></span>.
+For last-layer, reduced to <span><span class="MathJax_Preview">\theta_{last}</span><script type="math/tex">\theta_{last}</script></span></p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H_ef</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>EF <code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.gradients" href="#laplace.curvature.curvature.CurvatureInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.jacobians" href="#laplace.curvature.curvature.CurvatureInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace.curvature" href="index.html">laplace.curvature</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.curvature.curvature.CurvatureInterface" href="#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.jacobians" href="#laplace.curvature.curvature.CurvatureInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.gradients" href="#laplace.curvature.curvature.CurvatureInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.full" href="#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.curvature.GGNInterface" href="#laplace.curvature.curvature.GGNInterface">GGNInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.curvature.GGNInterface.full" href="#laplace.curvature.curvature.GGNInterface.full">full</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.curvature.EFInterface" href="#laplace.curvature.curvature.EFInterface">EFInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.curvature.EFInterface.full" href="#laplace.curvature.curvature.EFInterface.full">full</a></code></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/curvature/index.html b/docs/curvature/index.html
new file mode 100644
index 00000000..77c3799b
--- /dev/null
+++ b/docs/curvature/index.html
@@ -0,0 +1,704 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.curvature API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.curvature</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+<h2 class="section-title" id="header-submodules">Sub-modules</h2>
+<dl>
+<dt><code class="name"><a title="laplace.curvature.asdl" href="asdl.html">laplace.curvature.asdl</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.curvature.backpack" href="backpack.html">laplace.curvature.backpack</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.curvature.curvature" href="curvature.html">laplace.curvature.curvature</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+</dl>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.curvature.CurvatureInterface"><code class="flex name class">
+<span>class <span class="ident">CurvatureInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface to access curvature for a model and corresponding likelihood.
+A <code><a title="laplace.curvature.CurvatureInterface" href="#laplace.curvature.CurvatureInterface">CurvatureInterface</a></code> must inherit from this baseclass and implement the
+necessary functions <code>jacobians</code>, <code>full</code>, <code>kron</code>, and <code>diag</code>.
+The interface might be extended in the future to account for other curvature
+structures, for example, a block-diagonal one.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code> or <code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>torch model (neural network)</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>last_layer</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>only consider curvature of last layer</dd>
+</dl>
+<h2 id="attributes">Attributes</h2>
+<dl>
+<dt><strong><code>lossfunc</code></strong> :&ensp;<code>torch.nn.MSELoss</code> or <code>torch.nn.CrossEntropyLoss</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>factor</code></strong> :&ensp;<code>float</code></dt>
+<dd>conversion factor between torch losses and base likelihoods
+For example, <span><span class="MathJax_Preview">\frac{1}{2}</span><script type="math/tex">\frac{1}{2}</script></span> to get to <span><span class="MathJax_Preview">\mathcal{N}(f, 1)</span><script type="math/tex">\mathcal{N}(f, 1)</script></span> from MSELoss.</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlInterface" href="asdl.html#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></li>
+<li><a title="laplace.curvature.backpack.BackPackInterface" href="backpack.html#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></li>
+<li><a title="laplace.curvature.curvature.EFInterface" href="curvature.html#laplace.curvature.curvature.EFInterface">EFInterface</a></li>
+<li><a title="laplace.curvature.curvature.GGNInterface" href="curvature.html#laplace.curvature.curvature.GGNInterface">GGNInterface</a></li>
+</ul>
+<h3>Static methods</h3>
+<dl>
+<dt id="laplace.curvature.CurvatureInterface.jacobians"><code class="name flex">
+<span>def <span class="ident">jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_\theta f(x;\theta)</span><script type="math/tex">\nabla_\theta f(x;\theta)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.CurvatureInterface.last_layer_jacobians"><code class="name flex">
+<span>def <span class="ident">last_layer_jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})</span><script type="math/tex">\nabla_{\theta_\textrm{last}} f(x;\theta_\textrm{last})</script></span>
+only at current last-layer parameter <span><span class="MathJax_Preview">\theta_{\textrm{last}}</span><script type="math/tex">\theta_{\textrm{last}}</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, last-layer-parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.CurvatureInterface.gradients"><code class="name flex">
+<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute gradients <span><span class="MathJax_Preview">\nabla_\theta \ell(f(x;\theta, y)</span><script type="math/tex">\nabla_\theta \ell(f(x;\theta, y)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>Gs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>gradients <code>(batch, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.CurvatureInterface.full"><code class="name flex">
+<span>def <span class="ident">full</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute a dense curvature (approximation) in the form of a <span><span class="MathJax_Preview">P \times P</span><script type="math/tex">P \times P</script></span> matrix
+<span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> with respect to parameters <span><span class="MathJax_Preview">\theta \in \mathbb{R}^P</span><script type="math/tex">\theta \in \mathbb{R}^P</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Hessian approximation <code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.CurvatureInterface.kron"><code class="name flex">
+<span>def <span class="ident">kron</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute a Kronecker factored curvature approximation (such as KFAC).
+The approximation to <span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> takes the form of two Kronecker factors <span><span class="MathJax_Preview">Q, H</span><script type="math/tex">Q, H</script></span>,
+i.e., <span><span class="MathJax_Preview">H \approx Q \otimes H</span><script type="math/tex">H \approx Q \otimes H</script></span> for each Module in the neural network permitting
+such curvature.
+<span><span class="MathJax_Preview">Q</span><script type="math/tex">Q</script></span> is quadratic in the input-dimension of a module <span><span class="MathJax_Preview">p_{in} \times p_{in}</span><script type="math/tex">p_{in} \times p_{in}</script></span>
+and <span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> in the output-dimension <span><span class="MathJax_Preview">p_{out} \times p_{out}</span><script type="math/tex">p_{out} \times p_{out}</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H</code></strong> :&ensp;<code><a title="laplace.matrix.Kron" href="../matrix.html#laplace.matrix.Kron">Kron</a></code></dt>
+<dd>Kronecker factored Hessian approximation.</dd>
+</dl></div>
+</dd>
+<dt id="laplace.curvature.CurvatureInterface.diag"><code class="name flex">
+<span>def <span class="ident">diag</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute a diagonal Hessian approximation to <span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> and is represented as a
+vector of the dimensionality of parameters <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>vector representing the diagonal of H</dd>
+</dl></div>
+</dd>
+</dl>
+</dd>
+<dt id="laplace.curvature.GGNInterface"><code class="flex name class">
+<span>class <span class="ident">GGNInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False, stochastic=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Generalized Gauss-Newton or Fisher Curvature Interface.
+The GGN is equal to the Fisher information for the available likelihoods.
+In addition to <code><a title="laplace.curvature.CurvatureInterface" href="#laplace.curvature.CurvatureInterface">CurvatureInterface</a></code>, methods for Jacobians are required by subclasses.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code> or <code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>torch model (neural network)</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>last_layer</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>only consider curvature of last layer</dd>
+<dt><strong><code>stochastic</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>Fisher if stochastic else GGN</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlGGN" href="asdl.html#laplace.curvature.asdl.AsdlGGN">AsdlGGN</a></li>
+<li><a title="laplace.curvature.backpack.BackPackGGN" href="backpack.html#laplace.curvature.backpack.BackPackGGN">BackPackGGN</a></li>
+</ul>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.GGNInterface.full"><code class="name flex">
+<span>def <span class="ident">full</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute the full GGN <span><span class="MathJax_Preview">P \times P</span><script type="math/tex">P \times P</script></span> matrix as Hessian approximation
+<span><span class="MathJax_Preview">H_{ggn}</span><script type="math/tex">H_{ggn}</script></span> with respect to parameters <span><span class="MathJax_Preview">\theta \in \mathbb{R}^P</span><script type="math/tex">\theta \in \mathbb{R}^P</script></span>.
+For last-layer, reduced to <span><span class="MathJax_Preview">\theta_{last}</span><script type="math/tex">\theta_{last}</script></span></p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H_ggn</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>GGN <code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.gradients" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.EFInterface"><code class="flex name class">
+<span>class <span class="ident">EFInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface for Empirical Fisher as Hessian approximation.
+In addition to <code><a title="laplace.curvature.CurvatureInterface" href="#laplace.curvature.CurvatureInterface">CurvatureInterface</a></code>, methods for gradients are required by subclasses.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code> or <code><a title="laplace.feature_extractor.FeatureExtractor" href="../feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>torch model (neural network)</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>last_layer</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>only consider curvature of last layer</dd>
+</dl>
+<h2 id="attributes">Attributes</h2>
+<dl>
+<dt><strong><code>lossfunc</code></strong> :&ensp;<code>torch.nn.MSELoss</code> or <code>torch.nn.CrossEntropyLoss</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>factor</code></strong> :&ensp;<code>float</code></dt>
+<dd>conversion factor between torch losses and base likelihoods
+For example, <span><span class="MathJax_Preview">\frac{1}{2}</span><script type="math/tex">\frac{1}{2}</script></span> to get to <span><span class="MathJax_Preview">\mathcal{N}(f, 1)</span><script type="math/tex">\mathcal{N}(f, 1)</script></span> from MSELoss.</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlEF" href="asdl.html#laplace.curvature.asdl.AsdlEF">AsdlEF</a></li>
+<li><a title="laplace.curvature.backpack.BackPackEF" href="backpack.html#laplace.curvature.backpack.BackPackEF">BackPackEF</a></li>
+</ul>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.EFInterface.full"><code class="name flex">
+<span>def <span class="ident">full</span></span>(<span>self, x, y, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute the full EF <span><span class="MathJax_Preview">P \times P</span><script type="math/tex">P \times P</script></span> matrix as Hessian approximation
+<span><span class="MathJax_Preview">H_{ef}</span><script type="math/tex">H_{ef}</script></span> with respect to parameters <span><span class="MathJax_Preview">\theta \in \mathbb{R}^P</span><script type="math/tex">\theta \in \mathbb{R}^P</script></span>.
+For last-layer, reduced to <span><span class="MathJax_Preview">\theta_{last}</span><script type="math/tex">\theta_{last}</script></span></p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code></dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>labels <code>(batch, label_shape)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>H_ef</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>EF <code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.gradients" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.BackPackInterface"><code class="flex name class">
+<span>class <span class="ident">BackPackInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface for Backpack backend.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.backpack.BackPackEF" href="backpack.html#laplace.curvature.backpack.BackPackEF">BackPackEF</a></li>
+<li><a title="laplace.curvature.backpack.BackPackGGN" href="backpack.html#laplace.curvature.backpack.BackPackGGN">BackPackGGN</a></li>
+</ul>
+<h3>Static methods</h3>
+<dl>
+<dt id="laplace.curvature.BackPackInterface.jacobians"><code class="name flex">
+<span>def <span class="ident">jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_{\theta} f(x;\theta)</span><script type="math/tex">\nabla_{\theta} f(x;\theta)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>
+using backpack's BatchGrad per output dimension.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.BackPackInterface.gradients"><code class="name flex">
+<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute gradients <span><span class="MathJax_Preview">\nabla_\theta \ell(f(x;\theta, y)</span><script type="math/tex">\nabla_\theta \ell(f(x;\theta, y)</script></span> at current parameter
+<span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span> using Backpack's BatchGrad.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>Gs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>gradients <code>(batch, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.BackPackGGN"><code class="flex name class">
+<span>class <span class="ident">BackPackGGN</span></span>
+<span>(</span><span>model, likelihood, last_layer=False, stochastic=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of the <code><a title="laplace.curvature.GGNInterface" href="#laplace.curvature.GGNInterface">GGNInterface</a></code> using Backpack.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.backpack.BackPackInterface" href="backpack.html#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></li>
+<li><a title="laplace.curvature.curvature.GGNInterface" href="curvature.html#laplace.curvature.curvature.GGNInterface">GGNInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.backpack.BackPackInterface" href="backpack.html#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.gradients" href="backpack.html#laplace.curvature.backpack.BackPackInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.jacobians" href="backpack.html#laplace.curvature.backpack.BackPackInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.BackPackEF"><code class="flex name class">
+<span>class <span class="ident">BackPackEF</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of <code><a title="laplace.curvature.EFInterface" href="#laplace.curvature.EFInterface">EFInterface</a></code> using Backpack.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.backpack.BackPackInterface" href="backpack.html#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></li>
+<li><a title="laplace.curvature.curvature.EFInterface" href="curvature.html#laplace.curvature.curvature.EFInterface">EFInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.backpack.BackPackInterface" href="backpack.html#laplace.curvature.backpack.BackPackInterface">BackPackInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.gradients" href="backpack.html#laplace.curvature.backpack.BackPackInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.jacobians" href="backpack.html#laplace.curvature.backpack.BackPackInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.backpack.BackPackInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.AsdlInterface"><code class="flex name class">
+<span>class <span class="ident">AsdlInterface</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Interface for asdfghjkl backend.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlEF" href="asdl.html#laplace.curvature.asdl.AsdlEF">AsdlEF</a></li>
+<li><a title="laplace.curvature.asdl.AsdlGGN" href="asdl.html#laplace.curvature.asdl.AsdlGGN">AsdlGGN</a></li>
+</ul>
+<h3>Static methods</h3>
+<dl>
+<dt id="laplace.curvature.AsdlInterface.jacobians"><code class="name flex">
+<span>def <span class="ident">jacobians</span></span>(<span>model, x)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute Jacobians <span><span class="MathJax_Preview">\nabla_\theta f(x;\theta)</span><script type="math/tex">\nabla_\theta f(x;\theta)</script></span> at current parameter <span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span>
+using asdfghjkl's gradient per output dimension.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>Js</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians <code>(batch, parameters, outputs)</code></dd>
+<dt><strong><code>f</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output function <code>(batch, outputs)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.curvature.AsdlInterface.gradients"><code class="name flex">
+<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute gradients <span><span class="MathJax_Preview">\nabla_\theta \ell(f(x;\theta, y)</span><script type="math/tex">\nabla_\theta \ell(f(x;\theta, y)</script></span> at current parameter
+<span><span class="MathJax_Preview">\theta</span><script type="math/tex">\theta</script></span> using asdfghjkl's backend.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch, input_shape)</code> on compatible device with model.</dd>
+<dt><strong><code>y</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>loss</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>Gs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>gradients <code>(batch, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.AsdlGGN"><code class="flex name class">
+<span>class <span class="ident">AsdlGGN</span></span>
+<span>(</span><span>model, likelihood, last_layer=False, stochastic=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of the <code><a title="laplace.curvature.GGNInterface" href="#laplace.curvature.GGNInterface">GGNInterface</a></code> using asdfghjkl.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlInterface" href="asdl.html#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></li>
+<li><a title="laplace.curvature.curvature.GGNInterface" href="curvature.html#laplace.curvature.curvature.GGNInterface">GGNInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.asdl.AsdlInterface" href="asdl.html#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.gradients" href="asdl.html#laplace.curvature.asdl.AsdlInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.jacobians" href="asdl.html#laplace.curvature.asdl.AsdlInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.curvature.AsdlEF"><code class="flex name class">
+<span>class <span class="ident">AsdlEF</span></span>
+<span>(</span><span>model, likelihood, last_layer=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Implementation of the <code><a title="laplace.curvature.EFInterface" href="#laplace.curvature.EFInterface">EFInterface</a></code> using asdfghjkl.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.curvature.asdl.AsdlInterface" href="asdl.html#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></li>
+<li><a title="laplace.curvature.curvature.EFInterface" href="curvature.html#laplace.curvature.curvature.EFInterface">EFInterface</a></li>
+<li><a title="laplace.curvature.curvature.CurvatureInterface" href="curvature.html#laplace.curvature.curvature.CurvatureInterface">CurvatureInterface</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.curvature.asdl.AsdlInterface" href="asdl.html#laplace.curvature.asdl.AsdlInterface">AsdlInterface</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.diag" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.diag">diag</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.full" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.gradients" href="asdl.html#laplace.curvature.asdl.AsdlInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.jacobians" href="asdl.html#laplace.curvature.asdl.AsdlInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.kron" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.asdl.AsdlInterface.last_layer_jacobians" href="curvature.html#laplace.curvature.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace" href="../index.html">laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-submodules">Sub-modules</a></h3>
+<ul>
+<li><code><a title="laplace.curvature.asdl" href="asdl.html">laplace.curvature.asdl</a></code></li>
+<li><code><a title="laplace.curvature.backpack" href="backpack.html">laplace.curvature.backpack</a></code></li>
+<li><code><a title="laplace.curvature.curvature" href="curvature.html">laplace.curvature.curvature</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.curvature.CurvatureInterface" href="#laplace.curvature.CurvatureInterface">CurvatureInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.CurvatureInterface.jacobians" href="#laplace.curvature.CurvatureInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.CurvatureInterface.last_layer_jacobians" href="#laplace.curvature.CurvatureInterface.last_layer_jacobians">last_layer_jacobians</a></code></li>
+<li><code><a title="laplace.curvature.CurvatureInterface.gradients" href="#laplace.curvature.CurvatureInterface.gradients">gradients</a></code></li>
+<li><code><a title="laplace.curvature.CurvatureInterface.full" href="#laplace.curvature.CurvatureInterface.full">full</a></code></li>
+<li><code><a title="laplace.curvature.CurvatureInterface.kron" href="#laplace.curvature.CurvatureInterface.kron">kron</a></code></li>
+<li><code><a title="laplace.curvature.CurvatureInterface.diag" href="#laplace.curvature.CurvatureInterface.diag">diag</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.GGNInterface" href="#laplace.curvature.GGNInterface">GGNInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.GGNInterface.full" href="#laplace.curvature.GGNInterface.full">full</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.EFInterface" href="#laplace.curvature.EFInterface">EFInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.EFInterface.full" href="#laplace.curvature.EFInterface.full">full</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.BackPackInterface" href="#laplace.curvature.BackPackInterface">BackPackInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.BackPackInterface.jacobians" href="#laplace.curvature.BackPackInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.BackPackInterface.gradients" href="#laplace.curvature.BackPackInterface.gradients">gradients</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.BackPackGGN" href="#laplace.curvature.BackPackGGN">BackPackGGN</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.BackPackEF" href="#laplace.curvature.BackPackEF">BackPackEF</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.AsdlInterface" href="#laplace.curvature.AsdlInterface">AsdlInterface</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.curvature.AsdlInterface.jacobians" href="#laplace.curvature.AsdlInterface.jacobians">jacobians</a></code></li>
+<li><code><a title="laplace.curvature.AsdlInterface.gradients" href="#laplace.curvature.AsdlInterface.gradients">gradients</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.AsdlGGN" href="#laplace.curvature.AsdlGGN">AsdlGGN</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.curvature.AsdlEF" href="#laplace.curvature.AsdlEF">AsdlEF</a></code></h4>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/feature_extractor.html b/docs/feature_extractor.html
new file mode 100644
index 00000000..db8a7509
--- /dev/null
+++ b/docs/feature_extractor.html
@@ -0,0 +1,164 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.feature_extractor API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.feature_extractor</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.feature_extractor.FeatureExtractor"><code class="flex name class">
+<span>class <span class="ident">FeatureExtractor</span></span>
+<span>(</span><span>model: torch.nn.modules.module.Module, last_layer_name: Optional[str] = None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Feature extractor for a PyTorch neural network.
+A wrapper which can return the output of the penultimate layer in addition to
+the output of the last layer for each forward pass. If the name of the last
+layer is not known, it can determine it automatically. It assumes that the
+last layer is linear and that for every forward pass the last layer is the same.
+If the name of the last layer is known, it can be passed as a parameter at
+initilization; this is the safest way to use this class.
+Based on <a href="https://gist.github.com/fkodom/27ed045c9051a39102e8bcf4ce31df76.">https://gist.github.com/fkodom/27ed045c9051a39102e8bcf4ce31df76.</a></p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>PyTorch model</dd>
+<dt><strong><code>last_layer_name</code></strong> :&ensp;<code>str</code>, default=<code>None</code></dt>
+<dd>if the name of the last layer is already known, otherwise it will
+be determined automatically.</dd>
+</dl>
+<p>Initializes internal Module state, shared by both nn.Module and ScriptModule.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>torch.nn.modules.module.Module</li>
+</ul>
+<h3>Class variables</h3>
+<dl>
+<dt id="laplace.feature_extractor.FeatureExtractor.dump_patches"><code class="name">var <span class="ident">dump_patches</span> : bool</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.feature_extractor.FeatureExtractor.training"><code class="name">var <span class="ident">training</span> : bool</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.feature_extractor.FeatureExtractor.forward"><code class="name flex">
+<span>def <span class="ident">forward</span></span>(<span>self, x: torch.Tensor) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Forward pass. If the last layer is not known yet, it will be
+determined when this function is called for the first time.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>one batch of data to use as input for the forward pass</dd>
+</dl></div>
+</dd>
+<dt id="laplace.feature_extractor.FeatureExtractor.forward_with_features"><code class="name flex">
+<span>def <span class="ident">forward_with_features</span></span>(<span>self, x: torch.Tensor) ‑> Tuple[torch.Tensor, torch.Tensor]</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Forward pass which returns the output of the penultimate layer along
+with the output of the last layer. If the last layer is not known yet,
+it will be determined when this function is called for the first time.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>one batch of data to use as input for the forward pass</dd>
+</dl></div>
+</dd>
+<dt id="laplace.feature_extractor.FeatureExtractor.set_last_layer"><code class="name flex">
+<span>def <span class="ident">set_last_layer</span></span>(<span>self, last_layer_name: str) ‑> NoneType</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Set the last layer of the model by its name. This sets the forward
+hook to get the output of the penultimate layer.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>last_layer_name</code></strong> :&ensp;<code>str</code></dt>
+<dd>the name of the last layer (fixed in <code>model.named_modules()</code>).</dd>
+</dl></div>
+</dd>
+<dt id="laplace.feature_extractor.FeatureExtractor.find_last_layer"><code class="name flex">
+<span>def <span class="ident">find_last_layer</span></span>(<span>self, x: torch.Tensor) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Automatically determines the last layer of the model with one
+forward pass. It assumes that the last layer is the same for every
+forward pass and that it is an instance of <code>torch.nn.Linear</code>.
+Might not work with every architecture, but is tested with all PyTorch
+torchvision classification models (besides SqueezeNet, which has no
+linear last layer).</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>one batch of data to use as input for the forward pass</dd>
+</dl></div>
+</dd>
+</dl>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace" href="index.html">laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.feature_extractor.FeatureExtractor" href="#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.feature_extractor.FeatureExtractor.forward" href="#laplace.feature_extractor.FeatureExtractor.forward">forward</a></code></li>
+<li><code><a title="laplace.feature_extractor.FeatureExtractor.forward_with_features" href="#laplace.feature_extractor.FeatureExtractor.forward_with_features">forward_with_features</a></code></li>
+<li><code><a title="laplace.feature_extractor.FeatureExtractor.set_last_layer" href="#laplace.feature_extractor.FeatureExtractor.set_last_layer">set_last_layer</a></code></li>
+<li><code><a title="laplace.feature_extractor.FeatureExtractor.find_last_layer" href="#laplace.feature_extractor.FeatureExtractor.find_last_layer">find_last_layer</a></code></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/index.html b/docs/index.html
new file mode 100644
index 00000000..fdf416d1
--- /dev/null
+++ b/docs/index.html
@@ -0,0 +1,1073 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace API documentation</title>
+<meta name="description" content="Laplace …" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Package <code>laplace</code></h1>
+</header>
+<section id="section-intro">
+<h1 id="laplace">Laplace</h1>
+<p>Main: <a href="https://travis-ci.com/AlexImmer/Laplace"><img alt="Main" src="https://travis-ci.com/AlexImmer/Laplace.svg?token=rpuRxEjQS6cCZi7ptL9y&amp;branch=main"></a>
+Dev: <a href="https://travis-ci.com/AlexImmer/Laplace"><img alt="Dev" src="https://travis-ci.com/AlexImmer/Laplace.svg?token=rpuRxEjQS6cCZi7ptL9y&amp;branch=dev"></a></p>
+<p>The documentation is available at <a href="https://aleximmer.github.io/Laplace/index.html">https://aleximmer.github.io/Laplace/index.html</a>.</p>
+<h2 id="setup">Setup</h2>
+<p>We assume <code>python3.8</code> since the package was developed with that version.</p>
+<pre><code class="language-bash">pip install -r requirements.txt
+# for development
+pip install -e .
+# for &quot;production&quot;
+pip install .
+
+
+# run tests
+pip install -r tests/requirements.txt
+pytest tests/
+</code></pre>
+<h2 id="structure">Structure</h2>
+<p>The laplace package consists of two main components: </p>
+<ol>
+<li>The subclasses of <a href="laplace/baselaplace.py"><code>laplace.BaseLaplace</code></a> that implement different sparsity structures: different subsets of weights (<code>'all'</code> and <code>'last_layer'</code>) and different structures of the Hessian approximation (<code>'full'</code>, <code>'kron'</code>, and <code>'diag'</code>). This results in six currently available options: <code><a title="laplace.FullLaplace" href="#laplace.FullLaplace">FullLaplace</a></code>, <code><a title="laplace.KronLaplace" href="#laplace.KronLaplace">KronLaplace</a></code>, <code><a title="laplace.DiagLaplace" href="#laplace.DiagLaplace">DiagLaplace</a></code>, and the corresponding last-layer variations <code><a title="laplace.FullLLLaplace" href="#laplace.FullLLLaplace">FullLLLaplace</a></code>, <code><a title="laplace.KronLLLaplace" href="#laplace.KronLLLaplace">KronLLLaplace</a></code>,
+and <code><a title="laplace.DiagLLLaplace" href="#laplace.DiagLLLaplace">DiagLLLaplace</a></code>, which are all subclasses of <a href="laplace/lllaplace.py"><code>laplace.LLLaplace</code></a>. All of these can be conveniently accessed via the <a href="laplace/laplace.py"><code>laplace.Laplace</code></a> function.</li>
+<li>The backends in <a href="laplace/curvature/"><code>laplace.curvature</code></a> which provide access to Hessian approximations of
+the corresponding sparsity structures, for example, the diagonal GGN.</li>
+</ol>
+<p>Additionally, the package provides utilities for
+decomposing a neural network into feature extractor and last layer for <code><a title="laplace.LLLaplace" href="#laplace.LLLaplace">LLLaplace</a></code> subclasses (<a href="laplace/feature_extractor.py"><code>laplace.feature_extractor</code></a>)
+and
+effectively dealing with Kronecker factors (<a href="laplace/matrix.py"><code>laplace.matrix</code></a>).</p>
+<h2 id="extendability">Extendability</h2>
+<p>To extend the laplace package, new <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> subclasses can be designed, for example,
+a block-diagonal structure or subset-of-weights Laplace.
+Alternatively, extending or integrating backends (subclasses of <a href="laplace/curvature/curvature.py"><code>curvature.curvature</code></a>) allows to provide different Hessian
+approximations to the Laplace approximations.
+For example, currently the <a href="laplace/curvature/backpack.py"><code>curvature.BackPackInterface</code></a> based on <a href="https://github.com/f-dangel/backpack/">BackPACK</a> and <a href="laplace/curvature/asdl.py"><code>curvature.AsdlInterface</code></a> based on <a href="https://github.com/kazukiosawa/asdfghjkl">ASDL</a> are available.
+The <code><a title="laplace.curvature.AsdlInterface" href="curvature/index.html#laplace.curvature.AsdlInterface">AsdlInterface</a></code> provides a Kronecker factored empirical Fisher while the <code><a title="laplace.curvature.BackPackInterface" href="curvature/index.html#laplace.curvature.BackPackInterface">BackPackInterface</a></code>
+does not, and only the <code><a title="laplace.curvature.BackPackInterface" href="curvature/index.html#laplace.curvature.BackPackInterface">BackPackInterface</a></code> provides access to Hessian approximations
+for a regression (MSELoss) loss function.</p>
+<h2 id="example-usage">Example usage</h2>
+<h3 id="post-hoc-prior-precision-tuning-of-last-layer-la"><em>Post-hoc</em> prior precision tuning of last-layer LA</h3>
+<p>In the following example, a pre-trained model is loaded,
+then the Laplace approximation is fit to the training data,
+and the prior precision is optimized with cross-validation <code>'CV'</code>.
+After that, the resulting LA is used for prediction with
+the <code>'probit'</code> predictive for classification. </p>
+<pre><code class="language-python">from laplace import Laplace
+
+# pre-trained model
+model = load_map_model()  
+
+# User-specified LA flavor
+la = Laplace(model, 'classification',
+             subset_of_weights='all', 
+             hessian_structure='diag')
+la.fit(train_loader)
+la.optimize_prior_precision(method='CV', val_loader=val_loader)
+
+# User-specified predictive approx.
+pred = la(x, link_approx='probit')
+</code></pre>
+<h3 id="differentiating-the-log-marginal-likelihood-wrt-hyperparameters">Differentiating the log marginal likelihood w.r.t. hyperparameters</h3>
+<p>The marginal likelihood can be used for model selection and is differentiable
+for continuous hyperparameters like the prior precision or observation noise.
+Here, we fit the library default, KFAC last-layer LA and differentiate
+the log marginal likelihood.</p>
+<pre><code class="language-python">from laplace import Laplace
+
+# Un- or pre-trained model
+model = load_model()  
+
+# Default to recommended last-layer KFAC LA:
+la = Laplace(model, likelihood='regression')
+la.fit(train_loader)
+
+# ML w.r.t. prior precision and observation noise
+ml = la.log_marginal_likelihood(prior_prec, obs_noise)
+ml.backward()
+</code></pre>
+<h2 id="documentation"><a href="https://aleximmer.github.io/Laplace/index.html">Documentation</a></h2>
+<pre><code class="language-bash">pip install pdoc3 matplotlib
+# create docs and write to html
+bash update_docs.sh
+# .. or serve the docs directly
+pdoc --http 0.0.0.0:8080 laplace --template-dir template
+</code></pre>
+<h3 id="full-example-post-hoc-optimization-of-the-marginal-likelihood-and-prediction">Full example: <em>post-hoc</em> optimization of the marginal likelihood and prediction</h3>
+<h4 id="sinusoidal-toy-data">Sinusoidal toy data</h4>
+<p>We show how the marginal likelihood can be used after training a MAP network on a simple sinusoidal regression task.
+Subsequently, we use the optimized LA to predict which provides uncertainty on top of the MAP prediction.
+First, we set up the training data for the problem with observation noise \(\sigma=0.3\):</p>
+<pre><code class="language-python">import numpy as np
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from laplace import Laplace
+
+n_epochs = 1000
+batch_size = 150  # full batch
+true_sigma_noise = 0.3
+
+# create simple sinusoid data set
+X_train = (torch.rand(150) * 8).unsqueeze(-1)
+y_train = torch.sin(X_train) + torch.randn_like(X_train) * true_sigma_noise
+train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=batch_size)
+X_test = torch.linspace(-5, 13, 500).unsqueeze(-1)  # +-5 on top of the training X-range
+</code></pre>
+<h4 id="training-a-map">Training a MAP</h4>
+<p>We now use <code>pytorch</code> to train a neural network with single hidden layer and Tanh activation.
+This is standard so nothing new here, yet:</p>
+<pre><code class="language-python"># create and train MAP model
+model = torch.nn.Sequential(torch.nn.Linear(1, 50),
+                            torch.nn.Tanh(),
+                            torch.nn.Linear(50, 1))
+
+criterion = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), weight_decay=5e-4, lr=1e-2)
+for i in range(n_epochs):
+    for X, y in train_loader:
+        optimizer.zero_grad()
+        loss = criterion(model(X), y)
+        loss.backward()
+        optimizer.step()
+</code></pre>
+<h4 id="fitting-and-optimizing-the-laplace-approximation-using-empirical-bayes">Fitting and optimizing the Laplace approximation using empirical Bayes</h4>
+<p>With the MAP-trained model at hand, we can estimate the prior precision and observation noise
+using empirical Bayes after training.
+The <code><a title="laplace.Laplace" href="#laplace.Laplace">Laplace()</a></code> method is called to construct a LA for <code>'regression'</code> with <code>'all'</code> weights.
+As default <code><a title="laplace.Laplace" href="#laplace.Laplace">Laplace()</a></code> returns a Kronecker factored LA.
+We fit the LA to the training data and initialize <code>log_prior</code> and <code>log_sigma</code>.
+Using Adam, we minimize the negative log marginal likelihood for <code>n_epochs</code>.</p>
+<pre><code class="language-python">la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
+la.fit(train_loader)
+log_prior, log_sigma = torch.ones(1, requires_grad=True), torch.ones(1, requires_grad=True)
+hyper_optimizer = torch.optim.Adam([log_prior, log_sigma], lr=1e-1)
+for i in range(n_epochs):
+    hyper_optimizer.zero_grad()
+    neg_marglik = - la.log_marginal_likelihood(log_prior.exp(), log_sigma.exp())
+    neg_marglik.backward()
+    hyper_optimizer.step()
+</code></pre>
+<p>The obtained observation noise is close to the ground truth with a value of \(\sigma \approx 0.28\)
+without the need for any validation data.
+The resulting prior precision is \(\delta \approx 0.18\).</p>
+<h4 id="bayesian-predictive">Bayesian predictive</h4>
+<p>Lastly, we compare the MAP prediction to the obtained LA prediction.
+For LA, we have a closed-form predictive distribution on the output \(f\) which is a Gaussian
+\(\mathcal{N}(f(x;\theta_{MAP}), \mathbb{V}[f] + \sigma^2)\):</p>
+<pre><code class="language-python">x = X_test.flatten().cpu().numpy()
+f_mu, f_var = la(X_test)
+f_mu = f_mu.squeeze().detach().cpu().numpy()
+f_sigma = f_var.squeeze().sqrt().cpu().numpy()
+pred_std = np.sqrt(f_sigma**2 + la.sigma_noise.item()**2)
+</code></pre>
+<p>In comparison to the MAP, the predictive shows useful uncertainties:</p>
+<pre><code class="language-python">import matplotlib.pyplot as plt
+fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, sharey=True,
+                               figsize=(4.5, 2.8))
+ax1.set_title('MAP')
+ax1.scatter(X_train.flatten(), y_train.flatten(), alpha=0.7, color='tab:orange')
+ax1.plot(x, f_mu, color='black', label='$f_{MAP}$')
+ax1.legend()
+
+ax2.set_title('LA')
+ax2.scatter(X_train.flatten(), y_train.flatten(), alpha=0.7, color='tab:orange')
+ax2.plot(x, f_mu, label='$\mathbb{E}[f]$')
+ax2.fill_between(x, f_mu-pred_std*2, f_mu+pred_std*2, 
+                 alpha=0.3, color='tab:blue', label='$2\sqrt{\mathbb{V}\,[f]}$')
+ax2.legend()
+ax1.set_ylim([-4, 6])
+ax1.set_xlim([x.min(), x.max()])
+ax2.set_xlim([x.min(), x.max()])
+ax1.set_ylabel('$y$')
+ax1.set_xlabel('$x$')
+ax2.set_xlabel('$x$')
+plt.tight_layout()
+plt.show()
+</code></pre>
+<p><img alt=":align: center" src="regression_example.png"></p>
+</section>
+<section>
+<h2 class="section-title" id="header-submodules">Sub-modules</h2>
+<dl>
+<dt><code class="name"><a title="laplace.baselaplace" href="baselaplace.html">laplace.baselaplace</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.curvature" href="curvature/index.html">laplace.curvature</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.feature_extractor" href="feature_extractor.html">laplace.feature_extractor</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.laplace" href="laplace.html">laplace.laplace</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.lllaplace" href="lllaplace.html">laplace.lllaplace</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.matrix" href="matrix.html">laplace.matrix</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt><code class="name"><a title="laplace.utils" href="utils.html">laplace.utils</a></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+</dl>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-functions">Functions</h2>
+<dl>
+<dt id="laplace.Laplace"><code class="name flex">
+<span>def <span class="ident">Laplace</span></span>(<span>model, likelihood, subset_of_weights='last_layer', hessian_structure='kron', *args, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Simplified Laplace access using strings instead of different classes.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>subset_of_weights</code></strong> :&ensp;<code>{'last_layer', 'all'}</code>, default=<code>'last_layer'</code></dt>
+<dd>subset of weights to consider for inference</dd>
+<dt><strong><code>hessian_structure</code></strong> :&ensp;<code>{'diag', 'kron', 'full'}</code>, default=<code>'kron'</code></dt>
+<dd>structure of the Hessian approximation</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>laplace</code></strong> :&ensp;<code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code></dt>
+<dd>chosen subclass of BaseLaplace instantiated with additional arguments</dd>
+</dl></div>
+</dd>
+</dl>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.BaseLaplace"><code class="flex name class">
+<span>class <span class="ident">BaseLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Baseclass for all Laplace approximations in this library.
+Subclasses need to specify how the Hessian approximation is initialized,
+how to add up curvature over training data, how to sample from the
+Laplace approximation, and how to compute the functional variance.</p>
+<p>A Laplace approximation is represented by a MAP which is given by the
+<code>model</code> parameter and a posterior precision or covariance specifying
+a Gaussian distribution <span><span class="MathJax_Preview">\mathcal{N}(\theta_{MAP}, P^{-1})</span><script type="math/tex">\mathcal{N}(\theta_{MAP}, P^{-1})</script></span>.
+The goal of this class is to compute the posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>
+which sums as
+<span><span class="MathJax_Preview">
+P = \sum_{n=1}^N \nabla^2_\theta \log p(\mathcal{D}_n \mid \theta)
+\vert_{\theta_{MAP}} + \nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}}.
+</span><script type="math/tex; mode=display">
+P = \sum_{n=1}^N \nabla^2_\theta \log p(\mathcal{D}_n \mid \theta)
+\vert_{\theta_{MAP}} + \nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}}.
+</script></span>
+Every subclass implements different approximations to the log likelihood Hessians,
+for example, a diagonal one. The prior is assumed to be Gaussian and therefore we have
+a simple form for <span><span class="MathJax_Preview">\nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}} = P_0 </span><script type="math/tex">\nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}} = P_0 </script></span>.
+In particular, we assume a scalar, layer-wise, or diagonal prior precision so that in
+all cases <span><span class="MathJax_Preview">P_0 = \textrm{diag}(p_0)</span><script type="math/tex">P_0 = \textrm{diag}(p_0)</script></span> and the structure of <span><span class="MathJax_Preview">p_0</span><script type="math/tex">p_0</script></span> can be varied.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>determines the log likelihood Hessian approximation</dd>
+<dt><strong><code>sigma_noise</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>1</code></dt>
+<dd>observation noise for the regression setting; must be 1 for classification</dd>
+<dt><strong><code>prior_precision</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>1</code></dt>
+<dd>prior precision of a Gaussian prior (= weight decay);
+can be scalar, per-layer, or diagonal in the most general case</dd>
+<dt><strong><code>prior_mean</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>0</code></dt>
+<dd>prior mean of a Gaussian prior, useful for continual learning</dd>
+<dt><strong><code>temperature</code></strong> :&ensp;<code>float</code>, default=<code>1</code></dt>
+<dd>temperature of the likelihood; lower temperature leads to more
+concentrated posterior and vice versa.</dd>
+<dt><strong><code>backend</code></strong> :&ensp;<code>subclasses</code> of <code><a title="laplace.curvature.CurvatureInterface" href="curvature/index.html#laplace.curvature.CurvatureInterface">CurvatureInterface</a></code></dt>
+<dd>backend for access to curvature/Hessian approximations</dd>
+<dt><strong><code>backend_kwargs</code></strong> :&ensp;<code>dict</code>, default=<code>None</code></dt>
+<dd>arguments passed to the backend on initialization, for example to
+set the number of MC samples for stochastic approximations.</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.DiagLaplace" href="baselaplace.html#laplace.baselaplace.DiagLaplace">DiagLaplace</a></li>
+<li><a title="laplace.baselaplace.FullLaplace" href="baselaplace.html#laplace.baselaplace.FullLaplace">FullLaplace</a></li>
+<li><a title="laplace.baselaplace.KronLaplace" href="baselaplace.html#laplace.baselaplace.KronLaplace">KronLaplace</a></li>
+<li>laplace.lllaplace.LLLaplace</li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.BaseLaplace.backend"><code class="name">var <span class="ident">backend</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.BaseLaplace.log_likelihood"><code class="name">var <span class="ident">log_likelihood</span></code></dt>
+<dd>
+<div class="desc"><p>Compute log likelihood on the training data after <code>.fit()</code> has been called.
+The log likelihood is computed on-demand based on the loss and, for example,
+the observation noise which makes it differentiable in the latter for
+iterative updates.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_likelihood</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.scatter"><code class="name">var <span class="ident">scatter</span></code></dt>
+<dd>
+<div class="desc"><p>Computes the <em>scatter</em>, a term of the log marginal likelihood that
+corresponds to L-2 regularization:
+<code>scatter</code> = <span><span class="MathJax_Preview">(\theta_{MAP} - \mu_0)^{T} P_0 (\theta_{MAP} - \mu_0) </span><script type="math/tex">(\theta_{MAP} - \mu_0)^{T} P_0 (\theta_{MAP} - \mu_0) </script></span>.</p>
+<h2 id="returns">Returns</h2>
+<p>[type]
+[description]</p></div>
+</dd>
+<dt id="laplace.BaseLaplace.log_det_prior_precision"><code class="name">var <span class="ident">log_det_prior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Compute log determinant of the prior precision
+<span><span class="MathJax_Preview">\log \det P_0</span><script type="math/tex">\log \det P_0</script></span></p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_det</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.log_det_posterior_precision"><code class="name">var <span class="ident">log_det_posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Compute log determinant of the posterior precision
+<span><span class="MathJax_Preview">\log \det P</span><script type="math/tex">\log \det P</script></span> which depends on the subclasses structure
+used for the Hessian approximation.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_det</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.log_det_ratio"><code class="name">var <span class="ident">log_det_ratio</span></code></dt>
+<dd>
+<div class="desc"><p>Compute the log determinant ratio, a part of the log marginal likelihood.
+<span><span class="MathJax_Preview">
+\log \frac{\det P}{\det P_0} = \log \det P - \log \det P_0
+</span><script type="math/tex; mode=display">
+\log \frac{\det P}{\det P_0} = \log \det P - \log \det P_0
+</script></span></p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_det_ratio</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.prior_precision_diag"><code class="name">var <span class="ident">prior_precision_diag</span></code></dt>
+<dd>
+<div class="desc"><p>Obtain the diagonal prior precision <span><span class="MathJax_Preview">p_0</span><script type="math/tex">p_0</script></span> constructed from either
+a scalar, layer-wise, or diagonal prior precision.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>prior_precision_diag</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.prior_mean"><code class="name">var <span class="ident">prior_mean</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.BaseLaplace.prior_precision"><code class="name">var <span class="ident">prior_precision</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.BaseLaplace.sigma_noise"><code class="name">var <span class="ident">sigma_noise</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.BaseLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Compute or return the posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>posterior_prec</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.BaseLaplace.fit"><code class="name flex">
+<span>def <span class="ident">fit</span></span>(<span>self, train_loader)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Fit the local Laplace approximation at the parameters of the model.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>train_loader</code></strong> :&ensp;<code>torch.data.utils.DataLoader</code></dt>
+<dd>each iterate is a training batch (X, y);
+<code>train_loader.dataset</code> needs to be set to access <span><span class="MathJax_Preview">N</span><script type="math/tex">N</script></span>, size of the data set</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.log_marginal_likelihood"><code class="name flex">
+<span>def <span class="ident">log_marginal_likelihood</span></span>(<span>self, prior_precision=None, sigma_noise=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute the Laplace approximation to the log marginal likelihood subject
+to specific Hessian approximations that subclasses implement.
+Requires that the Laplace approximation has been fit before.
+The resulting torch.Tensor is differentiable in <code>prior_precision</code> and
+<code>sigma_noise</code> if these have gradients enabled.
+By passing <code>prior_precision</code> or <code>sigma_noise</code>, the current value is
+overwritten. This is useful for iterating on the log marginal likelihood.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>prior_precision</code></strong> :&ensp;<code>torch.Tensor</code>, optional</dt>
+<dd>prior precision if should be changed from current <code>prior_precision</code> value</dd>
+<dt><strong><code>sigma_noise</code></strong> :&ensp;<code>[type]</code>, optional</dt>
+<dd>observation noise standard deviation if should be changed</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>log_marglik</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.predictive"><code class="name flex">
+<span>def <span class="ident">predictive</span></span>(<span>self, x, pred_type='glm', link_approx='mc', n_samples=100)</span>
+</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.BaseLaplace.predictive_samples"><code class="name flex">
+<span>def <span class="ident">predictive_samples</span></span>(<span>self, x, pred_type='glm', n_samples=100)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Sample from the posterior predictive on input data <code>x</code>.
+Can be used, for example, for Thompson sampling.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>x</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>input data <code>(batch_size, input_shape)</code></dd>
+<dt><strong><code>pred_type</code></strong> :&ensp;<code>{'glm', 'nn'}</code>, default=<code>'glm'</code></dt>
+<dd>type of posterior predictive, linearized GLM predictive or neural
+network sampling predictive. The GLM predictive is consistent with
+the curvature approximations used here.</dd>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</code></dt>
+<dd>number of samples</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>samples</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>samples <code>(n_samples, batch_size, output_shape)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.functional_variance"><code class="name flex">
+<span>def <span class="ident">functional_variance</span></span>(<span>self, Jacs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute functional variance for the <code>'glm'</code> predictive:
+<code>f_var[i] = Jacs[i] @ P.inv() @ Jacs[i].T</code>, which is a output x output
+predictive covariance matrix.
+Mathematically, we have for a single Jacobian
+<span><span class="MathJax_Preview">\mathcal{J} = \nabla_\theta f(x;\theta)\vert_{\theta_{MAP}}</span><script type="math/tex">\mathcal{J} = \nabla_\theta f(x;\theta)\vert_{\theta_{MAP}}</script></span>
+the output covariance matrix
+<span><span class="MathJax_Preview"> \mathcal{J} P^{-1} \mathcal{J}^T </span><script type="math/tex"> \mathcal{J} P^{-1} \mathcal{J}^T </script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>Jacs</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>Jacobians of model output wrt parameters
+<code>(batch, outputs, parameters)</code></dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>f_var</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>output covariance <code>(batch, outputs, outputs)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.sample"><code class="name flex">
+<span>def <span class="ident">sample</span></span>(<span>self, n_samples=100)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Sample from the Laplace posterior approximation, i.e.,
+<span><span class="MathJax_Preview"> \theta \sim \mathcal{N}(\theta_{MAP}, P^{-1})</span><script type="math/tex"> \theta \sim \mathcal{N}(\theta_{MAP}, P^{-1})</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>number of samples</dd>
+</dl></div>
+</dd>
+<dt id="laplace.BaseLaplace.optimize_prior_precision"><code class="name flex">
+<span>def <span class="ident">optimize_prior_precision</span></span>(<span>self, method='marglik', n_steps=100, lr=0.1, init_prior_prec=1.0, val_loader=None, loss=&lt;function get_nll&gt;, log_prior_prec_min=-4, log_prior_prec_max=4, grid_size=100, pred_type='glm', link_approx='probit', n_samples=100, verbose=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Optimize the prior precision post-hoc using the <code>method</code>
+specified by the user.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>method</code></strong> :&ensp;<code>{'marglik', 'CV'}</code>, default=<code>'marglik'</code></dt>
+<dd>specifies how the prior precision should be optimized.</dd>
+<dt><strong><code>n_steps</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>the number of gradient descent steps to take.</dd>
+<dt><strong><code>lr</code></strong> :&ensp;<code>float</code>, default=<code>1e-1</code></dt>
+<dd>the learning rate to use for gradient descent.</dd>
+<dt><strong><code>init_prior_prec</code></strong> :&ensp;<code>float</code>, default=<code>1.0</code></dt>
+<dd>initial prior precision before the first optimization step.</dd>
+<dt><strong><code>val_loader</code></strong> :&ensp;<code>torch.data.utils.DataLoader</code>, default=<code>None</code></dt>
+<dd>DataLoader for the validation set; each iterate is a training batch (X, y).</dd>
+<dt><strong><code>loss</code></strong> :&ensp;<code>callable</code>, default=<code>get_nll</code></dt>
+<dd>loss function to use for CV.</dd>
+<dt><strong><code>log_prior_prec_min</code></strong> :&ensp;<code>float</code>, default=<code>-4</code></dt>
+<dd>lower bound of gridsearch interval for CV.</dd>
+<dt><strong><code>log_prior_prec_max</code></strong> :&ensp;<code>float</code>, default=<code>4</code></dt>
+<dd>upper bound of gridsearch interval for CV.</dd>
+<dt><strong><code>grid_size</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>number of values to consider inside the gridsearch interval for CV.</dd>
+<dt><strong><code>pred_type</code></strong> :&ensp;<code>{'glm', 'nn'}</code>, default=<code>'glm'</code></dt>
+<dd>type of posterior predictive, linearized GLM predictive or neural
+network sampling predictive. The GLM predictive is consistent with
+the curvature approximations used here.</dd>
+<dt><strong><code>link_approx</code></strong> :&ensp;<code>{'mc', 'probit', 'bridge'}</code>, default=<code>'probit'</code></dt>
+<dd>how to approximate the classification link function for the <code>'glm'</code>.
+For <code>pred_type='nn'</code>, only <code>'mc'</code> is possible.</dd>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</code>, default=<code>100</code></dt>
+<dd>number of samples for <code>link_approx='mc'</code>.</dd>
+<dt><strong><code>verbose</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>if true, the optimized prior precision will be printed
+(can be a large tensor if the prior has a diagonal covariance).</dd>
+</dl></div>
+</dd>
+</dl>
+</dd>
+<dt id="laplace.FullLaplace"><code class="flex name class">
+<span>class <span class="ident">FullLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Laplace approximation with full, i.e., dense, log likelihood Hessian approximation
+and hence posterior precision. Based on the chosen <code>backend</code> parameter, the full
+approximation can be, for example, a generalized Gauss-Newton matrix.
+Mathematically, we have <span><span class="MathJax_Preview">P \in \mathbb{R}^{P \times P}</span><script type="math/tex">P \in \mathbb{R}^{P \times P}</script></span>.
+See <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.lllaplace.FullLLLaplace" href="lllaplace.html#laplace.lllaplace.FullLLLaplace">FullLLLaplace</a></li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.FullLaplace.posterior_scale"><code class="name">var <span class="ident">posterior_scale</span></code></dt>
+<dd>
+<div class="desc"><p>Posterior scale (square root of the covariance), i.e.,
+<span><span class="MathJax_Preview">P^{-\frac{1}{2}}</span><script type="math/tex">P^{-\frac{1}{2}}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>scale</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.FullLaplace.posterior_covariance"><code class="name">var <span class="ident">posterior_covariance</span></code></dt>
+<dd>
+<div class="desc"><p>Posterior covariance, i.e., <span><span class="MathJax_Preview">P^{-1}</span><script type="math/tex">P^{-1}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>covariance</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.FullLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters, parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.KronLaplace"><code class="flex name class">
+<span>class <span class="ident">KronLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, damping=False, **backend_kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Laplace approximation with Kronecker factored log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have for each parameter group, e.g., torch.nn.Module,
+that \P\approx Q \otimes H.
+See <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> for the full interface and see
+<code><a title="laplace.matrix.Kron" href="matrix.html#laplace.matrix.Kron">Kron</a></code> and <code><a title="laplace.matrix.KronDecomposed" href="matrix.html#laplace.matrix.KronDecomposed">KronDecomposed</a></code> for the structure of
+the Kronecker factors. <code>Kron</code> is used to aggregate factors by summing up and
+<code>KronDecomposed</code> is used to add the prior, a Hessian factor (e.g. temperature),
+and computing posterior covariances, marginal likelihood, etc.
+Damping can be enabled by setting <code>damping=True</code>.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.lllaplace.KronLLLaplace" href="lllaplace.html#laplace.lllaplace.KronLLLaplace">KronLLLaplace</a></li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.KronLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Kronecker factored Posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code><a title="laplace.matrix.KronDecomposed" href="matrix.html#laplace.matrix.KronDecomposed">KronDecomposed</a></code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.KronLaplace.prior_precision"><code class="name">var <span class="ident">prior_precision</span></code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.DiagLaplace"><code class="flex name class">
+<span>class <span class="ident">DiagLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Laplace approximation with diagonal log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have <span><span class="MathJax_Preview">P \approx \textrm{diag}(P)</span><script type="math/tex">P \approx \textrm{diag}(P)</script></span>.
+See <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.lllaplace.DiagLLLaplace" href="lllaplace.html#laplace.lllaplace.DiagLLLaplace">DiagLLLaplace</a></li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.DiagLaplace.posterior_precision"><code class="name">var <span class="ident">posterior_precision</span></code></dt>
+<dd>
+<div class="desc"><p>Diagonal posterior precision <span><span class="MathJax_Preview">p</span><script type="math/tex">p</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.DiagLaplace.posterior_scale"><code class="name">var <span class="ident">posterior_scale</span></code></dt>
+<dd>
+<div class="desc"><p>Diagonal posterior scale <span><span class="MathJax_Preview">\sqrt{p^{-1}}</span><script type="math/tex">\sqrt{p^{-1}}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.DiagLaplace.posterior_variance"><code class="name">var <span class="ident">posterior_variance</span></code></dt>
+<dd>
+<div class="desc"><p>Diagonal posterior variance <span><span class="MathJax_Preview">p^{-1}</span><script type="math/tex">p^{-1}</script></span>.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>precision</code></strong> :&ensp;<code>torch.tensor</code></dt>
+<dd><code>(parameters)</code></dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.LLLaplace"><code class="flex name class">
+<span>class <span class="ident">LLLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, last_layer_name=None, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Baseclass for all last-layer Laplace approximations in this library.
+Subclasses specify the structure of the Hessian approximation.
+See <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> for the full interface.</p>
+<p>A Laplace approximation is represented by a MAP which is given by the
+<code>model</code> parameter and a posterior precision or covariance specifying
+a Gaussian distribution <span><span class="MathJax_Preview">\mathcal{N}(\theta_{MAP}, P^{-1})</span><script type="math/tex">\mathcal{N}(\theta_{MAP}, P^{-1})</script></span>.
+Here, only the parameters of the last layer of the neural network
+are treated probabilistically.
+The goal of this class is to compute the posterior precision <span><span class="MathJax_Preview">P</span><script type="math/tex">P</script></span>
+which sums as
+<span><span class="MathJax_Preview">
+P = \sum_{n=1}^N \nabla^2_\theta \log p(\mathcal{D}_n \mid \theta)
+\vert_{\theta_{MAP}} + \nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}}.
+</span><script type="math/tex; mode=display">
+P = \sum_{n=1}^N \nabla^2_\theta \log p(\mathcal{D}_n \mid \theta)
+\vert_{\theta_{MAP}} + \nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}}.
+</script></span>
+Every subclass implements different approximations to the log likelihood Hessians,
+for example, a diagonal one. The prior is assumed to be Gaussian and therefore we have
+a simple form for <span><span class="MathJax_Preview">\nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}} = P_0 </span><script type="math/tex">\nabla^2_\theta \log p(\theta) \vert_{\theta_{MAP}} = P_0 </script></span>.
+In particular, we assume a scalar or diagonal prior precision so that in
+all cases <span><span class="MathJax_Preview">P_0 = \textrm{diag}(p_0)</span><script type="math/tex">P_0 = \textrm{diag}(p_0)</script></span> and the structure of <span><span class="MathJax_Preview">p_0</span><script type="math/tex">p_0</script></span> can be varied.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code> or <code><a title="laplace.feature_extractor.FeatureExtractor" href="feature_extractor.html#laplace.feature_extractor.FeatureExtractor">FeatureExtractor</a></code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>determines the log likelihood Hessian approximation</dd>
+<dt><strong><code>sigma_noise</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>1</code></dt>
+<dd>observation noise for the regression setting; must be 1 for classification</dd>
+<dt><strong><code>prior_precision</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>1</code></dt>
+<dd>prior precision of a Gaussian prior (= weight decay);
+can be scalar, per-layer, or diagonal in the most general case</dd>
+<dt><strong><code>prior_mean</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code>, default=<code>0</code></dt>
+<dd>prior mean of a Gaussian prior, useful for continual learning</dd>
+<dt><strong><code>temperature</code></strong> :&ensp;<code>float</code>, default=<code>1</code></dt>
+<dd>temperature of the likelihood; lower temperature leads to more
+concentrated posterior and vice versa.</dd>
+<dt><strong><code>backend</code></strong> :&ensp;<code>subclasses</code> of <code><a title="laplace.curvature.CurvatureInterface" href="curvature/index.html#laplace.curvature.CurvatureInterface">CurvatureInterface</a></code></dt>
+<dd>backend for access to curvature/Hessian approximations</dd>
+<dt><strong><code>last_layer_name</code></strong> :&ensp;<code>str</code>, default=<code>None</code></dt>
+<dd>name of the model's last layer, if None it will be determined automatically</dd>
+<dt><strong><code>backend_kwargs</code></strong> :&ensp;<code>dict</code>, default=<code>None</code></dt>
+<dd>arguments passed to the backend on initialization, for example to
+set the number of MC samples for stochastic approximations.</dd>
+</dl></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Subclasses</h3>
+<ul class="hlist">
+<li><a title="laplace.lllaplace.DiagLLLaplace" href="lllaplace.html#laplace.lllaplace.DiagLLLaplace">DiagLLLaplace</a></li>
+<li><a title="laplace.lllaplace.FullLLLaplace" href="lllaplace.html#laplace.lllaplace.FullLLLaplace">FullLLLaplace</a></li>
+<li><a title="laplace.lllaplace.KronLLLaplace" href="lllaplace.html#laplace.lllaplace.KronLLLaplace">KronLLLaplace</a></li>
+</ul>
+<h3>Instance variables</h3>
+<dl>
+<dt id="laplace.LLLaplace.prior_precision_diag"><code class="name">var <span class="ident">prior_precision_diag</span></code></dt>
+<dd>
+<div class="desc"><p>Obtain the diagonal prior precision <span><span class="MathJax_Preview">p_0</span><script type="math/tex">p_0</script></span> constructed from either
+a scalar or diagonal prior precision.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>prior_precision_diag</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.BaseLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.posterior_precision">posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.BaseLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.FullLLLaplace"><code class="flex name class">
+<span>class <span class="ident">FullLLLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, last_layer_name=None, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Last-layer Laplace approximation with full, i.e., dense, log likelihood Hessian approximation
+and hence posterior precision. Based on the chosen <code>backend</code> parameter, the full
+approximation can be, for example, a generalized Gauss-Newton matrix.
+Mathematically, we have <span><span class="MathJax_Preview">P \in \mathbb{R}^{P \times P}</span><script type="math/tex">P \in \mathbb{R}^{P \times P}</script></span>.
+See <code><a title="laplace.FullLaplace" href="#laplace.FullLaplace">FullLaplace</a></code>, <code><a title="laplace.LLLaplace" href="#laplace.LLLaplace">LLLaplace</a></code>, and <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>laplace.lllaplace.LLLaplace</li>
+<li><a title="laplace.baselaplace.FullLaplace" href="baselaplace.html#laplace.baselaplace.FullLaplace">FullLaplace</a></li>
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.FullLaplace" href="baselaplace.html#laplace.baselaplace.FullLaplace">FullLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.FullLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.posterior_covariance" href="baselaplace.html#laplace.baselaplace.FullLaplace.posterior_covariance">posterior_covariance</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.posterior_precision" href="baselaplace.html#laplace.baselaplace.FullLaplace.posterior_precision">posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.posterior_scale" href="baselaplace.html#laplace.baselaplace.FullLaplace.posterior_scale">posterior_scale</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.KronLLLaplace"><code class="flex name class">
+<span>class <span class="ident">KronLLLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, last_layer_name=None, damping=False, **backend_kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Last-layer Laplace approximation with Kronecker factored log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have for the last parameter group, i.e., torch.nn.Linear,
+that \P\approx Q \otimes H.
+See <code><a title="laplace.KronLaplace" href="#laplace.KronLaplace">KronLaplace</a></code>, <code><a title="laplace.LLLaplace" href="#laplace.LLLaplace">LLLaplace</a></code>, and <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> for the full interface and see
+<code><a title="laplace.matrix.Kron" href="matrix.html#laplace.matrix.Kron">Kron</a></code> and <code><a title="laplace.matrix.KronDecomposed" href="matrix.html#laplace.matrix.KronDecomposed">KronDecomposed</a></code> for the structure of
+the Kronecker factors. <code>Kron</code> is used to aggregate factors by summing up and
+<code>KronDecomposed</code> is used to add the prior, a Hessian factor (e.g. temperature),
+and computing posterior covariances, marginal likelihood, etc.
+Use of <code>damping</code> is possible by initializing or setting <code>damping=True</code>.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>laplace.lllaplace.LLLaplace</li>
+<li><a title="laplace.baselaplace.KronLaplace" href="baselaplace.html#laplace.baselaplace.KronLaplace">KronLaplace</a></li>
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.KronLaplace" href="baselaplace.html#laplace.baselaplace.KronLaplace">KronLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.KronLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.posterior_precision" href="baselaplace.html#laplace.baselaplace.KronLaplace.posterior_precision">posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.DiagLLLaplace"><code class="flex name class">
+<span>class <span class="ident">DiagLLLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, last_layer_name=None, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Last-layer Laplace approximation with diagonal log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have <span><span class="MathJax_Preview">P \approx \textrm{diag}(P)</span><script type="math/tex">P \approx \textrm{diag}(P)</script></span>.
+See <code><a title="laplace.DiagLaplace" href="#laplace.DiagLaplace">DiagLaplace</a></code>, <code><a title="laplace.LLLaplace" href="#laplace.LLLaplace">LLLaplace</a></code>, and <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>laplace.lllaplace.LLLaplace</li>
+<li><a title="laplace.baselaplace.DiagLaplace" href="baselaplace.html#laplace.baselaplace.DiagLaplace">DiagLaplace</a></li>
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.DiagLaplace" href="baselaplace.html#laplace.baselaplace.DiagLaplace">DiagLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.DiagLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.posterior_precision" href="baselaplace.html#laplace.baselaplace.DiagLaplace.posterior_precision">posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.posterior_scale" href="baselaplace.html#laplace.baselaplace.DiagLaplace.posterior_scale">posterior_scale</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.posterior_variance" href="baselaplace.html#laplace.baselaplace.DiagLaplace.posterior_variance">posterior_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul>
+<li><a href="#laplace">Laplace</a><ul>
+<li><a href="#setup">Setup</a></li>
+<li><a href="#structure">Structure</a></li>
+<li><a href="#extendability">Extendability</a></li>
+<li><a href="#example-usage">Example usage</a><ul>
+<li><a href="#post-hoc-prior-precision-tuning-of-last-layer-la">Post-hoc prior precision tuning of last-layer LA</a></li>
+<li><a href="#differentiating-the-log-marginal-likelihood-wrt-hyperparameters">Differentiating the log marginal likelihood w.r.t. hyperparameters</a></li>
+</ul>
+</li>
+<li><a href="#documentation">Documentation</a><ul>
+<li><a href="#full-example-post-hoc-optimization-of-the-marginal-likelihood-and-prediction">Full example: post-hoc optimization of the marginal likelihood and prediction</a><ul>
+<li><a href="#sinusoidal-toy-data">Sinusoidal toy data</a></li>
+<li><a href="#training-a-map">Training a MAP</a></li>
+<li><a href="#fitting-and-optimizing-the-laplace-approximation-using-empirical-bayes">Fitting and optimizing the Laplace approximation using empirical Bayes</a></li>
+<li><a href="#bayesian-predictive">Bayesian predictive</a></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+</div>
+<ul id="index">
+<li><h3><a href="#header-submodules">Sub-modules</a></h3>
+<ul>
+<li><code><a title="laplace.baselaplace" href="baselaplace.html">laplace.baselaplace</a></code></li>
+<li><code><a title="laplace.curvature" href="curvature/index.html">laplace.curvature</a></code></li>
+<li><code><a title="laplace.feature_extractor" href="feature_extractor.html">laplace.feature_extractor</a></code></li>
+<li><code><a title="laplace.laplace" href="laplace.html">laplace.laplace</a></code></li>
+<li><code><a title="laplace.lllaplace" href="lllaplace.html">laplace.lllaplace</a></code></li>
+<li><code><a title="laplace.matrix" href="matrix.html">laplace.matrix</a></code></li>
+<li><code><a title="laplace.utils" href="utils.html">laplace.utils</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-functions">Functions</a></h3>
+<ul class="">
+<li><code><a title="laplace.Laplace" href="#laplace.Laplace">Laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.BaseLaplace.fit" href="#laplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.BaseLaplace.log_marginal_likelihood" href="#laplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.BaseLaplace.predictive" href="#laplace.BaseLaplace.predictive">predictive</a></code></li>
+<li><code><a title="laplace.BaseLaplace.predictive_samples" href="#laplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.BaseLaplace.functional_variance" href="#laplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.BaseLaplace.sample" href="#laplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.BaseLaplace.optimize_prior_precision" href="#laplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.FullLaplace" href="#laplace.FullLaplace">FullLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.KronLaplace" href="#laplace.KronLaplace">KronLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.DiagLaplace" href="#laplace.DiagLaplace">DiagLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.LLLaplace" href="#laplace.LLLaplace">LLLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.FullLLLaplace" href="#laplace.FullLLLaplace">FullLLLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.KronLLLaplace" href="#laplace.KronLLLaplace">KronLLLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.DiagLLLaplace" href="#laplace.DiagLLLaplace">DiagLLLaplace</a></code></h4>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/laplace.html b/docs/laplace.html
new file mode 100644
index 00000000..cc38555f
--- /dev/null
+++ b/docs/laplace.html
@@ -0,0 +1,84 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.laplace API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.laplace</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-functions">Functions</h2>
+<dl>
+<dt id="laplace.laplace.Laplace"><code class="name flex">
+<span>def <span class="ident">Laplace</span></span>(<span>model, likelihood, subset_of_weights='last_layer', hessian_structure='kron', *args, **kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Simplified Laplace access using strings instead of different classes.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>likelihood</code></strong> :&ensp;<code>{'classification', 'regression'}</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>subset_of_weights</code></strong> :&ensp;<code>{'last_layer', 'all'}</code>, default=<code>'last_layer'</code></dt>
+<dd>subset of weights to consider for inference</dd>
+<dt><strong><code>hessian_structure</code></strong> :&ensp;<code>{'diag', 'kron', 'full'}</code>, default=<code>'kron'</code></dt>
+<dd>structure of the Hessian approximation</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>laplace</code></strong> :&ensp;<code>BaseLaplace</code></dt>
+<dd>chosen subclass of BaseLaplace instantiated with additional arguments</dd>
+</dl></div>
+</dd>
+</dl>
+</section>
+<section>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace" href="index.html">laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-functions">Functions</a></h3>
+<ul class="">
+<li><code><a title="laplace.laplace.Laplace" href="#laplace.laplace.Laplace">Laplace</a></code></li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/lllaplace.html b/docs/lllaplace.html
new file mode 100644
index 00000000..36df7d3d
--- /dev/null
+++ b/docs/lllaplace.html
@@ -0,0 +1,192 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.lllaplace API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.lllaplace</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.lllaplace.FullLLLaplace"><code class="flex name class">
+<span>class <span class="ident">FullLLLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, last_layer_name=None, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Last-layer Laplace approximation with full, i.e., dense, log likelihood Hessian approximation
+and hence posterior precision. Based on the chosen <code>backend</code> parameter, the full
+approximation can be, for example, a generalized Gauss-Newton matrix.
+Mathematically, we have <span><span class="MathJax_Preview">P \in \mathbb{R}^{P \times P}</span><script type="math/tex">P \in \mathbb{R}^{P \times P}</script></span>.
+See <code>FullLaplace</code>, <code>LLLaplace</code>, and <code>BaseLaplace</code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>laplace.lllaplace.LLLaplace</li>
+<li><a title="laplace.baselaplace.FullLaplace" href="baselaplace.html#laplace.baselaplace.FullLaplace">FullLaplace</a></li>
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.FullLaplace" href="baselaplace.html#laplace.baselaplace.FullLaplace">FullLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.FullLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.posterior_covariance" href="baselaplace.html#laplace.baselaplace.FullLaplace.posterior_covariance">posterior_covariance</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.posterior_precision" href="baselaplace.html#laplace.baselaplace.FullLaplace.posterior_precision">posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.posterior_scale" href="baselaplace.html#laplace.baselaplace.FullLaplace.posterior_scale">posterior_scale</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.FullLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.lllaplace.KronLLLaplace"><code class="flex name class">
+<span>class <span class="ident">KronLLLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, last_layer_name=None, damping=False, **backend_kwargs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Last-layer Laplace approximation with Kronecker factored log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have for the last parameter group, i.e., torch.nn.Linear,
+that \P\approx Q \otimes H.
+See <code>KronLaplace</code>, <code>LLLaplace</code>, and <code>BaseLaplace</code> for the full interface and see
+<code><a title="laplace.matrix.Kron" href="matrix.html#laplace.matrix.Kron">Kron</a></code> and <code><a title="laplace.matrix.KronDecomposed" href="matrix.html#laplace.matrix.KronDecomposed">KronDecomposed</a></code> for the structure of
+the Kronecker factors. <code>Kron</code> is used to aggregate factors by summing up and
+<code>KronDecomposed</code> is used to add the prior, a Hessian factor (e.g. temperature),
+and computing posterior covariances, marginal likelihood, etc.
+Use of <code>damping</code> is possible by initializing or setting <code>damping=True</code>.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>laplace.lllaplace.LLLaplace</li>
+<li><a title="laplace.baselaplace.KronLaplace" href="baselaplace.html#laplace.baselaplace.KronLaplace">KronLaplace</a></li>
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.KronLaplace" href="baselaplace.html#laplace.baselaplace.KronLaplace">KronLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.KronLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.posterior_precision" href="baselaplace.html#laplace.baselaplace.KronLaplace.posterior_precision">posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.KronLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+<dt id="laplace.lllaplace.DiagLLLaplace"><code class="flex name class">
+<span>class <span class="ident">DiagLLLaplace</span></span>
+<span>(</span><span>model, likelihood, sigma_noise=1.0, prior_precision=1.0, prior_mean=0.0, temperature=1.0, backend=laplace.curvature.backpack.BackPackGGN, last_layer_name=None, backend_kwargs=None)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Last-layer Laplace approximation with diagonal log likelihood Hessian approximation
+and hence posterior precision.
+Mathematically, we have <span><span class="MathJax_Preview">P \approx \textrm{diag}(P)</span><script type="math/tex">P \approx \textrm{diag}(P)</script></span>.
+See <code>DiagLaplace</code>, <code>LLLaplace</code>, and <code>BaseLaplace</code> for the full interface.</p></div>
+<h3>Ancestors</h3>
+<ul class="hlist">
+<li>laplace.lllaplace.LLLaplace</li>
+<li><a title="laplace.baselaplace.DiagLaplace" href="baselaplace.html#laplace.baselaplace.DiagLaplace">DiagLaplace</a></li>
+<li><a title="laplace.baselaplace.BaseLaplace" href="baselaplace.html#laplace.baselaplace.BaseLaplace">BaseLaplace</a></li>
+<li>abc.ABC</li>
+</ul>
+<h3>Inherited members</h3>
+<ul class="hlist">
+<li><code><b><a title="laplace.baselaplace.DiagLaplace" href="baselaplace.html#laplace.baselaplace.DiagLaplace">DiagLaplace</a></b></code>:
+<ul class="hlist">
+<li><code><a title="laplace.baselaplace.DiagLaplace.fit" href="baselaplace.html#laplace.baselaplace.BaseLaplace.fit">fit</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.functional_variance" href="baselaplace.html#laplace.baselaplace.BaseLaplace.functional_variance">functional_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_det_posterior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_posterior_precision">log_det_posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_det_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_prior_precision">log_det_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_det_ratio" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_det_ratio">log_det_ratio</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_likelihood">log_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.log_marginal_likelihood" href="baselaplace.html#laplace.baselaplace.BaseLaplace.log_marginal_likelihood">log_marginal_likelihood</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.optimize_prior_precision" href="baselaplace.html#laplace.baselaplace.BaseLaplace.optimize_prior_precision">optimize_prior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.posterior_precision" href="baselaplace.html#laplace.baselaplace.DiagLaplace.posterior_precision">posterior_precision</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.posterior_scale" href="baselaplace.html#laplace.baselaplace.DiagLaplace.posterior_scale">posterior_scale</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.posterior_variance" href="baselaplace.html#laplace.baselaplace.DiagLaplace.posterior_variance">posterior_variance</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.predictive_samples" href="baselaplace.html#laplace.baselaplace.BaseLaplace.predictive_samples">predictive_samples</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.prior_precision_diag" href="baselaplace.html#laplace.baselaplace.BaseLaplace.prior_precision_diag">prior_precision_diag</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.sample" href="baselaplace.html#laplace.baselaplace.BaseLaplace.sample">sample</a></code></li>
+<li><code><a title="laplace.baselaplace.DiagLaplace.scatter" href="baselaplace.html#laplace.baselaplace.BaseLaplace.scatter">scatter</a></code></li>
+</ul>
+</li>
+</ul>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace" href="index.html">laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.lllaplace.FullLLLaplace" href="#laplace.lllaplace.FullLLLaplace">FullLLLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.lllaplace.KronLLLaplace" href="#laplace.lllaplace.KronLLLaplace">KronLLLaplace</a></code></h4>
+</li>
+<li>
+<h4><code><a title="laplace.lllaplace.DiagLLLaplace" href="#laplace.lllaplace.DiagLLLaplace">DiagLLLaplace</a></code></h4>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/matrix.html b/docs/matrix.html
new file mode 100644
index 00000000..a2b6d565
--- /dev/null
+++ b/docs/matrix.html
@@ -0,0 +1,287 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.matrix API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.matrix</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-classes">Classes</h2>
+<dl>
+<dt id="laplace.matrix.Kron"><code class="flex name class">
+<span>class <span class="ident">Kron</span></span>
+<span>(</span><span>kfacs)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Kronecker factored approximate curvature representation for a corresponding
+neural network.
+Each element in <code>kfacs</code> is either a tuple or single matrix.
+A tuple represents two Kronecker factors <span><span class="MathJax_Preview">Q</span><script type="math/tex">Q</script></span>, and <span><span class="MathJax_Preview">H</span><script type="math/tex">H</script></span> and a single element
+is just a full block Hessian approximation.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>kfacs</code></strong> :&ensp;<code>list[Tuple]</code></dt>
+<dd>each element in the list is a Tuple of two Kronecker factors Q, H
+or a single matrix approximating the Hessian (in case of bias, for example)</dd>
+</dl></div>
+<h3>Static methods</h3>
+<dl>
+<dt id="laplace.matrix.Kron.init_from_model"><code class="name flex">
+<span>def <span class="ident">init_from_model</span></span>(<span>model, device)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Initialize Kronecker factors based on a models architecture.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>device</code></strong> :&ensp;<code>torch.device</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>kron</code></strong> :&ensp;<code><a title="laplace.matrix.Kron" href="#laplace.matrix.Kron">Kron</a></code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+</dl>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.matrix.Kron.decompose"><code class="name flex">
+<span>def <span class="ident">decompose</span></span>(<span>self, damping=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Eigendecompose Kronecker factors and turn into <code><a title="laplace.matrix.KronDecomposed" href="#laplace.matrix.KronDecomposed">KronDecomposed</a></code>.
+Parameters</p>
+<hr>
+<dl>
+<dt><strong><code>damping</code></strong> :&ensp;<code>bool</code></dt>
+<dd>use damping</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>kron_decomposed</code></strong> :&ensp;<code><a title="laplace.matrix.KronDecomposed" href="#laplace.matrix.KronDecomposed">KronDecomposed</a></code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.matrix.Kron.bmm"><code class="name flex">
+<span>def <span class="ident">bmm</span></span>(<span>self, W: torch.Tensor, exponent: float = 1) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Batched matrix multiplication with the Kronecker factors.
+If Kron is <code>H</code>, we compute <code>H @ W</code>.
+This is useful for computing the predictive or a regularization
+based on Kronecker factors as in continual learning.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>W</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>matrix <code>(batch, classes, params)</code></dd>
+<dt><strong><code>exponent</code></strong> :&ensp;<code>float</code>, default=<code>1</code></dt>
+<dd>only can be <code>1</code> for Kron, requires <code><a title="laplace.matrix.KronDecomposed" href="#laplace.matrix.KronDecomposed">KronDecomposed</a></code> for other
+exponent values of the Kronecker factors.</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>SW</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>result <code>(batch, classes, params)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.matrix.Kron.logdet"><code class="name flex">
+<span>def <span class="ident">logdet</span></span>(<span>self) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute log determinant of the Kronecker factors and sums them up.
+This corresponds to the log determinant of the entire Hessian approximation.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>logdet</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.matrix.Kron.diag"><code class="name flex">
+<span>def <span class="ident">diag</span></span>(<span>self) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Extract diagonal of the entire Kronecker factorization.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>diag</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.matrix.Kron.to_matrix"><code class="name flex">
+<span>def <span class="ident">to_matrix</span></span>(<span>self) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Make the Kronecker factorization dense by computing the kronecker product.
+Warning: this should only be used for testing purposes as it will allocate
+large amounts of memory for big architectures.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>block_diag</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+</dl>
+</dd>
+<dt id="laplace.matrix.KronDecomposed"><code class="flex name class">
+<span>class <span class="ident">KronDecomposed</span></span>
+<span>(</span><span>eigenvectors, eigenvalues, deltas=None, damping=False)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Decomposed Kronecker factored approximate curvature representation
+for a corresponding neural network.
+Each matrix in <code><a title="laplace.matrix.Kron" href="#laplace.matrix.Kron">Kron</a></code> is decomposed to obtain <code><a title="laplace.matrix.KronDecomposed" href="#laplace.matrix.KronDecomposed">KronDecomposed</a></code>.
+Front-loading decomposition allows cheap repeated computation
+of inverses and log determinants.
+In contrast to <code><a title="laplace.matrix.Kron" href="#laplace.matrix.Kron">Kron</a></code>, we can add scalar or layerwise scalars but
+we cannot add other <code><a title="laplace.matrix.Kron" href="#laplace.matrix.Kron">Kron</a></code> or <code><a title="laplace.matrix.KronDecomposed" href="#laplace.matrix.KronDecomposed">KronDecomposed</a></code> anymore.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>eigenvectors</code></strong> :&ensp;<code>list[Tuple[torch.Tensor]]</code></dt>
+<dd>eigenvectors corresponding to matrices in a corresponding <code><a title="laplace.matrix.Kron" href="#laplace.matrix.Kron">Kron</a></code></dd>
+<dt><strong><code>eigenvalues</code></strong> :&ensp;<code>list[Tuple[torch.Tensor]]</code></dt>
+<dd>eigenvalues corresponding to matrices in a corresponding <code><a title="laplace.matrix.Kron" href="#laplace.matrix.Kron">Kron</a></code></dd>
+<dt><strong><code>deltas</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>addend for each group of Kronecker factors representing, for example,
+a prior precision</dd>
+<dt><strong><code>dampen</code></strong> :&ensp;<code>bool</code>, default=<code>False</code></dt>
+<dd>use dampen approximation mixing prior and Kron partially multiplicatively</dd>
+</dl></div>
+<h3>Methods</h3>
+<dl>
+<dt id="laplace.matrix.KronDecomposed.detach"><code class="name flex">
+<span>def <span class="ident">detach</span></span>(<span>self)</span>
+</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.matrix.KronDecomposed.logdet"><code class="name flex">
+<span>def <span class="ident">logdet</span></span>(<span>self) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute log determinant of the Kronecker factors and sums them up.
+This corresponds to the log determinant of the entire Hessian approximation.
+In contrast to <code><a title="laplace.matrix.Kron.logdet" href="#laplace.matrix.Kron.logdet">Kron.logdet()</a></code>, additive <code>deltas</code> corresponding to prior
+precisions are added.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>logdet</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.matrix.KronDecomposed.inv_square_form"><code class="name flex">
+<span>def <span class="ident">inv_square_form</span></span>(<span>self, W: torch.Tensor) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.matrix.KronDecomposed.bmm"><code class="name flex">
+<span>def <span class="ident">bmm</span></span>(<span>self, W: torch.Tensor, exponent: float = -1) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Batched matrix multiplication with the decomposed Kronecker factors.
+This is useful for computing the predictive or a regularization loss.
+Compared to <code><a title="laplace.matrix.Kron.bmm" href="#laplace.matrix.Kron.bmm">Kron.bmm()</a></code>, a prior can be added here in form of <code>deltas</code>
+and the exponent can be other than just 1.
+Computes <span><span class="MathJax_Preview">H^{exponent} W</span><script type="math/tex">H^{exponent} W</script></span>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>W</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>matrix <code>(batch, classes, params)</code></dd>
+<dt><strong><code>exponent</code></strong> :&ensp;<code>float</code>, default=<code>1</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>SW</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>result <code>(batch, classes, params)</code></dd>
+</dl></div>
+</dd>
+<dt id="laplace.matrix.KronDecomposed.to_matrix"><code class="name flex">
+<span>def <span class="ident">to_matrix</span></span>(<span>self, exponent: float = 1) ‑> torch.Tensor</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Make the Kronecker factorization dense by computing the kronecker product.
+Warning: this should only be used for testing purposes as it will allocate
+large amounts of memory for big architectures.</p>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>block_diag</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+</dl>
+</dd>
+</dl>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace" href="index.html">laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-classes">Classes</a></h3>
+<ul>
+<li>
+<h4><code><a title="laplace.matrix.Kron" href="#laplace.matrix.Kron">Kron</a></code></h4>
+<ul class="two-column">
+<li><code><a title="laplace.matrix.Kron.init_from_model" href="#laplace.matrix.Kron.init_from_model">init_from_model</a></code></li>
+<li><code><a title="laplace.matrix.Kron.decompose" href="#laplace.matrix.Kron.decompose">decompose</a></code></li>
+<li><code><a title="laplace.matrix.Kron.bmm" href="#laplace.matrix.Kron.bmm">bmm</a></code></li>
+<li><code><a title="laplace.matrix.Kron.logdet" href="#laplace.matrix.Kron.logdet">logdet</a></code></li>
+<li><code><a title="laplace.matrix.Kron.diag" href="#laplace.matrix.Kron.diag">diag</a></code></li>
+<li><code><a title="laplace.matrix.Kron.to_matrix" href="#laplace.matrix.Kron.to_matrix">to_matrix</a></code></li>
+</ul>
+</li>
+<li>
+<h4><code><a title="laplace.matrix.KronDecomposed" href="#laplace.matrix.KronDecomposed">KronDecomposed</a></code></h4>
+<ul class="">
+<li><code><a title="laplace.matrix.KronDecomposed.detach" href="#laplace.matrix.KronDecomposed.detach">detach</a></code></li>
+<li><code><a title="laplace.matrix.KronDecomposed.logdet" href="#laplace.matrix.KronDecomposed.logdet">logdet</a></code></li>
+<li><code><a title="laplace.matrix.KronDecomposed.inv_square_form" href="#laplace.matrix.KronDecomposed.inv_square_form">inv_square_form</a></code></li>
+<li><code><a title="laplace.matrix.KronDecomposed.bmm" href="#laplace.matrix.KronDecomposed.bmm">bmm</a></code></li>
+<li><code><a title="laplace.matrix.KronDecomposed.to_matrix" href="#laplace.matrix.KronDecomposed.to_matrix">to_matrix</a></code></li>
+</ul>
+</li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/regression_example.png b/docs/regression_example.png
new file mode 100644
index 00000000..0869691a
Binary files /dev/null and b/docs/regression_example.png differ
diff --git a/docs/utils.html b/docs/utils.html
new file mode 100644
index 00000000..ddb90e62
--- /dev/null
+++ b/docs/utils.html
@@ -0,0 +1,184 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+<meta name="generator" content="pdoc 0.9.2" />
+<title>laplace.utils API documentation</title>
+<meta name="description" content="" />
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/sanitize.min.css" integrity="sha256-PK9q560IAAa6WVRRh76LtCaI8pjTJ2z11v0miyNNjrs=" crossorigin>
+<link rel="preload stylesheet" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/11.0.1/typography.min.css" integrity="sha256-7l/o7C8jubJiy74VsKTidCy1yBkRtiUGbVkYBylBqUg=" crossorigin>
+<link rel="stylesheet preload" as="style" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/styles/github.min.css" crossorigin>
+<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:30px;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:1em 0 .50em 0}h3{font-size:1.4em;margin:25px 0 10px 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .3s ease-in-out}a:hover{color:#e82}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900}pre code{background:#f8f8f8;font-size:.8em;line-height:1.4em}code{background:#f2f2f1;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{background:#f8f8f8;border:0;border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0;padding:1ex}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-weight:bold;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em .5em;margin-bottom:1em}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
+<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.item .name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul{padding-left:1.5em}.toc > ul > li{margin-top:.5em}}</style>
+<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/latest.js?config=TeX-AMS_CHTML" integrity="sha256-kZafAc6mZvK3W3v1pHOcUix30OHQN6pU/NO2oFkqZVw=" crossorigin></script>
+<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/10.1.1/highlight.min.js" integrity="sha256-Uv3H6lx7dJmRfRvH8TH6kJD1TSK1aFcwgx+mdg3epi8=" crossorigin></script>
+<script>window.addEventListener('DOMContentLoaded', () => hljs.initHighlighting())</script>
+</head>
+<body>
+<main>
+<article id="content">
+<header>
+<h1 class="title">Module <code>laplace.utils</code></h1>
+</header>
+<section id="section-intro">
+</section>
+<section>
+</section>
+<section>
+</section>
+<section>
+<h2 class="section-title" id="header-functions">Functions</h2>
+<dl>
+<dt id="laplace.utils.get_nll"><code class="name flex">
+<span>def <span class="ident">get_nll</span></span>(<span>out_dist, targets)</span>
+</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.utils.validate"><code class="name flex">
+<span>def <span class="ident">validate</span></span>(<span>laplace, val_loader, pred_type='glm', link_approx='probit', n_samples=100)</span>
+</code></dt>
+<dd>
+<div class="desc"></div>
+</dd>
+<dt id="laplace.utils.parameters_per_layer"><code class="name flex">
+<span>def <span class="ident">parameters_per_layer</span></span>(<span>model)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Get number of parameters per layer.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>model</code></strong> :&ensp;<code>torch.nn.Module</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>params_per_layer</code></strong> :&ensp;<code>list[int]</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.utils.invsqrt_precision"><code class="name flex">
+<span>def <span class="ident">invsqrt_precision</span></span>(<span>M)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compute <code>M^{-0.5}</code> as a tridiagonal matrix.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>M</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>M_invsqrt</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.utils.kron"><code class="name flex">
+<span>def <span class="ident">kron</span></span>(<span>t1, t2)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Computes the Kronecker product between two tensors.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>t1</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>t2</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>kron_product</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.utils.diagonal_add_scalar"><code class="name flex">
+<span>def <span class="ident">diagonal_add_scalar</span></span>(<span>X, value)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Add scalar value <code>value</code> to diagonal of <code>X</code>.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>X</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+<dt><strong><code>value</code></strong> :&ensp;<code>torch.Tensor</code> or <code>float</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>X_add_scalar</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+<dt id="laplace.utils.symeig"><code class="name flex">
+<span>def <span class="ident">symeig</span></span>(<span>M)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Symetric eigendecomposition avoiding failure cases by
+adding and removing jitter to the diagonal.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>M</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>L</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>eigenvalues</dd>
+<dt><strong><code>W</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>eigenvectors</dd>
+</dl></div>
+</dd>
+<dt id="laplace.utils.block_diag"><code class="name flex">
+<span>def <span class="ident">block_diag</span></span>(<span>blocks)</span>
+</code></dt>
+<dd>
+<div class="desc"><p>Compose block-diagonal matrix of individual blocks.</p>
+<h2 id="parameters">Parameters</h2>
+<dl>
+<dt><strong><code>blocks</code></strong> :&ensp;<code>list[torch.Tensor]</code></dt>
+<dd>&nbsp;</dd>
+</dl>
+<h2 id="returns">Returns</h2>
+<dl>
+<dt><strong><code>M</code></strong> :&ensp;<code>torch.Tensor</code></dt>
+<dd>&nbsp;</dd>
+</dl></div>
+</dd>
+</dl>
+</section>
+<section>
+</section>
+</article>
+<nav id="sidebar">
+<h1>Index</h1>
+<div class="toc">
+<ul></ul>
+</div>
+<ul id="index">
+<li><h3>Super-module</h3>
+<ul>
+<li><code><a title="laplace" href="index.html">laplace</a></code></li>
+</ul>
+</li>
+<li><h3><a href="#header-functions">Functions</a></h3>
+<ul class="">
+<li><code><a title="laplace.utils.get_nll" href="#laplace.utils.get_nll">get_nll</a></code></li>
+<li><code><a title="laplace.utils.validate" href="#laplace.utils.validate">validate</a></code></li>
+<li><code><a title="laplace.utils.parameters_per_layer" href="#laplace.utils.parameters_per_layer">parameters_per_layer</a></code></li>
+<li><code><a title="laplace.utils.invsqrt_precision" href="#laplace.utils.invsqrt_precision">invsqrt_precision</a></code></li>
+<li><code><a title="laplace.utils.kron" href="#laplace.utils.kron">kron</a></code></li>
+<li><code><a title="laplace.utils.diagonal_add_scalar" href="#laplace.utils.diagonal_add_scalar">diagonal_add_scalar</a></code></li>
+<li><code><a title="laplace.utils.symeig" href="#laplace.utils.symeig">symeig</a></code></li>
+<li><code><a title="laplace.utils.block_diag" href="#laplace.utils.block_diag">block_diag</a></code></li>
+</ul>
+</li>
+</ul>
+</nav>
+</main>
+<footer id="footer">
+<p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> 0.9.2</a>.</p>
+</footer>
+</body>
+</html>
\ No newline at end of file
diff --git a/laplace/__init__.py b/laplace/__init__.py
new file mode 100644
index 00000000..90dd0cec
--- /dev/null
+++ b/laplace/__init__.py
@@ -0,0 +1,17 @@
+"""
+.. include:: ../README.md
+
+.. include:: ../regression_example.md
+"""
+REGRESSION = 'regression'
+CLASSIFICATION = 'classification'
+
+from laplace.baselaplace import BaseLaplace, FullLaplace, KronLaplace, DiagLaplace
+from laplace.lllaplace import LLLaplace, FullLLLaplace, KronLLLaplace, DiagLLLaplace
+from laplace.laplace import Laplace
+
+__all__ = ['Laplace',  # direct access to all Laplace classes via unified interface
+           'BaseLaplace',  # base-class
+           'FullLaplace', 'KronLaplace', 'DiagLaplace',  # all-weights
+           'LLLaplace',  # base-class last-layer
+           'FullLLLaplace', 'KronLLLaplace', 'DiagLLLaplace'] # last-layer
diff --git a/laplace/baselaplace.py b/laplace/baselaplace.py
new file mode 100644
index 00000000..d273ce1e
--- /dev/null
+++ b/laplace/baselaplace.py
@@ -0,0 +1,800 @@
+from abc import ABC, abstractmethod, abstractproperty
+from math import sqrt, pi
+import numpy as np
+import torch
+from torch.nn.utils import parameters_to_vector, vector_to_parameters
+from torch.distributions import MultivariateNormal, Dirichlet, Normal
+
+from laplace.utils import parameters_per_layer, invsqrt_precision, get_nll, validate
+from laplace.matrix import Kron
+from laplace.curvature import BackPackGGN
+
+
+__all__ = ['BaseLaplace', 'FullLaplace', 'KronLaplace', 'DiagLaplace']
+
+
+class BaseLaplace(ABC):
+    """Baseclass for all Laplace approximations in this library.
+    Subclasses need to specify how the Hessian approximation is initialized,
+    how to add up curvature over training data, how to sample from the
+    Laplace approximation, and how to compute the functional variance.
+
+    A Laplace approximation is represented by a MAP which is given by the
+    `model` parameter and a posterior precision or covariance specifying
+    a Gaussian distribution \\(\\mathcal{N}(\\theta_{MAP}, P^{-1})\\).
+    The goal of this class is to compute the posterior precision \\(P\\)
+    which sums as
+    \\[
+        P = \\sum_{n=1}^N \\nabla^2_\\theta \\log p(\\mathcal{D}_n \\mid \\theta)
+        \\vert_{\\theta_{MAP}} + \\nabla^2_\\theta \\log p(\\theta) \\vert_{\\theta_{MAP}}.
+    \\]
+    Every subclass implements different approximations to the log likelihood Hessians,
+    for example, a diagonal one. The prior is assumed to be Gaussian and therefore we have
+    a simple form for \\(\\nabla^2_\\theta \\log p(\\theta) \\vert_{\\theta_{MAP}} = P_0 \\).
+    In particular, we assume a scalar, layer-wise, or diagonal prior precision so that in
+    all cases \\(P_0 = \\textrm{diag}(p_0)\\) and the structure of \\(p_0\\) can be varied.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+    likelihood : {'classification', 'regression'}
+        determines the log likelihood Hessian approximation
+    sigma_noise : torch.Tensor or float, default=1
+        observation noise for the regression setting; must be 1 for classification
+    prior_precision : torch.Tensor or float, default=1
+        prior precision of a Gaussian prior (= weight decay);
+        can be scalar, per-layer, or diagonal in the most general case
+    prior_mean : torch.Tensor or float, default=0
+        prior mean of a Gaussian prior, useful for continual learning
+    temperature : float, default=1
+        temperature of the likelihood; lower temperature leads to more
+        concentrated posterior and vice versa.
+    backend : subclasses of `laplace.curvature.CurvatureInterface`
+        backend for access to curvature/Hessian approximations
+    backend_kwargs : dict, default=None
+        arguments passed to the backend on initialization, for example to
+        set the number of MC samples for stochastic approximations.
+    """
+    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
+                 prior_mean=0., temperature=1., backend=BackPackGGN, backend_kwargs=None):
+        if likelihood not in ['classification', 'regression']:
+            raise ValueError(f'Invalid likelihood type {likelihood}')
+
+        self.model = model
+        self._device = next(model.parameters()).device
+        # initialize state #
+        # posterior mean/mode
+        self.mean = parameters_to_vector(self.model.parameters()).detach()
+        self.n_params = len(self.mean)
+        self.n_layers = len(list(self.model.parameters()))
+        self.prior_precision = prior_precision
+        self.prior_mean = prior_mean
+        if sigma_noise != 1 and likelihood != 'regression':
+            raise ValueError('Sigma noise != 1 only available for regression.')
+        self.likelihood = likelihood
+        self.sigma_noise = sigma_noise
+        self.temperature = temperature
+        self._backend = None
+        self._backend_cls = backend
+        self._backend_kwargs = dict() if backend_kwargs is None else backend_kwargs
+        self.H = None
+
+        # log likelihood = g(loss)
+        self.loss = 0.
+        self.n_outputs = None
+        self.n_data = None
+
+    @property
+    def backend(self):
+        if self._backend is None:
+            self._backend = self._backend_cls(self.model, self.likelihood,
+                                              **self._backend_kwargs)
+        return self._backend
+
+    @abstractmethod
+    def _init_H(self):
+        pass
+
+    @abstractmethod
+    def _curv_closure(self, X, y, N):
+        pass
+
+    def _check_fit(self):
+        if self.H is None:
+            raise AttributeError('Laplace not fitted. Run fit() first.')
+
+    def fit(self, train_loader):
+        """Fit the local Laplace approximation at the parameters of the model.
+
+        Parameters
+        ----------
+        train_loader : torch.data.utils.DataLoader
+            each iterate is a training batch (X, y);
+            `train_loader.dataset` needs to be set to access \\(N\\), size of the data set
+        """
+        if self.H is not None:
+            raise ValueError('Already fit.')
+
+        self._init_H()
+
+        self.model.eval()
+
+        X, _ = next(iter(train_loader))
+        with torch.no_grad():
+            self.n_outputs = self.model(X[:1].to(self._device)).shape[-1]
+        setattr(self.model, 'output_size', self.n_outputs)
+
+        N = len(train_loader.dataset)
+        for X, y in train_loader:
+            self.model.zero_grad()
+            X, y = X.to(self._device), y.to(self._device)
+            loss_batch, H_batch = self._curv_closure(X, y, N)
+            self.loss += loss_batch
+            self.H += H_batch
+
+        self.n_data = N
+
+    def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
+        """Compute the Laplace approximation to the log marginal likelihood subject
+        to specific Hessian approximations that subclasses implement.
+        Requires that the Laplace approximation has been fit before.
+        The resulting torch.Tensor is differentiable in `prior_precision` and
+        `sigma_noise` if these have gradients enabled.
+        By passing `prior_precision` or `sigma_noise`, the current value is
+        overwritten. This is useful for iterating on the log marginal likelihood.
+
+        Parameters
+        ----------
+        prior_precision : torch.Tensor, optional
+            prior precision if should be changed from current `prior_precision` value
+        sigma_noise : [type], optional
+            observation noise standard deviation if should be changed
+
+        Returns
+        -------
+        log_marglik : torch.Tensor
+        """
+        # make sure we can differentiate wrt prior and sigma_noise for regression
+        self._check_fit()
+
+        # update prior precision (useful when iterating on marglik)
+        if prior_precision is not None:
+            self.prior_precision = prior_precision
+
+        # update sigma_noise (useful when iterating on marglik)
+        if sigma_noise is not None:
+            if self.likelihood != 'regression':
+                raise ValueError('Can only change sigma_noise for regression.')
+            self.sigma_noise = sigma_noise
+
+        return self.log_likelihood - 0.5 * (self.log_det_ratio + self.scatter)
+
+    @property
+    def log_likelihood(self):
+        """Compute log likelihood on the training data after `.fit()` has been called.
+        The log likelihood is computed on-demand based on the loss and, for example,
+        the observation noise which makes it differentiable in the latter for
+        iterative updates.
+
+        Returns
+        -------
+        log_likelihood : torch.Tensor
+        """
+        self._check_fit()
+
+        factor = - self._H_factor
+        if self.likelihood == 'regression':
+            # loss used is just MSE, need to add normalizer for gaussian likelihood
+            c = self.n_data * self.n_outputs * torch.log(self.sigma_noise * sqrt(2 * pi))
+            return factor * self.loss - c
+        else:
+            # for classification Xent == log Cat
+            return factor * self.loss
+
+    def __call__(self, x, pred_type='glm', link_approx='probit', n_samples=100):
+        """Compute the posterior predictive on input data `X`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            `(batch_size, input_shape)`
+
+        pred_type : {'glm', 'nn'}, default='glm'
+            type of posterior predictive, linearized GLM predictive or neural
+            network sampling predictive. The GLM predictive is consistent with
+            the curvature approximations used here.
+
+        link_approx : {'mc', 'probit', 'bridge'}
+            how to approximate the classification link function for the `'glm'`.
+            For `pred_type='nn'`, only 'mc' is possible.
+
+        n_samples : int
+            number of samples for `link_approx='mc'`.
+
+        Returns
+        -------
+        predictive: torch.Tensor or Tuple[torch.Tensor]
+            For `likelihood='classification'`, a torch.Tensor is returned with
+            a distribution over classes (similar to a Softmax).
+            For `likelihood='regression'`, a tuple of torch.Tensor is returned
+            with the mean and the predictive variance.
+        """
+        self._check_fit()
+
+        if pred_type not in ['glm', 'nn']:
+            raise ValueError('Only glm and nn supported as prediction types.')
+
+        if link_approx not in ['mc', 'probit', 'bridge']:
+            raise ValueError(f'Unsupported link approximation {link_approx}.')
+
+        if pred_type == 'glm':
+            f_mu, f_var = self._glm_predictive_distribution(x)
+            # regression
+            if self.likelihood == 'regression':
+                return f_mu, f_var
+            # classification
+            if link_approx == 'mc':
+                try:
+                    dist = MultivariateNormal(f_mu, f_var)
+                except:
+                    dist = Normal(f_mu, torch.diagonal(f_var, dim1=1, dim2=2).sqrt())
+                return torch.softmax(dist.sample((n_samples,)), dim=-1).mean(dim=0)
+            elif link_approx == 'probit':
+                kappa = 1 / torch.sqrt(1. + np.pi / 8 * f_var.diagonal(dim1=1, dim2=2))
+                return torch.softmax(kappa * f_mu, dim=-1)
+            elif link_approx == 'bridge':
+                _, K = f_mu.size(0), f_mu.size(-1)
+                f_var_diag = torch.diagonal(f_var, dim1=1, dim2=2)
+                sum_exp = torch.sum(torch.exp(-f_mu), dim=1).unsqueeze(-1)
+                alpha = 1/f_var_diag * (1 - 2/K + torch.exp(f_mu)/(K**2) * sum_exp)
+                dist = Dirichlet(alpha)
+                return torch.nan_to_num(dist.mean, nan=1.0)
+        else:
+            samples = self._nn_predictive_samples(x, n_samples)
+            if self.likelihood == 'regression':
+                return samples.mean(dim=0), samples.var(dim=0)
+            return samples.mean(dim=0)
+
+    def predictive(self, x, pred_type='glm', link_approx='mc', n_samples=100):
+        return self(x, pred_type, link_approx, n_samples)
+
+    def predictive_samples(self, x, pred_type='glm', n_samples=100):
+        """Sample from the posterior predictive on input data `x`.
+        Can be used, for example, for Thompson sampling.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch_size, input_shape)`
+
+        pred_type : {'glm', 'nn'}, default='glm'
+            type of posterior predictive, linearized GLM predictive or neural
+            network sampling predictive. The GLM predictive is consistent with
+            the curvature approximations used here.
+
+        n_samples : int
+            number of samples
+
+        Returns
+        -------
+        samples : torch.Tensor
+            samples `(n_samples, batch_size, output_shape)`
+        """
+        self._check_fit()
+
+        if pred_type not in ['glm', 'nn']:
+            raise ValueError('Only glm and nn supported as prediction types.')
+
+        if pred_type == 'glm':
+            f_mu, f_var = self._glm_predictive_distribution(x)
+            assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
+            dist = MultivariateNormal(f_mu, f_var)
+            samples = dist.sample((n_samples,))
+            if self.likelihood == 'regression':
+                return samples
+            return torch.softmax(samples, dim=-1)
+
+        else:  # 'nn'
+            return self._nn_predictive_samples(x, n_samples)
+
+    @torch.enable_grad()
+    def _glm_predictive_distribution(self, X):
+        Js, f_mu = self.backend.jacobians(self.model, X)
+        f_var = self.functional_variance(Js)
+        return f_mu.detach(), f_var.detach()
+
+    def _nn_predictive_samples(self, X, n_samples=100):
+        fs = list()
+        for sample in self.sample(n_samples):
+            vector_to_parameters(sample, self.model.parameters())
+            fs.append(self.model(X.to(self._device)).detach())
+        vector_to_parameters(self.mean, self.model.parameters())
+        fs = torch.stack(fs)
+        if self.likelihood == 'classification':
+            fs = torch.softmax(fs, dim=-1)
+        return fs
+
+    @abstractmethod
+    def functional_variance(self, Jacs):
+        """Compute functional variance for the `'glm'` predictive:
+        `f_var[i] = Jacs[i] @ P.inv() @ Jacs[i].T`, which is a output x output
+        predictive covariance matrix.
+        Mathematically, we have for a single Jacobian
+        \\(\\mathcal{J} = \\nabla_\\theta f(x;\\theta)\\vert_{\\theta_{MAP}}\\)
+        the output covariance matrix
+        \\( \\mathcal{J} P^{-1} \\mathcal{J}^T \\).
+
+        Parameters
+        ----------
+        Jacs : torch.Tensor
+            Jacobians of model output wrt parameters
+            `(batch, outputs, parameters)`
+
+        Returns
+        -------
+        f_var : torch.Tensor
+            output covariance `(batch, outputs, outputs)`
+        """
+        pass
+
+    def _check_jacobians(self, Js):
+        if not isinstance(Js, torch.Tensor):
+            raise ValueError('Jacobians have to be torch.Tensor.')
+        if not Js.device == self._device:
+            raise ValueError('Jacobians need to be on the same device as Laplace.')
+        m, k, p = Js.size()
+        if p != self.n_params:
+            raise ValueError('Invalid Jacobians shape for Laplace posterior approx.')
+
+    @abstractmethod
+    def sample(self, n_samples=100):
+        """Sample from the Laplace posterior approximation, i.e.,
+        \\( \\theta \\sim \\mathcal{N}(\\theta_{MAP}, P^{-1})\\).
+
+        Parameters
+        ----------
+        n_samples : int, default=100
+            number of samples
+        """
+        pass
+
+    @property
+    def scatter(self):
+        """Computes the _scatter_, a term of the log marginal likelihood that
+        corresponds to L-2 regularization:
+        `scatter` = \\((\\theta_{MAP} - \\mu_0)^{T} P_0 (\\theta_{MAP} - \\mu_0) \\).
+
+        Returns
+        -------
+        [type]
+            [description]
+        """
+        delta = (self.mean - self.prior_mean)
+        return (delta * self.prior_precision_diag) @ delta
+
+    @property
+    def log_det_prior_precision(self):
+        """Compute log determinant of the prior precision
+        \\(\\log \\det P_0\\)
+
+        Returns
+        -------
+        log_det : torch.Tensor
+        """
+        return self.prior_precision_diag.log().sum()
+
+    @abstractproperty
+    def log_det_posterior_precision(self):
+        """Compute log determinant of the posterior precision
+        \\(\\log \\det P\\) which depends on the subclasses structure
+        used for the Hessian approximation.
+
+        Returns
+        -------
+        log_det : torch.Tensor
+        """
+        pass
+
+    @property
+    def log_det_ratio(self):
+        """Compute the log determinant ratio, a part of the log marginal likelihood.
+        \\[
+            \\log \\frac{\\det P}{\\det P_0} = \\log \\det P - \\log \\det P_0
+        \\]
+
+        Returns
+        -------
+        log_det_ratio : torch.Tensor
+        """
+        return self.log_det_posterior_precision - self.log_det_prior_precision
+
+    @property
+    def prior_precision_diag(self):
+        """Obtain the diagonal prior precision \\(p_0\\) constructed from either
+        a scalar, layer-wise, or diagonal prior precision.
+
+        Returns
+        -------
+        prior_precision_diag : torch.Tensor
+        """
+        if len(self.prior_precision) == 1:  # scalar
+            return self.prior_precision * torch.ones_like(self.mean, device=self._device)
+
+        elif len(self.prior_precision) == self.n_params:  # diagonal
+            return self.prior_precision
+
+        elif len(self.prior_precision) == self.n_layers:  # per layer
+            n_params_per_layer = parameters_per_layer(self.model)
+            return torch.cat([prior * torch.ones(n_params, device=self._device) for prior, n_params
+                              in zip(self.prior_precision, n_params_per_layer)])
+
+        else:
+            raise ValueError('Mismatch of prior and model. Diagonal, scalar, or per-layer prior.')
+
+    @property
+    def prior_mean(self):
+        return self._prior_mean
+
+    @prior_mean.setter
+    def prior_mean(self, prior_mean):
+        if np.isscalar(prior_mean) and np.isreal(prior_mean):
+            self._prior_mean = torch.tensor(prior_mean, device=self._device)
+        elif torch.is_tensor(prior_mean):
+            if prior_mean.ndim == 0:
+                self._prior_mean = prior_mean.reshape(-1).to(self._device)
+            elif prior_mean.ndim == 1:
+                if not len(prior_mean) in [1, self.n_params]:
+                    raise ValueError('Invalid length of prior mean.')
+                self._prior_mean = prior_mean
+            else:
+                raise ValueError('Prior mean has too many dimensions!')
+        else:
+            raise ValueError('Invalid argument type of prior mean.')
+
+    @property
+    def prior_precision(self):
+        return self._prior_precision
+
+    @prior_precision.setter
+    def prior_precision(self, prior_precision):
+        self._posterior_scale = None
+        if np.isscalar(prior_precision) and np.isreal(prior_precision):
+            self._prior_precision = torch.tensor([prior_precision], device=self._device)
+        elif torch.is_tensor(prior_precision):
+            if prior_precision.ndim == 0:
+                # make dimensional
+                self._prior_precision = prior_precision.reshape(-1).to(self._device)
+            elif prior_precision.ndim == 1:
+                if len(prior_precision) not in [1, self.n_layers, self.n_params]:
+                    raise ValueError('Length of prior precision does not align with architecture.')
+                self._prior_precision = prior_precision.to(self._device)
+            else:
+                raise ValueError('Prior precision needs to be at most one-dimensional tensor.')
+        else:
+            raise ValueError('Prior precision either scalar or torch.Tensor up to 1-dim.')
+
+    def optimize_prior_precision(self, method='marglik', n_steps=100, lr=1e-1,
+                                 init_prior_prec=1., val_loader=None, loss=get_nll,
+                                 log_prior_prec_min=-4, log_prior_prec_max=4, grid_size=100,
+                                 pred_type='glm', link_approx='probit', n_samples=100,
+                                 verbose=False):
+        """Optimize the prior precision post-hoc using the `method`
+        specified by the user.
+
+        Parameters
+        ----------
+        method : {'marglik', 'CV'}, default='marglik'
+            specifies how the prior precision should be optimized.
+        n_steps : int, default=100
+            the number of gradient descent steps to take.
+        lr : float, default=1e-1
+            the learning rate to use for gradient descent.
+        init_prior_prec : float, default=1.0
+            initial prior precision before the first optimization step.
+        val_loader : torch.data.utils.DataLoader, default=None
+            DataLoader for the validation set; each iterate is a training batch (X, y).
+        loss : callable, default=get_nll
+            loss function to use for CV.
+        log_prior_prec_min : float, default=-4
+            lower bound of gridsearch interval for CV.
+        log_prior_prec_max : float, default=4
+            upper bound of gridsearch interval for CV.
+        grid_size : int, default=100
+            number of values to consider inside the gridsearch interval for CV.
+        pred_type : {'glm', 'nn'}, default='glm'
+            type of posterior predictive, linearized GLM predictive or neural
+            network sampling predictive. The GLM predictive is consistent with
+            the curvature approximations used here.
+        link_approx : {'mc', 'probit', 'bridge'}, default='probit'
+            how to approximate the classification link function for the `'glm'`.
+            For `pred_type='nn'`, only `'mc'` is possible.
+        n_samples : int, default=100
+            number of samples for `link_approx='mc'`.
+        verbose : bool, default=False
+            if true, the optimized prior precision will be printed
+            (can be a large tensor if the prior has a diagonal covariance).
+        """
+        if method == 'marglik':
+            self.prior_precision = init_prior_prec
+            log_prior_prec = self.prior_precision.log()
+            log_prior_prec.requires_grad = True
+            optimizer = torch.optim.Adam([log_prior_prec], lr=lr)
+            for _ in range(n_steps):
+                optimizer.zero_grad()
+                prior_prec = log_prior_prec.exp()
+                neg_log_marglik = -self.log_marginal_likelihood(prior_precision=prior_prec)
+                neg_log_marglik.backward()
+                optimizer.step()
+            self.prior_precision = log_prior_prec.detach().exp()
+        elif method == 'CV':
+            if val_loader is None:
+                raise ValueError('CV requires a validation set DataLoader')
+            interval = torch.logspace(
+                log_prior_prec_min, log_prior_prec_max, grid_size
+            )
+            self.prior_precision = self._gridsearch(
+                loss, interval, val_loader, pred_type=pred_type,
+                link_approx=link_approx, n_samples=n_samples
+            )
+        else:
+            raise ValueError('For now only marglik and CV is implemented.')
+        if verbose:
+            print(f'Optimized prior precision is {self.prior_precision}.')
+
+    def _gridsearch(self, loss, interval, val_loader, pred_type='glm',
+                    link_approx='probit', n_samples=100):
+        results = list()
+        prior_precs = list()
+        for prior_prec in interval:
+            self.prior_precision = prior_prec
+            try:
+                out_dist, targets = validate(
+                    self, val_loader, pred_type=pred_type,
+                    link_approx=link_approx, n_samples=n_samples
+                )
+                result = loss(out_dist, targets)
+            except RuntimeError:
+                result = np.inf
+            results.append(result)
+            prior_precs.append(prior_prec)
+        return prior_precs[np.argmin(results)]
+
+    @property
+    def sigma_noise(self):
+        return self._sigma_noise
+
+    @sigma_noise.setter
+    def sigma_noise(self, sigma_noise):
+        self._posterior_scale = None
+        if np.isscalar(sigma_noise) and np.isreal(sigma_noise):
+            self._sigma_noise = torch.tensor(sigma_noise, device=self._device)
+        elif torch.is_tensor(sigma_noise):
+            if sigma_noise.ndim == 0:
+                self._sigma_noise = sigma_noise.to(self._device)
+            elif sigma_noise.ndim == 1:
+                if len(sigma_noise) > 1:
+                    raise ValueError('Only homoscedastic output noise supported.')
+                self._sigma_noise = sigma_noise[0].to(self._device)
+            else:
+                raise ValueError('Sigma noise needs to be scalar or 1-dimensional.')
+        else:
+            raise ValueError('Invalid type: sigma noise needs to be torch.Tensor or scalar.')
+
+    @property
+    def _H_factor(self):
+        sigma2 = self.sigma_noise.square()
+        return 1 / sigma2 / self.temperature
+
+    @abstractproperty
+    def posterior_precision(self):
+        """Compute or return the posterior precision \\(P\\).
+
+        Returns
+        -------
+        posterior_prec : torch.Tensor
+        """
+        pass
+
+
+class FullLaplace(BaseLaplace):
+    """Laplace approximation with full, i.e., dense, log likelihood Hessian approximation
+    and hence posterior precision. Based on the chosen `backend` parameter, the full
+    approximation can be, for example, a generalized Gauss-Newton matrix.
+    Mathematically, we have \\(P \\in \\mathbb{R}^{P \\times P}\\).
+    See `BaseLaplace` for the full interface.
+    """
+    # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
+    _key = ('all', 'full')
+
+    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
+                 prior_mean=0., temperature=1., backend=BackPackGGN, backend_kwargs=None):
+        super().__init__(model, likelihood, sigma_noise, prior_precision,
+                         prior_mean, temperature, backend, backend_kwargs)
+        self._posterior_scale = None
+
+    def _init_H(self):
+        self.H = torch.zeros(self.n_params, self.n_params, device=self._device)
+
+    def _curv_closure(self, X, y, N):
+        return self.backend.full(X, y, N=N)
+
+    def _compute_scale(self):
+        self._posterior_scale = invsqrt_precision(self.posterior_precision)
+
+    @property
+    def posterior_scale(self):
+        """Posterior scale (square root of the covariance), i.e.,
+        \\(P^{-\\frac{1}{2}}\\).
+
+        Returns
+        -------
+        scale : torch.tensor
+            `(parameters, parameters)`
+        """
+        if self._posterior_scale is None:
+            self._compute_scale()
+        return self._posterior_scale
+
+    @property
+    def posterior_covariance(self):
+        """Posterior covariance, i.e., \\(P^{-1}\\).
+
+        Returns
+        -------
+        covariance : torch.tensor
+            `(parameters, parameters)`
+        """
+        scale = self.posterior_scale
+        return scale @ scale.T
+
+    @property
+    def posterior_precision(self):
+        """Posterior precision \\(P\\).
+
+        Returns
+        -------
+        precision : torch.tensor
+            `(parameters, parameters)`
+        """
+        self._check_fit()
+        return self._H_factor * self.H + torch.diag(self.prior_precision_diag)
+
+    @property
+    def log_det_posterior_precision(self):
+        return self.posterior_precision.logdet()
+
+    def functional_variance(self, Js):
+        return torch.einsum('ncp,pq,nkq->nck', Js, self.posterior_covariance, Js)
+
+    def sample(self, n_samples=100):
+        dist = MultivariateNormal(loc=self.mean, scale_tril=self.posterior_scale)
+        return dist.sample((n_samples,))
+
+
+class KronLaplace(BaseLaplace):
+    """Laplace approximation with Kronecker factored log likelihood Hessian approximation
+    and hence posterior precision.
+    Mathematically, we have for each parameter group, e.g., torch.nn.Module,
+    that \\P\\approx Q \\otimes H\\.
+    See `BaseLaplace` for the full interface and see
+    `laplace.matrix.Kron` and `laplace.matrix.KronDecomposed` for the structure of
+    the Kronecker factors. `Kron` is used to aggregate factors by summing up and
+    `KronDecomposed` is used to add the prior, a Hessian factor (e.g. temperature),
+    and computing posterior covariances, marginal likelihood, etc.
+    Damping can be enabled by setting `damping=True`.
+    """
+    # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
+    _key = ('all', 'kron')
+
+    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
+                 prior_mean=0., temperature=1., backend=BackPackGGN, damping=False,
+                 **backend_kwargs):
+        self.damping = damping
+        super().__init__(model, likelihood, sigma_noise, prior_precision,
+                         prior_mean, temperature, backend, **backend_kwargs)
+
+    def _init_H(self):
+        self.H = Kron.init_from_model(self.model, self._device)
+
+    def _curv_closure(self, X, y, N):
+        return self.backend.kron(X, y, N=N)
+
+    def fit(self, train_loader, keep_factors=False):
+        super().fit(train_loader)
+        # Kron requires postprocessing as all quantities depend on the decomposition.
+        if keep_factors:
+            self.H_facs = self.H
+        self.H = self.H.decompose(damping=self.damping)
+
+    @property
+    def posterior_precision(self):
+        """Kronecker factored Posterior precision \\(P\\).
+
+        Returns
+        -------
+        precision : `laplace.matrix.KronDecomposed`
+        """
+        self._check_fit()
+        return self.H * self._H_factor + self.prior_precision
+
+    @property
+    def log_det_posterior_precision(self):
+        return self.posterior_precision.logdet()
+
+    def functional_variance(self, Js):
+        return self.posterior_precision.inv_square_form(Js)
+
+    def sample(self, n_samples=100):
+        samples = torch.randn(n_samples, self.n_params, device=self._device)
+        samples = self.posterior_precision.bmm(samples, exponent=-0.5)
+        return self.mean.reshape(1, self.n_params) + samples.reshape(n_samples, self.n_params)
+
+    @BaseLaplace.prior_precision.setter
+    def prior_precision(self, prior_precision):
+        # Extend setter from Laplace to restrict prior precision structure.
+        super(KronLaplace, type(self)).prior_precision.fset(self, prior_precision)
+        if len(self.prior_precision) not in [1, self.n_layers]:
+            raise ValueError('Prior precision for Kron either scalar or per-layer.')
+
+
+class DiagLaplace(BaseLaplace):
+    """Laplace approximation with diagonal log likelihood Hessian approximation
+    and hence posterior precision.
+    Mathematically, we have \\(P \\approx \\textrm{diag}(P)\\).
+    See `BaseLaplace` for the full interface.
+    """
+    # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
+    _key = ('all', 'diag')
+
+    def _init_H(self):
+        self.H = torch.zeros(self.n_params, device=self._device)
+
+    def _curv_closure(self, X, y, N):
+        return self.backend.diag(X, y, N=N)
+
+    @property
+    def posterior_precision(self):
+        """Diagonal posterior precision \\(p\\).
+
+        Returns
+        -------
+        precision : torch.tensor
+            `(parameters)`
+        """
+        self._check_fit()
+        return self._H_factor * self.H + self.prior_precision_diag
+
+    @property
+    def posterior_scale(self):
+        """Diagonal posterior scale \\(\\sqrt{p^{-1}}\\).
+
+        Returns
+        -------
+        precision : torch.tensor
+            `(parameters)`
+        """
+        return 1 / self.posterior_precision.sqrt()
+
+    @property
+    def posterior_variance(self):
+        """Diagonal posterior variance \\(p^{-1}\\).
+
+        Returns
+        -------
+        precision : torch.tensor
+            `(parameters)`
+        """
+        return 1 / self.posterior_precision
+
+    @property
+    def log_det_posterior_precision(self):
+        return self.posterior_precision.log().sum()
+
+    def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
+        self._check_jacobians(Js)
+        return torch.einsum('ncp,p,nkp->nck', Js, self.posterior_variance, Js)
+
+    def sample(self, n_samples=100):
+        samples = torch.randn(n_samples, self.n_params, device=self._device)
+        samples = samples * self.posterior_scale.reshape(1, self.n_params)
+        return self.mean.reshape(1, self.n_params) + samples
diff --git a/laplace/curvature/__init__.py b/laplace/curvature/__init__.py
new file mode 100644
index 00000000..93e6cc8f
--- /dev/null
+++ b/laplace/curvature/__init__.py
@@ -0,0 +1,17 @@
+import logging
+
+from laplace.curvature.curvature import CurvatureInterface, GGNInterface, EFInterface
+
+try:
+    from laplace.curvature.backpack import BackPackGGN, BackPackEF, BackPackInterface
+except ModuleNotFoundError:
+    logging.info('Backpack not available.')
+
+try:
+    from laplace.curvature.asdl import AsdlGGN, AsdlEF, AsdlInterface
+except ModuleNotFoundError:
+    logging.info('asdfghjkl backend not available.')
+
+__all__ = ['CurvatureInterface', 'GGNInterface', 'EFInterface',
+           'BackPackInterface', 'BackPackGGN', 'BackPackEF',
+           'AsdlInterface', 'AsdlGGN', 'AsdlEF']
diff --git a/laplace/curvature/asdl.py b/laplace/curvature/asdl.py
new file mode 100644
index 00000000..94a2580b
--- /dev/null
+++ b/laplace/curvature/asdl.py
@@ -0,0 +1,174 @@
+from abc import abstractproperty
+import warnings
+import numpy as np
+import torch
+
+from asdfghjkl import FISHER_EXACT, FISHER_MC, COV
+from asdfghjkl import SHAPE_KRON, SHAPE_DIAG
+from asdfghjkl import fisher_for_cross_entropy
+from asdfghjkl.gradient import batch_gradient
+
+from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.matrix import Kron
+from laplace.utils import _is_batchnorm
+
+
+class AsdlInterface(CurvatureInterface):
+    """Interface for asdfghjkl backend.
+    """
+    def __init__(self, model, likelihood, last_layer=False):
+        if likelihood != 'classification':
+            raise ValueError('This backend only supports classification currently.')
+        super().__init__(model, likelihood, last_layer)
+
+    @staticmethod
+    def jacobians(model, x):
+        """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
+        using asdfghjkl's gradient per output dimension.
+
+        Parameters
+        ----------
+        model : torch.nn.Module
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+
+        Returns
+        -------
+        Js : torch.Tensor
+            Jacobians `(batch, parameters, outputs)`
+        f : torch.Tensor
+            output function `(batch, outputs)`
+        """
+        Js = list()
+        for i in range(model.output_size):
+            def loss_fn(outputs, targets):
+                return outputs[:, i].sum()
+
+            f = batch_gradient(model, loss_fn, x, None).detach()
+            Js.append(_get_batch_grad(model))
+        Js = torch.stack(Js, dim=1)
+        return Js, f
+
+    def gradients(self, x, y):
+        """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
+        \\(\\theta\\) using asdfghjkl's backend.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+        y : torch.Tensor
+
+        Returns
+        -------
+        loss : torch.Tensor
+        Gs : torch.Tensor
+            gradients `(batch, parameters)`
+        """
+        f = batch_gradient(self.model, self.lossfunc, x, y).detach()
+        Gs = _get_batch_grad(self._model)
+        loss = self.lossfunc(f, y)
+        return Gs, loss
+
+    @abstractproperty
+    def _ggn_type(self):
+        raise NotImplementedError()
+
+    def _get_kron_factors(self, curv, M):
+        kfacs = list()
+        for module in curv._model.modules():
+            if _is_batchnorm(module):
+                warnings.warn('BatchNorm unsupported for Kron, ignore.')
+                continue
+
+            stats = getattr(module, self._ggn_type, None)
+            if stats is None:
+                continue
+            if hasattr(module, 'bias') and module.bias is not None:
+                # split up bias and weights
+                kfacs.append([stats.kron.B, stats.kron.A[:-1, :-1]])
+                kfacs.append([stats.kron.B * stats.kron.A[-1, -1] / M])
+            elif hasattr(module, 'weight'):
+                p, q = np.prod(stats.kron.B.shape), np.prod(stats.kron.A.shape)
+                if p == q == 1:
+                    kfacs.append([stats.kron.B * stats.kron.A])
+                else:
+                    kfacs.append([stats.kron.B, stats.kron.A])
+            else:
+                raise ValueError(f'Whats happening with {module}?')
+        return Kron(kfacs)
+
+    @staticmethod
+    def _rescale_kron_factors(kron, N):
+        for F in kron.kfacs:
+            if len(F) == 2:
+                F[1] *= 1/N
+        return kron
+
+    def diag(self, X, y, **kwargs):
+        with torch.no_grad():
+            if self.last_layer:
+                f, X = self.model.forward_with_features(X)
+            else:
+                f = self.model(X)
+            loss = self.lossfunc(f, y)
+        curv = fisher_for_cross_entropy(self._model, self._ggn_type, SHAPE_DIAG,
+                                        inputs=X, targets=y)
+        diag_ggn = curv.matrices_to_vector(None)
+        return self.factor * loss, self.factor * diag_ggn
+
+    def kron(self, X, y, N, **wkwargs) -> [torch.Tensor, Kron]:
+        with torch.no_grad():
+            if self.last_layer:
+                f, X = self.model.forward_with_features(X)
+            else:
+                f = self.model(X)
+            loss = self.lossfunc(f, y)
+        curv = fisher_for_cross_entropy(self._model, self._ggn_type, SHAPE_KRON,
+                                        inputs=X, targets=y)
+        M = len(y)
+        kron = self._get_kron_factors(curv, M)
+        kron = self._rescale_kron_factors(kron, N)
+        return self.factor * loss, self.factor * kron
+
+
+class AsdlGGN(AsdlInterface, GGNInterface):
+    """Implementation of the `GGNInterface` using asdfghjkl.
+    """
+    def __init__(self, model, likelihood, last_layer=False, stochastic=False):
+        super().__init__(model, likelihood, last_layer)
+        self.stochastic = stochastic
+
+    @property
+    def _ggn_type(self):
+        return FISHER_MC if self.stochastic else FISHER_EXACT
+
+
+class AsdlEF(AsdlInterface, EFInterface):
+    """Implementation of the `EFInterface` using asdfghjkl.
+    """
+
+    @property
+    def _ggn_type(self):
+        return COV
+
+
+def _flatten_after_batch(tensor: torch.Tensor):
+    if tensor.ndim == 1:
+        return tensor.unsqueeze(-1)
+    else:
+        return tensor.flatten(start_dim=1)
+
+
+def _get_batch_grad(model):
+    batch_grads = list()
+    for module in model.modules():
+        if hasattr(module, 'op_results'):
+            res = module.op_results['batch_grads']
+            if 'weight' in res:
+                batch_grads.append(_flatten_after_batch(res['weight']))
+            if 'bias' in res:
+                batch_grads.append(_flatten_after_batch(res['bias']))
+            if len(set(res.keys()) - {'weight', 'bias'}) > 0:
+                raise ValueError(f'Invalid parameter keys {res.keys()}')
+    return torch.cat(batch_grads, dim=1)
diff --git a/laplace/curvature/backpack.py b/laplace/curvature/backpack.py
new file mode 100644
index 00000000..885ee2b9
--- /dev/null
+++ b/laplace/curvature/backpack.py
@@ -0,0 +1,162 @@
+import torch
+
+from backpack import backpack, extend, memory_cleanup
+from backpack.extensions import DiagGGNExact, DiagGGNMC, KFAC, KFLR, SumGradSquared, BatchGrad
+from backpack.context import CTX
+
+from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.matrix import Kron
+
+
+class BackPackInterface(CurvatureInterface):
+    """Interface for Backpack backend.
+    """
+    def __init__(self, model, likelihood, last_layer=False):
+        super().__init__(model, likelihood, last_layer)
+        extend(self._model)
+        extend(self.lossfunc)
+
+    @staticmethod
+    def jacobians(model, x):
+        """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\)
+        using backpack's BatchGrad per output dimension.
+
+        Parameters
+        ----------
+        model : torch.nn.Module
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+
+        Returns
+        -------
+        Js : torch.Tensor
+            Jacobians `(batch, parameters, outputs)`
+        f : torch.Tensor
+            output function `(batch, outputs)`
+        """
+        model = extend(model)
+        to_stack = []
+        for i in range(model.output_size):
+            model.zero_grad()
+            out = model(x)
+            with backpack(BatchGrad()):
+                if model.output_size > 1:
+                    out[:, i].sum().backward()
+                else:
+                    out.sum().backward()
+                to_cat = []
+                for param in model.parameters():
+                    to_cat.append(param.grad_batch.detach().reshape(x.shape[0], -1))
+                    delattr(param, 'grad_batch')
+                Jk = torch.cat(to_cat, dim=1)
+            to_stack.append(Jk)
+            if i == 0:
+                f = out.detach()
+
+        model.zero_grad()
+        CTX.remove_hooks()
+        _cleanup(model)
+        if model.output_size > 1:
+            return torch.stack(to_stack, dim=2).transpose(1, 2), f
+        else:
+            return Jk.unsqueeze(-1).transpose(1, 2), f
+
+    def gradients(self, x, y):
+        """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
+        \\(\\theta\\) using Backpack's BatchGrad.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+        y : torch.Tensor
+
+        Returns
+        -------
+        loss : torch.Tensor
+        Gs : torch.Tensor
+            gradients `(batch, parameters)`
+        """
+        f = self.model(x)
+        loss = self.lossfunc(f, y)
+        with backpack(BatchGrad()):
+            loss.backward()
+        Gs = torch.cat([p.grad_batch.data.flatten(start_dim=1)
+                        for p in self._model.parameters()], dim=1)
+        return Gs, loss
+
+
+class BackPackGGN(BackPackInterface, GGNInterface):
+    """Implementation of the `GGNInterface` using Backpack.
+    """
+    def __init__(self, model, likelihood, last_layer=False, stochastic=False):
+        super().__init__(model, likelihood, last_layer)
+        self.stochastic = stochastic
+
+    def _get_diag_ggn(self):
+        if self.stochastic:
+            return torch.cat([p.diag_ggn_mc.data.flatten() for p in self._model.parameters()])
+        else:
+            return torch.cat([p.diag_ggn_exact.data.flatten() for p in self._model.parameters()])
+
+    def _get_kron_factors(self):
+        if self.stochastic:
+            return Kron([p.kfac for p in self._model.parameters()])
+        else:
+            return Kron([p.kflr for p in self._model.parameters()])
+
+    @staticmethod
+    def _rescale_kron_factors(kron, M, N):
+        # Renormalize Kronecker factor to sum up correctly over N data points with batches of M
+        # for M=N (full-batch) just M/N=1
+        for F in kron.kfacs:
+            if len(F) == 2:
+                F[1] *= M/N
+        return kron
+
+    def diag(self, X, y, **kwargs):
+        context = DiagGGNMC if self.stochastic else DiagGGNExact
+        f = self.model(X)
+        loss = self.lossfunc(f, y)
+        with backpack(context()):
+            loss.backward()
+        dggn = self._get_diag_ggn()
+
+        return self.factor * loss.detach(), self.factor * dggn
+
+    def kron(self, X, y, N, **kwargs) -> [torch.Tensor, Kron]:
+        context = KFAC if self.stochastic else KFLR
+        f = self.model(X)
+        loss = self.lossfunc(f, y)
+        with backpack(context()):
+            loss.backward()
+        kron = self._get_kron_factors()
+        kron = self._rescale_kron_factors(kron, len(y), N)
+
+        return self.factor * loss.detach(), self.factor * kron
+
+
+class BackPackEF(BackPackInterface, EFInterface):
+    """Implementation of `EFInterface` using Backpack.
+    """
+
+    def diag(self, X, y, **kwargs):
+        f = self.model(X)
+        loss = self.lossfunc(f, y)
+        with backpack(SumGradSquared()):
+            loss.backward()
+        diag_EF = torch.cat([p.sum_grad_squared.data.flatten()
+                             for p in self._model.parameters()])
+
+        return self.factor * loss.detach(), self.factor * diag_EF
+
+    def kron(self, X, y, **kwargs):
+        raise NotImplementedError('Unavailable through Backpack.')
+
+
+def _cleanup(module):
+    for child in module.children():
+        _cleanup(child)
+
+    setattr(module, "_backpack_extend", False)
+    memory_cleanup(module)
diff --git a/laplace/curvature/curvature.py b/laplace/curvature/curvature.py
new file mode 100644
index 00000000..485fd94f
--- /dev/null
+++ b/laplace/curvature/curvature.py
@@ -0,0 +1,292 @@
+from abc import ABC, abstractmethod, abstractstaticmethod
+import torch
+from torch.nn import MSELoss, CrossEntropyLoss
+
+
+class CurvatureInterface(ABC):
+    """Interface to access curvature for a model and corresponding likelihood.
+    A `CurvatureInterface` must inherit from this baseclass and implement the
+    necessary functions `jacobians`, `full`, `kron`, and `diag`.
+    The interface might be extended in the future to account for other curvature
+    structures, for example, a block-diagonal one.
+
+    Parameters
+    ----------
+    model : torch.nn.Module or `laplace.feature_extractor.FeatureExtractor`
+        torch model (neural network)
+    likelihood : {'classification', 'regression'}
+    last_layer : bool, default=False
+        only consider curvature of last layer
+
+    Attributes
+    ----------
+    lossfunc : torch.nn.MSELoss or torch.nn.CrossEntropyLoss
+    factor : float
+        conversion factor between torch losses and base likelihoods
+        For example, \\(\\frac{1}{2}\\) to get to \\(\\mathcal{N}(f, 1)\\) from MSELoss.
+    """
+    def __init__(self, model, likelihood, last_layer=False):
+        assert likelihood in ['regression', 'classification']
+        self.likelihood = likelihood
+        self.model = model
+        self.last_layer = last_layer
+        if likelihood == 'regression':
+            self.lossfunc = MSELoss(reduction='sum')
+            self.factor = 0.5
+        else:
+            self.lossfunc = CrossEntropyLoss(reduction='sum')
+            self.factor = 1.
+
+    @property
+    def _model(self):
+        return self.model.last_layer if self.last_layer else self.model
+
+    @abstractstaticmethod
+    def jacobians(model, x):
+        """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\).
+
+        Parameters
+        ----------
+        model : torch.nn.Module
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+
+        Returns
+        -------
+        Js : torch.Tensor
+            Jacobians `(batch, parameters, outputs)`
+        f : torch.Tensor
+            output function `(batch, outputs)`
+        """
+        pass
+
+    @staticmethod
+    def last_layer_jacobians(model, x):
+        """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\) 
+        only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).
+
+        Parameters
+        ----------
+        model : laplace.feature_extractor.FeatureExtractor
+        x : torch.Tensor
+
+        Returns
+        -------
+        Js : torch.Tensor
+            Jacobians `(batch, last-layer-parameters, outputs)`
+        f : torch.Tensor
+            output function `(batch, outputs)`
+        """
+        f, phi = model.forward_with_features(x)
+        bsize = len(x)
+        output_size = f.shape[-1]
+
+        # calculate Jacobians using the feature vector 'phi'
+        identity = torch.eye(output_size, device=x.device).unsqueeze(0).tile(bsize, 1, 1)
+        # Jacobians are batch x output x params
+        Js = torch.einsum('kp,kij->kijp', phi, identity).reshape(bsize, output_size, -1)
+        if model.last_layer.bias is not None:
+            Js = torch.cat([Js, identity], dim=2)
+
+        return Js, f.detach()
+
+    @abstractstaticmethod
+    def gradients(self, x, y):
+        """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter \\(\\theta\\).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+        y : torch.Tensor
+
+        Returns
+        -------
+        loss : torch.Tensor
+        Gs : torch.Tensor
+            gradients `(batch, parameters)`
+        """
+        pass
+
+    @abstractmethod
+    def full(self, x, y, **kwargs):
+        """Compute a dense curvature (approximation) in the form of a \\(P \\times P\\) matrix
+        \\(H\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)`
+        y : torch.Tensor
+            labels `(batch, label_shape)`
+
+        Returns
+        -------
+        loss : torch.Tensor
+        H : torch.Tensor
+            Hessian approximation `(parameters, parameters)`
+        """
+        pass
+
+    @abstractmethod
+    def kron(self, x, y, **kwargs):
+        """Compute a Kronecker factored curvature approximation (such as KFAC).
+        The approximation to \\(H\\) takes the form of two Kronecker factors \\(Q, H\\),
+        i.e., \\(H \\approx Q \\otimes H\\) for each Module in the neural network permitting 
+        such curvature.
+        \\(Q\\) is quadratic in the input-dimension of a module \\(p_{in} \\times p_{in}\\)
+        and \\(H\\) in the output-dimension \\(p_{out} \\times p_{out}\\).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)`
+        y : torch.Tensor
+            labels `(batch, label_shape)`
+
+        Returns
+        -------
+        loss : torch.Tensor
+        H : `laplace.matrix.Kron`
+            Kronecker factored Hessian approximation.
+        """
+
+    @abstractmethod
+    def diag(self, x, y, **kwargs):
+        """Compute a diagonal Hessian approximation to \\(H\\) and is represented as a 
+        vector of the dimensionality of parameters \\(\\theta\\).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)`
+        y : torch.Tensor
+            labels `(batch, label_shape)`
+
+        Returns
+        -------
+        loss : torch.Tensor
+        H : torch.Tensor
+            vector representing the diagonal of H
+        """
+        pass
+
+
+class GGNInterface(CurvatureInterface):
+    """Generalized Gauss-Newton or Fisher Curvature Interface.
+    The GGN is equal to the Fisher information for the available likelihoods.
+    In addition to `CurvatureInterface`, methods for Jacobians are required by subclasses.
+
+    Parameters
+    ----------
+    model : torch.nn.Module or `laplace.feature_extractor.FeatureExtractor`
+        torch model (neural network)
+    likelihood : {'classification', 'regression'}
+    last_layer : bool, default=False
+        only consider curvature of last layer
+    stochastic : bool, default=False
+        Fisher if stochastic else GGN
+    """
+    def __init__(self, model, likelihood, last_layer=False, stochastic=False):
+        self.stochastic = stochastic
+        super().__init__(model, likelihood, last_layer)
+
+    def _get_full_ggn(self, Js, f, y):
+        """Compute full GGN from Jacobians.
+
+        Parameters
+        ----------
+        Js : torch.Tensor
+            Jacobians `(batch, parameters, outputs)`
+        f : torch.Tensor
+            functions `(batch, outputs)`
+        y : torch.Tensor
+            labels compatible with loss
+
+        Returns
+        -------
+        loss : torch.Tensor
+        H_ggn : torch.Tensor
+            full GGN approximation `(parameters, parameters)`
+        """
+        loss = self.factor * self.lossfunc(f, y)
+        if self.likelihood == 'regression':
+            H_ggn = torch.einsum('mkp,mkq->pq', Js, Js)
+        else:
+            # second derivative of log lik is diag(p) - pp^T
+            ps = torch.softmax(f, dim=-1)
+            H_lik = torch.diag_embed(ps) - torch.einsum('mk,mc->mck', ps, ps)
+            H_ggn = torch.einsum('mcp,mck,mkq->pq', Js, H_lik, Js)
+        return loss.detach(), H_ggn
+
+    def full(self, x, y, **kwargs):
+        """Compute the full GGN \\(P \\times P\\) matrix as Hessian approximation
+        \\(H_{ggn}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
+        For last-layer, reduced to \\(\\theta_{last}\\)
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)`
+        y : torch.Tensor
+            labels `(batch, label_shape)`
+
+        Returns
+        -------
+        loss : torch.Tensor
+        H_ggn : torch.Tensor
+            GGN `(parameters, parameters)`
+        """
+        if self.stochastic:
+            raise ValueError('Stochastic approximation not implemented for full GGN.')
+
+        if self.last_layer:
+            Js, f = self.last_layer_jacobians(self.model, x)
+        else:
+            Js, f = self.jacobians(self.model, x)
+        loss, H_ggn = self._get_full_ggn(Js, f, y)
+
+        return loss, H_ggn
+
+
+class EFInterface(CurvatureInterface):
+    """Interface for Empirical Fisher as Hessian approximation.
+    In addition to `CurvatureInterface`, methods for gradients are required by subclasses.
+
+    Parameters
+    ----------
+    model : torch.nn.Module or `laplace.feature_extractor.FeatureExtractor`
+        torch model (neural network)
+    likelihood : {'classification', 'regression'}
+    last_layer : bool, default=False
+        only consider curvature of last layer
+
+    Attributes
+    ----------
+    lossfunc : torch.nn.MSELoss or torch.nn.CrossEntropyLoss
+    factor : float
+        conversion factor between torch losses and base likelihoods
+        For example, \\(\\frac{1}{2}\\) to get to \\(\\mathcal{N}(f, 1)\\) from MSELoss.
+    """
+
+    def full(self, x, y, **kwargs):
+        """Compute the full EF \\(P \\times P\\) matrix as Hessian approximation
+        \\(H_{ef}\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
+        For last-layer, reduced to \\(\\theta_{last}\\)
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)`
+        y : torch.Tensor
+            labels `(batch, label_shape)`
+
+        Returns
+        -------
+        loss : torch.Tensor
+        H_ef : torch.Tensor
+            EF `(parameters, parameters)`
+        """
+        Gs, loss = self.gradients(x, y)
+        H_ef = Gs.T @ Gs
+        return self.factor * loss.detach(), self.factor * H_ef
diff --git a/laplace/feature_extractor.py b/laplace/feature_extractor.py
new file mode 100644
index 00000000..784d6370
--- /dev/null
+++ b/laplace/feature_extractor.py
@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+from typing import Tuple, Callable, Optional
+
+
+__all__ = ['FeatureExtractor']
+
+
+class FeatureExtractor(nn.Module):
+    """Feature extractor for a PyTorch neural network.
+    A wrapper which can return the output of the penultimate layer in addition to
+    the output of the last layer for each forward pass. If the name of the last
+    layer is not known, it can determine it automatically. It assumes that the
+    last layer is linear and that for every forward pass the last layer is the same.
+    If the name of the last layer is known, it can be passed as a parameter at
+    initilization; this is the safest way to use this class.
+    Based on https://gist.github.com/fkodom/27ed045c9051a39102e8bcf4ce31df76.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        PyTorch model
+    last_layer_name : str, default=None
+        if the name of the last layer is already known, otherwise it will
+        be determined automatically.
+    """
+    def __init__(self, model: nn.Module, last_layer_name: Optional[str] = None) -> None:
+        super().__init__()
+        self.model = model
+        self._features = dict()
+        if last_layer_name is None:
+            self.last_layer = None
+        else:
+            self.set_last_layer(last_layer_name)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass. If the last layer is not known yet, it will be
+        determined when this function is called for the first time.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            one batch of data to use as input for the forward pass
+        """
+        if self.last_layer is None:
+            # if this is the first forward pass and last layer is unknown
+            out = self.find_last_layer(x)
+        else:
+            # if last and penultimate layers are already known
+            out = self.model(x)
+        return out
+
+    def forward_with_features(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Forward pass which returns the output of the penultimate layer along
+        with the output of the last layer. If the last layer is not known yet,
+        it will be determined when this function is called for the first time.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            one batch of data to use as input for the forward pass
+        """
+        out = self.forward(x)
+        features = self._features[self._last_layer_name]
+        return out, features
+
+    def set_last_layer(self, last_layer_name: str) -> None:
+        """Set the last layer of the model by its name. This sets the forward
+        hook to get the output of the penultimate layer.
+
+        Parameters
+        ----------
+        last_layer_name : str
+            the name of the last layer (fixed in `model.named_modules()`).
+        """
+        # set last_layer attributes and check if it is linear
+        self._last_layer_name = last_layer_name
+        self.last_layer = dict(self.model.named_modules())[last_layer_name]
+        if not isinstance(self.last_layer, nn.Linear):
+            raise ValueError('Use model with a linear last layer.')
+
+        # set forward hook to extract features in future forward passes
+        self.last_layer.register_forward_hook(self._get_hook(last_layer_name))
+
+    def _get_hook(self, name: str) -> Callable:
+        def hook(_, input, __):
+            # only accepts one input (expects linear layer)
+            self._features[name] = input[0].detach()
+        return hook
+
+    def find_last_layer(self, x: torch.Tensor) -> torch.Tensor:
+        """Automatically determines the last layer of the model with one
+        forward pass. It assumes that the last layer is the same for every
+        forward pass and that it is an instance of `torch.nn.Linear`.
+        Might not work with every architecture, but is tested with all PyTorch
+        torchvision classification models (besides SqueezeNet, which has no
+        linear last layer).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            one batch of data to use as input for the forward pass
+        """
+        if self.last_layer is not None:
+            raise ValueError('Last layer is already known.')
+
+        act_out = dict()
+        def get_act_hook(name):
+            def act_hook(_, input, __):
+                # only accepts one input (expects linear layer)
+                try:
+                    act_out[name] = input[0].detach()
+                except (IndexError, AttributeError):
+                    act_out[name] = None
+                # remove hook
+                handles[name].remove()
+            return act_hook
+
+        # set hooks for all modules
+        handles = dict()
+        for name, module in self.model.named_modules():
+            handles[name] = module.register_forward_hook(get_act_hook(name))
+
+        # check if model has more than one module
+        # (there might be pathological exceptions)
+        if len(handles) <= 2:
+            raise ValueError('The model only has one module.')
+
+        # forward pass to find execution order
+        out = self.model(x)
+
+        # find the last layer, store features, return output of forward pass
+        keys = list(act_out.keys())
+        for key in reversed(keys):
+            layer = dict(self.model.named_modules())[key]
+            if len(list(layer.children())) == 0:
+                self.set_last_layer(key)
+
+                # save features from first forward pass
+                self._features[key] = act_out[key]
+
+                return out
+
+        raise ValueError('Something went wrong (all modules have children).')
diff --git a/laplace/laplace.py b/laplace/laplace.py
new file mode 100644
index 00000000..508631cb
--- /dev/null
+++ b/laplace/laplace.py
@@ -0,0 +1,31 @@
+from laplace.baselaplace import BaseLaplace
+from laplace import *
+
+
+def Laplace(model, likelihood, subset_of_weights='last_layer', hessian_structure='kron',
+            *args, **kwargs):
+    """Simplified Laplace access using strings instead of different classes.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+    likelihood : {'classification', 'regression'}
+    subset_of_weights : {'last_layer', 'all'}, default='last_layer'
+        subset of weights to consider for inference
+    hessian_structure : {'diag', 'kron', 'full'}, default='kron'
+        structure of the Hessian approximation
+
+    Returns
+    -------
+    laplace : BaseLaplace
+        chosen subclass of BaseLaplace instantiated with additional arguments
+    """
+    laplace_map = {subclass._key: subclass for subclass in _all_subclasses(BaseLaplace)
+                   if hasattr(subclass, '_key')}
+    laplace_class = laplace_map[(subset_of_weights, hessian_structure)]
+    return laplace_class(model, likelihood, *args, **kwargs)
+
+
+def _all_subclasses(cls):
+    return set(cls.__subclasses__()).union(
+        [s for c in cls.__subclasses__() for s in _all_subclasses(c)])
diff --git a/laplace/lllaplace.py b/laplace/lllaplace.py
new file mode 100644
index 00000000..79fd460b
--- /dev/null
+++ b/laplace/lllaplace.py
@@ -0,0 +1,199 @@
+import torch
+from torch.nn.utils import parameters_to_vector, vector_to_parameters
+
+from laplace.baselaplace import BaseLaplace, FullLaplace, KronLaplace, DiagLaplace
+from laplace.feature_extractor import FeatureExtractor
+
+from laplace.matrix import Kron
+from laplace.curvature import BackPackGGN
+
+
+__all__ = ['FullLLLaplace', 'KronLLLaplace', 'DiagLLLaplace']
+
+
+class LLLaplace(BaseLaplace):
+    """Baseclass for all last-layer Laplace approximations in this library.
+    Subclasses specify the structure of the Hessian approximation.
+    See `BaseLaplace` for the full interface.
+
+    A Laplace approximation is represented by a MAP which is given by the
+    `model` parameter and a posterior precision or covariance specifying
+    a Gaussian distribution \\(\\mathcal{N}(\\theta_{MAP}, P^{-1})\\).
+    Here, only the parameters of the last layer of the neural network
+    are treated probabilistically.
+    The goal of this class is to compute the posterior precision \\(P\\)
+    which sums as
+    \\[
+        P = \\sum_{n=1}^N \\nabla^2_\\theta \\log p(\\mathcal{D}_n \\mid \\theta)
+        \\vert_{\\theta_{MAP}} + \\nabla^2_\\theta \\log p(\\theta) \\vert_{\\theta_{MAP}}.
+    \\]
+    Every subclass implements different approximations to the log likelihood Hessians,
+    for example, a diagonal one. The prior is assumed to be Gaussian and therefore we have
+    a simple form for \\(\\nabla^2_\\theta \\log p(\\theta) \\vert_{\\theta_{MAP}} = P_0 \\).
+    In particular, we assume a scalar or diagonal prior precision so that in
+    all cases \\(P_0 = \\textrm{diag}(p_0)\\) and the structure of \\(p_0\\) can be varied.
+
+    Parameters
+    ----------
+    model : torch.nn.Module or `laplace.feature_extractor.FeatureExtractor`
+    likelihood : {'classification', 'regression'}
+        determines the log likelihood Hessian approximation
+    sigma_noise : torch.Tensor or float, default=1
+        observation noise for the regression setting; must be 1 for classification
+    prior_precision : torch.Tensor or float, default=1
+        prior precision of a Gaussian prior (= weight decay);
+        can be scalar, per-layer, or diagonal in the most general case
+    prior_mean : torch.Tensor or float, default=0
+        prior mean of a Gaussian prior, useful for continual learning
+    temperature : float, default=1
+        temperature of the likelihood; lower temperature leads to more
+        concentrated posterior and vice versa.
+    backend : subclasses of `laplace.curvature.CurvatureInterface`
+        backend for access to curvature/Hessian approximations
+    last_layer_name: str, default=None
+        name of the model's last layer, if None it will be determined automatically
+    backend_kwargs : dict, default=None
+        arguments passed to the backend on initialization, for example to
+        set the number of MC samples for stochastic approximations.
+    """
+    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
+                 prior_mean=0., temperature=1., backend=BackPackGGN, last_layer_name=None,
+                 backend_kwargs=None):
+        super().__init__(model, likelihood, sigma_noise=sigma_noise, prior_precision=1.,
+                         prior_mean=0., temperature=temperature, backend=backend,
+                         backend_kwargs=backend_kwargs)
+        self.model = FeatureExtractor(model, last_layer_name=last_layer_name)
+        if self.model.last_layer is None:
+            self.mean = None
+            self.n_params = None
+            self.n_layers = None
+            # ignore checks of prior mean setter temporarily, check on .fit()
+            self._prior_precision = prior_precision
+            self._prior_mean = prior_mean
+        else:
+            self.mean = parameters_to_vector(self.model.last_layer.parameters()).detach()
+            self.n_params = len(self.mean)
+            self.n_layers = len(list(self.model.last_layer.parameters()))
+            self.prior_precision = prior_precision
+            self.prior_mean = prior_mean
+        self._backend_kwargs['last_layer'] = True
+
+    def fit(self, train_loader):
+        """Fit the local Laplace approximation at the parameters of the model.
+
+        Parameters
+        ----------
+        train_loader : torch.data.utils.DataLoader
+            each iterate is a training batch (X, y);
+            `train_loader.dataset` needs to be set to access \\(N\\), size of the data set
+        """
+        if self.H is not None:
+            raise ValueError('Already fit.')
+
+        self.model.eval()
+
+        if self.model.last_layer is None:
+            X, _ = next(iter(train_loader))
+            with torch.no_grad():
+                self.model.find_last_layer(X[:1].to(self._device))
+            self.mean = parameters_to_vector(self.model.last_layer.parameters()).detach()
+            self.n_params = len(self.mean)
+            self.n_layers = len(list(self.model.last_layer.parameters()))
+            # here, check the already set prior precision again
+            self.prior_precision = self._prior_precision
+            self.prior_mean = self._prior_mean
+
+        super().fit(train_loader)
+
+    def _glm_predictive_distribution(self, X):
+        Js, f_mu = self.backend.last_layer_jacobians(self.model, X)
+        f_var = self.functional_variance(Js)
+        return f_mu.detach(), f_var.detach()
+
+    def _nn_predictive_samples(self, X, n_samples=100):
+        fs = list()
+        for sample in self.sample(n_samples):
+            vector_to_parameters(sample, self.model.last_layer.parameters())
+            fs.append(self.model(X.to(self._device)).detach())
+        vector_to_parameters(self.mean, self.model.last_layer.parameters())
+        fs = torch.stack(fs)
+        if self.likelihood == 'classification':
+            fs = torch.softmax(fs, dim=-1)
+        return fs
+
+    @property
+    def prior_precision_diag(self):
+        """Obtain the diagonal prior precision \\(p_0\\) constructed from either
+        a scalar or diagonal prior precision.
+
+        Returns
+        -------
+        prior_precision_diag : torch.Tensor
+        """
+        if len(self.prior_precision) == 1:  # scalar
+            return self.prior_precision * torch.ones_like(self.mean)
+
+        elif len(self.prior_precision) == self.n_params:  # diagonal
+            return self.prior_precision
+
+        else:
+            raise ValueError('Mismatch of prior and model. Diagonal or scalar prior.')
+
+
+class FullLLLaplace(LLLaplace, FullLaplace):
+    """Last-layer Laplace approximation with full, i.e., dense, log likelihood Hessian approximation
+    and hence posterior precision. Based on the chosen `backend` parameter, the full
+    approximation can be, for example, a generalized Gauss-Newton matrix.
+    Mathematically, we have \\(P \\in \\mathbb{R}^{P \\times P}\\).
+    See `FullLaplace`, `LLLaplace`, and `BaseLaplace` for the full interface.
+    """
+    # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
+    _key = ('last_layer', 'full')
+
+    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
+                 prior_mean=0., temperature=1., backend=BackPackGGN, last_layer_name=None,
+                 backend_kwargs=None):
+        super().__init__(model, likelihood, sigma_noise, prior_precision,
+                         prior_mean, temperature, backend, last_layer_name, backend_kwargs)
+
+
+class KronLLLaplace(LLLaplace, KronLaplace):
+    """Last-layer Laplace approximation with Kronecker factored log likelihood Hessian approximation
+    and hence posterior precision.
+    Mathematically, we have for the last parameter group, i.e., torch.nn.Linear,
+    that \\P\\approx Q \\otimes H\\.
+    See `KronLaplace`, `LLLaplace`, and `BaseLaplace` for the full interface and see
+    `laplace.matrix.Kron` and `laplace.matrix.KronDecomposed` for the structure of
+    the Kronecker factors. `Kron` is used to aggregate factors by summing up and
+    `KronDecomposed` is used to add the prior, a Hessian factor (e.g. temperature),
+    and computing posterior covariances, marginal likelihood, etc.
+    Use of `damping` is possible by initializing or setting `damping=True`.
+    """
+    # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
+    _key = ('last_layer', 'kron')
+
+    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
+                 prior_mean=0., temperature=1., backend=BackPackGGN, last_layer_name=None,
+                 damping=False, **backend_kwargs):
+        self.damping = damping
+        super().__init__(model, likelihood, sigma_noise, prior_precision,
+                         prior_mean, temperature, backend, last_layer_name, backend_kwargs)
+
+    def _init_H(self):
+        self.H = Kron.init_from_model(self.model.last_layer, self._device)
+
+
+class DiagLLLaplace(LLLaplace, DiagLaplace):
+    """Last-layer Laplace approximation with diagonal log likelihood Hessian approximation
+    and hence posterior precision.
+    Mathematically, we have \\(P \\approx \\textrm{diag}(P)\\).
+    See `DiagLaplace`, `LLLaplace`, and `BaseLaplace` for the full interface.
+    """
+    # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
+    _key = ('last_layer', 'diag')
+
+    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
+                 prior_mean=0., temperature=1., backend=BackPackGGN, last_layer_name=None,
+                 backend_kwargs=None):
+        super().__init__(model, likelihood, sigma_noise, prior_precision,
+                         prior_mean, temperature, backend, last_layer_name, backend_kwargs)
diff --git a/laplace/matrix.py b/laplace/matrix.py
new file mode 100644
index 00000000..61c07ab5
--- /dev/null
+++ b/laplace/matrix.py
@@ -0,0 +1,461 @@
+from math import pow
+import torch
+import numpy as np
+from typing import Union
+
+from laplace.utils import _is_valid_scalar, symeig, kron, block_diag
+
+
+class Kron:
+    """Kronecker factored approximate curvature representation for a corresponding
+    neural network.
+    Each element in `kfacs` is either a tuple or single matrix.
+    A tuple represents two Kronecker factors \\(Q\\), and \\(H\\) and a single element
+    is just a full block Hessian approximation.
+
+    Parameters
+    ----------
+    kfacs : list[Tuple]
+        each element in the list is a Tuple of two Kronecker factors Q, H
+        or a single matrix approximating the Hessian (in case of bias, for example)
+    """
+    def __init__(self, kfacs):
+        self.kfacs = kfacs
+
+    @classmethod
+    def init_from_model(cls, model, device):
+        """Initialize Kronecker factors based on a models architecture.
+
+        Parameters
+        ----------
+        model : torch.nn.Module
+        device : torch.device
+
+        Returns
+        -------
+        kron : Kron
+        """
+        kfacs = list()
+        for p in model.parameters():
+            if p.ndim == 1:  # bias
+                P = p.size(0)
+                kfacs.append([torch.zeros(P, P, device=device)])
+            elif 4 >= p.ndim >= 2:  # fully connected or conv
+                if p.ndim == 2:  # fully connected
+                    P_in, P_out = p.size()
+                elif p.ndim > 2:
+                    P_in, P_out = p.shape[0], np.prod(p.shape[1:])
+
+                kfacs.append([
+                    torch.zeros(P_in, P_in, device=device),
+                    torch.zeros(P_out, P_out, device=device)
+                ])
+            else:
+                raise ValueError('Invalid parameter shape in network.')
+        return cls(kfacs)
+
+    def __add__(self, other):
+        """Add up Kronecker factors `self` and `other`.
+
+        Parameters
+        ----------
+        other : Kron
+
+        Returns
+        -------
+        kron : Kron
+        """
+        if not isinstance(other, Kron):
+            raise ValueError('Can only add Kron to Kron.')
+
+        kfacs = [[Hi.add(Hj) for Hi, Hj in zip(Fi, Fj)]
+                 for Fi, Fj in zip(self.kfacs, other.kfacs)]
+        return Kron(kfacs)
+
+    def __mul__(self, scalar: Union[float, torch.Tensor]):
+        """Multiply all Kronecker factors by scalar.
+        The multiplication is distributed across the number of factors
+        using `pow(scalar, 1 / len(F))`. `len(F)` is either `1` or `2`.
+
+        Parameters
+        ----------
+        scalar : float, torch.Tensor
+
+        Returns
+        -------
+        kron : Kron
+        """
+        if not _is_valid_scalar(scalar):
+            raise ValueError('Input not valid python or torch scalar.')
+
+        # distribute factors evenly so that each group is multiplied by factor
+        kfacs = [[pow(scalar, 1/len(F)) * Hi for Hi in F] for F in self.kfacs]
+        return Kron(kfacs)
+
+    def __len__(self):
+        return len(self.kfacs)
+
+    def decompose(self, damping=False):
+        """Eigendecompose Kronecker factors and turn into `KronDecomposed`.
+        Parameters
+        ----------
+        damping : bool
+            use damping
+
+        Returns
+        -------
+        kron_decomposed : KronDecomposed
+        """
+        eigvecs, eigvals = list(), list()
+        for F in self.kfacs:
+            Qs, ls = list(), list()
+            for Hi in F:
+                l, Q = symeig(Hi)
+                Qs.append(Q)
+                ls.append(l)
+            eigvecs.append(Qs)
+            eigvals.append(ls)
+        return KronDecomposed(eigvecs, eigvals, damping=damping)
+
+    def _bmm(self, W: torch.Tensor) -> torch.Tensor:
+        """Implementation of `bmm` which casts the parameters to the right shape.
+
+        Parameters
+        ----------
+        W : torch.Tensor
+            matrix `(batch, classes, params)`
+
+        Returns
+        -------
+        SW : torch.Tensor
+            result `(batch, classes, params)`
+        """
+        # self @ W[batch, k, params]
+        assert len(W.size()) == 3
+        B, K, P = W.size()
+        W = W.reshape(B * K, P)
+        cur_p = 0
+        SW = list()
+        for Fs in self.kfacs:
+            if len(Fs) == 1:
+                Q = Fs[0]
+                p = len(Q)
+                W_p = W[:, cur_p:cur_p+p].T
+                SW.append((Q @ W_p).T)
+                cur_p += p
+            elif len(Fs) == 2:
+                Q, H = Fs
+                p_in, p_out = len(Q), len(H)
+                p = p_in * p_out
+                W_p = W[:, cur_p:cur_p+p].reshape(B * K, p_in, p_out)
+                SW.append((Q @ W_p @ H.T).reshape(B * K, p_in * p_out))
+                cur_p += p
+            else:
+                raise AttributeError('Shape mismatch')
+        SW = torch.cat(SW, dim=1).reshape(B, K, P)
+        return SW
+
+    def bmm(self, W: torch.Tensor, exponent: float = 1) -> torch.Tensor:
+        """Batched matrix multiplication with the Kronecker factors.
+        If Kron is `H`, we compute `H @ W`.
+        This is useful for computing the predictive or a regularization
+        based on Kronecker factors as in continual learning.
+
+        Parameters
+        ----------
+        W : torch.Tensor
+            matrix `(batch, classes, params)`
+        exponent: float, default=1
+            only can be `1` for Kron, requires `KronDecomposed` for other
+            exponent values of the Kronecker factors.
+
+        Returns
+        -------
+        SW : torch.Tensor
+            result `(batch, classes, params)`
+        """
+        if exponent != 1:
+            raise ValueError('Only supported after decomposition.')
+        if W.ndim == 1:
+            return self._bmm(W.unsqueeze(0).unsqueeze(0)).squeeze()
+        elif W.ndim == 2:
+            return self._bmm(W.unsqueeze(1)).squeeze()
+        elif W.ndim == 3:
+            return self._bmm(W)
+        else:
+            raise ValueError('Invalid shape for W')
+
+    def logdet(self) -> torch.Tensor:
+        """Compute log determinant of the Kronecker factors and sums them up.
+        This corresponds to the log determinant of the entire Hessian approximation.
+
+        Returns
+        -------
+        logdet : torch.Tensor
+        """
+        logdet = 0
+        for F in self.kfacs:
+            if len(F) == 1:
+                logdet += F[0].logdet()
+            else:  # len(F) == 2
+                Hi, Hj = F
+                p_in, p_out = len(Hi), len(Hj)
+                logdet += p_out * Hi.logdet() + p_in * Hj.logdet()
+        return logdet
+
+    def diag(self) -> torch.Tensor:
+        """Extract diagonal of the entire Kronecker factorization.
+
+        Returns
+        -------
+        diag : torch.Tensor
+        """
+        diags = list()
+        for F in self.kfacs:
+            if len(F) == 1:
+                diags.append(F[0].diagonal())
+            else:
+                diags.append(torch.ger(F[0].diagonal(), F[1].diagonal()).flatten())
+        return torch.cat(diags)
+
+    def to_matrix(self) -> torch.Tensor:
+        """Make the Kronecker factorization dense by computing the kronecker product.
+        Warning: this should only be used for testing purposes as it will allocate
+        large amounts of memory for big architectures.
+
+        Returns
+        -------
+        block_diag : torch.Tensor
+        """
+        blocks = list()
+        for F in self.kfacs:
+            if len(F) == 1:
+                blocks.append(F[0])
+            else:
+                blocks.append(kron(F[0], F[1]))
+        return block_diag(blocks)
+
+    # for commutative operations
+    __radd__ = __add__
+    __rmul__ = __mul__
+
+
+class KronDecomposed:
+    """Decomposed Kronecker factored approximate curvature representation
+    for a corresponding neural network.
+    Each matrix in `Kron` is decomposed to obtain `KronDecomposed`.
+    Front-loading decomposition allows cheap repeated computation
+    of inverses and log determinants.
+    In contrast to `Kron`, we can add scalar or layerwise scalars but
+    we cannot add other `Kron` or `KronDecomposed` anymore.
+
+    Parameters
+    ----------
+    eigenvectors : list[Tuple[torch.Tensor]]
+        eigenvectors corresponding to matrices in a corresponding `Kron`
+    eigenvalues : list[Tuple[torch.Tensor]]
+        eigenvalues corresponding to matrices in a corresponding `Kron`
+    deltas : torch.Tensor
+        addend for each group of Kronecker factors representing, for example,
+        a prior precision
+    dampen : bool, default=False
+        use dampen approximation mixing prior and Kron partially multiplicatively
+    """
+
+    def __init__(self, eigenvectors, eigenvalues, deltas=None, damping=False):
+        self.eigenvectors = eigenvectors
+        self.eigenvalues = eigenvalues
+        device = eigenvectors[0][0].device
+        if deltas is None:
+            self.deltas = torch.zeros(len(self), device=device)
+        else:
+            self._check_deltas(deltas)
+            self.deltas = deltas
+        self.damping = damping
+
+    def detach(self):
+        self.deltas = self.deltas.detach()
+        return self
+
+    def _check_deltas(self, deltas: torch.Tensor):
+        if not isinstance(deltas, torch.Tensor):
+            raise ValueError('Can only add torch.Tensor to KronDecomposed.')
+
+        if (deltas.ndim == 0  # scalar
+            or (deltas.ndim == 1  # vector of length 1 or len(self)
+                and (len(deltas) == 1 or len(deltas) == len(self)))):
+            return
+        else:
+            raise ValueError('Invalid shape of delta added to KronDecomposed.')
+
+    def __add__(self, deltas: torch.Tensor):
+        """Add scalar per layer or only scalar to Kronecker factors.
+
+        Parameters
+        ----------
+        deltas : torch.Tensor
+            either same length as `eigenvalues` or scalar.
+
+        Returns
+        -------
+        kron : KronDecomposed
+        """
+        self._check_deltas(deltas)
+        return KronDecomposed(self.eigenvectors, self.eigenvalues, self.deltas + deltas)
+
+    def __mul__(self, scalar):
+        """Multiply by a scalar by changing the eigenvalues.
+        Same as for the case of `Kron`.
+
+        Parameters
+        ----------
+        scalar : torch.Tensor or float
+
+        Returns
+        -------
+        kron : KronDecomposed
+        """
+        if not _is_valid_scalar(scalar):
+            raise ValueError('Invalid argument, can only multiply Kron with scalar.')
+
+        eigenvalues = [[pow(scalar, 1/len(ls)) * l for l in ls] for ls in self.eigenvalues]
+        return KronDecomposed(self.eigenvectors, eigenvalues, self.deltas)
+
+    def __len__(self) -> int:
+        return len(self.eigenvalues)
+
+    def logdet(self) -> torch.Tensor:
+        """Compute log determinant of the Kronecker factors and sums them up.
+        This corresponds to the log determinant of the entire Hessian approximation.
+        In contrast to `Kron.logdet()`, additive `deltas` corresponding to prior
+        precisions are added.
+
+        Returns
+        -------
+        logdet : torch.Tensor
+        """
+        logdet = 0
+        for ls, delta in zip(self.eigenvalues, self.deltas):
+            if len(ls) == 1:  # not KFAC just full
+                logdet += torch.log(ls[0] + delta).sum()
+            elif len(ls) == 2:
+                l1, l2 = ls
+                if self.damping:
+                    l1d, l2d = l1 + torch.sqrt(delta), l2 + torch.sqrt(delta)
+                    logdet += torch.log(torch.ger(l1d, l2d)).sum()
+                else:
+                    logdet += torch.log(torch.ger(l1, l2) + delta).sum()
+            else:
+                raise ValueError('Too many Kronecker factors. Something went wrong.')
+        return logdet
+
+    def _bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
+        """Implementation of `bmm`, i.e., `self ** exponent @ W`.
+
+        Parameters
+        ----------
+        W : torch.Tensor
+            matrix `(batch, classes, params)`
+        exponent : float
+            exponent on `self`
+
+        Returns
+        -------
+        SW : torch.Tensor
+            result `(batch, classes, params)`
+        """
+        # self @ W[batch, k, params]
+        assert len(W.size()) == 3
+        B, K, P = W.size()
+        W = W.reshape(B * K, P)
+        cur_p = 0
+        SW = list()
+        for ls, Qs, delta in zip(self.eigenvalues, self.eigenvectors, self.deltas):
+            if len(ls) == 1:
+                Q, l, p = Qs[0], ls[0], len(ls[0])
+                ldelta_exp = torch.pow(l + delta, exponent).reshape(-1, 1)
+                W_p = W[:, cur_p:cur_p+p].T
+                SW.append((Q @ (ldelta_exp * (Q.T @ W_p))).T)
+                cur_p += p
+            elif len(ls) == 2:
+                Q1, Q2 = Qs
+                l1, l2 = ls
+                p = len(l1) * len(l2)
+                if self.damping:
+                    l1d, l2d = l1 + torch.sqrt(delta), l2 + torch.sqrt(delta)
+                    ldelta_exp = torch.pow(torch.ger(l1d, l2d), exponent).unsqueeze(0)
+                else:
+                    ldelta_exp = torch.pow(torch.ger(l1, l2) + delta, exponent).unsqueeze(0)
+                p_in, p_out = len(l1), len(l2)
+                W_p = W[:, cur_p:cur_p+p].reshape(B * K, p_in, p_out)
+                W_p = (Q1.T @ W_p @ Q2) * ldelta_exp
+                W_p = Q1 @ W_p @ Q2.T
+                SW.append(W_p.reshape(B * K, p_in * p_out))
+                cur_p += p
+            else:
+                raise AttributeError('Shape mismatch')
+        SW = torch.cat(SW, dim=1).reshape(B, K, P)
+        return SW
+
+    def inv_square_form(self, W: torch.Tensor) -> torch.Tensor:
+        # W either Batch x K x params or Batch x params
+        SW = self._bmm(W, exponent=-1)
+        return torch.bmm(W, SW.transpose(1, 2))
+
+    def bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
+        """Batched matrix multiplication with the decomposed Kronecker factors.
+        This is useful for computing the predictive or a regularization loss.
+        Compared to `Kron.bmm`, a prior can be added here in form of `deltas`
+        and the exponent can be other than just 1.
+        Computes \\(H^{exponent} W\\).
+
+        Parameters
+        ----------
+        W : torch.Tensor
+            matrix `(batch, classes, params)`
+        exponent: float, default=1
+
+        Returns
+        -------
+        SW : torch.Tensor
+            result `(batch, classes, params)`
+        """
+        if W.ndim == 1:
+            return self._bmm(W.unsqueeze(0).unsqueeze(0), exponent).squeeze()
+        elif W.ndim == 2:
+            return self._bmm(W.unsqueeze(1), exponent).squeeze()
+        elif W.ndim == 3:
+            return self._bmm(W, exponent)
+        else:
+            raise ValueError('Invalid shape for W')
+
+    def to_matrix(self, exponent: float = 1) -> torch.Tensor:
+        """Make the Kronecker factorization dense by computing the kronecker product.
+        Warning: this should only be used for testing purposes as it will allocate
+        large amounts of memory for big architectures.
+
+        Returns
+        -------
+        block_diag : torch.Tensor
+        """
+        blocks = list()
+        for Qs, ls, delta in zip(self.eigenvectors, self.eigenvalues, self.deltas):
+            if len(ls) == 1:
+                Q, l = Qs[0], ls[0]
+                blocks.append(Q @ torch.diag(torch.pow(l + delta, exponent)) @ Q.T)
+            else:
+                Q1, Q2 = Qs
+                l1, l2 = ls
+                Q = kron(Q1, Q2)
+                if self.damping:
+                    delta_sqrt = torch.sqrt(delta)
+                    l = torch.pow(torch.ger(l1 + delta_sqrt, l2 + delta_sqrt), exponent)
+                else:
+                    l = torch.pow(torch.ger(l1, l2) + delta, exponent)
+                L = torch.diag(l.flatten())
+                blocks.append(Q @ L @ Q.T)
+        return block_diag(blocks)
+
+    # for commutative operations
+    __radd__ = __add__
+    __rmul__ = __mul__
diff --git a/laplace/utils.py b/laplace/utils.py
new file mode 100644
index 00000000..e096cd18
--- /dev/null
+++ b/laplace/utils.py
@@ -0,0 +1,175 @@
+import logging
+from typing import Union
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.nn import BatchNorm1d, BatchNorm2d, BatchNorm3d
+from torch.distributions.multivariate_normal import _precision_to_scale_tril
+
+
+def get_nll(out_dist, targets):
+    return F.nll_loss(torch.log(out_dist), targets)
+
+
+@torch.no_grad()
+def validate(laplace, val_loader, pred_type='glm', link_approx='probit', n_samples=100):
+    laplace.model.eval()
+    outputs = list()
+    targets = list()
+    for X, y in val_loader:
+        X, y = X.to(laplace._device), y.to(laplace._device)
+        out = laplace(X, pred_type=pred_type, link_approx=link_approx, n_samples=n_samples)
+        outputs.append(out)
+        targets.append(y)
+    return torch.cat(outputs, dim=0), torch.cat(targets, dim=0)
+
+
+def parameters_per_layer(model):
+    """Get number of parameters per layer.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+
+    Returns
+    -------
+    params_per_layer : list[int]
+    """
+    return [np.prod(p.shape) for p in model.parameters()]
+
+
+def invsqrt_precision(M):
+    """Compute ``M^{-0.5}`` as a tridiagonal matrix.
+
+    Parameters
+    ----------
+    M : torch.Tensor
+
+    Returns
+    -------
+    M_invsqrt : torch.Tensor
+    """
+    return _precision_to_scale_tril(M)
+
+
+def _is_batchnorm(module):
+    if isinstance(module, BatchNorm1d) or \
+        isinstance(module, BatchNorm2d) or \
+            isinstance(module, BatchNorm3d):
+        return True
+    return False
+
+
+def _is_valid_scalar(scalar: Union[float, int, torch.Tensor]) -> bool:
+    if np.isscalar(scalar) and np.isreal(scalar):
+        return True
+    elif torch.is_tensor(scalar) and scalar.ndim <= 1:
+        if scalar.ndim == 1 and len(scalar) != 1:
+            return False
+        return True
+    return False
+
+
+def kron(t1, t2):
+    """Computes the Kronecker product between two tensors.
+
+    Parameters
+    ----------
+    t1 : torch.Tensor
+    t2 : torch.Tensor
+
+    Returns
+    -------
+    kron_product : torch.Tensor
+    """
+    t1_height, t1_width = t1.size()
+    t2_height, t2_width = t2.size()
+    out_height = t1_height * t2_height
+    out_width = t1_width * t2_width
+
+    tiled_t2 = t2.repeat(t1_height, t1_width)
+    expanded_t1 = (
+        t1.unsqueeze(2)
+          .unsqueeze(3)
+          .repeat(1, t2_height, t2_width, 1)
+          .view(out_height, out_width)
+    )
+
+    return expanded_t1 * tiled_t2
+
+
+def diagonal_add_scalar(X, value):
+    """Add scalar value `value` to diagonal of `X`.
+
+    Parameters
+    ----------
+    X : torch.Tensor
+    value : torch.Tensor or float
+
+    Returns
+    -------
+    X_add_scalar : torch.Tensor
+    """
+    if not X.device == torch.device('cpu'):
+        indices = torch.cuda.LongTensor([[i, i] for i in range(X.shape[0])])
+    else:
+        indices = torch.LongTensor([[i, i] for i in range(X.shape[0])])
+    values = X.new_ones(X.shape[0]).mul(value)
+    return X.index_put(tuple(indices.t()), values, accumulate=True)
+
+
+def symeig(M):
+    """Symetric eigendecomposition avoiding failure cases by
+    adding and removing jitter to the diagonal.
+
+    Parameters
+    ----------
+    M : torch.Tensor
+
+    Returns
+    -------
+    L : torch.Tensor
+        eigenvalues
+    W : torch.Tensor
+        eigenvectors
+    """
+    try:
+        L, W = torch.symeig(M, eigenvectors=True)
+    except RuntimeError:  # did not converge
+        logging.info('SYMEIG: adding jitter, did not converge.')
+        # use W L W^T + I = W (L + I) W^T
+        M = diagonal_add_scalar(M, value=1.)
+        try:
+            L, W = torch.symeig(M, eigenvectors=True)
+            L -= 1.
+        except RuntimeError:
+            stats = f'diag: {M.diagonal()}, max: {M.abs().max()}, '
+            stats = stats + f'min: {M.abs().min()}, mean: {M.abs().mean()}'
+            logging.info(f'SYMEIG: adding jitter failed. Stats: {stats}')
+            exit()
+    # eigenvalues of symeig at least 0
+    L = L.clamp(min=0.0)
+    L[torch.isnan(L)] = 0.0
+    W[torch.isnan(W)] = 0.0
+    return L, W
+
+
+def block_diag(blocks):
+    """Compose block-diagonal matrix of individual blocks.
+
+    Parameters
+    ----------
+    blocks : list[torch.Tensor]
+
+    Returns
+    -------
+    M : torch.Tensor
+    """
+    P = sum([b.shape[0] for b in blocks])
+    M = torch.zeros(P, P)
+    p_cur = 0
+    for block in blocks:
+        p_block = block.shape[0]
+        M[p_cur:p_cur+p_block, p_cur:p_cur+p_block] = block
+        p_cur += p_block
+    return M
diff --git a/regression_example.md b/regression_example.md
new file mode 100644
index 00000000..ac8727aa
--- /dev/null
+++ b/regression_example.md
@@ -0,0 +1,109 @@
+### Full example: *post-hoc* optimization of the marginal likelihood and prediction
+
+#### Sinusoidal toy data
+We show how the marginal likelihood can be used after training a MAP network on a simple sinusoidal regression task. 
+Subsequently, we use the optimized LA to predict which provides uncertainty on top of the MAP prediction.
+First, we set up the training data for the problem with observation noise \\(\\sigma=0.3\\):
+```python
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from laplace import Laplace
+
+n_epochs = 1000
+batch_size = 150  # full batch
+true_sigma_noise = 0.3
+
+# create simple sinusoid data set
+X_train = (torch.rand(150) * 8).unsqueeze(-1)
+y_train = torch.sin(X_train) + torch.randn_like(X_train) * true_sigma_noise
+train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=batch_size)
+X_test = torch.linspace(-5, 13, 500).unsqueeze(-1)  # +-5 on top of the training X-range
+```
+
+#### Training a MAP
+We now use `pytorch` to train a neural network with single hidden layer and Tanh activation.
+This is standard so nothing new here, yet:
+```python
+# create and train MAP model
+model = torch.nn.Sequential(torch.nn.Linear(1, 50),
+                            torch.nn.Tanh(),
+                            torch.nn.Linear(50, 1))
+
+criterion = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), weight_decay=5e-4, lr=1e-2)
+for i in range(n_epochs):
+    for X, y in train_loader:
+        optimizer.zero_grad()
+        loss = criterion(model(X), y)
+        loss.backward()
+        optimizer.step()
+```
+
+#### Fitting and optimizing the Laplace approximation using empirical Bayes
+With the MAP-trained model at hand, we can estimate the prior precision and observation noise
+using empirical Bayes after training.
+The `Laplace` method is called to construct a LA for `'regression'` with `'all'` weights.
+As default `Laplace` returns a Kronecker factored LA.
+We fit the LA to the training data and initialize `log_prior` and `log_sigma`.
+Using Adam, we minimize the negative log marginal likelihood for `n_epochs`.
+
+```python
+la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
+la.fit(train_loader)
+log_prior, log_sigma = torch.ones(1, requires_grad=True), torch.ones(1, requires_grad=True)
+hyper_optimizer = torch.optim.Adam([log_prior, log_sigma], lr=1e-1)
+for i in range(n_epochs):
+    hyper_optimizer.zero_grad()
+    neg_marglik = - la.log_marginal_likelihood(log_prior.exp(), log_sigma.exp())
+    neg_marglik.backward()
+    hyper_optimizer.step()
+```
+
+The obtained observation noise is close to the ground truth with a value of \\(\\sigma \\approx 0.28\\)
+without the need for any validation data.
+The resulting prior precision is \\(\\delta \\approx 0.18\\).
+
+#### Bayesian predictive
+Lastly, we compare the MAP prediction to the obtained LA prediction.
+For LA, we have a closed-form predictive distribution on the output \\(f\\) which is a Gaussian
+\\(\\mathcal{N}(f(x;\\theta_{MAP}), \\mathbb{V}[f] + \\sigma^2)\\):
+
+```python
+x = X_test.flatten().cpu().numpy()
+f_mu, f_var = la(X_test)
+f_mu = f_mu.squeeze().detach().cpu().numpy()
+f_sigma = f_var.squeeze().sqrt().cpu().numpy()
+pred_std = np.sqrt(f_sigma**2 + la.sigma_noise.item()**2)
+```
+
+In comparison to the MAP, the predictive shows useful uncertainties:
+
+```python
+import matplotlib.pyplot as plt
+fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, sharey=True,
+                               figsize=(4.5, 2.8))
+ax1.set_title('MAP')
+ax1.scatter(X_train.flatten(), y_train.flatten(), alpha=0.7, color='tab:orange')
+ax1.plot(x, f_mu, color='black', label='$f_{MAP}$')
+ax1.legend()
+
+ax2.set_title('LA')
+ax2.scatter(X_train.flatten(), y_train.flatten(), alpha=0.7, color='tab:orange')
+ax2.plot(x, f_mu, label='$\mathbb{E}[f]$')
+ax2.fill_between(x, f_mu-pred_std*2, f_mu+pred_std*2, 
+                 alpha=0.3, color='tab:blue', label='$2\sqrt{\mathbb{V}\,[f]}$')
+ax2.legend()
+ax1.set_ylim([-4, 6])
+ax1.set_xlim([x.min(), x.max()])
+ax2.set_xlim([x.min(), x.max()])
+ax1.set_ylabel('$y$')
+ax1.set_xlabel('$x$')
+ax2.set_xlabel('$x$')
+plt.tight_layout()
+plt.show()
+```
+
+.. image:: regression_example.png
+   :align: center
\ No newline at end of file
diff --git a/regression_example.py b/regression_example.py
new file mode 100644
index 00000000..9dc93ec1
--- /dev/null
+++ b/regression_example.py
@@ -0,0 +1,70 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from laplace import Laplace
+
+n_epochs = 1000
+batch_size = 150  # full batch
+true_sigma_noise = 0.3
+torch.manual_seed(711)
+
+# create simple sinusoid data set
+X_train = (torch.rand(150) * 8).unsqueeze(-1)
+y_train = torch.sin(X_train) + torch.randn_like(X_train) * true_sigma_noise
+train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=batch_size)
+X_test = torch.linspace(-5, 13, 500).unsqueeze(-1)  # +-5 on top of the training X-range
+
+# create and train MAP model
+model = torch.nn.Sequential(torch.nn.Linear(1, 50),
+                            torch.nn.Tanh(),
+                            torch.nn.Linear(50, 1))
+
+criterion = torch.nn.MSELoss()
+optimizer = torch.optim.Adam(model.parameters(), weight_decay=5e-4, lr=1e-2)
+for i in range(n_epochs):
+    for X, y in train_loader:
+        optimizer.zero_grad()
+        loss = criterion(model(X), y)
+        loss.backward()
+        optimizer.step()
+
+la = Laplace(model, 'regression', subset_of_weights='all', hessian_structure='full')
+la.fit(train_loader)
+log_prior, log_sigma = torch.ones(1, requires_grad=True), torch.ones(1, requires_grad=True)
+hyper_optimizer = torch.optim.Adam([log_prior, log_sigma], lr=1e-1)
+for i in range(n_epochs):
+    hyper_optimizer.zero_grad()
+    neg_marglik = - la.log_marginal_likelihood(log_prior.exp(), log_sigma.exp())
+    neg_marglik.backward()
+    hyper_optimizer.step()
+print('sigma:', log_sigma.exp().item(), '; prior precision:', log_prior.exp().item())
+
+x = X_test.flatten().cpu().numpy()
+f_mu, f_var = la(X_test)
+f_mu = f_mu.squeeze().detach().cpu().numpy()
+f_sigma = f_var.squeeze().sqrt().cpu().numpy()
+pred_std = np.sqrt(f_sigma**2 + la.sigma_noise.item()**2)
+
+fig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2, sharey=True,
+                               figsize=(4.5, 2.8))
+ax1.set_title('MAP')
+ax1.scatter(X_train.flatten(), y_train.flatten(), alpha=0.7, color='tab:orange')
+ax1.plot(x, f_mu, color='black', label='$f_{MAP}$')
+ax1.legend()
+
+ax2.set_title('LA')
+ax2.scatter(X_train.flatten(), y_train.flatten(), alpha=0.7, color='tab:orange')
+ax2.plot(x, f_mu, label='$\mathbb{E}[f]$')
+ax2.fill_between(x, f_mu-pred_std*2, f_mu+pred_std*2, 
+                 alpha=0.3, color='tab:blue', label='$2\sqrt{\mathbb{V}\,[f]}$')
+ax2.legend()
+ax1.set_ylim([-4, 6])
+ax1.set_xlim([x.min(), x.max()])
+ax2.set_xlim([x.min(), x.max()])
+ax1.set_ylabel('$y$')
+ax1.set_xlabel('$x$')
+ax2.set_xlabel('$x$')
+plt.tight_layout()
+plt.savefig('docs/regression_example.png', dpi=300)
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 00000000..38b22452
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+torch
+torchvision
+torchaudio
+backpack-for-pytorch
+asdfghjkl @ git+https://github.com/kazukiosawa/asdfghjkl.git
diff --git a/setup.py b/setup.py
new file mode 100644
index 00000000..d4e4488a
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,43 @@
+from os import path
+from setuptools import find_packages, setup
+
+
+CURRENT_DIR = path.abspath(path.dirname(__file__))
+
+
+def read_me(filename):
+    with open(path.join(CURRENT_DIR, filename), encoding='utf-8') as f:
+        return f.read()
+
+
+def requirements(filename):
+    with open(path.join(CURRENT_DIR, filename)) as f:
+        return f.read().splitlines()
+
+
+AUTHORS = ""
+NAME = "laplace"
+PACKAGES = find_packages()
+DESCR = ""
+LONG_DESCR = ""
+LONG_DESCR_TYPE = 'text/markdown'
+REQUIREMENTS = requirements('requirements.txt')
+VERSION = "0.1"
+URL = ""
+LICENSE = ""
+
+
+setup(
+    author=AUTHORS,
+    name=NAME,
+    version=VERSION,
+    description=DESCR,
+    long_description=LONG_DESCR,
+    long_description_content_type=LONG_DESCR_TYPE,
+    install_requires=REQUIREMENTS,
+    url=URL,
+    license=LICENSE,
+    packages=PACKAGES,
+    zip_safe=False,
+    python_requires=">=3.7",
+)
diff --git a/template/config.mako b/template/config.mako
new file mode 100644
index 00000000..106f0bf3
--- /dev/null
+++ b/template/config.mako
@@ -0,0 +1,54 @@
+<%!
+    # Template configuration. Copy over in your template directory
+    # (used with `--template-dir`) and adapt as necessary.
+    # Note, defaults are loaded from this distribution file, so your
+    # config.mako only needs to contain values you want overridden.
+    # You can also run pdoc with `--config KEY=VALUE` to override
+    # individual values.
+    html_lang = 'en'
+    show_inherited_members = False
+    extract_module_toc_into_sidebar = True
+    list_class_variables_in_index = False
+    sort_identifiers = False
+    show_type_annotations = True
+    # Show collapsed source code block next to each item.
+    # Disabling this can improve rendering speed of large modules.
+    show_source_code = False
+    # If set, format links to objects in online source code repository
+    # according to this template. Supported keywords for interpolation
+    # are: commit, path, start_line, end_line.
+    #git_link_template = 'https://github.com/USER/PROJECT/blob/{commit}/{path}#L{start_line}-L{end_line}'
+    #git_link_template = 'https://gitlab.com/USER/PROJECT/blob/{commit}/{path}#L{start_line}-L{end_line}'
+    #git_link_template = 'https://bitbucket.org/USER/PROJECT/src/{commit}/{path}#lines-{start_line}:{end_line}'
+    #git_link_template = 'https://CGIT_HOSTNAME/PROJECT/tree/{path}?id={commit}#n{start-line}'
+    git_link_template = None
+    # A prefix to use for every HTML hyperlink in the generated documentation.
+    # No prefix results in all links being relative.
+    link_prefix = ''
+    # Enable syntax highlighting for code/source blocks by including Highlight.js
+    syntax_highlighting = True
+    # Set the style keyword such as 'atom-one-light' or 'github-gist'
+    #     Options: https://github.com/highlightjs/highlight.js/tree/master/src/styles
+    #     Demo: https://highlightjs.org/static/demo/
+    hljs_style = 'github'
+    # If set, insert Google Analytics tracking code. Value is GA
+    # tracking id (UA-XXXXXX-Y).
+    google_analytics = ''
+    # If set, insert Google Custom Search search bar widget above the sidebar index.
+    # The whitespace-separated tokens represent arbitrary extra queries (at least one
+    # must match) passed to regular Google search. Example:
+    #google_search_query = 'inurl:github.com/USER/PROJECT  site:PROJECT.github.io  site:PROJECT.website'
+    google_search_query = ''
+    # Enable offline search using Lunr.js. For explanation of 'fuzziness' parameter, which is
+    # added to every query word, see: https://lunrjs.com/guides/searching.html#fuzzy-matches
+    # If 'index_docstrings' is False, a shorter index is built, indexing only
+    # the full object reference names.
+    #lunr_search = {'fuzziness': 1, 'index_docstrings': True}
+    lunr_search = None
+    # If set, render LaTeX math syntax within \(...\) (inline equations),
+    # or within \[...\] or $$...$$ or `.. math::` (block equations)
+    # as nicely-formatted math formulas using MathJax.
+    # Note: in Python docstrings, either all backslashes need to be escaped (\\)
+    # or you need to use raw r-strings.
+    latex_math = True
+%>
\ No newline at end of file
diff --git a/update_docs.sh b/update_docs.sh
new file mode 100644
index 00000000..84972df6
--- /dev/null
+++ b/update_docs.sh
@@ -0,0 +1,5 @@
+rm -rf docs/
+pdoc --html -o docs --template-dir template laplace --force
+python regression_example.py
+mv docs/laplace/* docs/
+rm -rf docs/laplace/


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/requirements.txt b/tests/requirements.txt
new file mode 100644
index 00000000..d44283bd
--- /dev/null
+++ b/tests/requirements.txt
@@ -0,0 +1,4 @@
+pytest
+pytest-cov 
+coveralls
+scipy
diff --git a/tests/test_baselaplace.py b/tests/test_baselaplace.py
new file mode 100644
index 00000000..8bff47fd
--- /dev/null
+++ b/tests/test_baselaplace.py
@@ -0,0 +1,334 @@
+import pytest
+from itertools import product
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+from torch.utils.data import DataLoader, TensorDataset
+from torch.distributions import Normal, Categorical
+
+from laplace.laplace import Laplace, FullLaplace, KronLaplace, DiagLaplace
+from tests.utils import jacobians_naive
+
+
+torch.manual_seed(240)
+torch.set_default_tensor_type(torch.DoubleTensor)
+flavors = [FullLaplace, KronLaplace, DiagLaplace]
+
+
+@pytest.fixture
+def model():
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    return model
+
+
+@pytest.fixture
+def class_loader():
+    X = torch.randn(10, 3)
+    y = torch.randint(2, (10,))
+    return DataLoader(TensorDataset(X, y), batch_size=3)
+
+
+@pytest.fixture
+def reg_loader():
+    X = torch.randn(10, 3)
+    y = torch.randn(10, 2)
+    return DataLoader(TensorDataset(X, y), batch_size=3)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init(laplace, model):
+    lap = laplace(model, 'classification')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_invalid_likelihood(laplace, model):
+    with pytest.raises(ValueError):
+        lap = laplace(model, 'otherlh')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_noise(laplace, model):
+    # float
+    sigma_noise = 1.2
+    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+    # torch.tensor 0-dim
+    sigma_noise = torch.tensor(1.2)
+    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+    # torch.tensor 1-dim
+    sigma_noise = torch.tensor(1.2).reshape(-1)
+    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+
+    # for classification should fail
+    sigma_noise = 1.2
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='classification', sigma_noise=sigma_noise)
+
+    # other than that should fail
+    # higher dim
+    sigma_noise = torch.tensor(1.2).reshape(1, 1)
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+    # other datatype, only reals supported
+    sigma_noise = '1.2'
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_precision(laplace, model):
+    # float
+    precision = 10.6
+    lap = laplace(model, likelihood='regression', prior_precision=precision)
+    # torch.tensor 0-dim
+    precision = torch.tensor(10.6)
+    lap = laplace(model, likelihood='regression', prior_precision=precision)
+    # torch.tensor 1-dim
+    precision = torch.tensor(10.7).reshape(-1)
+    lap = laplace(model, likelihood='regression', prior_precision=precision)
+    # torch.tensor 1-dim param-shape
+    precision = torch.tensor(10.7).reshape(-1).repeat(model.n_params)
+    if laplace == KronLaplace:
+        # Kron should not accept per parameter prior precision
+        with pytest.raises(ValueError):
+            lap = laplace(model, likelihood='regression', prior_precision=precision)
+    else:
+        lap = laplace(model, likelihood='regression', prior_precision=precision)
+    # torch.tensor 1-dim layer-shape
+    precision = torch.tensor(10.7).reshape(-1).repeat(model.n_layers)
+    lap = laplace(model, likelihood='regression', prior_precision=precision)
+
+    # other than that should fail
+    # higher dim
+    precision = torch.tensor(10.6).reshape(1, 1)
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', prior_precision=precision)
+    # unmatched dim
+    precision = torch.tensor(10.6).reshape(-1).repeat(17)
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', prior_precision=precision)
+    # other datatype, only reals supported
+    precision = '1.5'
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', prior_precision=precision)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_prior_mean_and_scatter(laplace, model):
+    mean = parameters_to_vector(model.parameters())
+    P = len(mean)
+    lap_scalar_mean = laplace(model, 'classification',
+                              prior_precision=1e-2, prior_mean=1.)
+    assert torch.allclose(lap_scalar_mean.prior_mean, torch.tensor([1.]))
+    lap_tensor_mean = laplace(model, 'classification',
+                              prior_precision=1e-2, prior_mean=torch.ones(1))
+    assert torch.allclose(lap_tensor_mean.prior_mean, torch.tensor([1.]))
+    lap_tensor_scalar_mean = laplace(model, 'classification',
+                                     prior_precision=1e-2, prior_mean=torch.ones(1)[0])
+    assert torch.allclose(lap_tensor_scalar_mean.prior_mean, torch.tensor(1.))
+    lap_tensor_full_mean = laplace(model, 'classification',
+                                   prior_precision=1e-2, prior_mean=torch.ones(P))
+    assert torch.allclose(lap_tensor_full_mean.prior_mean, torch.ones(P))
+    expected = ((mean - 1) * 1e-2) @ (mean - 1)
+    assert expected.ndim == 0
+    assert torch.allclose(lap_scalar_mean.scatter, expected)
+    assert lap_scalar_mean.scatter.shape == expected.shape
+    assert torch.allclose(lap_tensor_mean.scatter, expected)
+    assert lap_tensor_mean.scatter.shape == expected.shape
+    assert torch.allclose(lap_tensor_scalar_mean.scatter, expected)
+    assert lap_tensor_scalar_mean.scatter.shape == expected.shape
+    assert torch.allclose(lap_tensor_full_mean.scatter, expected)
+    assert lap_tensor_full_mean.scatter.shape == expected.shape
+
+    # too many dims
+    with pytest.raises(ValueError):
+        prior_mean = torch.ones(P).unsqueeze(-1)
+        laplace(model, 'classification', prior_precision=1e-2, prior_mean=prior_mean)
+
+    # unmatched dim
+    with pytest.raises(ValueError):
+        prior_mean = torch.ones(P-3)
+        laplace(model, 'classification', prior_precision=1e-2, prior_mean=prior_mean)
+
+    # invalid argument type
+    with pytest.raises(ValueError):
+        laplace(model, 'classification', prior_precision=1e-2, prior_mean='72')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_temperature(laplace, model):
+    # valid float
+    T = 1.1
+    lap = laplace(model, likelihood='classification', temperature=T)
+    assert lap.temperature == T
+
+
+@pytest.mark.parametrize('laplace,lh', product(flavors, ['classification', 'regression']))
+def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
+    if lh == 'classification':
+        loader = class_loader
+        sigma_noise = 1.
+    else:
+        loader = reg_loader
+        sigma_noise = 0.3
+    lap = laplace(model, lh, sigma_noise=sigma_noise, prior_precision=0.7)
+    lap.fit(loader)
+    assert lap.n_data == len(loader.dataset)
+    assert lap.n_outputs == model.output_size
+    f = model(loader.dataset.tensors[0])
+    y = loader.dataset.tensors[1]
+    assert f.shape == torch.Size([10, 2])
+
+    # Test log likelihood (Train)
+    log_lik = lap.log_likelihood
+    # compute true log lik
+    if lh == 'classification':
+        log_lik_true = Categorical(logits=f).log_prob(y).sum()
+        assert torch.allclose(log_lik, log_lik_true)
+    else:
+        assert y.size() == f.size()
+        log_lik_true = Normal(loc=f, scale=sigma_noise).log_prob(y).sum()
+        assert torch.allclose(log_lik, log_lik_true)
+        # change likelihood and test again
+        lap.sigma_noise = 0.72
+        log_lik = lap.log_likelihood
+        log_lik_true = Normal(loc=f, scale=0.72).log_prob(y).sum()
+        assert torch.allclose(log_lik, log_lik_true)
+
+    # Test marginal likelihood
+    # lml = log p(y|f) - 1/2 theta @ prior_prec @ theta
+    #       + 1/2 logdet prior_prec - 1/2 log det post_prec
+    lml = log_lik_true
+    theta = parameters_to_vector(model.parameters()).detach()
+    assert torch.allclose(theta, lap.mean)
+    prior_prec = torch.diag(lap.prior_precision_diag)
+    assert prior_prec.shape == torch.Size([len(theta), len(theta)])
+    lml = lml - 1/2 * theta @ prior_prec @ theta
+    Sigma_0 = torch.inverse(prior_prec)
+    if laplace == DiagLaplace:
+        log_det_post_prec = lap.posterior_precision.log().sum()
+    else:
+        log_det_post_prec = lap.posterior_precision.logdet()
+    lml = lml + 1/2 * (prior_prec.logdet() - log_det_post_prec)
+    assert torch.allclose(lml, lap.log_marginal_likelihood())
+
+    # test sampling
+    torch.manual_seed(61)
+    samples = lap.sample(n_samples=1)
+    assert samples.shape == torch.Size([1, len(theta)])
+    samples = lap.sample(n_samples=1000000)
+    assert samples.shape == torch.Size([1000000, len(theta)])
+    mu_comp = samples.mean(dim=0)
+    mu_true = lap.mean
+    assert torch.allclose(mu_comp, mu_true, rtol=1)
+
+    # test functional variance
+    if laplace == FullLaplace:
+        Sigma = lap.posterior_covariance
+    elif laplace == KronLaplace:
+        Sigma = lap.posterior_precision.to_matrix(exponent=-1)
+    elif laplace == DiagLaplace:
+        Sigma = torch.diag(lap.posterior_variance)
+    Js, f = jacobians_naive(model, loader.dataset.tensors[0])
+    true_f_var = torch.einsum('mkp,pq,mcq->mkc', Js, Sigma, Js)
+    comp_f_var = lap.functional_variance(Js)
+    assert torch.allclose(true_f_var, comp_f_var, rtol=1e-4)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_regression_predictive(laplace, model, reg_loader):
+    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap.fit(reg_loader)
+    X, y = reg_loader.dataset.tensors
+    f = model(X)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive, functional variance tested already above.
+    f_mu, f_var = lap(X, pred_type='glm')
+    assert torch.allclose(f_mu, f)
+    assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
+    assert len(f_mu) == len(X)
+
+    # NN predictive (only diagonal variance estimation)
+    f_mu, f_var = lap(X, pred_type='nn')
+    assert f_mu.shape == f_var.shape
+    assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1]])
+    assert len(f_mu) == len(X)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_classification_predictive(laplace, model, class_loader):
+    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap.fit(class_loader)
+    X, y = class_loader.dataset.tensors
+    f = torch.softmax(model(X), dim=-1)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive
+    f_pred = lap(X, pred_type='glm', link_approx='mc', n_samples=100)
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    f_pred = lap(X, pred_type='glm', link_approx='probit')
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    f_pred = lap(X, pred_type='glm', link_approx='bridge')
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+
+
+    # NN predictive
+    f_pred = lap(X, pred_type='nn', n_samples=100)
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_regression_predictive_samples(laplace, model, reg_loader):
+    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap.fit(reg_loader)
+    X, y = reg_loader.dataset.tensors
+    f = model(X)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive, functional variance tested already above.
+    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+
+    # NN predictive (only diagonal variance estimation)
+    fsamples = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_classification_predictive_samples(laplace, model, class_loader):
+    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap.fit(class_loader)
+    X, y = class_loader.dataset.tensors
+    f = torch.softmax(model(X), dim=-1)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive
+    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+    assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
+
+    # NN predictive
+    f_pred = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+    assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
diff --git a/tests/test_curv_backends_asdl.py b/tests/test_curv_backends_asdl.py
new file mode 100644
index 00000000..696c8969
--- /dev/null
+++ b/tests/test_curv_backends_asdl.py
@@ -0,0 +1,216 @@
+import pytest
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+
+from asdfghjkl.operations import Bias, Scale
+
+from laplace.curvature import AsdlGGN, AsdlEF, BackPackGGN, BackPackEF
+
+
+@pytest.fixture
+def model():
+    torch.manual_seed(711)
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    return model
+
+
+@pytest.fixture
+def class_Xy():
+    torch.manual_seed(711)
+    X = torch.randn(10, 3)
+    y = torch.randint(2, (10,))
+    return X, y
+
+
+@pytest.fixture
+def complex_model():
+    torch.manual_seed(711)
+    model = torch.nn.Sequential(nn.Conv2d(3, 4, 2, 2), nn.Flatten(), nn.Tanh(),
+                                nn.Linear(16, 20), nn.Tanh(), Scale(), Bias(), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    return model
+
+
+@pytest.fixture
+def complex_class_Xy():
+    torch.manual_seed(711)
+    X = torch.randn(10, 3, 5, 5)
+    y = torch.randint(2, (10,))
+    return X, y
+
+
+def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
+    X, y = class_Xy
+    backend = AsdlGGN(model, 'classification', stochastic=False)
+    loss, dggn = backend.diag(X[:5], y[:5])
+    loss2, dggn2 = backend.diag(X[5:], y[5:])
+    loss += loss2
+    dggn += dggn2
+
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss_f, H_ggn = backend.full(X, y)
+    assert torch.allclose(loss, loss_f)
+    assert torch.allclose(dggn, H_ggn.diagonal())
+
+
+def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
+    X, y = class_Xy
+    backend = AsdlEF(model, 'classification')
+    loss, dggn = backend.diag(X[:5], y[:5])
+    loss2, dggn2 = backend.diag(X[5:], y[5:])
+    loss += loss2
+    dggn += dggn2
+
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackEF(model, 'classification')
+    loss_f, H_ggn = backend.full(X, y)
+    assert torch.allclose(loss, loss_f)
+    assert torch.allclose(dggn, H_ggn.diagonal())
+
+
+def test_diag_ggn_stoch_cls_kazuki(class_Xy, model):
+    X, y = class_Xy
+    backend = AsdlGGN(model, 'classification', stochastic=True)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # same order of magnitude os non-stochastic.
+    backend = AsdlGGN(model, 'classification', stochastic=False)
+    loss_ns, dggn_ns = backend.diag(X, y)
+    assert loss_ns == loss
+    assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
+
+
+@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+def test_kron_kazuki_vs_diag_class(class_Xy, model, Backend):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = class_Xy
+    backend = Backend(model, 'classification')
+    loss, dggn = backend.diag(X[:1], y[:1], N=1)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+    loss, kron = backend.kron(X[:1], y[:1], N=1)
+    assert torch.allclose(kron.diag(), dggn)
+
+
+@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
+    X, y = class_Xy
+    backend = Backend(model, 'classification')
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert len(kron.diag()) == model.n_params
+
+    N = len(X)
+    M = 3
+    loss1, kron1 = backend.kron(X[:M], y[:M], N=N)
+    loss2, kron2 = backend.kron(X[M:], y[M:], N=N)
+    kron_two = kron1 + kron2
+    loss_two = loss1 + loss2
+    assert torch.allclose(kron.diag(), kron_two.diag())
+    assert torch.allclose(loss, loss_two)
+
+
+@pytest.mark.parametrize('Backend', [AsdlGGN, AsdlEF])
+def test_kron_summing_up_vs_diag_kazuki(class_Xy, model, Backend):
+    # For a single data point, Kron is exact and should equal diag class_Xy
+    X, y = class_Xy
+    backend = Backend(model, 'classification')
+    loss, dggn = backend.diag(X, y, N=len(X))
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
+
+
+def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
+    X, y = complex_class_Xy
+    backend = AsdlGGN(complex_model, 'classification', stochastic=True)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == complex_model.n_params
+
+    # same order of magnitude os non-stochastic.
+    loss_ns, dggn_ns = backend.diag(X, y)
+    assert loss_ns == loss
+    assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1)
+
+
+@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+def test_complex_kron_kazuki_vs_diag_kazuki(complex_class_Xy, complex_model, Backend):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = complex_class_Xy
+    backend = Backend(complex_model, 'classification')
+    loss, dggn = backend.diag(X[:1], y[:1], N=1)
+    # sanity check size of diag ggn
+    assert len(dggn) == complex_model.n_params
+    loss, kron = backend.kron(X[:1], y[:1], N=1)
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
+
+
+@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+def test_complex_kron_batching_correction_kazuki(complex_class_Xy, complex_model, Backend):
+    X, y = complex_class_Xy
+    backend = Backend(complex_model, 'classification')
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert len(kron.diag()) == complex_model.n_params
+
+    N = len(X)
+    M = 3
+    loss1, kron1 = backend.kron(X[:M], y[:M], N=N)
+    loss2, kron2 = backend.kron(X[M:], y[M:], N=N)
+    kron_two = kron1 + kron2
+    loss_two = loss1 + loss2
+    assert torch.allclose(kron.diag(), kron_two.diag())
+    assert torch.allclose(loss, loss_two)
+
+
+@pytest.mark.parametrize('Backend', [AsdlGGN, AsdlEF])
+def test_complex_kron_summing_up_vs_diag_class_kazuki(complex_class_Xy, complex_model, Backend):
+    # For a single data point, Kron is exact and should equal diag class_Xy
+    X, y = complex_class_Xy
+    backend = Backend(complex_model, 'classification')
+    loss, dggn = backend.diag(X, y, N=len(X))
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-2)
+
+
+def test_kron_normalization_ggn_class(class_Xy, model):
+    X, y = class_Xy
+    xi, yi = X[:1], y[:1]
+    backend = AsdlGGN(model, 'classification', stochastic=False)
+    loss, kron = backend.kron(xi, yi, N=1)
+    kron_true = 7 * kron
+    loss_true = 7 * loss
+    X = torch.repeat_interleave(xi, 7, 0)
+    y = torch.repeat_interleave(yi, 7, 0)
+    loss_test, kron_test  = backend.kron(X, y, N=7)
+    assert torch.allclose(kron_true.diag(), kron_test.diag())
+    assert torch.allclose(loss_true, loss_test)
+
+
+def test_kron_normalization_ef_class(class_Xy, model):
+    X, y = class_Xy
+    xi, yi = X[:1], y[:1]
+    backend = AsdlEF(model, 'classification')
+    loss, kron = backend.kron(xi, yi, N=1)
+    kron_true = 7 * kron
+    loss_true = 7 * loss
+    X = torch.repeat_interleave(xi, 7, 0)
+    y = torch.repeat_interleave(yi, 7, 0)
+    loss_test, kron_test  = backend.kron(X, y, N=7)
+    assert torch.allclose(kron_true.diag(), kron_test.diag())
+    assert torch.allclose(loss_true, loss_test)
diff --git a/tests/test_curv_backends_backpack.py b/tests/test_curv_backends_backpack.py
new file mode 100644
index 00000000..772dc542
--- /dev/null
+++ b/tests/test_curv_backends_backpack.py
@@ -0,0 +1,242 @@
+import pytest
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+
+from laplace.curvature import BackPackGGN, BackPackEF
+
+
+@pytest.fixture
+def model():
+    torch.manual_seed(711)
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    return model
+
+
+@pytest.fixture
+def class_Xy():
+    torch.manual_seed(711)
+    X = torch.randn(10, 3)
+    y = torch.randint(2, (10,))
+    return X, y
+
+
+@pytest.fixture
+def reg_Xy():
+    torch.manual_seed(711)
+    X = torch.randn(10, 3)
+    y = torch.randn(10, 2)
+    return X, y
+
+
+def test_full_ggn_backpack_reg_integration(reg_Xy, model):
+    X, y = reg_Xy
+    backend = BackPackGGN(model, 'regression', stochastic=True)
+    with pytest.raises(ValueError):
+        loss, fggn = backend.full(X, y)
+
+    # cannot test, its implemented based on Jacobians.
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss, H_ggn = backend.full(X, y)
+    assert H_ggn.size() == torch.Size((model.n_params, model.n_params))
+
+
+def test_full_ggn_backpack_cls_integration(class_Xy, model):
+    X, y = class_Xy
+    backend = BackPackGGN(model, 'classification', stochastic=True)
+    with pytest.raises(ValueError):
+        loss, fggn = backend.full(X, y)
+
+    # cannot test, its implemented based on Jacobians.
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss, H_ggn = backend.full(X, y)
+    assert H_ggn.size() == torch.Size((model.n_params, model.n_params))
+
+
+def test_diag_ggn_cls_backpack(class_Xy, model):
+    X, y = class_Xy
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss_f, H_ggn = backend.full(X, y)
+    assert loss == loss_f
+    assert torch.allclose(dggn, H_ggn.diagonal())
+
+
+def test_diag_ggn_reg_backpack(reg_Xy, model):
+    X, y = reg_Xy
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss_f, H_ggn = backend.full(X, y)
+    assert loss == loss_f
+    assert torch.allclose(dggn, H_ggn.diagonal())
+
+
+def test_diag_ggn_stoch_cls_backpack(class_Xy, model):
+    X, y = class_Xy
+    backend = BackPackGGN(model, 'classification', stochastic=True)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # same order of magnitude os non-stochastic.
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss_ns, dggn_ns = backend.diag(X, y)
+    assert loss_ns == loss
+    assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
+
+
+def test_diag_ggn_stoch_reg_backpack(reg_Xy, model):
+    X, y = reg_Xy
+    backend = BackPackGGN(model, 'regression', stochastic=True)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # same order of magnitude os non-stochastic.
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss_ns, dggn_ns = backend.diag(X, y)
+    assert loss_ns == loss
+    assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
+
+
+def test_kron_ggn_reg_backpack_vs_diag_reg(reg_Xy, model):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = reg_Xy
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss, dggn = backend.diag(X[:1], y[:1], N=1)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+    loss, kron = backend.kron(X[:1], y[:1], N=1)
+    assert torch.allclose(kron.diag(), dggn)
+
+
+def test_kron_batching_correction_reg(reg_Xy, model):
+    X, y = reg_Xy
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert len(kron.diag()) == model.n_params
+
+    N = len(X)
+    M = int(N / 2)
+    loss1, kron1 = backend.kron(X[:M], y[:M], N=N)
+    loss2, kron2 = backend.kron(X[M:], y[M:], N=N)
+    kron_two = kron1 + kron2
+    loss_two = loss1 + loss2
+    assert torch.allclose(kron.diag(), kron_two.diag())
+    assert torch.allclose(loss, loss_two)
+
+
+def test_kron_summing_up_vs_diag_reg(reg_Xy, model):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = reg_Xy
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss, dggn = backend.diag(X, y, N=len(X))
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
+
+
+def test_kron_ggn_reg_backpack_vs_diag_class(class_Xy, model):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = class_Xy
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss, dggn = backend.diag(X[:1], y[:1], N=1)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+    loss, kron = backend.kron(X[:1], y[:1], N=1)
+    assert torch.allclose(kron.diag(), dggn)
+
+
+def test_kron_batching_correction_class(class_Xy, model):
+    X, y = class_Xy
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert len(kron.diag()) == model.n_params
+
+    N = len(X)
+    M = int(N / 2)
+    loss1, kron1 = backend.kron(X[:M], y[:M], N=N)
+    loss2, kron2 = backend.kron(X[M:], y[M:], N=N)
+    kron_two = kron1 + kron2
+    loss_two = loss1 + loss2
+    assert torch.allclose(kron.diag(), kron_two.diag())
+    assert torch.allclose(loss, loss_two)
+
+
+def test_kron_summing_up_vs_diag_class(class_Xy, model):
+    # For a single data point, Kron is exact and should equal diag class_Xy
+    X, y = class_Xy
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss, dggn = backend.diag(X, y, N=len(X))
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
+
+
+def test_full_vs_diag_ef_cls_backpack(class_Xy, model):
+    X, y = class_Xy
+    backend = BackPackEF(model, 'classification')
+    loss, diag_ef = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(diag_ef) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackEF(model, 'classification')
+    loss_f, H_ef = backend.full(X, y)
+    assert loss == loss_f
+    assert torch.allclose(diag_ef, H_ef.diagonal())
+
+
+def test_full_vs_diag_ef_reg_backpack(reg_Xy, model):
+    X, y = reg_Xy
+    backend = BackPackEF(model, 'regression')
+    loss, diag_ef = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(diag_ef) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackEF(model, 'regression')
+    loss_f, H_ef = backend.full(X, y)
+    assert loss == loss_f
+    assert torch.allclose(diag_ef, H_ef.diagonal())
+
+
+def test_kron_normalization_reg(reg_Xy, model):
+    X, y = reg_Xy
+    xi, yi = X[:1], y[:1]
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss, kron = backend.kron(xi, yi, N=1)
+    kron_true = 7 * kron
+    loss_true = 7 * loss
+    X = torch.repeat_interleave(xi, 7, 0)
+    y = torch.repeat_interleave(yi, 7, 0)
+    loss_test, kron_test  = backend.kron(X, y, N=7)
+    assert torch.allclose(kron_true.diag(), kron_test.diag())
+    assert torch.allclose(loss_true, loss_test)
+
+
+def test_kron_normalization_class(class_Xy, model):
+    X, y = class_Xy
+    xi, yi = X[:1], y[:1]
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss, kron = backend.kron(xi, yi, N=1)
+    kron_true = 7 * kron
+    loss_true = 7 * loss
+    X = torch.repeat_interleave(xi, 7, 0)
+    y = torch.repeat_interleave(yi, 7, 0)
+    loss_test, kron_test  = backend.kron(X, y, N=7)
+    assert torch.allclose(kron_true.diag(), kron_test.diag())
+    assert torch.allclose(loss_true, loss_test)
+
diff --git a/tests/test_feature_extractor.py b/tests/test_feature_extractor.py
new file mode 100644
index 00000000..37494d76
--- /dev/null
+++ b/tests/test_feature_extractor.py
@@ -0,0 +1,137 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+from laplace.feature_extractor import FeatureExtractor
+
+
+class CNN(nn.Module):
+    def __init__(self, num_classes):
+        nn.Module.__init__(self)
+        self.conv1 = nn.Sequential(
+            # Input shape (3, 64, 64)
+            nn.Conv2d(
+                in_channels=3,
+                out_channels=6,
+                kernel_size=5,
+                stride=1,
+                padding=2
+            ),
+            # Output shape (6, 60, 60)
+            nn.ReLU(),
+            # Output shape (6, 30, 30)
+            nn.MaxPool2d(kernel_size=2)
+        )
+
+        self.fc = nn.Sequential(
+            nn.Linear(in_features=16 * 16 * 16,
+                      out_features=300),
+            nn.ReLU(),
+            nn.Linear(in_features=300,
+                      out_features=84),
+            nn.ReLU(),
+            nn.Linear(in_features=84,
+                      out_features=num_classes)
+        )
+
+        self.conv2 = nn.Sequential(
+            # Input shape (6, 30, 30)
+            nn.Conv2d(
+                in_channels=6,
+                out_channels=16,
+                kernel_size=5,
+                stride=1,
+                padding=2
+            ),
+            # Output shape (16, 26, 26)
+            nn.ReLU(),
+            # Output shape (16, 13, 13)
+            nn.MaxPool2d(kernel_size=2)
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        # print("x = self.conv1(x)   ", x.shape, "  ", torch.Size)
+        x = self.conv2(x)
+        x = x.view(x.size()[0], -1)
+        x = self.fc(x)
+        return x
+
+
+def get_model(model_name):
+    if model_name == 'resnet18':
+        model = models.resnet18()
+    elif model_name == 'alexnet':
+        model = models.alexnet()
+    elif model_name == 'vgg16':
+        model = models.vgg16()
+    elif model_name == 'squeezenet':
+        model = models.squeezenet1_0()
+    elif model_name == 'densenet':
+        model = models.densenet161()
+    elif model_name == 'inception':
+        model = models.inception_v3(init_weights=True)
+    elif model_name == 'googlenet':
+        model = models.googlenet(init_weights=True)
+    elif model_name == 'shufflenet':
+        model = models.shufflenet_v2_x1_0()
+    elif model_name == 'mobilenet_v2':
+        model = models.mobilenet_v2()
+    elif model_name == 'mobilenet_v3_large':
+        model = models.mobilenet_v3_large()
+    elif model_name == 'mobilenet_v3_small':
+        model = models.mobilenet_v3_small()
+    elif model_name == 'resnext50_32x4d':
+        model = models.resnext50_32x4d()
+    elif model_name == 'wide_resnet50_2':
+        model = models.wide_resnet50_2()
+    elif model_name == 'mnasnet':
+        model = models.mnasnet1_0()
+    elif model_name == 'switchedCNN':
+        model = CNN(10)
+    elif model_name == 'sequential':
+        model = nn.Sequential(
+            nn.Conv2d(3, 6, 3, 1, 1),
+            nn.ReLU(),
+            nn.Flatten(),
+            nn.Linear(6*64*64, 10),
+            nn.ReLU(),
+            nn.Linear(10, 10)
+        )
+    else:
+        raise ValueError(f'{model_name} is not supported.')
+    return model
+
+
+@torch.no_grad()
+def test_feature_extractor():
+    # all torchvision classifcation models but 'squeezenet' (no linear last layer)
+    # + model where modules are initilaized in wrong order + nn.Sequential model
+    model_names = ['resnet18', 'alexnet', 'vgg16', 'densenet', 'inception',
+                   'googlenet', 'shufflenet', 'mobilenet_v2', 'mobilenet_v3_large',
+                   'mobilenet_v3_small', 'resnext50_32x4d', 'wide_resnet50_2',
+                   'mnasnet', 'switchedCNN', 'sequential']
+
+    # to test the last_layer_name argument
+    # last_layer_names = ['fc', 'classifier.6', 'classifier.6', 'classifier', 'fc',
+    #                     'fc', 'fc', 'classifier.1', 'classifier.3', 'classifier.3',
+    #                     'fc', 'fc', 'classifier.1', 'fc.4', '5']
+
+    for model_name in model_names:
+        # generate random test input
+        if model_name == 'inception':
+            x = torch.randn(1, 3, 299, 299)
+        else:
+            x = torch.randn(1, 3, 64, 64)
+
+        model = get_model(model_name)
+        model.eval()
+
+        # extract features and get output
+        feature_extractor = FeatureExtractor(model)
+        out, features = feature_extractor.forward_with_features(x)
+
+        # test if it worked
+        last_layer = feature_extractor.last_layer
+        out2 = last_layer(features)
+        assert (out == out2).all().item()
diff --git a/tests/test_jacobians.py b/tests/test_jacobians.py
new file mode 100644
index 00000000..7a5a22ef
--- /dev/null
+++ b/tests/test_jacobians.py
@@ -0,0 +1,92 @@
+import pytest
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+
+from laplace.curvature import AsdlInterface, BackPackInterface
+from laplace.feature_extractor import FeatureExtractor
+from tests.utils import jacobians_naive
+
+
+@pytest.fixture
+def multioutput_model():
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    return model
+
+
+@pytest.fixture
+def singleoutput_model():
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 1))
+    setattr(model, 'output_size', 1)
+    return model
+
+
+@pytest.fixture
+def linear_model():
+    model = torch.nn.Sequential(nn.Linear(3, 1, bias=False))
+    setattr(model, 'output_size', 1)
+    return model
+
+
+@pytest.fixture
+def X():
+    torch.manual_seed(15)
+    return torch.randn(200, 3)
+
+
+@pytest.mark.parametrize('backend', [AsdlInterface, BackPackInterface])
+def test_linear_jacobians(linear_model, X, backend):
+    # jacobian of linear model is input X.
+    Js, f = backend.jacobians(linear_model, X)
+    # into Jacs shape (batch_size, output_size, params)
+    true_Js = X.reshape(len(X), 1, -1)
+    assert true_Js.shape == Js.shape
+    assert torch.allclose(true_Js, Js, atol=1e-5)
+    assert torch.allclose(f, linear_model(X), atol=1e-5)
+
+
+@pytest.mark.parametrize('backend', [AsdlInterface, BackPackInterface])
+def test_jacobians_singleoutput(singleoutput_model, X, backend):
+    model = singleoutput_model
+    Js, f = backend.jacobians(model, X)
+    Js_naive, f_naive = jacobians_naive(model, X)
+    assert Js.shape == Js_naive.shape
+    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.allclose(model(X), f_naive)
+    assert torch.allclose(f, f_naive)
+
+
+@pytest.mark.parametrize('backend', [AsdlInterface, BackPackInterface])
+def test_jacobians_multioutput(multioutput_model, X, backend):
+    model = multioutput_model
+    Js, f = backend.jacobians(model, X)
+    Js_naive, f_naive = jacobians_naive(model, X)
+    assert Js.shape == Js_naive.shape
+    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.allclose(model(X), f_naive)
+    assert torch.allclose(f, f_naive)
+
+
+@pytest.mark.parametrize('backend', [AsdlInterface, BackPackInterface])
+def test_last_layer_jacobians_singleoutput(singleoutput_model, X, backend):
+    model = FeatureExtractor(singleoutput_model)
+    Js, f = backend.last_layer_jacobians(model, X)
+    _, phi = model.forward_with_features(X)
+    Js_naive, f_naive = jacobians_naive(model.last_layer, phi)
+    assert Js.shape == Js_naive.shape
+    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.allclose(model(X), f_naive)
+    assert torch.allclose(f, f_naive)
+
+
+@pytest.mark.parametrize('backend', [AsdlInterface, BackPackInterface])
+def test_last_layer_jacobians_multioutput(multioutput_model, X, backend):
+    model = FeatureExtractor(multioutput_model)
+    Js, f = backend.last_layer_jacobians(model, X)
+    _, phi = model.forward_with_features(X)
+    Js_naive, f_naive = jacobians_naive(model.last_layer, phi)
+    assert Js.shape == Js_naive.shape
+    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.allclose(model(X), f_naive)
+    assert torch.allclose(f, f_naive)
diff --git a/tests/test_laplace.py b/tests/test_laplace.py
new file mode 100644
index 00000000..c40be165
--- /dev/null
+++ b/tests/test_laplace.py
@@ -0,0 +1,49 @@
+import pytest
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+
+from laplace.laplace import Laplace
+from laplace.baselaplace import FullLaplace, KronLaplace, DiagLaplace
+from laplace.lllaplace import FullLLLaplace, KronLLLaplace, DiagLLLaplace
+
+
+torch.manual_seed(240)
+torch.set_default_tensor_type(torch.DoubleTensor)
+flavors = [FullLaplace, KronLaplace, DiagLaplace,
+           FullLLLaplace, KronLLLaplace, DiagLLLaplace]
+all_keys = [('all', 'full'), ('all', 'kron'), ('all', 'diag'),
+            ('last_layer', 'full'), ('last_layer', 'kron'), ('last_layer', 'diag')]
+
+
+@pytest.fixture
+def model():
+    model = nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
+    model_params = list(model.parameters())
+    return model
+
+
+def test_default_init(model, likelihood='classification'):
+    # test if default initialization works, id=(last-layer, kron)
+    lap = Laplace(model, likelihood)
+    assert isinstance(lap, KronLLLaplace)
+
+
+@pytest.mark.parametrize('laplace, key', zip(flavors, all_keys))
+def test_all_init(laplace, key, model, likelihood='classification'):
+    # test if all flavors are correctly initialized
+    w, s = key
+    lap = Laplace(model, likelihood, subset_of_weights=w, hessian_structure=s)
+    assert isinstance(lap, laplace)
+
+
+@pytest.mark.parametrize('key', all_keys)
+def test_opt_keywords(key, model, likelihood='classification'):
+    # test if optional keywords are correctly passed on
+    w, s = key
+    prior_mean = torch.zeros_like(parameters_to_vector(model.parameters()))
+    lap = Laplace(model, likelihood, subset_of_weights=w, hessian_structure=s,
+                  prior_precision=0.01, prior_mean=prior_mean, temperature=10.)
+    assert torch.allclose(lap.prior_mean, prior_mean)
+    assert lap.prior_precision == 0.01
+    assert lap.temperature == 10.
diff --git a/tests/test_lllaplace.py b/tests/test_lllaplace.py
new file mode 100644
index 00000000..102f429b
--- /dev/null
+++ b/tests/test_lllaplace.py
@@ -0,0 +1,357 @@
+import pytest
+from itertools import product
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+from torch.utils.data import DataLoader, TensorDataset
+from torch.distributions import Normal, Categorical
+
+from laplace.lllaplace import LLLaplace, FullLLLaplace, KronLLLaplace, DiagLLLaplace
+from laplace.feature_extractor import FeatureExtractor
+from tests.utils import jacobians_naive
+
+
+torch.manual_seed(240)
+torch.set_default_tensor_type(torch.DoubleTensor)
+flavors = [FullLLLaplace, KronLLLaplace, DiagLLLaplace]
+
+
+@pytest.fixture
+def model():
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    return model
+
+
+@pytest.fixture
+def class_loader():
+    X = torch.randn(10, 3)
+    y = torch.randint(2, (10,))
+    return DataLoader(TensorDataset(X, y), batch_size=3)
+
+
+@pytest.fixture
+def reg_loader():
+    X = torch.randn(10, 3)
+    y = torch.randn(10, 2)
+    return DataLoader(TensorDataset(X, y), batch_size=3)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init(laplace, model):
+    lap = laplace(model, 'classification', last_layer_name='1')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_invalid_likelihood(laplace, model):
+    with pytest.raises(ValueError):
+        lap = laplace(model, 'otherlh', last_layer_name='1')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_noise(laplace, model):
+    # float
+    sigma_noise = 1.2
+    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise,
+                  last_layer_name='1')
+    # torch.tensor 0-dim
+    sigma_noise = torch.tensor(1.2)
+    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise,
+                  last_layer_name='1')
+    # torch.tensor 1-dim
+    sigma_noise = torch.tensor(1.2).reshape(-1)
+    lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise,
+                  last_layer_name='1')
+
+    # for classification should fail
+    sigma_noise = 1.2
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='classification',
+                      sigma_noise=sigma_noise, last_layer_name='1')
+
+    # other than that should fail
+    # higher dim
+    sigma_noise = torch.tensor(1.2).reshape(1, 1)
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise,
+                      last_layer_name='1')
+    # other datatype, only reals supported
+    sigma_noise = '1.2'
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', sigma_noise=sigma_noise,
+                      last_layer_name='1')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_precision(laplace, model):
+    feature_extractor = FeatureExtractor(model, last_layer_name='1')
+    model_params = list(feature_extractor.last_layer.parameters())
+    setattr(model, 'n_layers', 2)  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    # float
+    precision = 10.6
+    lap = laplace(model, likelihood='regression', prior_precision=precision,
+                  last_layer_name='1')
+    # torch.tensor 0-dim
+    precision = torch.tensor(10.6)
+    lap = laplace(model, likelihood='regression', prior_precision=precision,
+                  last_layer_name='1')
+    # torch.tensor 1-dim
+    precision = torch.tensor(10.7).reshape(-1)
+    lap = laplace(model, likelihood='regression', prior_precision=precision,
+                  last_layer_name='1')
+    # torch.tensor 1-dim param-shape
+    if laplace == KronLLLaplace:  # kron only supports per layer
+        with pytest.raises(ValueError):
+            precision = torch.tensor(10.7).reshape(-1).repeat(model.n_params)
+            lap = laplace(model, likelihood='regression', prior_precision=precision,
+                        last_layer_name='1')
+    else:
+        precision = torch.tensor(10.7).reshape(-1).repeat(model.n_params)
+        lap = laplace(model, likelihood='regression', prior_precision=precision,
+                    last_layer_name='1')
+    # torch.tensor 1-dim layer-shape
+    precision = torch.tensor(10.7).reshape(-1).repeat(model.n_layers)
+    lap = laplace(model, likelihood='regression', prior_precision=precision,
+                  last_layer_name='1')
+
+    # other than that should fail
+    # higher dim
+    precision = torch.tensor(10.6).reshape(1, 1)
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', prior_precision=precision,
+                      last_layer_name='1')
+    # unmatched dim
+    precision = torch.tensor(10.6).reshape(-1).repeat(17)
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', prior_precision=precision,
+                      last_layer_name='1')
+    # other datatype, only reals supported
+    precision = '1.5'
+    with pytest.raises(ValueError):
+        lap = laplace(model, likelihood='regression', prior_precision=precision,
+                      last_layer_name='1')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_prior_mean_and_scatter(laplace, model):
+    lap_scalar_mean = laplace(model, 'classification', last_layer_name='1',
+                              prior_precision=1e-2, prior_mean=1.)
+    assert torch.allclose(lap_scalar_mean.prior_mean, torch.tensor([1.]))
+    lap_tensor_mean = laplace(model, 'classification', last_layer_name='1',
+                              prior_precision=1e-2, prior_mean=torch.ones(1))
+    assert torch.allclose(lap_tensor_mean.prior_mean, torch.tensor([1.]))
+    lap_tensor_scalar_mean = laplace(model, 'classification', last_layer_name='1',
+                                     prior_precision=1e-2, prior_mean=torch.ones(1)[0])
+    assert torch.allclose(lap_tensor_scalar_mean.prior_mean, torch.tensor(1.))
+    lap_tensor_full_mean = laplace(model, 'classification', last_layer_name='1',
+                                   prior_precision=1e-2, prior_mean=torch.ones(20*2+2))
+    assert torch.allclose(lap_tensor_full_mean.prior_mean, torch.ones(20*2+2))
+    expected = lap_scalar_mean.scatter
+    assert expected.ndim == 0
+    assert torch.allclose(lap_tensor_mean.scatter, expected)
+    assert lap_tensor_mean.scatter.shape == expected.shape
+    assert torch.allclose(lap_tensor_scalar_mean.scatter, expected)
+    assert lap_tensor_scalar_mean.scatter.shape == expected.shape
+    assert torch.allclose(lap_tensor_full_mean.scatter, expected)
+    assert lap_tensor_full_mean.scatter.shape == expected.shape
+
+    # too many dims
+    with pytest.raises(ValueError):
+        prior_mean = torch.ones(20*2+2).unsqueeze(-1)
+        laplace(model, 'classification', last_layer_name='1',
+                prior_precision=1e-2, prior_mean=prior_mean)
+
+    # unmatched dim
+    with pytest.raises(ValueError):
+        prior_mean = torch.ones(20*2-3)
+        laplace(model, 'classification', last_layer_name='1',
+                prior_precision=1e-2, prior_mean=prior_mean)
+
+    # invalid argument type
+    with pytest.raises(ValueError):
+        laplace(model, 'classification', last_layer_name='1',
+                prior_precision=1e-2, prior_mean='72')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_laplace_init_temperature(laplace, model):
+    # valid float
+    T = 1.1
+    lap = laplace(model, likelihood='classification', temperature=T,
+                  last_layer_name='1')
+    assert lap.temperature == T
+
+
+@pytest.mark.parametrize('laplace,lh', product(flavors, ['classification', 'regression']))
+def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
+    if lh == 'classification':
+        loader = class_loader
+        sigma_noise = 1.
+    else:
+        loader = reg_loader
+        sigma_noise = 0.3
+    lap = laplace(model, lh, sigma_noise=sigma_noise, prior_precision=0.7)
+    lap.fit(loader)
+    assert lap.n_data == len(loader.dataset)
+    assert lap.n_outputs == model.output_size
+    X = loader.dataset.tensors[0]
+    f = model(X)
+    y = loader.dataset.tensors[1]
+    assert f.shape == torch.Size([10, 2])
+
+    # Test log likelihood (Train)
+    log_lik = lap.log_likelihood
+    # compute true log lik
+    if lh == 'classification':
+        log_lik_true = Categorical(logits=f).log_prob(y).sum()
+        assert torch.allclose(log_lik, log_lik_true)
+    else:
+        assert y.size() == f.size()
+        log_lik_true = Normal(loc=f, scale=sigma_noise).log_prob(y).sum()
+        assert torch.allclose(log_lik, log_lik_true)
+        # change likelihood and test again
+        lap.sigma_noise = 0.72
+        log_lik = lap.log_likelihood
+        log_lik_true = Normal(loc=f, scale=0.72).log_prob(y).sum()
+        assert torch.allclose(log_lik, log_lik_true)
+
+    # Test marginal likelihood
+    # lml = log p(y|f) - 1/2 theta @ prior_prec @ theta
+    #       + 1/2 logdet prior_prec - 1/2 log det post_prec
+    lml = log_lik_true
+    feature_extractor = FeatureExtractor(model, last_layer_name='1')
+    theta = parameters_to_vector(feature_extractor.last_layer.parameters()).detach()
+    assert torch.allclose(theta, lap.mean)
+    prior_prec = torch.diag(lap.prior_precision_diag)
+    assert prior_prec.shape == torch.Size([len(theta), len(theta)])
+    lml = lml - 1/2 * theta @ prior_prec @ theta
+    Sigma_0 = torch.inverse(prior_prec)
+    if laplace == DiagLLLaplace:
+        log_det_post_prec = lap.posterior_precision.log().sum()
+    else:
+        log_det_post_prec = lap.posterior_precision.logdet()
+    lml = lml + 1/2 * (prior_prec.logdet() - log_det_post_prec)
+    assert torch.allclose(lml, lap.log_marginal_likelihood())
+
+    # test sampling
+    torch.manual_seed(61)
+    samples = lap.sample(n_samples=1)
+    assert samples.shape == torch.Size([1, len(theta)])
+    samples = lap.sample(n_samples=1000000)
+    assert samples.shape == torch.Size([1000000, len(theta)])
+    mu_comp = samples.mean(dim=0)
+    mu_true = lap.mean
+    assert torch.allclose(mu_comp, mu_true, rtol=1)
+
+    # test functional variance
+    if laplace == FullLLLaplace:
+        Sigma = lap.posterior_covariance
+    elif laplace == KronLLLaplace:
+        Sigma = lap.posterior_precision.to_matrix(exponent=-1)
+    elif laplace == DiagLLLaplace:
+        Sigma = torch.diag(lap.posterior_variance)
+    _, phi = feature_extractor.forward_with_features(X)
+    Js, f = jacobians_naive(feature_extractor.last_layer, phi)
+    true_f_var = torch.einsum('mkp,pq,mcq->mkc', Js, Sigma, Js)
+    # test last-layer Jacobians
+    comp_Js, comp_f = lap.backend.last_layer_jacobians(lap.model, X)
+    assert torch.allclose(Js, comp_Js)
+    assert torch.allclose(f, comp_f)
+    comp_f_var = lap.functional_variance(comp_Js)
+    assert torch.allclose(true_f_var, comp_f_var, rtol=1e-4)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_regression_predictive(laplace, model, reg_loader):
+    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap.fit(reg_loader)
+    X, y = reg_loader.dataset.tensors
+    f = model(X)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive, functional variance tested already above.
+    f_mu, f_var = lap(X, pred_type='glm')
+    assert torch.allclose(f_mu, f)
+    assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
+    assert len(f_mu) == len(X)
+
+    # NN predictive (only diagonal variance estimation)
+    f_mu, f_var = lap(X, pred_type='nn')
+    assert f_mu.shape == f_var.shape
+    assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1]])
+    assert len(f_mu) == len(X)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_classification_predictive(laplace, model, class_loader):
+    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap.fit(class_loader)
+    X, y = class_loader.dataset.tensors
+    f = torch.softmax(model(X), dim=-1)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive
+    f_pred = lap(X, pred_type='glm', link_approx='mc', n_samples=100)
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    f_pred = lap(X, pred_type='glm', link_approx='probit')
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    f_pred = lap(X, pred_type='glm', link_approx='bridge')
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+
+    # NN predictive
+    f_pred = lap(X, pred_type='nn', n_samples=100)
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_regression_predictive_samples(laplace, model, reg_loader):
+    lap = laplace(model, 'regression', sigma_noise=0.3, prior_precision=0.7)
+    lap.fit(reg_loader)
+    X, y = reg_loader.dataset.tensors
+    f = model(X)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive, functional variance tested already above.
+    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+
+    # NN predictive (only diagonal variance estimation)
+    fsamples = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_classification_predictive_samples(laplace, model, class_loader):
+    lap = laplace(model, 'classification', prior_precision=0.7)
+    lap.fit(class_loader)
+    X, y = class_loader.dataset.tensors
+    f = torch.softmax(model(X), dim=-1)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(X, pred_type='linear')
+
+    # GLM predictive
+    fsamples = lap.predictive_samples(X, pred_type='glm', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+    assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
+
+    # NN predictive
+    f_pred = lap.predictive_samples(X, pred_type='nn', n_samples=100)
+    assert fsamples.shape == torch.Size([100, f.shape[0], f.shape[1]])
+    assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
new file mode 100644
index 00000000..bd5a07f3
--- /dev/null
+++ b/tests/test_matrix.py
@@ -0,0 +1,173 @@
+import pytest
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+
+from laplace.matrix import Kron, KronDecomposed
+from laplace.utils import kron as kron_prod
+from laplace.curvature import BackPackGGN
+from laplace.utils import block_diag
+from tests.utils import get_psd_matrix, jacobians_naive
+
+
+torch.set_default_tensor_type(torch.DoubleTensor)
+
+
+@pytest.fixture
+def model():
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    return model
+
+
+@pytest.fixture
+def small_model():
+    model = torch.nn.Sequential(nn.Linear(3, 5), nn.Tanh(), nn.Linear(5, 2))
+    setattr(model, 'output_size', 2)
+    return model
+
+
+def test_init_from_model(model):
+    kron = Kron.init_from_model(model, 'cpu')
+    expected_sizes = [[20*20, 3*3], [20*20], [2*2, 20*20], [2*2]]
+    for facs, exp_facs in zip(kron.kfacs, expected_sizes):
+        for fi, exp_fi in zip(facs, exp_facs):
+            assert torch.all(fi == 0)
+            assert np.prod(fi.shape) == exp_fi
+
+
+def test_addition(model):
+    kron = Kron.init_from_model(model, 'cpu')
+    expected_sizes = [[20, 3], [20], [2, 20], [2]]
+    to_add = Kron([[torch.randn(i, i) for i in sizes] for sizes in expected_sizes])
+    kron += to_add
+    for facs, exp_facs in zip(kron.kfacs, to_add.kfacs):
+        for fi, exp_fi in zip(facs, exp_facs):
+            assert torch.allclose(fi, exp_fi)
+
+def test_multiplication():
+    # kron * x should be the same as the expanded kronecker product * x
+    expected_sizes = [[20, 3], [20], [2, 20], [2]]
+    kfacs = [[torch.randn(i, i) for i in sizes] for sizes in expected_sizes]
+    kron = Kron(kfacs)
+    kron *= 1.5
+    for facs, exp_facs in zip(kron.kfacs, kfacs):
+        if len(facs) == 1:
+            assert torch.allclose(facs[0], 1.5 * exp_facs[0])
+        else:  # len(facs) == 2
+            exp = 1.5 * kron_prod(*exp_facs)
+            facs = kron_prod(*facs)
+            assert torch.allclose(exp, facs)
+
+def test_decompose():
+    expected_sizes = [[20, 3], [20], [2, 20], [2]]
+    P = 20 * 3 + 20 + 2 * 20 + 2
+    torch.manual_seed(7171)
+    kfacs = [[get_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
+    kron = Kron(kfacs)
+    kron_decomp = kron.decompose()
+    for facs, Qs, ls in zip(kron.kfacs, kron_decomp.eigenvectors, kron_decomp.eigenvalues):
+        if len(facs) == 1:
+            H, Q, l = facs[0], Qs[0], ls[0]
+            reconstructed = Q @ torch.diag(l) @ Q.T
+            assert torch.allclose(H, reconstructed, rtol=1e-3)
+        if len(facs) == 2:
+            gtruth = kron_prod(facs[0], facs[1])
+            rec_1 = Qs[0] @ torch.diag(ls[0]) @ Qs[0].T
+            rec_2 = Qs[1] @ torch.diag(ls[1]) @ Qs[1].T
+            reconstructed = kron_prod(rec_1, rec_2)
+            assert torch.allclose(gtruth, reconstructed, rtol=1e-2)
+    W = torch.randn(P)
+    SW_kron = kron.bmm(W)
+    SW_kron_decomp = kron_decomp.bmm(W, exponent=1)
+    assert torch.allclose(SW_kron, SW_kron_decomp)
+
+
+def test_logdet_consistent():
+    expected_sizes = [[20, 3], [20], [2, 20], [2]]
+    torch.manual_seed(7171)
+    kfacs = [[get_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
+    kron = Kron(kfacs)
+    kron_decomp = kron.decompose()
+    assert torch.allclose(kron.logdet(), kron_decomp.logdet())
+
+
+def test_bmm(small_model):
+    model = small_model
+    # model = single_output_model
+    X = torch.randn(5, 3)
+    y = torch.randn(5, 2)
+    backend = BackPackGGN(model, 'regression', stochastic=False)
+    loss, kron = backend.kron(X, y, N=5)
+    kron_decomp = kron.decompose()
+    Js, f = jacobians_naive(model, X)
+    blocks = list()
+    for F in kron.kfacs:
+        if len(F) == 1:
+            blocks.append(F[0])
+        else:
+            blocks.append(kron_prod(*F))
+    S = block_diag(blocks)
+    assert torch.allclose(S, S.T)
+    assert torch.allclose(S.diagonal(), kron.diag())
+
+    # test J @ Kron_decomp @ Jt (square form)
+    JS = kron_decomp.bmm(Js, exponent=1)
+    JS_true = Js @ S
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
+    JSJ = torch.bmm(JS, Js.transpose(1,2))
+    assert torch.allclose(JSJ, JSJ_true)
+    assert torch.allclose(JS, JS_true)
+
+    # test J @ Kron @ Jt (square form)
+    JS_nodecomp = kron.bmm(Js)
+    JSJ_nodecomp = torch.bmm(JS_nodecomp, Js.transpose(1,2))
+    assert torch.allclose(JSJ_nodecomp, JSJ)
+    assert torch.allclose(JS_nodecomp, JS)
+
+    # test J @ S_inv @ J (funcitonal variance)
+    JSJ = kron_decomp.inv_square_form(Js)
+    S_inv = S.inverse()
+    JSJ_true = torch.bmm(Js @ S_inv, Js.transpose(1,2))
+    assert torch.allclose(JSJ, JSJ_true)
+
+    # test J @ S^-1/2  (sampling)
+    JS = kron_decomp.bmm(Js, exponent=-1/2)
+    JSJ = torch.bmm(JS, Js.transpose(1,2))
+    l, Q = S_inv.symeig(eigenvectors=True)
+    JS_true = Js @ Q @ torch.diag(torch.sqrt(l)) @ Q.T
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
+    assert torch.allclose(JS, JS_true)
+    assert torch.allclose(JSJ, JSJ_true)
+
+    # test different Js shapes:
+    # 2 - dimensional
+    JS = kron_decomp.bmm(Js[:, 0, :].squeeze(), exponent=1)
+    JS_nodecomp = kron.bmm(Js[:, 0, :].squeeze())
+    JS_true = Js[:, 0, :].squeeze() @ S
+    assert torch.allclose(JS, JS_true)
+    assert torch.allclose(JS, JS_nodecomp)
+    # 1 - dimensional
+    JS = kron_decomp.bmm(Js[0, 0, :].squeeze(), exponent=1)
+    JS_nodecomp = kron.bmm(Js[0, 0, :].squeeze())
+    JS_true = Js[0, 0, :].squeeze() @ S
+    assert torch.allclose(JS, JS_true)
+    assert torch.allclose(JS, JS_nodecomp)
+
+
+def test_matrix_consistent():
+    expected_sizes = [[20, 3], [20], [2, 20], [2]]
+    torch.manual_seed(7171)
+    kfacs = [[get_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
+    kron = Kron(kfacs)
+    kron_decomp = kron.decompose()
+    assert torch.allclose(kron.to_matrix(), kron_decomp.to_matrix(exponent=1))
+    assert torch.allclose(kron.to_matrix().inverse(), kron_decomp.to_matrix(exponent=-1))
+    M_true = kron.to_matrix()
+    M_true.diagonal().add_(3.4)
+    kron_decomp += torch.tensor(3.4)
+    assert torch.allclose(M_true, kron_decomp.to_matrix(exponent=1))
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 00000000..2bf56262
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,36 @@
+import torch
+from laplace.utils import invsqrt_precision, diagonal_add_scalar, symeig
+
+
+def test_sqrt_precision():
+    X = torch.randn(20, 100)
+    M = X @ X.T
+    Scale = invsqrt_precision(M)
+    torch.allclose(Scale @ Scale.T, torch.inverse(M))
+
+
+def test_diagonal_add_scalar():
+    M = torch.randn(50, 50)
+    M_plus_diag = diagonal_add_scalar(M, 1.98)
+    diag = 1.98 * torch.eye(len(M))
+    torch.allclose(M_plus_diag, M + diag)
+
+
+def test_symeig_custom():
+    X = torch.randn(20, 100)
+    M = X @ X.T
+    l1, W1 = torch.symeig(M, eigenvectors=True)
+    l2, W2 = symeig(M)
+    assert torch.allclose(l1, l2)
+    assert torch.allclose(W1, W2)
+    
+
+def test_symeig_custom_low_rank():
+    X = torch.randn(1000, 10)
+    M = X @ X.T
+    l1, W1 = torch.symeig(M, eigenvectors=True)
+    l2, W2 = symeig(M)
+    # symeig should fail for low-rank
+    assert not torch.all(l1 >= 0.0)
+    # test clamping to zeros
+    assert torch.all(l2 >= 0.0)
\ No newline at end of file
diff --git a/tests/utils.py b/tests/utils.py
new file mode 100644
index 00000000..9488db17
--- /dev/null
+++ b/tests/utils.py
@@ -0,0 +1,33 @@
+import torch
+import numpy as np
+
+
+def get_psd_matrix(dim):
+    X = torch.randn(dim, dim*3)
+    return X @ X.T / (dim * 3)
+
+def grad(model):
+    return torch.cat([p.grad.data.flatten() for p in model.parameters()]).detach()
+
+def jacobians_naive(model, data):
+    model.zero_grad()
+    f = model(data)
+    Jacs = list()
+    for i in range(f.shape[0]):
+        if len(f.shape) > 1:
+            jacs = list()
+            for j in range(f.shape[1]):
+                rg = (i != (f.shape[0] - 1) or j != (f.shape[1] - 1))
+                f[i, j].backward(retain_graph=rg)
+                Jij = grad(model)
+                jacs.append(Jij)
+                model.zero_grad()
+            jacs = torch.stack(jacs).t()
+        else:
+            rg = (i != (f.shape[0] - 1))
+            f[i].backward(retain_graph=rg)
+            jacs = grad(model)
+            model.zero_grad()
+        Jacs.append(jacs)
+    Jacs = torch.stack(Jacs).transpose(1, 2)
+    return Jacs.detach(), f.detach()