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).