Active return values with automatic pullback (differential return value) deduction only supported for floating-like values and not type EnzymeCore.Active{Float64}

[gdalle]

Hi there! I'm trying to optimize DifferentiationInterface.gradient by using autodiff everywhere, as discussed with @wsmoses in TuringLang/AdvancedVI.jl#98 (comment).

In the PR JuliaDiff/DifferentiationInterface.jl#515, going from Enzyme.gradient to Enzyme.autodiff in DI caused new bugs for second-order differentiation (both forward-over-reverse and reverse-over-reverse).

Error message:

Active return values with automatic pullback (differential return value)
deduction only supported for floating-like values and not type
EnzymeCore.Active{Float64}. If mutable memory, please use Duplicated.
Otherwise, you can explicitly specify a pullback by using split mode,
e.g. autodiff_thunk(ReverseSplitWithPrimal, ...)

Stacktrace

  Stacktrace:
    [1] default_adjoint
      @ ~/.julia/packages/Enzyme/xD7hH/src/compiler.jl:8385 [inlined]
    [2] autodiff_deferred
      @ ~/.julia/packages/Enzyme/xD7hH/src/Enzyme.jl:770 [inlined]
    [3] autodiff
      @ ~/.julia/packages/Enzyme/xD7hH/src/Enzyme.jl:537 [inlined]
    [4] autodiff
      @ ~/.julia/packages/Enzyme/xD7hH/src/Enzyme.jl:504 [inlined]
    [5] gradient
      @ ~/work/DifferentiationInterface.jl/DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl:223
    [6] macro expansion
      @ ~/.julia/packages/Enzyme/xD7hH/src/compiler.jl:8839 [inlined]
    [7] enzyme_call
      @ ~/.julia/packages/Enzyme/xD7hH/src/compiler.jl:8405 [inlined]
    [8] PrimalErrorThunk
      @ ~/.julia/packages/Enzyme/xD7hH/src/compiler.jl:8162 [inlined]
    [9] batch_seeded_autodiff_thunk(rmode::EnzymeCore.ReverseModeSplit{true, true, false, 0, true, EnzymeCore.FFIABI, false, true, false}, dresults::NTuple{6, Matrix{Float64}}, f::EnzymeCore.Const{DifferentiationInterface.var"#inner_gradient#46"{typeof(DifferentiationInterfaceTest.arr_to_num_linalg), AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, DifferentiationInterface.Rewrap{0, Tuple{}}}}, ::Type{EnzymeCore.BatchDuplicated}, args::EnzymeCore.BatchDuplicated{Matrix{Float64}, 6})
      @ DifferentiationInterfaceEnzymeExt ~/work/DifferentiationInterface.jl/DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl:33
   [10] value_and_pullback!(::DifferentiationInterface.var"#inner_gradient#46"{typeof(DifferentiationInterfaceTest.arr_to_num_linalg), AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, DifferentiationInterface.Rewrap{0, Tuple{}}}, ::NTuple{6, Matrix{Float64}}, ::DifferentiationInterface.NoPullbackPrep, ::AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, ::Matrix{Float64}, ::NTuple{6, Matrix{Float64}})
      @ DifferentiationInterfaceEnzymeExt ~/work/DifferentiationInterface.jl/DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl:187
   [11] pullback!(::DifferentiationInterface.var"#inner_gradient#46"{typeof(DifferentiationInterfaceTest.arr_to_num_linalg), AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, DifferentiationInterface.Rewrap{0, Tuple{}}}, ::NTuple{6, Matrix{Float64}}, ::DifferentiationInterface.NoPullbackPrep, ::AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, ::Matrix{Float64}, ::NTuple{6, Matrix{Float64}})
      @ DifferentiationInterfaceEnzymeExt ~/work/DifferentiationInterface.jl/DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl:208
   [12] hvp!(::typeof(DifferentiationInterfaceTest.arr_to_num_linalg), ::NTuple{6, Matrix{Float64}}, ::DifferentiationInterface.ReverseOverReverseHVPPrep{DifferentiationInterface.var"#inner_gradient#46"{typeof(DifferentiationInterfaceTest.arr_to_num_linalg), AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, DifferentiationInterface.Rewrap{0, Tuple{}}}, DifferentiationInterface.NoPullbackPrep}, ::AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, ::Matrix{Float64}, ::NTuple{6, Matrix{Float64}})
      @ DifferentiationInterface ~/work/DifferentiationInterface.jl/DifferentiationInterface.jl/DifferentiationInterface/src/second_order/hvp.jl:271
   [13] hessian!
      @ ~/work/DifferentiationInterface.jl/DifferentiationInterface.jl/DifferentiationInterface/src/second_order/hessian.jl:123 [inlined]
   [14] value_gradient_and_hessian!(::typeof(DifferentiationInterfaceTest.arr_to_num_linalg), ::Matrix{Float64}, ::Matrix{Float64}, ::DifferentiationInterface.HVPGradientHessianPrep{6, NTuple{6, Matrix{Float64}}, NTuple{6, Matrix{Float64}}, DifferentiationInterface.ReverseOverReverseHVPPrep{DifferentiationInterface.var"#inner_gradient#46"{typeof(DifferentiationInterfaceTest.arr_to_num_linalg), AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, DifferentiationInterface.Rewrap{0, Tuple{}}}, DifferentiationInterface.NoPullbackPrep}, DifferentiationInterfaceEnzymeExt.EnzymeGradientPrep{Matrix{Float64}}}, ::AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, ::Matrix{Float64})
      @ DifferentiationInterface ~/work/DifferentiationInterface.jl/DifferentiationInterface.jl/DifferentiationInterface/src/second_order/hessian.jl:159
   [15] value_gradient_and_hessian!(::typeof(DifferentiationInterfaceTest.arr_to_num_linalg), ::Matrix{Float64}, ::Matrix{Float64}, ::AutoEnzyme{EnzymeCore.ReverseMode{false, false, EnzymeCore.FFIABI, false, false}, Nothing}, ::Matrix{Float64})

