Merge pull request #208 from aleximmer/v0.2

v0.2

README.md

@@ -60,11 +60,11 @@ pip install laplace-torch
 For development purposes, clone the repository and then install:
 
 ```bash
-# or after cloning the repository for development
-pip install -e .
-# run tests
-pip install -e .[tests]
-pytest tests/
+# first install the build system:
+pip install --upgrade pip wheel packaging
+
+# then install the develop 
+pip install -e ".[all]"
 ```
 
 ## Example usage
@@ -392,7 +392,7 @@ The documentation is available [here](https://aleximmer.github.io/Laplace) or ca
 
 ```bash
 # assuming the repository was cloned
-pip install -e .[docs]
+pip install -e ".[docs]"
 # create docs and write to html
 bash update_docs.sh
 # .. or serve the docs directly

docs/curvature/asdfghjkl.html

@@ -35,7 +35,7 @@ <h2 class="section-title" id="header-classes">Classes</h2>
 <dl>
 <dt id="laplace.curvature.asdfghjkl.AsdfghjklInterface"><code class="flex name class">
 <span>class <span class="ident">AsdfghjklInterface</span></span>
-<span>(</span><span>model, likelihood, last_layer=False, subnetwork_indices=None, dict_key_x='input_ids', dict_key_y='labels')</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Interface for asdfghjkl backend.</p></div>
@@ -52,7 +52,7 @@ <h3>Subclasses</h3>
 <h3>Methods</h3>
 <dl>
 <dt id="laplace.curvature.asdfghjkl.AsdfghjklInterface.jacobians"><code class="name flex">
-<span>def <span class="ident">jacobians</span></span>(<span>self, x, enable_backprop=False)</span>
+<span>def <span class="ident">jacobians</span></span>(<span>self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], enable_backprop: bool = False) ‑> tuple[torch.Tensor, torch.Tensor]</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>
@@ -73,7 +73,7 @@ <h2 id="returns">Returns</h2>
 </dl></div>
 </dd>
 <dt id="laplace.curvature.asdfghjkl.AsdfghjklInterface.gradients"><code class="name flex">
-<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+<span>def <span class="ident">gradients</span></span>(<span>self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]</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
@@ -109,7 +109,7 @@ <h3>Inherited members</h3>
 </dd>
 <dt id="laplace.curvature.asdfghjkl.AsdfghjklHessian"><code class="flex name class">
 <span>class <span class="ident">AsdfghjklHessian</span></span>
-<span>(</span><span>model, likelihood, last_layer=False, dict_key_x='input_ids', dict_key_y='labels', low_rank=10)</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels', low_rank: int = 10)</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Interface for asdfghjkl backend.</p></div>
@@ -121,7 +121,7 @@ <h3>Ancestors</h3>
 <h3>Methods</h3>
 <dl>
 <dt id="laplace.curvature.asdfghjkl.AsdfghjklHessian.eig_lowrank"><code class="name flex">
-<span>def <span class="ident">eig_lowrank</span></span>(<span>self, data_loader)</span>
+<span>def <span class="ident">eig_lowrank</span></span>(<span>self, data_loader: DataLoader) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor]</span>
 </code></dt>
 <dd>
 <div class="desc"></div>
@@ -144,7 +144,7 @@ <h3>Inherited members</h3>
 </dd>
 <dt id="laplace.curvature.asdfghjkl.AsdfghjklGGN"><code class="flex name class">
 <span>class <span class="ident">AsdfghjklGGN</span></span>
-<span>(</span><span>model, likelihood, last_layer=False, subnetwork_indices=None, dict_key_x='input_ids', dict_key_y='labels', stochastic=False)</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels', stochastic: bool = False)</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Implementation of the <code>GGNInterface</code> using asdfghjkl.</p></div>
@@ -171,7 +171,7 @@ <h3>Inherited members</h3>
 </dd>
 <dt id="laplace.curvature.asdfghjkl.AsdfghjklEF"><code class="flex name class">
 <span>class <span class="ident">AsdfghjklEF</span></span>
-<span>(</span><span>model, likelihood, last_layer=False, dict_key_x='input_ids', dict_key_y='labels')</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | None, last_layer: bool = False, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Implementation of the <code>EFInterface</code> using asdfghjkl.</p></div>

docs/curvature/asdl.html

@@ -35,7 +35,7 @@ <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, subnetwork_indices=None, dict_key_x='input_ids', dict_key_y='labels')</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Interface for asdfghjkl backend.</p></div>
@@ -51,15 +51,15 @@ <h3>Subclasses</h3>
 </ul>
 <h3>Instance variables</h3>
 <dl>
-<dt id="laplace.curvature.asdl.AsdlInterface.loss_type"><code class="name">var <span class="ident">loss_type</span></code></dt>
+<dt id="laplace.curvature.asdl.AsdlInterface.loss_type"><code class="name">var <span class="ident">loss_type</span> : str</code></dt>
 <dd>
 <div class="desc"></div>
 </dd>
 </dl>
 <h3>Methods</h3>
 <dl>
 <dt id="laplace.curvature.asdl.AsdlInterface.jacobians"><code class="name flex">
-<span>def <span class="ident">jacobians</span></span>(<span>self, x, enable_backprop=False)</span>
+<span>def <span class="ident">jacobians</span></span>(<span>self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], enable_backprop: bool = False) ‑> tuple[torch.Tensor, torch.Tensor]</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>
@@ -82,7 +82,7 @@ <h2 id="returns">Returns</h2>
 </dl></div>
 </dd>
 <dt id="laplace.curvature.asdl.AsdlInterface.gradients"><code class="name flex">
-<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+<span>def <span class="ident">gradients</span></span>(<span>self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]</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
@@ -118,7 +118,7 @@ <h3>Inherited members</h3>
 </dd>
 <dt id="laplace.curvature.asdl.AsdlHessian"><code class="flex name class">
 <span>class <span class="ident">AsdlHessian</span></span>
-<span>(</span><span>model, likelihood, last_layer=False, dict_key_x='input_ids', dict_key_y='labels')</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Interface for asdfghjkl backend.</p></div>
@@ -144,7 +144,7 @@ <h3>Inherited members</h3>
 </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, subnetwork_indices=None, dict_key_x='input_ids', dict_key_y='labels', stochastic=False)</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels', stochastic: bool = False)</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Implementation of the <code>GGNInterface</code> using asdfghjkl.</p></div>
@@ -171,7 +171,7 @@ <h3>Inherited members</h3>
 </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, dict_key_x='input_ids', dict_key_y='labels')</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Implementation of the <code>EFInterface</code> using asdfghjkl.</p></div>

docs/curvature/backpack.html

@@ -35,7 +35,7 @@ <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, subnetwork_indices=None, dict_key_x='input_ids', dict_key_y='labels')</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Interface for Backpack backend.</p></div>
@@ -51,7 +51,7 @@ <h3>Subclasses</h3>
 <h3>Methods</h3>
 <dl>
 <dt id="laplace.curvature.backpack.BackPackInterface.jacobians"><code class="name flex">
-<span>def <span class="ident">jacobians</span></span>(<span>self, x, enable_backprop=False)</span>
+<span>def <span class="ident">jacobians</span></span>(<span>self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], enable_backprop: bool = False) ‑> tuple[torch.Tensor, torch.Tensor]</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>
@@ -73,7 +73,7 @@ <h2 id="returns">Returns</h2>
 </dl></div>
 </dd>
 <dt id="laplace.curvature.backpack.BackPackInterface.gradients"><code class="name flex">
-<span>def <span class="ident">gradients</span></span>(<span>self, x, y)</span>
+<span>def <span class="ident">gradients</span></span>(<span>self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]</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
@@ -110,7 +110,7 @@ <h3>Inherited members</h3>
 </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, subnetwork_indices=None, dict_key_x='input_ids', dict_key_y='labels', stochastic=False)</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels', stochastic: bool = False)</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Implementation of the <code>GGNInterface</code> using Backpack.</p></div>
@@ -137,7 +137,7 @@ <h3>Inherited members</h3>
 </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, subnetwork_indices=None, dict_key_x='input_ids', dict_key_y='labels')</span>
+<span>(</span><span>model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')</span>
 </code></dt>
 <dd>
 <div class="desc"><p>Implementation of <code>EFInterface</code> using Backpack.</p></div>