`DI.hessian`: Tuple field type cannot be `Union{}`

[timholy]

This example arose from using GAMSFiles to parse a bunch of files used as a test suite for optimization problems (specifically convex3_10_1.gms). I've set this up so you don't need GAMSFiles to run this demo, although there are commented-out references to that package. (In my local checkout of that package I defined a @zero_adjoint for GAMSFiles.checkaxes, and that fixed one problem, woohoo. But this is different.)

Unless I'm missing something, the only operations in this objective are +, -, and *. This is with Mooncake v0.4.129 and DI v0.7.1.

using Mooncake
using DifferentiationInterface

module MooncakeDemo
    # using GAMSFiles
    function objective(x)
    #   begin
        #   (GAMSFiles.checkaxes)(x, (Base.OneTo(10),))
          objvar = (((((((((((((((((((((((((((x[1] * x[1] + x[10] * x[10]) * (x[1] * x[1] + x[10] * x[10]) - 4 * x[1]) + 3) + (x[2] * x[2] + x[10] * x[10]) * (x[2] * x[2] + x[10] * x[10])) - 4 * x[2]) + 3) + (x[3] * x[3] + x[10] * x[10]) * (x[3] * x[3] + x[10] * x[10])) - 4 * x[3]) + 3) + (x[4] * x[4] + x[10] * x[10]) * (x[4] * x[4] + x[10] * x[10])) - 4 * x[4]) + 3) + (x[5] * x[5] + x[10] * x[10]) * (x[5] * x[5] + x[10] * x[10])) - 4 * x[5]) + 3) + (x[6] * x[6] + x[10] * x[10]) * (x[6] * x[6] + x[10] * x[10])) - 4 * x[6]) + 3) + (x[7] * x[7] + x[10] * x[10]) * (x[7] * x[7] + x[10] * x[10])) - 4 * x[7]) + 3) + (x[8] * x[8] + x[10] * x[10]) * (x[8] * x[8] + x[10] * x[10])) - 4 * x[8]) + 3) + (x[9] * x[9] + x[10] * x[10]) * (x[9] * x[9] + x[10] * x[10])) - 4 * x[9]) + 3) - 0
          return objvar
    #   end
  end
end
x0 = [0.0; fill(1.0, 9)]
f = MooncakeDemo.objective
f(x0)  # works

backend = AutoMooncake(; config=nothing)
prepg = prepare_gradient(f, backend, x0)
∇f(x) = gradient(f, prepg, backend, x)
@show isa(∇f(x0), AbstractVector{<:Real})

preph = prepare_hessian(f, backend, x0)
∇²f(x) = hessian(f, preph, backend, x)
@show isa(∇²f(x0), AbstractMatrix{<:Real})

This fails with

julia> include("mooncakebug.jl")
WARNING: replacing module MooncakeDemo.
∇f(x0) isa AbstractVector{<:Real} = true
ERROR: LoadError: MooncakeRuleCompilationError: an error occured while Mooncake was compiling a rule to differentiate something. If the `caused by` error message below does not make it clear to you how the problem can be fixed, please open an issue at github.com/chalk-lab/Mooncake.jl describing your problem.
To replicate this error run the following:

Mooncake.build_rrule(Mooncake.MooncakeInterpreter(), Tuple{typeof(Mooncake.tuple_splat), typeof(Core.kwcall), Tuple{@NamedTuple{debug_mode::Bool, silence_debug_messages::Bool}, typeof(Mooncake.build_rrule), typeof(Main.MooncakeDemo.objective), Vector{Float64}}}; debug_mode=false)

Note that you may need to `using` some additional packages if not all of the names printed in the above signature are available currently in your environment.

