Fix enzyme extension

Author: ChrisRackauckas

Project.toml

@@ -46,8 +46,9 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [extensions]
+DiffEqBaseChainRulesCoreExt = "ChainRulesCore"
 DiffEqBaseDistributionsExt = "Distributions"
-DiffEqBaseEnzymeExt = "Enzyme"
+DiffEqBaseEnzymeExt = ["ChainRulesCore","Enzyme"]
 DiffEqBaseGeneralizedGeneratedExt = "GeneralizedGenerated"
 DiffEqBaseMPIExt = "MPI"
 DiffEqBaseMeasurementsExt = "Measurements"

ext/DiffEqBaseChainRulesCoreExt.jl

@@ -1,27 +1,28 @@
 module DiffEqBaseChainRulesCoreExt
 
 using DiffEqBase
-import DiffEqBase: numargs
+using DiffEqBase.SciMLBase
+import DiffEqBase: numargs, AbstractSensitivityAlgorithm, AbstractDEProblem
 
 import ChainRulesCore
 import ChainRulesCore: NoTangent
 
 ChainRulesCore.rrule(::typeof(numargs), f) = (numargs(f), df -> (NoTangent(), NoTangent()))
-ChainRulesCore.@non_differentiable checkkwargs(kwargshandle)
+ChainRulesCore.@non_differentiable DiffEqBase.checkkwargs(kwargshandle)
 
-function ChainRulesCore.frule(::typeof(solve_up), prob,
+function ChainRulesCore.frule(::typeof(DiffEqBase.solve_up), prob,
     sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
     u0, p, args...;
     kwargs...)
-    _solve_forward(prob, sensealg, u0, p, SciMLBase.ChainRulesOriginator(), args...;
+    DiffEqBase._solve_forward(prob, sensealg, u0, p, SciMLBase.ChainRulesOriginator(), args...;
         kwargs...)
 end
 
-function ChainRulesCore.rrule(::typeof(solve_up), prob::AbstractDEProblem,
+function ChainRulesCore.rrule(::typeof(DiffEqBase.solve_up), prob::AbstractDEProblem,
     sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
     u0, p, args...;
     kwargs...)
-    _solve_adjoint(prob, sensealg, u0, p, SciMLBase.ChainRulesOriginator(), args...;
+    DiffEqBase._solve_adjoint(prob, sensealg, u0, p, SciMLBase.ChainRulesOriginator(), args...;
         kwargs...)
 end
 

ext/DiffEqBaseEnzymeExt.jl

@@ -2,14 +2,13 @@ module DiffEqBaseEnzymeExt
 
 using DiffEqBase
 import DiffEqBase: value
-isdefined(Base, :get_extension) ? (import Enzyme) : (import ..Enzyme)
-
+using Enzyme
+import Enzyme: Const
 using ChainRulesCore
-using EnzymeCore
 
-function EnzymeCore.EnzymeRules.augmented_primal(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, func::Const{typeof(DiffEqBase.solve_up)}, ::Type{Duplicated{RT}}, prob, sensealg::Union{Const{Nothing}, Const{<:DiffEqBase.AbstractSensitivityAlgorithm}}, u0, p, args...; kwargs...) where RT
+function Enzyme.EnzymeRules.augmented_primal(config::Enzyme.EnzymeRules.ConfigWidth{1}, func::Const{typeof(DiffEqBase.solve_up)}, ::Type{Duplicated{RT}}, prob, sensealg::Union{Const{Nothing}, Const{<:DiffEqBase.AbstractSensitivityAlgorithm}}, u0, p, args...; kwargs...) where RT
     @inline function copy_or_reuse(val, idx)
-        if EnzymeCore.EnzymeRules.overwritten(config)[idx] && ismutable(val)
+        if Enzyme.EnzymeRules.overwritten(config)[idx] && ismutable(val)
             return deepcopy(val)
         else
             return val
@@ -28,15 +27,15 @@ function EnzymeCore.EnzymeRules.augmented_primal(config::EnzymeCore.EnzymeRules.
         v.= 0
     end
     tup = (dres, res[2])
-    return EnzymeCore.EnzymeRules.AugmentedReturn{RT, RT, Any}(res[1], dres, tup::Any)
+    return Enzyme.EnzymeRules.AugmentedReturn{RT, RT, Any}(res[1], dres, tup::Any)
 end
 
-function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, func::Const{typeof(DiffEqBase.solve_up)}, ::Type{<:Duplicated{RT}}, tape, prob, sensealg::Union{Const{Nothing}, Const{<:DiffEqBase.AbstractSensitivityAlgorithm}}, u0, p, args...; kwargs...) where RT
+function Enzyme.reverse(config::Enzyme.EnzymeRules.ConfigWidth{1}, func::Const{typeof(DiffEqBase.solve_up)}, ::Type{<:Duplicated{RT}}, tape, prob, sensealg::Union{Const{Nothing}, Const{<:DiffEqBase.AbstractSensitivityAlgorithm}}, u0, p, args...; kwargs...) where RT
 	dres, clos = tape
     dres = dres::RT
 	dargs = clos(dres)
     for (darg, ptr) in zip(dargs, (func, prob, sensealg, u0, p, args...))
-        if ptr isa EnzymeCore.Const
+        if ptr isa Enzyme.Const
             continue
         end
         if darg == ChainRulesCore.NoTangent()