Repository: aleximmer/Laplace (language: python)

Issue:
Temperature scaling naming issue
We call it temperature but its actually `1/temperature`
Adjust to public ASDL library
All changes necessary to adjust the library to the public [ASDL library](https://github.com/kazukiosawa/asdfghjkl): Correct install instructions for travis and in the README, remove the `dependencies` directory, and rename backend.


Discussion / hints:

Fix patch:
diff --git a/.travis.yml b/.travis.yml
index 098b550b..309bcab3 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -4,7 +4,6 @@ python:
 install:
 - pip install -r requirements.txt
 - pip install -r tests/requirements.txt
-- pip install dependencies/asdfghjkl-0.0.1-py3-none-any.whl
 - python setup.py install
 script:
 - pytest -vx --cov=laplace/ tests/
diff --git a/README.md b/README.md
index 8122e2c2..11976b0f 100644
--- a/README.md
+++ b/README.md
@@ -3,8 +3,12 @@
 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://d2ca8dae-c541-11eb-9203-acde48001122.github.io/lapras/index.html](https://d2ca8dae-c541-11eb-9203-acde48001122.github.io/lapras/index.html).
+
 ## Setup
 
+We assume `python3.8` since the package was developed with that version.
+
 ```bash
 pip install -r requirements.txt
 # for development
@@ -12,13 +16,18 @@ pip install -e .
 # for "production"
 pip install .
 
-# for Kazuki's backend
-cd dependencies
-pip install asdfghjkl-0.0.1-py3-none-any.whl
-
 
 # run tests
-pip install pytest
+pip install -r tests/requirements.txt
 pytest tests/
 ```
 
+## [Documentation](https://d2ca8dae-c541-11eb-9203-acde48001122.github.io/lapras/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/dependencies/asdfghjkl-0.0.1-py3-none-any.whl b/dependencies/asdfghjkl-0.0.1-py3-none-any.whl
deleted file mode 100644
index 0b8166d6..00000000
Binary files a/dependencies/asdfghjkl-0.0.1-py3-none-any.whl and /dev/null differ
diff --git a/docs/baselaplace.html b/docs/baselaplace.html
new file mode 100644
index 00000000..4ef53e54
--- /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 (if applicable).</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 (if applicable).</dd>
+<dt>init_prior_prec : float, default=1.</dt>
+<dt>initial prior precision before the first optimization step (if applicable).</dt>
+<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></dt>
+<dd>how to approximate the classification link function for the <code>'glm'</code>.
+For <code>pred_type='nn'</code>, only 'mc' is possible.</dd>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</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..07a42e6e
--- /dev/null
+++ b/docs/index.html
@@ -0,0 +1,1056 @@
+<!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="Setup …" />
+<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">
+<h2 id="setup">Setup</h2>
+<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 <code><a title="laplace.BaseLaplace" href="#laplace.BaseLaplace">BaseLaplace</a></code> that implement different sparsity structures
+(last-layer <code><a title="laplace.LLLaplace" href="#laplace.LLLaplace">LLLaplace</a></code> vs. <code><a title="laplace.FullLaplace" href="#laplace.FullLaplace">FullLaplace</a></code> and <code><a title="laplace.DiagLaplace" href="#laplace.DiagLaplace">DiagLaplace</a></code> vs. <code><a title="laplace.KronLaplace" href="#laplace.KronLaplace">KronLaplace</a></code>) all of which
+can be conveniently accessed with the method <code><a title="laplace.Laplace" href="#laplace.Laplace">Laplace()</a></code>.</li>
+<li>The backends in <code><a title="laplace.curvature" href="curvature/index.html">laplace.curvature</a></code> which provide access to Hessian approximations of
+the corresponding sparsity structures, for example, the diagonal GGN.</li>
+</ol>
+<p>Additionally, the package provies utilities for
+decomposing a neural network into feature extractor and last layer for <code><a title="laplace.LLLaplace" href="#laplace.LLLaplace">LLLaplace</a></code> (<code><a title="laplace.feature_extractor" href="feature_extractor.html">laplace.feature_extractor</a></code>)
+and
+effectively dealing with Kronecker factors (<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>).</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 allows to provide different Hessian
+approximations to the Laplace approximations.
+For example, currently the <code><a title="laplace.curvature.BackPackInterface" href="curvature/index.html#laplace.curvature.BackPackInterface">BackPackInterface</a></code> and <code><a title="laplace.curvature.AsdlInterface" href="curvature/index.html#laplace.curvature.AsdlInterface">AsdlInterface</a></code> 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>
+<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 (if applicable).</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 (if applicable).</dd>
+<dt>init_prior_prec : float, default=1.</dt>
+<dt>initial prior precision before the first optimization step (if applicable).</dt>
+<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></dt>
+<dd>how to approximate the classification link function for the <code>'glm'</code>.
+For <code>pred_type='nn'</code>, only 'mc' is possible.</dd>
+<dt><strong><code>n_samples</code></strong> :&ensp;<code>int</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="#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>
+<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>
+</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/documentation.md b/documentation.md
new file mode 100644
index 00000000..e752dbc1
--- /dev/null
+++ b/documentation.md
@@ -0,0 +1,87 @@
+## Setup
+
+```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` that implement different sparsity structures
+(last-layer `LLLaplace` vs. `FullLaplace` and `DiagLaplace` vs. `KronLaplace`) all of which
+can be conveniently accessed with the method `Laplace`.
+2. The backends in `laplace.curvature` which provide access to Hessian approximations of
+the corresponding sparsity structures, for example, the diagonal GGN.
+
+Additionally, the package provies utilities for
+decomposing a neural network into feature extractor and last layer for `LLLaplace` (`laplace.feature_extractor`)
+and
+effectively dealing with Kronecker factors (`matrix.Kron` and `matrix.KronDecomposed`).
+
+## 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 allows to provide different Hessian
+approximations to the Laplace approximations.
+For example, currently the `curvature.BackPackInterface` based on [BackPACK](https://github.com/f-dangel/backpack/) and `curvature.AsdlInterface` 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()
+```
diff --git a/laplace/__init__.py b/laplace/__init__.py
index b4ed2c25..b9f21838 100644
--- a/laplace/__init__.py
+++ b/laplace/__init__.py
@@ -1,6 +1,17 @@
+"""
+.. include:: ../documentation.md
+
+.. include:: ../regression_example.md
+"""
 REGRESSION = 'regression'
 CLASSIFICATION = 'classification'
 
-from laplace.baselaplace import *
-from laplace.lllaplace import *
+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
index 803ea876..9e99e2bc 100644
--- a/laplace/baselaplace.py
+++ b/laplace/baselaplace.py
@@ -1,4 +1,4 @@
-from abc import ABC, abstractmethod
+from abc import ABC, abstractmethod, abstractproperty
 from math import sqrt, pi
 import numpy as np
 import torch
@@ -10,50 +10,53 @@
 from laplace.curvature import BackPackGGN
 
 
-__all__ = ['FullLaplace', 'KronLaplace', 'DiagLaplace']
+__all__ = ['BaseLaplace', 'FullLaplace', 'KronLaplace', 'DiagLaplace']
 
 
 class BaseLaplace(ABC):
-    """Laplace approximation for a pytorch neural network.
-    The Laplace approximation is a Gaussian distribution but can have different
-    sparsity structures. Further, it provides an approximation to the marginal
-    likelihood.
+    """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
-        torch model
-
-    likelihood : str
-        'classification' or 'regression' are supported
-
-    sigma_noise : float
-        observation noise for likelihood = 'regression'
-
-    prior_precision : one-dimensional torch.Tensor, str, default='auto'
-        prior precision of a Gaussian prior corresponding to weight decay
-        'auto' determines the prior automatically during fitting
-
+    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
-        posterior temperature scaling affects the posterior covariance as
-        `Sigma' = temperature * Sigma`, so low temperatures lead to a more
-        concentrated posterior.
-
-    backend : CurvatureInterface
-        provides access to curvature/second-order quantities.
-
-    Attributes
-    ----------
-
-    Examples
-    --------
-
-    Notes
-    -----
+        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):
+                 prior_mean=0., temperature=1., backend=BackPackGGN, backend_kwargs=None):
         if likelihood not in ['classification', 'regression']:
             raise ValueError(f'Invalid likelihood type {likelihood}')
 
