Repository: aleximmer/Laplace (language: python)

Issue:
Add support for diagonal Kronecker factors in `Kron` matrix class
To support diagonal Kronecker factors in this library we have to adopt the methods of the `Kron` matrix class to be able to handle this correctly.
Remove default `loss` for `optimize_prior_precision`
When using `method=gridsearch` the user is supposed to specify a `loss` which will be used to evaluate the performance. Currently `loss` defaults to a classification loss, independent of the likelihood used for the Laplace approximation itself. I suggest to not set any default and raise an informative error if `gridsearch` is selected without specifying a `loss`. Alternatively, we could choose the default based on `likelihood`.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 6139719c..88e089ae 100644
--- a/README.md
+++ b/README.md
@@ -22,13 +22,20 @@ The [code](https://github.com/runame/laplace-redux) to reproduce the experiments
 
 ## Setup
 
-We assume `python3.8` since the package was developed with that version.
+For full compatibility, install this package in a fresh virtual env.
+We assume Python >= 3.8 (last tested on Python 3.10.12).
 PyTorch version 2.0 and up is also required for full compatibility.
 To install laplace with `pip`, run the following:
 ```bash
 pip install laplace-torch
 ```
 
+
+Note that, if you would like to use ASDL as a backend, please install it via the following, for full compatibility:
+```bash
+pip install git+https://github.com/wiseodd/asdl.git@dev
+```
+
 For development purposes, clone the repository and then install:
 ```bash
 # or after cloning the repository for development
@@ -88,6 +95,93 @@ ml = la.log_marginal_likelihood(prior_prec, obs_noise)
 ml.backward()
 ```
 
+### Laplace on foundation models like LLMs
+
+This library also supports Huggingface models and parameter-efficient fine-tuning.
+See `examples/huggingface_examples.py` for a runnable example.
+
+```python
+class MyGPT2(nn.Module):
+    def __init__(self, tokenizer: PreTrainedTokenizer) -> None:
+        super().__init__()
+        config = GPT2Config.from_pretrained('gpt2')
+        config.pad_token_id = tokenizer.pad_token_id
+        config.num_labels = 2
+        self.hf_model = GPT2ForSequenceClassification.from_pretrained(
+            'gpt2', config=config
+        )
+
+    def forward(self, data: MutableMapping) -> torch.Tensor:
+        '''
+        Custom forward function. Handles things like moving the
+        input tensor to the correct device inside.
+
+        Args:
+            data: A dict-like data structure with `input_ids` inside.
+                This is the default data structure assumed by Huggingface
+                dataloaders.
+
+        Returns:
+            logits: An `(batch_size, n_classes)`-sized tensor of logits.
+        '''
+        device = next(self.parameters()).device
+        input_ids = data['input_ids'].to(device)
+        attn_mask = data['attention_mask'].to(device)
+        output_dict = self.hf_model(input_ids=input_ids, attention_mask=attn_mask)
+        return output_dict.logits
+
+# Last-layer Laplace on the foundation model itself
+# -------------------------------------------------
+model = MyGPT2(tokenizer)
+model.eval()
+
+# Enable grad only for the last layer
+for p in model.hf_model.parameters():
+    p.requires_grad = False
+for p in model.hf_model.score.parameters():
+    p.requires_grad = True
+
+la = Laplace(
+    model,
+    likelihood='classification',
+    # Will only hit the last-layer since it's the only one that is grad-enabled
+    subset_of_weights='all',
+    hessian_structure='diag',
+)
+la.fit(dataloader)
+la.optimize_prior_precision()
+
+pred = la(next(iter(dataloader))
+
+# Laplace on the LoRA-attached LLM
+# --------------------------------
+def get_lora_model():
+    model = MyGPT2(tokenizer)  # Note we don't disable grad
+    config = LoraConfig(
+        r=4,
+        lora_alpha=16,
+        target_modules=['c_attn'],  # LoRA on the attention weights
+        lora_dropout=0.1,
+        bias='none',
+    )
+    lora_model = get_peft_model(model, config)
+    return lora_model
+
+lora_model = get_lora_model()
+# Train it as usual here...
+lora_model.eval()
+
+lora_la = Laplace(
+    lora_model,
+    likelihood='classification',
+    subset_of_weights='all',
+    hessian_structure='diag',
+    backend=AsdlGGN,
+)
+
+lora_pred = lora_la(next(iter(dataloader)))
+```
+
 ### Applying the LA over only a subset of the model parameters
 
 This example shows how to fit the Laplace approximation over only
diff --git a/examples/huggingface_example.md b/examples/huggingface_example.md
new file mode 100644
index 00000000..589aea65
--- /dev/null
+++ b/examples/huggingface_example.md
@@ -0,0 +1,235 @@
+# Applying Laplace on a Huggingface LLM model
+
+In this example, we will see how to apply Laplace on a GPT2 Huggingface (HF) model.
+Laplace only has lightweight requirements for this; namely that the model's `forward`
+method must only take a single dict-like object (`dict`, `UserDict`, or in general,
+`collections.abc.MutableMapping`). This is entirely compatible with HF since HF's
+data loaders are assumed to emit an object derived from `UserDict`. However, you
+need to ensure this yourself --- you need to wrap the standard HF model to conform
+to that requirement. Also, you need to e.g. do `torch.to(device)` _inside_ the
+said `forward` method.
+
+Let's start with as usual with importing stuff.
+
+```python
+from collections.abc import MutableMapping
+from collections import UserDict
+import numpy
+import torch
+from torch import nn
+import torch.utils.data as data_utils
+
+from laplace import Laplace
+
+import logging
+import warnings
+
+logging.basicConfig(level='ERROR')
+warnings.filterwarnings('ignore')
+
+from transformers import ( # noqa: E402
+    GPT2Config,
+    GPT2ForSequenceClassification,
+    GPT2Tokenizer,
+    DataCollatorWithPadding,
+    PreTrainedTokenizer,
+)
+from peft import LoraConfig, get_peft_model # noqa: E402
+from datasets import Dataset # noqa: E402
+
+# make deterministic
+
+torch.manual_seed(0)
+numpy.random.seed(0)
+```
+
+Next, we create a toy dataset. You can use any HF datasets or your own, of course.
+
+```python
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+tokenizer.pad_token_id = tokenizer.eos_token_id
+
+data = [
+{'text': 'Today is hot, but I will manage!!!!', 'label': 1},
+{'text': 'Tomorrow is cold', 'label': 0},
+{'text': 'Carpe diem', 'label': 1},
+{'text': 'Tempus fugit', 'label': 1},
+]
+dataset = Dataset.from_list(data)
+
+def tokenize(row):
+return tokenizer(row['text'])
+
+dataset = dataset.map(tokenize, remove_columns=['text'])
+dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+dataloader = data_utils.DataLoader(
+dataset, batch_size=100, collate_fn=DataCollatorWithPadding(tokenizer)
+)
+
+data = next(iter(dataloader))
+print(
+    f'Huggingface data defaults to UserDict, which is a MutableMapping? {isinstance(data, UserDict)}'
+)
+for k, v in data.items():
+    print(k, v.shape)
+```
+
+This is the output:
+
+```
+Huggingface data defaults to UserDict, which is a MutableMapping? True
+input_ids torch.Size([4, 9])
+attention_mask torch.Size([4, 9])
+labels torch.Size([4])
+```
+
+## Last-layer Laplace on a LLM
+
+Now, let's do the main "meat" of this example: Wrapping the HF model into a model that is
+compatible with Laplace. Notice that this wrapper just wraps the HF model and nothing else.
+Notice also we do `inputs.to(device)` inside `self.forward()`.
+
+```python
+class MyGPT2(nn.Module):
+    """
+    Huggingface LLM wrapper.
+
+    Args:
+        tokenizer: The tokenizer used for preprocessing the text data. Needed
+            since the model needs to know the padding token id.
+    """
+
+    def __init__(self, tokenizer: PreTrainedTokenizer) -> None:
+        super().__init__()
+        config = GPT2Config.from_pretrained('gpt2')
+        config.pad_token_id = tokenizer.pad_token_id
+        config.num_labels = 2
+        self.hf_model = GPT2ForSequenceClassification.from_pretrained(
+            'gpt2', config=config
+        )
+
+    def forward(self, data: MutableMapping) -> torch.Tensor:
+        '''
+        Custom forward function. Handles things like moving the
+        input tensor to the correct device inside.
+
+        Args:
+            data: A dict-like data structure with `input_ids` inside.
+                This is the default data structure assumed by Huggingface
+                dataloaders.
+
+        Returns:
+            logits: An `(batch_size, n_classes)`-sized tensor of logits.
+        '''
+        device = next(self.parameters()).device
+        input_ids = data['input_ids'].to(device)
+        attn_mask = data['attention_mask'].to(device)
+        output_dict = self.hf_model(input_ids=input_ids, attention_mask=attn_mask)
+        return output_dict.logits
+
+model = MyGPT2(tokenizer)
+```
+
+Now, let's apply Laplace. Let's do a last-layer Laplace first. We do so by switching off the
+gradients of all layers except the top layer. Laplace will automatically only compute the
+Hessian (and Jacobians) of the parameters in which `requires_grad` is `True`.
+
+Notice that you can "mix-and-match" this gradient switching. You can do a subnetwork Laplace
+easily by doing so!
+
+```python
+model.eval()
+
+# Enable grad only for the last layer
+
+for p in model.hf_model.parameters():
+    p.requires_grad = False
+
+for p in model.hf_model.score.parameters():
+    p.requires_grad = True
+
+la = Laplace(
+    model,
+    # Will only hit the last-layer since it's the only one that is grad-enabled
+    likelihood='classification',
+    subset_of_weights='all',
+    hessian_structure='diag',
+)
+la.fit(dataloader)
+la.optimize_prior_precision()
+
+X_test = next(iter(dataloader))
+print(f'[Foundation Model] The predictive tensor is of shape: {la(X_test).shape}.')
+```
+
+Here are the outputs to validate that Laplace works:
+
+```
+[Foundation Model] The predictive tensor is of shape: torch.Size([4, 2]).
+```
+
+## Laplace on LoRA parameters only
+
+Of course, you can also apply Laplace on the parameter-efficient fine tuning weights (like LoRA).
+To do this, simply extend your LLM with LoRA, using HF's `peft` library, and apply Laplace as
+usual. Note that `peft` automatically switches off the non-LoRA weights.
+
+```python
+def get_lora_model():
+    model = MyGPT2(tokenizer) # Note we don't disable grad
+        config = LoraConfig(
+        r=4,
+        lora_alpha=16,
+        target_modules=['c_attn'], # LoRA on the attention weights
+        lora_dropout=0.1,
+        bias='none',
+    )
+    lora_model = get_peft_model(model, config)
+    return lora_model
+
+lora_model = get_lora_model()
+
+# Train it as usual
+
+lora_model.eval()
+
+lora_la = Laplace(
+    lora_model,
+    likelihood='classification',
+    subset_of_weights='all',
+    hessian_structure='diag',
+)
+
+X_test = next(iter(dataloader))
+print(f'[LoRA-LLM] The predictive tensor is of shape: {lora_la(X_test).shape}.')
+```
+
+Here is the output, as expected:
+
+```
+[LoRA-LLM] The predictive tensor is of shape: torch.Size([4, 2]).
+```
+
+As a final note, the dict-like input requirement of Laplace is very flexible. It can essentially
+be applicable to any tasks and any models. You just need to wrap the said model and make sure
+that your data loaders emit dict-like objects, where the input tensors are the dicts' values.
+
+## Caveats
+
+Currently, diagonal EF with the Curvlinops backend is unsupported for dict-based inputs.
+This is because we use `torch.func`'s `vmap` to compute the diag-EF, and it only accepts
+tensor input in the model's `forward`.
+See [this issue](https://github.com/pytorch/functorch/issues/159).
+So, if you can write down your Huggingface model's `forward` to accept only a single tensor,
+this is much preferable.
+
+For instance, in the case of causal LLM like GPTs, only `input_ids`
+tensor is necessary.
+Then, any backend and any hessian factorization can be used in this case.
+
+Otherwise, if you must use dict-based inputs, choose the following backends:
+
+- `CurvlinopsGGN` for `hessian_factorization = {'kron', 'diag'}`
+- `CurvlinopsEF` for `hessian_factorization = {'kron'}`
+- `AsdlGGN` for `hessian_factorization = {'kron', 'diag'}`
+- `AsdlEF` for `hessian_factorization = {'kron', 'diag'}`
diff --git a/examples/huggingface_example.py b/examples/huggingface_example.py
new file mode 100644
index 00000000..8d8e6519
--- /dev/null
+++ b/examples/huggingface_example.py
@@ -0,0 +1,159 @@
+from collections.abc import MutableMapping
+from collections import UserDict
+import numpy
+import torch
+from torch import nn
+import torch.utils.data as data_utils
+
+from laplace import Laplace
+
+import logging
+import warnings
+
+logging.basicConfig(level='ERROR')
+warnings.filterwarnings('ignore')
+
+from transformers import (  # noqa: E402
+    GPT2Config,
+    GPT2ForSequenceClassification,
+    GPT2Tokenizer,
+    DataCollatorWithPadding,
+    PreTrainedTokenizer,
+)
+from peft import LoraConfig, get_peft_model  # noqa: E402
+from datasets import Dataset  # noqa: E402
+
+
+# make deterministic
+torch.manual_seed(0)
+numpy.random.seed(0)
+
+tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+tokenizer.pad_token_id = tokenizer.eos_token_id
+
+data = [
+    {'text': 'Today is hot, but I will manage!!!!', 'label': 1},
+    {'text': 'Tomorrow is cold', 'label': 0},
+    {'text': 'Carpe diem', 'label': 1},
+    {'text': 'Tempus fugit', 'label': 1},
+]
+dataset = Dataset.from_list(data)
+
+
+def tokenize(row):
+    return tokenizer(row['text'])
+
+
+dataset = dataset.map(tokenize, remove_columns=['text'])
+dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
+dataloader = data_utils.DataLoader(
+    dataset, batch_size=100, collate_fn=DataCollatorWithPadding(tokenizer)
+)
+
+data = next(iter(dataloader))
+print(
+    f'Huggingface data defaults to UserDict, which is a MutableMapping? {isinstance(data, UserDict)}'
+)
+for k, v in data.items():
+    print(k, v.shape)
+
+
+class MyGPT2(nn.Module):
+    """
+    Huggingface LLM wrapper.
+
+    Args:
+        tokenizer: The tokenizer used for preprocessing the text data. Needed
+            since the model needs to know the padding token id.
+    """
+
+    def __init__(self, tokenizer: PreTrainedTokenizer) -> None:
+        super().__init__()
+        config = GPT2Config.from_pretrained('gpt2')
+        config.pad_token_id = tokenizer.pad_token_id
+        config.num_labels = 2
+        self.hf_model = GPT2ForSequenceClassification.from_pretrained(
+            'gpt2', config=config
+        )
+
+    def forward(self, data: MutableMapping) -> torch.Tensor:
+        """
+        Custom forward function. Handles things like moving the
+        input tensor to the correct device inside.
+
+        Args:
+            data: A dict-like data structure with `input_ids` inside.
+                This is the default data structure assumed by Huggingface
+                dataloaders.
+
+        Returns:
+            logits: An `(batch_size, n_classes)`-sized tensor of logits.
+        """
+        device = next(self.parameters()).device
+        input_ids = data['input_ids'].to(device)
+        attn_mask = data['attention_mask'].to(device)
+        output_dict = self.hf_model(input_ids=input_ids, attention_mask=attn_mask)
+        return output_dict.logits
+
+
+model = MyGPT2(tokenizer)
+
+# Last-layer Laplace on the foundation model itself
+# -------------------------------------------------
+model.eval()
+
+# Enable grad only for the last layer
+for p in model.hf_model.parameters():
+    p.requires_grad = False
+
+for p in model.hf_model.score.parameters():
+    p.requires_grad = True
+
+la = Laplace(
+    model,
+    likelihood='classification',
+    # Will only hit the last-layer since it's the only one that is grad-enabled
+    subset_of_weights='all',
+    hessian_structure='diag',
+)
+la.fit(dataloader)
+la.optimize_prior_precision()
+
+X_test = next(iter(dataloader))
+print(f'[Foundation Model] The predictive tensor is of shape: {la(X_test).shape}.')
+
+del model
+del la
+
+
+# Laplace on the LoRA-attached LLM
+# --------------------------------
+
+
+def get_lora_model():
+    model = MyGPT2(tokenizer)  # Note we don't disable grad
+    config = LoraConfig(
+        r=4,
+        lora_alpha=16,
+        target_modules=['c_attn'],  # LoRA on the attention weights
+        lora_dropout=0.1,
+        bias='none',
+    )
+    lora_model = get_peft_model(model, config)
+    return lora_model
+
+
+lora_model = get_lora_model()
+# Train it as usual
+
+lora_model.eval()
+
+lora_la = Laplace(
+    lora_model,
+    likelihood='classification',
+    subset_of_weights='all',
+    hessian_structure='diag',
+)
+
+X_test = next(iter(dataloader))
+print(f'[LoRA-LLM] The predictive tensor is of shape: {lora_la(X_test).shape}.')
diff --git a/laplace/baselaplace.py b/laplace/baselaplace.py
index 7f146513..3bb26a54 100644
--- a/laplace/baselaplace.py
+++ b/laplace/baselaplace.py
@@ -2,17 +2,32 @@
 import numpy as np
 import torch
 from torch.nn.utils import parameters_to_vector, vector_to_parameters
-from torch.distributions import MultivariateNormal
+import tqdm
+from collections.abc import MutableMapping
+from laplace.curvature.asdfghjkl import AsdfghjklHessian
+from laplace.curvature.curvlinops import CurvlinopsEF
 import warnings
-
-from laplace.utils import (parameters_per_layer, invsqrt_precision,
-                           get_nll, validate, Kron, normal_samples,
-                           fix_prior_prec_structure)
+from torchmetrics import MeanSquaredError
+
+from laplace.utils import (
+    invsqrt_precision,
+    validate,
+    Kron,
+    normal_samples,
+    fix_prior_prec_structure,
+    RunningNLLMetric,
+)
 from laplace.curvature import AsdlHessian, CurvlinopsGGN
 
 
-__all__ = ['BaseLaplace', 'ParametricLaplace',
-           'FullLaplace', 'KronLaplace', 'DiagLaplace', 'LowRankLaplace']
+__all__ = [
+    'BaseLaplace',
+    'ParametricLaplace',
+    'FullLaplace',
+    'KronLaplace',
+    'DiagLaplace',
+    'LowRankLaplace',
+]
 
 
 class BaseLaplace:
@@ -21,8 +36,13 @@ class BaseLaplace:
     Parameters
     ----------
     model : torch.nn.Module
-    likelihood : {'classification', 'regression'}
-        determines the log likelihood Hessian approximation
+    likelihood : {'classification', 'regression', 'reward_modeling'}
+        determines the log likelihood Hessian approximation.
+        In the case of 'reward_modeling', it fits Laplace in using the classification likelihood,
+        then do prediction as in regression likelihood. The model needs to be defined accordingly:
+        The forward pass during training takes `x.shape == (batch_size, 2, dim)` with
+        `y.shape = (batch_size,)`. Meanwhile, during evaluation `x.shape == (batch_size, dim)`.
+        Note that 'reward_modeling' only supports `KronLaplace` and `DiagLaplace`.
     sigma_noise : torch.Tensor or float, default=1
         observation noise for the regression setting; must be 1 for classification
     prior_precision : torch.Tensor or float, default=1
@@ -41,32 +61,72 @@ class BaseLaplace:
     backend_kwargs : dict, default=None
         arguments passed to the backend on initialization, for example to
         set the number of MC samples for stochastic approximations.
+    asdl_fisher_kwargs : dict, default=None
+        arguments passed to the ASDL backend specifically on initialization.
     """
-    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., enable_backprop=False,
-                 backend=None, backend_kwargs=None):
-        if likelihood not in ['classification', 'regression']:
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        sigma_noise=1.0,
+        prior_precision=1.0,
+        prior_mean=0.0,
+        temperature=1.0,
+        enable_backprop=False,
+        backend=None,
+        backend_kwargs=None,
+        asdl_fisher_kwargs=None,
+    ):
+        if likelihood not in ['classification', 'regression', 'reward_modeling']:
             raise ValueError(f'Invalid likelihood type {likelihood}')
 
         self.model = model
 
-        self.n_params = len(parameters_to_vector(self.model.parameters()).detach())
-        self.n_layers = len(list(self.model.parameters()))
+        # Only do Laplace on params that require grad
+        self.params = []
+        self.is_subset_params = 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':
             raise ValueError('Sigma noise != 1 only available for regression.')
-        self.likelihood = likelihood
+
+        self.reward_modeling = likelihood == 'reward_modeling'
+        if self.reward_modeling:
+            # For fitting only. After it's done, self.likelihood = 'regression', see self.fit()
+            self.likelihood = 'classification'
+        else:
+            self.likelihood = likelihood
+
         self.sigma_noise = sigma_noise
         self.temperature = temperature
         self.enable_backprop = enable_backprop
 
+        if backend is None:
+            backend = CurvlinopsGGN
+        else:
+            if self.is_subset_params and 'backpack' in backend.__name__.lower():
+                raise ValueError(
+                    'If some grad are switched off, the BackPACK backend is not supported.'
+                )
+
         self._backend = None
-        self._backend_cls = backend if backend is not None else CurvlinopsGGN
+        self._backend_cls = backend
         self._backend_kwargs = dict() if backend_kwargs is None else backend_kwargs
+        self._asdl_fisher_kwargs = (
+            dict() if asdl_fisher_kwargs is None else asdl_fisher_kwargs
+        )
 
         # log likelihood = g(loss)
-        self.loss = 0.
+        self.loss = 0.0
         self.n_outputs = None
         self.n_data = 0
 
@@ -77,8 +137,9 @@ def _device(self):
     @property
     def backend(self):
         if self._backend is None:
-            self._backend = self._backend_cls(self.model, self.likelihood,
-                                              **self._backend_kwargs)
+            self._backend = self._backend_cls(
+                self.model, self.likelihood, **self._backend_kwargs
+            )
         return self._backend
 
     def _curv_closure(self, X, y, N):
@@ -101,10 +162,14 @@ def log_likelihood(self):
         -------
         log_likelihood : torch.Tensor
         """
-        factor = - self._H_factor
+        factor = -self._H_factor
         if self.likelihood == 'regression':
             # loss used is just MSE, need to add normalizer for gaussian likelihood
-            c = self.n_data * self.n_outputs * torch.log(self.sigma_noise * sqrt(2 * pi))
+            c = (
+                self.n_data
+                * self.n_outputs
+                * torch.log(self.sigma_noise * sqrt(2 * pi))
+            )
             return factor * self.loss - c
         else:
             # for classification Xent == log Cat
@@ -141,12 +206,18 @@ def prior_precision_diag(self):
             return self.prior_precision
 
         elif len(self.prior_precision) == self.n_layers:  # per layer
-            n_params_per_layer = parameters_per_layer(self.model)
-            return torch.cat([prior * torch.ones(n_params, device=self._device) for prior, n_params
-                              in zip(self.prior_precision, n_params_per_layer)])
+            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)
+                ]
+            )
 
         else:
-            raise ValueError('Mismatch of prior and model. Diagonal, scalar, or per-layer prior.')
+            raise ValueError(
+                'Mismatch of prior and model. Diagonal, scalar, or per-layer prior.'
+            )
 
     @property
     def prior_mean(self):
@@ -160,7 +231,7 @@ def prior_mean(self, prior_mean):
             if prior_mean.ndim == 0:
                 self._prior_mean = prior_mean.reshape(-1).to(self._device)
             elif prior_mean.ndim == 1:
-                if not len(prior_mean) in [1, self.n_params]:
+                if len(prior_mean) not in [1, self.n_params]:
                     raise ValueError('Invalid length of prior mean.')
                 self._prior_mean = prior_mean
             else:
@@ -183,12 +254,18 @@ def prior_precision(self, prior_precision):
                 self._prior_precision = prior_precision.reshape(-1).to(self._device)
             elif prior_precision.ndim == 1:
                 if len(prior_precision) not in [1, self.n_layers, self.n_params]:
-                    raise ValueError('Length of prior precision does not align with architecture.')
+                    raise ValueError(
+                        'Length of prior precision does not align with architecture.'
+                    )
                 self._prior_precision = prior_precision.to(self._device)
             else:
-                raise ValueError('Prior precision needs to be at most one-dimensional tensor.')
+                raise ValueError(
+                    'Prior precision needs to be at most one-dimensional tensor.'
+                )
         else:
