added rules

Author: AstitvaAggarwal

DifferentiationInterface/Project.toml

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

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -15,9 +15,9 @@ using Mooncake:
     @is_primitive,
     zero_fcodual,
     MinimalCtx,
-    NoRData
-
-using ChainRulesCore: ChainRulesCore, rrule
+    NoRData,
+    fdata,
+    primal
 
 DI.check_available(::AutoMooncake) = true
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl

@@ -1,21 +1,31 @@
-@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)
+@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 = f(primal_args...)
+    y = f(primal_x)
+
+    function pullback!!(dy)
+        tx = DI.pullback(f, backend, primal_x, (dy,))
+        return NoRData(), only(tx)
+    end
 
-    prep_same = DI.prepare_pullback_same_point_nokwarg(
-        Val(true), f, backend, primal_args..., (y,)
-    )
+    return zero_fcodual(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 = f(primal_x)
 
     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...
+        tx = DI.pullback(f, backend, primal_x, (dy,))
+        fdata_arg .+= only(tx)
+        return NoRData(), NoRData()
     end
 
     return zero_fcodual(y), pullback!!

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()],
     differentiatewith_scenarios();
     excluded=SECOND_ORDER,
     logging=LOGGING,