Compatibility with Enzyme

[gdalle]

As of today, Enzyme supports custom rules, which means we could make ImplicitDifferentation compatible with this promising AD backend

[gdalle]

Done by #135 for forward mode

[tpapp]

Is support for reverse mode planned?

[gdalle]

Actually I removed even forward mode in the latest releases because Enzyme was giving me too many headaches (and also because no one is using this package anyway).
If you want to use Enzyme, the best way right now is to import the chain rule from ChainRules.
Do you have an application where ImplicitDifferentiation would be of use?

[tpapp]

If you want to use Enzyme, the best way right now is to import the chain rule from ChainRules.

I am confused, does ChainRules handle implicit differentiation?

Do you have an application where ImplicitDifferentiation would be of use?

Plenty (I am an economist, a lot of my models are defined by implicit systems). I am already using this package, but with ForwardDiff. In my experience Enzyme is the most robust and fastest reverse AD solution so support would be great, but I guess I can code up a rule myself.

Or, unless you deem it impossible, contribute a solution with Enzyme for this package?

[gdalle]

I am confused, does ChainRules handle implicit differentiation?

No, that is why ImplicitDifferentiation.jl defines a ChainRulesCore.rrule for its callable ImplicitFunction type. What I meant is that there is it doesn't also define a native EnzymeRule at the moment.

If you want to use Enzyme as the outer differentiator, my idea was that you could import the corresponding chain rule with Enzyme.@import_rrule. I just tried it myself however and I got an Enzyme error I don't understand. Here's the code I had in mind:

using ChainRulesCore
using Enzyme
using ImplicitDifferentiation

struct CustomError{C} <: Exception
    content::C
end

function solver(x)
    y = sqrt.(x)
    # make Enzyme error to check it doesn't differentiate through the solver
    try
        throw(CustomError(y))
    catch e
        if e isa CustomError
            return e.content
        end
    end
end

x0 = [1.0, 2.0]

solver(x0)
Enzyme.gradient(Enzyme.Reverse, x -> sum(solver(x)), x0)  # fails

conditions(x, y, z) = y .^ 2 .- x
implicit = ImplicitFunction(solver, conditions)

implicitfirst = first ∘ implicit

Enzyme.@import_rrule typeof(implicitfirst) AbstractArray

Enzyme.gradient(Enzyme.Reverse, x -> sum(implicitfirst(x)), x0)  # should succeed?

And here's the error, in case you have any clue. I'm on vacation right now so I won't investigate it for some time:

Details

