descent.utils.loss

Utilities for defining loss functions.

Functions

approximate_hessian(x, y_pred)

Compute the outer product approximation of the hessian of a least squares loss function of the sum sum((y_pred - y_ref)**2).

combine_closures(closures[, weights, verbose])

Combine multiple closures into a single closure.

to_closure(loss_fn, *args, **kwargs)

Convert a loss function to a closure function used by second-order optimizers.

descent.utils.loss.to_closure(loss_fn: ~typing.Callable[[~typing.Concatenate[~torch.Tensor, ~P]], ~torch.Tensor], *args: ~typing.~P, **kwargs: ~typing.~P) Callable[[Tensor, bool, bool], tuple[Tensor, Tensor | None, Tensor | None]][source]

Convert a loss function to a closure function used by second-order optimizers.

Args:
loss_fn: The loss function to convert. This should take in a tensor of

parameters with shape=(n,), and optionally a set of args and kwargs.

*args: Positional arguments passed to loss_fn. **kwargs: Keyword arguments passed to loss_fn.

Returns:

A closure function that takes in a tensor of parameters with shape=(n,), a boolean flag indicating whether to compute the gradient, and a boolean flag indicating whether to compute the Hessian. It returns a tuple of the loss value, the gradient, and the Hessian.

descent.utils.loss.combine_closures(closures: dict[str, Callable[[Tensor, bool, bool], tuple[Tensor, Tensor | None, Tensor | None]]], weights: dict[str, float] | None = None, verbose: bool = False) Callable[[Tensor, bool, bool], tuple[Tensor, Tensor | None, Tensor | None]][source]

Combine multiple closures into a single closure.

Args:

closures: A dictionary of closure functions. weights: Optional dictionary of weights for each closure function. verbose: Whether to log the loss of each closure function.

Returns:

A combined closure function.

descent.utils.loss.approximate_hessian(x: Tensor, y_pred: Tensor)[source]

Compute the outer product approximation of the hessian of a least squares loss function of the sum sum((y_pred - y_ref)**2).

Args:

x: The parameter tensor with shape=(n_parameters,). y_pred: The values predicted using x with shape=(n_predications,).

Returns:

The outer product approximation of the hessian with ``shape=n_parameters