add the ChainRulesCore extension

Author: jClugstor

lib/NonlinearSolveBase/Project.toml

@@ -29,6 +29,7 @@ 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"
@@ -44,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"

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