optax.second_order.hessian_diag

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

Computes the diagonal hessian of loss at (inputs, targets).

Parameters:

• loss – the loss function.

• params – model parameters.

• inputs – inputs at which loss is evaluated.

• targets – targets at which loss is evaluated.

Returns:

A DeviceArray corresponding to the product to the Hessian of loss evaluated at (params, inputs, targets).