Computing hessian of unnamed functions throws error

[Vaibhavdixit02]

I am trying to implement a function to compute hessians of an array of functions (it's actually a single function with vector output but functions are created to output individual indexes to allow the hessian calculation)

I tried 2 approaches these are the MWE for them

fncs = [x -> x[1]^2 + x[2]^2, x -> x[2] * sin(x[1]) - x[1]]
cons_h = function (res, θ)
    for i in 1:2
        res[i] .= Enzyme.jacobian(Forward, (x) -> Enzyme.gradient(Reverse, fncs[i], x), θ)
    end
end
x0 = zeros(2)
H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
cons_h(H3, x0)

and

cons_h = function (res, θ)
    for i in 1:2
        y = zeros(1)

        vdy = Tuple(zeros(1) for i in eachindex(θ))
        vdθ = Tuple((Array(r) for r in eachrow(I(length(θ)) * 1.0)))

        bθ = zeros(length(θ))
        by = [1.0]
        vdbθ = Tuple(zeros(length(θ)) for i in eachindex(θ))
        vdby = Tuple(zeros(1) for i in eachindex(θ))

        Enzyme.autodiff(Forward,
                        (θ, y) -> Enzyme.autodiff((x, y) -> (y .= fncs[i](x); return nothing),
                                                  θ, y),
                        BatchDuplicated(Duplicated(θ, bθ), Duplicated.(vdθ, vdbθ)),
                        BatchDuplicated(Duplicated(y, by), Duplicated.(vdy, vdby)))

    end
end
H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
cons_h(H3, x0)

Both give the below error which makes me suspect that the way I am defining the individual function might be the issue

ERROR: Attempting to call an indirect active function whose runtime value is inactive:
Backtrace

Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/Enzyme/DIkTv/src/compiler.jl:6819
 [2] enzyme_call
   @ ~/.julia/packages/Enzyme/DIkTv/src/compiler.jl:6544
 [3] CombinedAdjointThunk
   @ ~/.julia/packages/Enzyme/DIkTv/src/compiler.jl:6507
 [4] autodiff
   @ ~/.julia/packages/Enzyme/DIkTv/src/Enzyme.jl:217
 [5] autodiff
   @ ~/.julia/packages/Enzyme/DIkTv/src/Enzyme.jl:248
 [6] autodiff
   @ ~/.julia/packages/Enzyme/DIkTv/src/Enzyme.jl:354
 [7] #18
   @ ~/Optimization.jl/test/ADtests.jl:92
 [8] #18
   @ ~/Optimization.jl/test/ADtests.jl:0

The library code this is intended for involves a bit more of indirection (shown below) and throws a different stacktrace (very large)

function con2_c(res, x, p)
    res .= [x[1]^2 + x[2]^2, x[2] * sin(x[1]) - x[1]]
end
cons = (res, θ) -> (con2_c(res, θ, nothing); return nothing)
cons_oop = (x) -> (_res = zeros(eltype(x), num_cons); cons(_res, x); _res)
fncs = [(x) -> cons_oop(x)[i] for i in 1:2]
cons_h = function (res, θ)
    for i in 1:2
        res[i] .= Enzyme.jacobian(Forward, (x) -> Enzyme.gradient(Reverse, fncs[i], x), θ)
    end
end
x0 = zeros(2)
H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
cons_h(H3, x0)

error

error.txt

[wsmoses]

@Vaibhavdixit02 you must use autodiff_deferred on the inner functions (thus the gradient syntactic sugar wrapper should not be used).

[Vaibhavdixit02]

I used that actually, I had changed it to experiment. It gives the same error as in the file above

[wsmoses]

julia> function f2(fnc, x)
           dx = zeros(length(x))
           Enzyme.autodiff_deferred(Reverse, fnc, Duplicated(x, dx))
           dx
       end
f2 (generic function with 1 method)

julia> Enzyme.autodiff(Forward, f2, fncs[1], Duplicated(x0, similar(x0)))
(var"1" = [1.38758355690365e-309, 1.38758355690365e-309],)

julia> Enzyme.jacobian(Forward, (x) -> f2(fncs[1], x), x0)
2×2 Matrix{Float64}:
 2.0  0.0
 0.0  2.0

julia> Enzyme.jacobian(Forward, (x) -> f2(fncs[2], x), x0)
2×2 Matrix{Float64}:
 0.0  1.0
 1.0  0.0

[Vaibhavdixit02]

fncs = [(x) -> x[1]^2 + x[2]^2, (x) -> x[2] * sin(x[1]) - x[1]]
cons_h = function (res, θ)
    for i in 1:2
        function f2(fnc, x)
           dx = zeros(length(x))
           Enzyme.autodiff_deferred(Reverse, fnc, Duplicated(x, dx))
           dx
        end
        Enzyme.jacobian(Forward, (x) -> f2(fncs[i], x), θ)
    end
end
H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
cons_h(H3, x0)

gives me

ERROR: Return type inferred to be Union{}. Giving up.
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] #s883#169
   @ ~/.julia/packages/Enzyme/DIkTv/src/compiler.jl:6943 [inlined]
 [3] var"#s883#169"(F::Any, Fn::Any, DF::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, specid::Any, ReturnPrimal::Any, ShadowInit::Any, ::Any, #unused#::Type, f::Any, df::Any, #unused#::Type, tt::Any, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Any)
   @ Enzyme.Compiler ./none:0
 [4] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
   @ Core ./boot.jl:582
 [5] thunk
   @ ~/.julia/packages/Enzyme/DIkTv/src/compiler.jl:7001 [inlined]
 [6] thunk(f::var"#f2#43", df::Nothing, ::Type{Const{Union{}}}, tt::Type{Tuple{Const{var"#36#38"}, BatchDuplicated{Vector{Float64}, 2}}}, ::Val{Enzyme.API.DEM_ForwardMode}, ::Val{2}, ::Val{false}, ::Val{true})
   @ Enzyme.Compiler ~/.julia/packages/Enzyme/DIkTv/src/compiler.jl:6994

[wsmoses]

What version of julia are you using? Also regardless, you should also use autodiff instead of jacobian so you can pass in multiple parameters rather than forcing a closure.

[Vaibhavdixit02]

I am on 1.8.5. I was able to remove the above Union{} error by not specifying the mode in the autodiff_deferred call.

Below is a reduced MWE of two cases, one that passes

x0 = zeros(2)
fncs = [x -> x[1]^2 + x[2]^2, x-> x[2] * sin(x[1]) - x[1]]
function f2(fnc, x)
    dx = zeros(length(x))
    Enzyme.autodiff_deferred(fnc, Duplicated(x, dx))
    dx
end
Enzyme.jacobian(Forward, x -> f2(fncs[1], x), x0)

And this one causes the attached error

function con2_c(x)
    [x[1]^2 + x[2]^2, x[2] * sin(x[1]) - x[1]]
end
fncs = [x -> con2_c(x)[i] for i in 1:2]
Enzyme.jacobian(Forward, x -> f2(fncs[1], x), x0)

error.txt

[Vaibhavdixit02]

julia> Enzyme.autodiff(Forward, f2, fncs[1], Duplicated(x0, similar(x0))) (var"1" = [1.38758355690365e-309, 1.38758355690365e-309],)

In your comment doesn't this look incorrect?

[Vaibhavdixit02]

Also regardless, you should also use autodiff instead of jacobian so you can pass in multiple parameters rather than forcing a closure.

I have this working for the objective hessian, both cases work for that actually. But in the case with the array of anonymous functions generated in the comprehension, I have been unable to get it to work using either of the approaches.

[wsmoses]

In your comment doesn't this look incorrect?

No that's to be expected. That was jsut testing the function ran, since I passed in uninitialized memory via similar(x0) instead of a particular shadow.

[wsmoses]

Can you use the latest main? That might be the fix and potential source of your bug. On main we changed the ABI to take autodiff_deferred(Mode, ...) rather than the previous one which did not. Moreover, main has many many fixes.

[Vaibhavdixit02]

I had been intermittently trying with master. It does work now! Awesome! 🎉