@@ -73,7 +76,7 @@ def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
         self.temperature = temperature
         self._backend = None
         self._backend_cls = backend
-        self._backend_kwargs = backend_kwargs
+        self._backend_kwargs = dict() if backend_kwargs is None else backend_kwargs
         self.H = None
 
         # log likelihood = g(loss)
@@ -105,8 +108,9 @@ def fit(self, train_loader):
 
         Parameters
         ----------
-        train_loader : iterator
-            each iterate is a training batch (X, y)
+        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.')
@@ -130,10 +134,25 @@ def fit(self, train_loader):
 
         self.n_data = N
 
-    def marginal_likelihood(self, prior_precision=None, sigma_noise=None):
-        """Compute the Laplace approximation to the marginal likelihood.
-        The resulting value is differentiable in differentiable likelihood
-        and prior parameters.
+    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()
@@ -148,13 +167,22 @@ def marginal_likelihood(self, prior_precision=None, sigma_noise=None):
                 raise ValueError('Can only change sigma_noise for regression.')
             self.sigma_noise = sigma_noise
 
-        return self.log_lik - 0.5 * (self.log_det_ratio + self.scatter)
+        return self.log_likelihood - 0.5 * (self.log_det_ratio + self.scatter)
 
     @property
-    def log_lik(self):
+    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
+        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))
@@ -163,31 +191,44 @@ def log_lik(self):
             # for classification Xent == log Cat
             return factor * self.loss
 
-    def __call__(self, X, pred_type='glm', link_approx='mc', n_samples=100):
+    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_size)
+        x : torch.Tensor
+            `(batch_size, input_shape)`
 
-        pred_type : str 'lin' or 'nn'
-            type of posterior predictive, linearized GLM predictive or
-            neural network sampling predictive.
+        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 : str 'mc' or 'probit'
-            how to approximate the classification link function for GLM.
-            For NN, only 'mc' is possible.
+        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 in case necessary
+            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)
+            f_mu, f_var = self._glm_predictive_distribution(x)
             # regression
             if self.likelihood == 'regression':
                 return f_mu, f_var
@@ -209,27 +250,35 @@ def __call__(self, X, pred_type='glm', link_approx='mc', n_samples=100):
                 dist = Dirichlet(alpha)
                 return torch.nan_to_num(dist.mean, nan=1.0)
         else:
-            samples = self.nn_predictive_samples(X, n_samples)
+            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(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):
-        """Compute the posterior predictive on input data `X`.
+    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 (batch_size, *input_size)
+        x : torch.Tensor
+            input data `(batch_size, input_shape)`
 
-        pred_type : str 'lin' or 'nn'
-            type of posterior predictive, linearized GLM predictive or
-            neural network sampling predictive.
+        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()
 
@@ -237,7 +286,7 @@ def predictive_samples(self, X, pred_type='glm', n_samples=100):
             raise ValueError('Only glm and nn supported as prediction types.')
 
         if pred_type == 'glm':
-            f_mu, f_var = self.glm_predictive_distribution(X)
+            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,))
@@ -246,15 +295,15 @@ def predictive_samples(self, X, pred_type='glm', n_samples=100):
             return torch.softmax(samples, dim=-1)
 
         else:  # 'nn'
-            return self.nn_predictive_samples(X, n_samples)
+            return self._nn_predictive_samples(x, n_samples)
 
     @torch.enable_grad()
-    def glm_predictive_distribution(self, X):
+    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):
+    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())
@@ -265,34 +314,26 @@ def nn_predictive_samples(self, X, n_samples=100):
             fs = torch.softmax(fs, dim=-1)
         return fs
 
-    @abstractmethod
-    def sample(self, n_samples=100):
-        """Sample from the Laplace posterior torch.Tensor (n_samples, P)"""
-        pass
-
-    @property
-    def scatter(self):
-        """Computes the scatter used for the marginal likelihood `m^T P_0 m`
-        """
-        delta = (self.mean - self.prior_mean)
-        return (delta * self.prior_precision_diag) @ delta
-
-    @property
-    def log_det_prior_precision(self):
-        """Computes log determinant of the prior precision `log det P_0`
-        """
-        return self.prior_precision_diag.log().sum()
-
     @abstractmethod
     def functional_variance(self, Jacs):
-        """Compute functional variance for the predictive:
-        `f_var[i] = Jacs[i] @ Sigma @ Jacs[i].T`, which is a output x output
+        """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 batch x outputs x parameters
-            Jacobians of model output wrt parameters.
+        Jacs : torch.Tensor
+            Jacobians of model output wrt parameters
+            `(batch, outputs, parameters)`
+
+        Returns
+        -------
+        f_var : torch.Tensor
+            output covariance `(batch, outputs, outputs)`
         """
         pass
 
@@ -305,21 +346,77 @@ def _check_jacobians(self, Js):
         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 log_det_ratio(self):
-        """Computes the log of the determinant ratios for the marginal likelihood
-        `log (det P / det P_0) = log det P - log det P_0`
+    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]
         """
-        return self.log_det_posterior_precision - self.log_det_prior_precision
+        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):
-        """Computes log determinant of the posterior precision `log det P`
+        """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
         """
-        raise NotImplementedError()
+        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)
 
@@ -354,7 +451,6 @@ def prior_mean(self, prior_mean):
         else:
             raise ValueError('Invalid argument type of prior mean.')
 
-    # TODO: protect prior precision and sigma updates when covariance computed/update covariance?
     @property
     def prior_precision(self):
         return self._prior_precision
