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 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

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