Can you tell us what this message means? I'm assuming there are tricks in Enzyme.gradient that give it better compatibility with nested autodiff? How can I reproduce them?

Related issues:

• Don't rely on Enzyme.gradient, use Enzyme.autodiff directly JuliaDiff/DifferentiationInterface.jl#512
• confusing error message in forward over reverse #1870
• Error in Forward over Reverse Mode for Hessian #1628

[wsmoses]

So it says that the return type of the function you're differentiating is not a float, and an Active{Float}. Only float-like operations are supported as return types for reverse mode.

[Red-Portal]

Is this a green light for using DI on everything in AdvancedVI?

[wsmoses]

No, this was just helping @gdalle understand what an error message means

[Red-Portal]

Ah, I should have paid more attention. My question should have been: "Can we use DI on everything once DI calls autodiff for gradients too?"

[wsmoses]

I’m not sure, it depends on what and how DI implements things.

In general I’m somewhat skeptical that this will be the case. In losing out on the fine grain control of a direct autodiff in user code, you’re probably going to have to trade off compatibility, performance, or both.

For codes that want to quickly adopt or test out a bunch of AD backends, DI is the right tool for the job.

However, in the case that either performance or ensuring as many codes as possible are supported is critical, I’d probably recommend directly calling Enzyme.

Especially for places that already have that direct Enzyme support added, I don’t think it’s worth the potential loss by getting rid of it at this time.

[gdalle]

So it says that the return type of the function you're differentiating is not a float, and an Active{Float}. Only float-like operations are supported as return types for reverse mode.

I'm sorry I still don't understand what this means in concrete terms. Would you know what I need to change in the code below to make it work for higher-order like Enzyme.gradient does?

https://github.com/gdalle/DifferentiationInterface.jl/blob/26e692de5c343fc89de5990065edeff936732477/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl#L215-L230