@@ -382,6 +478,42 @@ def optimize_prior_precision(self, method='marglik', n_steps=100, lr=1e-1,
                                  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 (if applicable).
+        lr : float, default=1e-1
+            the learning rate to use for gradient descent (if applicable).
+        init_prior_prec : float, default=1.
+            initial prior precision before the first optimization step (if applicable).
+        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'}
+            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'`.
+        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()
@@ -390,8 +522,8 @@ def optimize_prior_precision(self, method='marglik', n_steps=100, lr=1e-1,
             for _ in range(n_steps):
                 optimizer.zero_grad()
                 prior_prec = log_prior_prec.exp()
-                neg_marglik = -self.marginal_likelihood(prior_precision=prior_prec)
-                neg_marglik.backward()
+                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':
@@ -433,6 +565,7 @@ def sigma_noise(self):
 
     @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):
@@ -448,23 +581,35 @@ def sigma_noise(self, sigma_noise):
             raise ValueError('Invalid type: sigma noise needs to be torch.Tensor or scalar.')
 
     @property
-    def H_factor(self):
+    def _H_factor(self):
         sigma2 = self.sigma_noise.square()
-        return 1 / sigma2 * self.temperature
+        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):
-    # TODO: list additional attributes
-    # TODO: recompute scale once prior precision or sigma noise change?
-    #       do in lazy way with a flag probably.
-
+    """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')
+    _key = ('all', 'full')
 
     def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., backend=BackPackGGN, **backend_kwargs):
+                 prior_mean=0., temperature=1., backend=BackPackGGN, backend_kwargs=None):
         super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, backend, **backend_kwargs)
+                         prior_mean, temperature, backend, backend_kwargs)
         self._posterior_scale = None
 
     def _init_H(self):
@@ -478,27 +623,44 @@ def _compute_scale(self):
 
     @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):
-        """Computes log determinant of the posterior precision `log det P`
+        """Posterior precision \\(P\\).
+
+        Returns
+        -------
+        precision : torch.tensor
+            `(parameters, parameters)`
         """
         self._check_fit()
-        return self.H_factor * self.H + torch.diag(self.prior_precision_diag)
+        return self._H_factor * self.H + torch.diag(self.prior_precision_diag)
 
     @property
     def log_det_posterior_precision(self):
-        """Computes log determinant of the posterior precision `log det P`
-        """
-        # TODO: could make more efficient for scalar prior precision.
         return self.posterior_precision.logdet()
 
     def functional_variance(self, Js):
@@ -510,10 +672,19 @@ def sample(self, n_samples=100):
 
 
 class KronLaplace(BaseLaplace):
-    # TODO list additional attributes
-
+    """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')
+    _key = ('all', 'kron')
 
     def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
                  prior_mean=0., temperature=1., backend=BackPackGGN, damping=False,
@@ -537,13 +708,17 @@ def fit(self, train_loader, keep_factors=False):
 
     @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
+        return self.H * self._H_factor + self.prior_precision
 
     @property
     def log_det_posterior_precision(self):
-        """Computes log determinant of the posterior precision `log det P`
-        """
         return self.posterior_precision.logdet()
 
     def functional_variance(self, Js):
@@ -563,16 +738,13 @@ def prior_precision(self, prior_precision):
 
 
 class DiagLaplace(BaseLaplace):
-    # TODO: list additional attributes
-    # TODO: caching prior_precision_diag for fast lazy computation?
-
+    """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__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., backend=BackPackGGN, **backend_kwargs):
-        super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, backend, **backend_kwargs)
+    _key = ('all', 'diag')
 
     def _init_H(self):
         self.H = torch.zeros(self.n_params, device=self._device)
@@ -582,23 +754,40 @@ def _curv_closure(self, X, y, N):
 
     @property
     def posterior_precision(self):
-        """Computes log determinant of the posterior precision `log det P`
+        """Diagonal posterior precision \\(p\\).
+
+        Returns
+        -------
+        precision : torch.tensor
+            `(parameters)`
         """
         self._check_fit()
-        return self.H_factor * self.H + self.prior_precision_diag
+        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):
-        """Computes log determinant of the posterior precision `log det P`
-        """
         return self.posterior_precision.log().sum()
 
     def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
diff --git a/laplace/curvature/__init__.py b/laplace/curvature/__init__.py
index 66462d9f..93e6cc8f 100644
--- a/laplace/curvature/__init__.py
+++ b/laplace/curvature/__init__.py
@@ -1,6 +1,6 @@
 import logging
 
-from laplace.curvature.curvature import CurvatureInterface
+from laplace.curvature.curvature import CurvatureInterface, GGNInterface, EFInterface
 
 try:
     from laplace.curvature.backpack import BackPackGGN, BackPackEF, BackPackInterface
@@ -8,6 +8,10 @@
     logging.info('Backpack not available.')
 
 try:
-    from laplace.curvature.kazuki import KazukiGGN, KazukiEF, KazukiInterface
+    from laplace.curvature.asdl import AsdlGGN, AsdlEF, AsdlInterface
 except ModuleNotFoundError:
-    logging.info('Kazuki backend not available.')
+    logging.info('asdfghjkl backend not available.')
+
+__all__ = ['CurvatureInterface', 'GGNInterface', 'EFInterface',
+           'BackPackInterface', 'BackPackGGN', 'BackPackEF',
+           'AsdlInterface', 'AsdlGGN', 'AsdlEF']
diff --git a/laplace/curvature/kazuki.py b/laplace/curvature/asdl.py
similarity index 57%
rename from laplace/curvature/kazuki.py
rename to laplace/curvature/asdl.py
index 61716d3d..94a2580b 100644
--- a/laplace/curvature/kazuki.py
+++ b/laplace/curvature/asdl.py
@@ -1,3 +1,4 @@
+from abc import abstractproperty
 import warnings
 import numpy as np
 import torch
@@ -7,42 +8,80 @@
 from asdfghjkl import fisher_for_cross_entropy
 from asdfghjkl.gradient import batch_gradient
 
-from laplace.curvature import CurvatureInterface
+from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
 from laplace.matrix import Kron
-from laplace.utils import is_batchnorm
+from laplace.utils import _is_batchnorm
 
 
-class KazukiInterface(CurvatureInterface):
-
+class AsdlInterface(CurvatureInterface):
+    """Interface for asdfghjkl backend.
+    """
     def __init__(self, model, likelihood, last_layer=False):
         if likelihood != 'classification':
-            raise ValueError('This backend does only support classification currently.')
+            raise ValueError('This backend only supports classification currently.')
         super().__init__(model, likelihood, last_layer)
 
     @staticmethod
-    def jacobians(model, X):
+    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()
+            f = batch_gradient(model, loss_fn, x, None).detach()
             Js.append(_get_batch_grad(model))
         Js = torch.stack(Js, dim=1)
         return Js, f
 
-    @property
-    def ggn_type(self):
+    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):
+            if _is_batchnorm(module):
                 warnings.warn('BatchNorm unsupported for Kron, ignore.')
                 continue
 
-            stats = getattr(module, self.ggn_type, None)
+            stats = getattr(module, self._ggn_type, None)
             if stats is None:
                 continue
             if hasattr(module, 'bias') and module.bias is not None:
@@ -57,7 +96,7 @@ def _get_kron_factors(self, curv, M):
                     kfacs.append([stats.kron.B, stats.kron.A])
             else:
                 raise ValueError(f'Whats happening with {module}?')
-        return kfacs
+        return Kron(kfacs)
 
     @staticmethod
     def _rescale_kron_factors(kron, N):
@@ -73,7 +112,8 @@ def diag(self, X, y, **kwargs):
             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)
+        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
 
@@ -84,35 +124,36 @@ def kron(self, X, y, N, **wkwargs) -> [torch.Tensor, Kron]:
             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)
+        curv = fisher_for_cross_entropy(self._model, self._ggn_type, SHAPE_KRON,
+                                        inputs=X, targets=y)
         M = len(y)
-        kron = Kron(self._get_kron_factors(curv, M))
+        kron = self._get_kron_factors(curv, M)
         kron = self._rescale_kron_factors(kron, N)
         return self.factor * loss, self.factor * kron
 
-    def full(self, X, y, **kwargs):
-        raise NotImplementedError()
-
-
-class KazukiGGN(KazukiInterface):
 
+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):
+    def _ggn_type(self):
         return FISHER_MC if self.stochastic else FISHER_EXACT
 
 
-class KazukiEF(KazukiInterface):
+class AsdlEF(AsdlInterface, EFInterface):
+    """Implementation of the `EFInterface` using asdfghjkl.
+    """
 
     @property
-    def ggn_type(self):
+    def _ggn_type(self):
         return COV
 
 
-def flatten_after_batch(tensor: torch.Tensor):
+def _flatten_after_batch(tensor: torch.Tensor):
     if tensor.ndim == 1:
         return tensor.unsqueeze(-1)
     else:
@@ -125,9 +166,9 @@ def _get_batch_grad(model):
         if hasattr(module, 'op_results'):
             res = module.op_results['batch_grads']
             if 'weight' in res:
-                batch_grads.append(flatten_after_batch(res['weight']))
+                batch_grads.append(_flatten_after_batch(res['weight']))
             if 'bias' in res:
-                batch_grads.append(flatten_after_batch(res['bias']))
+                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
index 5e280ae9..885ee2b9 100644
--- a/laplace/curvature/backpack.py
+++ b/laplace/curvature/backpack.py
@@ -4,33 +4,41 @@
 from backpack.extensions import DiagGGNExact, DiagGGNMC, KFAC, KFLR, SumGradSquared, BatchGrad
 from backpack.context import CTX
 
-from laplace.curvature import CurvatureInterface
+from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
 from laplace.matrix import Kron
 
 
-def cleanup(module):
-    for child in module.children():
-        cleanup(child)
-
-    setattr(module, "_backpack_extend", False)
-    memory_cleanup(module)
-
-
 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):
-        # Jacobians are batch x output x params
+    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)
+            out = model(x)
             with backpack(BatchGrad()):
                 if model.output_size > 1:
                     out[:, i].sum().backward()
@@ -38,31 +46,49 @@ def jacobians(model, X):
                     out.sum().backward()
                 to_cat = []
                 for param in model.parameters():
-                    to_cat.append(param.grad_batch.detach().reshape(X.shape[0], -1))
+                    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()
 
-        # cleanup
         model.zero_grad()
         CTX.remove_hooks()
-        cleanup(model)
+        _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):
-    """[summary]
 
