Migrate solve dispatches from DiffEqBase to NonlinearSolveBase

lib/NonlinearSolveBase/Project.toml

@@ -22,12 +22,14 @@ 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"
 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"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b"
@@ -43,13 +45,15 @@ NonlinearSolveBaseLineSearchExt = "LineSearch"
 NonlinearSolveBaseLinearSolveExt = "LinearSolve"
 NonlinearSolveBaseSparseArraysExt = "SparseArrays"
 NonlinearSolveBaseSparseMatrixColoringsExt = "SparseMatrixColorings"
+NonlinearSolveBaseChainRulesCoreExt = "ChainRulesCore"
 
 [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"
@@ -71,6 +75,7 @@ RecursiveArrayTools = "3"
 SciMLBase = "2.92"
 SciMLJacobianOperators = "0.1.1"
 SciMLOperators = "0.4, 1.0"
+SciMLStructures = "1.5"
 SparseArrays = "1.10"
 SparseMatrixColorings = "0.4.5"
 StaticArraysCore = "1.4"

lib/NonlinearSolveBase/ext/NonlinearSolveBaseChainRulesCoreExt.jl

@@ -0,0 +1,33 @@
+module NonlinearSolveBaseChainRulesCoreExt
+
+using NonlinearSolveBase
+using NonlinearSolveBase: AbstractNonlinearProblem
+using SciMLBase
+using SciMLBase: AbstractSensitivityAlgorithm
+
+import ChainRulesCore
+import ChainRulesCore: NoTangent
+
+ChainRulesCore.@non_differentiable NonlinearSolveBase.checkkwargs(kwargshandle)
+
+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/NonlinearSolveBaseForwardDiffExt.jl

@@ -195,4 +195,13 @@ NonlinearSolveBase.nodual_value(x::AbstractArray{<:Dual}) = map(ForwardDiff.valu
 @inline NonlinearSolveBase.pickchunksize(x) = pickchunksize(length(x))
 @inline NonlinearSolveBase.pickchunksize(x::Int) = ForwardDiff.pickchunksize(x)
 
+eltypedual(x) = eltype(x) <: ForwardDiff.Dual
+isdualtype(::Type{<:ForwardDiff.Dual}) = true
+
+function anyeltypedual(
+        prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem},
+        ::Type{Val{counter}} = Val{0}) where {counter}
+    anyeltypedual((prob.u0, prob.p))
+end
+
 end

lib/NonlinearSolveBase/src/NonlinearSolveBase.jl

@@ -15,15 +15,17 @@ using StaticArraysCore: StaticArray, SMatrix, SArray, MArray
 using CommonSolve: CommonSolve, init
 using EnzymeCore: EnzymeCore
 using MaybeInplace: @bb
-using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition
+using RecursiveArrayTools: RecursiveArrayTools, AbstractVectorOfArray, ArrayPartition
 using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinearProblem,
                  AbstractNonlinearAlgorithm,
                  NonlinearProblem, NonlinearLeastSquaresProblem,
                  NonlinearFunction, NLStats, LinearProblem,
-                 LinearAliasSpecifier, ImmutableNonlinearProblem
+                 LinearAliasSpecifier, ImmutableNonlinearProblem, NonlinearAliasSpecifier
+import SciMLBase: solve, init, solve!, __init, __solve, wrap_sol, get_root_indp, isinplace, remake
 using SciMLJacobianOperators: JacobianOperator, StatefulJacobianOperator
 using SciMLOperators: AbstractSciMLOperator, IdentityOperator
 using SymbolicIndexingInterface: SymbolicIndexingInterface
+import SciMLStructures
 
 using LinearAlgebra: LinearAlgebra, Diagonal, norm, ldiv!, diagind, mul!
 using Markdown: @doc_str
@@ -32,6 +34,11 @@ using Printf: @printf
 const DI = DifferentiationInterface
 const SII = SymbolicIndexingInterface
 
+# Extension Functions
+eltypedual(x) = false
+promote_u0(::Nothing, p, t0) = nothing
+isdualtype(::Type{T}) where {T} = false
+
 include("public.jl")
 include("utils.jl")