overlay, more calls to set_mooncakeoriginator_if_mooncake

Author: AstitvaAggarwal

ext/DiffEqBaseMooncakeExt.jl

@@ -34,6 +34,8 @@ import Mooncake: rrule!!, CoDual, zero_fcodual, @is_primitive,
     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(DiffEqBase.set_mooncakeoriginator_if_mooncake)},
         X::CoDual{SciMLBase.ChainRulesOriginator}

src/solve.jl

@@ -1093,9 +1093,9 @@ function solve(prob::AbstractDEProblem, args...; sensealg = nothing,
     p = p !== nothing ? p : prob.p
 
     if wrap isa Val{true}
-        wrap_sol(solve_up(prob, sensealg, u0, p, args...; kwargs...))
+        wrap_sol(solve_up(prob, sensealg, u0, p, args...; originator=set_mooncakeoriginator_if_mooncake(SciMLBase.ChainRulesOriginator()), kwargs...))
     else
-        solve_up(prob, sensealg, u0, p, args...; kwargs...)
+        solve_up(prob, sensealg, u0, p, args...; originator=set_mooncakeoriginator_if_mooncake(SciMLBase.ChainRulesOriginator()), kwargs...)
     end
 end