-            raise ValueError('Prior precision either scalar or torch.Tensor up to 1-dim.')
+            raise ValueError(
+                'Prior precision either scalar or torch.Tensor up to 1-dim.'
+            )
 
     def optimize_prior_precision_base(
         self,
@@ -196,10 +273,10 @@ def optimize_prior_precision_base(
         method='marglik',
         n_steps=100,
         lr=1e-1,
-        init_prior_prec=1.,
+        init_prior_prec=1.0,
         prior_structure='scalar',
         val_loader=None,
-        loss=get_nll,
+        loss=None,
         log_prior_prec_min=-4,
         log_prior_prec_max=4,
         grid_size=100,
@@ -207,6 +284,7 @@ def optimize_prior_precision_base(
         n_samples=100,
         verbose=False,
         cv_loss_with_var=False,
+        progress_bar=False,
     ):
         """Optimize the prior precision post-hoc using the `method`
         specified by the user.
@@ -230,8 +308,11 @@ def optimize_prior_precision_base(
             otherwise, the structure of init_prior_prec is maintained.
         val_loader : torch.data.utils.DataLoader, default=None
             DataLoader for the validation set; each iterate is a training batch (X, y).
-        loss : callable, default=get_nll
-            loss function to use for gridsearch.
+        loss : callable or torchmetrics.Metric, default=None
+            loss function to use for CV. If callable, the loss is computed offline (memory intensive).
+            If torchmetrics.Metric, running loss is computed (efficient). The default
+            depends on the likelihood: `RunningNLLMetric()` for classification and
+            reward modeling, running `MeanSquaredError()` for regression.
         cv_loss_with_var: bool, default=False
             if true, `loss` takes three arguments `loss(output_mean, output_var, target)`,
             otherwise, `loss` takes two arguments `loss(output_mean, target)`
@@ -249,60 +330,103 @@ def optimize_prior_precision_base(
         verbose : bool, default=False
             if true, the optimized prior precision will be printed
             (can be a large tensor if the prior has a diagonal covariance).
+        progress_bar : bool, default=False
+            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':
                 self.prior_precision = fix_prior_prec_structure(
-                    self.prior_precision.item(), prior_structure,
-                    self.n_layers, self.n_params, self._device
+                    self.prior_precision.item(),
+                    prior_structure,
+                    self.n_layers,
+                    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)
-            for _ in range(n_steps):
+
+            if progress_bar:
+                pbar = tqdm.trange(n_steps)
+                pbar.set_description('[Optimizing marginal likelihood]')
+            else:
+                pbar = range(n_steps)
+
+            for _ in pbar:
                 optimizer.zero_grad()
                 prior_prec = log_prior_prec.exp()
-                neg_log_marglik = -self.log_marginal_likelihood(prior_precision=prior_prec)
+                neg_log_marglik = -self.log_marginal_likelihood(
+                    prior_precision=prior_prec
+                )
                 neg_log_marglik.backward()
                 optimizer.step()
             self.prior_precision = log_prior_prec.detach().exp()
         elif method == 'gridsearch':
             if val_loader is None:
                 raise ValueError('gridsearch requires a validation set DataLoader')
-            interval = torch.logspace(
-                log_prior_prec_min, log_prior_prec_max, grid_size
-            )
+
+            interval = torch.logspace(log_prior_prec_min, log_prior_prec_max, grid_size)
+
+            if loss is None:
+                loss = (
+                    MeanSquaredError(num_outputs=self.n_outputs)
+                    if self.likelihood == 'regression'
+                    else RunningNLLMetric()
+                )
+
             self.prior_precision = self._gridsearch(
-                loss, interval, val_loader, pred_type=pred_type,
-                link_approx=link_approx, n_samples=n_samples, loss_with_var=cv_loss_with_var
+                loss,
+                interval,
+                val_loader,
+                pred_type=pred_type,
+                link_approx=link_approx,
+                n_samples=n_samples,
+                loss_with_var=cv_loss_with_var,
+                progress_bar=progress_bar,
             )
         else:
             raise ValueError('For now only marglik and gridsearch is implemented.')
         if verbose:
             print(f'Optimized prior precision is {self.prior_precision}.')
 
-    def _gridsearch(self, loss, interval, val_loader, pred_type,
-                    link_approx='probit', n_samples=100, loss_with_var=False):
+    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)
+
         results = list()
         prior_precs = list()
-        for prior_prec in interval:
+        pbar = tqdm.tqdm(interval) if progress_bar else interval
+        for prior_prec in pbar:
             self.prior_precision = prior_prec
             try:
-                out_dist, targets = validate(
-                    self, val_loader, pred_type=pred_type,
-                    link_approx=link_approx, n_samples=n_samples
+                result = validate(
+                    self,
+                    val_loader,
+                    loss,
+                    pred_type=pred_type,
+                    link_approx=link_approx,
+                    n_samples=n_samples,
+                    loss_with_var=loss_with_var,
                 )
-                if self.likelihood == 'regression':
-                    out_mean, out_var = out_dist
-                    if loss_with_var:
-                        result = loss(out_mean, out_var, targets).item()
-                    else:
-                        result = loss(out_mean, targets).item()
-                else:
-                    result = loss(out_dist, targets).item()
             except RuntimeError:
                 result = np.inf
+
+            if progress_bar:
+                pbar.set_description(
+                    f'[Grid search | prior_prec: {prior_prec:.3e}, loss: {result:.3f}]'
+                )
+
             results.append(result)
             prior_precs.append(prior_prec)
         return prior_precs[np.argmin(results)]
@@ -326,7 +450,9 @@ def sigma_noise(self, sigma_noise):
             else:
                 raise ValueError('Sigma noise needs to be scalar or 1-dimensional.')
         else:
-            raise ValueError('Invalid type: sigma noise needs to be torch.Tensor or scalar.')
+            raise ValueError(
+                'Invalid type: sigma noise needs to be torch.Tensor or scalar.'
+            )
 
     @property
     def _H_factor(self):
@@ -358,11 +484,31 @@ class ParametricLaplace(BaseLaplace):
     all cases \\(P_0 = \\textrm{diag}(p_0)\\) and the structure of \\(p_0\\) can be varied.
     """
 
-    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., enable_backprop=False,
-                 backend=None, backend_kwargs=None):
-        super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, enable_backprop, backend, backend_kwargs)
+    def __init__(
+        self,
+        model,
+        likelihood,
+        sigma_noise=1.0,
+        prior_precision=1.0,
+        prior_mean=0.0,
+        temperature=1.0,
+        enable_backprop=False,
+        backend=None,
+        backend_kwargs=None,
+        asdl_fisher_kwargs=None,
+    ):
+        super().__init__(
+            model,
+            likelihood,
+            sigma_noise,
+            prior_precision,
+            prior_mean,
+            temperature,
+            enable_backprop,
+            backend,
+            backend_kwargs,
+            asdl_fisher_kwargs,
+        )
         if not hasattr(self, 'H'):
             self._init_H()
             # posterior mean/mode
@@ -375,7 +521,7 @@ def _check_H_init(self):
         if self.H is None:
             raise AttributeError('Laplace not fitted. Run fit() first.')
 
-    def fit(self, train_loader, override=True):
+    def fit(self, train_loader, override=True, progress_bar=False):
         """Fit the local Laplace approximation at the parameters of the model.
 
         Parameters
@@ -386,6 +532,9 @@ 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
+            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 override:
             self._init_H()
@@ -393,24 +542,46 @@ def fit(self, train_loader, override=True):
             self.n_data = 0
 
         self.model.eval()
-        self.mean = parameters_to_vector(self.model.parameters())
 
+        self.mean = parameters_to_vector(self.params)
         if not self.enable_backprop:
             self.mean = self.mean.detach()
 
-        X, _ = next(iter(train_loader))
+        data = next(iter(train_loader))
         with torch.no_grad():
-            try:
-                out = self.model(X[:1].to(self._device))
-            except (TypeError, AttributeError):
-                out = self.model(X.to(self._device))
+            if isinstance(data, MutableMapping):  # To support Huggingface dataset
+                if isinstance(self, DiagLaplace) and self._backend_cls == CurvlinopsEF:
+                    raise ValueError(
+                        'Currently DiagEF is not supported under CurvlinopsEF backend '
+                        + 'for custom models with non-tensor inputs '
+                        + '(https://github.com/pytorch/functorch/issues/159). Consider '
+                        + 'using AsdlEF backend instead.'
+                    )
+
+                out = self.model(data)
+            else:
+                X = data[0]
+                try:
+                    out = self.model(X[:1].to(self._device))
+                except (TypeError, AttributeError):
+                    out = self.model(X.to(self._device))
         self.n_outputs = out.shape[-1]
         setattr(self.model, 'output_size', self.n_outputs)
 
         N = len(train_loader.dataset)
-        for X, y in train_loader:
+        if progress_bar:
+            pbar = tqdm.tqdm(train_loader)
+            pbar.set_description('[Computing Hessian]')
+        else:
+            pbar = train_loader
+
+        for data in pbar:
+            if isinstance(data, MutableMapping):  # To support Huggingface dataset
+                X, y = data, data['labels'].to(self._device)
+            else:
+                X, y = data
+                X, y = X.to(self._device), y.to(self._device)
             self.model.zero_grad()
-            X, y = X.to(self._device), y.to(self._device)
             loss_batch, H_batch = self._curv_closure(X, y, N)
             self.loss += loss_batch
             self.H += H_batch
@@ -428,7 +599,7 @@ def scatter(self):
         [type]
             [description]
         """
-        delta = (self.mean - self.prior_mean)
+        delta = self.mean - self.prior_mean
         return (delta * self.prior_precision_diag) @ delta
 
     @property
@@ -492,8 +663,10 @@ def log_prob(self, value, normalized=True):
         log_prob : torch.Tensor
         """
         if not normalized:
-            return - self.square_norm(value) / 2
-        log_prob = - self.n_params / 2 * log(2 * pi) + self.log_det_posterior_precision / 2
+            return -self.square_norm(value) / 2
+        log_prob = (
+            -self.n_params / 2 * log(2 * pi) + self.log_det_posterior_precision / 2
+        )
         log_prob -= self.square_norm(value) / 2
         return log_prob
 
@@ -529,19 +702,31 @@ def log_marginal_likelihood(self, prior_precision=None, sigma_noise=None):
 
         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):
+    def __call__(
+        self,
+        x,
+        pred_type='glm',
+        joint=False,
+        link_approx='probit',
+        n_samples=100,
+        diagonal_output=False,
+        generator=None,
+        **model_kwargs,
+    ):
         """Compute the posterior predictive on input data `x`.
 
         Parameters
         ----------
-        x : torch.Tensor
-            `(batch_size, input_shape)`
+        x : torch.Tensor or MutableMapping
+            `(batch_size, input_shape)` if tensor. If MutableMapping, must contain
+            the said tensor.
 
         pred_type : {'glm', 'nn'}, default='glm'
             type of posterior predictive, linearized GLM predictive or neural
             network sampling predictive. The GLM predictive is consistent with
-            the curvature approximations used here.
+            the curvature approximations used here. When Laplace is done only
+            on subset of parameters (i.e. some grad are disabled),
+            only `nn` predictive is supported.
 
         link_approx : {'mc', 'probit', 'bridge', 'bridge_norm'}
             how to approximate the classification link function for the `'glm'`.
@@ -581,12 +766,22 @@ def __call__(self, x, pred_type='glm', joint=False, link_approx='probit',
             raise ValueError(f'Unsupported link approximation {link_approx}.')
 
         if pred_type == 'nn' and link_approx != 'mc':
-            raise ValueError('Only mc link approximation is supported for nn prediction type.')
+            raise ValueError(
+                'Only mc link approximation is supported for nn prediction type.'
+            )
 
         if generator is not None:
-            if not isinstance(generator, torch.Generator) or generator.device != x.device:
+            if (
+                not isinstance(generator, torch.Generator)
+                or generator.device != x.device
+            ):
                 raise ValueError('Invalid random generator (check type and device).')
 
+        # For reward modeling, replace the likelihood to regression and override model state
+        if self.reward_modeling and self.likelihood == 'classification':
+            self.likelihood = 'regression'
+            self.model.output_size = 1
+
         if pred_type == 'glm':
             f_mu, f_var = self._glm_predictive_distribution(
                 x, joint=joint and self.likelihood == 'regression'
@@ -596,37 +791,49 @@ def __call__(self, x, pred_type='glm', joint=False, link_approx='probit',
                 return f_mu, f_var
             # classification
             if link_approx == 'mc':
-                return self.predictive_samples(x, pred_type='glm', n_samples=n_samples,
-                                               diagonal_output=diagonal_output).mean(dim=0)
+                return self.predictive_samples(
+                    x,
+                    pred_type='glm',
+                    n_samples=n_samples,
+                    diagonal_output=diagonal_output,
+                ).mean(dim=0)
             elif link_approx == 'probit':
-                kappa = 1 / torch.sqrt(1. + np.pi / 8 * f_var.diagonal(dim1=1, dim2=2))
+                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:
                 # zero mean correction
-                f_mu -= (f_var.sum(-1) * f_mu.sum(-1).reshape(-1, 1) /
-                         f_var.sum(dim=(1, 2)).reshape(-1, 1))
-                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))
+                f_mu -= (
+                    f_var.sum(-1)
+                    * f_mu.sum(-1).reshape(-1, 1)
+                    / f_var.sum(dim=(1, 2)).reshape(-1, 1)
+                )
+                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':
                     f_var_diag_mean = f_var_diag.mean(dim=1)
-                    f_var_diag_mean /= torch.as_tensor([K/2], device=self._device).sqrt()
+                    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
+                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:
-            samples = self._nn_predictive_samples(x, n_samples)
             if self.likelihood == 'regression':
+                samples = self._nn_predictive_samples(x, n_samples, **model_kwargs)
                 return samples.mean(dim=0), samples.var(dim=0)
-            return samples.mean(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):
+    def predictive_samples(
+        self, x, pred_type='glm', n_samples=100, diagonal_output=False, generator=None
+    ):
         """Sample from the posterior predictive on input data `x`.
         Can be used, for example, for Thompson sampling.
 
@@ -660,20 +867,30 @@ def predictive_samples(self, x, pred_type='glm', n_samples=100,
 
         if pred_type == 'glm':
             f_mu, f_var = self._glm_predictive_distribution(x)
-            assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
+            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':
                 return f_samples
-            return torch.softmax(f_samples, dim=-1)
+            else:
+                return torch.softmax(f_samples, dim=-1)
 
         else:  # 'nn'
-            return self._nn_predictive_samples(x, n_samples)
+            return self._nn_predictive_samples(x, n_samples, generator)
 
     @torch.enable_grad()
     def _glm_predictive_distribution(self, X, joint=False):
-        Js, f_mu = self.backend.jacobians(X, enable_backprop=self.enable_backprop)
+        if 'backpack' in self._backend_cls.__name__.lower():
+            # BackPACK supports backprop through Jacobians, but it interferes with functorch
+            Js, f_mu = self.backend.jacobians(X, enable_backprop=self.enable_backprop)
+        else:
+            # For ASDL and Curvlinops, we use functorch
+            Js, f_mu = self.backend.functorch_jacobians(
+                X, enable_backprop=self.enable_backprop
+            )
 
         if joint:
             f_mu = f_mu.flatten()  # (batch*out)
@@ -681,20 +898,37 @@ def _glm_predictive_distribution(self, X, joint=False):
         else:
             f_var = self.functional_variance(Js)
 
-        return (f_mu.detach(), f_var.detach()) if not self.enable_backprop else (f_mu, f_var)
+        return (
+            (f_mu.detach(), f_var.detach())
+            if not self.enable_backprop
+            else (f_mu, f_var)
+        )
 
-    def _nn_predictive_samples(self, X, n_samples=100):
+    def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwargs):
         fs = list()
-        for sample in self.sample(n_samples):
-            vector_to_parameters(sample, self.model.parameters())
-            f = self.model(X.to(self._device))
-            fs.append(f.detach() if not self.enable_backprop else f)
-        vector_to_parameters(self.mean, self.model.parameters())
+        for sample in self.sample(n_samples, generator):
+            vector_to_parameters(sample, self.params)
+            logits = self.model(
+                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':
             fs = torch.softmax(fs, dim=-1)
         return fs
 
+    def _nn_predictive_classification(self, X, n_samples=100, **model_kwargs):
+        py = 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):
         """Compute functional variance for the `'glm'` predictive:
         `f_var[i] = Jacs[i] @ P.inv() @ Jacs[i].T`, which is a output x output
@@ -738,7 +972,7 @@ def functional_covariance(self, Jacs):
         """
         raise NotImplementedError
 
-    def sample(self, n_samples=100):
+    def sample(self, n_samples=100, generator=None):
         """Sample from the Laplace posterior approximation, i.e.,
         \\( \\theta \\sim \\mathcal{N}(\\theta_{MAP}, P^{-1})\\).
 
@@ -746,21 +980,50 @@ def sample(self, n_samples=100):
         ----------
         n_samples : int, default=100
             number of samples
+
+        generator : torch.Generator, optional
+            random number generator to control the samples
         """
         raise NotImplementedError
 
-    def optimize_prior_precision(self, method='marglik', pred_type='glm', n_steps=100, lr=1e-1,
-                                 init_prior_prec=1., prior_structure='scalar', val_loader=None,
-                                 loss=get_nll, 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):
+    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, method, n_steps, lr,
-                                           init_prior_prec, prior_structure,
-                                           val_loader, loss, log_prior_prec_min,
-                                           log_prior_prec_max, grid_size,
-                                           link_approx, n_samples, verbose,
-                                           cv_loss_with_var)
+        self.optimize_prior_precision_base(
+            pred_type,
+            method,
+            n_steps,
+            lr,
+            init_prior_prec,
+            prior_structure,
+            val_loader,
+            loss,
+            log_prior_prec_min,
+            log_prior_prec_max,
+            grid_size,
+            link_approx,
+            n_samples,
+            verbose,
+            cv_loss_with_var,
+            progress_bar,
+        )
 
     @property
     def posterior_precision(self):
@@ -786,7 +1049,7 @@ def state_dict(self) -> dict:
             'likelihood': self.likelihood,
             'temperature': self.temperature,
             'enable_backprop': self.enable_backprop,
-            'cls_name': self.__class__.__name__
+            'cls_name': self.__class__.__name__,
         }
         return state_dict
 
@@ -808,9 +1071,13 @@ def load_state_dict(self, state_dict: dict):
 
         # Ignorable warnings
         if self.prior_mean is None and state_dict['prior_mean'] is not None:
-            warnings.warn('Loading non-`None` prior mean into a `None` prior mean. You might get wrong results.')
+            warnings.warn(
+                'Loading non-`None` prior mean into a `None` prior mean. You might get wrong results.'
+            )
         if self.temperature != state_dict['temperature']:
-            warnings.warn('Different `temperature` parameters detected. Some calculation might be off!')
+            warnings.warn(
+                'Different `temperature` parameters detected. Some calculation might be off!'
+            )
         if self.enable_backprop != state_dict['enable_backprop']:
             warnings.warn(
                 'Different `enable_backprop` values. You might encounter error when differentiating'
@@ -838,13 +1105,33 @@ class FullLaplace(ParametricLaplace):
     Mathematically, we have \\(P \\in \\mathbb{R}^{P \\times P}\\).
     See `BaseLaplace` for the full interface.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ('all', 'full')
 
-    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., enable_backprop=False, backend=None, backend_kwargs=None):
-        super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, enable_backprop, backend, backend_kwargs)
+    def __init__(
+        self,
+        model,
+        likelihood,
+        sigma_noise=1.0,
+        prior_precision=1.0,
+        prior_mean=0.0,
+        temperature=1.0,
+        enable_backprop=False,
+        backend=None,
+        backend_kwargs=None,
+    ):
+        super().__init__(
+            model,
+            likelihood,
+            sigma_noise,
+            prior_precision,
+            prior_mean,
+            temperature,
+            enable_backprop,
+            backend,
+            backend_kwargs,
+        )
         self._posterior_scale = None
 
     def _init_H(self):
@@ -853,9 +1140,9 @@ def _init_H(self):
     def _curv_closure(self, X, y, N):
         return self.backend.full(X, y, N=N)
 
-    def fit(self, train_loader, override=True):
+    def fit(self, train_loader, override=True, progress_bar=False):
         self._posterior_scale = None
-        return super().fit(train_loader, override=override)
+        return super().fit(train_loader, override=override, progress_bar=progress_bar)
 
     def _compute_scale(self):
         self._posterior_scale = invsqrt_precision(self.posterior_precision)
@@ -911,12 +1198,16 @@ def functional_variance(self, Js):
 
     def functional_covariance(self, Js):
         n_batch, n_outs, n_params = Js.shape
-        Js = Js.reshape(n_batch*n_outs, n_params)
+        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):
-        dist = MultivariateNormal(loc=self.mean, scale_tril=self.posterior_scale)
-        return dist.sample((n_samples,))
+    def sample(self, n_samples=100, generator=None):
+        samples = torch.randn(
+            n_samples, self.n_params, device=self._device, generator=generator
+        )
+        # (n_samples, n_params) x (n_params, n_params) -> (n_samples, n_params)
+        samples = samples @ self.posterior_scale
+        return self.mean.reshape(1, self.n_params) + samples
 
 
 class KronLaplace(ParametricLaplace):
@@ -931,22 +1222,44 @@ class KronLaplace(ParametricLaplace):
     and computing posterior covariances, marginal likelihood, etc.
     Damping can be enabled by setting `damping=True`.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ('all', 'kron')
 
-    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., enable_backprop=False, backend=None,
-                 damping=False, **backend_kwargs):
+    def __init__(
+        self,
+        model,
+        likelihood,
+        sigma_noise=1.0,
+        prior_precision=1.0,
+        prior_mean=0.0,
+        temperature=1.0,
+        enable_backprop=False,
+        backend=None,
+        damping=False,
+        backend_kwargs=None,
+        asdl_fisher_kwargs=None,
+    ):
         self.damping = damping
         self.H_facs = None
-        super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, enable_backprop, backend, **backend_kwargs)
+        super().__init__(
+            model,
+            likelihood,
+            sigma_noise,
+            prior_precision,
+            prior_mean,
+            temperature,
+            enable_backprop,
+            backend,
+            backend_kwargs,
+            asdl_fisher_kwargs,
+        )
 
     def _init_H(self):
-        self.H = Kron.init_from_model(self.model, self._device)
+        self.H = Kron.init_from_model(self.params, self._device)
 
     def _curv_closure(self, X, y, N):
-        return self.backend.kron(X, y, N=N)
+        return self.backend.kron(X, y, N=N, **self._asdl_fisher_kwargs)
 
     @staticmethod
     def _rescale_factors(kron, factor):
@@ -955,7 +1268,7 @@ def _rescale_factors(kron, factor):
                 F[1] *= factor
         return kron
 
-    def fit(self, train_loader, override=True):
+    def fit(self, train_loader, override=True, progress_bar=False):
         if override:
             self.H_facs = None
 
@@ -964,15 +1277,19 @@ def fit(self, train_loader, override=True):
             n_data_new = 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(self.H_facs, n_data_old / (n_data_old + n_data_new))
+            self.H_facs = self._rescale_factors(
+                self.H_facs, n_data_old / (n_data_old + n_data_new)
+            )
 
-        super().fit(train_loader, override=override)
+        super().fit(train_loader, override=override, progress_bar=progress_bar)
 
         if self.H_facs is None:
             self.H_facs = self.H
         else:
             # discount new factors that were computed assuming N = n_data_new
-            self.H = self._rescale_factors(self.H, n_data_new / (n_data_new + n_data_old))
+            self.H = self._rescale_factors(
+                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)
@@ -1006,15 +1323,19 @@ def functional_variance(self, Js):
     def functional_covariance(self, Js):
         self._check_jacobians(Js)
         n_batch, n_outs, n_params = Js.shape
-        Js = Js.reshape(n_batch*n_outs, n_params).unsqueeze(0)
+        Js = Js.reshape(n_batch * n_outs, n_params).unsqueeze(0)
         cov = self.posterior_precision.inv_square_form(Js).squeeze(0)
-        assert cov.shape == (n_batch*n_outs, n_batch*n_outs)
+        assert cov.shape == (n_batch * n_outs, n_batch * n_outs)
         return cov
 
-    def sample(self, n_samples=100):
-        samples = torch.randn(n_samples, self.n_params, device=self._device)
+    def sample(self, n_samples=100, generator=None):
+        samples = torch.randn(
+            n_samples, self.n_params, device=self._device, generator=generator
+        )
         samples = self.posterior_precision.bmm(samples, exponent=-0.5)
-        return self.mean.reshape(1, self.n_params) + samples.reshape(n_samples, self.n_params)
+        return self.mean.reshape(1, self.n_params) + samples.reshape(
+            n_samples, self.n_params
+        )
 
     @BaseLaplace.prior_precision.setter
     def prior_precision(self, prior_precision):
@@ -1049,13 +1370,32 @@ class LowRankLaplace(ParametricLaplace):
 
     See `BaseLaplace` for the full interface.
     """
+
     _key = ('all', 'lowrank')
-    def __init__(self, model, likelihood, sigma_noise=1, prior_precision=1, prior_mean=0,
-                 temperature=1, enable_backprop=False, backend=AsdlHessian, backend_kwargs=None):
-        super().__init__(model, likelihood, sigma_noise=sigma_noise,
-                         prior_precision=prior_precision, prior_mean=prior_mean,
-                         temperature=temperature, enable_backprop=enable_backprop,
-                         backend=backend, backend_kwargs=backend_kwargs)
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        sigma_noise=1,
+        prior_precision=1,
+        prior_mean=0,
+        temperature=1,
+        enable_backprop=False,
+        backend=AsdfghjklHessian,
+        backend_kwargs=None,
+    ):
+        super().__init__(
+            model,
+            likelihood,
+            sigma_noise=sigma_noise,
+            prior_precision=prior_precision,
+            prior_mean=prior_mean,
+            temperature=temperature,
+            enable_backprop=enable_backprop,
+            backend=backend,
+            backend_kwargs=backend_kwargs,
+        )
 
     def _init_H(self):
         self.H = None
@@ -1120,16 +1460,16 @@ def functional_variance(self, Jacs):
 
     def functional_covariance(self, Jacs):
         n_batch, n_outs, n_params = Jacs.shape
-        Jacs = Jacs.reshape(n_batch*n_outs, n_params)
+        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)
         cov = prior_cov - info_gain
-        assert cov.shape == (n_batch*n_outs, n_batch*n_outs)
+        assert cov.shape == (n_batch * n_outs, n_batch * n_outs)
         return cov
 
-    def sample(self, n_samples):
-        samples = torch.randn(self.n_params, n_samples)
+    def sample(self, n_samples, generator=None):
+        samples = torch.randn(self.n_params, n_samples, generator=generator)
         d = self.prior_precision_diag
         Vs = self.V * d.sqrt().reshape(-1, 1)
         VtV = Vs.T @ Vs
@@ -1156,6 +1496,7 @@ class DiagLaplace(ParametricLaplace):
     Mathematically, we have \\(P \\approx \\textrm{diag}(P)\\).
     See `BaseLaplace` for the full interface.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ('all', 'diag')
 
