Enzyme Hessian bug

[gdalle]

See EnzymeAD/Enzyme.jl#1628

[gdalle]

@dleather the issue you noticed is fixed on the latest version of everything, presumably thanks to the introduction of Constant arguments in DI.

Setup code

import DifferentiationInterface as AD
using Enzyme: Enzyme
using ForwardDiff: ForwardDiff

function compute_EI1uv_aug(z, t0, T, θz, θy)
    wu = compute_wu(t0, T, θz, θy)
    wv = compute_wv(t0, T, θz, θy)

    return wu * z[2] + wv * z[1] + 2 * z[1]^2
end

function compute_wu(t0, T, θz, θy)
    if θz ≠ θy
        numerator =
            -θy - θz * coth((T - t0) * θz) + exp((T - t0) * θy) * θz * csch((T - t0) * θz)
        denominator = (θy - θz) * (θy + θz)
        return numerator / denominator
    else
        numerator =
            exp((T + t0) * θz) *
            (exp(2 * t0 * θz) + exp(2 * T * θz) * (-1 + 2 * T * θz - 2 * t0 * θz)) *
            csch((T - t0) * θz)
        denominator = 4 * θz
        return numerator / denominator
    end
end

function compute_wv(t0, T, θz, θy)
    if θz ≠ θy
        numerator =
            exp((T - t0) * θy) * θy - exp((T - t0) * θy) * θz * coth((T - t0) * θz) +
            θz * csch((T - t0) * θz)
        denominator = (θy - θz) * (θy + θz)
        return numerator / denominator
    else
        numerator =
            exp(-(T + t0) * θz) *
            (exp((3 * T - t0) * θz) - exp((T + t0) * θz) * (1 + 2 * T * θz - 2 * t0 * θz)) *
            csch((T - t0) * θz)
        denominator = 4 * θz
        return numerator / denominator
    end
end

t0 = 0.0
T = 0.25
θz = 0.5
θy = 2.0
u = 0.1
v = 0.4
z = [v, u, v^2, u * v, u * v, u^2]

julia> compute_EI1uv_aug(z, t0, T, θz, θy)
0.4050218388870948

julia> AD.gradient(
           compute_EI1uv_aug,
           AD.AutoEnzyme(),
           z,
           AD.Constant(t0),
           AD.Constant(T),
           AD.Constant(θz),
           AD.Constant(θy),
       )
6-element Vector{Float64}:
 1.7754258975969854
 0.14851479848300594
 0.0
 0.0
 0.0
 0.0

julia> enzyme_hess = AD.hessian(
           compute_EI1uv_aug,
           AD.AutoEnzyme(),
           z,
           AD.Constant(t0),
           AD.Constant(T),
           AD.Constant(θz),
           AD.Constant(θy),
       )
6×6 Matrix{Float64}:
 4.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0

julia> R(z) = compute_EI1uv_aug(z, t0, T, θz, θy)
R (generic function with 1 method)

julia> ForwardDiff.hessian(R, z)
6×6 Matrix{Float64}:
 4.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0  0.0  0.0