-    MSELoss = |y-f|_2^2 -> d/df = -2(y-f)
-    log N(y|f,1) propto 1/2|y-f|_2^2 -> d/df = -(y-f)
-    --> factor for regression is 0.5 for loss and ggn
+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
@@ -98,7 +124,7 @@ def diag(self, X, y, **kwargs):
 
         return self.factor * loss.detach(), self.factor * dggn
 
-    def kron(self, X, y, N, **wkwargs) -> [torch.Tensor, Kron]:
+    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)
@@ -109,25 +135,10 @@ def kron(self, X, y, N, **wkwargs) -> [torch.Tensor, Kron]:
 
         return self.factor * loss.detach(), self.factor * kron
 
-    def full(self, X, y, **kwargs):
-        # TODO: put in shared GGN interaface for both backends
-        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 BackPackEF(BackPackInterface):
 
-    def _get_individual_gradients(self):
-        return torch.cat([p.grad_batch.data.flatten(start_dim=1)
-                          for p in self._model.parameters()], dim=1)
+class BackPackEF(BackPackInterface, EFInterface):
+    """Implementation of `EFInterface` using Backpack.
+    """
 
     def diag(self, X, y, **kwargs):
         f = self.model(X)
@@ -140,15 +151,12 @@ def diag(self, X, y, **kwargs):
         return self.factor * loss.detach(), self.factor * diag_EF
 
     def kron(self, X, y, **kwargs):
-        raise NotImplementedError()
+        raise NotImplementedError('Unavailable through Backpack.')
 
-    def full(self, X, y, **kwargs):
-        # TODO: put in shared EF interface for both kazuki and backpack
-        f = self.model(X)
-        loss = self.lossfunc(f, y)
-        with backpack(BatchGrad()):
-            loss.backward()
-        # TODO: implement similarly to jacobians as staticmethod
-        Gs = self._get_individual_gradients()
-        H_ef = Gs.T @ Gs
-        return self.factor * loss.detach(), self.factor * H_ef
+
+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
index d69e4aa1..485fd94f 100644
--- a/laplace/curvature/curvature.py
+++ b/laplace/curvature/curvature.py
@@ -4,7 +4,27 @@
 
 
 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
@@ -12,7 +32,7 @@ def __init__(self, model, likelihood, last_layer=False):
         self.last_layer = last_layer
         if likelihood == 'regression':
             self.lossfunc = MSELoss(reduction='sum')
-            self.factor = 0.5  # convert to standard Gauss. log N(y|f,1)
+            self.factor = 0.5
         else:
             self.lossfunc = CrossEntropyLoss(reduction='sum')
             self.factor = 1.
@@ -22,17 +42,47 @@ def _model(self):
         return self.model.last_layer if self.last_layer else self.model
 
     @abstractstaticmethod
-    def jacobians(model, X):
-        raise NotImplementedError()
+    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):
-        f, phi = model.forward_with_features(X)
-        bsize = len(X)
+    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)
+        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:
@@ -40,19 +90,125 @@ def last_layer_jacobians(model, X):
 
         return Js, f.detach()
 
-    @abstractmethod
-    def full(self, X, y, **kwargs):
+    @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 kron(self, X, y, **kwargs):
+    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 diag(self, X, y, **kwargs):
+    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)
@@ -62,3 +218,75 @@ def _get_full_ggn(self, Js, f, y):
             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
index b44503d7..784d6370 100644
--- a/laplace/feature_extractor.py
+++ b/laplace/feature_extractor.py
@@ -3,64 +3,76 @@
 from typing import Tuple, Callable, Optional
 
 
+__all__ = ['FeatureExtractor']
+
+
 class FeatureExtractor(nn.Module):
     """Feature extractor for a PyTorch neural network.
-    A wrapper which returns the output of the penultimate layer in addition to
-    the output of the last layer for each forward pass. It assumes that the
-    last layer is linear.
+    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.
 
-    Arguments
+    Parameters
     ----------
     model : torch.nn.Module
         PyTorch model
-
-    last_layer_name (optional) : str, default=None
-        If the user already knows the name of the last layer, otherwise it will
+    last_layer_name : str, default=None
+        if the name of the last layer is already known, otherwise it will
         be determined automatically.
-
-    Attributes
-    ----------
-    model : torch.nn.Module
-        The underlying PyTorch model.
-
-    last_layer : torch.nn.module
-        The torch module corresponding to the last layer (has to be instance
-        of torch.nn.Linear).
-
-    Examples
-    --------
-
-    Notes
-    -----
-    Limitations:
-        - Assumes that the last layer is always the same for any forward pass
-        - Assumes that the last layer is an instance of torch.nn.Linear
     """
     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._found = False
