Module laplace.curvature.asdl
Classes
class AsdlInterface (model, likelihood, last_layer=False, subnetwork_indices=None)
-
Interface for asdfghjkl backend.
Ancestors
Subclasses
Instance variables
var loss_type
Methods
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
orUserDict
- 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
Returns
Js
:torch.Tensor
- Jacobians
(batch, parameters, outputs)
f
:torch.Tensor
- output function
(batch, outputs)
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)
Inherited members
class AsdlHessian (model, likelihood, last_layer=False, low_rank=10)
-
Interface for asdfghjkl backend.
Ancestors
Inherited members
class AsdlGGN (model, likelihood, last_layer=False, subnetwork_indices=None, stochastic=False)
-
Implementation of the
GGNInterface
using asdfghjkl.Ancestors
Inherited members
class AsdlEF (model, likelihood, last_layer=False)
-
Implementation of the
EFInterface
using asdfghjkl.Ancestors
Inherited members