Core

Core#

Core second-order operations module.

torch_secorder.core.approximate_hvp(func: Callable[[], Tensor], params: List[Tensor], v: Tensor | List[Tensor], num_samples: int = 1, damping: float = 0.0) → Tensor | List[Tensor][source]#

Compute an approximate Hessian-vector product using finite differences.

This method uses a finite difference approximation of the Hessian-vector product, which can be more memory efficient than the exact computation.

Parameters:

func – A callable that returns a scalar tensor (loss).
params – List of parameters with respect to which to compute the Hessian.
v – Vector to multiply with the Hessian. Can be a single tensor or a list of tensors matching the structure of params.
num_samples – Number of samples to use for the approximation.
damping – Damping term to add to the diagonal of the Hessian (lambda * I).

Returns:

The approximate Hessian-vector product Hv. Returns a single tensor if v is a single tensor, otherwise returns a list of tensors matching the structure of params.

torch_secorder.core.batch_jvp(func: Callable[[], Tensor], params: List[Tensor], vs: Tensor | List[Tensor], create_graph: bool = False) → Tensor[source]#

Compute a batch of Jacobian-vector products (JVPs).

Parameters:

func – A callable that returns a tensor output (can be vector-valued).
params – List of parameters with respect to which to compute the Jacobian.
vs – Batch of vectors to multiply with the Jacobian. Should be a tensor of shape (batch, …) or a list of such tensors.
create_graph – If True, graph of the derivative will be constructed.

Returns:

Tensor of shape (batch, …) with the JVPs for each vector in the batch.

torch_secorder.core.batch_vjp(func: Callable[[], Tensor], params: List[Tensor], vs: Tensor, create_graph: bool = False) → List[Tensor][source]#

Compute a batch of vector-Jacobian products (VJPs).

Parameters:

func – A callable that returns a tensor output (can be vector-valued).
params – List of parameters with respect to which to compute the Jacobian.
vs – Batch of vectors to multiply with the Jacobian (should match the output shape of func, with batch dimension first).
create_graph – If True, graph of the derivative will be constructed.

Returns:

List of tensors, each of shape (batch, …) matching the structure of params.

torch_secorder.core.exact_hvp(func: Callable[[], Tensor], params: List[Tensor], v: Tensor | List[Tensor], create_graph: bool = False) → Tensor | List[Tensor][source]#

Compute the exact Hessian-vector product Hv using double backpropagation.

Parameters:

func – A callable that returns a scalar tensor (loss).
params – List of parameters with respect to which to compute the Hessian.
v – Vector to multiply with the Hessian. Can be a single tensor or a list of tensors matching the structure of params.
create_graph – If True, the graph used to compute the grad will be constructed, allowing to compute higher order derivative products.

Returns:

The Hessian-vector product Hv. Returns a single tensor if v is a single tensor, otherwise returns a list of tensors matching the structure of params.

torch_secorder.core.flatten_params(params: Iterable[Tensor]) → Tensor[source]#

Flattens a list of parameter tensors into a single concatenated tensor.

Parameters:: params – An iterable of PyTorch parameter tensors.
Returns:: A single 1D tensor containing all flattened parameters.

torch_secorder.core.full_jacobian(func: Callable[[], Tensor], params: List[Tensor], create_graph: bool = False) → List[Tensor][source]#

Compute the full Jacobian matrix of func with respect to params.

Parameters:

func – A callable that returns a tensor output (can be vector-valued).
params – List of parameters with respect to which to compute the Jacobian.
create_graph – If True, graph of the derivative will be constructed.

Returns:

List of Jacobian tensors, one for each parameter, with shape (output_dim, param_shape).

torch_secorder.core.gauss_newton_product(model: Module, loss_fn: Callable[[Tensor, Tensor], Tensor], x: Tensor, y: Tensor, v: Tensor | List[Tensor], create_graph: bool = False) → Tensor | List[Tensor][source]#

Compute the Gauss-Newton matrix-vector product for a model’s loss.

Parameters:

model – The PyTorch model.
loss_fn – The loss function (should be MSE or cross-entropy for GN to be valid).
x – Input tensor.
y – Target tensor.
v – Vector to multiply with the Gauss-Newton matrix.
create_graph – If True, graph of the derivative will be constructed.

Returns:

The Gauss-Newton matrix-vector product (same structure as v).

torch_secorder.core.get_layer_curvature_stats(layer_curvatures: Dict[str, Tensor]) → Dict[str, Dict[str, float]][source]#

