Merge pull request #1158 from SciML/mooncake

Setup Mooncake rrule

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "DiffEqBase"
 uuid = "2b5f629d-d688-5b77-993f-72d75c75574e"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "6.174.0"
+version = "6.175.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -45,6 +45,7 @@ GeneralizedGenerated = "6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
@@ -61,6 +62,7 @@ DiffEqBaseGeneralizedGeneratedExt = "GeneralizedGenerated"
 DiffEqBaseMPIExt = "MPI"
 DiffEqBaseMeasurementsExt = "Measurements"
 DiffEqBaseMonteCarloMeasurementsExt = "MonteCarloMeasurements"
+DiffEqBaseMooncakeExt = "Mooncake"
 DiffEqBaseReverseDiffExt = "ReverseDiff"
 DiffEqBaseSparseArraysExt = "SparseArrays"
 DiffEqBaseTrackerExt = "Tracker"
@@ -91,14 +93,15 @@ MPI = "0.20"
 Markdown = "1.9"
 Measurements = "2"
 MonteCarloMeasurements = "1"
+Mooncake = "0.4"
 MuladdMacro = "0.2.1"
 Parameters = "0.12.0"
 PrecompileTools = "1"
 Printf = "1.9"
 RecursiveArrayTools = "3"
 Reexport = "1.0"
 ReverseDiff = "1"
-SciMLBase = "2.60.0"
+SciMLBase = "2.94.0"
 SciMLOperators = "0.3, 0.4, 1"
 SciMLStructures = "1.5"
 Setfield = "1"

ext/DiffEqBaseChainRulesCoreExt.jl

@@ -2,7 +2,7 @@ module DiffEqBaseChainRulesCoreExt
 
 using DiffEqBase
 using DiffEqBase.SciMLBase
-import DiffEqBase: numargs, AbstractSensitivityAlgorithm, AbstractDEProblem
+import DiffEqBase: numargs, AbstractSensitivityAlgorithm, AbstractDEProblem, set_mooncakeoriginator_if_mooncake
 
 import ChainRulesCore
 import ChainRulesCore: NoTangent
@@ -15,7 +15,7 @@ function ChainRulesCore.frule(::typeof(DiffEqBase.solve_up), prob,
         u0, p, args...;
         kwargs...)
     DiffEqBase._solve_forward(
-        prob, sensealg, u0, p, SciMLBase.ChainRulesOriginator(), args...;
+        prob, sensealg, u0, p, set_mooncakeoriginator_if_mooncake(SciMLBase.ChainRulesOriginator()), args...;
         kwargs...)
 end
 
@@ -24,7 +24,7 @@ function ChainRulesCore.rrule(::typeof(DiffEqBase.solve_up), prob::AbstractDEPro
         u0, p, args...;
         kwargs...)
     DiffEqBase._solve_adjoint(
-        prob, sensealg, u0, p, SciMLBase.ChainRulesOriginator(), args...;
+        prob, sensealg, u0, p, set_mooncakeoriginator_if_mooncake(SciMLBase.ChainRulesOriginator()), args...;
         kwargs...)
 end
 

ext/DiffEqBaseMooncakeExt.jl

@@ -0,0 +1,22 @@
+module DiffEqBaseMooncakeExt
+
+using DiffEqBase, Mooncake
+using DiffEqBase: SciMLBase
+using SciMLBase: ADOriginator, MooncakeOriginator
+Mooncake.@from_rrule(
+    Mooncake.MinimalCtx,
+    Tuple{
+        typeof(DiffEqBase.solve_up),
+        DiffEqBase.AbstractDEProblem,
+        Union{Nothing, DiffEqBase.AbstractSensitivityAlgorithm},
+        Any,
+        Any,
+        Any,
+    },
+    true,
+    )
+
+Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{typeof(DiffEqBase.numargs), Any}
+Mooncake.@mooncake_overlay DiffEqBase.set_mooncakeoriginator_if_mooncake(x::ADOriginator) = MooncakeOriginator
+
+end

src/utils.jl

@@ -55,6 +55,8 @@ timedepentdtmin(::Any, dtmin) = abs(dtmin)
 
 maybe_with_logger(f, logger) = logger === nothing ? f() : Logging.with_logger(f, logger)
 
+set_mooncakeoriginator_if_mooncake(x::SciMLBase.ADOriginator) = x
+
 function default_logger(logger)
     Logging.min_enabled_level(logger) ≤ ProgressLogging.ProgressLevel && return nothing