optax.second_order.hvp

optax.second_order.hvp(loss: LossFn, v: Array, params: Any, inputs: Array, targets: Array) → Array

Performs an efficient vector-Hessian (of loss) product.

Parameters:

• loss – the loss function.

• v – a vector of size ravel(params).

• params – model parameters.

• inputs – inputs at which loss is evaluated.

• targets – targets at which loss is evaluated.

Returns:

An Array corresponding to the product of v and the Hessian of loss evaluated at (params, inputs, targets).