Gradient of a Hessian dependent loss function

Specific Domains Machine Learning
1

For my application it would be useful to have a loss function that depends on the Hessian of a model.
Computing the Hessian with respect to the input variable (using Tracker) works fine, but trying to get the gradient with respect to the parameters doesn’t work if matrix multiplication (such as for Dense layers) is involved.

Consider the following example:

W = param(rand(1, 2))
fn = z -> sum(W * z)^2
hess = z -> Tracker.hessian(fn, z)
Tracker.gradient(() -> sum(hess([1.0, 2.0])), Flux.Params([W]))

This throws the error:
MethodError: Cannot `convert` an object of type Base.ReshapedArray{Float64,2,Transpose{Float64,Array{Float64,2}},Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}} to an object of type Transpose{Float64,Array{Float64,2}}

Any thoughts or advice appreciated.

I suspect this might be a bug somewhere, but it’s hard to track down where. It would already help to know which package is responsible, to submit an issue.

2

Unfortunately I don’t have anything to add here except that I’m also interested in being able to get gradients of a Hessian-dependent function. In this case it might be worth filing an issue with Tracker or trying Zygote which is intended to replace Tracker.