Repository: aleximmer/Laplace (language: python)

Issue:
Add support for type hints
A straight forward way to improve the code quality is to enable runtime support for type hints via `typing`.
Bringing laplace-torch to foundation-model era
Main features of this pull request:

1. Support only doing Laplace on params that require grad. Use case: PEFT (like LoRA) on top of a frozen foundation model. This is more efficient than SubnetLaplace since the latter still computes the full Jacobians.
2. Add support to multiple leading dims in `classification` likelihood. E.g. the logits is `(batch_size, seq_len, n_classes)`. Useful for language modeling and reward modeling.
    1. This PR also contains the `integrate-latest-asdl` changes. I tested it with my ASDL fork (only a couple of light changes to support weight-sharing dim and `ignore_index`; so crucial in language modeling): <https://github.com/wiseodd/asdl/commits/dev/>. Please also check this and let me know what's the most elegant way. 
4. Support Huggingface dataset. The assumption is that `x` is a `UserDict` containing `input_ids`, `attention_mask`, etc., things that are produced by HF dataloader.
5. Add a new likelihood called `reward_modeling` where the `classification` likelihood is used during training and the `regression` likelihood is used during prediction.
6. Add support to `torchmetrics` for gridsearch. The benefit is that it supports running metrics => less memory overhead (vis-a-vis gathering all the logits first).
7. Add Jacobian computation with `torch.func` (`functorch`, really) as a general Jacobian computation for GLM predictive. Useful for Bayesian optimization/invariance learning where you need to backprop through the variance. Much more elegant than to change ASDL. 

Relevant unit tests are provided. All tests passed; the only ones failed are the old `LowRankLaplace` issues.


Discussion / hints:

Thanks Agustinus for this PR, this is very useful! Just took a glance at the parts that I'm familiar with, which looks good to me. I'll defer to Alex and Runa for a more in-depth review due to their better familiarity with the core library.
This PR is currently blocked by https://github.com/f-dangel/curvlinops/issues/86. Need to make `curvlinops` supports Huggingface datasets first.
Regarding the addition of `logit_class_idx` argument: I did a "smoke test" in implementing this, and realized that the changes would be substantial---even without the present PR, `laplace-torch` always assumes that `logit_class_idx == -1`. 

So, I think this is better done in a separate PR. This will also make the scope of the PRs cleaner---in this PR, we focus on HF LLM models.
So, the remaining TODOs for this PR are:

1. Decide the conversations about `matrix.py`. @runame & @aleximmer could you please handle this?
2. Wait until https://github.com/f-dangel/curvlinops/pull/100 is merged