Compute basic statistics for each layer’s curvature information.

Parameters:: layer_curvatures – Dictionary mapping layer names to their curvature tensors (Hessian or Fisher diagonals).
Returns:: Dictionary mapping layer names to their curvature statistics (mean, std, max, min).

torch_secorder.core.get_param_shapes(params: Iterable[Tensor]) → List[Size][source]#

Retrieves the shapes of an iterable of parameter tensors.

Parameters:: params – An iterable of PyTorch parameter tensors.
Returns:: A list of torch.Size objects, each representing the shape of a parameter.

torch_secorder.core.get_params_by_module_type(model: Module, module_type: type | List[type] | Tuple[type, ...]) → Dict[str, List[Tensor]][source]#

Extracts parameters belonging to specific module types.

Parameters:

model – The PyTorch model to inspect.
module_type – The type(s) of torch.nn.Module to filter parameters by. Can be a single type (e.g., torch.nn.Linear) or a list/tuple of types.

Returns:

Dictionary mapping module names to lists of their parameter tensors.

torch_secorder.core.get_params_by_name_pattern(model: Module, pattern: str) → List[Tensor][source]#

Extracts parameters whose names match a given pattern.

This is useful for selecting parameters based on their hierarchical names within the model (e.g., “layer1.0.weight”).

Parameters:

model – The PyTorch model to inspect.
pattern – A regex pattern string to match against parameter names.

Returns:

A list of parameter tensors whose names match the pattern.

torch_secorder.core.hessian_diagonal(func: Callable[[], Tensor], params: List[Tensor], v: List[Tensor] | None = None, create_graph: bool = False, strict: bool = False) → List[Tensor][source]#

Compute the diagonal elements of the Hessian matrix.

This function computes the diagonal elements of the Hessian matrix by using double backward for each parameter element.

Parameters:

func – A callable that returns a scalar tensor (the loss).
params – List of parameters with respect to which the Hessian is computed.
v – Optional list of vectors to use for computing the diagonal. If None, computes the true diagonal. If provided, computes v_i * H_ii for each i.
create_graph – If True, the computational graph will be constructed, allowing for higher-order derivatives.
strict – If True, an error will be raised if any parameter requires grad but has no gradient.

Returns:

List of tensors containing the diagonal elements of the Hessian for each parameter.

torch_secorder.core.hessian_trace(func: Callable[[], Tensor], params: List[Tensor], num_samples: int = 1, create_graph: bool = False, strict: bool = False) → Tensor[source]#

Compute the trace of the Hessian matrix by summing the diagonal elements.

Parameters:

func – A callable that returns a scalar tensor (the loss).
params – List of parameters with respect to which the Hessian is computed.
num_samples – Ignored (kept for API compatibility).
create_graph – If True, the computational graph will be constructed, allowing for higher-order derivatives.
strict – If True, an error will be raised if any parameter requires grad but has no gradient.

Returns:

A tensor containing the trace of the Hessian.

torch_secorder.core.jvp(func: Callable[[], Tensor], params: List[Tensor], v: Tensor | List[Tensor], create_graph: bool = False) → Tensor | List[Tensor][source]#

Compute the Jacobian-vector product (JVP): J v.

Parameters:

func – A callable that returns a tensor output (can be vector-valued).
params – List of parameters with respect to which to compute the Jacobian.
v – Vector to multiply with the Jacobian. Can be a single tensor or a list of tensors matching the structure of params.
create_graph – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products.

Returns:

The JVP (same shape as the output of func).

torch_secorder.core.model_hessian_diagonal(model: Module, loss_fn: Callable[[Tensor, Tensor], Tensor], inputs: Tensor, targets: Tensor, create_graph: bool = False, strict: bool = False) → List[Tensor][source]#

torch_secorder.core.model_hessian_trace(model: Module, loss_fn: Callable[[Tensor, Tensor], Tensor], inputs: Tensor, targets: Tensor, num_samples: int = 1, create_graph: bool = False, strict: bool = False) → Tensor[source]#

Compute the trace of the Hessian for a model’s loss function.

A convenience function to estimate the trace of the Hessian of a model’s loss with respect to its parameters.

Parameters:

model – The PyTorch model.
loss_fn – The loss function, e.g., nn.MSELoss() or F.cross_entropy.
inputs – Input tensor to the model.
targets – Target tensor for the loss function.
num_samples – Number of random vectors for Hutchinson’s method (if used). Ignored for diagonal-based trace.
create_graph – If True, the computational graph will be constructed, allowing for higher-order derivatives.
strict – If True, an error will be raised if any parameter requires grad but has no gradient.