julia> Enzyme.gradient(Enzyme.Reverse, x -> sum(implicitfirst(x)), x0)
ERROR: AssertionError: Enzyme : mismatch between innerTy LLVM.PointerType({} addrspace(10)**) and tape type LLVM.PointerType({} addrspace(10)*)
tape_idx=1
true_idx=2
isKWCall=false
kwtup=nothing
funcTy=Const{ComposedFunction{typeof(first), ImplicitFunction{typeof(solver), typeof(conditions), IterativeLinearSolver{Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, OperatorRepresentation, Nothing, true}}}
isghostty(funcTy)=true
miRT=Tuple{Nothing}
sret=nothing
returnRoots=nothing
swiftself=false
RT=Active{Float64}
rev_RT=Tuple{Nothing}
applicablefn=true
tape=LLVM.Argument(0x0000000335ba6dc8)
llvmf={} addrspace(10)* ({} addrspace(10)*, {} addrspace(10)**, i32, {} addrspace(10)**)
TapeT=Any
mi=MethodInstance for (::ComposedFunction{typeof(first), ImplicitFunction{typeof(solver), typeof(conditions), IterativeLinearSolver{Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, OperatorRepresentation, Nothing, true}})(::Vector{Float64})
ami=MethodInstance for EnzymeCore.EnzymeRules.augmented_primal(::EnzymeCore.EnzymeRules.RevConfigWidth{1, false, false, (false, true), false, false}, ::Const{ComposedFunction{typeof(first), ImplicitFunction{typeof(solver), typeof(conditions), IterativeLinearSolver{Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, OperatorRepresentation, Nothing, true}}}, ::Type{Active{Float64}}, ::Duplicated{Vector{Float64}})
rev_TT =Tuple{EnzymeCore.EnzymeRules.RevConfigWidth{1, false, false, (false, true), false, false}, Const{ComposedFunction{typeof(first), ImplicitFunction{typeof(solver), typeof(conditions), IterativeLinearSolver{Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}}, OperatorRepresentation, Nothing, true}}}, Active{Float64}, Any, Duplicated{Vector{Float64}}}

Stacktrace:
  [1] enzyme_custom_common_rev(forward::Bool, B::LLVM.IRBuilder, orig::LLVM.CallInst, gutils::Enzyme.Compiler.GradientUtils, normalR::Ptr{…}, shadowR::Ptr{…}, tape::LLVM.Argument)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/rules/customrules.jl:1096
  [2] enzyme_custom_rev(B::LLVM.IRBuilder, orig::LLVM.CallInst, gutils::Enzyme.Compiler.GradientUtils, tape::Union{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/rules/customrules.jl:1519
  [3] enzyme_custom_rev_cfunc(B::Ptr{…}, OrigCI::Ptr{…}, gutils::Ptr{…}, tape::Ptr{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/rules/llvmrules.jl:48
  [4] EnzymeCreatePrimalAndGradient(logic::Enzyme.Logic, todiff::LLVM.Function, retType::Enzyme.API.CDIFFE_TYPE, constant_args::Vector{…}, TA::Enzyme.TypeAnalysis, returnValue::Bool, dretUsed::Bool, mode::Enzyme.API.CDerivativeMode, runtimeActivity::Bool, strongZero::Bool, width::Int64, additionalArg::Ptr{…}, forceAnonymousTape::Bool, typeInfo::Enzyme.FnTypeInfo, uncacheable_args::Vector{…}, augmented::Ptr{…}, atomicAdd::Bool)
    @ Enzyme.API ~/.julia/packages/Enzyme/4HcV9/src/api.jl:270
  [5] enzyme!(job::GPUCompiler.CompilerJob{…}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::NTuple{…} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{…}, boxedArgs::Set{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:1777
  [6] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:4822
  [7] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
    @ GPUCompiler ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:67
  [8] compile
    @ ~/.julia/packages/GPUCompiler/Ecaql/src/driver.jl:55 [inlined]
  [9] _thunk(job::GPUCompiler.CompilerJob{…}, postopt::Bool)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:5685
 [10] _thunk
    @ ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:5683 [inlined]
 [11] cached_compilation
    @ ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:5737 [inlined]
 [12] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, edges::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:5851
 [13] thunk_generator(world::UInt64, source::Union{…}, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type, strongzero::Type)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:6044
 [14] macro expansion
    @ ~/.julia/packages/Enzyme/4HcV9/src/rules/jitrules.jl:464 [inlined]
 [15] runtime_generic_augfwd(::Type{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::ComposedFunction{…}, ::Nothing, ::Vector{…}, ::Vector{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/4HcV9/src/rules/jitrules.jl:574
 [16] #47
    @ ~/Documents/GitHub/Julia/_Scratchpad/mwe.jl:32 [inlined]
 [17] augmented_julia__47_43892wrap
    @ ~/Documents/GitHub/Julia/_Scratchpad/mwe.jl:0
 [18] macro expansion
    @ ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:5613 [inlined]
 [19] enzyme_call
    @ ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:5147 [inlined]
 [20] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/4HcV9/src/compiler.jl:5086 [inlined]
 [21] autodiff
    @ ~/.julia/packages/Enzyme/4HcV9/src/Enzyme.jl:408 [inlined]
 [22] autodiff
    @ ~/.julia/packages/Enzyme/4HcV9/src/Enzyme.jl:538 [inlined]
 [23] macro expansion
    @ ~/.julia/packages/Enzyme/4HcV9/src/sugar.jl:324 [inlined]
 [24] gradient(::ReverseMode{false, false, false, FFIABI, false, false}, ::var"#47#48", ::Vector{Float64})
    @ Enzyme ~/.julia/packages/Enzyme/4HcV9/src/sugar.jl:262
 [25] top-level scope
    @ ~/Documents/GitHub/Julia/_Scratchpad/mwe.jl:32
Some type information was truncated. Use `show(err)` to see complete types.

Note however that you can still use Enzyme as the inner differentiator (inside the implicit function theorem), since it doesn't require a custom rule, it only requires selecting the right backends while constructing your ImplicitFunction.

[gdalle]

In any case I'd very much welcome a contribution for forward and reverse Enzyme rules!

[wsmoses]

hm this is definitely weird (and seems that the Julia Compiler is giving us Ir that does not correspond to the corresponding julia calling convention).

Can you open an issue on Enzyme with a mwe of this (ideally reproducing without the extra package, if you can find such a reproducer)?

Also in the future definitely feel free to tag me so I have visibility/can offer support

[gdalle]

Thanks! I opened EnzymeAD/Enzyme.jl#2513