+            self.last_layer = None
         else:
             self.set_last_layer(last_layer_name)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self._found:
+        """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)
-        else:
-            # if this is the first forward pass
-            out = self.find_last_layer(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]
@@ -70,9 +82,6 @@ def set_last_layer(self, last_layer_name: str) -> None:
         # set forward hook to extract features in future forward passes
         self.last_layer.register_forward_hook(self._get_hook(last_layer_name))
 
-        # last layer is now identified and hook is set
-        self._found = True
-
     def _get_hook(self, name: str) -> Callable:
         def hook(_, input, __):
             # only accepts one input (expects linear layer)
@@ -80,16 +89,28 @@ def hook(_, input, __):
         return hook
 
     def find_last_layer(self, x: torch.Tensor) -> torch.Tensor:
-        if self._found:
+        """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)
-                if isinstance(input[0], torch.Tensor):
+                try:
                     act_out[name] = input[0].detach()
-                else:
+                except (IndexError, AttributeError):
                     act_out[name] = None
                 # remove hook
                 handles[name].remove()
diff --git a/laplace/laplace.py b/laplace/laplace.py
index fe67fd20..508631cb 100644
--- a/laplace/laplace.py
+++ b/laplace/laplace.py
@@ -2,9 +2,26 @@
 from laplace import *
 
 
-def Laplace(model, likelihood, subset_of_weights='last_layer', hessian_structure='kron', *args, **kwargs):
-    laplace_map = {subclass.key: subclass for subclass in _all_subclasses(BaseLaplace)
-                   if hasattr(subclass, 'key')}
+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)
 
diff --git a/laplace/lllaplace.py b/laplace/lllaplace.py
index 3bd6a7ea..79fd460b 100644
--- a/laplace/lllaplace.py
+++ b/laplace/lllaplace.py
@@ -12,63 +12,70 @@
 
 
 class LLLaplace(BaseLaplace):
-    """Last-Layer Laplace approximation for a pytorch neural network.
-    The Laplace approximation is a Gaussian distribution but can have different
-    sparsity structures. Further, it provides an approximation to the marginal
-    likelihood.
+    """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
-        torch model
-
-    likelihood : str
-        'classification' or 'regression' are supported
-
-    sigma_noise : float
-        observation noise for likelihood = 'regression'
-
-    prior_precision : one-dimensional torch.Tensor, str, default='auto'
-        prior precision of a Gaussian prior corresponding to weight decay
-        'auto' determines the prior automatically during fitting
-
+    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
-        posterior temperature scaling affects the posterior covariance as
-        `Sigma' = temperature * Sigma`, so low temperatures lead to a more
-        concentrated posterior.
-
-    backend : CurvatureInterface
-        provides access to curvature/second-order quantities.
-
-    Attributes
-    ----------
-
-    Examples
-    --------
-
-    Notes
-    -----
+        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):
+                 backend_kwargs=None):
         super().__init__(model, likelihood, sigma_noise=sigma_noise, prior_precision=1.,
-                         prior_mean=0., temperature=temperature, backend=backend, **backend_kwargs)
+                         prior_mean=0., temperature=temperature, backend=backend,
+                         backend_kwargs=backend_kwargs)
         self.model = FeatureExtractor(model, last_layer_name=last_layer_name)
-        if self.model._found:
-            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
-        else:
+        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):
@@ -76,15 +83,16 @@ def fit(self, train_loader):
 
         Parameters
         ----------
-        train_loader : iterator
-            each iterate is a training batch (X, y)
+        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 not self.model._found:
+        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))
@@ -97,12 +105,12 @@ def fit(self, train_loader):
 
         super().fit(train_loader)
 
-    def glm_predictive_distribution(self, X):
+    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):
+    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())
@@ -115,6 +123,13 @@ def nn_predictive_samples(self, X, n_samples=100):
 
     @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)
 
@@ -126,43 +141,59 @@ def prior_precision_diag(self):
 
 
 class FullLLLaplace(LLLaplace, FullLaplace):
