Using DiffEqBase in Mooncake extension

ErikQQY · ErikQQY · commit 63ced6918c9a · 2025-08-26T11:52:49.000+08:00
diff --git a/Project.toml b/Project.toml
@@ -36,6 +36,7 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 [weakdeps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
@@ -58,7 +59,7 @@ SciMLBaseForwardDiffExt = "ForwardDiff"
 SciMLBaseMLStyleExt = "MLStyle"
 SciMLBaseMakieExt = "Makie"
 SciMLBaseMonteCarloMeasurementsExt = "MonteCarloMeasurements"
-SciMLBaseMooncakeExt = "Mooncake"
+SciMLBaseMooncakeExt = ["Mooncake", "DiffEqBase"]
 SciMLBasePartialFunctionsExt = "PartialFunctions"
 SciMLBasePyCallExt = "PyCall"
 SciMLBasePythonCallExt = "PythonCall"
diff --git a/ext/SciMLBaseDistributionsExt.jl b/ext/SciMLBaseDistributionsExt.jl
@@ -5,4 +5,4 @@ using Distributions, SciMLBase
 SciMLBase.handle_distribution_u0(_u0::Distributions.Sampleable) = rand(_u0)
 SciMLBase.isdistribution(_u0::Distributions.Sampleable) = true
 
-end
+end
diff --git a/ext/SciMLBaseForwardDiffExt.jl b/ext/SciMLBaseForwardDiffExt.jl
@@ -3,11 +3,12 @@ module SciMLBaseForwardDiffExt
 using SciMLBase, ForwardDiff
 using ArrayInterface
 
-import SciMLBase: 
-    wrapfun_oop, wrapfun_iip, isdualtype, value, DualEltypeChecker,
-    AbstractTimeseriesSolution, NonlinearProblem, NonlinearLeastSquaresProblem,
-    ODEProblem, SDEProblem, RODEProblem, DDEProblem, PDEProblem, DAEProblem,
-    RecursiveArrayTools, totallength, sse, anyeltypedual, reduce_tup
+import SciMLBase:
+                  wrapfun_oop, wrapfun_iip, isdualtype, value, DualEltypeChecker,
+                  AbstractTimeseriesSolution, NonlinearProblem,
+                  NonlinearLeastSquaresProblem,
+                  ODEProblem, SDEProblem, RODEProblem, DDEProblem, PDEProblem, DAEProblem,
+                  RecursiveArrayTools, totallength, sse, anyeltypedual, reduce_tup
 
 eltypedual(x) = eltype(x) <: ForwardDiff.Dual
 isdualtype(::Type{<:ForwardDiff.Dual}) = true
diff --git a/ext/SciMLBaseMonteCarloMeasurementsExt.jl b/ext/SciMLBaseMonteCarloMeasurementsExt.jl
@@ -32,4 +32,4 @@ function SciMLBase.unitfulvalue(x::Type{MonteCarloMeasurements.AbstractParticles
 end
 SciMLBase.unitfulvalue(x::MonteCarloMeasurements.AbstractParticles) = mean(x.particles)
 
-end
+end
diff --git a/ext/SciMLBaseMooncakeExt.jl b/ext/SciMLBaseMooncakeExt.jl
@@ -2,6 +2,7 @@ module SciMLBaseMooncakeExt
 
 using SciMLBase, Mooncake
 using SciMLBase: ADOriginator, ChainRulesOriginator, MooncakeOriginator
+using DiffEqBase: DiffEqBase
 import Mooncake: rrule!!, CoDual, zero_fcodual, @is_primitive,
                  @from_rrule, @zero_adjoint, @mooncake_overlay, MinimalCtx,
                  NoPullback
@@ -20,6 +21,4 @@ function rrule!!(
     return zero_fcodual(SciMLBase.MooncakeOriginator()), NoPullback(f, X)
 end
 
-
-
-end
+end
diff --git a/ext/SciMLBaseReverseDiffExt.jl b/ext/SciMLBaseReverseDiffExt.jl
@@ -54,4 +54,4 @@ function SciMLBase.promote_u0(
 end
 SciMLBase.promote_u0(u0, p::AbstractArray{<:ReverseDiff.TrackedReal}, t0) = eltype(p).(u0)
 
-end
+end
diff --git a/ext/SciMLBaseTrackerExt.jl b/ext/SciMLBaseTrackerExt.jl
@@ -33,5 +33,4 @@ SciMLBase.promote_u0(u0, p::AbstractArray{<:Tracker.TrackedReal}, t0) = eltype(p
 
 @inline Base.any(f::Function, x::Tracker.TrackedArray) = any(f, Tracker.data(x))
 
-
-end
+end
diff --git a/src/clock.jl b/src/clock.jl
@@ -130,7 +130,7 @@ end
     $(TYPEDEF)
 
 A struct representing the operation of indexing a clock to obtain a subset of the time
-points at which it ticked. The actual list of time points depends on the tick instances 
+points at which it ticked. The actual list of time points depends on the tick instances
 on which the clock was ticking, and can be obtained via `canonicalize_indexed_clock`
 by providing a timeseries solution object.
 
diff --git a/src/debug.jl b/src/debug.jl
@@ -58,7 +58,8 @@ expression. Two common reasons for this issue are:
 
 function __init__()
     Base.Experimental.register_error_hint(DomainError) do io, e
-        if e isa DomainError && occursin("will only return a complex result if called with a complex argument. Try ", e.msg)
+        if e isa DomainError &&
+           occursin("will only return a complex result if called with a complex argument. Try ", e.msg)
             println(io, DOMAINERROR_COMPLEX_MSG)
         end
     end
diff --git a/src/errors.jl b/src/errors.jl
@@ -90,7 +90,6 @@ const allowedkeywords = (:dense,
     # Parameter estimation with BVP
     :fit_parameters)
 
-
 const KWARGWARN_MESSAGE = """
                           Unrecognized keyword arguments found.
                           The only allowed keyword arguments to `solve` are:
@@ -469,4 +468,4 @@ struct LateBindingTstopsNotSupportedError <: Exception end
 function Base.showerror(io::IO, e::LateBindingTstopsNotSupportedError)
     println(io, LATE_BINDING_TSTOPS_ERROR_MESSAGE)
     println(io, TruncatedStacktraces.VERBOSE_MSG)
-end
+end
diff --git a/src/function_wrappers.jl b/src/function_wrappers.jl
@@ -5,27 +5,30 @@ Abstract base type for function wrappers used in automatic differentiation and s
 These wrappers provide specialized interfaces for computing derivatives with respect to different variables.
 
 # Type Parameter
-- `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
+
+  - `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
 """
 abstract type AbstractWrappedFunction{iip} end
 isinplace(f::AbstractWrappedFunction{iip}) where {iip} = iip
 
 """
     TimeGradientWrapper{iip, fType, uType, P} <: AbstractWrappedFunction{iip}
 
-Wraps functions to compute gradients with respect to time. This wrapper is particularly useful for 
+Wraps functions to compute gradients with respect to time. This wrapper is particularly useful for
 sensitivity analysis and optimization problems where the time dependence of the solution is critical.
 
 # Fields
-- `f`: The function to wrap
-- `uprev`: Previous state value
-- `p`: Parameters
+
+  - `f`: The function to wrap
+  - `uprev`: Previous state value
+  - `p`: Parameters
 
 # Type Parameters
-- `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
-- `fType`: Type of the wrapped function
-- `uType`: Type of the state variables
-- `P`: Type of the parameters
+
+  - `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
+  - `fType`: Type of the wrapped function
+  - `uType`: Type of the state variables
+  - `P`: Type of the parameters
 
 This wrapper enables automatic differentiation with respect to time by providing a consistent
 interface for computing `∂f/∂t` across different AD systems.
@@ -53,20 +56,22 @@ end
 """
     UJacobianWrapper{iip, fType, tType, P} <: AbstractWrappedFunction{iip}
 
-Wraps functions to compute Jacobians with respect to state variables `u`. This is one of the most 
-commonly used wrappers in the SciML ecosystem for computing the derivative of the right-hand side 
+Wraps functions to compute Jacobians with respect to state variables `u`. This is one of the most
+commonly used wrappers in the SciML ecosystem for computing the derivative of the right-hand side
 function with respect to the state variables.
 
 # Fields
-- `f`: The function to wrap
-- `t`: Time value
-- `p`: Parameters
+
+  - `f`: The function to wrap
+  - `t`: Time value
+  - `p`: Parameters
 
 # Type Parameters
-- `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
-- `fType`: Type of the wrapped function
-- `tType`: Type of the time variable
-- `P`: Type of the parameters
+
+  - `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
+  - `fType`: Type of the wrapped function
+  - `tType`: Type of the time variable
+  - `P`: Type of the parameters
 
 This wrapper enables efficient computation of `∂f/∂u` for Jacobian calculations in numerical solvers
 and automatic differentiation systems.
@@ -97,19 +102,21 @@ end
 """
     TimeDerivativeWrapper{iip, F, uType, P} <: AbstractWrappedFunction{iip}
 
-Wraps functions to compute derivatives with respect to time. This wrapper is used when you need to 
+Wraps functions to compute derivatives with respect to time. This wrapper is used when you need to
 compute `∂f/∂t` for sensitivity analysis or when the function has explicit time dependence.
 
 # Fields
-- `f`: The function to wrap
-- `u`: State variables
-- `p`: Parameters
+
+  - `f`: The function to wrap
+  - `u`: State variables
+  - `p`: Parameters
 
 # Type Parameters
-- `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
-- `F`: Type of the wrapped function
-- `uType`: Type of the state variables
-- `P`: Type of the parameters
+
+  - `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
+  - `F`: Type of the wrapped function
+  - `uType`: Type of the state variables
+  - `P`: Type of the parameters
 
 This wrapper provides a consistent interface for time derivative computations across different
 automatic differentiation backends.
@@ -134,19 +141,21 @@ end
 """
     UDerivativeWrapper{iip, F, tType, P} <: AbstractWrappedFunction{iip}
 
-Wraps functions to compute derivatives with respect to state variables. This wrapper is used for 
+Wraps functions to compute derivatives with respect to state variables. This wrapper is used for
 computing `∂f/∂u` and is fundamental for Jacobian computations in numerical solvers.
 
 # Fields
-- `f`: The function to wrap
-- `t`: Time value
-- `p`: Parameters
+
+  - `f`: The function to wrap
+  - `t`: Time value
+  - `p`: Parameters
 
 # Type Parameters
-- `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
-- `F`: Type of the wrapped function
-- `tType`: Type of the time variable
-- `P`: Type of the parameters
+
+  - `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
+  - `F`: Type of the wrapped function
+  - `tType`: Type of the time variable
+  - `P`: Type of the parameters
 
 This wrapper enables efficient state derivative computations for use in automatic differentiation
 and numerical analysis algorithms.
@@ -169,24 +178,27 @@ UDerivativeWrapper(f::F, t, p) where {F} = UDerivativeWrapper{isinplace(f, 4)}(f
 """
     ParamJacobianWrapper{iip, fType, tType, uType} <: AbstractWrappedFunction{iip}
 
-Wraps functions to compute Jacobians with respect to parameters `p`. This wrapper is essential for 
+Wraps functions to compute Jacobians with respect to parameters `p`. This wrapper is essential for
 parameter estimation, inverse problems, and sensitivity analysis with respect to model parameters.
 
 # Fields
-- `f`: The function to wrap
-- `t`: Time value
-- `u`: State variables
+
+  - `f`: The function to wrap
+  - `t`: Time value
+  - `u`: State variables
 
 # Type Parameters
-- `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
-- `fType`: Type of the wrapped function
-- `tType`: Type of the time variable
-- `uType`: Type of the state variables
+
+  - `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
+  - `fType`: Type of the wrapped function
+  - `tType`: Type of the time variable
+  - `uType`: Type of the state variables
 
 This wrapper enables efficient computation of `∂f/∂p` for parameter sensitivity analysis and
 optimization algorithms.
 """
-mutable struct ParamJacobianWrapper{iip, fType, tType, uType} <: AbstractWrappedFunction{iip}
+mutable struct ParamJacobianWrapper{iip, fType, tType, uType} <:
+               AbstractWrappedFunction{iip}
     f::fType
     t::tType
     u::uType
@@ -211,17 +223,19 @@ end
 """
     JacobianWrapper{iip, fType, pType} <: AbstractWrappedFunction{iip}
 
-A general-purpose Jacobian wrapper that can be configured for different types of Jacobian computations. 
+A general-purpose Jacobian wrapper that can be configured for different types of Jacobian computations.
 This wrapper provides a unified interface for various Jacobian calculations across the SciML ecosystem.
 
 # Fields
-- `f`: The function to wrap
-- `p`: Parameters
+
+  - `f`: The function to wrap
+  - `p`: Parameters
 
 # Type Parameters
-- `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
-- `fType`: Type of the wrapped function
-- `pType`: Type of the parameters
+
+  - `iip`: Boolean indicating if the function is in-place (`true`) or out-of-place (`false`)
+  - `fType`: Type of the wrapped function
+  - `pType`: Type of the parameters
 
 This wrapper provides a flexible interface for Jacobian computations that can adapt to different
 automatic differentiation backends and numerical methods.
diff --git a/src/integrator_interface.jl b/src/integrator_interface.jl
@@ -83,7 +83,7 @@ end
 """
     ratenoise_cache(integrator::DEIntegrator)
 
-Returns cache arrays for rate noise in stochastic differential equations. 
+Returns cache arrays for rate noise in stochastic differential equations.
 Returns an empty tuple by default for deterministic problems.
 """
 ratenoise_cache(i::DEIntegrator) = ()
diff --git a/src/scimlfunctions.jl b/src/scimlfunctions.jl
@@ -4966,13 +4966,16 @@ function __has_initializeprob(f)
     has_initialization_data(f) && hasfield(typeof(f.initialization_data), :initializeprob)
 end
 function __has_update_initializeprob!(f)
-    has_initialization_data(f) && hasfield(typeof(f.initialization_data), :update_initializeprob!)
+    has_initialization_data(f) &&
+        hasfield(typeof(f.initialization_data), :update_initializeprob!)
 end
 function __has_initializeprobmap(f)
-    has_initialization_data(f) && hasfield(typeof(f.initialization_data), :initializeprobmap)
+    has_initialization_data(f) &&
+        hasfield(typeof(f.initialization_data), :initializeprobmap)
 end
 function __has_initializeprobpmap(f)
-    has_initialization_data(f) && hasfield(typeof(f.initialization_data), :initializeprobpmap)
+    has_initialization_data(f) &&
+        hasfield(typeof(f.initialization_data), :initializeprobpmap)
 end
 __has_initialization_data(f) = hasfield(typeof(f), :initialization_data)
 __has_polynomialize(f) = hasfield(typeof(f), :polynomialize)
diff --git a/src/solve.jl b/src/solve.jl
@@ -248,7 +248,6 @@ function __solve(prob::OptimizationProblem, alg, args...; kwargs...)
     throw(OptimizerMissingError(alg))
 end
 
-
 # Functions used in solve dispatches
 
 eltypedual(x) = false
@@ -340,8 +339,8 @@ contain the `u0` and `p` passed as keyword arguments.
 
 # Keyword Arguments
 
-- `u0`, `p`: Override values for `state_values(prob)` and `parameter_values(prob)` which
-  should be used instead of the ones in `prob`.
+  - `u0`, `p`: Override values for `state_values(prob)` and `parameter_values(prob)` which
+    should be used instead of the ones in `prob`.
 """
 function get_updated_symbolic_problem(indp, prob; kw...)
     return prob
@@ -515,4 +514,4 @@ set_mooncakeoriginator_if_mooncake(x::SciMLBase.ADOriginator) = x
         W >>>= 1
     end
     q
-end
+end
diff --git a/test/downstream/splitodeproblem_cache.jl b/test/downstream/splitodeproblem_cache.jl
@@ -3,18 +3,18 @@ using OrdinaryDiffEq, Test
 # https://github.com/SciML/OrdinaryDiffEq.jl/issues/2719
 
 # set up functions
-function f1!(du, u , p, t)
-    du .= -u.^2
+function f1!(du, u, p, t)
+    du .= -u .^ 2
     return nothing
 end
 
-function f2!(du, u , p, t)
+function f2!(du, u, p, t)
     du .= 2u
     return nothing
 end
 
 function f!(du, u, p, t)
-    du .= -u.^2 .+ 2u
+    du .= -u .^ 2 .+ 2u
     return nothing
 end
 

-Original file line number
+Diff line change
 end
 SciMLBase.unitfulvalue(x::MonteCarloMeasurements.AbstractParticles) = mean(x.particles)
 -end
 +end