Module laplace.curvature.asdfghjkl

Classes

class AsdfghjklInterface (model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')

Interface for asdfghjkl backend.

Ancestors

Subclasses

Methods

def jacobians(self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], enable_backprop: bool = False) ‑> tuple[torch.Tensor, torch.Tensor]

Compute Jacobians \nabla_\theta f(x;\theta) at current parameter \theta using asdfghjkl's gradient per output dimension.

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)
def gradients(self, x: torch.Tensor | MutableMapping[str, torch.Tensor | Any], y: torch.Tensor) ‑> tuple[torch.Tensor, torch.Tensor]

Compute gradients \nabla_\theta \ell(f(x;\theta, y) at current parameter \theta using asdfghjkl's backend.

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)

Inherited members

class AsdfghjklHessian (model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels', low_rank: int = 10)

Interface for asdfghjkl backend.

Ancestors

Methods

def eig_lowrank(self, data_loader: DataLoader) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Inherited members

class AsdfghjklGGN (model: nn.Module, likelihood: Likelihood | str, last_layer: bool = False, subnetwork_indices: torch.LongTensor | None = None, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels', stochastic: bool = False)

Implementation of the GGNInterface using asdfghjkl.

Ancestors

Inherited members

class AsdfghjklEF (model: nn.Module, likelihood: Likelihood | None, last_layer: bool = False, dict_key_x: str = 'input_ids', dict_key_y: str = 'labels')

Implementation of the EFInterface using asdfghjkl.

Ancestors

Inherited members