Migrate solve dispatches from DiffEqBase to NonlinearSolveBase

Project.toml

@@ -9,7 +9,6 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BracketingNonlinearSolve = "70df07ce-3d50-431d-a3e7-ca6ddb60ac1e"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
-DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -44,24 +43,6 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 
-[sources.BracketingNonlinearSolve]
-path = "lib/BracketingNonlinearSolve"
-
-[sources.NonlinearSolveBase]
-path = "lib/NonlinearSolveBase"
-
-[sources.NonlinearSolveFirstOrder]
-path = "lib/NonlinearSolveFirstOrder"
-
-[sources.NonlinearSolveQuasiNewton]
-path = "lib/NonlinearSolveQuasiNewton"
-
-[sources.NonlinearSolveSpectralMethods]
-path = "lib/NonlinearSolveSpectralMethods"
-
-[sources.SimpleNonlinearSolve]
-path = "lib/SimpleNonlinearSolve"
-
 [extensions]
 NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
 NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
@@ -83,8 +64,8 @@ BenchmarkTools = "1.4"
 BracketingNonlinearSolve = "1"
 CommonSolve = "0.2.4"
 ConcreteStructs = "0.2.3"
-DiffEqBase = "6.188"
 DifferentiationInterface = "0.7.3"
+DiffEqBase = "6.188"
 ExplicitImports = "1.5"
 FastClosures = "0.3.2"
 FastLevenbergMarquardt = "0.1"
@@ -117,6 +98,7 @@ Preferences = "1.4.3"
 Random = "1.10"
 ReTestItems = "1.24"
 Reexport = "1.2.2"
+ReverseDiff = "1.15"
 SIAMFANLEquations = "1.0.1"
 SciMLBase = "2.116"
 SimpleNonlinearSolve = "2.1"
@@ -137,6 +119,7 @@ julia = "1.10"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
@@ -155,6 +138,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
@@ -163,7 +147,29 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[sources.BracketingNonlinearSolve]
+path = "lib/BracketingNonlinearSolve"
+
+[sources.NonlinearSolveBase]
+path = "lib/NonlinearSolveBase"
+
+[sources.NonlinearSolveFirstOrder]
+path = "lib/NonlinearSolveFirstOrder"
+
+[sources.NonlinearSolveQuasiNewton]
+path = "lib/NonlinearSolveQuasiNewton"
+
+[sources.NonlinearSolveSpectralMethods]
+path = "lib/NonlinearSolveSpectralMethods"
+
+[sources.SimpleNonlinearSolve]
+path = "lib/SimpleNonlinearSolve"
+
+[sources]
+DiffEqBase = {url = "https://github.com/jClugstor/DiffEqBase.jl", rev = "remove_nonlinear"}
+
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SparseMatrixColorings", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote"]
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "LineSearches", "MINPACK", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "PETSc", "Pkg", "PolyesterForwardDiff", "Random", "ReTestItems", "SIAMFANLEquations", "SparseConnectivityTracer", "SparseMatrixColorings", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Test", "Zygote", "ReverseDiff", "Tracker", "DiffEqBase"]

lib/NonlinearSolveBase/Project.toml

@@ -22,40 +22,53 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
-DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
+Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [extensions]
 NonlinearSolveBaseBandedMatricesExt = "BandedMatrices"
-NonlinearSolveBaseDiffEqBaseExt = "DiffEqBase"
+NonlinearSolveBaseChainRulesCoreExt = "ChainRulesCore"
+NonlinearSolveBaseEnzymeExt = ["ChainRulesCore", "Enzyme"]
 NonlinearSolveBaseForwardDiffExt = "ForwardDiff"
 NonlinearSolveBaseLineSearchExt = "LineSearch"
 NonlinearSolveBaseLinearSolveExt = "LinearSolve"
+NonlinearSolveBaseMooncakeExt = "Mooncake"
+NonlinearSolveBaseReverseDiffExt = "ReverseDiff"
 NonlinearSolveBaseSparseArraysExt = "SparseArrays"
 NonlinearSolveBaseSparseMatrixColoringsExt = "SparseMatrixColorings"
+NonlinearSolveBaseTrackerExt = "Tracker"
+
 
 [compat]
 ADTypes = "1.9"
 Adapt = "4.1.0"
 Aqua = "0.8.7"
 ArrayInterface = "7.9"
 BandedMatrices = "1.5"
+ChainRulesCore = "1"
 CommonSolve = "0.2.4"
 Compat = "4.15"
 ConcreteStructs = "0.2.3"
 DiffEqBase = "6.188"
 DifferentiationInterface = "0.7.3"
 EnzymeCore = "0.8"
+Enzyme = "0.13.12"
 ExplicitImports = "1.10.1"
 FastClosures = "0.3"
 ForwardDiff = "0.10.36, 1"
@@ -65,20 +78,28 @@ LinearAlgebra = "1.10"
 LinearSolve = "3.15"
 Markdown = "1.10"
 MaybeInplace = "0.1.4"
+Mooncake = "0.4"
 Preferences = "1.4"
 Printf = "1.10"
 RecursiveArrayTools = "3"
+ReverseDiff = "1.15"
 SciMLBase = "2.116"
 SciMLJacobianOperators = "0.1.1"
 SciMLOperators = "1.7"
