Calculating hessian of a NN w.r.t params

New to Julia
,
1

How would you calculate a hessian of a Neural Network w.r.t. it’s parameters?

For instance, a hessian of the loss function below

using Flux: Chain, Dense, σ, crossentropy, params
using Zygote
model = Chain(
    x -> reshape(x, :, size(x, 4)),
    Dense(2, 5),
    Dense(5, 1),
    x -> σ.(x)
)
n_data = 5
input = randn(2, 1, 1, n_data)
target = randn(1, n_data)
loss = model -> Flux.crossentropy(model(input), target)

I can get a gradient w.r.t parameters in two ways…

Zygote.gradient(model -> loss(model), model)

or

grad = Zygote.gradient(() -> loss(model), params(model))
grad[params(model)[1]]

However, I can’t find a way to get a hessian w.r.t its parameters. (I want to do something like Zygote.hessian(model -> loss(model), model), but I can’t)