Returns:

A tensor containing the estimated trace of the Hessian.

torch_secorder.core.model_hvp(model: Module, loss_fn: Callable[[Tensor, Tensor], Tensor], x: Tensor, y: Tensor, v: Tensor | List[Tensor], create_graph: bool = False) → Tensor | List[Tensor][source]#

Compute the Hessian-vector product for a model’s loss function.

This is a convenience wrapper around exact_hvp that handles the model’s forward pass and loss computation.

Parameters:

model – The PyTorch model.
loss_fn – The loss function that takes model output and target as arguments.
x – Input tensor.
y – Target tensor.
v – Vector to multiply with the Hessian.
create_graph – If True, the graph used to compute the grad will be constructed.

Returns:

The Hessian-vector product Hv.

torch_secorder.core.model_jvp(model: Module, x: Tensor, v: Tensor | List[Tensor], create_graph: bool = False) → Tensor[source]#

Compute the JVP for a model’s output with respect to its parameters.

Parameters:

model – The PyTorch model.
x – Input tensor.
v – Vector to multiply with the Jacobian (should match the structure of model.parameters()).
create_graph – If True, graph of the derivative will be constructed.

Returns:

The JVP (same shape as the model output).

torch_secorder.core.model_vjp(model: Module, x: Tensor, v: Tensor, create_graph: bool = False) → Tensor | List[Tensor][source]#

Compute the VJP for a model’s output with respect to its parameters.

Parameters:

model – The PyTorch model.
x – Input tensor.
v – Vector to multiply with the Jacobian (should match the output shape of model(x)).
create_graph – If True, graph of the derivative will be constructed.

Returns:

The VJP (list of tensors matching the structure of model.parameters()).

torch_secorder.core.per_layer_fisher_diagonal(model: Module, loss_fn: Module, inputs: Tensor, targets: Tensor, layer_types: List[type] | None = None, create_graph: bool = False) → Dict[str, Tensor][source]#

Compute the diagonal of the Fisher Information matrix for each layer separately.

This function computes the diagonal elements of the Fisher Information matrix for each layer in the model independently, treating other layers’ parameters as fixed. This provides a block-diagonal approximation of the full Fisher matrix.

Parameters:

model – The neural network model.
loss_fn – The loss function used for training.
inputs – Input tensor for the model.
targets – Target tensor for the loss function.
layer_types – List of layer types to include. If None, includes all layers.
create_graph – Whether to create the computational graph for higher-order derivatives.

Returns:

Dictionary mapping layer names to their Fisher diagonal tensors.

torch_secorder.core.per_layer_hessian_diagonal(model: Module, loss_fn: Module, inputs: Tensor, targets: Tensor, layer_types: List[type] | None = None, create_graph: bool = False) → Dict[str, Tensor][source]#

Compute the diagonal of the Hessian matrix for each layer separately.

This function computes the diagonal elements of the Hessian matrix for each layer in the model independently, treating other layers’ parameters as fixed. This provides a block-diagonal approximation of the full Hessian.

Parameters:

model – The neural network model.
loss_fn – The loss function used for training.
inputs – Input tensor for the model.
targets – Target tensor for the loss function.
layer_types – List of layer types to include. If None, includes all layers.
create_graph – Whether to create the computational graph for higher-order derivatives.

Returns:

Dictionary mapping layer names to their Hessian diagonal tensors.

torch_secorder.core.unflatten_params(flat_params: Tensor, param_shapes: List[Size]) → List[Tensor][source]#

Unflattens a single tensor of parameters back into a list of tensors with original shapes.

Parameters:

flat_params – A 1D tensor containing all flattened parameters.
param_shapes – A list of torch.Size objects, representing the original shapes of the parameters.

Returns:

A list of PyTorch tensors with their original shapes.

torch_secorder.core.vjp(func: Callable[[], Tensor], params: List[Tensor], v: Tensor, create_graph: bool = False) → Tensor | List[Tensor][source]#

Compute the vector-Jacobian product (VJP): v^T J.

Parameters:

func – A callable that returns a tensor output (can be vector-valued).
params – List of parameters with respect to which to compute the Jacobian.
v – Vector to multiply with the Jacobian (should match the output shape of func).
create_graph – If True, graph of the derivative will be constructed, allowing to compute higher order derivative products.

Returns:

The VJP (list of tensors matching the structure of params).

Core

Contents

Core#