+SciMLStructures = "1.5"
+Setfield = "1.1.2"
 SparseArrays = "1.10"
 SparseMatrixColorings = "0.4.5"
 StaticArraysCore = "1.4"
 SymbolicIndexingInterface = "0.3.43"
 Test = "1.10"
+Tracker = "0.2.35"
 TimerOutputs = "0.5.23"
 julia = "1.10"
 
+[sources]
+DiffEqBase = {url = "https://github.com/jClugstor/DiffEqBase.jl", rev = "remove_nonlinear"}
+
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"

lib/NonlinearSolveBase/ext/NonlinearSolveBaseChainRulesCoreExt.jl

@@ -0,0 +1,31 @@
+module NonlinearSolveBaseChainRulesCoreExt
+
+using NonlinearSolveBase
+using NonlinearSolveBase: AbstractNonlinearProblem
+using SciMLBase
+using SciMLBase: AbstractSensitivityAlgorithm
+
+import ChainRulesCore
+import ChainRulesCore: NoTangent
+
+function ChainRulesCore.frule(::typeof(NonlinearSolveBase.solve_up), prob,
+        sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
+        u0, p, args...; originator = SciMLBase.ChainRulesOriginator(),
+        kwargs...)
+    NonlinearSolveBase._solve_forward(
+        prob, sensealg, u0, p,
+        originator, args...;
+        kwargs...)
+end
+
+function ChainRulesCore.rrule(::typeof(NonlinearSolveBase.solve_up), prob::AbstractNonlinearProblem,
+        sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
+        u0, p, args...; originator = SciMLBase.ChainRulesOriginator(),
+        kwargs...)
+    NonlinearSolveBase._solve_adjoint(
+        prob, sensealg, u0, p,
+        originator, args...;
+        kwargs...)
+end
+
+end

lib/NonlinearSolveBase/ext/NonlinearSolveBaseEnzymeExt.jl

@@ -0,0 +1,61 @@
+module NonlinearSolveBaseEnzymeExt
+
+@static if isempty(VERSION.prerelease)
+    using NonlinearSolveBase
+    import SciMLBase: SciMLBase, value
+    using Enzyme
+    import Enzyme: Const
+    using ChainRulesCore
+
+    function Enzyme.EnzymeRules.augmented_primal(
+            config::Enzyme.EnzymeRules.RevConfigWidth{1},
+            func::Const{typeof(NonlinearSolveBase.solve_up)}, ::Type{Duplicated{RT}}, prob,
+            sensealg::Union{
+                Const{Nothing}, Const{<:SciMLBase.AbstractSensitivityAlgorithm}},
+            u0, p, args...; kwargs...) where {RT}
+        @inline function copy_or_reuse(val, idx)
+            if Enzyme.EnzymeRules.overwritten(config)[idx] && ismutable(val)
+                return deepcopy(val)
+            else
+                return val
+            end
+        end
+
+        @inline function arg_copy(i)
+            copy_or_reuse(args[i].val, i + 5)
+        end
+
+        res = NonlinearSolveBase._solve_adjoint(
+            copy_or_reuse(prob.val, 2), copy_or_reuse(sensealg.val, 3),
+            copy_or_reuse(u0.val, 4), copy_or_reuse(p.val, 5),
+            SciMLBase.EnzymeOriginator(), ntuple(arg_copy, Val(length(args)))...;
+            kwargs...)
+
+        dres = Enzyme.make_zero(res[1])::RT
+        tup = (dres, res[2])
+        return Enzyme.EnzymeRules.AugmentedReturn{RT, RT, Any}(res[1], dres, tup::Any)
+    end
+
+    function Enzyme.EnzymeRules.reverse(config::Enzyme.EnzymeRules.RevConfigWidth{1},
+            func::Const{typeof(NonlinearSolveBase.solve_up)}, ::Type{Duplicated{RT}}, tape, prob,
+            sensealg::Union{
+                Const{Nothing}, Const{<:SciMLBase.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 Enzyme.Const
+                continue
+            end
+            if darg == ChainRulesCore.NoTangent()
+                continue
+            end
+            ptr.dval .+= darg
+        end
+        Enzyme.make_zero!(dres.u)
+        return ntuple(_ -> nothing, Val(length(args) + 4))
+    end
+end
+
+end

lib/NonlinearSolveBase/ext/NonlinearSolveBaseMooncakeExt.jl

@@ -0,0 +1,31 @@
+module NonlinearSolveBaseMooncakeExt
+
+using NonlinearSolveBase, Mooncake
+using SciMLBase: SciMLBase
+import Mooncake: rrule!!, CoDual, zero_fcodual, @is_primitive,
+                 @from_rrule, @zero_adjoint, @mooncake_overlay, MinimalCtx,
+                 NoPullback
+
+@from_rrule(MinimalCtx,
+    Tuple{
+        typeof(NonlinearSolveBase.solve_up),
+        SciMLBase.AbstractDEProblem,
+        Union{Nothing, SciMLBase.AbstractSensitivityAlgorithm},
+        Any,
+        Any,
+        Any
+    },
+    true,)
+
+# Dispatch for auto-alg
+@from_rrule(MinimalCtx,
+    Tuple{
+        typeof(NonlinearSolveBase.solve_up),
+        SciMLBase.AbstractDEProblem,
+        Union{Nothing, SciMLBase.AbstractSensitivityAlgorithm},
+        Any,
+        Any
+    },
+    true,)
+
+end