-    # TODO list additional attributes
-
+    """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')
+    _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):
+                 backend_kwargs=None):
         super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, backend, last_layer_name, **backend_kwargs)
+                         prior_mean, temperature, backend, last_layer_name, backend_kwargs)
 
 
 class KronLLLaplace(LLLaplace, KronLaplace):
-    # TODO list additional attributes
-
+    """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')
+    _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)
+                         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):
-    # TODO list additional attributes
-
+    """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')
+    _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):
+                 backend_kwargs=None):
         super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, backend, last_layer_name, **backend_kwargs)
+                         prior_mean, temperature, backend, last_layer_name, backend_kwargs)
diff --git a/laplace/marglik.py b/laplace/marglik.py
deleted file mode 100644
index 451e8d1b..00000000
--- a/laplace/marglik.py
+++ /dev/null
@@ -1,246 +0,0 @@
-from copy import deepcopy
-import numpy as np
-import torch
-from torch.optim import Adam, SGD
-from torch.optim.lr_scheduler import ExponentialLR, CosineAnnealingLR
-from torch.nn import CrossEntropyLoss, MSELoss
-from torch.nn.utils import parameters_to_vector
-import logging
-
-from laplace import (DiagLaplace, KronLaplace, FullLaplace,
-                     DiagLLLaplace, KronLLLaplace, FullLLLaplace)
-from laplace.curvature import KazukiGGN
-
-
-def marglik_optimization(model,
-                         train_loader,
-                         valid_loader=None,
-                         likelihood='classification',
-                         prior_structure='layerwise',
-                         prior_prec_init=1.,
-                         sigma_noise_init=1.,
-                         temperature=1.,
-                         n_epochs=500,
-                         lr=1e-3,
-                         lr_min=None,
-                         optimizer='Adam',
-                         scheduler='exp',
-                         n_epochs_burnin=0,
-                         n_hypersteps=100,
-                         marglik_frequency=1,
-                         lr_hyp=1e-1,
-                         laplace=KronLaplace,
-                         backend=KazukiGGN,
-                         backend_kwargs=None):
-    """Runs marglik optimization training for a given model and training dataloader.
-
-    Parameters
-    ----------
-    model : torch.nn.Module
-        torch model
-    train_loader : DataLoader
-        pytorch training dataset loader
-    likelihood : str
-        'classification' or 'regression'
-    prior_structure : str
-        'scalar', 'layerwise', 'diagonal'
-    temperature : float default=1
-        factor for the likelihood for 'overcounting' data.
-        Often required when using data augmentation.
-    lr : float
-        learning rate for model optimizer
-    lr_min : float
-        minimum learning rate, defaults to lr and hence no decay
-        to have the learning rate decay from 1e-3 to 1e-6, set
-        lr=1e-3 and lr_min=1e-6.
-    optimizer : str
-        either 'Adam' or 'SGD'
-    scheduler : str
-        either 'exp' for exponential and 'cos' for cosine decay towards lr_min
-    """
-    # TODO: track deltas and sigmas
-    if lr_min is None:
-        lr_min = lr
-    if backend_kwargs is None:
-        backend_kwargs = dict()
-    device = parameters_to_vector(model.parameters()).device
-    N = len(train_loader.dataset)
-
-    last_layer = laplace in [DiagLLLaplace, KronLLLaplace, FullLLLaplace]
-    if last_layer:  # specific to last layer
-        assert prior_structure != 'layerwise', 'Not supported'
-        lap = laplace(model, likelihood, sigma_noise=1., prior_precision=1.,
-                      backend=backend, **backend_kwargs)
-        X, _ = next(iter(train_loader))
-        with torch.no_grad():
-            lap.model.find_last_layer(X.to(device))
-        last_layer_model = lap.model.last_layer
-        P = len(parameters_to_vector(last_layer_model.parameters()))
-    else:
-        H = len(list(model.parameters()))
-        P = len(parameters_to_vector(model.parameters()))
-
-    hyperparameters = list()
-    # set up prior precision
-    log_prior_prec_init = np.log(temperature * prior_prec_init)
-    if prior_structure == 'scalar':
-        log_prior_prec = log_prior_prec_init * torch.ones(1, device=device)
-    elif prior_structure == 'layerwise':
-        log_prior_prec = log_prior_prec_init * torch.ones(H, device=device)
-    elif prior_structure == 'diagonal':
-        log_prior_prec = log_prior_prec_init * torch.ones(P, device=device)
-    else:
-        raise ValueError(f'Invalid prior structure {prior_structure}')
-    log_prior_prec.requires_grad = True
-    hyperparameters.append(log_prior_prec)
-
-    # set up loss
-    if likelihood == 'classification':
-        criterion = CrossEntropyLoss(reduction='mean')
-        sigma_noise = 1.
-    elif likelihood == 'regression':
-        criterion = MSELoss(reduction='mean')
-        log_sigma_noise_init = np.log(sigma_noise_init)
-        log_sigma_noise = log_sigma_noise_init * torch.ones(1, device=device)
-        log_sigma_noise.requires_grad = True
-        hyperparameters.append(log_sigma_noise)
-
-    # set up model optimizer
-    # stochastic optimizers has exponentially decaying learning rate to min_lr
-    if optimizer == 'Adam':
-        optimizer = Adam(model.parameters(), lr=lr)
-    elif optimizer == 'SGD':
-        # fixup parameters should have 10x smaller learning rate
-        is_fixup = lambda param: param.size() == torch.Size([1])  # scalars
-        fixup_params = [p for p in model.parameters() if is_fixup(p)]
-        wrn_params = [p for p in model.parameters() if not is_fixup(p)]
-        params = [{'params': wrn_params}, {'params': fixup_params, 'lr': lr / 10.}]
-        optimizer = SGD(params, lr=lr, momentum=0.9, nesterov=True)
-    else:
-        raise ValueError(f'Invalid optimizer {optimizer}')
-
-    n_steps = n_epochs * len(train_loader)
-    if scheduler == 'exp':
-        min_lr_factor = lr_min / lr
-        gamma = np.exp(np.log(min_lr_factor) / n_steps)
-        scheduler = ExponentialLR(optimizer, gamma=gamma)
-    elif scheduler == 'cos':
-        scheduler = CosineAnnealingLR(optimizer, n_steps, eta_min=lr_min)
-    else:
-        raise ValueError(f'Invalid scheduler {scheduler}')
-
-    # set up hyperparameter optimizer
-    hyper_optimizer = Adam(hyperparameters, lr=lr_hyp)
-
-    best_marglik = np.inf
-    best_model_dict = None
-    best_precision = None
-    losses = list()
-    margliks = list()
-
-    for epoch in range(1, n_epochs + 1):
-        epoch_loss = 0
-        epoch_perf = 0
-        for X, y in train_loader:
-            X, y = X.to(device), y.to(device)
-            optimizer.zero_grad()
-            if likelihood == 'regression':
-                sigma_noise = torch.exp(log_sigma_noise).detach()
-                crit_factor = temperature / (2 * sigma_noise.square())
-            else:
-                crit_factor = temperature
-            prior_prec = torch.exp(log_prior_prec).detach()
-            if last_layer:
-                theta = parameters_to_vector(last_layer_model.parameters())
-                delta = expand_prior_precision(prior_prec, last_layer_model)
-            else:
-                theta = parameters_to_vector(model.parameters())
-                delta = expand_prior_precision(prior_prec, model)
-            f = model(X)
-            loss = criterion(f, y) + (0.5 * (delta * theta) @ theta) / N / crit_factor
-            loss.backward()
-            optimizer.step()
-            epoch_loss += loss.cpu().item() / len(train_loader)
-            if likelihood == 'regression':
-                epoch_perf += (f.detach() - y).square().sum() / N
-            else:
-                epoch_perf += torch.sum(torch.argmax(f.detach(), dim=-1) == y).item() / N
-            scheduler.step()
-        losses.append(epoch_loss)
-
-        if valid_loader is not None:
-            with torch.no_grad():
-                valid_perf = valid_performance(model, valid_loader, likelihood, device)
-            logging.info(f'MARGLIK[epoch={epoch}]: network training. Loss={losses[-1]:.3f}; '
-                         + f'Perf={epoch_perf:.3f}; Valid perf={valid_perf:.3f}; '
-                         + f'lr={scheduler.get_last_lr()[0]:.7f}')
-        else:
-            logging.info(f'MARGLIK[epoch={epoch}]: network training. Loss={losses[-1]:.3f}; '
-                         + f'Perf={epoch_perf:.3f}; lr={scheduler.get_last_lr()[0]:.7f}')
-
-        # only update hyperparameters every "Frequency" steps
-        if (epoch % marglik_frequency) != 0 or epoch < n_epochs_burnin:
-            continue
-
-        sigma_noise = 1 if likelihood == 'classification' else torch.exp(log_sigma_noise)
-        prior_prec = torch.exp(log_prior_prec)
-        lap = laplace(model, likelihood, sigma_noise=sigma_noise, prior_precision=prior_prec,
-                      temperature=temperature, backend=backend, **backend_kwargs)
-        lap.fit(train_loader)
-        for _ in range(n_hypersteps):
-            hyper_optimizer.zero_grad()
-            if likelihood == 'classification':
-                sigma_noise = None
-            elif likelihood == 'regression':
-                sigma_noise = torch.exp(log_sigma_noise)
-            prior_prec = torch.exp(log_prior_prec)
-            marglik = -lap.marginal_likelihood(prior_prec, sigma_noise)
-            marglik.backward()
-            hyper_optimizer.step()
-            margliks.append(marglik.item())
-
-        if margliks[-1] < best_marglik:
-            best_model_dict = deepcopy(model.state_dict())
-            best_precision = deepcopy(prior_prec.detach())
-            best_sigma = 1 if likelihood == 'classification' else deepcopy(sigma_noise.detach())
-            best_marglik = margliks[-1]
-            logging.info(f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={best_marglik:.2f}. '
-                         + 'Saving new best model.')
-        else:
-            logging.info(f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={margliks[-1]:.2f}.'
-                         + f'No improvement over {best_marglik:.2f}')
-
-    logging.info('MARGLIK: finished training. Recover best model and fit Lapras.')
-    if best_model_dict is not None:
-        model.load_state_dict(best_model_dict)
-        sigma_noise = best_sigma
-        prior_prec = best_precision
-    lap = laplace(model, likelihood, sigma_noise=sigma_noise, prior_precision=prior_prec,
-                  backend=backend, **backend_kwargs)
-    lap.fit(train_loader)
-    return lap, model, margliks, losses
-
-
-def expand_prior_precision(prior_prec, model):
-    theta = parameters_to_vector(model.parameters())
-    device, P = theta.device, len(theta)
-    assert prior_prec.ndim == 1
-    if len(prior_prec) == 1:  # scalar
-        return torch.ones(P, device=device) * prior_prec
-    elif len(prior_prec) == P:  # full diagonal
-        return prior_prec.to(device)
-    else:
-        return torch.cat([delta * torch.ones_like(m).flatten() for delta, m
-                          in zip(prior_prec, model.parameters())])
-
-
-def valid_performance(model, test_loader, likelihood, device):
-    N = len(test_loader.dataset)
-    perf = 0
-    for X, y in test_loader:
-        X, y = X.to(device), y.to(device)
-        if likelihood == 'classification':
-            perf += (torch.argmax(model(X), dim=-1) == y).sum() / N
-        else:
-            perf += (model(X) - y).square().sum() / N
-    return perf.item()
diff --git a/laplace/matrix.py b/laplace/matrix.py
index 1d2218c7..61c07ab5 100644
--- a/laplace/matrix.py
+++ b/laplace/matrix.py
@@ -7,14 +7,34 @@
 
 
 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):
-        # FIXME: this does not work for BatchNorm!!
-        # TODO: rewrite this functionality in terms of module types
+        """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
@@ -35,14 +55,36 @@ def init_from_model(cls, model, device):
         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)]
