Backend switching for Mooncake

DifferentiationInterface/Project.toml

@@ -38,7 +38,7 @@ DifferentiationInterfaceFiniteDiffExt = "FiniteDiff"
 DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences"
 DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
 DifferentiationInterfaceGTPSAExt = "GTPSA"
-DifferentiationInterfaceMooncakeExt = "Mooncake"
+DifferentiationInterfaceMooncakeExt = ["ChainRulesCore", "Mooncake"]
 DifferentiationInterfacePolyesterForwardDiffExt = ["PolyesterForwardDiff", "ForwardDiff", "DiffResults"]
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceSparseArraysExt = "SparseArrays"

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -3,14 +3,21 @@ 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
+    zero_tangent,
+    @is_primitive,
+    zero_fcodual,
+    MinimalCtx,
+    NoRData
+
+using ChainRulesCore: ChainRulesCore, rrule
 
 DI.check_available(::AutoMooncake) = true
 
@@ -26,5 +33,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,22 @@
+@is_primitive MinimalCtx Tuple{CoDual{<:DI.DifferentiateWith},CoDual{<:AbstractArray}}
+@is_primitive MinimalCtx Tuple{CoDual{<:DI.DifferentiateWith},CoDual{<:Number}}
+
+function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, args::CoDual...)
+    primal_func = Mooncake.primal(dw)
+    primal_args = map(arg -> Mooncake.primal(arg), args)
+
+    (; f, backend) = primal_func
+    y = f(primal_args...)
+
+    prep_same = DI.prepare_pullback_same_point_nokwarg(
+        Val(true), f, backend, primal_args..., (y,)
+    )
+
+    function pullback!!(dy)
+        tx = DI.pullback(f, prep_same, backend, primal_args, (dy,))
+        args_rdata = map((x) -> (x, Mooncake.zero_rdata(x)), only(tx))
+        return NoRData(), args_rdata...
+    end
+
+    return zero_fcodual(y), pullback!!
+end