@@ -1163,7 +1504,7 @@ def _init_H(self):
         self.H = torch.zeros(self.n_params, device=self._device)
 
     def _curv_closure(self, X, y, N):
-        return self.backend.diag(X, y, N=N)
+        return self.backend.diag(X, y, N=N, **self._asdl_fisher_kwargs)
 
     @property
     def posterior_precision(self):
@@ -1214,12 +1555,14 @@ def functional_variance(self, Js: torch.Tensor) -> torch.Tensor:
     def functional_covariance(self, Js):
         self._check_jacobians(Js)
         n_batch, n_outs, n_params = Js.shape
-        Js = Js.reshape(n_batch*n_outs, n_params)
+        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):
-        samples = torch.randn(n_samples, self.n_params, device=self._device)
+    def sample(self, n_samples=100, generator=None):
+        samples = torch.randn(
+            n_samples, self.n_params, device=self._device, generator=generator
+        )
         samples = samples * self.posterior_scale.reshape(1, self.n_params)
         return self.mean.reshape(1, self.n_params) + samples
 
diff --git a/laplace/curvature/__init__.py b/laplace/curvature/__init__.py
index c6f182fa..7bc83ab9 100644
--- a/laplace/curvature/__init__.py
+++ b/laplace/curvature/__init__.py
@@ -5,19 +5,50 @@
 try:
     from laplace.curvature.backpack import BackPackGGN, BackPackEF, BackPackInterface
 except ModuleNotFoundError:
-    logging.info('Backpack not available.')
+    logging.info('Backpack backend not available.')
+
+try:
+    from laplace.curvature.asdfghjkl import (
+        AsdfghjklHessian,
+        AsdfghjklGGN,
+        AsdfghjklEF,
+        AsdfghjklInterface,
+    )
+except ModuleNotFoundError:
+    logging.info('Asdfghjkl backend not available.')
 
 try:
     from laplace.curvature.asdl import AsdlHessian, AsdlGGN, AsdlEF, AsdlInterface
 except ModuleNotFoundError:
-    logging.info('asdfghjkl backend not available.')
+    logging.info('ASDL backend not available.')
 
 try:
-    from laplace.curvature.curvlinops import CurvlinopsHessian, CurvlinopsGGN, CurvlinopsEF, CurvlinopsInterface
+    from laplace.curvature.curvlinops import (
+        CurvlinopsHessian,
+        CurvlinopsGGN,
+        CurvlinopsEF,
+        CurvlinopsInterface,
+    )
 except ModuleNotFoundError:
-    logging.info('curvlinops backend not available.')
+    logging.info('Curvlinops backend not available.')
 
-__all__ = ['CurvatureInterface', 'GGNInterface', 'EFInterface',
-           'BackPackInterface', 'BackPackGGN', 'BackPackEF',
-           'AsdlInterface', 'AsdlGGN', 'AsdlEF', 'AsdlHessian',
-           'CurvlinopsInterface', 'CurvlinopsGGN', 'CurvlinopsEF', 'CurvlinopsHessian']
+__all__ = [
+    'CurvatureInterface',
+    'GGNInterface',
+    'EFInterface',
+    'BackPackInterface',
+    'BackPackGGN',
+    'BackPackEF',
+    'AsdfghjklInterface',
+    'AsdfghjklGGN',
+    'AsdfghjklEF',
+    'AsdfghjklHessian',
+    'AsdlInterface',
+    'AsdlGGN',
+    'AsdlEF',
+    'AsdlHessian',
+    'CurvlinopsInterface',
+    'CurvlinopsGGN',
+    'CurvlinopsEF',
+    'CurvlinopsHessian',
+]
diff --git a/laplace/curvature/asdfghjkl.py b/laplace/curvature/asdfghjkl.py
new file mode 100644
index 00000000..7b0d9a4b
--- /dev/null
+++ b/laplace/curvature/asdfghjkl.py
@@ -0,0 +1,233 @@
+import warnings
+import numpy as np
+import torch
+
+from asdfghjkl import FISHER_EXACT, FISHER_MC, COV
+from asdfghjkl import SHAPE_KRON, SHAPE_DIAG, SHAPE_FULL
+from asdfghjkl import fisher_for_cross_entropy
+from asdfghjkl.hessian import hessian_eigenvalues, hessian_for_loss
+from asdfghjkl.gradient import batch_gradient
+
+from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
+from laplace.utils import Kron, _is_batchnorm
+
+EPS = 1e-6
+
+
+class AsdfghjklInterface(CurvatureInterface):
+    """Interface for asdfghjkl backend."""
+
+    def jacobians(self, x, enable_backprop=False):
+        """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
+        using asdfghjkl's gradient per output dimension.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+        enable_backprop : bool, default = False
+            whether to enable backprop through the Js and f w.r.t. x
+
+        Returns
+        -------
+        Js : torch.Tensor
+            Jacobians `(batch, parameters, outputs)`
+        f : torch.Tensor
+            output function `(batch, outputs)`
+        """
+        Js = list()
+        for i in range(self.model.output_size):
+
+            def loss_fn(outputs, targets):
+                return outputs[:, i].sum()
+
+            f = batch_gradient(self.model, loss_fn, x, None).detach()
+            Jk = _get_batch_grad(self.model)
+            if self.subnetwork_indices is not None:
+                Jk = Jk[:, self.subnetwork_indices]
+            Js.append(Jk)
+        Js = torch.stack(Js, dim=1)
+        return Js, f
+
+    def gradients(self, x, y):
+        """Compute gradients \\(\\nabla_\\theta \\ell(f(x;\\theta, y)\\) at current parameter
+        \\(\\theta\\) using asdfghjkl's backend.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+        y : torch.Tensor
+
+        Returns
+        -------
+        loss : torch.Tensor
+        Gs : torch.Tensor
+            gradients `(batch, parameters)`
+        """
+        f = batch_gradient(self.model, self.lossfunc, x, y).detach()
+        Gs = _get_batch_grad(self._model)
+        if self.subnetwork_indices is not None:
+            Gs = Gs[:, self.subnetwork_indices]
+        loss = self.lossfunc(f, y)
+        return Gs, loss
+
+    @property
+    def _ggn_type(self):
+        raise NotImplementedError
+
+    def _get_kron_factors(self, curv, M):
+        kfacs = list()
+        for module in curv._model.modules():
+            if _is_batchnorm(module):
+                warnings.warn('BatchNorm unsupported for Kron, ignore.')
+                continue
+
+            stats = getattr(module, self._ggn_type, None)
+            if stats is None:
+                continue
+            if hasattr(module, 'bias') and module.bias is not None:
+                # split up bias and weights
+                kfacs.append([stats.kron.B, stats.kron.A[:-1, :-1]])
+                kfacs.append([stats.kron.B * stats.kron.A[-1, -1] / M])
+            elif hasattr(module, 'weight'):
+                p, q = stats.kron.B.numel(), stats.kron.A.numel()
+                if p == q == 1:
+                    kfacs.append([stats.kron.B * stats.kron.A])
+                else:
+                    kfacs.append([stats.kron.B, stats.kron.A])
+            else:
+                raise ValueError(f'Whats happening with {module}?')
+        return Kron(kfacs)
+
+    @staticmethod
+    def _rescale_kron_factors(kron, N):
+        for F in kron.kfacs:
+            if len(F) == 2:
+                F[1] *= 1 / N
+        return kron
+
+    def diag(self, X, y, **kwargs):
+        with torch.no_grad():
+            if self.last_layer:
+                f, X = self.model.forward_with_features(X)
+            else:
+                f = self.model(X)
+            loss = self.lossfunc(f, y)
+        curv = fisher_for_cross_entropy(
+            self._model, self._ggn_type, SHAPE_DIAG, inputs=X, targets=y
+        )
+        diag_ggn = curv.matrices_to_vector(None)
+        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):
+        with torch.no_grad():
+            if self.last_layer:
+                f, X = self.model.forward_with_features(X)
+            else:
+                f = self.model(X)
+            loss = self.lossfunc(f, y)
+        curv = fisher_for_cross_entropy(
+            self._model, self._ggn_type, SHAPE_KRON, inputs=X, targets=y
+        )
+        M = len(y)
+        kron = self._get_kron_factors(curv, M)
+        kron = self._rescale_kron_factors(kron, N)
+        return self.factor * loss, self.factor * kron
+
+
+class AsdfghjklHessian(AsdfghjklInterface):
+    def __init__(self, model, likelihood, last_layer=False, low_rank=10):
+        super().__init__(model, likelihood, last_layer)
+        self.low_rank = low_rank
+
+    @property
+    def _ggn_type(self):
+        raise NotImplementedError()
+
+    def full(self, x, y, **kwargs):
+        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):
+        # compute truncated eigendecomposition of the Hessian, only keep eigvals > EPS
+        eigvals, eigvecs = hessian_eigenvalues(
+            self.model,
+            self.lossfunc,
+            data_loader,
+            top_n=self.low_rank,
+            max_iters=self.low_rank * 10,
+        )
+        eigvals = torch.from_numpy(np.array(eigvals))
+        mask = eigvals > EPS
+        eigvecs = torch.stack(
+            [torch.cat([p.flatten() for p in params]) for params in eigvecs], dim=1
+        )[:, mask]
+        device = eigvecs.device
+        eigvals = eigvals[mask].to(eigvecs.dtype).to(device)
+        loss = sum(
+            [
+                self.lossfunc(self.model(x.to(device)).detach(), y.to(device))
+                for x, y in data_loader
+            ]
+        )
+        return eigvecs, self.factor * eigvals, self.factor * loss
+
+
+class AsdfghjklGGN(AsdfghjklInterface, GGNInterface):
+    """Implementation of the `GGNInterface` using asdfghjkl."""
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        last_layer=False,
+        subnetwork_indices=None,
+        stochastic=False,
+    ):
+        if likelihood != 'classification':
+            raise ValueError('This backend only supports classification currently.')
+        super().__init__(model, likelihood, last_layer, subnetwork_indices)
+        self.stochastic = stochastic
+
+    @property
+    def _ggn_type(self):
+        return FISHER_MC if self.stochastic else FISHER_EXACT
+
+
+class AsdfghjklEF(AsdfghjklInterface, EFInterface):
+    """Implementation of the `EFInterface` using asdfghjkl."""
+
+    def __init__(self, model, likelihood, last_layer=False):
+        if likelihood != 'classification':
+            raise ValueError('This backend only supports classification currently.')
+        super().__init__(model, likelihood, last_layer)
+
+    @property
+    def _ggn_type(self):
+        return COV
+
+
+def _flatten_after_batch(tensor: torch.Tensor):
+    if tensor.ndim == 1:
+        return tensor.unsqueeze(-1)
+    else:
+        return tensor.flatten(start_dim=1)
+
+
+def _get_batch_grad(model):
+    batch_grads = list()
+    for module in model.modules():
+        if hasattr(module, 'op_results'):
+            res = module.op_results['batch_grads']
+            if 'weight' in res:
+                batch_grads.append(_flatten_after_batch(res['weight']))
+            if 'bias' in res:
+                batch_grads.append(_flatten_after_batch(res['bias']))
+            if len(set(res.keys()) - {'weight', 'bias'}) > 0:
+                raise ValueError(f'Invalid parameter keys {res.keys()}')
+    return torch.cat(batch_grads, dim=1)
diff --git a/laplace/curvature/asdl.py b/laplace/curvature/asdl.py
index b7ac75de..f0a88397 100644
--- a/laplace/curvature/asdl.py
+++ b/laplace/curvature/asdl.py
@@ -1,22 +1,38 @@
 import warnings
+
 import numpy as np
 import torch
 
-from asdfghjkl import FISHER_EXACT, FISHER_MC, COV
-from asdfghjkl import SHAPE_KRON, SHAPE_DIAG, SHAPE_FULL
-from asdfghjkl import fisher_for_cross_entropy
-from asdfghjkl.hessian import hessian_eigenvalues, hessian_for_loss
-from asdfghjkl.gradient import batch_gradient
+from asdl.matrices import (
+    FISHER_EXACT,
+    FISHER_MC,
+    FISHER_EMP,
+    SHAPE_KRON,
+    SHAPE_DIAG,
+    SHAPE_FULL,
+)
+from asdl.grad_maker import LOSS_MSE, LOSS_CROSS_ENTROPY
+from asdl.fisher import FisherConfig, get_fisher_maker
+from asdl.hessian import HessianMaker, HessianConfig
+from asdl.gradient import batch_gradient
 
 from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
 from laplace.utils import Kron, _is_batchnorm
 
+from collections import UserDict
+
 EPS = 1e-6
 
 
 class AsdlInterface(CurvatureInterface):
-    """Interface for asdfghjkl backend.
-    """
+    """Interface for asdfghjkl backend."""
+
+    def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None):
+        super().__init__(model, likelihood, last_layer, subnetwork_indices)
+
+    @property
+    def loss_type(self):
+        return LOSS_MSE if self.likelihood == 'regression' else LOSS_CROSS_ENTROPY
 
     def jacobians(self, x, enable_backprop=False):
         """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\)
@@ -24,8 +40,10 @@ def jacobians(self, x, enable_backprop=False):
 
         Parameters
         ----------
-        x : torch.Tensor
-            input data `(batch, input_shape)` on compatible device with model.
+        x : torch.Tensor or UserDict
+            input data `(batch, input_shape)` on compatible device with model if torch.Tensor.
+            If UserDict, then at least contains key ['input_ids'] or ['input_ids_0', 'input_ids_1'].
+            The latter is specific for reward modeling.
         enable_backprop : bool, default = False
             whether to enable backprop through the Js and f w.r.t. x
 
@@ -38,14 +56,25 @@ def jacobians(self, x, enable_backprop=False):
         """
         Js = list()
         for i in range(self.model.output_size):
-            def loss_fn(outputs, targets):
-                return outputs[:, i].sum()
 
-            f = batch_gradient(self.model, loss_fn, x, None).detach()
-            Jk = _get_batch_grad(self.model)
+            def closure():
+                self.model.zero_grad()
+                f = self.model(x)
+                loss = f[:, i].sum()
+                loss.backward(
+                    create_graph=enable_backprop, retain_graph=enable_backprop
+                )
+                return f
+
+            Ji, f = batch_gradient(
+                self.model,
+                closure,
+                return_outputs=True,
+                batch_size=self._get_batch_size(x),
+            )
             if self.subnetwork_indices is not None:
-                Jk = Jk[:, self.subnetwork_indices]
-            Js.append(Jk)
+                Ji = Ji[:, self.subnetwork_indices]
+            Js.append(Ji)
         Js = torch.stack(Js, dim=1)
         return Js, f
 
@@ -65,31 +94,39 @@ def gradients(self, x, y):
         Gs : torch.Tensor
             gradients `(batch, parameters)`
         """
-        f = batch_gradient(self.model, self.lossfunc, x, y).detach()
-        Gs = _get_batch_grad(self._model)
+
+        def closure():
+            self.model.zero_grad()
+            loss = self.lossfunc(self.model(x), y)
+            loss.backward()
+            return loss
+
+        Gs, loss = batch_gradient(
+            self.model, closure, return_outputs=True, batch_size=self._get_batch_size(x)
+        )
         if self.subnetwork_indices is not None:
             Gs = Gs[:, self.subnetwork_indices]
-        loss = self.lossfunc(f, y)
         return Gs, loss
 
     @property
     def _ggn_type(self):
         raise NotImplementedError
 
-    def _get_kron_factors(self, curv, M):
+    def _get_kron_factors(self, M):
         kfacs = list()
-        for module in curv._model.modules():
+        for module in self.model.modules():
             if _is_batchnorm(module):
                 warnings.warn('BatchNorm unsupported for Kron, ignore.')
                 continue
 
-            stats = getattr(module, self._ggn_type, None)
+            stats = getattr(module, 'fisher', None)
             if stats is None:
                 continue
+
             if hasattr(module, 'bias') and module.bias is not None:
                 # split up bias and weights
-                kfacs.append([stats.kron.B, stats.kron.A[:-1, :-1]])
-                kfacs.append([stats.kron.B * stats.kron.A[-1, -1] / M])
+                kfacs.append([stats.kron.B, stats.kron.A])
+                kfacs.append([stats.kron.B])
             elif hasattr(module, 'weight'):
                 p, q = stats.kron.B.numel(), stats.kron.A.numel()
                 if p == q == 1:
@@ -104,40 +141,101 @@ def _get_kron_factors(self, curv, M):
     def _rescale_kron_factors(kron, N):
         for F in kron.kfacs:
             if len(F) == 2:
-                F[1] *= 1/N
+                F[1] *= 1 / N
         return kron
 
-    def diag(self, X, y, **kwargs):
-        with torch.no_grad():
-            if self.last_layer:
-                f, X = self.model.forward_with_features(X)
-            else:
-                f = self.model(X)
-            loss = self.lossfunc(f, y)
-        curv = fisher_for_cross_entropy(self._model, self._ggn_type, SHAPE_DIAG,
-                                        inputs=X, targets=y)
-        diag_ggn = curv.matrices_to_vector(None)
+    def diag(self, X, y, N=None, **kwargs):
+        del N
+        if self.last_layer:
+            _, X = self.model.forward_with_features(X)
+        cfg = FisherConfig(
+            fisher_type=self._ggn_type,
+            loss_type=self.loss_type,
+            fisher_shapes=[SHAPE_DIAG],
+            data_size=1,
+            **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 = (
+                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)
+        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))
+        loss = self.lossfunc(f.detach(), y)
+        vec = list()
+        for module in self.model.modules():
+            stats = getattr(module, 'fisher', None)
+            if stats is None:
+                continue
+            vec.extend(stats.to_vector())
+        diag_ggn = torch.cat(vec)
         if self.subnetwork_indices is not None:
             diag_ggn = diag_ggn[self.subnetwork_indices]
-        return self.factor * loss, self.factor * diag_ggn
+        if type(self) is AsdlEF and self.likelihood == 'regression':
+            curv_factor = 0.5  # correct scaling for diag ef
+        else:
+            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):
-        with torch.no_grad():
-            if self.last_layer:
-                f, X = self.model.forward_with_features(X)
-            else:
-                f = self.model(X)
-            loss = self.lossfunc(f, y)
-        curv = fisher_for_cross_entropy(self._model, self._ggn_type, SHAPE_KRON,
-                                        inputs=X, targets=y)
+        if self.last_layer:
+            _, X = self.model.forward_with_features(X)
+        cfg = FisherConfig(
+            fisher_type=self._ggn_type,
+            loss_type=self.loss_type,
+            fisher_shapes=[SHAPE_KRON],
+            data_size=1,
+            **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 = (
+                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)
+        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))
+        loss = self.lossfunc(f.detach(), y)
         M = len(y)
-        kron = self._get_kron_factors(curv, M)
+        kron = self._get_kron_factors(M)
         kron = self._rescale_kron_factors(kron, N)
-        return self.factor * loss, self.factor * kron
+        if type(self) is AsdlEF and self.likelihood == 'regression':
+            curv_factor = 0.5  # correct scaling for diag ef
+        else:
+            curv_factor = 1.0  # ASDL uses proper 1/2 * MSELoss
+        return self.factor * loss, curv_factor * kron
+
+    @staticmethod
+    def _get_batch_size(x):
+        """
+        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.
+        This is the case for LLMs.
+        """
+        # If x is UserDict, then it has weight-sharing dimension (from Huggingface datasets)
+        if isinstance(x, UserDict) or isinstance(x, dict):
+            try:
+                return x['input_ids'].shape[0]
+            except KeyError:
+                # The case of reward modeling; the UserDict contains ['input_ids_0', 'input_ids_1']
+                return x['input_ids_0'].shape[0]
+        else:
+            return None  # Use ASDL default behavior
 
 
 class AsdlHessian(AsdlInterface):
-
     def __init__(self, model, likelihood, last_layer=False, low_rank=10):
         super().__init__(model, likelihood, last_layer)
         self.low_rank = low_rank
@@ -147,31 +245,34 @@ def _ggn_type(self):
         raise NotImplementedError()
 
     def full(self, x, y, **kwargs):
-        hessian_for_loss(self.model, self.lossfunc, SHAPE_FULL, x, y)
+        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
-        loss = self.lossfunc(self.model(x), y).detach()
+        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
 
-    def eig_lowrank(self, data_loader):
-        # compute truncated eigendecomposition of the Hessian, only keep eigvals > EPS
-        eigvals, eigvecs = hessian_eigenvalues(self.model, self.lossfunc, data_loader,
-                                               top_n=self.low_rank, max_iters=self.low_rank*10)
-        eigvals = torch.from_numpy(np.array(eigvals))
-        mask = (eigvals > EPS)
-        eigvecs = torch.stack([torch.cat([p.flatten() for p in params])
-                               for params in eigvecs], dim=1)[:, mask]
-        device = eigvecs.device
-        eigvals = eigvals[mask].to(eigvecs.dtype).to(device)
-        loss = sum([self.lossfunc(self.model(x.to(device)).detach(), y.to(device)) for x, y in data_loader])
-        return eigvecs, self.factor * eigvals, self.factor * loss
-
 
 class AsdlGGN(AsdlInterface, GGNInterface):
-    """Implementation of the `GGNInterface` using asdfghjkl.
-    """
-    def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None, stochastic=False):
-        if likelihood != 'classification':
-            raise ValueError('This backend only supports classification currently.')
+    """Implementation of the `GGNInterface` using asdfghjkl."""
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        last_layer=False,
+        subnetwork_indices=None,
+        stochastic=False,
+    ):
         super().__init__(model, likelihood, last_layer, subnetwork_indices)
         self.stochastic = stochastic
 
@@ -181,34 +282,11 @@ def _ggn_type(self):
 
 
 class AsdlEF(AsdlInterface, EFInterface):
-    """Implementation of the `EFInterface` using asdfghjkl.
-    """
+    """Implementation of the `EFInterface` using asdfghjkl."""
+
     def __init__(self, model, likelihood, last_layer=False):
-        if likelihood != 'classification':
-            raise ValueError('This backend only supports classification currently.')
         super().__init__(model, likelihood, last_layer)
 
     @property
     def _ggn_type(self):
-        return COV
-
-
-def _flatten_after_batch(tensor: torch.Tensor):
-    if tensor.ndim == 1:
-        return tensor.unsqueeze(-1)
-    else:
-        return tensor.flatten(start_dim=1)
-
-
-def _get_batch_grad(model):
-    batch_grads = list()
-    for module in model.modules():
-        if hasattr(module, 'op_results'):
-            res = module.op_results['batch_grads']
-            if 'weight' in res:
-                batch_grads.append(_flatten_after_batch(res['weight']))
-            if 'bias' in res:
-                batch_grads.append(_flatten_after_batch(res['bias']))
-            if len(set(res.keys()) - {'weight', 'bias'}) > 0:
-                raise ValueError(f'Invalid parameter keys {res.keys()}')
-    return torch.cat(batch_grads, dim=1)
+        return FISHER_EMP
diff --git a/laplace/curvature/backpack.py b/laplace/curvature/backpack.py
index a78057ff..ee919417 100644
--- a/laplace/curvature/backpack.py
+++ b/laplace/curvature/backpack.py
@@ -2,7 +2,14 @@
 import torch
 
 from backpack import backpack, extend, memory_cleanup
-from backpack.extensions import DiagGGNExact, DiagGGNMC, KFAC, KFLR, SumGradSquared, BatchGrad
+from backpack.extensions import (
+    DiagGGNExact,
+    DiagGGNMC,
+    KFAC,
+    KFLR,
+    SumGradSquared,
+    BatchGrad,
+)
 from backpack.context import CTX
 
 from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
@@ -10,8 +17,8 @@
 
 
 class BackPackInterface(CurvatureInterface):
-    """Interface for Backpack backend.
-    """
+    """Interface for Backpack backend."""
+
     def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None):
         super().__init__(model, likelihood, last_layer, subnetwork_indices)
         extend(self._model)
@@ -44,13 +51,11 @@ def jacobians(self, x, enable_backprop=False):
             with backpack(BatchGrad()):
                 if model.output_size > 1:
                     out[:, i].sum().backward(
-                        create_graph=enable_backprop,
-                        retain_graph=enable_backprop
+                        create_graph=enable_backprop, retain_graph=enable_backprop
                     )
                 else:
                     out.sum().backward(
-                        create_graph=enable_backprop,
-                        retain_graph=enable_backprop
+                        create_graph=enable_backprop, retain_graph=enable_backprop
                     )
                 to_cat = []
                 for param in model.parameters():
@@ -92,25 +97,38 @@ def gradients(self, x, y):
         loss = self.lossfunc(f, y)
         with backpack(BatchGrad()):
             loss.backward()
-        Gs = torch.cat([p.grad_batch.data.flatten(start_dim=1)
-                        for p in self._model.parameters()], dim=1)
+        Gs = torch.cat(
+            [p.grad_batch.data.flatten(start_dim=1) for p in self._model.parameters()],
+            dim=1,
+        )
         if self.subnetwork_indices is not None:
             Gs = Gs[:, self.subnetwork_indices]
         return Gs, loss
 
 
 class BackPackGGN(BackPackInterface, GGNInterface):
-    """Implementation of the `GGNInterface` using Backpack.
-    """
-    def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None, stochastic=False):
+    """Implementation of the `GGNInterface` using Backpack."""
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        last_layer=False,
+        subnetwork_indices=None,
+        stochastic=False,
+    ):
         super().__init__(model, likelihood, last_layer, subnetwork_indices)
         self.stochastic = stochastic
 
     def _get_diag_ggn(self):
         if self.stochastic:
-            return torch.cat([p.diag_ggn_mc.data.flatten() for p in self._model.parameters()])
+            return torch.cat(
+                [p.diag_ggn_mc.data.flatten() for p in self._model.parameters()]
+            )
         else:
-            return torch.cat([p.diag_ggn_exact.data.flatten() for p in self._model.parameters()])
+            return torch.cat(
+                [p.diag_ggn_exact.data.flatten() for p in self._model.parameters()]
+            )
 
     def _get_kron_factors(self):
         if self.stochastic:
@@ -124,12 +142,15 @@ def _rescale_kron_factors(kron, M, N):
         # for M=N (full-batch) just M/N=1
         for F in kron.kfacs:
             if len(F) == 2:
-                F[1] *= M/N
+                F[1] *= M / N
         return kron
 
     def diag(self, X, y, **kwargs):
         context = DiagGGNMC if self.stochastic else DiagGGNExact
         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)
         loss = self.lossfunc(f, y)
         with backpack(context()):
             loss.backward()
@@ -142,6 +163,9 @@ def diag(self, X, y, **kwargs):
     def kron(self, X, y, N, **kwargs) -> Tuple[torch.Tensor, Kron]:
         context = KFAC if self.stochastic else KFLR
         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)
         loss = self.lossfunc(f, y)
         with backpack(context()):
             loss.backward()
@@ -152,16 +176,19 @@ def kron(self, X, y, N, **kwargs) -> Tuple[torch.Tensor, Kron]:
 
 
 class BackPackEF(BackPackInterface, EFInterface):
-    """Implementation of `EFInterface` using Backpack.
-    """
+    """Implementation of `EFInterface` using Backpack."""
 
     def diag(self, X, y, **kwargs):
         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)
         loss = self.lossfunc(f, y)
         with backpack(SumGradSquared()):
             loss.backward()
-        diag_EF = torch.cat([p.sum_grad_squared.data.flatten()
-                             for p in self._model.parameters()])
+        diag_EF = torch.cat(
+            [p.sum_grad_squared.data.flatten() for p in self._model.parameters()]
+        )
         if self.subnetwork_indices is not None:
             diag_EF = diag_EF[self.subnetwork_indices]
 
@@ -175,5 +202,5 @@ def _cleanup(module):
     for child in module.children():
         _cleanup(child)
 
-    setattr(module, "_backpack_extend", False)
+    setattr(module, '_backpack_extend', False)
     memory_cleanup(module)
diff --git a/laplace/curvature/curvature.py b/laplace/curvature/curvature.py
index 8c83300c..0dde26d3 100644
--- a/laplace/curvature/curvature.py
+++ b/laplace/curvature/curvature.py
@@ -28,6 +28,7 @@ class CurvatureInterface:
         conversion factor between torch losses and base likelihoods
         For example, \\(\\frac{1}{2}\\) to get to \\(\\mathcal{N}(f, 1)\\) from MSELoss.
     """
