Module laplace.curvature.backpack

Classes

class BackPackInterface (model, likelihood, last_layer=False)

Interface for Backpack backend.

Ancestors

Subclasses

Static methods

def jacobians(model, x)

Compute Jacobians \nabla_{\theta} f(x;\theta) at current parameter \theta using backpack's BatchGrad per output dimension.

Parameters

model : torch.nn.Module
 
x : torch.Tensor
input data (batch, input_shape) on compatible device with model.

Returns

Js : torch.Tensor
Jacobians (batch, parameters, outputs)
f : torch.Tensor
output function (batch, outputs)

Methods

def gradients(self, x, y)

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

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 BackPackGGN (model, likelihood, last_layer=False, stochastic=False)

Implementation of the GGNInterface using Backpack.

Ancestors

Inherited members

class BackPackEF (model, likelihood, last_layer=False)

Implementation of EFInterface using Backpack.

Ancestors

Inherited members