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