+
     def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None):
         assert likelihood in ['regression', 'classification']
         self.likelihood = likelihood
@@ -39,9 +40,11 @@ def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None)
             self.factor = 0.5
         else:
             self.lossfunc = CrossEntropyLoss(reduction='sum')
-            self.factor = 1.
+            self.factor = 1.0
         self.params = [p for p in self._model.parameters() if p.requires_grad]
-        self.params_dict = {k: v for k, v in self._model.named_parameters() if v.requires_grad}
+        self.params_dict = {
+            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()}
 
     @property
@@ -66,11 +69,14 @@ def jacobians(self, x, enable_backprop=False):
         f : torch.Tensor
             output function `(batch, outputs)`
         """
+
         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)
+        Js, f = torch.func.jacrev(model_fn_params_only, has_aux=True)(
+            self.params_dict, self.buffers_dict
+        )
 
         # Concatenate over flattened parameters
         Js = [
@@ -84,6 +90,46 @@ def model_fn_params_only(params_dict, buffers_dict):
 
         return (Js, f) if enable_backprop else (Js.detach(), f.detach())
 
+    def functorch_jacobians(self, x, enable_backprop=False):
+        """Compute Jacobians \\(\\nabla_\\theta f(x;\\theta)\\) at current parameter \\(\\theta\\).
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            input data `(batch, input_shape)` on compatible device with model.
+        enable_backprop : bool, default = False
+            whether to enable backprop through the Js and f w.r.t. x
+
+        Returns
+        -------
+        Js : torch.Tensor
+            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)
+        Js = [
+            j.flatten(start_dim=-p.dim())
+            for j, p in zip(Js.values(), params_dict.values())
+        ]
+        Js = torch.cat(Js, dim=-1)
+
+        if self.subnetwork_indices is not None:
+            Js = Js[:, :, self.subnetwork_indices]
+
+        return (Js, f) if enable_backprop else (Js.detach(), f.detach())
+
     def last_layer_jacobians(self, x, enable_backprop=False):
         """Compute Jacobians \\(\\nabla_{\\theta_\\textrm{last}} f(x;\\theta_\\textrm{last})\\)
         only at current last-layer parameter \\(\\theta_{\\textrm{last}}\\).
@@ -105,7 +151,11 @@ def last_layer_jacobians(self, x, enable_backprop=False):
         output_size = int(f.numel() / bsize)
 
         # calculate Jacobians using the feature vector 'phi'
-        identity = torch.eye(output_size, device=x.device).unsqueeze(0).tile(bsize, 1, 1)
+        identity = (
+            torch.eye(output_size, device=next(self.model.parameters()).device)
+            .unsqueeze(0)
+            .tile(bsize, 1, 1)
+        )
         # Jacobians are batch x output x params
         Js = torch.einsum('kp,kij->kijp', phi, identity).reshape(bsize, output_size, -1)
         if self.model.last_layer.bias is not None:
@@ -129,17 +179,22 @@ def gradients(self, x, y):
             gradients `(batch, parameters)`
         loss : torch.Tensor
         """
+
         def loss_single(x, y, params_dict, buffers_dict):
             """Compute the gradient for a single sample."""
             x, y = x.unsqueeze(0), y.unsqueeze(0)  # vmap removes the batch dimension
-            output = torch.func.functional_call(self.model, (params_dict, buffers_dict), x)
+            output = torch.func.functional_call(
+                self.model, (params_dict, buffers_dict), x
+            )
             loss = torch.func.functional_call(self.lossfunc, {}, (output, y))
             return loss, loss
 
         grad_fn = torch.func.grad(loss_single, argnums=2, has_aux=True)
         batch_grad_fn = torch.func.vmap(grad_fn, in_dims=(0, 0, None, None))
 
-        batch_grad, batch_loss = batch_grad_fn(x, y, self.params_dict, self.buffers_dict)
+        batch_grad, batch_loss = batch_grad_fn(
+            x, y, self.params_dict, self.buffers_dict
+        )
         Gs = torch.cat([bg.flatten(start_dim=1) for bg in batch_grad.values()], dim=1)
 
         if self.subnetwork_indices is not None:
@@ -231,13 +286,22 @@ class GGNInterface(CurvatureInterface):
     num_samples: int, default=100
         Number of samples used to approximate the stochastic Fisher
     """
-    def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None, stochastic=False, num_samples=1):
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        last_layer=False,
+        subnetwork_indices=None,
+        stochastic=False,
+        num_samples=1,
+    ):
         self.stochastic = stochastic
         self.num_samples = num_samples
         super().__init__(model, likelihood, last_layer, subnetwork_indices)
 
     def _get_mc_functional_fisher(self, f):
-        """ Approximate the Fisher's middle matrix (expected outer product of the functional gradient)
+        """Approximate the Fisher's middle matrix (expected outer product of the functional gradient)
         using MC integral with `self.num_samples` many samples.
         """
         F = 0
@@ -252,7 +316,11 @@ def _get_mc_functional_fisher(self, f):
                 p = torch.softmax(f, dim=-1)
                 grad_sample = p - y_sample
 
-            F += 1/self.num_samples * torch.einsum('bc,bk->bck', grad_sample, grad_sample)
+            F += (
+                1
+                / self.num_samples
+                * torch.einsum('bc,bk->bck', grad_sample, grad_sample)
+            )
 
         return F
 
@@ -284,7 +352,11 @@ def full(self, x, y, **kwargs):
             GGN `(parameters, parameters)`
         """
         Js, f = self.last_layer_jacobians(x) if self.last_layer else self.jacobians(x)
-        H_lik = self._get_mc_functional_fisher(f) if self.stochastic else self._get_functional_hessian(f)
+        H_lik = (
+            self._get_mc_functional_fisher(f)
+            if self.stochastic
+            else self._get_functional_hessian(f)
+        )
 
         if H_lik is not None:
             H = torch.einsum('bcp,bck,bkq->pq', Js, H_lik, Js)
@@ -298,7 +370,11 @@ def diag(self, X, y, **kwargs):
         Js, f = self.last_layer_jacobians(X) if self.last_layer else self.jacobians(X)
         loss = self.factor * self.lossfunc(f, y)
 
-        H_lik = self._get_mc_functional_fisher(f) if self.stochastic else self._get_functional_hessian(f)
+        H_lik = (
+            self._get_mc_functional_fisher(f)
+            if self.stochastic
+            else self._get_functional_hessian(f)
+        )
 
         if H_lik is not None:
             H = torch.einsum('bcp,bck,bkp->p', Js, H_lik, Js)
diff --git a/laplace/curvature/curvlinops.py b/laplace/curvature/curvlinops.py
index 805a7c7c..20b1e49e 100644
--- a/laplace/curvature/curvlinops.py
+++ b/laplace/curvature/curvlinops.py
@@ -2,17 +2,22 @@
 import numpy as np
 
 from curvlinops import (
-    HessianLinearOperator, GGNLinearOperator, FisherMCLinearOperator, EFLinearOperator,
-    KFACLinearOperator
+    HessianLinearOperator,
+    GGNLinearOperator,
+    FisherMCLinearOperator,
+    EFLinearOperator,
+    KFACLinearOperator,
 )
 
 from laplace.curvature import CurvatureInterface, GGNInterface, EFInterface
 from laplace.utils import Kron
 
+from collections import UserDict
+
 
 class CurvlinopsInterface(CurvatureInterface):
-    """Interface for Curvlinops backend. <https://github.com/f-dangel/curvlinops>
-    """
+    """Interface for Curvlinops backend. <https://github.com/f-dangel/curvlinops>"""
+
     def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None):
         super().__init__(model, likelihood, last_layer, subnetwork_indices)
 
@@ -30,7 +35,7 @@ def _rescale_kron_factors(kron, M, N):
         # for M=N (full-batch) just M/N=1
         for F in kron.kfacs:
             if len(F) == 2:
-                F[1] *= M/N
+                F[1] *= M / N
         return kron
 
     def _get_kron_factors(self, linop):
@@ -56,8 +61,13 @@ def _get_kron_factors(self, linop):
         return Kron(kfacs)
 
     def kron(self, X, y, N, **kwargs):
+        if isinstance(X, (dict, UserDict)):
+            kwargs['batch_size_fn'] = lambda x: x['input_ids'].shape[0]
         linop = KFACLinearOperator(
-            self.model, self.lossfunc, self.params, [(X, y)],
+            self.model,
+            self.lossfunc,
+            self.params,
+            [(X, y)],
             fisher_type=self._kron_fisher_type,
             loss_average=None,  # Since self.lossfunc is sum
             separate_weight_and_bias=True,
@@ -83,12 +93,20 @@ def full(self, X, y, **kwargs):
             return super().full(X, y, **kwargs)
 
         curvlinops_kwargs = {k: v for k, v in kwargs.items() if k != 'N'}
-        linop = self._linop_context(self.model, self.lossfunc, self.params, [(X, y)],
-                                    check_deterministic=False, **curvlinops_kwargs)
+        if isinstance(X, (dict, UserDict)):
+            curvlinops_kwargs['batch_size_fn'] = lambda x: x['input_ids'].shape[0]
+
+        linop = self._linop_context(
+            self.model,
+            self.lossfunc,
+            self.params,
+            [(X, y)],
+            check_deterministic=False,
+            **curvlinops_kwargs,
+        )
         H = torch.as_tensor(
-            linop @ np.eye(linop.shape[0]),
-            dtype=X.dtype,
-            device=X.device
+            linop @ torch.eye(linop.shape[0]),
+            device=next(self.model.parameters()).device,
         )
 
         f = self.model(X)
@@ -99,7 +117,15 @@ def full(self, X, y, **kwargs):
 
 class CurvlinopsGGN(CurvlinopsInterface, GGNInterface):
     """Implementation of the `GGNInterface` using Curvlinops."""
-    def __init__(self, model, likelihood, last_layer=False, subnetwork_indices=None, stochastic=False):
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        last_layer=False,
+        subnetwork_indices=None,
+        stochastic=False,
+    ):
         super().__init__(model, likelihood, last_layer, subnetwork_indices)
         self.stochastic = stochastic
 
diff --git a/laplace/laplace.py b/laplace/laplace.py
index bffc91c0..f8c45778 100644
--- a/laplace/laplace.py
+++ b/laplace/laplace.py
@@ -2,8 +2,14 @@
 from laplace import *
 
 
-def Laplace(model, likelihood, subset_of_weights='last_layer', hessian_structure='kron',
-            *args, **kwargs):
+def Laplace(
+    model,
+    likelihood,
+    subset_of_weights='last_layer',
+    hessian_structure='kron',
+    *args,
+    **kwargs,
+):
     """Simplified Laplace access using strings instead of different classes.
 
     Parameters
@@ -21,14 +27,20 @@ def Laplace(model, likelihood, subset_of_weights='last_layer', hessian_structure
         chosen subclass of ParametricLaplace instantiated with additional arguments
     """
     if subset_of_weights == 'subnetwork' and hessian_structure not in ['full', 'diag']:
-        raise ValueError('Subnetwork Laplace requires a full or diagonal Hessian approximation!')
+        raise ValueError(
+            'Subnetwork Laplace requires a full or diagonal Hessian approximation!'
+        )
 
-    laplace_map = {subclass._key: subclass for subclass in _all_subclasses(ParametricLaplace)
-                   if hasattr(subclass, '_key')}
+    laplace_map = {
+        subclass._key: subclass
+        for subclass in _all_subclasses(ParametricLaplace)
+        if hasattr(subclass, '_key')
+    }
     laplace_class = laplace_map[(subset_of_weights, hessian_structure)]
     return laplace_class(model, likelihood, *args, **kwargs)
 
 
 def _all_subclasses(cls):
     return set(cls.__subclasses__()).union(
-        [s for c in cls.__subclasses__() for s in _all_subclasses(c)])
+        [s for c in cls.__subclasses__() for s in _all_subclasses(c)]
+    )
diff --git a/laplace/lllaplace.py b/laplace/lllaplace.py
index ead2b6af..c08e0bb5 100644
--- a/laplace/lllaplace.py
+++ b/laplace/lllaplace.py
@@ -5,6 +5,9 @@
 from laplace.baselaplace import ParametricLaplace, FullLaplace, KronLaplace, DiagLaplace
 from laplace.utils import FeatureExtractor, Kron
 
+from collections.abc import MutableMapping
+from typing import Union
+
 
 __all__ = ['LLLaplace', 'FullLLLaplace', 'KronLLLaplace', 'DiagLLLaplace']
 
@@ -57,17 +60,39 @@ class LLLaplace(ParametricLaplace):
         arguments passed to the backend on initialization, for example to
         set the number of MC samples for stochastic approximations.
     """
-    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., enable_backprop=False, backend=None, last_layer_name=None,
-                 backend_kwargs=None):
+
+    def __init__(
+        self,
+        model,
+        likelihood,
+        sigma_noise=1.0,
+        prior_precision=1.0,
+        prior_mean=0.0,
+        temperature=1.0,
+        enable_backprop=False,
+        backend=None,
+        last_layer_name=None,
+        backend_kwargs=None,
+        asdl_fisher_kwargs=None,
+    ):
+        if asdl_fisher_kwargs is not None:
+            raise ValueError('Last-layer Laplace does not support asdl_fisher_kwargs.')
         self.H = None
-        super().__init__(model, likelihood, sigma_noise=sigma_noise, prior_precision=1.,
-                         prior_mean=0., temperature=temperature,
-                         enable_backprop=enable_backprop, backend=backend,
-                         backend_kwargs=backend_kwargs)
+        super().__init__(
+            model,
+            likelihood,
+            sigma_noise=sigma_noise,
+            prior_precision=1.0,
+            prior_mean=0.0,
+            temperature=temperature,
+            enable_backprop=enable_backprop,
+            backend=backend,
+            backend_kwargs=backend_kwargs,
+        )
         self.model = FeatureExtractor(
-            deepcopy(model), last_layer_name=last_layer_name,
-            enable_backprop=enable_backprop
+            deepcopy(model),
+            last_layer_name=last_layer_name,
+            enable_backprop=enable_backprop,
         )
         if self.model.last_layer is None:
             self.mean = None
@@ -77,7 +102,9 @@ def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
             self._prior_precision = prior_precision
             self._prior_mean = prior_mean
         else:
-            self.n_params = len(parameters_to_vector(self.model.last_layer.parameters()))
+            self.n_params = 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
@@ -99,18 +126,15 @@ def fit(self, train_loader, override=True):
             online learning settings to accumulate a sequential posterior approximation.
         """
         if not override:
-            raise ValueError('Last-layer Laplace approximations do not support `override=False`.')
+            raise ValueError(
+                'Last-layer Laplace approximations do not support `override=False`.'
+            )
 
         self.model.eval()
 
         if self.model.last_layer is None:
-            # Save an example batch for when loading the serialized Laplace
-            self.X, _ = next(iter(train_loader))
-            with torch.no_grad():
-                try:
-                    self.model.find_last_layer(self.X[:1].to(self._device))
-                except (TypeError, AttributeError):
-                    self.model.find_last_layer(self.X.to(self._device))
+            self.data = 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()))
@@ -134,13 +158,17 @@ def _glm_predictive_distribution(self, X, joint=False):
         else:
             f_var = self.functional_variance(Js)
 
-        return (f_mu.detach(), f_var.detach()) if not self.enable_backprop else (f_mu, f_var)
+        return (
+            (f_mu.detach(), f_var.detach())
+            if not self.enable_backprop
+            else (f_mu, f_var)
+        )
 
-    def _nn_predictive_samples(self, X, n_samples=100):
+    def _nn_predictive_samples(self, X, n_samples=100, generator=None, **model_kwargs):
         fs = list()
-        for sample in self.sample(n_samples):
+        for sample in self.sample(n_samples, generator):
             vector_to_parameters(sample, self.model.last_layer.parameters())
-            f = self.model(X.to(self._device))
+            f = self.model(X.to(self._device), **model_kwargs)
             fs.append(f.detach() if not self.enable_backprop else f)
         vector_to_parameters(self.mean, self.model.last_layer.parameters())
         fs = torch.stack(fs)
@@ -148,6 +176,18 @@ def _nn_predictive_samples(self, X, n_samples=100):
             fs = torch.softmax(fs, dim=-1)
         return fs
 
+    def _nn_predictive_classification(
+        self, X, n_samples=100, generator=None, **model_kwargs
+    ):
+        py = 0
+        for sample in self.sample(n_samples, generator):
+            vector_to_parameters(sample, self.model.last_layer.parameters())
+            # TODO: Implement with a single forward pass until last layer.
+            logits = self.model(X.to(self._device), **model_kwargs).detach()
+            py += torch.softmax(logits, dim=-1) / n_samples
+        vector_to_parameters(self.mean, self.model.last_layer.parameters())
+        return py
+
     @property
     def prior_precision_diag(self):
         """Obtain the diagonal prior precision \\(p_0\\) constructed from either
@@ -168,7 +208,7 @@ def prior_precision_diag(self):
 
     def state_dict(self) -> dict:
         state_dict = super().state_dict()
-        state_dict['X'] = getattr(self, 'X', None)
+        state_dict['data'] = getattr(self, 'data', None)  # None if not present
         state_dict['_last_layer_name'] = self._last_layer_name
         return state_dict
 
@@ -176,13 +216,9 @@ def load_state_dict(self, state_dict: dict):
         if self._last_layer_name != state_dict['_last_layer_name']:
             raise ValueError('Different `last_layer_name` detected!')
 
-        self.X = state_dict['X']
-        if self.X is not None:
-            with torch.no_grad():
-                try:
-                    self.model.find_last_layer(self.X[:1].to(self._device))
-                except (TypeError, AttributeError):
-                    self.model.find_last_layer(self.X.to(self._device))
+        self.data = state_dict['data']
+        if self.data is not None:
+            self._find_last_layer(self.data)
 
         super().load_state_dict(state_dict)
 
@@ -190,6 +226,18 @@ def load_state_dict(self, state_dict: dict):
         self.n_params = len(params)
         self.n_layers = len(list(self.model.last_layer.parameters()))
 
+    @torch.no_grad()
+    def _find_last_layer(self, data: Union[torch.Tensor, MutableMapping]) -> 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):
+                self.model.find_last_layer(X.to(self._device))
+
 
 class FullLLLaplace(LLLaplace, FullLaplace):
     """Last-layer Laplace approximation with full, i.e., dense, log likelihood Hessian approximation
@@ -198,6 +246,7 @@ class FullLLLaplace(LLLaplace, FullLaplace):
     Mathematically, we have \\(P \\in \\mathbb{R}^{P \\times P}\\).
     See `FullLaplace`, `LLLaplace`, and `BaseLaplace` for the full interface.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ('last_layer', 'full')
 
@@ -214,15 +263,37 @@ class KronLLLaplace(LLLaplace, KronLaplace):
     and computing posterior covariances, marginal likelihood, etc.
     Use of `damping` is possible by initializing or setting `damping=True`.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ('last_layer', 'kron')
 
-    def __init__(self, model, likelihood, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., enable_backprop=False, backend=None, last_layer_name=None,
-                 damping=False, **backend_kwargs):
+    def __init__(
+        self,
+        model,
+        likelihood,
+        sigma_noise=1.0,
+        prior_precision=1.0,
+        prior_mean=0.0,
+        temperature=1.0,
+        enable_backprop=False,
+        backend=None,
+        last_layer_name=None,
+        damping=False,
+        **backend_kwargs
+    ):
         self.damping = damping
-        super().__init__(model, likelihood, sigma_noise, prior_precision,
-                         prior_mean, temperature, enable_backprop, backend, last_layer_name, backend_kwargs)
+        super().__init__(
+            model,
+            likelihood,
+            sigma_noise,
+            prior_precision,
+            prior_mean,
+            temperature,
+            enable_backprop,
+            backend,
+            last_layer_name,
+            backend_kwargs,
+        )
 
     def _init_H(self):
         self.H = Kron.init_from_model(self.model.last_layer, self._device)