+                 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.')
 
@@ -54,6 +96,16 @@ 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()
@@ -66,6 +118,18 @@ def decompose(self, damping=False):
         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()
@@ -92,8 +156,24 @@ def _bmm(self, W: torch.Tensor) -> torch.Tensor:
         return SW
 
     def bmm(self, W: torch.Tensor, exponent: float = 1) -> torch.Tensor:
-        # self @ W with W[params], W[batch, params], W[batch, classes, params]
-        # returns SW[batch, classes, params]
+        """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:
@@ -106,6 +186,13 @@ def bmm(self, W: torch.Tensor, exponent: float = 1) -> torch.Tensor:
             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:
@@ -117,6 +204,12 @@ def logdet(self) -> torch.Tensor:
         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:
@@ -126,6 +219,14 @@ def diag(self) -> torch.Tensor:
         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:
@@ -134,14 +235,32 @@ def to_matrix(self) -> torch.Tensor:
                 blocks.append(kron(F[0], F[1]))
         return block_diag(blocks)
 
-    # inplace and permuted operations
+    # for commutative operations
     __radd__ = __add__
-    __iadd__ = __add__
     __rmul__ = __mul__
-    __imul__ = __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
@@ -170,21 +289,51 @@ def _check_deltas(self, deltas: torch.Tensor):
             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(l)) * l for l in ls] for ls in self.eigenvalues]
+        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 \sum_l log det (kron_l + delta I_l)
+        """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
@@ -201,6 +350,20 @@ def logdet(self) -> torch.Tensor:
         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()
@@ -209,14 +372,12 @@ def _bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
         SW = list()
         for ls, Qs, delta in zip(self.eigenvalues, self.eigenvectors, self.deltas):
             if len(ls) == 1:
-                # just Q (Lambda + delta) Q^T W_p
                 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:
-                # not so easy to explain...
                 Q1, Q2 = Qs
                 l1, l2 = ls
                 p = len(l1) * len(l2)
@@ -242,8 +403,23 @@ def inv_square_form(self, W: torch.Tensor) -> torch.Tensor:
         return torch.bmm(W, SW.transpose(1, 2))
 
     def bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
-        # self @ W with W[params], W[batch, params], W[batch, classes, params]
-        # returns SW[batch, classes, params]
+        """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:
@@ -254,6 +430,14 @@ def bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
             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:
@@ -272,8 +456,6 @@ def to_matrix(self, exponent: float = 1) -> torch.Tensor:
                 blocks.append(Q @ L @ Q.T)
         return block_diag(blocks)
 
-    # FIXME: iadd imul should change mutable types in principle.
+    # for commutative operations
     __radd__ = __add__
-    __iadd__ = __add__
     __rmul__ = __mul__
-    __imul__ = __mul__
diff --git a/laplace/utils.py b/laplace/utils.py
index dc5692a6..e096cd18 100644
--- a/laplace/utils.py
+++ b/laplace/utils.py
@@ -25,14 +25,34 @@ def validate(laplace, val_loader, pred_type='glm', link_approx='probit', n_sampl
 
 
 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):
+def _is_batchnorm(module):
     if isinstance(module, BatchNorm1d) or \
         isinstance(module, BatchNorm2d) or \
             isinstance(module, BatchNorm3d):
@@ -51,8 +71,16 @@ def _is_valid_scalar(scalar: Union[float, int, torch.Tensor]) -> bool:
 
 
 def kron(t1, t2):
-    """
-    Computes the Kronecker product between two tensors.
+    """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()
@@ -71,6 +99,17 @@ def kron(t1, 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:
@@ -81,13 +120,19 @@ def diagonal_add_scalar(X, value):
 
 def symeig(M):
     """Symetric eigendecomposition avoiding failure cases by
