optax.second_order.hvp

Contents

optax.second_order.hvp#

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

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

Contents