fix: config width

SciML · Sep 26, 2024 · b66347d · b66347d
1 parent 3e44128
commit b66347d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ext/DiffEqBaseEnzymeExt.jl b/ext/DiffEqBaseEnzymeExt.jl
@@ -6,7 +6,7 @@ using Enzyme
 import Enzyme: Const
 using ChainRulesCore
 
-function Enzyme.EnzymeRules.augmented_primal(config::Enzyme.EnzymeRules.ConfigWidth{1},
+function Enzyme.EnzymeRules.augmented_primal(config::Enzyme.EnzymeRules.RevConfigWidth{1},
         func::Const{typeof(DiffEqBase.solve_up)}, ::Type{Duplicated{RT}}, prob,
         sensealg::Union{Const{Nothing}, Const{<:DiffEqBase.AbstractSensitivityAlgorithm}},
         u0, p, args...; kwargs...) where {RT}
@@ -33,7 +33,7 @@ function Enzyme.EnzymeRules.augmented_primal(config::Enzyme.EnzymeRules.ConfigWi
     return Enzyme.EnzymeRules.AugmentedReturn{RT, RT, Any}(res[1], dres, tup::Any)
 end
 
-function Enzyme.EnzymeRules.reverse(config::Enzyme.EnzymeRules.ConfigWidth{1},
+function Enzyme.EnzymeRules.reverse(config::Enzyme.EnzymeRules.RevConfigWidth{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}