add Mooncake stuff

Author: jClugstor

Project.toml

@@ -43,6 +43,7 @@ MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
@@ -62,6 +63,7 @@ SciMLBaseRCallExt = "RCall"
 SciMLBaseZygoteExt = ["Zygote", "ChainRulesCore"]
 SciMLBaseDistributionsExt = "Distributions"
 SciMLBaseMonteCarloMeasurementsExt = "MonteCarloMeasurements"
+SciMLBaseMooncakeExt = "Mooncake"
 SciMLBaseReverseDiffExt = "ReverseDiff"
 SciMLBaseTrackerExt = "Tracker"
 SciMLBaseForwardDiffExt = "ForwardDiff"
@@ -88,6 +90,7 @@ MLStyle = "0.4.17"
 Makie = "0.20, 0.21, 0.22, 0.23, 0.24"
 Markdown = "1.10"
 Moshi = "0.3"
+Mooncake = "0.4"
 Measurements = "2"
 MonteCarloMeasurements = "1"
 PartialFunctions = "1.1"

ext/SciMLBaseMooncakeExt.jl

@@ -0,0 +1,25 @@
+module SciMLBaseMooncakeExt
+
+using SciMLBase, Mooncake
+using SciMLBase: ADOriginator, ChainRulesOriginator, MooncakeOriginator
+import Mooncake: rrule!!, CoDual, zero_fcodual, @is_primitive,
+                 @from_rrule, @zero_adjoint, @mooncake_overlay, MinimalCtx,
+                 NoPullback
+
+@zero_adjoint MinimalCtx Tuple{typeof(DiffEqBase.numargs), Any}
+@is_primitive MinimalCtx Tuple{
+    typeof(DiffEqBase.set_mooncakeoriginator_if_mooncake), SciMLBase.ChainRulesOriginator
+}
+
+@mooncake_overlay DiffEqBase.set_mooncakeoriginator_if_mooncake(x::SciMLBase.ADOriginator) = SciMLBase.MooncakeOriginator()
+
+function rrule!!(
+        f::CoDual{typeof(SciMLBase.set_mooncakeoriginator_if_mooncake)},
+        X::CoDual{SciMLBase.ChainRulesOriginator}
+)
+    return zero_fcodual(SciMLBase.MooncakeOriginator()), NoPullback(f, X)
+end
+
+
+
+end

src/solve.jl

@@ -484,4 +484,6 @@ totallength(x::AbstractArray) = __sum(totallength, x; init = 0)
 
 _reshape(v, siz) = reshape(v, siz)
 _reshape(v::Number, siz) = v
-_reshape(v::AbstractSciMLScalarOperator, siz) = v
+_reshape(v::AbstractSciMLScalarOperator, siz) = v
+
+set_mooncakeoriginator_if_mooncake(x::SciMLBase.ADOriginator) = x