@@ -234,5 +305,6 @@ class DiagLLLaplace(LLLaplace, DiagLaplace):
     Mathematically, we have \\(P \\approx \\textrm{diag}(P)\\).
     See `DiagLaplace`, `LLLaplace`, and `BaseLaplace` for the full interface.
     """
+
     # key to map to correct subclass of BaseLaplace, (subset of weights, Hessian structure)
     _key = ('last_layer', 'diag')
diff --git a/laplace/marglik_training.py b/laplace/marglik_training.py
index 23fa952b..b57ac520 100644
--- a/laplace/marglik_training.py
+++ b/laplace/marglik_training.py
@@ -6,6 +6,8 @@
 from torch.nn.utils import parameters_to_vector
 import warnings
 import logging
+from collections import UserDict
+import tqdm
 
 from laplace import Laplace
 from laplace.curvature import AsdlGGN
@@ -28,23 +30,25 @@ def marglik_training(
     n_epochs_burnin=0,
     n_hypersteps=10,
     marglik_frequency=1,
-    prior_prec_init=1.,
-    sigma_noise_init=1.,
-    temperature=1.,
-    enable_backprop=False
+    prior_prec_init=1.0,
+    sigma_noise_init=1.0,
+    temperature=1.0,
+    fix_sigma_noise=False,
+    progress_bar=False,
+    enable_backprop=False,
 ):
-    """Marginal-likelihood based training (Algorithm 1 in [1]). 
+    """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)
     while hyperparameters minimize negative log marginal likelihood.
 
     This method replaces standard neural network training and adds hyperparameter
-    optimization to the procedure. 
-    
-    The settings of standard training can be controlled by passing `train_loader`, 
+    optimization to the procedure.
+
+    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='classification'` or `'regression'`, one can choose between
     categorical likelihood (CrossEntropyLoss) and Gaussian likelihood (MSELoss).
 
     As in [1], we optimize prior precision and, for regression, observation noise
@@ -59,8 +63,8 @@ def marglik_training(
 
     References
     ----------
-    [1] Immer, A., Bauer, M., Fortuin, V., Rätsch, G., Khan, EM. 
-    [*Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning*](https://arxiv.org/abs/2104.04975). 
+    [1] Immer, A., Bauer, M., Fortuin, V., Rätsch, G., Khan, EM.
+    [*Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning*](https://arxiv.org/abs/2104.04975).
     ICML 2021.
 
     Parameters
@@ -96,7 +100,7 @@ def marglik_training(
         how many steps to take on the hyperparameters when marglik is estimated
     marglik_frequency : int
         how often to estimate (and differentiate) the marginal likelihood
-        `marglik_frequency=1` would be every epoch, 
+        `marglik_frequency=1` would be every epoch,
         `marglik_frequency=5` would be every 5 epochs.
     prior_prec_init : float, default=1.0
         initial prior precision
@@ -104,6 +108,11 @@ def marglik_training(
         initial observation noise (for regression only)
     temperature : float, default=1.0
         factor for the likelihood for 'overcounting' data. Might be required for data augmentation.
+    fix_sigma_noise: bool, default=False
+        if False, optimize observation noise via marglik otherwise use `sigma_noise_init` throughout.
+        Only works for regression.
+    progress_bar: bool, default=False
+        whether to show a progress bar (updated per epoch) or not
     enable_backprop : bool, default=False
         make the returned Laplace instance backpropable---useful for e.g. Bayesian optimization.
 
@@ -125,22 +134,24 @@ def marglik_training(
     # get device, data set size N, number of layers H, number of parameters P
     device = parameters_to_vector(model.parameters()).device
     N = len(train_loader.dataset)
-    H = len(list(model.parameters()))
-    P = len(parameters_to_vector(model.parameters()))
+    trainable_params = [p for p in model.parameters() if p.requires_grad]
+    H = len(trainable_params)
+    P = len(parameters_to_vector(trainable_params))
 
     # differentiable hyperparameters
     hyperparameters = list()
     # prior precision
     log_prior_prec_init = np.log(temperature * prior_prec_init)
     log_prior_prec = fix_prior_prec_structure(
-        log_prior_prec_init, prior_structure, H, P, device)
+        log_prior_prec_init, prior_structure, H, P, device
+    )
     log_prior_prec.requires_grad = True
     hyperparameters.append(log_prior_prec)
 
     # set up loss (and observation noise hyperparam)
     if likelihood == 'classification':
         criterion = CrossEntropyLoss(reduction='mean')
-        sigma_noise = 1.
+        sigma_noise = 1.0
     elif likelihood == 'regression':
         criterion = MSELoss(reduction='mean')
         log_sigma_noise_init = np.log(sigma_noise_init)
@@ -168,21 +179,41 @@ def marglik_training(
     losses = list()
     margliks = list()
 
-    for epoch in range(1, n_epochs + 1):
+    pbar = tqdm.trange(
+        1,
+        n_epochs + 1,
+        disable=not progress_bar,
+        position=1,
+        leave=False,
+        desc='[Training]',
+        colour='blue',
+    )
+    for epoch in pbar:
         epoch_loss = 0
         epoch_perf = 0
-        
+
         # standard NN training per batch
-        for X, y in train_loader:
-            X, y = X.to(device), y.to(device)
+        for data in train_loader:
+            if isinstance(data, UserDict) or isinstance(data, dict):
+                X, y = data, data['labels']
+                y = y.to(device, non_blocking=True)
+            else:
+                X, y = data
+                X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
             optimizer.zero_grad()
             if likelihood == 'regression':
-                sigma_noise = torch.exp(log_sigma_noise).detach()
-                crit_factor = temperature / (2 * sigma_noise.square())
+                sigma_noise = (
+                    torch.exp(log_sigma_noise).detach()
+                    if not fix_sigma_noise
+                    else sigma_noise_init
+                )
+                crit_factor = temperature / (2 * sigma_noise**2)
             else:
                 crit_factor = temperature
             prior_prec = torch.exp(log_prior_prec).detach()
-            theta = parameters_to_vector(model.parameters())
+            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
@@ -199,8 +230,10 @@ def marglik_training(
         losses.append(epoch_loss / N)
 
         # compute validation error to report during training
-        logging.info(f'MARGLIK[epoch={epoch}]: network training. Loss={losses[-1]:.3f}.' +
-                     f'Perf={epoch_perf/N:.3f}')
+        logging.info(
+            f'MARGLIK[epoch={epoch}]: network training. Loss={losses[-1]:.3f}.'
+            + f'Perf={epoch_perf/N:.3f}'
+        )
 
         # only update hyperparameters every marglik_frequency steps after burnin
         if (epoch % marglik_frequency) != 0 or epoch < n_epochs_burnin:
@@ -208,21 +241,31 @@ def marglik_training(
 
         # optimizer hyperparameters by differentiating marglik
         # 1. fit laplace approximation
-        sigma_noise = 1 if likelihood == 'classification' else torch.exp(log_sigma_noise)
+        if likelihood == 'classification':
+            sigma_noise = 1
+        else:
+            sigma_noise = (
+                torch.exp(log_sigma_noise) if not fix_sigma_noise else sigma_noise_init
+            )
         prior_prec = torch.exp(log_prior_prec)
         lap = Laplace(
-            model, likelihood, hessian_structure=hessian_structure, sigma_noise=sigma_noise, 
-            prior_precision=prior_prec, temperature=temperature, backend=backend,
-            subset_of_weights='all'
+            model,
+            likelihood,
+            hessian_structure=hessian_structure,
+            sigma_noise=sigma_noise,
+            prior_precision=prior_prec,
+            temperature=temperature,
+            backend=backend,
+            subset_of_weights='all',
         )
         lap.fit(train_loader)
 
         # 2. differentiate wrt. hyperparameters for n_hypersteps
         for _ in range(n_hypersteps):
             hyper_optimizer.zero_grad()
-            if likelihood == 'classification':
+            if likelihood == 'classification' or fix_sigma_noise:
                 sigma_noise = None
-            elif likelihood == 'regression':
+            else:
                 sigma_noise = torch.exp(log_sigma_noise)
             prior_prec = torch.exp(log_prior_prec)
             marglik = -lap.log_marginal_likelihood(prior_prec, sigma_noise)
@@ -234,13 +277,24 @@ def marglik_training(
         if margliks[-1] < best_marglik:
             best_model_dict = deepcopy(model.state_dict())
             best_precision = deepcopy(prior_prec.detach())
-            best_sigma = 1 if likelihood == 'classification' else deepcopy(sigma_noise.detach())
+            if likelihood == 'classification':
+                best_sigma = 1
+            else:
+                best_sigma = (
+                    deepcopy(sigma_noise.detach())
+                    if not fix_sigma_noise
+                    else sigma_noise_init
+                )
             best_marglik = margliks[-1]
-            logging.info(f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={best_marglik:.2f}. '
-                         + 'Saving new best model.')
+            logging.info(
+                f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={best_marglik:.2f}. '
+                + 'Saving new best model.'
+            )
         else:
-            logging.info(f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={margliks[-1]:.2f}.'
-                         + f'No improvement over {best_marglik:.2f}')
+            logging.info(
+                f'MARGLIK[epoch={epoch}]: marglik optimization. MargLik={margliks[-1]:.2f}.'
+                + f'No improvement over {best_marglik:.2f}'
+            )
 
     logging.info('MARGLIK: finished training. Recover best model and fit Laplace.')
     if best_model_dict is not None:
@@ -248,9 +302,15 @@ def marglik_training(
         sigma_noise = best_sigma
         prior_prec = best_precision
     lap = Laplace(
-        model, likelihood, hessian_structure=hessian_structure, sigma_noise=sigma_noise, 
-        prior_precision=prior_prec, temperature=temperature, backend=backend,
-        subset_of_weights='all', enable_backprop=enable_backprop
+        model,
+        likelihood,
+        hessian_structure=hessian_structure,
+        sigma_noise=sigma_noise,
+        prior_precision=prior_prec,
+        temperature=temperature,
+        backend=backend,
+        subset_of_weights='all',
+        enable_backprop=enable_backprop,
     )
     lap.fit(train_loader)
     return lap, model, margliks, losses
diff --git a/laplace/subnetlaplace.py b/laplace/subnetlaplace.py
index 8c2f7e1d..d4e61bd7 100644
--- a/laplace/subnetlaplace.py
+++ b/laplace/subnetlaplace.py
@@ -67,7 +67,9 @@ class SubnetLaplace(ParametricLaplace):
         set the number of MC samples for stochastic approximations.
     """
     def __init__(self, model, likelihood, subnetwork_indices, sigma_noise=1., prior_precision=1.,
-                 prior_mean=0., temperature=1., backend=None, backend_kwargs=None):
+                 prior_mean=0., temperature=1., backend=None, backend_kwargs=None, asdl_fisher_kwargs=None):
+        if asdl_fisher_kwargs is not None:
+            raise ValueError('Subnetwork Laplace does not support asdl_fisher_kwargs.')
         self.H = None
         super().__init__(model, likelihood, sigma_noise=sigma_noise,
                          prior_precision=prior_precision, prior_mean=prior_mean,
@@ -143,7 +145,7 @@ class FullSubnetLaplace(SubnetLaplace, FullLaplace):
     def _init_H(self):
         self.H = torch.zeros(self.n_params_subnet, self.n_params_subnet, device=self._device)
 
-    def sample(self, n_samples=100):
+    def sample(self, n_samples=100, generator=None):
         # 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,))
@@ -171,9 +173,9 @@ def _check_jacobians(self, Js):
         if p != self.n_params_subnet:
             raise ValueError('Invalid Jacobians shape for Laplace posterior approx.')
 
-    def sample(self, n_samples=100):
+    def sample(self, n_samples=100, generator=None):
         # 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)
+        samples = torch.randn(n_samples, self.n_params_subnet, device=self._device, generator=generator)
         samples = samples * self.posterior_scale.reshape(1, self.n_params_subnet)
         subnet_samples = self.mean_subnet.reshape(1, self.n_params_subnet) + samples
         return self.assemble_full_samples(subnet_samples)
diff --git a/laplace/utils/__init__.py b/laplace/utils/__init__.py
index acda1237..9e979119 100644
--- a/laplace/utils/__init__.py
+++ b/laplace/utils/__init__.py
@@ -3,6 +3,7 @@
 from laplace.utils.matrix import Kron, KronDecomposed
 from laplace.utils.swag import fit_diagonal_swag_var
 from laplace.utils.subnetmask import SubnetMask, RandomSubnetMask, LargestMagnitudeSubnetMask, LargestVarianceDiagLaplaceSubnetMask, LargestVarianceSWAGSubnetMask, ParamNameSubnetMask, ModuleNameSubnetMask, LastLayerSubnetMask
+from laplace.utils.metrics import RunningNLLMetric
 
 
 __all__ = ['get_nll', 'validate', 'parameters_per_layer', 'invsqrt_precision', 'kron',
@@ -13,4 +14,4 @@
            'Kron', 'KronDecomposed',
 		   'fit_diagonal_swag_var',
 		   'SubnetMask', 'RandomSubnetMask', 'LargestMagnitudeSubnetMask', 'LargestVarianceDiagLaplaceSubnetMask',
-		   'LargestVarianceSWAGSubnetMask', 'ParamNameSubnetMask', 'ModuleNameSubnetMask', 'LastLayerSubnetMask']
+		   'LargestVarianceSWAGSubnetMask', 'ParamNameSubnetMask', 'ModuleNameSubnetMask', 'LastLayerSubnetMask', 'RunningNLLMetric']
diff --git a/laplace/utils/matrix.py b/laplace/utils/matrix.py
index cddee587..e58d8058 100644
--- a/laplace/utils/matrix.py
+++ b/laplace/utils/matrix.py
@@ -23,6 +23,7 @@ class Kron:
         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
 
@@ -32,15 +33,20 @@ def init_from_model(cls, model, device):
 
         Parameters
         ----------
-        model : torch.nn.Module
+        model : nn.Module or iterable of parameters, e.g. model.parameters()
         device : torch.device
 
         Returns
         -------
         kron : Kron
         """
+        if isinstance(model, torch.nn.Module):
+            params = model.parameters()
+        else:
+            params = model
+
         kfacs = list()
-        for p in model.parameters():
+        for p in params:
             if p.ndim == 1:  # bias
                 P = p.size(0)
                 kfacs.append([torch.zeros(P, P, device=device)])
@@ -50,10 +56,12 @@ def init_from_model(cls, model, device):
                 elif p.ndim > 2:
                     P_in, P_out = p.shape[0], np.prod(p.shape[1:])
 
-                kfacs.append([
-                    torch.zeros(P_in, P_in, device=device),
-                    torch.zeros(P_out, P_out, device=device)
-                ])
+                kfacs.append(
+                    [
+                        torch.zeros(P_in, P_in, device=device),
+                        torch.zeros(P_out, P_out, device=device),
+                    ]
+                )
             else:
                 raise ValueError('Invalid parameter shape in network.')
         return cls(kfacs)
@@ -72,8 +80,11 @@ def __add__(self, other):
         if not isinstance(other, Kron):
             raise ValueError('Can only add Kron to Kron.')
 
-        kfacs = [[Hi.add(Hj) for Hi, Hj in zip(Fi, Fj)]
-                 for Fi, Fj in zip(self.kfacs, other.kfacs)]
+        kfacs = [
+            [Hi.add(Hj) for Hi, Hj in zip(Fi, Fj)]
+            for Fi, Fj in zip(self.kfacs, other.kfacs)
+        ]
+
         return Kron(kfacs)
 
     def __mul__(self, scalar: Union[float, torch.Tensor]):
@@ -93,7 +104,7 @@ def __mul__(self, scalar: Union[float, torch.Tensor]):
             raise ValueError('Input not valid python or torch scalar.')
 
         # distribute factors evenly so that each group is multiplied by factor
-        kfacs = [[pow(scalar, 1/len(F)) * Hi for Hi in F] for F in self.kfacs]
+        kfacs = [[pow(scalar, 1 / len(F)) * Hi for Hi in F] for F in self.kfacs]
         return Kron(kfacs)
 
     def __len__(self):
@@ -114,7 +125,14 @@ def decompose(self, damping=False):
         for F in self.kfacs:
             Qs, ls = list(), list()
             for Hi in F:
-                l, Q = symeig(Hi)
+                if Hi.ndim > 1:
+                    # Dense Kronecker factor.
+                    l, Q = symeig(Hi)
+                else:
+                    # Diagonal Kronecker factor.
+                    l = Hi
+                    # This might be too memory intensive since len(Hi) can be large.
+                    Q = torch.eye(len(Hi), dtype=Hi.dtype, device=Hi.device)
                 Qs.append(Q)
                 ls.append(l)
             eigvecs.append(Qs)
@@ -144,15 +162,17 @@ def _bmm(self, W: torch.Tensor) -> torch.Tensor:
             if len(Fs) == 1:
                 Q = Fs[0]
                 p = len(Q)
-                W_p = W[:, cur_p:cur_p+p].T
-                SW.append((Q @ W_p).T)
+                W_p = W[:, cur_p : cur_p + p].T
+                SW.append((Q @ W_p).T if Q.ndim > 1 else (Q.view(-1, 1) * W_p).T)
                 cur_p += p
             elif len(Fs) == 2:
                 Q, H = Fs
                 p_in, p_out = len(Q), len(H)
                 p = p_in * p_out
-                W_p = W[:, cur_p:cur_p+p].reshape(B * K, p_in, p_out)
-                SW.append((Q @ W_p @ H.T).reshape(B * K, p_in * p_out))
+                W_p = W[:, cur_p : cur_p + p].reshape(B * K, p_in, p_out)
+                QW_p = Q @ W_p if Q.ndim > 1 else Q.view(-1, 1) * W_p
+                QW_pHt = QW_p @ H.T if H.ndim > 1 else QW_p * H.view(1, -1)
+                SW.append(QW_pHt.reshape(B * K, p_in * p_out))
                 cur_p += p
             else:
                 raise AttributeError('Shape mismatch')
@@ -200,11 +220,12 @@ def logdet(self) -> torch.Tensor:
         logdet = 0
         for F in self.kfacs:
             if len(F) == 1:
-                logdet += F[0].logdet()
+                logdet += F[0].logdet() if F[0].ndim > 1 else F[0].log().sum()
             else:  # len(F) == 2
                 Hi, Hj = F
                 p_in, p_out = len(Hi), len(Hj)
-                logdet += p_out * Hi.logdet() + p_in * Hj.logdet()
+                logdet += p_out * Hi.logdet() if Hi.ndim > 1 else p_out * Hi.log().sum()
+                logdet += p_in * Hj.logdet() if Hj.ndim > 1 else p_in * Hj.log().sum()
         return logdet
 
     def diag(self) -> torch.Tensor:
@@ -216,10 +237,12 @@ def diag(self) -> torch.Tensor:
         """
         diags = list()
         for F in self.kfacs:
+            F0 = F[0].diag() if F[0].ndim > 1 else F[0]
             if len(F) == 1:
-                diags.append(F[0].diagonal())
+                diags.append(F0)
             else:
-                diags.append(torch.outer(F[0].diagonal(), F[1].diagonal()).flatten())
+                F1 = F[1].diag() if F[1].ndim > 1 else F[1]
+                diags.append(torch.outer(F0, F1).flatten())
         return torch.cat(diags)
 
     def to_matrix(self) -> torch.Tensor:
@@ -233,10 +256,12 @@ def to_matrix(self) -> torch.Tensor:
         """
         blocks = list()
         for F in self.kfacs:
+            F0 = F[0] if F[0].ndim > 1 else F[0].diag()
             if len(F) == 1:
-                blocks.append(F[0])
+                blocks.append(F0)
             else:
-                blocks.append(kron(F[0], F[1]))
+                F1 = F[1] if F[1].ndim > 1 else F[1].diag()
+                blocks.append(kron(F0, F1))
         return block_diag(blocks)
 
     # for commutative operations
@@ -285,9 +310,10 @@ def _check_deltas(self, deltas: torch.Tensor):
         if not isinstance(deltas, torch.Tensor):
             raise ValueError('Can only add torch.Tensor to KronDecomposed.')
 
-        if (deltas.ndim == 0  # scalar
-            or (deltas.ndim == 1  # vector of length 1 or len(self)
-                and (len(deltas) == 1 or len(deltas) == len(self)))):
+        if deltas.ndim == 0 or (  # scalar
+            deltas.ndim == 1  # vector of length 1 or len(self)
+            and (len(deltas) == 1 or len(deltas) == len(self))
+        ):
             return
         else:
             raise ValueError('Invalid shape of delta added to KronDecomposed.')
@@ -322,7 +348,9 @@ def __mul__(self, scalar):
         if not _is_valid_scalar(scalar):
             raise ValueError('Invalid argument, can only multiply Kron with scalar.')
 
-        eigenvalues = [[pow(scalar, 1/len(ls)) * l for l in ls] for ls in self.eigenvalues]
+        eigenvalues = [
+            [pow(scalar, 1 / len(ls)) * l for l in ls] for ls in self.eigenvalues
+        ]
         return KronDecomposed(self.eigenvectors, eigenvalues, self.deltas)
 
     def __len__(self) -> int:
@@ -374,11 +402,13 @@ def _bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
         W = W.reshape(B * K, P)
         cur_p = 0
         SW = list()
-        for ls, Qs, delta in zip(self.eigenvalues, self.eigenvectors, self.deltas):
+        for i, (ls, Qs, delta) in enumerate(
+            zip(self.eigenvalues, self.eigenvectors, self.deltas)
+        ):
             if len(ls) == 1:
                 Q, l, p = Qs[0], ls[0], len(ls[0])
                 ldelta_exp = torch.pow(l + delta, exponent).reshape(-1, 1)
-                W_p = W[:, cur_p:cur_p+p].T
+                W_p = W[:, cur_p : cur_p + p].T
                 SW.append((Q @ (ldelta_exp * (Q.T @ W_p))).T)
                 cur_p += p
             elif len(ls) == 2:
@@ -389,9 +419,11 @@ def _bmm(self, W: torch.Tensor, exponent: float = -1) -> torch.Tensor:
                     l1d, l2d = l1 + torch.sqrt(delta), l2 + torch.sqrt(delta)
                     ldelta_exp = torch.pow(torch.outer(l1d, l2d), exponent).unsqueeze(0)
                 else:
-                    ldelta_exp = torch.pow(torch.outer(l1, l2) + delta, exponent).unsqueeze(0)
+                    ldelta_exp = torch.pow(
+                        torch.outer(l1, l2) + delta, exponent
+                    ).unsqueeze(0)
                 p_in, p_out = len(l1), len(l2)
-                W_p = W[:, cur_p:cur_p+p].reshape(B * K, p_in, p_out)
+                W_p = W[:, cur_p : cur_p + p].reshape(B * K, p_in, p_out)
                 W_p = (Q1.T @ W_p @ Q2) * ldelta_exp
                 W_p = Q1 @ W_p @ Q2.T
                 SW.append(W_p.reshape(B * K, p_in * p_out))
@@ -449,14 +481,18 @@ def diag(self, exponent: float = 1) -> torch.Tensor:
         for Qs, ls, delta in zip(self.eigenvectors, self.eigenvalues, self.deltas):
             if len(ls) == 1:
                 Ql = Qs[0] * torch.pow(ls[0] + delta, exponent).reshape(1, -1)
-                d = torch.einsum('mp,mp->m', Ql, Qs[0])  # only compute inner products for diag
+                d = torch.einsum(
+                    'mp,mp->m', Ql, Qs[0]
+                )  # only compute inner products for diag
                 diags.append(d)
             else:
                 Q1, Q2 = Qs
                 l1, l2 = ls
                 if self.damping:
                     delta_sqrt = torch.sqrt(delta)
-                    l = torch.pow(torch.outer(l1 + delta_sqrt, l2 + delta_sqrt), exponent)
+                    l = torch.pow(
+                        torch.outer(l1 + delta_sqrt, l2 + delta_sqrt), exponent
+                    )
                 else:
                     l = torch.pow(torch.outer(l1, l2) + delta, exponent)
                 d = oe.contract('mp,nq,pq,mp,nq->mn', Q1, Q2, l, Q1, Q2).flatten()
@@ -488,7 +524,9 @@ def to_matrix(self, exponent: float = 1) -> torch.Tensor:
                 Q = kron(Q1, Q2)
                 if self.damping:
                     delta_sqrt = torch.sqrt(delta)
-                    l = torch.pow(torch.outer(l1 + delta_sqrt, l2 + delta_sqrt), exponent)
+                    l = torch.pow(
+                        torch.outer(l1 + delta_sqrt, l2 + delta_sqrt), exponent
+                    )
                 else:
                     l = torch.pow(torch.outer(l1, l2) + delta, exponent)
                 L = torch.diag(l.flatten())
diff --git a/laplace/utils/metrics.py b/laplace/utils/metrics.py
new file mode 100644
index 00000000..6a29d3e7
--- /dev/null
+++ b/laplace/utils/metrics.py
@@ -0,0 +1,40 @@
+import torch
+from torch.nn import functional as F
+from torchmetrics import Metric
+
+
+class RunningNLLMetric(Metric):
+    """
+    NLL metrics that
+
+    Parameters
+    ----------
+    ignore_index: int, default = -100
+        which class label to ignore when computing the NLL loss
+    """
+    def __init__(self, ignore_index=-100):
+        super().__init__()
+        self.add_state('nll_sum', default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.add_state('n_valid_labels', default=torch.tensor(0.), dist_reduce_fx='sum')
+        self.ignore_index = ignore_index
+
+    def update(self, probs: torch.Tensor, targets: torch.Tensor) -> None:
+        """
+        Parameters
+        ----------
+        probs: torch.Tensor
+            probability tensor of shape (..., n_classes)
+
+        targets: torch.Tensor
+            integer tensor of shape (...)
+        """
+        probs = probs.view(-1, probs.shape[-1])
+        targets = targets.view(-1)
+
+        self.nll_sum += F.nll_loss(
+            probs.log(), targets, ignore_index=self.ignore_index, reduction='sum'
+        )
+        self.n_valid_labels += (targets != self.ignore_index).sum()
+
+    def compute(self):
+        return self.nll_sum / self.n_valid_labels
diff --git a/laplace/utils/utils.py b/laplace/utils/utils.py
index 7ad3313e..315784fb 100644
--- a/laplace/utils/utils.py
+++ b/laplace/utils/utils.py
@@ -6,6 +6,9 @@
 from torch.nn.utils import parameters_to_vector
 from torch.nn import BatchNorm1d, BatchNorm2d, BatchNorm3d
 from torch.distributions.multivariate_normal import _precision_to_scale_tril
+from torchmetrics import Metric
+from collections import UserDict
+import math
 
 
 __all__ = ['get_nll', 'validate', 'parameters_per_layer', 'invsqrt_precision', 'kron',
@@ -17,29 +20,52 @@ def get_nll(out_dist, targets):
 
 
 @torch.no_grad()
-def validate(laplace, val_loader, pred_type='glm', link_approx='probit', n_samples=100):
+def validate(laplace, val_loader, loss, pred_type='glm', link_approx='probit', n_samples=100, loss_with_var=False) -> float:
     laplace.model.eval()
-    output_means, output_vars = list(), list()
-    targets = list()
-    for X, y in val_loader:
-        X, y = X.to(laplace._device), y.to(laplace._device)
+    assert callable(loss) or isinstance(loss, Metric)
+    is_offline = not isinstance(loss, Metric)
+
+    if is_offline:
+        output_means, output_vars = list(), list()
+        targets = list()
+
+    for data in val_loader:
+        # If x is UserDict, then it is a from Huggingface dataset
+        if isinstance(data, UserDict) or isinstance(data, dict):
+            X, y = data, data['labels']
+        else:
+            X, y = data
+            X = X.to(laplace._device)
+        y = y.to(laplace._device)
         out = laplace(
             X, pred_type=pred_type,
             link_approx=link_approx,
             n_samples=n_samples)
 
         if type(out) == tuple:
-            output_means.append(out[0])
-            output_vars.append(out[1])
+            if is_offline:
+                output_means.append(out[0])
+                output_vars.append(out[1])
+                targets.append(y)
+            else:
+                loss.update(*out, y)
         else:
-            output_means.append(out)
-
-        targets.append(y)
-
-    if len(output_vars) == 0:
-        return torch.cat(output_means, dim=0), torch.cat(targets, dim=0)
-    return ((torch.cat(output_means, dim=0), torch.cat(output_vars, dim=0)),
-            torch.cat(targets, dim=0))
+            if is_offline:
+                output_means.append(out)
+                targets.append(y)
+            else:
+                loss.update(out, y)
+
+    if is_offline:
+        if len(output_vars) == 0:
+            preds, targets = torch.cat(output_means, dim=0), torch.cat(targets, dim=0)
+            return loss(preds, targets).item()
+
+        means, variances = torch.cat(output_means, dim=0), torch.cat(output_vars, dim=0)
+        targets = torch.cat(targets, dim=0)
+        return loss(means, variances, targets).item()
+    else:
+        return loss.compute().item()
 
 
 def parameters_per_layer(model):
@@ -208,7 +234,8 @@ def expand_prior_precision(prior_prec, model):
     expanded_prior_prec : torch.Tensor
         expanded prior precision has the same shape as model parameters
     """
