Can't differentiate loopinfo expression when trying to compute second order derivatives using Flux and Distributions

Specific Domains Machine Learning
, ,
1

I am trying to compute the maximum a posterior (MAP) estimate and want to get the hessian of the log likelihood too. This is what I have tried with respect to computing the Hessian:

using Flux, Distributions

prior_mean = Normal(178, 20)
prior_sd = Uniform(0, 50)

heights = rand(Normal(150, 2), 100)

ll_data(μ, σ) = begin
    d = Normal(μ, σ)
    pdf.(d, heights) .|> log |> sum
end

l_param(θ, d) = pdf(d, θ) |> log
l_priors(μ, σ) = l_param(μ, prior_mean) + l_param(σ, prior_sd)

objective_fn(μ, σ) = ll_data(μ, σ) + l_priors(μ, σ)

df(a, b) = gradient(objective_fn, a, b)[1]
d2f(a, b) = gradient(df, a, b)

# this works
println(df(153.5, 7.0))

# this fails
println(d2f(153.5, 7.0))

The last line throws a Can't differentiate loopinfo expression exception.

What does this mean and how to resolve it?

Upgrading to Julia 1.4.0-rc1 didn’t work either.

Using ForwardDiff works (on Julia 1.3.1):

using ForwardDiff

f(arr) = begin
    a, b = arr
    objective_fn(a, b)
end

g(x) = ForwardDiff.hessian(f, x)

g([153.5, 7.0])