Then we can merge this PR and continue with further PRs that we deferred in this PR.
Sounds good, can you create issues for all comments that you want to address in future PRs?
I replaced `get_nll` in crossval with `RunningNLLMetrics()`. So this PR will close #160 .
(I thought before that I `RunningNLLMetics` hadn't been implemented, so would have involve major work.)
Might as well fixes #156 while we're at it.
All tasks finished!

Fix patch:
diff --git a/examples/bayesopt_example.py b/examples/bayesopt_example.py
index ff3fd3f..bb90743 100644
--- a/examples/bayesopt_example.py
+++ b/examples/bayesopt_example.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import warnings
 
 warnings.filterwarnings("ignore")
@@ -22,7 +24,7 @@
 from torch import nn, optim
 from torch.nn import functional as F
 
-from laplace import Laplace
+from laplace import BaseLaplace, Laplace
 
 
 class LaplaceBNN(Model):
@@ -35,7 +37,7 @@ def __init__(
         self,
         train_X: torch.Tensor,
         train_Y: torch.Tensor,
-        bnn: Laplace = None,
+        bnn: BaseLaplace | None = None,
         likelihood: str = "regression",
         batch_size: int = 1024,
     ):
diff --git a/laplace/__init__.py b/laplace/__init__.py
index 5446db2..7f8580e 100644
--- a/laplace/__init__.py
+++ b/laplace/__init__.py
@@ -6,9 +6,6 @@
 .. include:: ../examples/reward_modeling_example.md
 """
 
-REGRESSION = "regression"
-CLASSIFICATION = "classification"
-
 from laplace.baselaplace import (
     BaseLaplace,
     DiagLaplace,
@@ -21,6 +18,15 @@
 from laplace.lllaplace import DiagLLLaplace, FullLLLaplace, KronLLLaplace, LLLaplace
 from laplace.marglik_training import marglik_training
 from laplace.subnetlaplace import DiagSubnetLaplace, FullSubnetLaplace, SubnetLaplace
+from laplace.utils.enums import (
+    HessianStructure,
+    Likelihood,
+    LinkApprox,
+    PredType,
+    PriorStructure,
+    SubsetOfWeights,
+    TuningMethod,
+)
 
 __all__ = [
     "Laplace",  # direct access to all Laplace classes via unified interface
@@ -38,4 +44,12 @@
     "FullSubnetLaplace",
     "DiagSubnetLaplace",  # subnetwork
     "marglik_training",
-]  # methods
+    # Enums
+    "SubsetOfWeights",
+    "HessianStructure",
+    "Likelihood",
+    "PredType",
+    "LinkApprox",
+    "TuningMethod",
+    "PriorStructure",
+]
diff --git a/laplace/baselaplace.py b/laplace/baselaplace.py
index f8a848d..1127934 100644
--- a/laplace/baselaplace.py
+++ b/laplace/baselaplace.py
@@ -1,19 +1,31 @@
+from __future__ import annotations
+
 import warnings
 from collections.abc import MutableMapping
 from math import log, pi, sqrt
+from typing import Any, Callable
 
 import numpy as np
 import torch
-import torchmetrics as tm
+import torchmetrics
 import tqdm
+from torch import nn
 from torch.nn.utils import parameters_to_vector, vector_to_parameters
+from torch.utils.data import DataLoader
 
-from laplace.curvature import CurvlinopsGGN
 from laplace.curvature.asdfghjkl import AsdfghjklHessian
-from laplace.curvature.curvlinops import CurvlinopsEF
-from laplace.utils import (
-    Kron,
-    RunningNLLMetric,
+from laplace.curvature.curvature import CurvatureInterface
+from laplace.curvature.curvlinops import CurvlinopsEF, CurvlinopsGGN
+from laplace.utils.enums import (
+    Likelihood,
+    LinkApprox,
+    PredType,
+    PriorStructure,
+    TuningMethod,
+)
+from laplace.utils.matrix import Kron, KronDecomposed
+from laplace.utils.metrics import RunningNLLMetric
+from laplace.utils.utils import (
     fix_prior_prec_structure,
     invsqrt_precision,
     normal_samples,
@@ -36,7 +48,7 @@ class BaseLaplace:
     Parameters
     ----------
     model : torch.nn.Module
-    likelihood : {'classification', 'regression', 'reward_modeling'}
+    likelihood : Likelihood or str in {'classification', 'regression', 'reward_modeling'}
         determines the log likelihood Hessian approximation.
         In the case of 'reward_modeling', it fits Laplace using the classification likelihood,
         then does prediction as in regression likelihood. The model needs to be defined accordingly:
@@ -75,50 +87,51 @@ class BaseLaplace:
 
     def __init__(
         self,
-        model,
-        likelihood,
-        sigma_noise=1.0,
-        prior_precision=1.0,
-        prior_mean=0.0,
-        temperature=1.0,
-        enable_backprop=False,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        backend=None,
-        backend_kwargs=None,
-        asdl_fisher_kwargs=None,
-    ):
-        if likelihood not in ["classification", "regression", "reward_modeling"]:
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        sigma_noise: float | torch.Tensor = 1.0,
+        prior_precision: float | torch.Tensor = 1.0,
+        prior_mean: float | torch.Tensor = 0.0,
+        temperature: float = 1.0,
+        enable_backprop: bool = False,
+        dict_key_x: str = "input_ids",
+        dict_key_y: str = "labels",
+        backend: type[CurvatureInterface] | None = None,
+        backend_kwargs: dict[str, Any] | None = None,
+        asdl_fisher_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        if likelihood not in [lik.value for lik in Likelihood]:
             raise ValueError(f"Invalid likelihood type {likelihood}")
 
-        self.model = model
+        self.model: nn.Module = model
 
         # Only do Laplace on params that require grad
-        self.params = []
-        self.is_subset_params = False
+        self.params: list[torch.Tensor] = []
+        self.is_subset_params: bool = False
         for p in model.parameters():
             if p.requires_grad:
                 self.params.append(p)
             else:
                 self.is_subset_params = True
 
-        self.n_params = sum(p.numel() for p in self.params)
-        self.n_layers = len(self.params)
-        self.prior_precision = prior_precision
-        self.prior_mean = prior_mean
-        if sigma_noise != 1 and likelihood != "regression":
+        self.n_params: int = sum(p.numel() for p in self.params)
+        self.n_layers: int = len(self.params)
+        self.prior_precision: float | torch.Tensor = prior_precision
+        self.prior_mean: float | torch.Tensor = prior_mean
+        if sigma_noise != 1 and likelihood != Likelihood.REGRESSION:
             raise ValueError("Sigma noise != 1 only available for regression.")
 
-        self.reward_modeling = likelihood == "reward_modeling"
+        self.reward_modeling: bool = likelihood == Likelihood.REWARD_MODELING
+
         if self.reward_modeling:
             # For fitting only. After it's done, self.likelihood = 'regression', see self.fit()
-            self.likelihood = "classification"
+            self.likelihood = Likelihood.CLASSIFICATION
         else:
             self.likelihood = likelihood
 
-        self.sigma_noise = sigma_noise
-        self.temperature = temperature
-        self.enable_backprop = enable_backprop
+        self.sigma_noise: float | torch.Tensor = sigma_noise
+        self.temperature: float = temperature
+        self.enable_backprop: bool = enable_backprop
 
         # For models with dict-like inputs (e.g. Huggingface LLMs)
         self.dict_key_x = dict_key_x
@@ -136,24 +149,32 @@ def __init__(
                     " are not supported."
                 )
 
-        self._backend = None
-        self._backend_cls = backend
-        self._backend_kwargs = dict() if backend_kwargs is None else backend_kwargs
-        self._asdl_fisher_kwargs = (
+        self._backend: CurvatureInterface | None = None
+        self._backend_cls: type[CurvatureInterface] = backend
+        self._backend_kwargs: dict[str, Any] = (
+            dict() if backend_kwargs is None else backend_kwargs
+        )
+        self._asdl_fisher_kwargs: dict[str, Any] = (
             dict() if asdl_fisher_kwargs is None else asdl_fisher_kwargs
         )
 
         # log likelihood = g(loss)
-        self.loss = 0.0
-        self.n_outputs = None
-        self.n_data = 0
+        self.loss: float = 0.0
+        self.n_outputs: int = 0
+        self.n_data: int = 0
+
+        # Declare attributes
+        self._prior_mean: torch.Tensor
+        self._prior_precision: torch.Tensor
+        self._sigma_noise: torch.Tensor
+        self._posterior_scale: torch.Tensor | None
 
     @property
-    def _device(self):
+    def _device(self) -> torch.device:
         return next(self.model.parameters()).device
 
     @property
-    def backend(self):
+    def backend(self) -> CurvatureInterface:
         if self._backend is None:
             self._backend = self._backend_cls(
                 self.model,
@@ -164,17 +185,26 @@ def backend(self):
             )
         return self._backend
 
-    def _curv_closure(self, X, y, N):
+    def _curv_closure(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         raise NotImplementedError
 
-    def fit(self, train_loader):
+    def fit(self, train_loader: DataLoader) -> None:
         raise NotImplementedError
 
-    def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
+    def log_marginal_likelihood(
+        self,
+        prior_precision: torch.Tensor | None = None,
+        sigma_noise: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         raise NotImplementedError
 
     @property
-    def log_likelihood(self):
+    def log_likelihood(self) -> torch.Tensor:
         """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
@@ -190,20 +220,32 @@ def log_likelihood(self):
             c = (
                 self.n_data
                 * self.n_outputs
-                * torch.log(self.sigma_noise * sqrt(2 * pi))
+                * torch.log(torch.tensor(self.sigma_noise) * sqrt(2 * pi))
             )
             return factor * self.loss - c
         else:
             # for classification Xent == log Cat
             return factor * self.loss
 
-    def __call__(self, x, pred_type, link_approx, n_samples):
+    def __call__(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        pred_type: PredType | str,
+        link_approx: LinkApprox | str,
+        n_samples: int,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         raise NotImplementedError
 
-    def predictive(self, x, pred_type, link_approx, n_samples):
+    def predictive(
+        self,
+        x: torch.Tensor,
+        pred_type: PredType | str,
+        link_approx: LinkApprox | str,
+        n_samples: int,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         return self(x, pred_type, link_approx, n_samples)
 
-    def _check_jacobians(self, Js):
+    def _check_jacobians(self, Js: torch.Tensor) -> None:
         if not isinstance(Js, torch.Tensor):
             raise ValueError("Jacobians have to be torch.Tensor.")
         if not Js.device == self._device:
@@ -213,7 +255,7 @@ def _check_jacobians(self, Js):
             raise ValueError("Invalid Jacobians shape for Laplace posterior approx.")
 
     @property
-    def prior_precision_diag(self):
+    def prior_precision_diag(self) -> torch.Tensor:
         """Obtain the diagonal prior precision \\(p_0\\) constructed from either
         a scalar, layer-wise, or diagonal prior precision.
 
@@ -221,35 +263,38 @@ def prior_precision_diag(self):
         -------
         prior_precision_diag : torch.Tensor
         """
-        if len(self.prior_precision) == 1:  # scalar
-            return self.prior_precision * torch.ones(self.n_params, device=self._device)
-
-        elif len(self.prior_precision) == self.n_params:  # diagonal
-            return self.prior_precision
+        prior_prec: torch.Tensor = (
+            self.prior_precision
+            if isinstance(self.prior_precision, torch.Tensor)
+            else torch.tensor(self.prior_precision)
+        )
 
-        elif len(self.prior_precision) == self.n_layers:  # per layer
+        if prior_prec.ndim == 0 or len(prior_prec) == 1:  # scalar
+            return self.prior_precision * torch.ones(self.n_params, device=self._device)
+        elif len(prior_prec) == self.n_params:  # diagonal
+            return prior_prec
+        elif len(prior_prec) == self.n_layers:  # per layer
             n_params_per_layer = [p.numel() for p in self.params]
             return torch.cat(
                 [
                     prior * torch.ones(n_params, device=self._device)
-                    for prior, n_params in zip(self.prior_precision, n_params_per_layer)
+                    for prior, n_params in zip(prior_prec, n_params_per_layer)
                 ]
             )
-
         else:
             raise ValueError(
                 "Mismatch of prior and model. Diagonal, scalar, or per-layer prior."
             )
 
     @property
-    def prior_mean(self):
+    def prior_mean(self) -> torch.Tensor:
         return self._prior_mean
 
     @prior_mean.setter
-    def prior_mean(self, prior_mean):
+    def prior_mean(self, prior_mean: float | torch.Tensor) -> None:
         if np.isscalar(prior_mean) and np.isreal(prior_mean):
             self._prior_mean = torch.tensor(prior_mean, device=self._device)
-        elif torch.is_tensor(prior_mean):
+        elif isinstance(prior_mean, torch.Tensor):
             if prior_mean.ndim == 0:
                 self._prior_mean = prior_mean.reshape(-1).to(self._device)
             elif prior_mean.ndim == 1:
@@ -262,15 +307,16 @@ def prior_mean(self, prior_mean):
             raise ValueError("Invalid argument type of prior mean.")
 
     @property
-    def prior_precision(self):
+    def prior_precision(self) -> torch.Tensor:
         return self._prior_precision
 
     @prior_precision.setter
-    def prior_precision(self, prior_precision):
+    def prior_precision(self, prior_precision: float | torch.Tensor):
         self._posterior_scale = None
+
         if np.isscalar(prior_precision) and np.isreal(prior_precision):
             self._prior_precision = torch.tensor([prior_precision], device=self._device)
-        elif torch.is_tensor(prior_precision):
+        elif isinstance(prior_precision, torch.Tensor):
             if prior_precision.ndim == 0:
                 # make dimensional
                 self._prior_precision = prior_precision.reshape(-1).to(self._device)
@@ -289,35 +335,37 @@ def prior_precision(self, prior_precision):
                 "Prior precision either scalar or torch.Tensor up to 1-dim."
             )
 
-    def optimize_prior_precision_base(
+    def optimize_prior_precision(
         self,
-        pred_type,
-        method="marglik",
-        n_steps=100,
-        lr=1e-1,
-        init_prior_prec=1.0,
-        prior_structure="scalar",
-        val_loader=None,
-        loss=None,
-        log_prior_prec_min=-4,
-        log_prior_prec_max=4,
-        grid_size=100,
-        link_approx="probit",
-        n_samples=100,
-        verbose=False,
-        cv_loss_with_var=False,
-        progress_bar=False,
-    ):
+        pred_type: PredType | str,
+        method: TuningMethod | str = TuningMethod.MARGLIK,
+        n_steps: int = 100,
+        lr: float = 1e-1,
+        init_prior_prec: float | torch.Tensor = 1.0,
+        prior_structure: PriorStructure | str = PriorStructure.SCALAR,
+        val_loader: DataLoader | None = None,
+        loss: torchmetrics.Metric
+        | Callable[[torch.Tensor], torch.Tensor | float]
+        | None = None,
+        log_prior_prec_min: float = -4,
+        log_prior_prec_max: float = 4,
+        grid_size: int = 100,
+        link_approx: LinkApprox | str = LinkApprox.PROBIT,
+        n_samples: int = 100,
+        verbose: bool = False,
+        cv_loss_with_var: bool = False,
+        progress_bar: bool = False,
+    ) -> None:
         """Optimize the prior precision post-hoc using the `method`
         specified by the user.
 
         Parameters
         ----------
-        pred_type : {'glm', 'nn', 'gp'}, default='glm'
+        pred_type : PredType or str in {'glm', 'nn'}
             type of posterior predictive, linearized GLM predictive or neural
-            network sampling predictive or Gaussian Process (GP) inference.
-            The GLM predictive is consistent with the curvature approximations used here.
-        method : {'marglik', 'gridsearch'}, default='marglik'
+            network sampling predictiv. The GLM predictive is consistent with the
+            curvature approximations used here.
+        method : TuningMethod or str in {'marglik', 'gridsearch'}, default=PredType.MARGLIK
             specifies how the prior precision should be optimized.
         n_steps : int, default=100
             the number of gradient descent steps to take.
@@ -325,7 +373,7 @@ def optimize_prior_precision_base(
             the learning rate to use for gradient descent.
         init_prior_prec : float or tensor, default=1.0
             initial prior precision before the first optimization step.
-        prior_structure : {'scalar', 'layerwise', 'diag'}, default='scalar'
+        prior_structure : PriorStructure or str in {'scalar', 'layerwise', 'diag'}, default=PriorStructure.SCALAR
             if init_prior_prec is scalar, the prior precision is optimized with this structure.
             otherwise, the structure of init_prior_prec is maintained.
         val_loader : torch.data.utils.DataLoader, default=None
@@ -344,7 +392,7 @@ def optimize_prior_precision_base(
             upper bound of gridsearch interval.
         grid_size : int, default=100
             number of values to consider inside the gridsearch interval.
-        link_approx : {'mc', 'probit', 'bridge'}, default='probit'
+        link_approx : LinkApprox or str in {'mc', 'probit', 'bridge'}, default=LinkApprox.PROBIT
             how to approximate the classification link function for the `'glm'`.
             For `pred_type='nn'`, only `'mc'` is possible.
         n_samples : int, default=100
@@ -356,9 +404,17 @@ def optimize_prior_precision_base(
             whether to show a progress bar; updated at every batch-Hessian computation.
             Useful for very large model and large amount of data, esp. when `subset_of_weights='all'`.
         """
-        if method == "marglik":
-            self.prior_precision = init_prior_prec
-            if len(self.prior_precision) == 1 and prior_structure != "scalar":
+        if method == TuningMethod.MARGLIK:
+            self.prior_precision = (
+                init_prior_prec
+                if isinstance(init_prior_prec, torch.Tensor)
+                else torch.tensor(init_prior_prec)
+            )
+
+            if (
+                len(self.prior_precision) == 1
+                and prior_structure != PriorStructure.SCALAR
+            ):
                 self.prior_precision = fix_prior_prec_structure(
                     self.prior_precision.item(),
                     prior_structure,
@@ -366,6 +422,7 @@ def optimize_prior_precision_base(
                     self.n_params,
                     self._device,
                 )
+
             log_prior_prec = self.prior_precision.log()
             log_prior_prec.requires_grad = True
             optimizer = torch.optim.Adam([log_prior_prec], lr=lr)
@@ -384,8 +441,9 @@ def optimize_prior_precision_base(
                 )
                 neg_log_marglik.backward()
                 optimizer.step()
+
             self.prior_precision = log_prior_prec.detach().exp()
-        elif method == "gridsearch":
+        elif method == TuningMethod.GRIDSEARCH:
             if val_loader is None:
                 raise ValueError("gridsearch requires a validation set DataLoader")
 
@@ -393,7 +451,7 @@ def optimize_prior_precision_base(
 
             if loss is None:
                 loss = (
-                    tm.MeanSquaredError(num_outputs=self.n_outputs)
+                    torchmetrics.MeanSquaredError(num_outputs=self.n_outputs)
                     if self.likelihood == "regression"
                     else RunningNLLMetric()
                 )
@@ -415,20 +473,21 @@ def optimize_prior_precision_base(
 
     def _gridsearch(
         self,
-        loss,
-        interval,
-        val_loader,
-        pred_type,
-        link_approx="probit",
-        n_samples=100,
-        loss_with_var=False,
-        progress_bar=False,
-    ):
-        assert callable(loss) or isinstance(loss, tm.Metric)
+        loss: torchmetrics.Metric | Callable[[torch.Tensor], torch.Tensor | float],
+        interval: torch.Tensor,
+        val_loader: DataLoader,
+        pred_type: PredType | str,
+        link_approx: LinkApprox | str = LinkApprox.PROBIT,
+        n_samples: int = 100,
+        loss_with_var: bool = False,
+        progress_bar: bool = False,
+    ) -> torch.Tensor:
+        assert callable(loss) or isinstance(loss, torchmetrics.Metric)
+
+        results: list[float] = list()
+        prior_precs: list[torch.Tensor] = list()
+        pbar = tqdm.tqdm(interval, disable=not progress_bar)
 
-        results = list()
-        prior_precs = list()
-        pbar = tqdm.tqdm(interval) if progress_bar else interval
         for prior_prec in pbar:
             self.prior_precision = prior_prec
             try:
@@ -452,18 +511,20 @@ def _gridsearch(
 
             results.append(result)
             prior_precs.append(prior_prec)
+
         return prior_precs[np.argmin(results)]
 
     @property
-    def sigma_noise(self):
+    def sigma_noise(self) -> torch.Tensor:
         return self._sigma_noise
 
     @sigma_noise.setter
-    def sigma_noise(self, sigma_noise):
+    def sigma_noise(self, sigma_noise: float | torch.Tensor) -> None:
         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):
+        elif isinstance(sigma_noise, torch.Tensor):
             if sigma_noise.ndim == 0:
                 self._sigma_noise = sigma_noise.to(self._device)
             elif sigma_noise.ndim == 1:
@@ -478,7 +539,7 @@ def sigma_noise(self, sigma_noise):
             )
 
     @property
-    def _H_factor(self):
+    def _H_factor(self) -> torch.Tensor:
         sigma2 = self.sigma_noise.square()
         return 1 / sigma2 / self.temperature
 
@@ -509,18 +570,18 @@ class ParametricLaplace(BaseLaplace):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        sigma_noise=1.0,
-        prior_precision=1.0,
-        prior_mean=0.0,
-        temperature=1.0,
-        enable_backprop=False,
-        dict_key_x="inputs_id",
-        dict_key_y="labels",
-        backend=None,
-        backend_kwargs=None,
-        asdl_fisher_kwargs=None,
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        sigma_noise: float | torch.Tensor = 1.0,
+        prior_precision: float | torch.Tensor = 1.0,
+        prior_mean: float | torch.Tensor = 0.0,
+        temperature: float = 1.0,
+        enable_backprop: bool = False,
+        dict_key_x: str = "inputs_id",
+        dict_key_y: str = "labels",
+        backend: type[CurvatureInterface] | None = None,
+        backend_kwargs: dict[str, Any] | None = None,
+        asdl_fisher_kwargs: dict[str, Any] | None = None,
     ):
         super().__init__(
             model,
@@ -539,16 +600,21 @@ def __init__(
         if not hasattr(self, "H"):
             self._init_H()
             # posterior mean/mode
-            self.mean = self.prior_mean
+            self.mean: float | torch.Tensor = self.prior_mean
 
-    def _init_H(self):
+    def _init_H(self) -> None:
         raise NotImplementedError
 
-    def _check_H_init(self):
+    def _check_H_init(self) -> None:
         if getattr(self, "H", None) is None:
             raise AttributeError("Laplace not fitted. Run fit() first.")
 
-    def fit(self, train_loader, override=True, progress_bar=False):
+    def fit(
+        self,
+        train_loader: DataLoader,
+        override: bool = True,
+        progress_bar: bool = False,
+    ) -> None:
         """Fit the local Laplace approximation at the parameters of the model.
 
         Parameters
@@ -565,16 +631,19 @@ def fit(self, train_loader, override=True, progress_bar=False):
         """
         if override:
             self._init_H()
-            self.loss = 0
-            self.n_data = 0
+            self.loss: float | torch.Tensor = 0
+            self.n_data: int = 0
 
         self.model.eval()
 
-        self.mean = parameters_to_vector(self.params)
+        self.mean: torch.Tensor = parameters_to_vector(self.params)
         if not self.enable_backprop:
             self.mean = self.mean.detach()
 
-        data = next(iter(train_loader))
+        data: (
+            tuple[torch.Tensor, torch.Tensor] | MutableMapping[str, torch.Tensor | Any]
+        ) = next(iter(train_loader))
+
         with torch.no_grad():
             if isinstance(data, MutableMapping):  # To support Huggingface dataset
                 if "backpack" in self._backend_cls.__name__.lower() or (
@@ -599,11 +668,9 @@ def fit(self, train_loader, override=True, progress_bar=False):
         setattr(self.model, "output_size", self.n_outputs)
 
         N = len(train_loader.dataset)
-        if progress_bar:
-            pbar = tqdm.tqdm(train_loader)
-            pbar.set_description("[Computing Hessian]")
-        else:
-            pbar = train_loader
+
+        pbar = tqdm.tqdm(train_loader, disable=not progress_bar)
+        pbar.set_description("[Computing Hessian]")
 
         for data in pbar:
             if isinstance(data, MutableMapping):  # To support Huggingface dataset
@@ -612,28 +679,27 @@ def fit(self, train_loader, override=True, progress_bar=False):
                 X, y = data
                 X, y = X.to(self._device), y.to(self._device)
             self.model.zero_grad()
-            loss_batch, H_batch = self._curv_closure(X, y, N)
+            loss_batch, H_batch = self._curv_closure(X, y, N=N)
             self.loss += loss_batch
             self.H += H_batch
 
         self.n_data += N
 
     @property
-    def scatter(self):
+    def scatter(self) -> torch.Tensor:
         """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]
+        scatter: torch.Tensor
         """
         delta = self.mean - self.prior_mean
         return (delta * self.prior_precision_diag) @ delta
 
     @property
-    def log_det_prior_precision(self):
+    def log_det_prior_precision(self) -> torch.Tensor:
         """Compute log determinant of the prior precision
         \\(\\log \\det P_0\\)
 
@@ -644,7 +710,7 @@ def log_det_prior_precision(self):
         return self.prior_precision_diag.log().sum()
 
     @property
-    def log_det_posterior_precision(self):
+    def log_det_posterior_precision(self) -> torch.Tensor:
         """Compute log determinant of the posterior precision
         \\(\\log \\det P\\) which depends on the subclasses structure
         used for the Hessian approximation.
@@ -656,7 +722,7 @@ def log_det_posterior_precision(self):
         raise NotImplementedError
 
     @property
-    def log_det_ratio(self):
+    def log_det_ratio(self) -> torch.Tensor:
         """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
@@ -668,7 +734,7 @@ def log_det_ratio(self):
         """
         return self.log_det_posterior_precision - self.log_det_prior_precision
 
-    def square_norm(self, value):
+    def square_norm(self, value) -> torch.Tensor:
         """Compute the square norm under post. Precision with `value-self.mean` as 𝛥:
         \\[
             \\Delta^\top P \\Delta
@@ -679,11 +745,12 @@ def square_norm(self, value):
         """
         raise NotImplementedError
 
-    def log_prob(self, value, normalized=True):
+    def log_prob(self, value: torch.Tensor, normalized: bool = True) -> torch.Tensor:
         """Compute the log probability under the (current) Laplace approximation.
 
         Parameters
         ----------
+        value: torch.Tensor
         normalized : bool, default=True
             whether to return log of a properly normalized Gaussian or just the
             terms that depend on `value`.
@@ -700,7 +767,11 @@ def log_prob(self, value, normalized=True):
         log_prob -= self.square_norm(value) / 2
         return log_prob
 
-    def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
+    def log_marginal_likelihood(
+        self,
+        prior_precision: torch.Tensor | None = None,
+        sigma_noise: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         """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.
@@ -713,7 +784,7 @@ def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
         ----------
         prior_precision : torch.Tensor, optional
             prior precision if should be changed from current `prior_precision` value
-        sigma_noise : [type], optional
+        sigma_noise : torch.Tensor, optional
             observation noise standard deviation if should be changed
 
         Returns
@@ -726,23 +797,24 @@ def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
 
         # update sigma_noise (useful when iterating on marglik)
         if sigma_noise is not None:
-            if self.likelihood != "regression":
+            if self.likelihood != Likelihood.REGRESSION:
                 raise ValueError("Can only change sigma_noise for regression.")
+
             self.sigma_noise = sigma_noise
 
         return self.log_likelihood - 0.5 * (self.log_det_ratio + self.scatter)
 
     def __call__(
         self,
-        x,
-        pred_type="glm",
-        joint=False,
-        link_approx="probit",
-        n_samples=100,
-        diagonal_output=False,
-        generator=None,
-        **model_kwargs,
-    ):
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        pred_type: PredType | str = PredType.GLM,
+        joint: bool = False,
+        link_approx: LinkApprox | str = LinkApprox.PROBIT,
+        n_samples: int = 100,
+        diagonal_output: bool = False,
+        generator: torch.Generator | None = None,
+        **model_kwargs: dict[str, Any],
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """Compute the posterior predictive on input data `x`.
 
         Parameters
@@ -781,7 +853,7 @@ def __call__(
 
         Returns
         -------
-        predictive: torch.Tensor or Tuple[torch.Tensor]
+        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
@@ -789,13 +861,13 @@ def __call__(
             For `likelihood='regression'` and `joint=True`, a tuple of torch.Tensor
             is returned with the mean and the predictive covariance.
         """
-        if pred_type not in ["glm", "nn"]:
+        if pred_type not in [pred for pred in PredType]:
             raise ValueError("Only glm and nn supported as prediction types.")
 
-        if link_approx not in ["mc", "probit", "bridge", "bridge_norm"]:
+        if link_approx not in [la for la in LinkApprox]:
             raise ValueError(f"Unsupported link approximation {link_approx}.")
 
-        if pred_type == "nn" and link_approx != "mc":
+        if pred_type == PredType.NN and link_approx != LinkApprox.MC:
             raise ValueError(
                 "Only mc link approximation is supported for nn prediction type."
             )
@@ -803,30 +875,30 @@ def __call__(
         if generator is not None:
             if (
                 not isinstance(generator, torch.Generator)
-                or generator.device != x.device
+                or generator.device != self._device
             ):
                 raise ValueError("Invalid random generator (check type and device).")
 
-        # For reward modeling, replace the likelihood to regression
-        if self.reward_modeling and self.likelihood == "classification":
-            self.likelihood = "regression"
+        # For reward modeling, replace the likelihood to regression and override model state
+        if self.reward_modeling and self.likelihood == Likelihood.CLASSIFICATION:
+            self.likelihood = Likelihood.REGRESSION
 
-        if pred_type == "glm":
+        if pred_type == PredType.GLM:
             f_mu, f_var = self._glm_predictive_distribution(
                 x, joint=joint and self.likelihood == "regression"
             )
-            # regression
-            if self.likelihood == "regression":
+
+            if self.likelihood == Likelihood.REGRESSION:
                 return f_mu, f_var
-            # classification
-            if link_approx == "mc":
+
+            if link_approx == LinkApprox.MC:
                 return self.predictive_samples(
                     x,
                     pred_type="glm",
                     n_samples=n_samples,
                     diagonal_output=diagonal_output,
                 ).mean(dim=0)
-            elif link_approx == "probit":
+            elif link_approx == LinkApprox.PROBIT:
                 kappa = 1 / torch.sqrt(1.0 + np.pi / 8 * f_var.diagonal(dim1=1, dim2=2))
                 return torch.softmax(kappa * f_mu, dim=-1)
             elif "bridge" in link_approx:
@@ -839,36 +911,48 @@ def __call__(
                 f_var -= torch.einsum(
                     "bi,bj->bij", f_var.sum(-1), f_var.sum(-2)
                 ) / f_var.sum(dim=(1, 2)).reshape(-1, 1, 1)
+
                 # Laplace Bridge
                 _, K = f_mu.size(0), f_mu.size(-1)
                 f_var_diag = torch.diagonal(f_var, dim1=1, dim2=2)
+
                 # optional: variance correction
-                if link_approx == "bridge_norm":
+                if link_approx == LinkApprox.BRIDGE_NORM:
                     f_var_diag_mean = f_var_diag.mean(dim=1)
                     f_var_diag_mean /= torch.as_tensor(
                         [K / 2], device=self._device
                     ).sqrt()
                     f_mu /= f_var_diag_mean.sqrt().unsqueeze(-1)
                     f_var_diag /= f_var_diag_mean.unsqueeze(-1)
+
                 sum_exp = torch.exp(-f_mu).sum(dim=1).unsqueeze(-1)
                 alpha = (1 - 2 / K + f_mu.exp() / K**2 * sum_exp) / f_var_diag
                 return torch.nan_to_num(alpha / alpha.sum(dim=1).unsqueeze(-1), nan=1.0)
+            else:
+                raise ValueError(
+                    "Prediction path invalid. Check the likelihood, pred_type, link_approx combination!"
+                )
         else:
-            if self.likelihood == "regression":
+            if self.likelihood == Likelihood.REGRESSION:
                 samples = self._nn_predictive_samples(x, n_samples, **model_kwargs)
                 return samples.mean(dim=0), samples.var(dim=0)
             else:  # classification; the average is computed online
                 return self._nn_predictive_classification(x, n_samples, **model_kwargs)
 
     def predictive_samples(
-        self, x, pred_type="glm", n_samples=100, diagonal_output=False, generator=None
-    ):
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        pred_type: PredType | str = PredType.GLM,
+        n_samples: int = 100,
+        diagonal_output: bool = False,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
         """Sample from the posterior predictive on input data `x`.
         Can be used, for example, for Thompson sampling.
 
         Parameters
         ----------
-        x : torch.Tensor
+        x : torch.Tensor or MutableMapping
             input data `(batch_size, input_shape)`
 
         pred_type : {'glm', 'nn'}, default='glm'
@@ -891,18 +975,21 @@ def predictive_samples(
         samples : torch.Tensor
             samples `(n_samples, batch_size, output_shape)`
         """
-        if pred_type not in ["glm", "nn"]:
+        if pred_type not in PredType.__members__.values():
             raise ValueError("Only glm and nn supported as prediction types.")
 
-        if pred_type == "glm":
+        if pred_type == PredType.GLM:
             f_mu, f_var = self._glm_predictive_distribution(x)
             assert f_var.shape == torch.Size(
                 [f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]]
             )
+
             if diagonal_output:
                 f_var = torch.diagonal(f_var, dim1=1, dim2=2)
+
             f_samples = normal_samples(f_mu, f_var, n_samples, generator)
-            if self.likelihood == "regression":
+
+            if self.likelihood == Likelihood.REGRESSION:
                 return f_samples
             else:
                 return torch.softmax(f_samples, dim=-1)
@@ -911,7 +998,11 @@ def predictive_samples(
             return self._nn_predictive_samples(x, n_samples, generator)
 
     @torch.enable_grad()
-    def _glm_predictive_distribution(self, X, joint=False):
+    def _glm_predictive_distribution(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        joint: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         backend_name = self._backend_cls.__name__.lower()
         if self.enable_backprop and (
             "curvlinops" not in backend_name and "backpack" not in backend_name
@@ -935,7 +1026,13 @@ def _glm_predictive_distribution(self, X, joint=False):
             else (f_mu, f_var)
         )
 
-    def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwargs):
+    def _nn_predictive_samples(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        n_samples: int = 100,
+        generator: torch.Generator | None = None,
+        **model_kwargs: dict[str, Any],
+    ) -> torch.Tensor:
         fs = list()
         for sample in self.sample(n_samples, generator):
             vector_to_parameters(sample, self.params)
@@ -943,26 +1040,36 @@ def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwarg
                 X.to(self._device) if isinstance(X, torch.Tensor) else X, **model_kwargs
             )
             fs.append(logits.detach() if not self.enable_backprop else logits)
+
         vector_to_parameters(self.mean, self.params)
         fs = torch.stack(fs)
-        if self.likelihood == "classification":
+
+        if self.likelihood == Likelihood.CLASSIFICATION:
             fs = torch.softmax(fs, dim=-1)
+
         return fs
 
-    def _nn_predictive_classification(self, X, n_samples=100, **model_kwargs):
-        py = 0
+    def _nn_predictive_classification(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        n_samples: int = 100,
+        **model_kwargs: dict[str, Any],
+    ) -> torch.Tensor:
+        py = 0.0
         for sample in self.sample(n_samples):
             vector_to_parameters(sample, self.params)
             logits = self.model(
                 X.to(self._device) if isinstance(X, torch.Tensor) else X, **model_kwargs
             ).detach()
             py += torch.softmax(logits, dim=-1) / n_samples
+
         vector_to_parameters(self.mean, self.params)
+
         return py
 
-    def functional_variance(self, Jacs):
+    def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
         """Compute functional variance for the `'glm'` predictive:
-        `f_var[i] = Jacs[i] @ P.inv() @ Jacs[i].T`, which is a output x output
+        `f_var[i] = Js[i] @ P.inv() @ Js[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}}\\)
@@ -971,7 +1078,7 @@ def functional_variance(self, Jacs):
 
         Parameters
         ----------
-        Jacs : torch.Tensor
+        Js : torch.Tensor
             Jacobians of model output wrt parameters
             `(batch, outputs, parameters)`
 
@@ -982,9 +1089,9 @@ def functional_variance(self, Jacs):
         """
         raise NotImplementedError
 
-    def functional_covariance(self, Jacs):
+    def functional_covariance(self, Js: torch.Tensor) -> torch.Tensor:
         """Compute functional covariance for the `'glm'` predictive:
-        `f_cov = Jacs @ P.inv() @ Jacs.T`, which is a batch*output x batch*output
+        `f_cov = Js @ P.inv() @ Js.T`, which is a batch*output x batch*output
         predictive covariance matrix.
 
         This emulates the GP posterior covariance N([f(x1), ...,f(xm)], Cov[f(x1), ..., f(xm)]).
@@ -992,7 +1099,7 @@ def functional_covariance(self, Jacs):
 
         Parameters
         ----------
-        Jacs : torch.Tensor
+        Js : torch.Tensor
             Jacobians of model output wrt parameters
             `(batch*outputs, parameters)`
 
@@ -1003,7 +1110,9 @@ def functional_covariance(self, Jacs):
         """
         raise NotImplementedError
 
-    def sample(self, n_samples=100, generator=None):
+    def sample(
+        self, n_samples: int = 100, generator: torch.Generator | None = None
+    ) -> torch.Tensor:
         """Sample from the Laplace posterior approximation, i.e.,
         \\( \\theta \\sim \\mathcal{N}(\\theta_{MAP}, P^{-1})\\).
 
@@ -1014,30 +1123,37 @@ def sample(self, n_samples=100, generator=None):
 
         generator : torch.Generator, optional
             random number generator to control the samples
+
+        Returns
+        -------
+        samples: torch.Tensor
         """
         raise NotImplementedError
 
     def optimize_prior_precision(
         self,
-        method="marglik",
-        pred_type="glm",
-        n_steps=100,
-        lr=1e-1,
-        init_prior_prec=1.0,
-        prior_structure="scalar",
-        val_loader=None,
-        loss=None,
-        log_prior_prec_min=-4,
-        log_prior_prec_max=4,
-        grid_size=100,
-        link_approx="probit",
-        n_samples=100,
-        verbose=False,
-        cv_loss_with_var=False,
-        progress_bar=False,
-    ):
-        assert pred_type in ["glm", "nn"]
-        self.optimize_prior_precision_base(
+        pred_type: PredType | str = PredType.GLM,
+        method: TuningMethod | str = TuningMethod.MARGLIK,
+        n_steps: int = 100,
+        lr: float = 1e-1,
+        init_prior_prec: float | torch.Tensor = 1.0,
+        prior_structure: PriorStructure | str = PriorStructure.SCALAR,
+        val_loader: DataLoader | None = None,
+        loss: torchmetrics.Metric
+        | Callable[[torch.Tensor], torch.Tensor | float]
+        | None = None,
+        log_prior_prec_min: float = -4,
+        log_prior_prec_max: float = 4,
+        grid_size: int = 100,
+        link_approx: LinkApprox | str = LinkApprox.PROBIT,
+        n_samples: int = 100,
+        verbose: bool = False,
+        cv_loss_with_var: bool = False,
+        progress_bar: bool = False,
+    ) -> None:
+        assert pred_type in PredType.__members__.values()
+
+        super().optimize_prior_precision(
             pred_type,
             method,
             n_steps,
@@ -1057,7 +1173,7 @@ def optimize_prior_precision(
         )
 
     @property
-    def posterior_precision(self):
+    def posterior_precision(self) -> torch.Tensor:
         """Compute or return the posterior precision \\(P\\).
 
         Returns
@@ -1066,7 +1182,7 @@ def posterior_precision(self):
         """
         raise NotImplementedError
 
-    def state_dict(self) -> dict:
+    def state_dict(self) -> dict[str, Any]:
         self._check_H_init()
         state_dict = {
             "mean": self.mean,
@@ -1084,7 +1200,7 @@ def state_dict(self) -> dict:
         }
         return state_dict
 
-    def load_state_dict(self, state_dict: dict):
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         # Dealbreaker errors
         if self.__class__.__name__ != state_dict["cls_name"]:
             raise ValueError(
@@ -1142,17 +1258,17 @@ class FullLaplace(ParametricLaplace):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        sigma_noise=1.0,
-        prior_precision=1.0,
-        prior_mean=0.0,
-        temperature=1.0,
-        enable_backprop=False,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        backend=None,
-        backend_kwargs=None,
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        sigma_noise: float | torch.Tensor = 1.0,
+        prior_precision: float | torch.Tensor = 1.0,
+        prior_mean: float | torch.Tensor = 0.0,
+        temperature: float = 1.0,
+        enable_backprop: bool = False,
+        dict_key_x: str = "input_ids",
+        dict_key_y: str = "labels",
+        backend: type[CurvatureInterface] | None = None,
+        backend_kwargs: dict[str, Any] | None = None,
     ):
         super().__init__(
             model,
@@ -1167,23 +1283,35 @@ def __init__(
             backend,
             backend_kwargs,
         )
-        self._posterior_scale = None
+        self._posterior_scale: torch.Tensor | None = None
 
-    def _init_H(self):
-        self.H = torch.zeros(self.n_params, self.n_params, device=self._device)
+    def _init_H(self) -> None:
+        self.H: torch.Tensor = torch.zeros(
+            self.n_params, self.n_params, device=self._device
+        )
 
-    def _curv_closure(self, X, y, N):
+    def _curv_closure(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         return self.backend.full(X, y, N=N)
 
-    def fit(self, train_loader, override=True, progress_bar=False):
+    def fit(
+        self,
+        train_loader: DataLoader,
+        override: bool = True,
+        progress_bar: bool = False,
+    ) -> None:
         self._posterior_scale = None
-        return super().fit(train_loader, override=override, progress_bar=progress_bar)
+        super().fit(train_loader, override=override, progress_bar=progress_bar)
 
-    def _compute_scale(self):
+    def _compute_scale(self) -> None:
         self._posterior_scale = invsqrt_precision(self.posterior_precision)
 
     @property
-    def posterior_scale(self):
+    def posterior_scale(self) -> torch.Tensor:
         """Posterior scale (square root of the covariance), i.e.,
         \\(P^{-\\frac{1}{2}}\\).
 
@@ -1197,7 +1325,7 @@ def posterior_scale(self):
         return self._posterior_scale
 
     @property
-    def posterior_covariance(self):
+    def posterior_covariance(self) -> torch.Tensor:
         """Posterior covariance, i.e., \\(P^{-1}\\).
 
         Returns
@@ -1209,7 +1337,7 @@ def posterior_covariance(self):
         return scale @ scale.T
 
     @property
-    def posterior_precision(self):
+    def posterior_precision(self) -> torch.Tensor:
         """Posterior precision \\(P\\).
 
         Returns
@@ -1221,22 +1349,24 @@ def posterior_precision(self):
         return self._H_factor * self.H + torch.diag(self.prior_precision_diag)
 
     @property
-    def log_det_posterior_precision(self):
+    def log_det_posterior_precision(self) -> torch.Tensor:
         return self.posterior_precision.logdet()
 
-    def square_norm(self, value):
+    def square_norm(self, value: torch.Tensor) -> torch.Tensor:
         delta = value - self.mean
         return delta @ self.posterior_precision @ delta
 
-    def functional_variance(self, Js):
+    def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
         return torch.einsum("ncp,pq,nkq->nck", Js, self.posterior_covariance, Js)
 
-    def functional_covariance(self, Js):
+    def functional_covariance(self, Js: torch.Tensor) -> torch.Tensor:
         n_batch, n_outs, n_params = Js.shape
         Js = Js.reshape(n_batch * n_outs, n_params)
         return torch.einsum("np,pq,mq->nm", Js, self.posterior_covariance, Js)
 
-    def sample(self, n_samples=100, generator=None):
+    def sample(
+        self, n_samples: int = 100, generator: torch.Generator | None = None
+    ) -> torch.Tensor:
         samples = torch.randn(
             n_samples, self.n_params, device=self._device, generator=generator
         )
@@ -1263,22 +1393,22 @@ class KronLaplace(ParametricLaplace):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        sigma_noise=1.0,
-        prior_precision=1.0,
-        prior_mean=0.0,
-        temperature=1.0,
-        enable_backprop=False,
-        dict_key_x="inputs_id",
-        dict_key_y="labels",
-        backend=None,
-        damping=False,
-        backend_kwargs=None,
-        asdl_fisher_kwargs=None,
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        sigma_noise: float | torch.Tensor = 1.0,
+        prior_precision: float | torch.Tensor = 1.0,
+        prior_mean: float | torch.Tensor = 0.0,
+        temperature: float = 1.0,
+        enable_backprop: bool = False,
+        dict_key_x: str = "inputs_id",
+        dict_key_y: str = "labels",
+        backend: type[CurvatureInterface] | None = None,
+        damping: bool = False,
+        backend_kwargs: dict[str, Any] | None = None,
+        asdl_fisher_kwargs: dict[str, Any] | None = None,
     ):
-        self.damping = damping
-        self.H_facs = None
+        self.damping: bool = damping
+        self.H_facs: Kron | None = None
         super().__init__(
             model,
             likelihood,
@@ -1294,30 +1424,42 @@ def __init__(
             asdl_fisher_kwargs,
         )
 
-    def _init_H(self):
-        self.H = Kron.init_from_model(self.params, self._device)
+    def _init_H(self) -> None:
+        self.H: Kron | KronDecomposed | None = Kron.init_from_model(
+            self.params, self._device
+        )
 
     def _check_H_init(self):
         if getattr(self, "H_facs", None) is None:
             raise AttributeError("Laplace not fitted. Run fit() first.")
 
-    def _curv_closure(self, X, y, N):
+    def _curv_closure(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         return self.backend.kron(X, y, N=N, **self._asdl_fisher_kwargs)
 
     @staticmethod
-    def _rescale_factors(kron, factor):
+    def _rescale_factors(kron: Kron, factor: float) -> Kron:
         for F in kron.kfacs:
             if len(F) == 2:
                 F[1] *= factor
         return kron
 
-    def fit(self, train_loader, override=True, progress_bar=False):
+    def fit(
+        self,
+        train_loader: DataLoader,
+        override: bool = True,
+        progress_bar: bool = False,
+    ) -> None:
         if override:
             self.H_facs = None
 
         if self.H_facs is not None:
-            n_data_old = self.n_data
-            n_data_new = len(train_loader.dataset)
+            n_data_old: int = self.n_data
+            n_data_new: int = len(train_loader.dataset)
             self._init_H()  # re-init H non-decomposed
             # discount previous Kronecker factors to sum up properly together with new ones
             self.H_facs = self._rescale_factors(
@@ -1334,11 +1476,12 @@ def fit(self, train_loader, override=True, progress_bar=False):
                 self.H, n_data_new / (n_data_new + n_data_old)
             )
             self.H_facs += self.H
+
         # Decompose to self.H for all required quantities but keep H_facs for further inference
         self.H = self.H_facs.decompose(damping=self.damping)
 
     @property
-    def posterior_precision(self):
+    def posterior_precision(self) -> KronDecomposed:
         """Kronecker factored Posterior precision \\(P\\).
 
         Returns
@@ -1349,21 +1492,21 @@ def posterior_precision(self):
         return self.H * self._H_factor + self.prior_precision
 
     @property
-    def log_det_posterior_precision(self):
+    def log_det_posterior_precision(self) -> torch.Tensor:
         if type(self.H) is Kron:  # Fall back to diag prior
             return self.prior_precision_diag.log().sum()
         return self.posterior_precision.logdet()
 
-    def square_norm(self, value):
+    def square_norm(self, value: torch.Tensor) -> torch.Tensor:
         delta = value - self.mean
         if type(self.H) is Kron:  # fall back to prior
             return (delta * self.prior_precision_diag) @ delta
         return delta @ self.posterior_precision.bmm(delta, exponent=1)
 
-    def functional_variance(self, Js):
+    def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
         return self.posterior_precision.inv_square_form(Js)
 
-    def functional_covariance(self, Js):
+    def functional_covariance(self, Js: torch.Tensor) -> torch.Tensor:
         self._check_jacobians(Js)
         n_batch, n_outs, n_params = Js.shape
         Js = Js.reshape(n_batch * n_outs, n_params).unsqueeze(0)
@@ -1371,7 +1514,9 @@ def functional_covariance(self, Js):
         assert cov.shape == (n_batch * n_outs, n_batch * n_outs)
         return cov
 
-    def sample(self, n_samples=100, generator=None):
+    def sample(
+        self, n_samples: int = 100, generator: torch.Generator | None = None
+    ) -> torch.Tensor:
         samples = torch.randn(
             n_samples, self.n_params, device=self._device, generator=generator
         )
@@ -1381,20 +1526,22 @@ def sample(self, n_samples=100, generator=None):
         )
 
     @BaseLaplace.prior_precision.setter
-    def prior_precision(self, prior_precision):
+    def prior_precision(self, prior_precision: torch.Tensor) -> None:
         # Extend setter from Laplace to restrict prior precision structure.
         super(KronLaplace, type(self)).prior_precision.fset(self, prior_precision)
         if len(self.prior_precision) not in [1, self.n_layers]:
             raise ValueError("Prior precision for Kron either scalar or per-layer.")
 
-    def state_dict(self) -> dict:
+    def state_dict(self) -> dict[str, Any]:
         state_dict = super().state_dict()
+        assert isinstance(self.H_facs, Kron)
         state_dict["H"] = self.H_facs.kfacs
         return state_dict
 
-    def load_state_dict(self, state_dict: dict):
+    def load_state_dict(self, state_dict: dict[str, Any]):
         super().load_state_dict(state_dict)
         self._init_H()
+        assert isinstance(self.H, Kron)
         self.H_facs = self.H
         self.H_facs.kfacs = state_dict["H"]
         self.H = self.H_facs.decompose(damping=self.damping)
@@ -1418,17 +1565,17 @@ class LowRankLaplace(ParametricLaplace):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        sigma_noise=1,
-        prior_precision=1,
-        prior_mean=0,
-        temperature=1,
-        enable_backprop=False,
-        dict_key_x="inputs_id",
-        dict_key_y="labels",
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        sigma_noise: float | torch.Tensor = 1,
+        prior_precision: float | torch.Tensor = 1,
+        prior_mean: float | torch.Tensor = 0,
+        temperature: float = 1,
+        enable_backprop: bool = False,
+        dict_key_x: str = "inputs_id",
+        dict_key_y: str = "labels",
         backend=AsdfghjklHessian,
-        backend_kwargs=None,
+        backend_kwargs: dict[str, Any] | None = None,
     ):
         super().__init__(
             model,
@@ -1443,21 +1590,27 @@ def __init__(
             backend=backend,
             backend_kwargs=backend_kwargs,
         )
+        self.backend: AsdfghjklHessian
 
     def _init_H(self):
-        self.H = None
+        self.H: tuple[torch.Tensor, torch.Tensor] | None = None
 
     @property
-    def V(self):
-        (U, _), prior_prec_diag = self.posterior_precision
+    def V(self) -> torch.Tensor:
+        (U, eigvals), prior_prec_diag = self.posterior_precision
         return U / prior_prec_diag.reshape(-1, 1)
 
     @property
-    def Kinv(self):
-        (U, eigval), _ = self.posterior_precision
-        return torch.inverse(torch.diag(1 / eigval) + U.T @ self.V)
+    def Kinv(self) -> torch.Tensor:
+        (U, eigvals), _ = self.posterior_precision
+        return torch.inverse(torch.diag(1 / eigvals) + U.T @ self.V)
 
-    def fit(self, train_loader, override=True):
+    def fit(
+        self,
+        train_loader: DataLoader,
+        override: bool = True,
+        progress_bar: bool = False,
+    ) -> None:
         # override fit since output of eighessian not additive across batch
         if not override:
             # LowRankLA cannot be updated since eigenvalue representation not additive
@@ -1485,7 +1638,9 @@ def fit(self, train_loader, override=True):
         self.n_data = len(train_loader.dataset)
 
     @property
-    def posterior_precision(self):
+    def posterior_precision(
+        self,
+    ) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
         """Return correctly scaled posterior precision that would be constructed
         as H[0] @ diag(H[1]) @ H[0].T + self.prior_precision_diag.
 
@@ -1499,23 +1654,25 @@ def posterior_precision(self):
         self._check_H_init()
         return (self.H[0], self._H_factor * self.H[1]), self.prior_precision_diag
 
-    def functional_variance(self, Jacs):
-        prior_var = torch.einsum("ncp,nkp->nck", Jacs / self.prior_precision_diag, Jacs)
-        Jacs_V = torch.einsum("ncp,pl->ncl", Jacs, self.V)
-        info_gain = torch.einsum("ncl,nkl->nck", Jacs_V @ self.Kinv, Jacs_V)
+    def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
+        prior_var = torch.einsum("ncp,nkp->nck", Js / self.prior_precision_diag, Js)
+        Js_V = torch.einsum("ncp,pl->ncl", Js, self.V)
+        info_gain = torch.einsum("ncl,nkl->nck", Js_V @ self.Kinv, Js_V)
         return prior_var - info_gain
 
-    def functional_covariance(self, Jacs):
-        n_batch, n_outs, n_params = Jacs.shape
-        Jacs = Jacs.reshape(n_batch * n_outs, n_params)
-        prior_cov = torch.einsum("np,mp->nm", Jacs / self.prior_precision_diag, Jacs)
-        Jacs_V = torch.einsum("np,pl->nl", Jacs, self.V)
-        info_gain = torch.einsum("nl,ml->nm", Jacs_V @ self.Kinv, Jacs_V)
+    def functional_covariance(self, Js: torch.Tensor) -> torch.Tensor:
+        n_batch, n_outs, n_params = Js.shape
+        Js = Js.reshape(n_batch * n_outs, n_params)
+        prior_cov = torch.einsum("np,mp->nm", Js / self.prior_precision_diag, Js)
+        Js_V = torch.einsum("np,pl->nl", Js, self.V)
+        info_gain = torch.einsum("nl,ml->nm", Js_V @ self.Kinv, Js_V)
         cov = prior_cov - info_gain
         assert cov.shape == (n_batch * n_outs, n_batch * n_outs)
         return cov
 
-    def sample(self, n_samples, generator=None):
+    def sample(
+        self, n_samples: int = 100, generator: torch.Generator | None = None
+    ) -> torch.Tensor:
         samples = torch.randn(self.n_params, n_samples, generator=generator)
         d = self.prior_precision_diag
         Vs = self.V * d.sqrt().reshape(-1, 1)
@@ -1532,10 +1689,10 @@ def sample(self, n_samples, generator=None):
         return self.mean + (prior_sample - gain_sample).T
 
     @property
-    def log_det_posterior_precision(self):
-        (U, eigval), prior_prec_diag = self.posterior_precision
+    def log_det_posterior_precision(self) -> torch.Tensor:
+        (_, eigvals), prior_prec_diag = self.posterior_precision
         return (
-            eigval.log().sum() + prior_prec_diag.log().sum() - torch.logdet(self.Kinv)
+            eigvals.log().sum() + prior_prec_diag.log().sum() - torch.logdet(self.Kinv)
         )
 
 
@@ -1549,14 +1706,19 @@ class DiagLaplace(ParametricLaplace):
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ("all", "diag")
 
-    def _init_H(self):
-        self.H = torch.zeros(self.n_params, device=self._device)
+    def _init_H(self) -> None:
+        self.H: torch.Tensor = torch.zeros(self.n_params, device=self._device)
 
-    def _curv_closure(self, X, y, N):
+    def _curv_closure(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         return self.backend.diag(X, y, N=N, **self._asdl_fisher_kwargs)
 
     @property
-    def posterior_precision(self):
+    def posterior_precision(self) -> torch.Tensor:
         """Diagonal posterior precision \\(p\\).
 
         Returns
@@ -1568,7 +1730,7 @@ def posterior_precision(self):
         return self._H_factor * self.H + self.prior_precision_diag
 
     @property
-    def posterior_scale(self):
+    def posterior_scale(self) -> torch.Tensor:
         """Diagonal posterior scale \\(\\sqrt{p^{-1}}\\).
 
         Returns
@@ -1579,7 +1741,7 @@ def posterior_scale(self):
         return 1 / self.posterior_precision.sqrt()
 
     @property
-    def posterior_variance(self):
+    def posterior_variance(self) -> torch.Tensor:
         """Diagonal posterior variance \\(p^{-1}\\).
 
         Returns
@@ -1590,10 +1752,10 @@ def posterior_variance(self):
         return 1 / self.posterior_precision
 
     @property
-    def log_det_posterior_precision(self):
+    def log_det_posterior_precision(self) -> torch.Tensor:
         return self.posterior_precision.log().sum()
 
-    def square_norm(self, value):
+    def square_norm(self, value: torch.Tensor) -> torch.Tensor:
         delta = value - self.mean
         return delta @ (delta * self.posterior_precision)
 
@@ -1601,14 +1763,16 @@ def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
         self._check_jacobians(Js)
         return torch.einsum("ncp,p,nkp->nck", Js, self.posterior_variance, Js)
 
-    def functional_covariance(self, Js):
+    def functional_covariance(self, Js: torch.Tensor) -> torch.Tensor:
         self._check_jacobians(Js)
         n_batch, n_outs, n_params = Js.shape
         Js = Js.reshape(n_batch * n_outs, n_params)
         cov = torch.einsum("np,p,mp->nm", Js, self.posterior_variance, Js)
         return cov
 
-    def sample(self, n_samples=100, generator=None):
+    def sample(
+        self, n_samples: int = 100, generator: torch.Generator | None = None
+    ) -> torch.Tensor:
         samples = torch.randn(
             n_samples, self.n_params, device=self._device, generator=generator
         )
diff --git a/laplace/curvature/asdfghjkl.py b/laplace/curvature/asdfghjkl.py
index b83009d..ad7dacf 100644
--- a/laplace/curvature/asdfghjkl.py
+++ b/laplace/curvature/asdfghjkl.py
@@ -1,4 +1,8 @@
+from __future__ import annotations
+
 import warnings
+from collections.abc import MutableMapping
+from typing import Any
 
 import numpy as np
 import torch
@@ -13,9 +17,12 @@
 )
 from asdfghjkl.gradient import batch_gradient
 from asdfghjkl.hessian import hessian_eigenvalues, hessian_for_loss
+from torch import nn
+from torch.utils.data import DataLoader
 
 from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
 from laplace.utils import Kron, _is_batchnorm
+from laplace.utils.enums import Likelihood
 
 EPS = 1e-6
 
@@ -23,7 +30,11 @@
 class AsdfghjklInterface(CurvatureInterface):
     """Interface for asdfghjkl backend."""
 
-    def jacobians(self, x, enable_backprop=False):
+    def jacobians(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        enable_backprop: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
         using asdfghjkl's gradient per output dimension.
 
@@ -55,7 +66,11 @@ def loss_fn(outputs, targets):
         Js = torch.stack(Js, dim=1)
         return Js, f
 
-    def gradients(self, x, y):
+    def gradients(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
         \\(\\theta\\) using asdfghjkl's backend.
 
@@ -79,10 +94,10 @@ def gradients(self, x, y):
         return Gs, loss
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         raise NotImplementedError
 
-    def _get_kron_factors(self, curv, M):
+    def _get_kron_factors(self, curv, M: int) -> Kron:
         kfacs = list()
         for module in curv._model.modules():
             if _is_batchnorm(module):
@@ -107,36 +122,47 @@ def _get_kron_factors(self, curv, M):
         return Kron(kfacs)
 
     @staticmethod
-    def _rescale_kron_factors(kron, N):
+    def _rescale_kron_factors(kron: Kron, N: int) -> Kron:
         for F in kron.kfacs:
             if len(F) == 2:
                 F[1] *= 1 / N
         return kron
 
-    def diag(self, X, y, **kwargs):
+    def diag(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         with torch.no_grad():
             if self.last_layer:
-                f, X = self.model.forward_with_features(X)
+                f, x = self.model.forward_with_features(x)
             else:
-                f = self.model(X)
+                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
+            self._model, self._ggn_type, SHAPE_DIAG, inputs=x, targets=y
         )
         diag_ggn = curv.matrices_to_vector(None)
         if self.subnetwork_indices is not None:
             diag_ggn = diag_ggn[self.subnetwork_indices]
         return self.factor * loss, self.factor * diag_ggn
 
-    def kron(self, X, y, N, **kwargs):
+    def kron(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, Kron]:
         with torch.no_grad():
             if self.last_layer:
-                f, X = self.model.forward_with_features(X)
+                f, x = self.model.forward_with_features(x)
             else:
-                f = self.model(X)
+                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
+            self._model, self._ggn_type, SHAPE_KRON, inputs=x, targets=y
         )
         M = len(y)
         kron = self._get_kron_factors(curv, M)
@@ -147,13 +173,13 @@ def kron(self, X, y, N, **kwargs):
 class AsdfghjklHessian(AsdfghjklInterface):
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        low_rank=10,
-    ):
+        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,
+    ) -> None:
         super().__init__(
             model,
             likelihood,
@@ -165,16 +191,23 @@ def __init__(
         self.low_rank = low_rank
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         raise NotImplementedError()
 
-    def full(self, x, y, **kwargs):
+    def full(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         hessian_for_loss(self.model, self.lossfunc, SHAPE_FULL, x, y)
         H = self._model.hessian.data
         loss = self.lossfunc(self.model(x), y).detach()
         return self.factor * loss, self.factor * H
 
-    def eig_lowrank(self, data_loader):
+    def eig_lowrank(
+        self, data_loader: DataLoader
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         # compute truncated eigendecomposition of the Hessian, only keep eigvals > EPS
         eigvals, eigvecs = hessian_eigenvalues(
             self.model,
@@ -204,15 +237,15 @@ class AsdfghjklGGN(AsdfghjklInterface, GGNInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        stochastic=False,
-    ):
-        if likelihood != "classification":
+        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,
+    ) -> None:
+        if likelihood != Likelihood.CLASSIFICATION:
             raise ValueError("This backend only supports classification currently.")
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
@@ -220,7 +253,7 @@ def __init__(
         self.stochastic = stochastic
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         return FISHER_MC if self.stochastic else FISHER_EXACT
 
 
@@ -229,29 +262,30 @@ class AsdfghjklEF(AsdfghjklInterface, EFInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-    ):
-        if likelihood != "classification":
+        model: nn.Module,
+        likelihood: Likelihood | None,
+        last_layer: bool = False,
+        dict_key_x: str = "input_ids",
+        dict_key_y: str = "labels",
+    ) -> None:
+        if likelihood != Likelihood.CLASSIFICATION:
             raise ValueError("This backend only supports classification currently.")
+
         super().__init__(model, likelihood, last_layer, None, dict_key_x, dict_key_y)
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         return COV
 
 
-def _flatten_after_batch(tensor: torch.Tensor):
+def _flatten_after_batch(tensor: torch.Tensor) -> torch.Tensor:
     if tensor.ndim == 1:
         return tensor.unsqueeze(-1)
     else:
         return tensor.flatten(start_dim=1)
 
 
-def _get_batch_grad(model):
+def _get_batch_grad(model: nn.Module) -> torch.Tensor:
     batch_grads = list()
     for module in model.modules():
         if hasattr(module, "op_results"):
diff --git a/laplace/curvature/asdl.py b/laplace/curvature/asdl.py
index 8765855..9acc366 100644
--- a/laplace/curvature/asdl.py
+++ b/laplace/curvature/asdl.py
@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 import warnings
 from collections.abc import MutableMapping
+from typing import Any
 
 import torch
 from asdl.fisher import FisherConfig, get_fisher_maker
@@ -14,9 +17,10 @@
     SHAPE_FULL,
     SHAPE_KRON,
 )
+from torch import nn
 
 from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
-from laplace.utils import Kron, _is_batchnorm
+from laplace.utils import Kron, Likelihood, _is_batchnorm
 
 EPS = 1e-6
 
@@ -26,22 +30,28 @@ class AsdlInterface(CurvatureInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
+        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",
     ):
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
         )
 
     @property
-    def loss_type(self):
-        return LOSS_MSE if self.likelihood == "regression" else LOSS_CROSS_ENTROPY
+    def loss_type(self) -> str:
+        return (
+            LOSS_MSE if self.likelihood == Likelihood.REGRESSION else LOSS_CROSS_ENTROPY
+        )
 
-    def jacobians(self, x, enable_backprop=False):
+    def jacobians(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        enable_backprop: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
         using asdfghjkl's gradient per output dimension.
 
@@ -85,7 +95,9 @@ def closure():
         Js = torch.stack(Js, dim=1)
         return Js, f
 
-    def gradients(self, x, y):
+    def gradients(
+        self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
         \\(\\theta\\) using asdfghjkl's backend.
 
@@ -116,10 +128,10 @@ def closure():
         return Gs, loss
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         raise NotImplementedError
 
-    def _get_kron_factors(self, M):
+    def _get_kron_factors(self, M: int) -> Kron:
         kfacs = list()
         for module in self.model.modules():
             if _is_batchnorm(module):
@@ -145,16 +157,24 @@ def _get_kron_factors(self, M):
         return Kron(kfacs)
 
     @staticmethod
-    def _rescale_kron_factors(kron, N):
+    def _rescale_kron_factors(kron: Kron, N: int) -> Kron:
         for F in kron.kfacs:
             if len(F) == 2:
                 F[1] *= 1 / N
         return kron
 
-    def diag(self, X, y, N=None, **kwargs):
-        del N
+    def diag(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if "N" in kwargs:
+            del kwargs["N"]
+
         if self.last_layer:
-            _, X = self.model.forward_with_features(X)
+            _, x = self.model.forward_with_features(x)
+
         cfg = FisherConfig(
             fisher_type=self._ggn_type,
             loss_type=self.loss_type,
@@ -165,13 +185,13 @@ def diag(self, X, y, N=None, **kwargs):
         fisher_maker = get_fisher_maker(self.model, cfg)
         y = y if self.loss_type == LOSS_MSE else y.view(-1)
         if "emp" in self._ggn_type:
-            dummy = fisher_maker.setup_model_call(self._model, X)
+            dummy = fisher_maker.setup_model_call(self._model, x)
             dummy = (
                 dummy if self.loss_type == LOSS_MSE else dummy.view(-1, dummy.size(-1))
             )
             fisher_maker.setup_loss_call(self.lossfunc, dummy, y)
         else:
-            fisher_maker.setup_model_call(self._model, X)
+            fisher_maker.setup_model_call(self._model, x)
         f, _ = fisher_maker.forward_and_backward()
         # Assumes that the last dimension of f is of size outputs.
         f = f if self.loss_type == LOSS_MSE else f.view(-1, f.size(-1))
@@ -191,9 +211,15 @@ def diag(self, X, y, N=None, **kwargs):
             curv_factor = 1.0  # ASDL uses proper 1/2 * MSELoss
         return self.factor * loss, curv_factor * diag_ggn
 
-    def kron(self, X, y, N, **kwargs):
+    def kron(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, Kron]:
         if self.last_layer:
-            _, X = self.model.forward_with_features(X)
+            _, x = self.model.forward_with_features(x)
         cfg = FisherConfig(
             fisher_type=self._ggn_type,
             loss_type=self.loss_type,
@@ -204,13 +230,13 @@ def kron(self, X, y, N, **kwargs):
         fisher_maker = get_fisher_maker(self.model, cfg)
         y = y if self.loss_type == LOSS_MSE else y.view(-1)
         if "emp" in self._ggn_type:
-            dummy = fisher_maker.setup_model_call(self._model, X)
+            dummy = fisher_maker.setup_model_call(self._model, x)
             dummy = (
                 dummy if self.loss_type == LOSS_MSE else dummy.view(-1, dummy.size(-1))
             )
             fisher_maker.setup_loss_call(self.lossfunc, dummy, y)
         else:
-            fisher_maker.setup_model_call(self._model, X)
+            fisher_maker.setup_model_call(self._model, x)
         f, _ = fisher_maker.forward_and_backward()
         # Assumes that the last dimension of f is of size outputs.
         f = f if self.loss_type == LOSS_MSE else f.view(-1, f.size(-1))
@@ -224,7 +250,10 @@ def kron(self, X, y, N, **kwargs):
             curv_factor = 1.0  # ASDL uses proper 1/2 * MSELoss
         return self.factor * loss, curv_factor * kron
 
-    def _get_batch_size(self, x):
+    def _get_batch_size(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+    ) -> int | None:
         """
         ASDL assumes that all leading dimensions are the batch size by default (batch_size = None).
         Here, we want to specify that only the first dimension is the actual batch size.
@@ -239,12 +268,12 @@ def _get_batch_size(self, x):
 class AsdlHessian(AsdlInterface):
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-    ):
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        last_layer: bool = False,
+        dict_key_x: str = "input_ids",
+        dict_key_y: str = "labels",
+    ) -> None:
         super().__init__(
             model,
             likelihood,
@@ -255,24 +284,34 @@ def __init__(
         )
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         raise NotImplementedError()
 
-    def full(self, x, y, **kwargs):
+    def full(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         if self.last_layer:
             _, x = self.model.forward_with_features(x)
+
         cfg = HessianConfig(hessian_shapes=[SHAPE_FULL])
         hess_maker = HessianMaker(self.model, cfg)
+
         dummy = hess_maker.setup_model_call(self._model, x)
         dummy = dummy if self.loss_type == LOSS_MSE else dummy.view(-1, dummy.size(-1))
         y = y if self.loss_type == LOSS_MSE else y.view(-1)
+
         hess_maker.setup_loss_call(self.lossfunc, dummy, y)
         hess_maker.forward_and_backward()
+
         H = self._model.hessian.data
         f = self.model(x).detach()
         # Assumes that the last dimension of f is of size outputs.
         f = f if self.loss_type == LOSS_MSE else f.view(-1, f.size(-1))
         loss = self.lossfunc(f, y)
+
         return self.factor * loss, self.factor * H
 
 
@@ -281,13 +320,13 @@ class AsdlGGN(AsdlInterface, GGNInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        stochastic=False,
+        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,
     ):
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
@@ -295,7 +334,7 @@ def __init__(
         self.stochastic = stochastic
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         return FISHER_MC if self.stochastic else FISHER_EXACT
 
 
@@ -304,14 +343,14 @@ class AsdlEF(AsdlInterface, EFInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        last_layer: bool = False,
+        dict_key_x: str = "input_ids",
+        dict_key_y: str = "labels",
     ):
         super().__init__(model, likelihood, last_layer, None, dict_key_x, dict_key_y)
 
     @property
-    def _ggn_type(self):
+    def _ggn_type(self) -> str:
         return FISHER_EMP
diff --git a/laplace/curvature/backpack.py b/laplace/curvature/backpack.py
index b6c0195..40b03d8 100644
--- a/laplace/curvature/backpack.py
+++ b/laplace/curvature/backpack.py
@@ -1,4 +1,7 @@
-from typing import Tuple
+from __future__ import annotations
+
+from collections.abc import MutableMapping
+from typing import Any
 
 import torch
 from backpack import backpack, extend, memory_cleanup
@@ -11,9 +14,10 @@
     DiagGGNMC,
     SumGradSquared,
 )
+from torch import nn
 
 from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
-from laplace.utils import Kron
+from laplace.utils import Kron, Likelihood
 
 
 class BackPackInterface(CurvatureInterface):
@@ -21,13 +25,13 @@ class BackPackInterface(CurvatureInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-    ):
+        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",
+    ) -> None:
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
         )
@@ -35,7 +39,11 @@ def __init__(
         extend(self._model)
         extend(self.lossfunc)
 
-    def jacobians(self, x, enable_backprop=False):
+    def jacobians(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        enable_backprop: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\)
         using backpack's BatchGrad per output dimension. Note that BackPACK doesn't play well
         with torch.func, so this method has to be overridden.
@@ -54,6 +62,9 @@ def jacobians(self, x, enable_backprop=False):
         f : torch.Tensor
             output function `(batch, outputs)`
         """
+        if isinstance(x, MutableMapping):
+            raise ValueError("BackPACK backend does not support dict-like inputs!")
+
         model = extend(self.model)
         to_stack = []
         for i in range(model.output_size):
@@ -87,7 +98,9 @@ def jacobians(self, x, enable_backprop=False):
         else:
             return Jk.unsqueeze(-1).transpose(1, 2), f
 
-    def gradients(self, x, y):
+    def gradients(
+        self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
         \\(\\theta\\) using Backpack's BatchGrad. Note that BackPACK doesn't play well
         with torch.func, so this method has to be overridden.
@@ -122,20 +135,20 @@ class BackPackGGN(BackPackInterface, GGNInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        stochastic=False,
+        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,
     ):
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
         )
         self.stochastic = stochastic
 
-    def _get_diag_ggn(self):
+    def _get_diag_ggn(self) -> torch.Tensor:
         if self.stochastic:
             return torch.cat(
                 [p.diag_ggn_mc.data.flatten() for p in self._model.parameters()]
@@ -145,14 +158,14 @@ def _get_diag_ggn(self):
                 [p.diag_ggn_exact.data.flatten() for p in self._model.parameters()]
             )
 
-    def _get_kron_factors(self):
+    def _get_kron_factors(self) -> Kron:
         if self.stochastic:
             return Kron([p.kfac for p in self._model.parameters()])
         else:
             return Kron([p.kflr for p in self._model.parameters()])
 
     @staticmethod
-    def _rescale_kron_factors(kron, M, N):
+    def _rescale_kron_factors(kron: Kron, M: int, N: int) -> Kron:
         # Renormalize Kronecker factor to sum up correctly over N data points with batches of M
         # for M=N (full-batch) just M/N=1
         for F in kron.kfacs:
@@ -160,9 +173,14 @@ def _rescale_kron_factors(kron, M, N):
                 F[1] *= M / N
         return kron
 
-    def diag(self, X, y, **kwargs):
+    def diag(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         context = DiagGGNMC if self.stochastic else DiagGGNExact
-        f = self.model(X)
+        f = self.model(x)
         # Assumes that the last dimension of f is of size outputs.
         f = f if self.likelihood == "regression" else f.view(-1, f.size(-1))
         y = y if self.likelihood == "regression" else y.view(-1)
@@ -175,9 +193,15 @@ def diag(self, X, y, **kwargs):
 
         return self.factor * loss.detach(), self.factor * dggn
 
-    def kron(self, X, y, N, **kwargs) -> Tuple[torch.Tensor, Kron]:
+    def kron(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, Kron]:
         context = KFAC if self.stochastic else KFLR
-        f = self.model(X)
+        f = self.model(x)
         # Assumes that the last dimension of f is of size outputs.
         f = f if self.likelihood == "regression" else f.view(-1, f.size(-1))
         y = y if self.likelihood == "regression" else y.view(-1)
@@ -193,8 +217,13 @@ def kron(self, X, y, N, **kwargs) -> Tuple[torch.Tensor, Kron]:
 class BackPackEF(BackPackInterface, EFInterface):
     """Implementation of `EFInterface` using Backpack."""
 
-    def diag(self, X, y, **kwargs):
-        f = self.model(X)
+    def diag(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        f = self.model(x)
         # Assumes that the last dimension of f is of size outputs.
         f = f if self.likelihood == "regression" else f.view(-1, f.size(-1))
         y = y if self.likelihood == "regression" else y.view(-1)
@@ -209,11 +238,17 @@ def diag(self, X, y, **kwargs):
 
         return self.factor * loss.detach(), self.factor * diag_EF
 
-    def kron(self, X, y, **kwargs):
+    def kron(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, Kron]:
         raise NotImplementedError("Unavailable through Backpack.")
 
 
-def _cleanup(module):
+def _cleanup(module: nn.Module) -> None:
     for child in module.children():
         _cleanup(child)
 
diff --git a/laplace/curvature/curvature.py b/laplace/curvature/curvature.py
index 2c3a664..7b5fe60 100644
--- a/laplace/curvature/curvature.py
+++ b/laplace/curvature/curvature.py
@@ -1,6 +1,13 @@
+from __future__ import annotations
+
+from typing import Any, Callable, MutableMapping
+
 import torch
+from torch import nn
 from torch.nn import CrossEntropyLoss, MSELoss
 
+from laplace.utils import Kron, Likelihood
+
 
 class CurvatureInterface:
     """Interface to access curvature for a model and corresponding likelihood.
@@ -16,7 +23,7 @@ class CurvatureInterface:
     likelihood : {'classification', 'regression'}
     last_layer : bool, default=False
         only consider curvature of last layer
-    subnetwork_indices : torch.Tensor, default=None
+    subnetwork_indices : torch.LongTensor, default=None
         indices of the vectorized model parameters that define the subnetwork
         to apply the Laplace approximation over
     dict_key_x: str, default='input_ids'
@@ -38,39 +45,51 @@ class CurvatureInterface:
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
+        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",
     ):
-        assert likelihood in ["regression", "classification"]
-        self.likelihood = likelihood
-        self.model = model
-        self.last_layer = last_layer
-        self.subnetwork_indices = subnetwork_indices
+        assert likelihood in [Likelihood.REGRESSION, Likelihood.CLASSIFICATION]
+        self.likelihood: Likelihood | str = likelihood
+        self.model: nn.Module = model
+        self.last_layer: bool = last_layer
+        self.subnetwork_indices: torch.LongTensor | None = subnetwork_indices
         self.dict_key_x = dict_key_x
         self.dict_key_y = dict_key_y
 
         if likelihood == "regression":
-            self.lossfunc = MSELoss(reduction="sum")
-            self.factor = 0.5
+            self.lossfunc: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = (
+                MSELoss(reduction="sum")
+            )
+            self.factor: float = 0.5
         else:
-            self.lossfunc = CrossEntropyLoss(reduction="sum")
-            self.factor = 1.0
+            self.lossfunc: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = (
+                CrossEntropyLoss(reduction="sum")
+            )
+            self.factor: float = 1.0
 
-        self.params = [p for p in self._model.parameters() if p.requires_grad]
-        self.params_dict = {
+        self.params: list[nn.Parameter] = [
+            p for p in self._model.parameters() if p.requires_grad
+        ]
+        self.params_dict: dict[str, nn.Parameter] = {
             k: v for k, v in self._model.named_parameters() if v.requires_grad
         }
-        self.buffers_dict = {k: v for k, v in self.model.named_buffers()}
+        self.buffers_dict: dict[str, torch.Tensor] = {
+            k: v for k, v in self.model.named_buffers()
+        }
 
     @property
-    def _model(self):
+    def _model(self) -> nn.Module:
         return self.model.last_layer if self.last_layer else self.model
 
-    def jacobians(self, x, enable_backprop=False):
+    def jacobians(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        enable_backprop: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute Jacobians \\(\\nabla_{\\theta} f(x;\\theta)\\) at current parameter \\(\\theta\\),
         via torch.func.
 
@@ -84,25 +103,23 @@ def jacobians(self, x, enable_backprop=False):
         Returns
         -------
         Js : torch.Tensor
-            Jacobians `(batch, outputs, parameters)`
+            Jacobians `(batch, parameters, outputs)`
         f : torch.Tensor
             output function `(batch, outputs)`
         """
-        # Compute Js
-        # ------------------------
-        params = [p for p in self.model.parameters() if p.requires_grad]
-        name_dict = {p.data_ptr(): name for name, p in self.model.named_parameters()}
-        params_dict = {name_dict[p.data_ptr()]: p for p in params}
-
-        def model_fn_params_only(params_dict):
-            res = torch.func.functional_call(self.model, params_dict, x)
-            return res, res
-
-        # concatenate over flattened parameters
-        Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(params_dict)
+
+        def model_fn_params_only(params_dict, buffers_dict):
+            out = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
+            return out, out
+
+        Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
+            self.params_dict, self.buffers_dict
+        )
+
+        # Concatenate over flattened parameters
         Js = [
             j.flatten(start_dim=-p.dim())
-            for j, p in zip(Js.values(), params_dict.values())
+            for j, p in zip(Js.values(), self.params_dict.values())
         ]
         Js = torch.cat(Js, dim=-1)
 
@@ -111,7 +128,11 @@ def model_fn_params_only(params_dict):
 
         return (Js, f) if enable_backprop else (Js.detach(), f.detach())
 
-    def last_layer_jacobians(self, x, enable_backprop=False):
+    def last_layer_jacobians(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        enable_backprop: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
         only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).
 
@@ -144,7 +165,9 @@ def last_layer_jacobians(self, x, enable_backprop=False):
 
         return Js, f
 
-    def gradients(self, x, y):
+    def gradients(
+        self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute batch gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at
         current parameter \\(\\theta\\).
 
@@ -185,7 +208,12 @@ def loss_single(x, y, params_dict, buffers_dict):
 
         return Gs, loss
 
-    def full(self, x, y, **kwargs):
+    def full(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ):
         """Compute a dense curvature (approximation) in the form of a \\(P \\times P\\) matrix
         \\(H\\) with respect to parameters \\(\\theta \\in \\mathbb{R}^P\\).
 
@@ -204,7 +232,13 @@ def full(self, x, y, **kwargs):
         """
         raise NotImplementedError
 
-    def kron(self, x, y, **kwargs):
+    def kron(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, Kron]:
         """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
@@ -218,6 +252,8 @@ def kron(self, x, y, **kwargs):
             input data `(batch, input_shape)`
         y : torch.Tensor
             labels `(batch, label_shape)`
+        N : int
+            total number of data points
 
         Returns
         -------
@@ -227,7 +263,12 @@ def kron(self, x, y, **kwargs):
         """
         raise NotImplementedError
 
-    def diag(self, x, y, **kwargs):
+    def diag(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ):
         """Compute a diagonal Hessian approximation to \\(H\\) and is represented as a
         vector of the dimensionality of parameters \\(\\theta\\).
 
@@ -278,22 +319,23 @@ class GGNInterface(CurvatureInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        stochastic=False,
-        num_samples=1,
-    ):
-        self.stochastic = stochastic
-        self.num_samples = num_samples
+        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,
+        num_samples: int = 1,
+    ) -> None:
+        self.stochastic: bool = stochastic
+        self.num_samples: int = num_samples
+
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
         )
 
-    def _get_mc_functional_fisher(self, f):
+    def _get_mc_functional_fisher(self, f: torch.Tensor) -> torch.Tensor:
         """Approximate the Fisher's middle matrix (expected outer product of the functional gradient)
         using MC integral with `self.num_samples` many samples.
         """
@@ -317,7 +359,7 @@ def _get_mc_functional_fisher(self, f):
 
         return F
 
-    def _get_functional_hessian(self, f):
+    def _get_functional_hessian(self, f: torch.Tensor) -> torch.Tensor | None:
         if self.likelihood == "regression":
             return None
         else:
@@ -326,7 +368,12 @@ def _get_functional_hessian(self, f):
             G = torch.diag_embed(ps) - torch.einsum("mk,mc->mck", ps, ps)
             return G
 
-    def full(self, x, y, **kwargs):
+    def full(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """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}\\)
@@ -359,8 +406,13 @@ def full(self, x, y, **kwargs):
 
         return loss.detach(), H.detach()
 
-    def diag(self, X, y, **kwargs):
-        Js, f = self.last_layer_jacobians(X) if self.last_layer else self.jacobians(X)
+    def diag(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        Js, f = self.last_layer_jacobians(x) if self.last_layer else self.jacobians(x)
         loss = self.factor * self.lossfunc(f, y)
 
         H_lik = (
@@ -408,7 +460,12 @@ class EFInterface(CurvatureInterface):
         For example, \\(\\frac{1}{2}\\) to get to \\(\\mathcal{N}(f, 1)\\) from MSELoss.
     """
 
-    def full(self, x, y, **kwargs):
+    def full(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """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}\\)
@@ -431,9 +488,14 @@ def full(self, x, y, **kwargs):
         H_ef = torch.einsum("bp,bq->pq", Gs, Gs)
         return self.factor * loss.detach(), self.factor * H_ef
 
-    def diag(self, X, y, **kwargs):
+    def diag(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # Gs is (batchsize, n_params)
-        Gs, loss = self.gradients(X, y)
+        Gs, loss = self.gradients(x, y)
         Gs, loss = Gs.detach(), loss.detach()
         diag_ef = torch.einsum("bp,bp->p", Gs, Gs)
         return self.factor * loss, self.factor * diag_ef
diff --git a/laplace/curvature/curvlinops.py b/laplace/curvature/curvlinops.py
index 7801dd6..e47db0d 100644
--- a/laplace/curvature/curvlinops.py
+++ b/laplace/curvature/curvlinops.py
@@ -1,4 +1,7 @@
-from collections import UserDict
+from __future__ import annotations
+
+from collections.abc import MutableMapping
+from typing import Any
 
 import torch
 from curvlinops import (
@@ -9,9 +12,11 @@
     HessianLinearOperator,
     KFACLinearOperator,
 )
+from curvlinops._base import _LinearOperator
+from torch import nn
 
 from laplace.curvature import CurvatureInterface, EFInterface, GGNInterface
-from laplace.utils import Kron
+from laplace.utils import Kron, Likelihood
 
 
 class CurvlinopsInterface(CurvatureInterface):
@@ -19,35 +24,35 @@ class CurvlinopsInterface(CurvatureInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-    ):
+        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",
+    ) -> None:
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
         )
 
     @property
-    def _kron_fisher_type(self):
+    def _kron_fisher_type(self) -> str:
         raise NotImplementedError
 
     @property
-    def _linop_context(self):
+    def _linop_context(self) -> type[_LinearOperator]:
         raise NotImplementedError
 
     @staticmethod
-    def _rescale_kron_factors(kron, M, N):
-        # Renormalize Kronecker factor to sum up correctly over N data points with batches of M
-        # for M=N (full-batch) just M/N=1
+    def _rescale_kron_factors(kron: Kron, M: int, N: int) -> Kron:
+        # Renormalize Kronecker factor to sum up correctly over N data points with
+        # batches of M. For M=N (full-batch) just M/N=1
         for F in kron.kfacs:
             if len(F) == 2:
                 F[1] *= M / N
         return kron
 
-    def _get_kron_factors(self, linop):
+    def _get_kron_factors(self, linop: KFACLinearOperator) -> Kron:
         kfacs = list()
         for name, module in self.model.named_modules():
             if name not in linop._mapping.keys():
@@ -69,15 +74,21 @@ def _get_kron_factors(self, linop):
                 raise ValueError(f"Whats happening with {module}?")
         return Kron(kfacs)
 
-    def kron(self, X, y, N, **kwargs):
-        if isinstance(X, (dict, UserDict)):
+    def kron(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        N: int,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, Kron]:
+        if isinstance(x, (dict, MutableMapping)):
             kwargs["batch_size_fn"] = lambda x: x[self.dict_key_x].shape[0]
 
         linop = KFACLinearOperator(
             self.model,
             self.lossfunc,
             self.params,
-            [(X, y)],
+            [(x, y)],
             fisher_type=self._kron_fisher_type,
             separate_weight_and_bias=True,
             check_deterministic=False,  # To avoid overhead
@@ -92,24 +103,29 @@ def kron(self, X, y, N, **kwargs):
         kron = self._rescale_kron_factors(kron, len(y), N)
         kron *= self.factor
 
-        loss = self.lossfunc(self.model(X), y)
+        loss = self.lossfunc(self.model(x), y)
 
         return self.factor * loss.detach(), kron
 
-    def full(self, X, y, **kwargs):
+    def full(
+        self,
+        x: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        y: torch.Tensor,
+        **kwargs: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # Fallback to torch.func backend for SubnetLaplace
         if self.subnetwork_indices is not None:
-            return super().full(X, y, **kwargs)
+            return super().full(x, y, **kwargs)
 
         curvlinops_kwargs = {k: v for k, v in kwargs.items() if k != "N"}
-        if isinstance(X, (dict, UserDict)):
+        if isinstance(x, (dict, MutableMapping)):
             curvlinops_kwargs["batch_size_fn"] = lambda x: x[self.dict_key_x].shape[0]
 
         linop = self._linop_context(
             self.model,
             self.lossfunc,
             self.params,
-            [(X, y)],
+            [(x, y)],
             check_deterministic=False,
             **curvlinops_kwargs,
         )
@@ -118,7 +134,7 @@ def full(self, X, y, **kwargs):
             device=next(self.model.parameters()).device,
         )
 
-        f = self.model(X)
+        f = self.model(x)
         loss = self.lossfunc(f, y)
 
         return self.factor * loss.detach(), self.factor * H
@@ -129,25 +145,25 @@ class CurvlinopsGGN(CurvlinopsInterface, GGNInterface):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        last_layer=False,
-        subnetwork_indices=None,
-        dict_key_x="input_ids",
-        dict_key_y="labels",
-        stochastic=False,
-    ):
+        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,
+    ) -> None:
         super().__init__(
             model, likelihood, last_layer, subnetwork_indices, dict_key_x, dict_key_y
         )
         self.stochastic = stochastic
 
     @property
-    def _kron_fisher_type(self):
+    def _kron_fisher_type(self) -> FisherType:
         return FisherType.MC if self.stochastic else FisherType.TYPE2
 
     @property
-    def _linop_context(self):
+    def _linop_context(self) -> type[_LinearOperator]:
         return FisherMCLinearOperator if self.stochastic else GGNLinearOperator
 
 
@@ -155,11 +171,11 @@ class CurvlinopsEF(CurvlinopsInterface, EFInterface):
     """Implementation of `EFInterface` using Curvlinops."""
 
     @property
-    def _kron_fisher_type(self):
+    def _kron_fisher_type(self) -> FisherType:
         return FisherType.EMPIRICAL
 
     @property
-    def _linop_context(self):
+    def _linop_context(self) -> type[_LinearOperator]:
         return EFLinearOperator
 
 
@@ -167,5 +183,5 @@ class CurvlinopsHessian(CurvlinopsInterface):
     """Implementation of the full Hessian using Curvlinops."""
 
     @property
-    def _linop_context(self):
+    def _linop_context(self) -> type[_LinearOperator]:
         return HessianLinearOperator
diff --git a/laplace/laplace.py b/laplace/laplace.py
index 6a0d598..4308961 100644
--- a/laplace/laplace.py
+++ b/laplace/laplace.py
@@ -1,24 +1,32 @@
-from laplace import *  # noqa F403
+from __future__ import annotations
+
+import torch
+
 from laplace.baselaplace import ParametricLaplace
+from laplace.utils.enums import (
+    HessianStructure,
+    Likelihood,
+    SubsetOfWeights,
+)
 
 
 def Laplace(
-    model,
-    likelihood,
-    subset_of_weights="last_layer",
-    hessian_structure="kron",
+    model: torch.nn.Module,
+    likelihood: Likelihood | str,
+    subset_of_weights: SubsetOfWeights | str = SubsetOfWeights.LAST_LAYER,
+    hessian_structure: HessianStructure | str = HessianStructure.KRON,
     *args,
     **kwargs,
-):
+) -> ParametricLaplace:
     """Simplified Laplace access using strings instead of different classes.
 
     Parameters
     ----------
     model : torch.nn.Module
-    likelihood : {'classification', 'regression'}
-    subset_of_weights : {'last_layer', 'subnetwork', 'all'}, default='last_layer'
+    likelihood : Likelihood or str in {'classification', 'regression'}
+    subset_of_weights : SubsetofWeights or {'last_layer', 'subnetwork', 'all'}, default=SubsetOfWeights.LAST_LAYER
         subset of weights to consider for inference
-    hessian_structure : {'diag', 'kron', 'full', 'lowrank'}, default='kron'
+    hessian_structure : HessianStructure or str in {'diag', 'kron', 'full', 'lowrank'}, default=HessianStructure.KRON
         structure of the Hessian approximation
 
     Returns
@@ -40,7 +48,7 @@ def Laplace(
     return laplace_class(model, likelihood, *args, **kwargs)
 
 
-def _all_subclasses(cls):
+def _all_subclasses(cls) -> set:
     return set(cls.__subclasses__()).union(
         [s for c in cls.__subclasses__() for s in _all_subclasses(c)]
     )
diff --git a/laplace/lllaplace.py b/laplace/lllaplace.py
index 2bd87c8..bd7ddce 100644
--- a/laplace/lllaplace.py
+++ b/laplace/lllaplace.py
@@ -1,12 +1,24 @@
+from __future__ import annotations
+
 from collections.abc import MutableMapping
 from copy import deepcopy
-from typing import Union
+from typing import Any
 
 import torch
+from torch import nn
 from torch.nn.utils import parameters_to_vector, vector_to_parameters
-
-from laplace.baselaplace import DiagLaplace, FullLaplace, KronLaplace, ParametricLaplace
+from torch.utils.data import DataLoader
+
+from laplace.baselaplace import (
+    DiagLaplace,
+    FullLaplace,
+    KronLaplace,
+    Likelihood,
+    ParametricLaplace,
+)
+from laplace.curvature.curvature import CurvatureInterface
 from laplace.utils import FeatureExtractor, Kron
+from laplace.utils.feature_extractor import FeatureReduction
 
 __all__ = ["LLLaplace", "FullLLLaplace", "KronLLLaplace", "DiagLLLaplace"]
 
@@ -36,7 +48,7 @@ class LLLaplace(ParametricLaplace):
     Parameters
     ----------
     model : torch.nn.Module or `laplace.utils.feature_extractor.FeatureExtractor`
-    likelihood : {'classification', 'regression'}
+    likelihood : Likelihood or {'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
@@ -59,7 +71,6 @@ class LLLaplace(ParametricLaplace):
         Note: Make sure that this option faithfully reflects the reduction in the model
         definition. When inputting a string, available options are
         `{'pick_first', 'pick_last', 'average'}`.
-
     dict_key_x: str, default='input_ids'
         The dictionary key under which the input tensor `x` is stored. Only has effect
         when the model takes a `MutableMapping` as the input. Useful for Huggingface
@@ -79,20 +90,20 @@ class LLLaplace(ParametricLaplace):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        sigma_noise=1.0,
-        prior_precision=1.0,
-        prior_mean=0.0,
-        temperature=1.0,
-        enable_backprop=False,
-        feature_reduction=None,
-        dict_key_x="inputs_id",
-        dict_key_y="labels",
-        backend=None,
-        last_layer_name=None,
-        backend_kwargs=None,
-        asdl_fisher_kwargs=None,
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        sigma_noise: float | torch.Tensor = 1.0,
+        prior_precision: float | torch.Tensor = 1.0,
+        prior_mean: float | torch.Tensor = 0.0,
+        temperature: float = 1.0,
+        enable_backprop: bool = False,
+        feature_reduction: FeatureReduction | str | None = None,
+        dict_key_x: str = "inputs_id",
+        dict_key_y: str = "labels",
+        backend: type[CurvatureInterface] | None = None,
+        last_layer_name: str | None = None,
+        backend_kwargs: dict[str, Any] | None = None,
+        asdl_fisher_kwargs: dict[str, Any] | None = None,
     ):
         if asdl_fisher_kwargs is not None:
             raise ValueError("Last-layer Laplace does not support asdl_fisher_kwargs.")
@@ -117,26 +128,33 @@ def __init__(
             enable_backprop=enable_backprop,
             feature_reduction=feature_reduction,
         )
+
         if self.model.last_layer is None:
-            self.mean = None
-            self.n_params = None
-            self.n_layers = None
+            self.mean: torch.Tensor | None = None
+            self.n_params: int | None = None
+            self.n_layers: int | None = None
             # ignore checks of prior mean setter temporarily, check on .fit()
-            self._prior_precision = prior_precision
-            self._prior_mean = prior_mean
+            self._prior_precision: float | torch.Tensor = prior_precision
+            self._prior_mean: float | torch.Tensor = prior_mean
         else:
-            self.n_params = len(
+            self.n_params: int = len(
                 parameters_to_vector(self.model.last_layer.parameters())
             )
-            self.n_layers = len(list(self.model.last_layer.parameters()))
-            self.prior_precision = prior_precision
-            self.prior_mean = prior_mean
-            self.mean = self.prior_mean
+            self.n_layers: int | None = len(list(self.model.last_layer.parameters()))
+            self.prior_precision: float | torch.Tensor = prior_precision
+            self.prior_mean: float | torch.Tensor = prior_mean
+            self.mean: float | torch.Tensor = self.prior_mean
             self._init_H()
+
         self._backend_kwargs["last_layer"] = True
-        self._last_layer_name = last_layer_name
+        self._last_layer_name: str | None = last_layer_name
 
-    def fit(self, train_loader, override=True):
+    def fit(
+        self,
+        train_loader: DataLoader,
+        override: bool = True,
+        progress_bar: bool = False,
+    ) -> None:
         """Fit the local Laplace approximation at the parameters of the model.
 
         Parameters
@@ -147,6 +165,7 @@ def fit(self, train_loader, override=True):
         override : bool, default=True
             whether to initialize H, loss, and n_data again; setting to False is useful for
             online learning settings to accumulate a sequential posterior approximation.
+        progress_bar: bool, default=False
         """
         if not override:
             raise ValueError(
@@ -156,23 +175,31 @@ def fit(self, train_loader, override=True):
         self.model.eval()
 
         if self.model.last_layer is None:
-            self.data = next(iter(train_loader))
+            self.data: tuple[torch.Tensor, torch.Tensor] | MutableMapping = next(
+                iter(train_loader)
+            )
             self._find_last_layer(self.data)
-            params = parameters_to_vector(self.model.last_layer.parameters()).detach()
-            self.n_params = len(params)
-            self.n_layers = len(list(self.model.last_layer.parameters()))
+            params: torch.Tensor = parameters_to_vector(
+                self.model.last_layer.parameters()
+            ).detach()
+            self.n_params: int = len(params)
+            self.n_layers: int = len(list(self.model.last_layer.parameters()))
             # here, check the already set prior precision again
-            self.prior_precision = self._prior_precision
-            self.prior_mean = self._prior_mean
+            self.prior_precision: float | torch.Tensor = self._prior_precision
+            self.prior_mean: float | torch.Tensor = self._prior_mean
             self._init_H()
 
         super().fit(train_loader, override=override)
-        self.mean = parameters_to_vector(self.model.last_layer.parameters())
+        self.mean: torch.Tensor = parameters_to_vector(
+            self.model.last_layer.parameters()
+        )
 
         if not self.enable_backprop:
             self.mean = self.mean.detach()
 
-    def _glm_predictive_distribution(self, X, joint=False):
+    def _glm_predictive_distribution(
+        self, X: torch.Tensor | MutableMapping, joint: bool = False
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         Js, f_mu = self.backend.last_layer_jacobians(X)
 
         if joint:
@@ -187,7 +214,13 @@ def _glm_predictive_distribution(self, X, joint=False):
             else (f_mu, f_var)
         )
 
-    def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwargs):
+    def _nn_predictive_samples(
+        self,
+        X: torch.Tensor | MutableMapping[str, torch.Tensor | Any],
+        n_samples: int = 100,
+        generator: torch.Generator | None = None,
+        **model_kwargs,
+    ) -> torch.Tensor:
         fs = list()
 
         feats = None
@@ -208,14 +241,18 @@ def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwarg
         vector_to_parameters(self.mean, self.model.last_layer.parameters())
         fs = torch.stack(fs)
 
-        if self.likelihood == "classification":
+        if self.likelihood == Likelihood.CLASSIFICATION:
             fs = torch.softmax(fs, dim=-1)
 
         return fs
 
     def _nn_predictive_classification(
-        self, X, n_samples=100, generator=None, **model_kwargs
-    ):
+        self,
+        X: torch.Tensor | MutableMapping,
+        n_samples: int = 100,
+        generator: torch.Generator | None = None,
+        **model_kwargs,
+    ) -> torch.Tensor:
         py = 0
 
         feats = None
@@ -234,10 +271,11 @@ def _nn_predictive_classification(
             py += torch.softmax(logits.detach(), dim=-1) / n_samples
 
         vector_to_parameters(self.mean, self.model.last_layer.parameters())
+
         return py
 
     @property
-    def prior_precision_diag(self):
+    def prior_precision_diag(self) -> torch.Tensor:
         """Obtain the diagonal prior precision \\(p_0\\) constructed from either
         a scalar or diagonal prior precision.
 
@@ -245,22 +283,22 @@ def prior_precision_diag(self):
         -------
         prior_precision_diag : torch.Tensor
         """
-        if len(self.prior_precision) == 1:  # scalar
+        if (
+            isinstance(self.prior_precision, float) or len(self.prior_precision) == 1
+        ):  # scalar
             return self.prior_precision * torch.ones_like(self.mean)
-
         elif len(self.prior_precision) == self.n_params:  # diagonal
             return self.prior_precision
-
         else:
             raise ValueError("Mismatch of prior and model. Diagonal or scalar prior.")
 
-    def state_dict(self) -> dict:
+    def state_dict(self) -> dict[str, Any]:
         state_dict = super().state_dict()
         state_dict["data"] = getattr(self, "data", None)  # None if not present
         state_dict["_last_layer_name"] = self._last_layer_name
         return state_dict
 
-    def load_state_dict(self, state_dict: dict):
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         if self._last_layer_name != state_dict["_last_layer_name"]:
             raise ValueError("Different `last_layer_name` detected!")
 
@@ -275,12 +313,15 @@ def load_state_dict(self, state_dict: dict):
         self.n_layers = len(list(self.model.last_layer.parameters()))
 
     @torch.no_grad()
-    def _find_last_layer(self, data: Union[torch.Tensor, MutableMapping]) -> None:
+    def _find_last_layer(
+        self, data: torch.Tensor | MutableMapping[str, torch.Tensor | Any]
+    ) -> None:
         # To support Huggingface dataset
         if isinstance(data, MutableMapping):
             self.model.find_last_layer(data)
         else:
             X = data[0]
+
             try:
                 self.model.find_last_layer(X[:1].to(self._device))
             except (TypeError, AttributeError):
@@ -317,20 +358,21 @@ class KronLLLaplace(LLLaplace, KronLaplace):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        sigma_noise=1.0,
-        prior_precision=1.0,
-        prior_mean=0.0,
-        temperature=1.0,
-        enable_backprop=False,
-        feature_reduction=None,
-        dict_key_x="inputs_id",
-        dict_key_y="labels",
-        backend=None,
-        last_layer_name=None,
-        damping=False,
-        **backend_kwargs,
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        sigma_noise: float | torch.Tensor = 1.0,
+        prior_precision: float | torch.Tensor = 1.0,
+        prior_mean: float | torch.Tensor = 0.0,
+        temperature: float = 1.0,
+        enable_backprop: bool = False,
+        feature_reduction: FeatureReduction | str | None = None,
+        dict_key_x: str = "inputs_id",
+        dict_key_y: str = "labels",
+        backend: type[CurvatureInterface] | None = None,
+        last_layer_name: str | None = None,
+        damping: bool = False,
+        backend_kwargs: dict[str, Any] | None = None,
+        asdl_fisher_kwargs: dict[str, Any] | None = None,
     ):
         self.damping = damping
         super().__init__(
@@ -347,9 +389,10 @@ def __init__(
             backend,
             last_layer_name,
             backend_kwargs,
+            asdl_fisher_kwargs,
         )
 
-    def _init_H(self):
+    def _init_H(self) -> None:
         self.H = Kron.init_from_model(self.model.last_layer, self._device)
 
 
diff --git a/laplace/marglik_training.py b/laplace/marglik_training.py
index 9e8e2df..92b640f 100644
--- a/laplace/marglik_training.py
+++ b/laplace/marglik_training.py
@@ -1,45 +1,61 @@
+from __future__ import annotations
+
 import logging
 import warnings
 from collections.abc import MutableMapping
 from copy import deepcopy
+from typing import Type
 
 import numpy as np
 import torch
 import tqdm
+from torch import nn
 from torch.nn import CrossEntropyLoss, MSELoss
 from torch.nn.utils import parameters_to_vector
-from torch.optim import Adam
+from torch.optim import Adam, Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from torch.types import Number
+from torch.utils.data import DataLoader
 
 from laplace import Laplace
+from laplace.baselaplace import BaseLaplace
 from laplace.curvature import AsdlGGN
-from laplace.utils import expand_prior_precision, fix_prior_prec_structure
+from laplace.curvature.curvature import CurvatureInterface
+from laplace.utils import (
+    HessianStructure,
+    Likelihood,
+    PriorStructure,
+    SubsetOfWeights,
+    expand_prior_precision,
+    fix_prior_prec_structure,
+)
 
 
 def marglik_training(
-    model,
-    train_loader,
-    likelihood="classification",
-    hessian_structure="kron",
-    backend=AsdlGGN,
-    optimizer_cls=Adam,
-    optimizer_kwargs=None,
-    scheduler_cls=None,
-    scheduler_kwargs=None,
-    n_epochs=300,
-    lr_hyp=1e-1,
-    prior_structure="layerwise",
-    n_epochs_burnin=0,
-    n_hypersteps=10,
-    marglik_frequency=1,
-    prior_prec_init=1.0,
-    sigma_noise_init=1.0,
-    temperature=1.0,
-    fix_sigma_noise=False,
-    progress_bar=False,
-    enable_backprop=False,
-    dict_key_x="input_ids",
-    dict_key_y="labels",
-):
+    model: torch.nn.Module,
+    train_loader: DataLoader,
+    likelihood: Likelihood | str = Likelihood.CLASSIFICATION,
+    hessian_structure: HessianStructure | str = HessianStructure.KRON,
+    backend: Type[CurvatureInterface] = AsdlGGN,
+    optimizer_cls: Type[Optimizer] = Adam,
+    optimizer_kwargs: dict | None = None,
+    scheduler_cls: Type[LRScheduler] | None = None,
+    scheduler_kwargs: dict | None = None,
+    n_epochs: int = 300,
+    lr_hyp: float = 1e-1,
+    prior_structure: PriorStructure | str = PriorStructure.LAYERWISE,
+    n_epochs_burnin: int = 0,
+    n_hypersteps: int = 10,
+    marglik_frequency: int = 1,
+    prior_prec_init: float = 1.0,
+    sigma_noise_init: float = 1.0,
+    temperature: float = 1.0,
+    fix_sigma_noise: bool = False,
+    progress_bar: bool = False,
+    enable_backprop: bool = False,
+    dict_key_x: str = "input_ids",
+    dict_key_y: str = "labels",
+) -> tuple[BaseLaplace, nn.Module, list[Number], list[Number]]:
     """Marginal-likelihood based training (Algorithm 1 in [1]).
     Optimize model parameters and hyperparameters jointly.
     Model parameters are optimized to minimize negative log joint (train loss)
@@ -51,7 +67,7 @@ def marglik_training(
     The settings of standard training can be controlled by passing `train_loader`,
     `optimizer_cls`, `optimizer_kwargs`, `scheduler_cls`, `scheduler_kwargs`, and `n_epochs`.
     The `model` should return logits, i.e., no softmax should be applied.
-    With `likelihood='classification'` or `'regression'`, one can choose between
+    With `likelihood=Likelihood.CLASSIFICATION` or `Likelihood.REGRESSION`, one can choose between
     categorical likelihood (CrossEntropyLoss) and Gaussian likelihood (MSELoss).
 
     As in [1], we optimize prior precision and, for regression, observation noise
@@ -76,8 +92,8 @@ def marglik_training(
         torch neural network model (needs to comply with Backend choice)
     train_loader : DataLoader
         pytorch dataloader that implements `len(train_loader.dataset)` to obtain number of data points
-    likelihood : str, default='classification'
-        'classification' or 'regression'
+    likelihood : str, default=Likelihood.CLASSIFICATION
+        Likelihood.CLASSIFICATION or Likelihood.REGRESSION
     hessian_structure : {'diag', 'kron', 'full'}, default='kron'
         structure of the Hessian approximation
     backend : Backend, default=AsdlGGN
@@ -138,7 +154,7 @@ def marglik_training(
     losses : list
         list of losses (log joints) obtained during training (to monitor convergence)
     """
-    if "weight_decay" in optimizer_kwargs:
+    if optimizer_kwargs is not None and "weight_decay" in optimizer_kwargs:
         warnings.warn("Weight decay is handled and optimized. Will be set to 0.")
         optimizer_kwargs["weight_decay"] = 0.0
 
@@ -160,10 +176,10 @@ def marglik_training(
     hyperparameters.append(log_prior_prec)
 
     # set up loss (and observation noise hyperparam)
-    if likelihood == "classification":
+    if likelihood == Likelihood.CLASSIFICATION:
         criterion = CrossEntropyLoss(reduction="mean")
         sigma_noise = 1.0
-    elif likelihood == "regression":
+    elif likelihood == 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)
@@ -176,6 +192,7 @@ def marglik_training(
     optimizer = optimizer_cls(model.parameters(), **optimizer_kwargs)
 
     # set up learning rate scheduler
+    scheduler = None
     if scheduler_cls is not None:
         if scheduler_kwargs is None:
             scheduler_kwargs = dict()
@@ -199,6 +216,7 @@ def marglik_training(
         desc="[Training]",
         colour="blue",
     )
+
     for epoch in pbar:
         epoch_loss = 0
         epoch_perf = 0
@@ -211,8 +229,10 @@ def marglik_training(
             else:
                 X, y = data
                 X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
+
             optimizer.zero_grad()
-            if likelihood == "regression":
+
+            if likelihood == Likelihood.REGRESSION:
                 sigma_noise = (
                     torch.exp(log_sigma_noise).detach()
                     if not fix_sigma_noise
@@ -221,21 +241,25 @@ def marglik_training(
                 crit_factor = temperature / (2 * sigma_noise**2)
             else:
                 crit_factor = temperature
+
             prior_prec = torch.exp(log_prior_prec).detach()
             theta = parameters_to_vector(
                 [p for p in model.parameters() if p.requires_grad]
             )
             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(y)
-            if likelihood == "regression":
+
+            if likelihood == Likelihood.REGRESSION:
                 epoch_perf += (f.detach() - y).square().sum()
             else:
                 epoch_perf += torch.sum(torch.argmax(f.detach(), dim=-1) == y).item()
-            if scheduler_cls is not None:
+
+            if scheduler is not None:
                 scheduler.step()
 
         losses.append(epoch_loss / N)
@@ -252,7 +276,7 @@ def marglik_training(
 
         # optimizer hyperparameters by differentiating marglik
         # 1. fit laplace approximation
-        if likelihood == "classification":
+        if likelihood == Likelihood.CLASSIFICATION:
             sigma_noise = 1
         else:
             sigma_noise = (
@@ -276,7 +300,7 @@ def marglik_training(
         # 2. differentiate wrt. hyperparameters for n_hypersteps
         for _ in range(n_hypersteps):
             hyper_optimizer.zero_grad()
-            if likelihood == "classification" or fix_sigma_noise:
+            if likelihood == Likelihood.CLASSIFICATION or fix_sigma_noise:
                 sigma_noise = None
             else:
                 sigma_noise = torch.exp(log_sigma_noise)
@@ -290,7 +314,7 @@ def marglik_training(
         if margliks[-1] < best_marglik:
             best_model_dict = deepcopy(model.state_dict())
             best_precision = deepcopy(prior_prec.detach())
-            if likelihood == "classification":
+            if likelihood == Likelihood.CLASSIFICATION:
                 best_sigma = 1
             else:
                 best_sigma = (
@@ -310,10 +334,12 @@ def marglik_training(
             )
 
     logging.info("MARGLIK: finished training. Recover best model and fit Laplace.")
+
     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,
@@ -322,7 +348,7 @@ def marglik_training(
         prior_precision=prior_prec,
         temperature=temperature,
         backend=backend,
-        subset_of_weights="all",
+        subset_of_weights=SubsetOfWeights.ALL,
         enable_backprop=enable_backprop,
         dict_key_x=dict_key_x,
         dict_key_y=dict_key_y,
diff --git a/laplace/subnetlaplace.py b/laplace/subnetlaplace.py
index a074095..b801037 100644
--- a/laplace/subnetlaplace.py
+++ b/laplace/subnetlaplace.py
@@ -1,8 +1,14 @@
+from __future__ import annotations
+
+from typing import Type
+
 import torch
+from torch import nn
 from torch.distributions import MultivariateNormal
 
-from laplace.baselaplace import DiagLaplace, FullLaplace, ParametricLaplace
+from laplace.baselaplace import DiagLaplace, FullLaplace, Likelihood, ParametricLaplace
 from laplace.curvature import EFInterface, GGNInterface
+from laplace.curvature.curvature import CurvatureInterface
 
 __all__ = ["SubnetLaplace", "FullSubnetLaplace", "DiagSubnetLaplace"]
 
@@ -68,19 +74,20 @@ class SubnetLaplace(ParametricLaplace):
 
     def __init__(
         self,
-        model,
-        likelihood,
-        subnetwork_indices,
-        sigma_noise=1.0,
-        prior_precision=1.0,
-        prior_mean=0.0,
-        temperature=1.0,
-        backend=None,
-        backend_kwargs=None,
-        asdl_fisher_kwargs=None,
-    ):
+        model: nn.Module,
+        likelihood: Likelihood | str,
+        subnetwork_indices: torch.LongTensor,
+        sigma_noise: float | torch.Tensor = 1.0,
+        prior_precision: float | torch.Tensor = 1.0,
+        prior_mean: float | torch.Tensor = 0.0,
+        temperature: float = 1.0,
+        backend: Type[CurvatureInterface] | None = None,
+        backend_kwargs: dict | None = None,
+        asdl_fisher_kwargs: dict | None = None,
+    ) -> None:
         if asdl_fisher_kwargs is not None:
             raise ValueError("Subnetwork Laplace does not support asdl_fisher_kwargs.")
+
         self.H = None
         super().__init__(
             model,
@@ -92,28 +99,29 @@ def __init__(
             backend=backend,
             backend_kwargs=backend_kwargs,
         )
+
         if backend is not None:
             if not isinstance(backend, GGNInterface) and not isinstance(
                 backend, EFInterface
             ):
                 raise ValueError("SubnetLaplace can only be used with GGN and EF.")
+
         # check validity of subnetwork indices and pass them to backend
         self._check_subnetwork_indices(subnetwork_indices)
         self.backend.subnetwork_indices = subnetwork_indices
         self.n_params_subnet = len(subnetwork_indices)
         self._init_H()
 
-    def _check_subnetwork_indices(self, subnetwork_indices):
+    def _check_subnetwork_indices(
+        self, subnetwork_indices: torch.LongTensor | None
+    ) -> None:
         """Check that subnetwork indices are valid indices of the vectorized model parameters
         (i.e. `torch.nn.utils.parameters_to_vector(model.parameters())`).
         """
         if subnetwork_indices is None:
             raise ValueError("Subnetwork indices cannot be None.")
         elif not (
-            (
-                isinstance(subnetwork_indices, torch.LongTensor)
-                or isinstance(subnetwork_indices, torch.cuda.LongTensor)
-            )
+            isinstance(subnetwork_indices, torch.LongTensor)
             and subnetwork_indices.numel() > 0
             and len(subnetwork_indices.shape) == 1
         ):
@@ -131,7 +139,7 @@ def _check_subnetwork_indices(self, subnetwork_indices):
             raise ValueError("Subnetwork indices must not contain duplicate entries.")
 
     @property
-    def prior_precision_diag(self):
+    def prior_precision_diag(self) -> torch.Tensor:
         """Obtain the diagonal prior precision \\(p_0\\) constructed from either
         a scalar or diagonal prior precision.
 
@@ -139,27 +147,27 @@ def prior_precision_diag(self):
         -------
         prior_precision_diag : torch.Tensor
         """
-        if len(self.prior_precision) == 1:  # scalar
+        # scalar
+        if isinstance(self.prior_precision, float) or len(self.prior_precision) == 1:
             return self.prior_precision * torch.ones(
                 self.n_params_subnet, device=self._device
             )
 
         elif len(self.prior_precision) == self.n_params_subnet:  # diagonal
             return self.prior_precision
-
         else:
             raise ValueError("Mismatch of prior and model. Diagonal or scalar prior.")
 
     @property
-    def mean_subnet(self):
+    def mean_subnet(self) -> torch.Tensor:
         return self.mean[self.backend.subnetwork_indices]
 
     @property
-    def scatter(self):
+    def scatter(self) -> torch.Tensor:
         delta = self.mean_subnet - self.prior_mean
         return (delta * self.prior_precision_diag) @ delta
 
-    def assemble_full_samples(self, subnet_samples):
+    def assemble_full_samples(self, subnet_samples) -> torch.Tensor:
         full_samples = self.mean.repeat(subnet_samples.shape[0], 1)
         full_samples[:, self.backend.subnetwork_indices] = subnet_samples
         return full_samples
@@ -176,12 +184,14 @@ class FullSubnetLaplace(SubnetLaplace, FullLaplace):
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ("subnetwork", "full")
 
-    def _init_H(self):
+    def _init_H(self) -> None:
         self.H = torch.zeros(
             self.n_params_subnet, self.n_params_subnet, device=self._device
         )
 
-    def sample(self, n_samples=100, generator=None):
+    def sample(
+        self, n_samples: int = 100, generator: torch.Generator | None = None
+    ) -> torch.Tensor:
         # sample only subnetwork parameters and set all other parameters to their MAP estimates
         dist = MultivariateNormal(loc=self.mean_subnet, scale_tril=self.posterior_scale)
         subnet_samples = dist.sample((n_samples,))
@@ -201,16 +211,19 @@ class DiagSubnetLaplace(SubnetLaplace, DiagLaplace):
     def _init_H(self):
         self.H = torch.zeros(self.n_params_subnet, device=self._device)
 
-    def _check_jacobians(self, Js):
+    def _check_jacobians(self, Js: torch.Tensor) -> None:
         if not isinstance(Js, torch.Tensor):
             raise ValueError("Jacobians have to be torch.Tensor.")
         if not Js.device == self._device:
             raise ValueError("Jacobians need to be on the same device as Laplace.")
+
         m, k, p = Js.size()
         if p != self.n_params_subnet:
             raise ValueError("Invalid Jacobians shape for Laplace posterior approx.")
 
-    def sample(self, n_samples=100, generator=None):
+    def sample(
+        self, n_samples: int = 100, generator: torch.Generator | None = None
+    ) -> torch.Tensor:
         # sample only subnetwork parameters and set all other parameters to their MAP estimates
         samples = torch.randn(
             n_samples, self.n_params_subnet, device=self._device, generator=generator
diff --git a/laplace/utils/__init__.py b/laplace/utils/__init__.py
index 948250d..ba2b7a8 100644
--- a/laplace/utils/__init__.py
+++ b/laplace/utils/__init__.py
@@ -1,3 +1,12 @@
+from laplace.utils.enums import (
+    HessianStructure,
+    Likelihood,
+    LinkApprox,
+    PredType,
+    PriorStructure,
+    SubsetOfWeights,
+    TuningMethod,
+)
 from laplace.utils.feature_extractor import FeatureExtractor
 from laplace.utils.matrix import Kron, KronDecomposed
 from laplace.utils.metrics import RunningNLLMetric
@@ -55,4 +64,11 @@
     "ModuleNameSubnetMask",
     "LastLayerSubnetMask",
     "RunningNLLMetric",
+    "SubsetOfWeights",
+    "HessianStructure",
+    "Likelihood",
+    "PredType",
+    "LinkApprox",
+    "TuningMethod",
+    "PriorStructure",
 ]
diff --git a/laplace/utils/enums.py b/laplace/utils/enums.py
new file mode 100644
index 0000000..90dabf4
--- /dev/null
+++ b/laplace/utils/enums.py
@@ -0,0 +1,43 @@
+from enum import Enum
+
+
+class SubsetOfWeights(str, Enum):
+    ALL = "all"
+    LAST_LAYER = "last_layer"
+    SUBNETWORK = "subnetwork"
+
+
+class HessianStructure(str, Enum):
+    FULL = "full"
+    KRON = "kron"
+    DIAG = "diag"
+    LOWRANK = "lowrank"
+
+
+class Likelihood(str, Enum):
+    REGRESSION = "regression"
+    CLASSIFICATION = "classification"
+    REWARD_MODELING = "reward_modeling"
+
+
+class PredType(str, Enum):
+    GLM = "glm"
+    NN = "nn"
+
+
+class LinkApprox(str, Enum):
+    MC = "mc"
+    PROBIT = "probit"
+    BRIDGE = "bridge"
+    BRIDGE_NORM = "bridge_norm"
+
+
+class TuningMethod(str, Enum):
+    MARGLIK = "marglik"
+    GRIDSEARCH = "gridsearch"
+
+
+class PriorStructure(str, Enum):
+    SCALAR = "scalar"
+    DIAG = "diag"
+    LAYERWISE = "layerwise"
diff --git a/laplace/utils/feature_extractor.py b/laplace/utils/feature_extractor.py
index 92ec4f8..c389f58 100644
--- a/laplace/utils/feature_extractor.py
+++ b/laplace/utils/feature_extractor.py
@@ -1,10 +1,13 @@
+from __future__ import annotations
+
+from collections.abc import MutableMapping
 from enum import Enum
-from typing import Callable, Optional, Tuple
+from typing import Any, Callable
 
 import torch
 import torch.nn as nn
 
-__all__ = ["FeatureExtractor"]
+__all__ = ["FeatureReduction", "FeatureExtractor"]
 
 
 class FeatureReduction(str, Enum):
@@ -46,9 +49,9 @@ class FeatureExtractor(nn.Module):
     def __init__(
         self,
         model: nn.Module,
-        last_layer_name: Optional[str] = None,
+        last_layer_name: str | None = None,
         enable_backprop: bool = False,
-        feature_reduction: Optional[FeatureReduction] = None,
+        feature_reduction: FeatureReduction | str | None = None,
     ) -> None:
         if feature_reduction is not None and feature_reduction not in [
             fr.value for fr in FeatureReduction
@@ -59,23 +62,26 @@ def __init__(
             )
 
         super().__init__()
-        self.model = model
-        self._features = dict()
-        self.enable_backprop = enable_backprop
-        self.feature_reduction = feature_reduction
+        self.model: nn.Module = model
+        self._features: dict[str, torch.Tensor] = dict()
+        self.enable_backprop: bool = enable_backprop
+        self.feature_reduction: FeatureReduction | None = feature_reduction
 
+        self.last_layer: nn.Module | None
         if last_layer_name is None:
             self.last_layer = None
         else:
             self.set_last_layer(last_layer_name)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any]
+    ) -> torch.Tensor:
         """Forward pass. If the last layer is not known yet, it will be
         determined when this function is called for the first time.
 
         Parameters
         ----------
-        x : torch.Tensor
+        x : torch.Tensor or a dict-like object containing the input tensors
             one batch of data to use as input for the forward pass
         """
         if self.last_layer is None:
@@ -87,15 +93,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return out
 
     def forward_with_features(
-        self, x: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any]
+    ) -> 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
+        x : torch.Tensor or a dict-like object containing the input tensors
             one batch of data to use as input for the forward pass
         """
         out = self.forward(x)
@@ -148,7 +154,9 @@ def hook(_, input, __):
 
         return hook
 
-    def find_last_layer(self, x: torch.Tensor) -> torch.Tensor:
+    def find_last_layer(
+        self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any]
+    ) -> torch.Tensor:
         """Automatically determines the last layer of the model with one
         forward pass. It assumes that the last layer is the same for every
         forward pass and that it is an instance of `torch.nn.Linear`.
@@ -158,7 +166,7 @@ def find_last_layer(self, x: torch.Tensor) -> torch.Tensor:
 
         Parameters
         ----------
-        x : torch.Tensor
+        x : torch.Tensor or dict-like object containing the input tensors
             one batch of data to use as input for the forward pass
         """
         if self.last_layer is not None:
diff --git a/laplace/utils/matrix.py b/laplace/utils/matrix.py
index 4b1e301..6f218d8 100644
--- a/laplace/utils/matrix.py
+++ b/laplace/utils/matrix.py
@@ -1,9 +1,12 @@
+from __future__ import annotations
+
 from math import pow
-from typing import Union
+from typing import Iterable
 
 import numpy as np
 import opt_einsum as oe
 import torch
+from torch import nn
 
 from laplace.utils.utils import _is_valid_scalar, block_diag, kron, symeig
 
@@ -19,16 +22,18 @@ class Kron:
 
     Parameters
     ----------
-    kfacs : list[Tuple]
-        each element in the list is a Tuple of two Kronecker factors Q, H
+    kfacs : list[Iterable[torch.Tensor] | torch.Tensor]
+        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
+    def __init__(self, kfacs: list[tuple[torch.Tensor] | torch.Tensor]) -> None:
+        self.kfacs: list[tuple[torch.Tensor] | torch.Tensor] = kfacs
 
     @classmethod
-    def init_from_model(cls, model, device):
+    def init_from_model(
+        cls, model: nn.Module | Iterable[nn.Parameter], device: torch.device
+    ) -> Kron:
         """Initialize Kronecker factors based on a models architecture.
 
         Parameters
@@ -53,7 +58,7 @@ def init_from_model(cls, model, device):
             elif 4 >= p.ndim >= 2:  # fully connected or conv
                 if p.ndim == 2:  # fully connected
                     P_in, P_out = p.size()
-                elif p.ndim > 2:
+                else:
                     P_in, P_out = p.shape[0], np.prod(p.shape[1:])
 
                 kfacs.append(
@@ -66,7 +71,7 @@ def init_from_model(cls, model, device):
                 raise ValueError("Invalid parameter shape in network.")
         return cls(kfacs)
 
-    def __add__(self, other):
+    def __add__(self, other: Kron) -> Kron:
         """Add up Kronecker factors `self` and `other`.
 
         Parameters
@@ -87,7 +92,7 @@ def __add__(self, other):
 
         return Kron(kfacs)
 
-    def __mul__(self, scalar: Union[float, torch.Tensor]):
+    def __mul__(self, scalar: float | torch.Tensor) -> Kron:
         """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`.
@@ -107,10 +112,10 @@ def __mul__(self, scalar: Union[float, torch.Tensor]):
         kfacs = [[pow(scalar, 1 / len(F)) * Hi for Hi in F] for F in self.kfacs]
         return Kron(kfacs)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.kfacs)
 
-    def decompose(self, damping=False):
+    def decompose(self, damping: bool = False) -> KronDecomposed:
         """Eigendecompose Kronecker factors and turn into `KronDecomposed`.
         Parameters
         ----------
@@ -291,22 +296,28 @@ class KronDecomposed:
         use dampen approximation mixing prior and Kron partially multiplicatively
     """
 
-    def __init__(self, eigenvectors, eigenvalues, deltas=None, damping=False):
-        self.eigenvectors = eigenvectors
-        self.eigenvalues = eigenvalues
-        device = eigenvectors[0][0].device
+    def __init__(
+        self,
+        eigenvectors: list[tuple[torch.Tensor]],
+        eigenvalues: list[tuple[torch.Tensor]],
+        deltas: torch.Tensor | None = None,
+        damping: bool = False,
+    ):
+        self.eigenvectors: list[tuple[torch.Tensor]] = eigenvectors
+        self.eigenvalues: list[tuple[torch.Tensor]] = eigenvalues
+        device: torch.device = eigenvectors[0][0].device
         if deltas is None:
-            self.deltas = torch.zeros(len(self), device=device)
+            self.deltas: torch.Tensor = torch.zeros(len(self), device=device)
         else:
             self._check_deltas(deltas)
-            self.deltas = deltas
-        self.damping = damping
+            self.deltas: torch.Tensor = deltas
+        self.damping: bool = damping
 
-    def detach(self):
+    def detach(self) -> KronDecomposed:
         self.deltas = self.deltas.detach()
         return self
 
-    def _check_deltas(self, deltas: torch.Tensor):
+    def _check_deltas(self, deltas: torch.Tensor) -> None:
         if not isinstance(deltas, torch.Tensor):
             raise ValueError("Can only add torch.Tensor to KronDecomposed.")
 
@@ -318,7 +329,7 @@ def _check_deltas(self, deltas: torch.Tensor):
         else:
             raise ValueError("Invalid shape of delta added to KronDecomposed.")
 
-    def __add__(self, deltas: torch.Tensor):
+    def __add__(self, deltas: torch.Tensor) -> KronDecomposed:
         """Add scalar per layer or only scalar to Kronecker factors.
 
         Parameters
@@ -333,7 +344,7 @@ def __add__(self, deltas: torch.Tensor):
         self._check_deltas(deltas)
         return KronDecomposed(self.eigenvectors, self.eigenvalues, self.deltas + deltas)
 
-    def __mul__(self, scalar):
+    def __mul__(self, scalar: torch.Tensor | float) -> KronDecomposed:
         """Multiply by a scalar by changing the eigenvalues.
         Same as for the case of `Kron`.
 
diff --git a/laplace/utils/metrics.py b/laplace/utils/metrics.py
index 254aa7e..d4a5818 100644
--- a/laplace/utils/metrics.py
+++ b/laplace/utils/metrics.py
@@ -13,13 +13,13 @@ class RunningNLLMetric(Metric):
         which class label to ignore when computing the NLL loss
     """
 
-    def __init__(self, ignore_index=-100):
+    def __init__(self, ignore_index: int = -100) -> None:
         super().__init__()
         self.add_state("nll_sum", default=torch.tensor(0.0), dist_reduce_fx="sum")
         self.add_state(
             "n_valid_labels", default=torch.tensor(0.0), dist_reduce_fx="sum"
         )
-        self.ignore_index = ignore_index
+        self.ignore_index: int = ignore_index
 
     def update(self, probs: torch.Tensor, targets: torch.Tensor) -> None:
         """
@@ -39,5 +39,5 @@ def update(self, probs: torch.Tensor, targets: torch.Tensor) -> None:
         )
         self.n_valid_labels += (targets != self.ignore_index).sum()
 
-    def compute(self):
+    def compute(self) -> torch.Tensor:
         return self.nll_sum / self.n_valid_labels
diff --git a/laplace/utils/subnetmask.py b/laplace/utils/subnetmask.py
index 0d19679..e79ec3c 100644
--- a/laplace/utils/subnetmask.py
+++ b/laplace/utils/subnetmask.py
@@ -1,9 +1,15 @@
+from __future__ import annotations
+
 from copy import deepcopy
 
 import torch
+from torch import nn
 from torch.nn import CrossEntropyLoss, MSELoss
 from torch.nn.utils import parameters_to_vector
+from torch.utils.data import DataLoader
 
+import laplace
+from laplace.utils.enums import Likelihood
 from laplace.utils.feature_extractor import FeatureExtractor
 from laplace.utils.swag import fit_diagonal_swag_var
 
@@ -27,34 +33,38 @@ class SubnetMask:
     model : torch.nn.Module
     """
 
-    def __init__(self, model):
-        self.model = model
-        self.parameter_vector = parameters_to_vector(self.model.parameters()).detach()
-        self._n_params = len(self.parameter_vector)
-        self._indices = None
-        self._n_params_subnet = None
+    def __init__(self, model: nn.Module) -> None:
+        self.model: nn.Module = model
+        self.parameter_vector: torch.Tensor = parameters_to_vector(
+            self.model.parameters()
+        ).detach()
+        self._n_params: int = len(self.parameter_vector)
+        self._indices: torch.LongTensor | None = None
+        self._n_params_subnet: int | None = None
 
-    def _check_select(self):
+    def _check_select(self) -> None:
         if self._indices is None:
             raise AttributeError("Subnetwork mask not selected. Run select() first.")
 
     @property
-    def _device(self):
+    def _device(self) -> torch.device:
         return next(self.model.parameters()).device
 
     @property
-    def indices(self):
+    def indices(self) -> torch.LongTensor:
         self._check_select()
         return self._indices
 
     @property
-    def n_params_subnet(self):
+    def n_params_subnet(self) -> int:
         if self._n_params_subnet is None:
             self._check_select()
             self._n_params_subnet = len(self._indices)
         return self._n_params_subnet
 
-    def convert_subnet_mask_to_indices(self, subnet_mask):
+    def convert_subnet_mask_to_indices(
+        self, subnet_mask: torch.Tensor
+    ) -> torch.LongTensor:
         """Converts a subnetwork mask into subnetwork indices.
 
         Parameters
@@ -103,7 +113,7 @@ def convert_subnet_mask_to_indices(self, subnet_mask):
         subnet_mask_indices = subnet_mask.nonzero(as_tuple=True)[0]
         return subnet_mask_indices
 
-    def select(self, train_loader=None):
+    def select(self, train_loader: DataLoader | None = None) -> torch.LongTensor:
         """Select the subnetwork mask.
 
         Parameters
@@ -126,7 +136,7 @@ def select(self, train_loader=None):
         self._indices = self.convert_subnet_mask_to_indices(subnet_mask)
         return self._indices
 
-    def get_subnet_mask(self, train_loader):
+    def get_subnet_mask(self, train_loader: DataLoader) -> torch.Tensor:
         """Get the subnetwork mask.
 
         Parameters
@@ -156,7 +166,7 @@ class ScoreBasedSubnetMask(SubnetMask):
         number of parameters in the subnetwork (i.e. number of top-scoring parameters to select)
     """
 
-    def __init__(self, model, n_params_subnet):
+    def __init__(self, model: nn.Module, n_params_subnet: int) -> None:
         super().__init__(model)
 
         if n_params_subnet is None:
@@ -168,12 +178,13 @@ def __init__(self, model, n_params_subnet):
                 f"Subnetwork ({n_params_subnet}) cannot be larger than model ({self._n_params})."
             )
         self._n_params_subnet = n_params_subnet
-        self._param_scores = None
+        self._param_scores: torch.Tensor | None = None
 
-    def compute_param_scores(self, train_loader):
+    def compute_param_scores(self, train_loader: DataLoader) -> torch.Tensor:
         raise NotImplementedError
 
-    def _check_param_scores(self):
+    def _check_param_scores(self) -> None:
+        assert self._param_scores is not None
         if self._param_scores.shape != self.parameter_vector.shape:
             raise ValueError(
                 "Parameter scores need to be of same shape as parameter vector."
@@ -181,7 +192,6 @@ def _check_param_scores(self):
 
     def get_subnet_mask(self, train_loader):
         """Get the subnetwork mask by (descendingly) ranking parameters based on their scores."""
-
         if self._param_scores is None:
             self._param_scores = self.compute_param_scores(train_loader)
         self._check_param_scores()
@@ -198,14 +208,14 @@ def get_subnet_mask(self, train_loader):
 class RandomSubnetMask(ScoreBasedSubnetMask):
     """Subnetwork mask of parameters sampled uniformly at random."""
 
-    def compute_param_scores(self, train_loader):
+    def compute_param_scores(self, train_loader: DataLoader) -> torch.Tensor:
         return torch.rand_like(self.parameter_vector)
 
 
 class LargestMagnitudeSubnetMask(ScoreBasedSubnetMask):
     """Subnetwork mask identifying the parameters with the largest magnitude."""
 
-    def compute_param_scores(self, train_loader):
+    def compute_param_scores(self, train_loader: DataLoader) -> torch.Tensor:
         return self.parameter_vector.abs()
 
 
@@ -222,11 +232,16 @@ class LargestVarianceDiagLaplaceSubnetMask(ScoreBasedSubnetMask):
         diagonal Laplace model to use for variance estimation
     """
 
-    def __init__(self, model, n_params_subnet, diag_laplace_model):
+    def __init__(
+        self,
+        model: nn.Module,
+        n_params_subnet: int,
+        diag_laplace_model: laplace.baselaplace.DiagLaplace,
+    ):
         super().__init__(model, n_params_subnet)
-        self.diag_laplace_model = diag_laplace_model
+        self.diag_laplace_model: laplace.baselaplace.DiagLaplace = diag_laplace_model
 
-    def compute_param_scores(self, train_loader):
+    def compute_param_scores(self, train_loader: DataLoader) -> torch.Tensor:
         if train_loader is None:
             raise ValueError("Need to pass train loader for subnet selection.")
 
@@ -255,27 +270,32 @@ class LargestVarianceSWAGSubnetMask(ScoreBasedSubnetMask):
 
     def __init__(
         self,
-        model,
-        n_params_subnet,
-        likelihood="classification",
-        swag_n_snapshots=40,
-        swag_snapshot_freq=1,
-        swag_lr=0.01,
+        model: nn.Module,
+        n_params_subnet: int,
+        likelihood: Likelihood | str = Likelihood.CLASSIFICATION,
+        swag_n_snapshots: int = 40,
+        swag_snapshot_freq: int = 1,
+        swag_lr: float = 0.01,
     ):
+        if likelihood not in [Likelihood.CLASSIFICATION, Likelihood.REGRESSION]:
+            raise ValueError("Only available for classification and regression!")
+
         super().__init__(model, n_params_subnet)
-        self.likelihood = likelihood
-        self.swag_n_snapshots = swag_n_snapshots
-        self.swag_snapshot_freq = swag_snapshot_freq
-        self.swag_lr = swag_lr
 
-    def compute_param_scores(self, train_loader):
+        self.likelihood: Likelihood | str = likelihood
+        self.swag_n_snapshots: int = swag_n_snapshots
+        self.swag_snapshot_freq: int = swag_snapshot_freq
+        self.swag_lr: float = swag_lr
+
+    def compute_param_scores(self, train_loader: DataLoader) -> torch.Tensor:
         if train_loader is None:
             raise ValueError("Need to pass train loader for subnet selection.")
 
-        if self.likelihood == "classification":
+        if self.likelihood == Likelihood.CLASSIFICATION:
             criterion = CrossEntropyLoss(reduction="mean")
-        elif self.likelihood == "regression":
+        else:
             criterion = MSELoss(reduction="mean")
+
         param_variances = fit_diagonal_swag_var(
             self.model,
             train_loader,
@@ -298,12 +318,12 @@ class ParamNameSubnetMask(SubnetMask):
         that define the subnetwork
     """
 
-    def __init__(self, model, parameter_names):
+    def __init__(self, model: nn.Module, parameter_names: list[str]) -> None:
         super().__init__(model)
-        self._parameter_names = parameter_names
-        self._n_params_subnet = None
+        self._parameter_names: list[str] = parameter_names
+        self._n_params_subnet: int | None = None
 
-    def _check_param_names(self):
+    def _check_param_names(self) -> None:
         param_names = deepcopy(self._parameter_names)
         if len(param_names) == 0:
             raise ValueError("Parameter name list cannot be empty.")
@@ -314,7 +334,7 @@ def _check_param_names(self):
         if len(param_names) > 0:
             raise ValueError(f"Parameters {param_names} do not exist in model.")
 
-    def get_subnet_mask(self, train_loader):
+    def get_subnet_mask(self, train_loader: DataLoader) -> torch.Tensor:
         """Get the subnetwork mask identifying the specified parameters."""
 
         self._check_param_names()
@@ -341,12 +361,12 @@ class ModuleNameSubnetMask(SubnetMask):
         the modules cannot have children, i.e. need to be leaf modules
     """
 
-    def __init__(self, model, module_names):
+    def __init__(self, model: nn.Module, module_names: list[str]):
         super().__init__(model)
-        self._module_names = module_names
-        self._n_params_subnet = None
+        self._module_names: list[str] = module_names
+        self._n_params_subnet: int | None = None
 
-    def _check_module_names(self):
+    def _check_module_names(self) -> None:
         module_names = deepcopy(self._module_names)
         if len(module_names) == 0:
             raise ValueError("Module name list cannot be empty.")
@@ -364,7 +384,7 @@ def _check_module_names(self):
         if len(module_names) > 0:
             raise ValueError(f"Modules {module_names} do not exist in model.")
 
-    def get_subnet_mask(self, train_loader):
+    def get_subnet_mask(self, train_loader: DataLoader) -> torch.Tensor:
         """Get the subnetwork mask identifying the specified modules."""
 
         self._check_module_names()
@@ -394,16 +414,15 @@ class LastLayerSubnetMask(ModuleNameSubnetMask):
         name of the model's last layer, if None it will be determined automatically
     """
 
-    def __init__(self, model, last_layer_name=None):
-        super().__init__(model, None)
-        self._feature_extractor = FeatureExtractor(
+    def __init__(self, model: nn.Module, last_layer_name: str | None = None):
+        super().__init__(model, [])
+        self._feature_extractor: FeatureExtractor = FeatureExtractor(
             self.model, last_layer_name=last_layer_name
         )
-        self._n_params_subnet = None
+        self._n_params_subnet: int | None = None
 
-    def get_subnet_mask(self, train_loader):
+    def get_subnet_mask(self, train_loader: DataLoader) -> torch.Tensor:
         """Get the subnetwork mask identifying the last layer."""
-
         if train_loader is None:
             raise ValueError("Need to pass train loader for subnet selection.")
 
diff --git a/laplace/utils/swag.py b/laplace/utils/swag.py
index 5bac456..cdab626 100644
--- a/laplace/utils/swag.py
+++ b/laplace/utils/swag.py
@@ -1,26 +1,31 @@
+from __future__ import annotations
+
 from copy import deepcopy
 
 import torch
+from torch import nn
 from torch.nn.utils import parameters_to_vector
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
 
 __all__ = ["fit_diagonal_swag_var"]
 
 
-def _param_vector(model):
+def _param_vector(model: nn.Module) -> torch.Tensor:
     return parameters_to_vector(model.parameters()).detach()
 
 
 def fit_diagonal_swag_var(
-    model,
-    train_loader,
-    criterion,
-    n_snapshots_total=40,
-    snapshot_freq=1,
-    lr=0.01,
-    momentum=0.9,
-    weight_decay=3e-4,
-    min_var=1e-30,
-):
+    model: nn.Module,
+    train_loader: DataLoader,
+    criterion: nn.CrossEntropyLoss | nn.MSELoss,
+    n_snapshots_total: int = 40,
+    snapshot_freq: int = 1,
+    lr: float = 0.01,
+    momentum: float = 0.9,
+    weight_decay: float = 3e-4,
+    min_var: float = 1e-30,
+) -> torch.Tensor:
     """
     Fit diagonal SWAG [1], which estimates marginal variances of model parameters by
     computing the first and second moment of SGD iterates with a large learning rate.
@@ -62,20 +67,21 @@ def fit_diagonal_swag_var(
     """
 
     # create a copy of the model to avoid undesired changes to the original model parameters
-    _model = deepcopy(model)
+    _model: nn.Module = deepcopy(model)
     _model.train()
-    device = next(_model.parameters()).device
+    device: torch.device = next(_model.parameters()).device
 
     # initialize running estimates of first and second moment of model parameters
-    mean = torch.zeros_like(_param_vector(_model))
-    sq_mean = torch.zeros_like(_param_vector(_model))
-    n_snapshots = 0
+    mean: torch.Tensor = torch.zeros_like(_param_vector(_model))
+    sq_mean: torch.Tensor = torch.zeros_like(_param_vector(_model))
+    n_snapshots: int = 0
 
     # run SGD to collect model snapshots
-    optimizer = torch.optim.SGD(
+    optimizer: Optimizer = torch.optim.SGD(
         _model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay
     )
-    n_epochs = snapshot_freq * n_snapshots_total
+    n_epochs: int = snapshot_freq * n_snapshots_total
+
     for epoch in range(n_epochs):
         for inputs, targets in train_loader:
             inputs, targets = inputs.to(device), targets.to(device)
@@ -92,5 +98,5 @@ def fit_diagonal_swag_var(
             n_snapshots += 1
 
     # compute marginal parameter variances, Var[P] = E[P^2] - E[P]^2
-    param_variances = torch.clamp(sq_mean - mean**2, min_var)
+    param_variances: torch.Tensor = torch.clamp(sq_mean - mean**2, min_var)
     return param_variances
diff --git a/laplace/utils/utils.py b/laplace/utils/utils.py
index 745c616..dc26b5c 100644
--- a/laplace/utils/utils.py
+++ b/laplace/utils/utils.py
@@ -1,15 +1,23 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import MutableMapping
-from typing import Union
+from typing import Callable
 
 import numpy as np
 import torch
 import torch.nn.functional as F
+import torchmetrics
+from torch import nn
 from torch.distributions.multivariate_normal import _precision_to_scale_tril
 from torch.nn import BatchNorm1d, BatchNorm2d, BatchNorm3d
 from torch.nn.utils import parameters_to_vector
+from torch.utils.data import DataLoader
 from torchmetrics import Metric
 
+import laplace
+from laplace.utils.enums import LinkApprox, PredType, PriorStructure
+
 __all__ = [
     "get_nll",
     "validate",
@@ -23,20 +31,22 @@
 ]
 
 
-def get_nll(out_dist, targets):
+def get_nll(out_dist: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
     return F.nll_loss(torch.log(out_dist), targets)
 
 
 @torch.no_grad()
 def validate(
-    laplace,
-    val_loader,
-    loss,
-    pred_type="glm",
-    link_approx="probit",
-    n_samples=100,
-    loss_with_var=False,
-    dict_key_y="labels",
+    laplace: laplace.baselaplace.BaseLaplace,
+    val_loader: DataLoader,
+    loss: torchmetrics.Metric
+    | Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
+    | Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],
+    pred_type: PredType | str = PredType.GLM,
+    link_approx: LinkApprox | str = LinkApprox.PROBIT,
+    n_samples: int = 100,
+    loss_with_var: int = False,
+    dict_key_y: str = "labels",
 ) -> float:
     laplace.model.eval()
     assert callable(loss) or isinstance(loss, Metric)
@@ -83,7 +93,7 @@ def validate(
         return loss.compute().item()
 
 
-def parameters_per_layer(model):
+def parameters_per_layer(model: nn.Module) -> list[int]:
     """Get number of parameters per layer.
 
     Parameters
@@ -97,7 +107,7 @@ def parameters_per_layer(model):
     return [np.prod(p.shape) for p in model.parameters()]
 
 
-def invsqrt_precision(M):
+def invsqrt_precision(M: torch.Tensor) -> torch.Tensor:
     """Compute ``M^{-0.5}`` as a tridiagonal matrix.
 
     Parameters
@@ -111,13 +121,13 @@ def invsqrt_precision(M):
     return _precision_to_scale_tril(M)
 
 
-def _is_batchnorm(module):
+def _is_batchnorm(module: nn.Module) -> bool:
     if isinstance(module, (BatchNorm1d, BatchNorm2d, BatchNorm3d)):
         return True
     return False
 
 
-def _is_valid_scalar(scalar: Union[float, int, torch.Tensor]) -> bool:
+def _is_valid_scalar(scalar: float | int | torch.Tensor) -> bool:
     if np.isscalar(scalar) and np.isreal(scalar):
         return True
     elif torch.is_tensor(scalar) and scalar.ndim <= 1:
@@ -127,7 +137,7 @@ def _is_valid_scalar(scalar: Union[float, int, torch.Tensor]) -> bool:
     return False
 
 
-def kron(t1, t2):
+def kron(t1: torch.Tensor, t2: torch.Tensor) -> torch.Tensor:
     """Computes the Kronecker product between two tensors.
 
     Parameters
@@ -155,7 +165,7 @@ def kron(t1, t2):
     return expanded_t1 * tiled_t2
 
 
-def diagonal_add_scalar(X, value):
+def diagonal_add_scalar(X: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
     """Add scalar value `value` to diagonal of `X`.
 
     Parameters
@@ -167,15 +177,12 @@ def diagonal_add_scalar(X, value):
     -------
     X_add_scalar : torch.Tensor
     """
-    if not X.device == torch.device("cpu"):
-        indices = torch.cuda.LongTensor([[i, i] for i in range(X.shape[0])])
-    else:
-        indices = torch.LongTensor([[i, i] for i in range(X.shape[0])])
+    indices = torch.LongTensor([[i, i] for i in range(X.shape[0])], device=X.device)
     values = X.new_ones(X.shape[0]).mul(value)
     return X.index_put(tuple(indices.t()), values, accumulate=True)
 
 
-def symeig(M):
+def symeig(M: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Symetric eigendecomposition avoiding failure cases by
     adding and removing jitter to the diagonal.
 
@@ -211,7 +218,7 @@ def symeig(M):
     return L, W
 
 
-def block_diag(blocks):
+def block_diag(blocks: list[torch.Tensor]) -> torch.Tensor:
     """Compose block-diagonal matrix of individual blocks.
 
     Parameters
@@ -232,7 +239,7 @@ def block_diag(blocks):
     return M
 
 
-def expand_prior_precision(prior_prec, model):
+def expand_prior_precision(prior_prec: torch.Tensor, model: nn.Module) -> torch.Tensor:
     """Expand prior precision to match the shape of the model parameters.
 
     Parameters
@@ -265,20 +272,46 @@ def expand_prior_precision(prior_prec, model):
 
 
 def fix_prior_prec_structure(
-    prior_prec_init, prior_structure, n_layers, n_params, device
-):
-    if prior_structure == "scalar":
+    prior_prec_init: torch.Tensor,
+    prior_structure: PriorStructure | str,
+    n_layers: int,
+    n_params: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Create a tensor of prior precision with the correct shape, depending on the
+    choice of the prior structure type.
+
+    Parameters
+    ----------
+    prior_prec_init: torch.Tensor
+        the initial prior precision tensor (could be scalar)
+    prior_structure: PriorStructure | str
+        the choice of the prior structure type
+    n_layers: int
+    n_params: int
+    device: torch.device
+
+    Returns
+    -------
+    correct_prior_precision: torch.Tensor
+    """
+    if prior_structure == PriorStructure.SCALAR:
         prior_prec_init = torch.full((1,), prior_prec_init, device=device)
-    elif prior_structure == "layerwise":
+    elif prior_structure == PriorStructure.LAYERWISE:
         prior_prec_init = torch.full((n_layers,), prior_prec_init, device=device)
-    elif prior_structure == "diag":
+    elif prior_structure == PriorStructure.DIAG:
         prior_prec_init = torch.full((n_params,), prior_prec_init, device=device)
     else:
         raise ValueError(f"Invalid prior structure {prior_structure}.")
     return prior_prec_init
 
 
-def normal_samples(mean, var, n_samples, generator=None):
+def normal_samples(
+    mean: torch.Tensor,
+    var: torch.Tensor,
+    n_samples: int,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
     """Produce samples from a batch of Normal distributions either parameterized
     by a diagonal or full covariance given by `var`.
 


Test patch:
diff --git a/tests/test_baselaplace.py b/tests/test_baselaplace.py
index 5469ed9..35ccc3c 100644
--- a/tests/test_baselaplace.py
+++ b/tests/test_baselaplace.py
@@ -15,10 +15,10 @@
 from torch.utils.data import DataLoader, TensorDataset
 from torchvision.models import wide_resnet50_2
 
+from laplace import DiagLaplace, FullLaplace, KronLaplace, LowRankLaplace
 from laplace.curvature import AsdlEF, AsdlGGN, BackPackGGN
 from laplace.curvature.backpack import BackPackEF
 from laplace.curvature.curvlinops import CurvlinopsEF, CurvlinopsGGN
-from laplace.laplace import DiagLaplace, FullLaplace, KronLaplace, LowRankLaplace
 from laplace.utils import KronDecomposed
 from tests.utils import ListDataset, dict_data_collator, jacobians_naive
 
diff --git a/tests/test_lllaplace.py b/tests/test_lllaplace.py
index c1596f1..bbea36d 100644
--- a/tests/test_lllaplace.py
+++ b/tests/test_lllaplace.py
@@ -139,6 +139,7 @@ def test_laplace_invalid_likelihood(laplace, model):
 def test_laplace_init_noise(laplace, model):
     # float
     sigma_noise = 1.2
+
     laplace(
         model, likelihood="regression", sigma_noise=sigma_noise, last_layer_name="1"
     )
diff --git a/tests/test_subset_params.py b/tests/test_subset_params.py
index bc300ce..ab8a496 100644
--- a/tests/test_subset_params.py
+++ b/tests/test_subset_params.py
@@ -6,10 +6,10 @@
 from torch.nn.utils import parameters_to_vector
 from torch.utils.data import DataLoader, TensorDataset
 
+from laplace import DiagLaplace, FullLaplace, KronLaplace
 from laplace.curvature import AsdlEF, AsdlGGN, AsdlHessian, BackPackEF, BackPackGGN
 from laplace.curvature.asdfghjkl import AsdfghjklEF, AsdfghjklGGN, AsdfghjklHessian
 from laplace.curvature.curvlinops import CurvlinopsEF, CurvlinopsGGN, CurvlinopsHessian
-from laplace.laplace import DiagLaplace, FullLaplace, KronLaplace
 
 torch.manual_seed(240)
 torch.set_default_tensor_type(torch.DoubleTensor)