Handles backend switching for Mooncake using ChainRules

Author: AstitvaAggarwal

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -3,14 +3,19 @@ module DifferentiationInterfaceMooncakeExt
 using ADTypes: ADTypes, AutoMooncake
 import DifferentiationInterface as DI
 using Mooncake:
+    Mooncake,
     CoDual,
     Config,
     prepare_gradient_cache,
     prepare_pullback_cache,
     tangent_type,
     value_and_gradient!!,
     value_and_pullback!!,
-    zero_tangent
+    @from_rrule,
+    MinimalCtx,
+    NoFData
+
+using ChainRulesCore: ChainRulesCore, rrule
 
 DI.check_available(::AutoMooncake) = true
 
@@ -26,5 +31,6 @@ mycopy(x) = deepcopy(x)
 
 include("onearg.jl")
 include("twoarg.jl")
+include("differentiate_with.jl")
 
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl

@@ -0,0 +1,15 @@
+function define_rule!(primal_func, primal_args)
+    return eval(:(@from_rrule MinimalCtx Tuple{$primal_func,$primal_args...}))
+end
+
+function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, args::CoDual...)
+    primal_func = typeof(Mooncake.primal(dw))
+    primal_args = typeof.(map(arg -> Mooncake.primal(arg), args))
+    # use the DI.chainrule wrapper inside @from_rrule to create a custom rrule!!
+
+    # macro evaluation in global scope with more specialized types (@fromrrule requires non generic types)
+    define_rule!(primal_func, primal_args)
+
+    # Use the ChainRuleCore rrule mapping with backends, calling Mooncake rule!! that now wraps around that ChainRulesCore rrule.
+    return Base.invokelatest(Mooncake.rrule!!, CoDual(primal_func, dw.dx), args...)
+end