-    adding and removing jitter to the diagonal
-
-    returns eigenvalues (l) and eigenvectors (W)
+    adding and removing jitter to the diagonal.
+
+    Parameters
+    ----------
+    M : torch.Tensor
+
+    Returns
+    -------
+    L : torch.Tensor
+        eigenvalues
+    W : torch.Tensor
+        eigenvectors
     """
-    # could make double to get more precise computation
-    # M = M.double()
-    # and then below return L.float(), W.float()
     try:
         L, W = torch.symeig(M, eigenvectors=True)
     except RuntimeError:  # did not converge
@@ -110,6 +155,16 @@ def symeig(M):
 
 
 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
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
index 99baa1b9..38b22452 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,4 @@ torch
 torchvision
 torchaudio
 backpack-for-pytorch
+asdfghjkl @ git+https://github.com/kazukiosawa/asdfghjkl.git
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/test_baselaplace.py b/tests/test_baselaplace.py
index 08b3045a..8bff47fd 100644
--- a/tests/test_baselaplace.py
+++ b/tests/test_baselaplace.py
@@ -121,20 +121,20 @@ def test_laplace_init_precision(laplace, model):
 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', 
+    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', 
+    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', 
+    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', 
+    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 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)
@@ -184,7 +184,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     assert f.shape == torch.Size([10, 2])
 
     # Test log likelihood (Train)
-    log_lik = lap.log_lik
+    log_lik = lap.log_likelihood
     # compute true log lik
     if lh == 'classification':
         log_lik_true = Categorical(logits=f).log_prob(y).sum()
@@ -195,7 +195,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
         assert torch.allclose(log_lik, log_lik_true)
         # change likelihood and test again
         lap.sigma_noise = 0.72
-        log_lik = lap.log_lik
+        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)
 
@@ -214,7 +214,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     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.marginal_likelihood())
+    assert torch.allclose(lml, lap.log_marginal_likelihood())
 
     # test sampling
     torch.manual_seed(61)
diff --git a/tests/test_curv_backends_kazuki.py b/tests/test_curv_backends_asdl.py
similarity index 87%
rename from tests/test_curv_backends_kazuki.py
rename to tests/test_curv_backends_asdl.py
index b8a2ddd3..696c8969 100644
--- a/tests/test_curv_backends_kazuki.py
+++ b/tests/test_curv_backends_asdl.py
@@ -5,7 +5,7 @@
 
 from asdfghjkl.operations import Bias, Scale
 
-from laplace.curvature import KazukiGGN, KazukiEF, BackPackGGN, BackPackEF
+from laplace.curvature import AsdlGGN, AsdlEF, BackPackGGN, BackPackEF
 
 
 @pytest.fixture
@@ -30,7 +30,7 @@ def class_Xy():
 @pytest.fixture
 def complex_model():
     torch.manual_seed(711)
-    model = torch.nn.Sequential(nn.Conv2d(3, 4, 2, 2), nn.Flatten(), nn.Tanh(), 
+    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())
@@ -49,12 +49,12 @@ def complex_class_Xy():
 
 def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
     X, y = class_Xy
-    backend = KazukiGGN(model, 'classification', stochastic=False)
+    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
 
@@ -67,12 +67,12 @@ def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
 
 def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
     X, y = class_Xy
-    backend = KazukiEF(model, 'classification')
+    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
 
@@ -85,19 +85,19 @@ def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
 
 def test_diag_ggn_stoch_cls_kazuki(class_Xy, model):
     X, y = class_Xy
-    backend = KazukiGGN(model, 'classification', stochastic=True)
+    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 = KazukiGGN(model, 'classification', stochastic=False)
+    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', [KazukiEF, KazukiGGN])
+@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
@@ -109,7 +109,7 @@ def test_kron_kazuki_vs_diag_class(class_Xy, model, Backend):
     assert torch.allclose(kron.diag(), dggn)
 
 
-@pytest.mark.parametrize('Backend', [KazukiEF, KazukiGGN])
+@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
 def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
     X, y = class_Xy
     backend = Backend(model, 'classification')
@@ -126,7 +126,7 @@ def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [KazukiGGN, KazukiEF])
+@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
@@ -138,7 +138,7 @@ def test_kron_summing_up_vs_diag_kazuki(class_Xy, model, Backend):
 
 def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
     X, y = complex_class_Xy
-    backend = KazukiGGN(complex_model, 'classification', stochastic=True)
+    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
@@ -149,7 +149,7 @@ def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1)
 
 
-@pytest.mark.parametrize('Backend', [KazukiEF, KazukiGGN])
+@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
@@ -161,7 +161,7 @@ def test_complex_kron_kazuki_vs_diag_kazuki(complex_class_Xy, complex_model, Bac
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
 
 
-@pytest.mark.parametrize('Backend', [KazukiEF, KazukiGGN])
+@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')
@@ -178,7 +178,7 @@ def test_complex_kron_batching_correction_kazuki(complex_class_Xy, complex_model
     assert torch.allclose(loss, loss_two)
 
 
-@pytest.mark.parametrize('Backend', [KazukiGGN, KazukiEF])
+@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
@@ -191,7 +191,7 @@ def test_complex_kron_summing_up_vs_diag_class_kazuki(complex_class_Xy, complex_
 def test_kron_normalization_ggn_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = KazukiGGN(model, 'classification', stochastic=False)
+    backend = AsdlGGN(model, 'classification', stochastic=False)
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
@@ -205,7 +205,7 @@ def test_kron_normalization_ggn_class(class_Xy, model):
 def test_kron_normalization_ef_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
-    backend = KazukiEF(model, 'classification')
+    backend = AsdlEF(model, 'classification')
     loss, kron = backend.kron(xi, yi, N=1)
     kron_true = 7 * kron
     loss_true = 7 * loss
@@ -213,4 +213,4 @@ def test_kron_normalization_ef_class(class_Xy, model):
     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)
\ No newline at end of file
+    assert torch.allclose(loss_true, loss_test)
diff --git a/tests/test_jacobians.py b/tests/test_jacobians.py
index 6aac9e11..7a5a22ef 100644
--- a/tests/test_jacobians.py
+++ b/tests/test_jacobians.py
@@ -3,8 +3,7 @@
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
-from laplace.curvature import KazukiInterface
-from laplace.curvature import BackPackInterface
+from laplace.curvature import AsdlInterface, BackPackInterface
 from laplace.feature_extractor import FeatureExtractor
 from tests.utils import jacobians_naive
 
@@ -36,7 +35,7 @@ def X():
     return torch.randn(200, 3)
 
 
-@pytest.mark.parametrize('backend', [KazukiInterface, BackPackInterface])
+@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)
@@ -47,7 +46,7 @@ def test_linear_jacobians(linear_model, X, backend):
     assert torch.allclose(f, linear_model(X), atol=1e-5)
 
 
-@pytest.mark.parametrize('backend', [KazukiInterface, BackPackInterface])
+@pytest.mark.parametrize('backend', [AsdlInterface, BackPackInterface])
 def test_jacobians_singleoutput(singleoutput_model, X, backend):
     model = singleoutput_model
     Js, f = backend.jacobians(model, X)
@@ -58,7 +57,7 @@ def test_jacobians_singleoutput(singleoutput_model, X, backend):
     assert torch.allclose(f, f_naive)
 
 
-@pytest.mark.parametrize('backend', [KazukiInterface, BackPackInterface])
+@pytest.mark.parametrize('backend', [AsdlInterface, BackPackInterface])
 def test_jacobians_multioutput(multioutput_model, X, backend):
     model = multioutput_model
     Js, f = backend.jacobians(model, X)
@@ -69,7 +68,7 @@ def test_jacobians_multioutput(multioutput_model, X, backend):
     assert torch.allclose(f, f_naive)
 
 
-@pytest.mark.parametrize('backend', [KazukiInterface, BackPackInterface])
+@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)
@@ -81,7 +80,7 @@ def test_last_layer_jacobians_singleoutput(singleoutput_model, X, backend):
     assert torch.allclose(f, f_naive)
 
 
-@pytest.mark.parametrize('backend', [KazukiInterface, BackPackInterface])
+@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)
diff --git a/tests/test_lllaplace.py b/tests/test_lllaplace.py
index f59e77fd..102f429b 100644
--- a/tests/test_lllaplace.py
+++ b/tests/test_lllaplace.py
@@ -202,7 +202,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     assert f.shape == torch.Size([10, 2])
 
     # Test log likelihood (Train)
-    log_lik = lap.log_lik
+    log_lik = lap.log_likelihood
     # compute true log lik
     if lh == 'classification':
         log_lik_true = Categorical(logits=f).log_prob(y).sum()
@@ -213,7 +213,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
         assert torch.allclose(log_lik, log_lik_true)
         # change likelihood and test again
         lap.sigma_noise = 0.72
-        log_lik = lap.log_lik
+        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)
 
@@ -233,7 +233,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     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.marginal_likelihood())
+    assert torch.allclose(lml, lap.log_marginal_likelihood())
 
     # test sampling
     torch.manual_seed(61)