Autograd Utility: Hessian

Autograd Utility: Hessian#

Utilities for calculating Hessians via autograd.

Note

Batched Hessians are not supported yet (via vmap).

tad_mctc.autograd.hessian.hess_fn_rev(f, argnums=0)[source]#

Return the Hessian function using reverse-mode autodiff twice. (Functorch’s hessian uses forward and backward mode, but forward is not implemented for our custom autograd functions.)

Parameters:

  • f (Callable[[Any], Tensor]) – The function whose result is differentiated.

  • argnums (int or tuple[int], optional) – The variable w.r.t. which will be differentiated. Defaults to 0.

Returns:

A function that computes the Hessian of f with respect to the specified argument(s).

Return type:

Callable

tad_mctc.autograd.hessian.hessian(f, inputs, argnums, is_batched=False, **kwargs)[source]#

Wrapper for Hessian. The Hessian is the Jacobian of the gradient.

PyTorch, however, suggests calculating the Jacobian of the Jacobian, which does not yield the correct shape in this case.

Parameters:

  • f (Callable[[Any], Tensor]) – The function whose result is differentiated.

  • inputs (tuple[Any, …]) – The input parameters of f.

  • argnums (int, optional) – The variable w.r.t. which will be differentiated. Defaults to 0.

Returns:

The Hessian.

Return type:

Tensor

Raises:

RuntimeError – The parameter selected for differentiation (via argnums) is not a tensor.

Contents