-    theta = parameters_to_vector(model.parameters())
+    trainable_params = [p for p in model.parameters() if p.requires_grad]
+    theta = parameters_to_vector(trainable_params)
     device, P = theta.device, len(theta)
     assert prior_prec.ndim == 1
     if len(prior_prec) == 1:  # scalar
@@ -217,12 +244,12 @@ def expand_prior_precision(prior_prec, model):
         return prior_prec.to(device)
     else:
         return torch.cat([delta * torch.ones_like(m).flatten() for delta, m
-                          in zip(prior_prec, model.parameters())])
+                          in zip(prior_prec, trainable_params)])
 
 
 def fix_prior_prec_structure(prior_prec_init, prior_structure, n_layers, n_params, device):
     if prior_structure == 'scalar':
-        prior_prec_init = torch.full((1,), prior_prec_init, device=device)   
+        prior_prec_init = torch.full((1,), prior_prec_init, device=device)
     elif prior_structure == 'layerwise':
         prior_prec_init = torch.full((n_layers,), prior_prec_init, device=device)
     elif prior_structure == 'diag':
@@ -248,9 +275,9 @@ def normal_samples(mean, var, n_samples, generator=None):
     """
     assert mean.ndim == 2, 'Invalid input shape of mean, should be 2-dimensional.'
     _, output_dim = mean.shape
-    randn_samples = torch.randn((output_dim, n_samples), device=mean.device, 
+    randn_samples = torch.randn((output_dim, n_samples), device=mean.device,
                                 dtype=mean.dtype, generator=generator)
-    
+
     if mean.shape == var.shape:
         # diagonal covariance
         scaled_samples = var.sqrt().unsqueeze(-1) * randn_samples.unsqueeze(0)
diff --git a/pyproject.toml b/pyproject.toml
index 374b58cb..6f1f38a4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,3 +4,6 @@ requires = [
     "wheel"
 ]
 build-backend = "setuptools.build_meta"
+
+[tool.ruff.format]
+quote-style = "single"
diff --git a/setup.cfg b/setup.cfg
index e13542b6..da76b026 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -37,6 +37,8 @@ install_requires =
     torchaudio
     backpack-for-pytorch
     asdfghjkl
+    asdl @ git+https://github.com/wiseodd/asdl@dev
+    torchmetrics
     opt_einsum
     curvlinops-for-pytorch @ git+https://github.com/f-dangel/curvlinops
 # Require a specific Python version, e.g. Python 2.7 or >= 3.4


Test patch:
diff --git a/tests/test_baselaplace.py b/tests/test_baselaplace.py
index 4a509a5e..862b3d67 100644
--- a/tests/test_baselaplace.py
+++ b/tests/test_baselaplace.py
@@ -1,3 +1,4 @@
+from collections.abc import MutableMapping
 from math import sqrt, prod
 import pytest
 from itertools import product
@@ -9,11 +10,14 @@
 from torch.nn.utils import parameters_to_vector
 from torch.utils.data import DataLoader, TensorDataset
 from torch.distributions import Normal, Categorical
+from laplace.curvature.asdfghjkl import AsdfghjklGGN, AsdfghjklEF
+from laplace.curvature.curvlinops import CurvlinopsEF, CurvlinopsGGN
 from torchvision.models import wide_resnet50_2
 
 from laplace.laplace import FullLaplace, KronLaplace, DiagLaplace, LowRankLaplace
 from laplace.utils import KronDecomposed
-from tests.utils import jacobians_naive
+from laplace.curvature import AsdlGGN, BackPackGGN, AsdlEF
+from tests.utils import ListDataset, dict_data_collator, jacobians_naive
 
 
 torch.manual_seed(240)
@@ -38,6 +42,56 @@ def large_model():
     return model
 
 
+@pytest.fixture
+def custom_model():
+    class CustomModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.net = nn.Sequential(nn.Linear(5, 100), nn.ReLU(), nn.Linear(100, 2))
+
+        def forward(self, data: MutableMapping | torch.Tensor):
+            if isinstance(data, MutableMapping):
+                x = data['input_ids'].to(next(self.parameters()).device)
+            else:
+                x = data
+
+            logits = self.net(x)
+            return logits
+
+    return CustomModel()
+
+
+@pytest.fixture
+def reward_model():
+    class RewardModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.net = nn.Sequential(nn.Linear(3, 100), nn.ReLU(), nn.Linear(100, 1))
+
+        def forward(self, x):
+            """
+            x: torch.Tensor
+                If training == True then shape (batch_size, 2, dim)
+                Else shape (batch_size, dim)
+            """
+            if len(x.shape) == 3:
+                batch_size, _, dim = x.shape
+
+                # Flatten to (batch_size*2, dim)
+                flat_x = x.reshape(-1, dim)
+
+                # Forward
+                flat_logits = self.net(flat_x)  # (batch_size*2, 1)
+
+                # Reshape back to (batch_size, 2)
+                return flat_logits.reshape(batch_size, 2)
+            else:
+                logits = self.net(x)  # (batch_size, 1)
+                return logits
+
+    return RewardModel()
+
+
 @pytest.fixture
 def class_loader():
     X = torch.randn(10, 3)
@@ -52,6 +106,28 @@ def reg_loader():
     return DataLoader(TensorDataset(X, y), batch_size=3)
 
 
+@pytest.fixture
+def custom_loader():
+    data = []
+    for _ in range(10):
+        datum = {'input_ids': torch.randn(5), 'labels': torch.randint(2, (1,))}
+        data.append(datum)
+    return DataLoader(ListDataset(data), batch_size=3, collate_fn=dict_data_collator)
+
+
+@pytest.fixture
+def reward_loader():
+    X = torch.randn(10, 2, 3)
+    y = torch.randint(2, (10,))
+    return DataLoader(TensorDataset(X, y), batch_size=3)
+
+
+@pytest.fixture
+def reward_test_X():
+    X = torch.randn(10, 3)
+    return X
+
+
 @pytest.mark.parametrize('laplace', flavors)
 def test_laplace_init(laplace, model):
     lap = laplace(model, 'classification')
@@ -151,17 +227,21 @@ def test_laplace_init_precision(laplace, model):
 def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
     mean = parameters_to_vector(model.parameters())
     P = len(mean)
-    lap_scalar_mean = laplace(model, 'classification',
-                              prior_precision=1e-2, prior_mean=1.)
-    assert torch.allclose(lap_scalar_mean.prior_mean, torch.tensor([1.]))
-    lap_tensor_mean = laplace(model, 'classification',
-                              prior_precision=1e-2, prior_mean=torch.ones(1))
-    assert torch.allclose(lap_tensor_mean.prior_mean, torch.tensor([1.]))
-    lap_tensor_scalar_mean = laplace(model, 'classification',
-                                     prior_precision=1e-2, prior_mean=torch.ones(1)[0])
-    assert torch.allclose(lap_tensor_scalar_mean.prior_mean, torch.tensor(1.))
-    lap_tensor_full_mean = laplace(model, 'classification',
-                                   prior_precision=1e-2, prior_mean=torch.ones(P))
+    lap_scalar_mean = laplace(
+        model, 'classification', prior_precision=1e-2, prior_mean=1.0
+    )
+    assert torch.allclose(lap_scalar_mean.prior_mean, torch.tensor([1.0]))
+    lap_tensor_mean = laplace(
+        model, 'classification', prior_precision=1e-2, prior_mean=torch.ones(1)
+    )
+    assert torch.allclose(lap_tensor_mean.prior_mean, torch.tensor([1.0]))
+    lap_tensor_scalar_mean = laplace(
+        model, 'classification', prior_precision=1e-2, prior_mean=torch.ones(1)[0]
+    )
+    assert torch.allclose(lap_tensor_scalar_mean.prior_mean, torch.tensor(1.0))
+    lap_tensor_full_mean = laplace(
+        model, 'classification', prior_precision=1e-2, prior_mean=torch.ones(P)
+    )
     assert torch.allclose(lap_tensor_full_mean.prior_mean, torch.ones(P))
 
     lap_scalar_mean.fit(class_loader)
@@ -170,7 +250,7 @@ def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
     lap_tensor_full_mean.fit(class_loader)
     expected = torch.tensor(0).reshape(-1)
     # assert expected.ndim == 0
-    expected = ((mean -1) * 1e-2) @ (mean -1)
+    expected = ((mean - 1) * 1e-2) @ (mean - 1)
     assert torch.allclose(lap_scalar_mean.scatter, expected)
     assert lap_scalar_mean.scatter.shape == expected.shape
     assert torch.allclose(lap_tensor_mean.scatter, expected)
@@ -187,7 +267,7 @@ def test_laplace_init_prior_mean_and_scatter(laplace, model, class_loader):
 
     # unmatched dim
     with pytest.raises(ValueError):
-        prior_mean = torch.ones(P-3)
+        prior_mean = torch.ones(P - 3)
         laplace(model, 'classification', prior_precision=1e-2, prior_mean=prior_mean)
 
     # invalid argument type
@@ -203,11 +283,13 @@ def test_laplace_init_temperature(laplace, model):
     assert lap.temperature == T
 
 
-@pytest.mark.parametrize('laplace,lh', product(flavors, ['classification', 'regression']))
+@pytest.mark.parametrize(
+    'laplace,lh', product(flavors, ['classification', 'regression'])
+)
 def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     if lh == 'classification':
         loader = class_loader
-        sigma_noise = 1.
+        sigma_noise = 1.0
     else:
         loader = reg_loader
         sigma_noise = 0.3
@@ -243,7 +325,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
     assert torch.allclose(theta, lap.mean)
     prior_prec = torch.diag(lap.prior_precision_diag)
     assert prior_prec.shape == torch.Size([len(theta), len(theta)])
-    lml = lml - 1/2 * theta @ prior_prec @ theta
+    lml = lml - 1 / 2 * theta @ prior_prec @ theta
     if laplace == DiagLaplace:
         log_det_post_prec = lap.posterior_precision.log().sum()
     elif laplace == LowRankLaplace:
@@ -251,7 +333,7 @@ def test_laplace_functionality(laplace, lh, model, reg_loader, class_loader):
         log_det_post_prec = (U @ torch.diag(l) @ U.T + p0.diag()).logdet()
     else:
         log_det_post_prec = lap.posterior_precision.logdet()
-    lml = lml + 1/2 * (prior_prec.logdet() - log_det_post_prec)
+    lml = lml + 1 / 2 * (prior_prec.logdet() - log_det_post_prec)
     assert torch.allclose(lml, lap.log_marginal_likelihood())
 
     # test sampling
@@ -309,12 +391,12 @@ def test_online_fit(laplace, model, reg_loader):
         P = lap.H.to_matrix().clone()
     else:
         P = lap.H.clone()
-    loss, n_data = deepcopy(lap.loss), deepcopy(lap.n_data)
+    loss, n_data = deepcopy(lap.loss.item()), deepcopy(lap.n_data)
     # fit a second and third time but don't override
     lap.fit(reg_loader, override=False)
     lap.fit(reg_loader, override=False)
     # Hessian should be now roughly 3x the one before
-    assert torch.allclose(3 * loss, lap.loss)
+    assert torch.allclose(3 * torch.tensor(loss), lap.loss)
     assert (3 * n_data) == lap.n_data
     if type(lap.H) is KronDecomposed:
         assert torch.allclose(lap.H.to_matrix(), 3 * P)
@@ -326,20 +408,24 @@ def test_log_prob_full(model, class_loader):
     lap = FullLaplace(model, 'classification', prior_precision=0.7)
     theta = torch.randn_like(parameters_to_vector(model.parameters()))
     # posterior without fitting is just prior
-    posterior = Normal(loc=torch.zeros_like(theta), scale=sqrt(1/0.7))
+    posterior = Normal(loc=torch.zeros_like(theta), scale=sqrt(1 / 0.7))
     assert torch.allclose(lap.log_prob(theta), posterior.log_prob(theta).sum())
     lap.fit(class_loader)
-    posterior = MultivariateNormal(loc=lap.mean, precision_matrix=lap.posterior_precision)
+    posterior = MultivariateNormal(
+        loc=lap.mean, precision_matrix=lap.posterior_precision
+    )
     assert torch.allclose(lap.log_prob(theta), posterior.log_prob(theta))
 
-    
+
 def test_log_prob_kron(model, class_loader):
     lap = KronLaplace(model, 'classification', prior_precision=0.24)
     theta = torch.randn_like(parameters_to_vector(model.parameters()))
-    posterior = Normal(loc=lap.mean, scale=sqrt(1/0.24))
+    posterior = Normal(loc=lap.mean, scale=sqrt(1 / 0.24))
     assert torch.allclose(lap.log_prob(theta), posterior.log_prob(theta).sum())
     lap.fit(class_loader)
-    posterior = MultivariateNormal(loc=lap.mean, precision_matrix=lap.posterior_precision.to_matrix())
+    posterior = MultivariateNormal(
+        loc=lap.mean, precision_matrix=lap.posterior_precision.to_matrix()
+    )
     assert torch.allclose(lap.log_prob(theta), posterior.log_prob(theta))
 
 
@@ -357,7 +443,9 @@ def test_regression_predictive(laplace, model, reg_loader):
     # GLM predictive, functional variance tested already above.
     f_mu_glm, f_var_glm = lap(X, pred_type='glm')
     assert torch.allclose(f_mu_glm, f)
-    assert f_var_glm.shape == torch.Size([f_mu_glm.shape[0], f_mu_glm.shape[1], f_mu_glm.shape[1]])
+    assert f_var_glm.shape == torch.Size(
+        [f_mu_glm.shape[0], f_mu_glm.shape[1], f_mu_glm.shape[1]]
+    )
     assert len(f_mu_glm) == len(X)
 
     # NN predictive (only diagonal variance estimation)
@@ -366,7 +454,7 @@ def test_regression_predictive(laplace, model, reg_loader):
     assert f_var_nn.shape == torch.Size([f_mu_nn.shape[0], f_mu_nn.shape[1]])
     assert len(f_mu_nn) == len(X)
 
-    # Test joint prediction 
+    # Test joint prediction
     f_mu_joint, f_cov_joint = lap(X, pred_type='glm', joint=True)
     assert len(f_mu_joint.shape) == 1
     assert f_mu_joint.shape[0] == prod(f_mu_glm.shape)
@@ -393,21 +481,31 @@ def test_classification_predictive(laplace, model, class_loader):
     # GLM predictive
     f_pred = lap(X, pred_type='glm', link_approx='mc', n_samples=100)
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
     f_pred = lap(X, pred_type='glm', link_approx='probit')
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
     f_pred = lap(X, pred_type='glm', link_approx='bridge')
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
     f_pred = lap(X, pred_type='glm', link_approx='bridge_norm')
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
 
     # NN predictive
     f_pred = lap(X, pred_type='nn', link_approx='mc', n_samples=100)
     assert f_pred.shape == f.shape
-    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))  # sum up to 1
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
 
 
 @pytest.mark.parametrize('laplace', flavors)
@@ -452,19 +550,74 @@ def test_classification_predictive_samples(laplace, model, class_loader):
     assert np.allclose(fsamples.sum().item(), len(f) * 100)  # sum up to 1
 
 
-# TODO: Add LowRankLaplace
+@pytest.mark.parametrize('laplace', [KronLaplace, DiagLaplace])
+def test_reward_modeling(laplace, reward_model, reward_loader, reward_test_X):
+    lap = laplace(reward_model, 'reward_modeling')
+    lap.fit(reward_loader)
+    f = reward_model(reward_test_X)
+
+    # error
+    with pytest.raises(ValueError):
+        lap(reward_test_X, pred_type='linear')
+
+    # GLM predictive, functional variance tested already above.
+    f_mu, f_var = lap(reward_test_X, pred_type='glm')
+    assert torch.allclose(f_mu, f)
+    assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1], f_mu.shape[1]])
+    assert len(f_mu) == len(reward_test_X)
+
+    # NN predictive (only diagonal variance estimation)
+    f_mu, f_var = lap(reward_test_X, pred_type='nn', link_approx='mc')
+    assert f_mu.shape == f_var.shape
+    assert f_var.shape == torch.Size([f_mu.shape[0], f_mu.shape[1]])
+    assert len(f_mu) == len(reward_test_X)
+
+
+@pytest.mark.parametrize('laplace', [KronLaplace, DiagLaplace])
+@pytest.mark.parametrize('backend', [AsdlEF, AsdlGGN, CurvlinopsEF, CurvlinopsGGN])
+def test_dict_data(laplace, backend, custom_model, custom_loader):
+    if laplace == DiagLaplace and backend == CurvlinopsEF:
+        pytest.skip(
+            'DiagEF is unsupported with Curvlinops when the input is non-tensor.'
+        )
+
+    for data in custom_loader:
+        print(data)
+
+    lap = laplace(custom_model, 'classification', backend=backend)
+    lap.fit(custom_loader)
+
+    test_data = next(iter(custom_loader))
+    f = custom_model(test_data)
+
+    # GLM predictive
+    f_pred = lap(test_data, pred_type='glm')
+    assert f_pred.shape == f.shape
+    assert torch.allclose(
+        f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double)
+    )  # sum up to 1
+
+    # NN predictive
+    f_pred = lap(test_data, pred_type='nn', link_approx='mc')
+    assert f_pred.shape == f.shape
+    assert torch.allclose(f_pred.sum(), torch.tensor(len(f_pred), dtype=torch.double))
+
+
 @pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
-def test_backprop_glm(laplace, model, reg_loader):
+@pytest.mark.parametrize(
+    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+)
+def test_backprop_glm(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
     f_mu, f_var = lap(X, pred_type='glm')
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
-        grad_X_var = torch.autograd.grad(f_var.sum(), X)[0] 
+        grad_X_var = torch.autograd.grad(f_var.sum(), X)[0]
 
         assert grad_X_mu.shape == X.shape
         assert grad_X_var.shape == X.shape
@@ -472,19 +625,21 @@ def test_backprop_glm(laplace, model, reg_loader):
         assert False
 
 
-# TODO: Add LowRankLaplace
 @pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
-def test_backprop_glm_joint(laplace, model, reg_loader):
+@pytest.mark.parametrize(
+    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+)
+def test_backprop_glm_joint(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
     f_mu, f_cov = lap(X, pred_type='glm', joint=True)
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
-        grad_X_var = torch.autograd.grad(f_cov.sum(), X)[0] 
+        grad_X_var = torch.autograd.grad(f_cov.sum(), X)[0]
 
         assert grad_X_mu.shape == X.shape
         assert grad_X_var.shape == X.shape
@@ -492,19 +647,21 @@ def test_backprop_glm_joint(laplace, model, reg_loader):
         assert False
 
 
-# TODO: Add LowRankLaplace
 @pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
-def test_backprop_glm_mc(laplace, model, reg_loader):
+@pytest.mark.parametrize(
+    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+)
+def test_backprop_glm_mc(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
     f_mu, f_var = lap(X, pred_type='glm', link_approx='mc')
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
-        grad_X_var = torch.autograd.grad(f_var.sum(), X)[0] 
+        grad_X_var = torch.autograd.grad(f_var.sum(), X)[0]
 
         assert grad_X_mu.shape == X.shape
         assert grad_X_var.shape == X.shape
@@ -512,21 +669,31 @@ def test_backprop_glm_mc(laplace, model, reg_loader):
         assert False
 
 
-# TODO: Add LowRankLaplace
 @pytest.mark.parametrize('laplace', [FullLaplace, KronLaplace, DiagLaplace])
-def test_backprop_nn(laplace, model, reg_loader):
+@pytest.mark.parametrize(
+    'backend', [BackPackGGN, AsdlGGN, AsdlEF, CurvlinopsGGN, CurvlinopsEF]
+)
+def test_backprop_nn(laplace, model, reg_loader, backend):
     X, y = reg_loader.dataset.tensors
     X.requires_grad = True
 
-    lap = laplace(model, 'regression', enable_backprop=True)
+    lap = laplace(model, 'regression', enable_backprop=True, backend=backend)
     lap.fit(reg_loader)
     f_mu, f_var = lap(X, pred_type='nn', link_approx='mc', n_samples=10)
 
     try:
         grad_X_mu = torch.autograd.grad(f_mu.sum(), X, retain_graph=True)[0]
-        grad_X_var = torch.autograd.grad(f_var.sum(), X)[0] 
+        grad_X_var = torch.autograd.grad(f_var.sum(), X)[0]
 
         assert grad_X_mu.shape == X.shape
         assert grad_X_var.shape == X.shape
     except ValueError:
         assert False
+
+
+@pytest.mark.parametrize('likelihood', ['classification', 'regression'])
+def test_dict_data_diagEF_curvlinops_fails(custom_model, custom_loader, likelihood):
+    lap = DiagLaplace(custom_model, likelihood=likelihood, backend=CurvlinopsEF)
+
+    with pytest.raises(ValueError):
+        lap.fit(custom_loader)
diff --git a/tests/test_curv_backends_asdfghjkl.py b/tests/test_curv_backends_asdfghjkl.py
new file mode 100644
index 00000000..c22ce300
--- /dev/null
+++ b/tests/test_curv_backends_asdfghjkl.py
@@ -0,0 +1,228 @@
+import pytest
+import torch
+from torch import nn
+from torch.nn.utils import parameters_to_vector
+
+from asdfghjkl.operations import Bias, Scale
+
+from laplace.curvature import AsdfghjklGGN, AsdfghjklEF, BackPackGGN, BackPackEF
+
+
+@pytest.fixture
+def model():
+    torch.manual_seed(711)
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    return model
+
+
+@pytest.fixture
+def class_Xy():
+    torch.manual_seed(711)
+    X = torch.randn(10, 3)
+    y = torch.randint(2, (10,))
+    return X, y
+
+
+@pytest.fixture
+def complex_model():
+    torch.manual_seed(711)
+    model = torch.nn.Sequential(
+        nn.Conv2d(3, 4, 2, 2),
+        nn.Flatten(),
+        nn.Tanh(),
+        nn.Linear(16, 20),
+        nn.Tanh(),
+        Scale(),
+        Bias(),
+        nn.Linear(20, 2),
+    )
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+    return model
+
+
+@pytest.fixture
+def complex_class_Xy():
+    torch.manual_seed(711)
+    X = torch.randn(10, 3, 5, 5)
+    y = torch.randint(2, (10,))
+    return X, y
+
+
+def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
+    X, y = class_Xy
+    backend = AsdfghjklGGN(model, 'classification', stochastic=False)
+    loss, dggn = backend.diag(X[:5], y[:5])
+    loss2, dggn2 = backend.diag(X[5:], y[5:])
+    loss += loss2
+    dggn += dggn2
+
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackGGN(model, 'classification', stochastic=False)
+    loss_f, H_ggn = backend.full(X, y)
+    assert torch.allclose(loss, loss_f)
+    assert torch.allclose(dggn, H_ggn.diagonal())
+
+
+def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
+    X, y = class_Xy
+    backend = AsdfghjklEF(model, 'classification')
+    loss, dggn = backend.diag(X[:5], y[:5])
+    loss2, dggn2 = backend.diag(X[5:], y[5:])
+    loss += loss2
+    dggn += dggn2
+
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackEF(model, 'classification')
+    loss_f, H_ggn = backend.full(X, y)
+    assert torch.allclose(loss, loss_f)
+    assert torch.allclose(dggn, H_ggn.diagonal())
+
+
+def test_diag_ggn_stoch_cls_kazuki(class_Xy, model):
+    X, y = class_Xy
+    backend = AsdfghjklGGN(model, 'classification', stochastic=True)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # same order of magnitude os non-stochastic.
+    backend = AsdfghjklGGN(model, 'classification', stochastic=False)
+    loss_ns, dggn_ns = backend.diag(X, y)
+    assert loss_ns == loss
+    assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
+
+
+@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+def test_kron_kazuki_vs_diag_class(class_Xy, model, Backend):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = class_Xy
+    backend = Backend(model, 'classification')
+    loss, dggn = backend.diag(X[:1], y[:1], N=1)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+    loss, kron = backend.kron(X[:1], y[:1], N=1)
+    assert torch.allclose(kron.diag(), dggn)
+
+
+@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
+    X, y = class_Xy
+    backend = Backend(model, 'classification')
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert len(kron.diag()) == model.n_params
+
+    N = len(X)
+    M = 3
+    loss1, kron1 = backend.kron(X[:M], y[:M], N=N)
+    loss2, kron2 = backend.kron(X[M:], y[M:], N=N)
+    kron_two = kron1 + kron2
+    loss_two = loss1 + loss2
+    assert torch.allclose(kron.diag(), kron_two.diag())
+    assert torch.allclose(loss, loss_two)
+
+
+@pytest.mark.parametrize('Backend', [AsdfghjklGGN, AsdfghjklEF])
+def test_kron_summing_up_vs_diag_kazuki(class_Xy, model, Backend):
+    # For a single data point, Kron is exact and should equal diag class_Xy
+    X, y = class_Xy
+    backend = Backend(model, 'classification')
+    loss, dggn = backend.diag(X, y, N=len(X))
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
+
+
+def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
+    X, y = complex_class_Xy
+    backend = AsdfghjklGGN(complex_model, 'classification', stochastic=True)
+    loss, dggn = backend.diag(X, y)
+    # sanity check size of diag ggn
+    assert len(dggn) == complex_model.n_params
+
+    # same order of magnitude os non-stochastic.
+    loss_ns, dggn_ns = backend.diag(X, y)
+    assert loss_ns == loss
+    assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1)
+
+
+@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+def test_complex_kron_kazuki_vs_diag_kazuki(complex_class_Xy, complex_model, Backend):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = complex_class_Xy
+    backend = Backend(complex_model, 'classification')
+    loss, dggn = backend.diag(X[:1], y[:1], N=1)
+    # sanity check size of diag ggn
+    assert len(dggn) == complex_model.n_params
+    loss, kron = backend.kron(X[:1], y[:1], N=1)
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
+
+
+@pytest.mark.parametrize('Backend', [AsdfghjklEF, AsdfghjklGGN])
+def test_complex_kron_batching_correction_kazuki(
+    complex_class_Xy, complex_model, Backend
+):
+    X, y = complex_class_Xy
+    backend = Backend(complex_model, 'classification')
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert len(kron.diag()) == complex_model.n_params
+
+    N = len(X)
+    M = 3
+    loss1, kron1 = backend.kron(X[:M], y[:M], N=N)
+    loss2, kron2 = backend.kron(X[M:], y[M:], N=N)
+    kron_two = kron1 + kron2
+    loss_two = loss1 + loss2
+    assert torch.allclose(kron.diag(), kron_two.diag())
+    assert torch.allclose(loss, loss_two)
+
+
+@pytest.mark.parametrize('Backend', [AsdfghjklGGN, AsdfghjklEF])
+def test_complex_kron_summing_up_vs_diag_class_kazuki(
+    complex_class_Xy, complex_model, Backend
+):
+    # For a single data point, Kron is exact and should equal diag class_Xy
+    X, y = complex_class_Xy
+    backend = Backend(complex_model, 'classification')
+    loss, dggn = backend.diag(X, y, N=len(X))
+    loss, kron = backend.kron(X, y, N=len(X))
+    assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-2)
+
+
+def test_kron_normalization_ggn_class(class_Xy, model):
+    X, y = class_Xy
+    xi, yi = X[:1], y[:1]
+    backend = AsdfghjklGGN(model, 'classification', stochastic=False)
+    loss, kron = backend.kron(xi, yi, N=1)
+    kron_true = 7 * kron
+    loss_true = 7 * loss
+    X = torch.repeat_interleave(xi, 7, 0)
+    y = torch.repeat_interleave(yi, 7, 0)
+    loss_test, kron_test = backend.kron(X, y, N=7)
+    assert torch.allclose(kron_true.diag(), kron_test.diag())
+    assert torch.allclose(loss_true, loss_test)
+
+
+def test_kron_normalization_ef_class(class_Xy, model):
+    X, y = class_Xy
+    xi, yi = X[:1], y[:1]
+    backend = AsdfghjklEF(model, 'classification')
+    loss, kron = backend.kron(xi, yi, N=1)
+    kron_true = 7 * kron
+    loss_true = 7 * loss
+    X = torch.repeat_interleave(xi, 7, 0)
+    y = torch.repeat_interleave(yi, 7, 0)
+    loss_test, kron_test = backend.kron(X, y, N=7)
+    assert torch.allclose(kron_true.diag(), kron_test.diag())
+    assert torch.allclose(loss_true, loss_test)
diff --git a/tests/test_curv_backends_asdl.py b/tests/test_curv_backends_asdl.py
index 696c8969..4c9d7fa1 100644
--- a/tests/test_curv_backends_asdl.py
+++ b/tests/test_curv_backends_asdl.py
@@ -2,8 +2,7 @@
 import torch
 from torch import nn
 from torch.nn.utils import parameters_to_vector
-
-from asdfghjkl.operations import Bias, Scale
+from copy import deepcopy
 
 from laplace.curvature import AsdlGGN, AsdlEF, BackPackGGN, BackPackEF
 
@@ -11,7 +10,7 @@
 @pytest.fixture
 def model():
     torch.manual_seed(711)
-    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Tanh(), nn.Linear(20, 2))
+    model = torch.nn.Sequential(nn.Linear(3, 3), nn.Tanh(), nn.Linear(3, 2))
     setattr(model, 'output_size', 2)
     model_params = list(model.parameters())
     setattr(model, 'n_layers', len(model_params))  # number of parameter groups
@@ -27,11 +26,25 @@ def class_Xy():
     return X, y
 
 
+@pytest.fixture
+def reg_Xy():
+    torch.manual_seed(711)
+    X = torch.randn(10, 3)
+    y = torch.randn(10, 2)
+    return X, y
+
+
 @pytest.fixture
 def complex_model():
     torch.manual_seed(711)
-    model = torch.nn.Sequential(nn.Conv2d(3, 4, 2, 2), nn.Flatten(), nn.Tanh(),
-                                nn.Linear(16, 20), nn.Tanh(), Scale(), Bias(), nn.Linear(20, 2))
+    model = torch.nn.Sequential(
+        nn.Conv2d(3, 4, 2, 2),
+        nn.Flatten(),
+        nn.Tanh(),
+        nn.Linear(16, 20),
+        nn.Tanh(),
+        nn.Linear(20, 2),
+    )
     setattr(model, 'output_size', 2)
     model_params = list(model.parameters())
     setattr(model, 'n_layers', len(model_params))  # number of parameter groups
@@ -47,8 +60,9 @@ def complex_class_Xy():
     return X, y
 
 
-def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
+def test_diag_ggn_cls_asdl_against_backpack_full(class_Xy, model):
     X, y = class_Xy
+    model2 = deepcopy(model)
     backend = AsdlGGN(model, 'classification', stochastic=False)
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
@@ -59,14 +73,34 @@ def test_diag_ggn_cls_kazuki_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackGGN(model, 'classification', stochastic=False)
+    backend = BackPackGGN(model2, 'classification', stochastic=False)
+    loss_f, H_ggn = backend.full(X, y)
+    assert torch.allclose(loss, loss_f)
+    assert torch.allclose(dggn, H_ggn.diagonal())
+
+
+def test_diag_ggn_reg_asdl_against_backpack_full(reg_Xy, model):
+    X, y = reg_Xy
+    model2 = deepcopy(model)
+    backend = AsdlGGN(model, 'regression', stochastic=False)
+    loss, dggn = backend.diag(X[:5], y[:5])
+    loss2, dggn2 = backend.diag(X[5:], y[5:])
+    loss += loss2
+    dggn += dggn2
+
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackGGN(model2, 'regression', stochastic=False)
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diagonal())
 
 
-def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
+def test_diag_ef_cls_asdl_against_backpack_full(class_Xy, model):
     X, y = class_Xy
+    model2 = deepcopy(model)
     backend = AsdlEF(model, 'classification')
     loss, dggn = backend.diag(X[:5], y[:5])
     loss2, dggn2 = backend.diag(X[5:], y[5:])
@@ -77,13 +111,32 @@ def test_diag_ef_cls_kazuki_against_backpack_full(class_Xy, model):
     assert len(dggn) == model.n_params
 
     # check against manually computed full GGN:
-    backend = BackPackEF(model, 'classification')
+    backend = BackPackEF(model2, 'classification')
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(dggn, H_ggn.diagonal())
 
 
-def test_diag_ggn_stoch_cls_kazuki(class_Xy, model):
+def test_diag_ef_reg_asdl_against_backpack_full(reg_Xy, model):
+    X, y = reg_Xy
+    model2 = deepcopy(model)
+    backend = AsdlEF(model, 'regression')
+    loss, dggn = backend.diag(X[:5], y[:5])
+    loss2, dggn2 = backend.diag(X[5:], y[5:])
+    loss += loss2
+    dggn += dggn2
+
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+
+    # check against manually computed full GGN:
+    backend = BackPackEF(model2, 'regression')
+    loss_bp, dggn_bp = backend.diag(X, y)
+    assert torch.allclose(loss, loss_bp)
+    assert torch.allclose(dggn, dggn_bp)  # H_ggn.diagonal())
+
+
+def test_diag_ggn_stoch_cls_asdl(class_Xy, model):
     X, y = class_Xy
     backend = AsdlGGN(model, 'classification', stochastic=True)
     loss, dggn = backend.diag(X, y)
@@ -97,8 +150,43 @@ def test_diag_ggn_stoch_cls_kazuki(class_Xy, model):
     assert torch.allclose(dggn, dggn_ns, atol=1e-8, rtol=1e1)
 
 
+def test_kron_ggn_cls_asdl_against_backpack(class_Xy, model):
+    X, y = class_Xy
+    model2 = deepcopy(model)
+    backend = AsdlGGN(model, 'classification', stochastic=False)
+    loss, kron = backend.kron(X[:5], y[:5], len(y))
+    loss2, kron2 = backend.kron(X[5:], y[5:], len(y))
+    loss += loss2
+    kron += kron2
+    backend_bp = BackPackGGN(model2, 'classification', stochastic=False)
+    loss_bp, kron_bp = backend_bp.kron(X, y, len(y))
+
+    # sanity check size of diag ggn
+    assert len(kron.diag()) == model.n_params == len(kron_bp.diag())
+    assert torch.allclose(loss, loss_bp)
+    assert torch.allclose(kron.diag(), kron_bp.diag())
+
+
+def test_kron_ggn_reg_asdl_against_backpack(reg_Xy, model):
+    X, y = reg_Xy
+    model2 = deepcopy(model)
+    backend = AsdlGGN(model, 'regression', stochastic=False)
+    # loss, kron = backend.kron(X, y, len(y))
+    loss, kron = backend.kron(X[:5], y[:5], len(y))
+    loss2, kron2 = backend.kron(X[5:], y[5:], len(y))
+    loss += loss2
+    kron += kron2
+    backend_bp = BackPackGGN(model2, 'regression', stochastic=False)
+    loss_bp, kron_bp = backend_bp.kron(X, y, len(y))
+
+    # sanity check size of diag ggn
+    assert len(kron.diag()) == model.n_params == len(kron_bp.diag())
+    assert torch.allclose(loss, loss_bp)
+    assert torch.allclose(kron.diag(), kron_bp.diag())
+
+
 @pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
-def test_kron_kazuki_vs_diag_class(class_Xy, model, Backend):
+def test_kron_asdl_vs_diag_class(class_Xy, model, Backend):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = class_Xy
     backend = Backend(model, 'classification')
@@ -110,7 +198,19 @@ def test_kron_kazuki_vs_diag_class(class_Xy, model, Backend):
 
 
 @pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
-def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
+def test_kron_asdl_vs_diag_reg(reg_Xy, model, Backend):
+    # For a single data point, Kron is exact and should equal diag GGN
+    X, y = reg_Xy
+    backend = Backend(model, 'regression')
+    loss, dggn = backend.diag(X[:1], y[:1], N=1)
+    # sanity check size of diag ggn
+    assert len(dggn) == model.n_params
+    loss, kron = backend.kron(X[:1], y[:1], N=1)
+    assert torch.allclose(kron.diag(), dggn)
+
+
+@pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
+def test_kron_batching_correction_asdl(class_Xy, model, Backend):
     X, y = class_Xy
     backend = Backend(model, 'classification')
     loss, kron = backend.kron(X, y, N=len(X))
@@ -127,7 +227,7 @@ def test_kron_batching_correction_kazuki(class_Xy, model, Backend):
 
 
 @pytest.mark.parametrize('Backend', [AsdlGGN, AsdlEF])
-def test_kron_summing_up_vs_diag_kazuki(class_Xy, model, Backend):
+def test_kron_summing_up_vs_diag_asdl(class_Xy, model, Backend):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = class_Xy
     backend = Backend(model, 'classification')
@@ -136,7 +236,7 @@ def test_kron_summing_up_vs_diag_kazuki(class_Xy, model, Backend):
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-1)
 
 
-def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
+def test_complex_diag_ggn_stoch_cls_asdl(complex_class_Xy, complex_model):
     X, y = complex_class_Xy
     backend = AsdlGGN(complex_model, 'classification', stochastic=True)
     loss, dggn = backend.diag(X, y)
@@ -150,7 +250,7 @@ def test_complex_diag_ggn_stoch_cls_kazuki(complex_class_Xy, complex_model):
 
 
 @pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
-def test_complex_kron_kazuki_vs_diag_kazuki(complex_class_Xy, complex_model, Backend):
+def test_complex_kron_asdl_vs_diag_asdl(complex_class_Xy, complex_model, Backend):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = complex_class_Xy
     backend = Backend(complex_model, 'classification')
@@ -162,7 +262,9 @@ def test_complex_kron_kazuki_vs_diag_kazuki(complex_class_Xy, complex_model, Bac
 
 
 @pytest.mark.parametrize('Backend', [AsdlEF, AsdlGGN])
-def test_complex_kron_batching_correction_kazuki(complex_class_Xy, complex_model, Backend):
+def test_complex_kron_batching_correction_asdl(
+    complex_class_Xy, complex_model, Backend
+):
     X, y = complex_class_Xy
     backend = Backend(complex_model, 'classification')
     loss, kron = backend.kron(X, y, N=len(X))
@@ -179,7 +281,9 @@ def test_complex_kron_batching_correction_kazuki(complex_class_Xy, complex_model
 
 
 @pytest.mark.parametrize('Backend', [AsdlGGN, AsdlEF])
-def test_complex_kron_summing_up_vs_diag_class_kazuki(complex_class_Xy, complex_model, Backend):
+def test_complex_kron_summing_up_vs_diag_class_asdl(
+    complex_class_Xy, complex_model, Backend
+):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = complex_class_Xy
     backend = Backend(complex_model, 'classification')
@@ -188,6 +292,28 @@ def test_complex_kron_summing_up_vs_diag_class_kazuki(complex_class_Xy, complex_
     assert torch.allclose(kron.diag().norm(), dggn.norm(), rtol=1e-2)
 
 
+def test_kron_summation_ggn_class(class_Xy, model):
+    X, y = class_Xy
+    X = torch.repeat_interleave(X[:1], 7, 0)
+    y = torch.repeat_interleave(y[:1], 7, 0)
+    backend = AsdlGGN(model, 'classification', stochastic=False)
+    loss_dg, diag = backend.diag(X, y)
+    loss, kron = backend.kron(X, y, N=7)
+    assert torch.allclose(loss_dg, loss)
+    assert torch.allclose(kron.diag(), diag)
+
+
+def test_kron_summation_ggn_reg(reg_Xy, model):
+    X, y = reg_Xy
+    X = torch.repeat_interleave(X[:1], 7, 0)
+    y = torch.repeat_interleave(y[:1], 7, 0)
+    backend = AsdlGGN(model, 'regression', stochastic=False)
+    loss_dg, diag = backend.diag(X, y)
+    loss, kron = backend.kron(X, y, N=7)
+    assert torch.allclose(loss_dg, loss)
+    assert torch.allclose(kron.diag(), diag)
+
+
 def test_kron_normalization_ggn_class(class_Xy, model):
     X, y = class_Xy
     xi, yi = X[:1], y[:1]
@@ -197,7 +323,7 @@ def test_kron_normalization_ggn_class(class_Xy, model):
     loss_true = 7 * loss
     X = torch.repeat_interleave(xi, 7, 0)
     y = torch.repeat_interleave(yi, 7, 0)
-    loss_test, kron_test  = backend.kron(X, y, N=7)
+    loss_test, kron_test = backend.kron(X, y, N=7)
     assert torch.allclose(kron_true.diag(), kron_test.diag())
     assert torch.allclose(loss_true, loss_test)
 
@@ -211,6 +337,20 @@ def test_kron_normalization_ef_class(class_Xy, model):
     loss_true = 7 * loss
     X = torch.repeat_interleave(xi, 7, 0)
     y = torch.repeat_interleave(yi, 7, 0)
-    loss_test, kron_test  = backend.kron(X, y, N=7)
+    loss_test, kron_test = backend.kron(X, y, N=7)
+    assert torch.allclose(kron_true.diag(), kron_test.diag())
+    assert torch.allclose(loss_true, loss_test)
+
+
+def test_kron_normalization_ggn_reg(reg_Xy, model):
+    X, y = reg_Xy
+    xi, yi = X[:1], y[:1]
+    backend = AsdlGGN(model, 'regression', stochastic=False)
+    loss, kron = backend.kron(xi, yi, N=1)
+    kron_true = 7 * kron
+    loss_true = 7 * loss
+    X = torch.repeat_interleave(xi, 7, 0)
+    y = torch.repeat_interleave(yi, 7, 0)
+    loss_test, kron_test = backend.kron(X, y, N=7)
     assert torch.allclose(kron_true.diag(), kron_test.diag())
     assert torch.allclose(loss_true, loss_test)
diff --git a/tests/test_curv_backends_curvlinops.py b/tests/test_curv_backends_curvlinops.py
index ff041c35..2baaf10b 100644
--- a/tests/test_curv_backends_curvlinops.py
+++ b/tests/test_curv_backends_curvlinops.py
@@ -3,7 +3,16 @@
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
-from laplace.curvature import CurvlinopsGGN, CurvlinopsEF, CurvlinopsHessian, BackPackGGN, AsdlHessian, AsdlEF, AsdlGGN
+from laplace.curvature import (
+    CurvlinopsGGN,
+    CurvlinopsEF,
+    CurvlinopsHessian,
+    BackPackGGN,
+    AsdlHessian,
+    AsdlEF,
+    AsdlGGN,
+)
+
 
 @pytest.fixture(autouse=True)
 def run_around_tests():
@@ -33,8 +42,14 @@ def class_Xy():
 @pytest.fixture
 def complex_model():
     torch.manual_seed(711)
-    model = torch.nn.Sequential(nn.Conv2d(3, 4, 2, 2), nn.Flatten(), nn.Tanh(),
-                                nn.Linear(16, 20), nn.Tanh(), nn.Linear(20, 2))
+    model = torch.nn.Sequential(
+        nn.Conv2d(3, 4, 2, 2),
+        nn.Flatten(),
+        nn.Tanh(),
+        nn.Linear(16, 20),
+        nn.Tanh(),
+        nn.Linear(20, 2),
+    )
     setattr(model, 'output_size', 2)
     model_params = list(model.parameters())
     setattr(model, 'n_layers', len(model_params))  # number of parameter groups
@@ -77,7 +92,7 @@ def test_full_hess_curvlinops_vs_asdl(class_Xy, reg_Xy, model, loss_type):
     loss_ref, H_ref = backend.full(X, y)
 
     assert torch.allclose(loss, loss_ref)
-    assert torch.allclose(H, H_ref, rtol=5e-5)
+    assert torch.allclose(H, H_ref, rtol=5e-4)
 
 
 def test_full_ggn_curvlinops_vs_asdl(class_Xy, model):
@@ -162,7 +177,9 @@ def test_kron_ggn_stochastic(class_Xy, reg_Xy, model, loss_type):
 
 
 @pytest.mark.parametrize('loss_type', ['classification', 'regression'])
-def test_kron_ggn_set_kfac_approx(complex_class_Xy, complex_reg_Xy, complex_model, loss_type):
+def test_kron_ggn_set_kfac_approx(
+    complex_class_Xy, complex_reg_Xy, complex_model, loss_type
+):
     X, y = complex_class_Xy if loss_type == 'classification' else complex_reg_Xy
 
     backend = CurvlinopsGGN(complex_model, loss_type, stochastic=False)
@@ -244,7 +261,9 @@ def test_complex_diag_ggn_stoch_cls_curvlinops(complex_class_Xy, complex_model):
 
 
 @pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
-def test_complex_kron_curvlinops_vs_diag_curvlinops(complex_class_Xy, complex_model, Backend):
+def test_complex_kron_curvlinops_vs_diag_curvlinops(
+    complex_class_Xy, complex_model, Backend
+):
     # For a single data point, Kron is exact and should equal diag GGN
     X, y = complex_class_Xy
     backend = Backend(complex_model, 'classification')
@@ -256,7 +275,9 @@ def test_complex_kron_curvlinops_vs_diag_curvlinops(complex_class_Xy, complex_mo
 
 
 @pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
-def test_complex_kron_batching_correction_curvlinops(complex_class_Xy, complex_model, Backend):
+def test_complex_kron_batching_correction_curvlinops(
+    complex_class_Xy, complex_model, Backend
+):
     X, y = complex_class_Xy
     backend = Backend(complex_model, 'classification')
     loss, kron = backend.kron(X, y, N=len(X))
@@ -273,7 +294,9 @@ def test_complex_kron_batching_correction_curvlinops(complex_class_Xy, complex_m
 
 
 @pytest.mark.parametrize('Backend', [CurvlinopsEF, CurvlinopsGGN])
-def test_complex_kron_summing_up_vs_diag_class_curvlinops(complex_class_Xy, complex_model, Backend):
+def test_complex_kron_summing_up_vs_diag_class_curvlinops(
+    complex_class_Xy, complex_model, Backend
+):
     # For a single data point, Kron is exact and should equal diag class_Xy
     X, y = complex_class_Xy
     backend = Backend(complex_model, 'classification')
@@ -291,7 +314,7 @@ def test_kron_normalization_ggn_class(class_Xy, model):
     loss_true = 7 * loss
     X = torch.repeat_interleave(xi, 7, 0)
     y = torch.repeat_interleave(yi, 7, 0)
-    loss_test, kron_test  = backend.kron(X, y, N=7)
+    loss_test, kron_test = backend.kron(X, y, N=7)
     assert torch.allclose(kron_true.diag(), kron_test.diag())
     assert torch.allclose(loss_true, loss_test)
 
@@ -305,6 +328,6 @@ def test_kron_normalization_ef_class(class_Xy, model):
     loss_true = 7 * loss
     X = torch.repeat_interleave(xi, 7, 0)
     y = torch.repeat_interleave(yi, 7, 0)
-    loss_test, kron_test  = backend.kron(X, y, N=7)
+    loss_test, kron_test = backend.kron(X, y, N=7)
     assert torch.allclose(kron_true.diag(), kron_test.diag())
     assert torch.allclose(loss_true, loss_test)
diff --git a/tests/test_curv_backends_interface.py b/tests/test_curv_backends_interface.py
index c5499705..d3270b94 100644
--- a/tests/test_curv_backends_interface.py
+++ b/tests/test_curv_backends_interface.py
@@ -4,9 +4,19 @@
 from torch import nn
 from torch.nn.utils import parameters_to_vector
 
-from asdfghjkl.operations import Bias, Scale
+from asdl.operations import Bias, Scale
 
-from laplace.curvature import CurvlinopsGGN, CurvlinopsEF, BackPackGGN, BackPackEF, GGNInterface, EFInterface, AsdlGGN, AsdlEF, CurvatureInterface
+from laplace.curvature import (
+    CurvlinopsGGN,
+    CurvlinopsEF,
+    BackPackGGN,
+    BackPackEF,
+    GGNInterface,
+    EFInterface,
+    AsdlGGN,
+    AsdlEF,
+    CurvatureInterface,
+)
 
 
 @pytest.fixture
@@ -39,8 +49,14 @@ def reg_Xy():
 @pytest.fixture
 def complex_model():
     torch.manual_seed(711)
-    model = torch.nn.Sequential(nn.Conv2d(3, 4, 2, 2), nn.Flatten(), nn.Tanh(),
-                                nn.Linear(16, 20), nn.Tanh(), Scale(), Bias(), nn.Linear(20, 2))
+    model = torch.nn.Sequential(
+        nn.Conv2d(3, 4, 2, 2),
+        nn.Flatten(),
+        nn.Tanh(),
+        nn.Linear(16, 20),
+        nn.Tanh(),
+        nn.Linear(20, 2),
+    )
     setattr(model, 'output_size', 2)
     model_params = list(model.parameters())
     setattr(model, 'n_layers', len(model_params))  # number of parameter groups
@@ -109,7 +125,9 @@ def test_diag_ggn_cls_curvlinops_against_backpack_full(class_Xy, model):
     assert torch.allclose(dggn, H_ggn.diag())
 
 
-def test_diag_ggn_complex_cls_curvlinops_against_backpack_full(complex_class_Xy, complex_model):
+def test_diag_ggn_complex_cls_curvlinops_against_backpack_full(
+    complex_class_Xy, complex_model
+):
     torch.manual_seed(0)
     np.random.seed(0)
     X, y = complex_class_Xy
@@ -347,4 +365,3 @@ def test_full_ef_reg_curvlinops_against_backpack_full(reg_Xy, model):
     loss_f, H_ggn = backend.full(X, y)
     assert torch.allclose(loss, loss_f)
     assert torch.allclose(fggn, H_ggn, atol=0.0001)
-
diff --git a/tests/test_jacobians.py b/tests/test_jacobians.py
index 2fb64977..b32bfa40 100644
--- a/tests/test_jacobians.py
+++ b/tests/test_jacobians.py
@@ -34,7 +34,9 @@ def X():
     return torch.randn(200, 3)
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface])
+@pytest.mark.parametrize(
+    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+)
 def test_linear_jacobians(linear_model, X, backend_cls):
     # jacobian of linear model is input X.
     backend = backend_cls(linear_model, 'classification')
@@ -46,31 +48,37 @@ def test_linear_jacobians(linear_model, X, backend_cls):
     assert torch.allclose(f, linear_model(X), atol=1e-5)
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface])
+@pytest.mark.parametrize(
+    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+)
 def test_jacobians_singleoutput(singleoutput_model, X, backend_cls):
     model = singleoutput_model
     backend = backend_cls(model, 'classification')
     Js, f = backend.jacobians(X)
     Js_naive, f_naive = jacobians_naive(model, X)
     assert Js.shape == Js_naive.shape
-    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.abs(Js - Js_naive).max() < 1e-6
     assert torch.allclose(model(X), f_naive)
     assert torch.allclose(f, f_naive)
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface])
+@pytest.mark.parametrize(
+    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+)
 def test_jacobians_multioutput(multioutput_model, X, backend_cls):
     model = multioutput_model
     backend = backend_cls(model, 'classification')
     Js, f = backend.jacobians(X)
     Js_naive, f_naive = jacobians_naive(model, X)
     assert Js.shape == Js_naive.shape
-    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.abs(Js - Js_naive).max() < 1e-6
     assert torch.allclose(model(X), f_naive)
     assert torch.allclose(f, f_naive)
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface])
+@pytest.mark.parametrize(
+    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+)
 def test_last_layer_jacobians_singleoutput(singleoutput_model, X, backend_cls):
     model = FeatureExtractor(singleoutput_model)
     backend = backend_cls(model, 'classification')
@@ -78,12 +86,14 @@ def test_last_layer_jacobians_singleoutput(singleoutput_model, X, backend_cls):
     _, phi = model.forward_with_features(X)
     Js_naive, f_naive = jacobians_naive(model.last_layer, phi)
     assert Js.shape == Js_naive.shape
-    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.abs(Js - Js_naive).max() < 1e-6
     assert torch.allclose(model(X), f_naive)
     assert torch.allclose(f, f_naive)
 
 
-@pytest.mark.parametrize('backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface])
+@pytest.mark.parametrize(
+    'backend_cls', [CurvatureInterface, AsdlInterface, BackPackInterface]
+)
 def test_last_layer_jacobians_multioutput(multioutput_model, X, backend_cls):
     model = FeatureExtractor(multioutput_model)
     backend = backend_cls(model, 'classification')
@@ -91,7 +101,7 @@ def test_last_layer_jacobians_multioutput(multioutput_model, X, backend_cls):
     _, phi = model.forward_with_features(X)
     Js_naive, f_naive = jacobians_naive(model.last_layer, phi)
     assert Js.shape == Js_naive.shape
-    assert torch.abs(Js-Js_naive).max() < 1e-6
+    assert torch.abs(Js - Js_naive).max() < 1e-6
     assert torch.allclose(model(X), f_naive)
     assert torch.allclose(f, f_naive)
 
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index 73d00a0b..757130c7 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -6,8 +6,8 @@
 
 from laplace.utils import Kron, block_diag
 from laplace.utils import kron as kron_prod
-from laplace.curvature import BackPackGGN
-from tests.utils import get_psd_matrix, jacobians_naive
+from laplace.curvature import BackPackGGN, AsdlGGN
+from tests.utils import get_psd_matrix, get_diag_psd_matrix, jacobians_naive
 
 
 torch.set_default_tensor_type(torch.DoubleTensor)
@@ -39,8 +39,17 @@ def test_init_from_model(model):
             assert np.prod(fi.shape) == exp_fi
 
 
+def test_init_from_iterable(model):
+    kron = Kron.init_from_model(model.parameters(), 'cpu')
+    expected_sizes = [[20*20, 3*3], [20*20], [2*2, 20*20], [2*2]]
+    for facs, exp_facs in zip(kron.kfacs, expected_sizes):
+        for fi, exp_fi in zip(facs, exp_facs):
+            assert torch.all(fi == 0)
+            assert np.prod(fi.shape) == exp_fi
+
+
 def test_addition(model):
-    kron = Kron.init_from_model(model, 'cpu')
+    kron = Kron.init_from_model(model.parameters(), 'cpu')
     expected_sizes = [[20, 3], [20], [2, 20], [2]]
     to_add = Kron([[torch.randn(i, i) for i in sizes] for sizes in expected_sizes])
     kron += to_add
@@ -48,6 +57,7 @@ def test_addition(model):
         for fi, exp_fi in zip(facs, exp_facs):
             assert torch.allclose(fi, exp_fi)
 
+
 def test_multiplication():
     # kron * x should be the same as the expanded kronecker product * x
     expected_sizes = [[20, 3], [20], [2, 20], [2]]
@@ -62,10 +72,12 @@ def test_multiplication():
             facs = kron_prod(*facs)
             assert torch.allclose(exp, facs)
 
+
 def test_decompose():
     expected_sizes = [[20, 3], [20], [2, 20], [2]]
     P = 20 * 3 + 20 + 2 * 20 + 2
     torch.manual_seed(7171)
+    # Dense Kronecker factors.
     kfacs = [[get_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
     kron = Kron(kfacs)
     kron_decomp = kron.decompose()
@@ -84,6 +96,25 @@ def test_decompose():
     SW_kron = kron.bmm(W)
     SW_kron_decomp = kron_decomp.bmm(W, exponent=1)
     assert torch.allclose(SW_kron, SW_kron_decomp)
+    # Diagonal Kronecker factors.
+    diag_kfacs = [[get_diag_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
+    kron = Kron(diag_kfacs)
+    kron_decomp = kron.decompose()
+    for facs, Qs, ls in zip(kron.kfacs, kron_decomp.eigenvectors, kron_decomp.eigenvalues):
+        if len(facs) == 1:
+            H, Q, l = facs[0], Qs[0], ls[0]
+            reconstructed = (Q @ torch.diag(l) @ Q.T).diag()
+            assert torch.allclose(H, reconstructed, rtol=1e-3)
+        if len(facs) == 2:
+            gtruth = kron_prod(facs[0].diag(), facs[1].diag())
+            rec_1 = Qs[0] @ torch.diag(ls[0]) @ Qs[0].T
+            rec_2 = Qs[1] @ torch.diag(ls[1]) @ Qs[1].T
+            reconstructed = kron_prod(rec_1, rec_2)
+            assert torch.allclose(gtruth, reconstructed, rtol=1e-2)
+    W = torch.randn(P)
+    SW_kron = kron.bmm(W)
+    SW_kron_decomp = kron_decomp.bmm(W, exponent=1)
+    assert torch.allclose(SW_kron, SW_kron_decomp)
 
 
 def test_logdet_consistent():
@@ -93,17 +124,23 @@ def test_logdet_consistent():
     kron = Kron(kfacs)
     kron_decomp = kron.decompose()
     assert torch.allclose(kron.logdet(), kron_decomp.logdet())
+    diag_kfacs = [[get_diag_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
+    kron = Kron(diag_kfacs)
+    kron_decomp = kron.decompose()
+    assert torch.allclose(kron.logdet(), kron_decomp.logdet())
 
 
-def test_bmm(small_model):
+def test_bmm_dense(small_model):
     model = small_model
     # model = single_output_model
     X = torch.randn(5, 3)
     y = torch.randn(5, 2)
+
+    # Dense Kronecker factors.
     backend = BackPackGGN(model, 'regression', stochastic=False)
-    loss, kron = backend.kron(X, y, N=5)
+    _, kron = backend.kron(X, y, N=5)
     kron_decomp = kron.decompose()
-    Js, f = jacobians_naive(model, X)
+    Js, _ = jacobians_naive(model, X)
     blocks = list()
     for F in kron.kfacs:
         if len(F) == 1:
@@ -158,9 +195,78 @@ def test_bmm(small_model):
     assert torch.allclose(JS, JS_nodecomp)
 
 
+@pytest.mark.skip(reason='For compatibility with ASDL master')
+def test_bmm_diag(small_model):
+    model = small_model
+    # model = single_output_model
+    X = torch.randn(5, 3)
+    y = torch.randn(5, 2)
+
+    # Diagonal Kronecker factors.
+    backend = AsdlGGN(model, 'regression', stochastic=False)
+    _, kron = backend.kron(X, y, N=5, diag_A=True, diag_B=True)
+    kron_decomp = kron.decompose()
+    Js, _ = jacobians_naive(model, X)
+    blocks = list()
+    for F in kron.kfacs:
+        F0 = F[0] if F[0].ndim > 1 else F[0].diag()
+        if len(F) == 1:
+            blocks.append(F0)
+        else:
+            F1 = F[1] if F[1].ndim > 1 else F[1].diag()
+            blocks.append(kron_prod(F0, F1))
+    S = block_diag(blocks)
+    assert torch.allclose(S, S.T)
+    assert torch.allclose(S.diagonal(), kron.diag())
+
+    # test J @ Kron_decomp @ Jt (square form)
+    JS = kron_decomp.bmm(Js, exponent=1)
+    JS_true = Js @ S
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
+    JSJ = torch.bmm(JS, Js.transpose(1,2))
+    assert torch.allclose(JSJ, JSJ_true)
+    assert torch.allclose(JS, JS_true)
+
+    # test J @ Kron @ Jt (square form)
+    JS_nodecomp = kron.bmm(Js)
+    JSJ_nodecomp = torch.bmm(JS_nodecomp, Js.transpose(1,2))
+    assert torch.allclose(JSJ_nodecomp, JSJ)
+    assert torch.allclose(JS_nodecomp, JS)
+
+    # test J @ S_inv @ J (funcitonal variance)
+    JSJ = kron_decomp.inv_square_form(Js)
+    S_inv = S.inverse()
+    JSJ_true = torch.bmm(Js @ S_inv, Js.transpose(1,2))
+    assert torch.allclose(JSJ, JSJ_true)
+
+    # test J @ S^-1/2  (sampling)
+    JS = kron_decomp.bmm(Js, exponent=-1/2)
+    JSJ = torch.bmm(JS, Js.transpose(1,2))
+    l, Q = torch.linalg.eigh(S_inv, UPLO='U')
+    JS_true = Js @ Q @ torch.diag(torch.sqrt(l)) @ Q.T
+    JSJ_true = torch.bmm(JS_true, Js.transpose(1,2))
+    assert torch.allclose(JS, JS_true)
+    assert torch.allclose(JSJ, JSJ_true)
+
+    # test different Js shapes:
+    # 2 - dimensional
+    JS = kron_decomp.bmm(Js[:, 0, :].squeeze(), exponent=1)
+    JS_nodecomp = kron.bmm(Js[:, 0, :].squeeze())
+    JS_true = Js[:, 0, :].squeeze() @ S
+    assert torch.allclose(JS, JS_true)
+    assert torch.allclose(JS, JS_nodecomp)
+    # 1 - dimensional
+    JS = kron_decomp.bmm(Js[0, 0, :].squeeze(), exponent=1)
+    JS_nodecomp = kron.bmm(Js[0, 0, :].squeeze())
+    JS_true = Js[0, 0, :].squeeze() @ S
+    assert torch.allclose(JS, JS_true)
+    assert torch.allclose(JS, JS_nodecomp)
+
+
 def test_matrix_consistent():
     expected_sizes = [[20, 3], [20], [2, 20], [2]]
     torch.manual_seed(7171)
+    # Dense Kronecker factors.
     kfacs = [[get_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
     kron = Kron(kfacs)
     kron_decomp = kron.decompose()
@@ -170,15 +276,13 @@ def test_matrix_consistent():
     M_true.diagonal().add_(3.4)
     kron_decomp += torch.tensor(3.4)
     assert torch.allclose(M_true, kron_decomp.to_matrix(exponent=1))
-
-    
-def test_diag():
-    expected_sizes = [[20, 3], [20], [2, 20], [2]]
-    torch.manual_seed(7171)
-    kfacs = [[get_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
-    kron = Kron(kfacs)
+    # Diagonal Kronecker factors.
+    diag_kfacs = [[get_diag_psd_matrix(i) for i in sizes] for sizes in expected_sizes]
+    kron = Kron(diag_kfacs)
     kron_decomp = kron.decompose()
-    assert torch.allclose(kron.diag(), kron_decomp.diag(exponent=1)) 
-    assert torch.allclose(kron.diag(), torch.diag(kron.to_matrix()))
-    assert torch.allclose(kron_decomp.diag(), torch.diag(kron_decomp.to_matrix()))
-    assert torch.allclose(kron_decomp.diag(exponent=-1), torch.diag(kron_decomp.to_matrix(exponent=-1)))
+    assert torch.allclose(kron.to_matrix(), kron_decomp.to_matrix(exponent=1))
+    assert torch.allclose(kron.to_matrix().inverse(), kron_decomp.to_matrix(exponent=-1))
+    M_true = kron.to_matrix()
+    M_true.diagonal().add_(3.4)
+    kron_decomp += torch.tensor(3.4)
+    assert torch.allclose(M_true, kron_decomp.to_matrix(exponent=1))
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
new file mode 100644
index 00000000..2edb72e0
--- /dev/null
+++ b/tests/test_metrics.py
@@ -0,0 +1,47 @@
+import torch
+from torch.nn import functional as F
+from laplace.utils import RunningNLLMetric
+import math
+
+
+def test_running_nll_metric():
+    metric = RunningNLLMetric()
+    all_probs, all_targets = [], []
+
+    for _ in range(10):
+        probs = torch.softmax(torch.randn(3, 5, 10), dim=-1)
+        targets = torch.randint(10, size=(3, 5))
+        metric.update(probs, targets)
+        all_probs.append(probs)
+        all_targets.append(targets)
+
+    all_probs, all_targets = torch.cat(all_probs, 0), torch.cat(all_targets, 0)
+
+    nll_running = metric.compute().item()
+    nll_offline = F.nll_loss(all_probs.log().flatten(end_dim=-2), all_targets.flatten()).item()
+
+    assert math.isclose(nll_running, nll_offline, rel_tol=1e-7)
+
+
+def test_running_nll_metric_ignore_idx():
+    ignore_idx = -1232
+    metric = RunningNLLMetric(ignore_index=ignore_idx)
+    all_probs, all_targets = [], []
+
+    for _ in range(10):
+        probs = torch.softmax(torch.randn(3, 5, 10), dim=-1)
+        targets = torch.randint(10, size=(3, 5))
+        mask = torch.FloatTensor(*targets.shape).uniform_() > 0.5  # ~50% zeros
+        targets[mask] = ignore_idx  # ~50% changed to ignore_idx
+        all_probs.append(probs)
+        all_targets.append(targets)
+        metric.update(probs, targets)
+
+    all_probs, all_targets = torch.cat(all_probs, 0), torch.cat(all_targets, 0)
+
+    nll_running = metric.compute().item()
+    nll_offline = F.nll_loss(all_probs.log().flatten(end_dim=-2), all_targets.flatten(), ignore_index=ignore_idx).item()
+
+    print(nll_running, nll_offline)
+
+    assert math.isclose(nll_running, nll_offline, rel_tol=1e-7)
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
index b2bac6b4..612a433c 100644
--- a/tests/test_serialization.py
+++ b/tests/test_serialization.py
@@ -26,7 +26,9 @@
 
 torch.manual_seed(240)
 torch.set_default_tensor_type(torch.DoubleTensor)
-flavors_no_llla = [FullLaplace, KronLaplace, DiagLaplace, LowRankLaplace]
+lrlaplace_param = pytest.param(LowRankLaplace, marks=pytest.mark.xfail(reason='Unimplemented in the new ASDL'))
+flavors = [FullLaplace, KronLaplace, DiagLaplace, lrlaplace_param, FullLLLaplace, KronLLLaplace, DiagLLLaplace]
+flavors_no_llla = [FullLaplace, KronLaplace, DiagLaplace, lrlaplace_param]
 flavors_llla = [FullLLLaplace, KronLLLaplace, DiagLLLaplace]
 flavors_subnet = [DiagSubnetLaplace, FullSubnetLaplace]
 flavors = flavors_llla + flavors_no_llla
@@ -102,7 +104,7 @@ def test_serialize(laplace, model, reg_loader):
     assert torch.allclose(f_var, f_var2)
 
 
-@pytest.mark.parametrize('laplace', set(flavors_no_llla) - {LowRankLaplace})
+@pytest.mark.parametrize('laplace', flavors_no_llla[:-1])
 def test_serialize_override(laplace, model, reg_loader):
     la = laplace(model, 'regression')
     la.fit(reg_loader)
diff --git a/tests/test_subnetlaplace.py b/tests/test_subnetlaplace.py
index d854e8a2..92c888b3 100644
--- a/tests/test_subnetlaplace.py
+++ b/tests/test_subnetlaplace.py
@@ -332,7 +332,7 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
         lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
                       subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
         lap.fit(loader)
-        assert torch.allclose(lllap.H, lap.H, rtol=1e-4)
+        assert torch.allclose(lllap.H, lap.H, rtol=1e-3)
 
         # define valid parameter name subnet mask
         subnetmask_kwargs.update(parameter_names=['0.weight', '1.bias'])
@@ -377,7 +377,7 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
         lap = Laplace(model, likelihood=likelihood, subset_of_weights='subnetwork',
                       subnetwork_indices=subnetmask.indices, hessian_structure=hessian_structure)
         lap.fit(loader)
-        assert torch.allclose(lllap.H, lap.H, rtol=1e-4)
+        assert torch.allclose(lllap.H, lap.H, rtol=1e-3)
 
         # define valid parameter name subnet mask
         subnetmask_kwargs.update(module_names=['0'])
@@ -418,7 +418,7 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
         assert isinstance(lap, SubnetLaplace)
 
         # check that Hessian is identical to that of a full LLLaplace model
-        assert torch.allclose(lllap.H, lap.H, rtol=1e-4)
+        assert torch.allclose(lllap.H, lap.H, rtol=1e-3)
 
         # define valid last-layer subnet mask (with passing the last-layer name)
         subnetmask_kwargs.update(last_layer_name='1')
@@ -437,7 +437,7 @@ def test_layer_subnet_masks(model, likelihood, subnetwork_mask, class_loader, re
         assert isinstance(lap, SubnetLaplace)
 
         # check that Hessian is identical to that of a full LLLaplace model
-        assert torch.allclose(lllap.H, lap.H, rtol=1e-4)
+        assert torch.allclose(lllap.H, lap.H, rtol=1e-3)
 
     # check some parameters
     assert subnetmask.indices.equal(lap.backend.subnetwork_indices)
@@ -479,7 +479,7 @@ def get_subnet_mask(self, train_loader):
     full_lap = Laplace(model, likelihood=likelihood, subset_of_weights='all',
                        hessian_structure=hessian_structure)
     full_lap.fit(loader)
-    assert torch.allclose(full_lap.H, lap.H, rtol=1e-4)
+    assert torch.allclose(full_lap.H, lap.H, rtol=1e-3)
 
 
 @pytest.mark.parametrize('subnetwork_mask,hessian_structure', product(all_subnet_masks, hessian_structures))
diff --git a/tests/test_subset_params.py b/tests/test_subset_params.py
new file mode 100644
index 00000000..01053273
--- /dev/null
+++ b/tests/test_subset_params.py
@@ -0,0 +1,120 @@
+from math import sqrt
+import pytest
+from itertools import product
+import numpy as np
+from copy import deepcopy
+import torch
+from torch import nn
+from torch.distributions.multivariate_normal import MultivariateNormal
+from torch.nn.utils import parameters_to_vector
+from torch.utils.data import DataLoader, TensorDataset
+from torch.distributions import Normal, Categorical
+from torchvision.models import wide_resnet50_2
+
+from laplace.laplace import FullLaplace, KronLaplace, DiagLaplace, LowRankLaplace
+from laplace.utils import KronDecomposed
+from laplace.curvature import AsdlGGN, AsdlHessian, AsdlEF, BackPackEF, BackPackGGN
+from tests.utils import jacobians_naive
+
+
+torch.manual_seed(240)
+torch.set_default_tensor_type(torch.DoubleTensor)
+flavors = [KronLaplace, DiagLaplace, FullLaplace]
+
+
+@pytest.fixture
+def model():
+    model = torch.nn.Sequential(nn.Linear(3, 20), nn.Linear(20, 2))
+    setattr(model, 'output_size', 2)
+    model_params = list(model.parameters())
+    setattr(model, 'n_layers', len(model_params))  # number of parameter groups
+    setattr(model, 'n_params', len(parameters_to_vector(model_params)))
+
+    # Subset of params
+    for p in model.parameters():
+        p.requires_grad = False
+    model[0].weight.requires_grad = True
+
+    return model
+
+
+@pytest.fixture
+def class_loader():
+    X = torch.randn(10, 3)
+    y = torch.randint(2, (10,))
+    return DataLoader(TensorDataset(X, y), batch_size=3)
+
+
+@pytest.fixture
+def reg_loader():
+    X = torch.randn(10, 3)
+    y = torch.randn(10, 2)
+    return DataLoader(TensorDataset(X, y), batch_size=3)
+
+
+@pytest.mark.parametrize('laplace,lh', product(flavors, ['classification', 'regression']))
+def test_incompatible_backend(laplace, lh, model):
+    lap = laplace(model, lh, backend=AsdlEF)
+    lap = laplace(model, lh, backend=AsdlGGN)
+    lap = laplace(model, lh, backend=AsdlHessian)
+
+
+@pytest.mark.parametrize('laplace,lh', product(flavors, ['classification', 'regression']))
+def test_incompatible_backend(laplace, lh, model):
+    with pytest.raises(ValueError):
+        lap = laplace(model, lh, backend=BackPackGGN)
+
+    with pytest.raises(ValueError):
+        lap = laplace(model, lh, backend=BackPackEF)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_mean_clf(laplace, model, class_loader):
+    n_params = model[0].weight.numel()
+    lap = laplace(model, 'classification')
+    lap.fit(class_loader)
+    assert lap.mean.shape == (n_params,)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_mean_reg(laplace, model, reg_loader):
+    n_params = model[0].weight.numel()
+    lap = laplace(model, 'regression')
+    lap.fit(reg_loader)
+    assert lap.mean.shape == (n_params,)
+
+
+def test_post_precision_diag(model, class_loader):
+    n_params = model[0].weight.numel()
+    lap = DiagLaplace(model, 'classification')
+    lap.fit(class_loader)
+    assert lap.posterior_precision.shape == (n_params,)
+
+
+def test_post_precision_kron(model, class_loader):
+    n_params = model[0].weight.numel()
+    lap = KronLaplace(model, 'classification')
+    lap.fit(class_loader)
+    assert lap.posterior_precision.to_matrix().shape == (n_params, n_params)
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_predictive(laplace, model, class_loader):
+    lap = laplace(model, 'classification')
+    lap.fit(class_loader)
+    lap(torch.randn(5, 3), pred_type='nn', link_approx='mc')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_marglik(laplace, model, class_loader):
+    lap = laplace(model, 'classification')
+    lap.fit(class_loader)
+    lap.optimize_prior_precision(method='marglik')
+
+
+@pytest.mark.parametrize('laplace', flavors)
+def test_marglik(laplace, model, class_loader):
+    lap = laplace(model, 'classification')
+    lap.fit(class_loader)
+    lap.optimize_prior_precision(method='gridsearch', val_loader=class_loader,
+                                 pred_type='nn', link_approx='mc')
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 22f01541..d6f9165c 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,5 +1,8 @@
 import torch
-from laplace.utils import invsqrt_precision, diagonal_add_scalar, symeig, normal_samples
+from torch.utils.data import TensorDataset, DataLoader
+from laplace import Laplace
+from laplace.utils import invsqrt_precision, diagonal_add_scalar, symeig, normal_samples, validate, get_nll, RunningNLLMetric
+import math
 
 
 def test_sqrt_precision():
@@ -23,7 +26,7 @@ def test_symeig_custom():
     l2, W2 = symeig(M)
     assert torch.allclose(l1, l2)
     assert torch.allclose(W1, W2)
-    
+
 
 def test_symeig_custom_low_rank():
     X = torch.randn(1000, 10)
@@ -35,7 +38,7 @@ def test_symeig_custom_low_rank():
     # test clamping to zeros
     assert torch.all(l2 >= 0.0)
 
-    
+
 def test_diagonal_normal_samples():
     mean = torch.randn(10, 2)
     var = torch.exp(torch.randn(10, 2))
@@ -48,7 +51,7 @@ def test_diagonal_normal_samples():
     same_samples = normal_samples(mean, var, n_samples=100, generator=generator)
     assert torch.allclose(samples, same_samples)
 
-    
+
 def test_multivariate_normal_samples():
     mean = torch.randn(10, 2)
     rndns = torch.randn(10, 2, 10) / 100
@@ -61,3 +64,28 @@ def test_multivariate_normal_samples():
     generator.set_state(gen_state)
     same_samples = normal_samples(mean, var, n_samples=100, generator=generator)
     assert torch.allclose(samples, same_samples)
+
+
+def test_validate():
+    X = torch.randn(50, 10)
+    y = torch.randint(3, size=(50,))
+    dataloader = DataLoader(TensorDataset(X, y), batch_size=10)
+
+    model = torch.nn.Sequential(torch.nn.Linear(10, 20), torch.nn.ReLU(), torch.nn.Linear(20, 3))
+    la = Laplace(model, 'classification', 'all')
+    la.fit(dataloader)
+
+    res = validate(
+        la, dataloader, get_nll, pred_type='nn', link_approx='mc', n_samples=10
+    )
+    assert res != math.nan
+    assert isinstance(res, float)
+    assert res > 0
+
+    res = validate(
+        la, dataloader, RunningNLLMetric(), pred_type='nn', link_approx='mc', n_samples=10
+    )
+    assert res != math.nan
+    assert isinstance(res, float)
+    assert res > 0
+
diff --git a/tests/utils.py b/tests/utils.py
index ccaef943..c9e136bf 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,13 +1,23 @@
+from collections import UserDict
+from collections.abc import MutableMapping
+from typing import Any, List
 import torch
+from torch.utils.data import Dataset
 
 
 def get_psd_matrix(dim):
-    X = torch.randn(dim, dim*3)
+    X = torch.randn(dim, dim * 3)
     return X @ X.T / (dim * 3)
 
+
+def get_diag_psd_matrix(dim):
+    return torch.randn(dim) ** 2
+
+
 def grad(model):
     return torch.cat([p.grad.data.flatten() for p in model.parameters()]).detach()
 
+
 def jacobians_naive(model, data):
     model.zero_grad()
     f = model(data)
@@ -16,17 +26,38 @@ def jacobians_naive(model, data):
         if len(f.shape) > 1:
             jacs = list()
             for j in range(f.shape[1]):
-                rg = (i != (f.shape[0] - 1) or j != (f.shape[1] - 1))
+                rg = i != (f.shape[0] - 1) or j != (f.shape[1] - 1)
                 f[i, j].backward(retain_graph=rg)
                 Jij = grad(model)
                 jacs.append(Jij)
                 model.zero_grad()
             jacs = torch.stack(jacs).t()
         else:
-            rg = (i != (f.shape[0] - 1))
+            rg = i != (f.shape[0] - 1)
             f[i].backward(retain_graph=rg)
             jacs = grad(model)
             model.zero_grad()
         Jacs.append(jacs)
     Jacs = torch.stack(Jacs).transpose(1, 2)
     return Jacs.detach(), f.detach()
+
+
+class ListDataset(Dataset):
+    def __init__(self, data: List[Any]):
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]
+
+
+def dict_data_collator(batch: List[MutableMapping]) -> UserDict:
+    ret = UserDict({})
+
+    for k in batch[0].keys():
+        vals = torch.stack([v[k] for v in batch]).squeeze(-1)
+        ret[k] = vals
+
+    return ret