feat: backend switching for Mooncake

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,
+    fdata,
+    primal
 
 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,49 @@
+@is_primitive MinimalCtx Tuple{DI.DifferentiateWith,<:AbstractArray}
+@is_primitive MinimalCtx Tuple{DI.DifferentiateWith,<:Number}
+
+function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:Number})
+    primal_func = primal(dw)
+    primal_x = primal(x)
+    (; f, backend) = primal_func
+    y = zero_fcodual(f(primal_x))
+
+    # output is a vector, so we need to use the vector pullback
+    function pullback!!(dy::NoRData)
+        tx = DI.pullback(f, backend, primal_x, (fdata(y.dx),))
+        return NoRData(), only(tx)
+    end
+
+    # output is a scalar, so we can use the scalar pullback
+    function pullback!!(dy)
+        tx = DI.pullback(f, backend, primal_x, (dy,))
+        return NoRData(), only(tx)
+    end
+
+    return y, pullback!!
+end
+
+function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:AbstractArray})
+    primal_func = primal(dw)
+    primal_x = primal(x)
+    fdata_arg = fdata(x.dx)
+    (; f, backend) = primal_func
+    y = zero_fcodual(f(primal_x))
+
+    # output is a vector, so we need to use the vector pullback
+    function pullback!!(dy::NoRData)
+        tx = DI.pullback(f, backend, primal_x, (fdata(y.dx),))
+        fdata_arg .+= only(tx)
+        return NoRData(), dy
+    end
+
+    # output is a scalar, so we can use the scalar pullback
+    function pullback!!(dy)
+        tx = DI.pullback(f, backend, primal_x, (dy,))
+        fdata_arg .+= only(tx)
+        return NoRData(), NoRData()
+    end
+
+    # in case x is mutated when passed into f
+    x = CoDual(primal_x, x.dx)
+    return y, pullback!!
+end

DifferentiationInterface/test/Back/DifferentiateWith/test.jl

@@ -1,11 +1,12 @@
 using Pkg
-Pkg.add(["FiniteDiff", "ForwardDiff", "Zygote"])
+Pkg.add(["FiniteDiff", "ForwardDiff", "Zygote", "Mooncake"])
 
 using DifferentiationInterface, DifferentiationInterfaceTest
 import DifferentiationInterfaceTest as DIT
 using FiniteDiff: FiniteDiff
 using ForwardDiff: ForwardDiff
 using Zygote: Zygote
+using Mooncake: Mooncake
 using Test
 
 LOGGING = get(ENV, "CI", "false") == "false"
@@ -24,7 +25,7 @@ function differentiatewith_scenarios()
 end
 
 test_differentiation(
-    [AutoForwardDiff(), AutoZygote()],
+    [AutoForwardDiff(), AutoZygote(), AutoMooncake(; config=nothing)],
     differentiatewith_scenarios();
     excluded=SECOND_ORDER,
     logging=LOGGING,