Stacktrace:
  [1] build_rrule(interp::Mooncake.MooncakeInterpreter{Mooncake.DefaultCtx}, sig_or_mi::Type; debug_mode::Bool, silence_debug_messages::Bool)
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1136
  [2] build_rrule
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1077 [inlined]
  [3] (::Mooncake.DynamicDerivedRule{Dict{…}})(::Mooncake.CoDual{typeof(Mooncake.tuple_splat), Mooncake.NoFData}, ::Mooncake.CoDual{typeof(Core.kwcall), Mooncake.NoFData}, ::Mooncake.CoDual{Tuple{…}, Tuple{…}})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1736
  [4] (::Mooncake.RRuleZeroWrapper{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:302
  [5] _apply_iterate_equivalent
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\rrules\builtins.jl:508 [inlined]
  [6] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
    @ Base.Experimental .\<missing>:0
  [7] (::Mooncake.DerivedRule{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:966
  [8] (::Mooncake.DynamicDerivedRule{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1739
  [9] (::Mooncake.RRuleZeroWrapper{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:302
 [10] #prepare_gradient_cache#668
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interface.jl:483 [inlined]
 [11] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual)
    @ Base.Experimental .\<missing>:0
 [12] DerivedRule
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:966 [inlined]
 [13] _build_rule!(rule::Mooncake.LazyDerivedRule{Tuple{…}, Mooncake.DerivedRule{…}}, args::Tuple{Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1827
 [14] LazyDerivedRule
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1822 [inlined]
 [15] RRuleZeroWrapper
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:302 [inlined]
 [16] prepare_gradient_nokwarg
    @ C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\ext\DifferentiationInterfaceMooncakeExt\onearg.jl:112 [inlined]
 [17] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
    @ Base.Experimental .\<missing>:0
 [18] DerivedRule
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:966 [inlined]
 [19] _build_rule!(rule::Mooncake.LazyDerivedRule{Tuple{…}, Mooncake.DerivedRule{…}}, args::Tuple{Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1827
 [20] LazyDerivedRule
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1822 [inlined]
 [21] RRuleZeroWrapper
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:302 [inlined]
 [22] gradient
    @ C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\first_order\gradient.jl:62 [inlined]
 [23] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
    @ Base.Experimental .\<missing>:0
 [24] DerivedRule
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:966 [inlined]
 [25] _build_rule!(rule::Mooncake.LazyDerivedRule{Tuple{…}, Mooncake.DerivedRule{…}}, args::Tuple{Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}, Mooncake.CoDual{…}})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1827
 [26] LazyDerivedRule
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:1822 [inlined]
 [27] RRuleZeroWrapper
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:302 [inlined]
 [28] shuffled_gradient
    @ C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\first_order\gradient.jl:158 [inlined]
 [29] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
    @ Base.Experimental .\<missing>:0
 [30] (::Mooncake.DerivedRule{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interpreter\s2s_reverse_mode_ad.jl:966
 [31] prepare_pullback_cache(::Function, ::Vararg{Any}; kwargs::@Kwargs{debug_mode::Bool, silence_debug_messages::Bool})
    @ Mooncake C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interface.jl:414
 [32] prepare_pullback_cache
    @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\interface.jl:407 [inlined]
 [33] prepare_pullback_nokwarg(::Val{…}, ::typeof(DifferentiationInterface.shuffled_gradient), ::AutoMooncake{…}, ::Vector{…}, ::Tuple{…}, ::DifferentiationInterface.FunctionContext{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterfaceMooncakeExt C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\ext\DifferentiationInterfaceMooncakeExt\onearg.jl:14
 [34] _prepare_hvp_aux(::Val{…}, ::DifferentiationInterface.ReverseOverReverse, ::DifferentiationInterface.DontPrepareInner, ::typeof(Main.MooncakeDemo.objective), ::AutoMooncake{…}, ::Vector{…}, ::Tuple{…})
    @ DifferentiationInterface C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\second_order\hvp.jl:663
 [35] prepare_hvp_nokwarg
    @ C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\second_order\hvp.jl:125 [inlined]
 [36] _prepare_hessian_aux(::Val{true}, ::DifferentiationInterface.BatchSizeSettings{1, false, true}, ::typeof(Main.MooncakeDemo.objective), ::AutoMooncake{Nothing}, ::Vector{Float64})
    @ DifferentiationInterface C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\second_order\hessian.jl:122
 [37] prepare_hessian_nokwarg
    @ C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\second_order\hessian.jl:104 [inlined]
 [38] #prepare_hessian#78
    @ C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\second_order\hessian.jl:11 [inlined]
 [39] prepare_hessian(::typeof(Main.MooncakeDemo.objective), ::AutoMooncake{Nothing}, ::Vector{Float64})
    @ DifferentiationInterface C:\Users\timho\.julia\packages\DifferentiationInterface\alBlj\src\second_order\hessian.jl:8
 [40] top-level scope
    @ C:\Users\timho\OneDrive\Documents\publications\diquadraticoptim\src\testproblems\gams\mooncakebug.jl:23
 [41] include(fname::String)
    @ Base.MainInclude .\client.jl:494
 [42] top-level scope
    @ REPL[10]:1
in expression starting at C:\Users\timho\OneDrive\Documents\publications\diquadraticoptim\src\testproblems\gams\mooncakebug.jl:23

caused by: StackOverflowError:
Stacktrace:
      [1] macro expansion
        @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\tangents.jl:439 [inlined]
      [2] macro expansion
        @ .\none:0 [inlined]
      [3] tangent_type(::Type{Core.CodeInstance})
        @ Mooncake .\none:0
--- the last 3 lines are repeated 79982 more times ---
 [239950] macro expansion
        @ C:\Users\timho\.julia\packages\Mooncake\Bo7Z0\src\tangents.jl:439 [inlined]
 [239951] macro expansion
        @ .\none:0 [inlined]
Some type information was truncated. Use `show(err)` to see complete types.

I suspect the tangent_type(::Type{Core.CodeInstance}) is unexpected?

[willtebbutt]

Hi Tim! Thanks for opening this issue. We don't (currently) support higher-order AD, so what you're seeing here is just one of (presumably) various possible bugs that will come up if you try to apply Mooncake to code generated to perform AD using Mooncake.

@gdalle am I remembering correctly that there's not a clean way to get DI to provide an informative error message if someone tries to eg. compute a Hessian in the above manner?

[willtebbutt]

Looking at our docs, I can't see where we make this limitation as clear. We should probably add something to the Known Limitation section of the docs to make this clear to users.

[gdalle]

@gdalle am I remembering correctly that there's not a clean way to get DI to provide an informative error message if someone tries to eg. compute a Hessian in the above manner?

Actually there is, we could just throw an error on DI.hessian(f, ::Mooncake, x) and friends. Of course this wouldn't stop people from trying to nest derivatives manually, but it would catch basis uses like this one.

[gdalle]

However it is debatable whether we should do that or just merge #389 and get the Hessian actually working properly?

[willtebbutt]

Good point. Let's just get #389 merged (I'm going to spend some time on it this afternoon).

[yebai]

This now throws a different error with Mooncake@v0.4.150. A rule seems to be missing somewhere.

julia> using Mooncake

julia> import DifferentiationInterface as DI

julia> module MooncakeDemo
           # using GAMSFiles
           function objective(x)
           #   begin
               #   (GAMSFiles.checkaxes)(x, (Base.OneTo(10),))
                 objvar = (((((((((((((((((((((((((((x[1] * x[1] + x[10] * x[10]) * (x[1] * x[1] + x[10] * x[10]) - 4 * x[1]) + 3) + (x[2] * x[2] + x[10] * x[10]) * (x[2] * x[2] + x[10] * x[10])) - 4 * x[2]) + 3) + (x[3] * x[3] + x[10] * x[10]) * (x[3] * x[3] + x[10] * x[10])) - 4 * x[3]) + 3) + (x[4] * x[4] + x[10] * x[10]) * (x[4] * x[4] + x[10] * x[10])) - 4 * x[4]) + 3) + (x[5] * x[5] + x[10] * x[10]) * (x[5] * x[5] + x[10] * x[10])) - 4 * x[5]) + 3) + (x[6] * x[6] + x[10] * x[10]) * (x[6] * x[6] + x[10] * x[10])) - 4 * x[6]) + 3) + (x[7] * x[7] + x[10] * x[10]) * (x[7] * x[7] + x[10] * x[10])) - 4 * x[7]) + 3) + (x[8] * x[8] + x[10] * x[10]) * (x[8] * x[8] + x[10] * x[10])) - 4 * x[8]) + 3) + (x[9] * x[9] + x[10] * x[10]) * (x[9] * x[9] + x[10] * x[10])) - 4 * x[9]) + 3) - 0
                 return objvar
           #   end
         end
       end
Main.MooncakeDemo

julia> x0 = [0.0; fill(1.0, 9)]

julia> f = MooncakeDemo.objective
objective (generic function with 1 method)

julia> f(x0)  # works


julia> backend = DI.SecondOrder(DI.AutoMooncakeForward(), DI.AutoMooncake())

julia> preph = DI.prepare_hessian(f, backend, x0) # works 

julia> ∇²f(x) = DI.hessian(f, preph, backend, x)

julia> @show isa(∇²f(x0), AbstractMatrix{<:Real}) # errors
ERROR: StackOverflowError:
Stacktrace:
      [1] tangent_type(::Type{Core.MethodInstance})
        @ Mooncake ./none:0
      [2] macro expansion
        @ ~/.julia/packages/Mooncake/6uvjA/src/tangents.jl:466 [inlined]
      [3] macro expansion
        @ ./none:0 [inlined]--- the above 3 lines are repeated 79982 more times ---

[gdalle]

Copied from our discussion on Slack with @willtebbutt: MethodInstance and CodeInstance seem to have mutually recursive definitions. It is unclear whether they can contain differentiable data in usual cases (but if differentiable data is passed through type parameters, they can, although all hell breaks loose in other places too).
Thus, setting the tangent type to NoTangent for both of them should be enough?

[yebai]

Witth Mooncake v0.4.166, we get #747 (comment), which is helpful for future debugging.