Add inplace and cached ForwardDiff rules

Author: Avik Pal

src/ad.jl

@@ -4,20 +4,78 @@ function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, <:AbstractArray}
         kwargs...) where {T, V, P, iip}
     sol, partials = __nlsolve_ad(prob, alg, args...; kwargs...)
     dual_soln = __nlsolve_dual_soln(sol.u, partials, prob.p)
-    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode)
+    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats,
+        sol.original)
 end
 
-# Differentiate Out-of-Place Nonlinear Root Finding Problems
-function __nlsolve_ad(prob::NonlinearProblem{uType, false}, alg, args...;
-        kwargs...) where {uType}
+@concrete mutable struct NonlinearSolveForwardDiffCache
+    cache
+    prob
+    alg
+    p
+    values_p
+    partials_p
+end
+
+@inline function __has_duals(::Union{<:Dual{T, V, P},
+        <:AbstractArray{<:Dual{T, V, P}}}) where {T, V, P}
+    return true
+end
+@inline __has_duals(::Any) = false
+
+function SciMLBase.reinit!(cache::NonlinearSolveForwardDiffCache; p = cache.p,
+        u0 = get_u(cache.cache), kwargs...)
+    inner_cache = SciMLBase.reinit!(cache.cache; p = value(p), u0 = value(u0), kwargs...)
+    cache.cache = inner_cache
+    cache.p = p
+    cache.values_p = value(p)
+    cache.partials_p = ForwardDiff.partials(p)
+    return cache
+end
+
+function SciMLBase.init(prob::NonlinearProblem{<:Union{Number, <:AbstractArray},
+            iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}},
+        alg::Union{Nothing, AbstractNonlinearAlgorithm}, args...;
+        kwargs...) where {T, V, P, iip}
+    p = value(prob.p)
+    newprob = NonlinearProblem(prob.f, value(prob.u0), p; prob.kwargs...)
+    cache = init(newprob, alg, args...; kwargs...)
+    return NonlinearSolveForwardDiffCache(cache, newprob, alg, prob.p, p,
+        ForwardDiff.partials(prob.p))
+end
+
+function SciMLBase.solve!(cache::NonlinearSolveForwardDiffCache)
+    sol = solve!(cache.cache)
+    prob = cache.prob
+
+    uu = sol.u
+    f_p = __nlsolve_∂f_∂p(prob, prob.f, uu, cache.values_p)
+    f_x = __nlsolve_∂f_∂u(prob, prob.f, uu, cache.values_p)
+
+    z_arr = -f_x \ f_p
+
+    sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z)
+    if cache.p isa Number
+        partials = sumfun((z_arr, cache.p))
+    else
+        partials = sum(sumfun, zip(eachcol(z_arr), cache.p))
+    end
+
+    dual_soln = __nlsolve_dual_soln(sol.u, partials, cache.p)
+    return SciMLBase.build_solution(prob, cache.alg, dual_soln, sol.resid; sol.retcode,
+        sol.stats, sol.original)
+end
+
+function __nlsolve_ad(prob::NonlinearProblem{uType, iip}, alg, args...;
+        kwargs...) where {uType, iip}
     p = value(prob.p)
     newprob = NonlinearProblem(prob.f, value(prob.u0), p; prob.kwargs...)
 
     sol = solve(newprob, alg, args...; kwargs...)
 
     uu = sol.u
-    f_p = __nlsolve_∂f_∂p(prob.f, uu, p)
-    f_x = __nlsolve_∂f_∂u(prob.f, uu, p)
+    f_p = __nlsolve_∂f_∂p(prob, prob.f, uu, p)
+    f_x = __nlsolve_∂f_∂u(prob, prob.f, uu, p)
 
     z_arr = -f_x \ f_p
 
@@ -34,8 +92,16 @@ function __nlsolve_ad(prob::NonlinearProblem{uType, false}, alg, args...;
     return sol, partials
 end
 
-@inline function __nlsolve_∂f_∂p(f::F, u, p) where {F}
-    __f = Base.Fix1(f, u)
+@inline function __nlsolve_∂f_∂p(prob, f::F, u, p) where {F}
+    if isinplace(prob)
+        __f = p -> begin
+            du = similar(u, promote_type(eltype(u), eltype(p)))
+            f(du, u, p)
+            return du
+        end
+    else
+        __f = Base.Fix1(f, u)
+    end
     if p isa Number
         return __reshape(ForwardDiff.derivative(__f, p), :, 1)
     elseif u isa Number
@@ -45,12 +111,18 @@ end
     end
 end
 
-@inline function __nlsolve_∂f_∂u(f::F, u, p) where {F}
-    __f = Base.Fix2(f, p)
-    if u isa Number
-        return ForwardDiff.derivative(__f, u)
+@inline function __nlsolve_∂f_∂u(prob, f::F, u, p) where {F}
+    if isinplace(prob)
+        du = similar(u)
+        __f = (du, u) -> f(du, u, p)
+        ForwardDiff.jacobian(__f, du, u)
     else
-        return ForwardDiff.jacobian(__f, u)
+        __f = Base.Fix2(f, p)
+        if u isa Number
+            return ForwardDiff.derivative(__f, u)
+        else
+            return ForwardDiff.jacobian(__f, u)
+        end
     end
 end
 

test/forward_ad.jl

@@ -8,18 +8,33 @@ jacobian_f(::Number, p::Number) = 1 / (2 * √p)
 jacobian_f(u, p::Number) = one.(u) .* (1 / (2 * √p))
 jacobian_f(u, p::AbstractArray) = diagm(vec(@. 1 / (2 * √p)))
 
-function solve_with(::Val{iip}, u, alg) where {iip}
-    f = if iip
+function solve_with(::Val{mode}, u, alg) where {mode}
+    f = if mode === :iip
         solve_iip(p) = solve(NonlinearProblem(test_f!, u, p), alg).u
-    else
+    elseif mode === :iip_cache
+        function solve_iip_init(p)
+            cache = SciMLBase.init(NonlinearProblem(test_f!, u, p), alg)
+            return SciMLBase.solve!(cache).u
+        end
+    elseif mode === :oop
         solve_oop(p) = solve(NonlinearProblem(test_f, u, p), alg).u
+    elseif mode === :oop_cache
+        function solve_oop_init(p)
+            cache = SciMLBase.init(NonlinearProblem(test_f, u, p), alg)
+            return SciMLBase.solve!(cache).u
+        end
     end
     return f
 end
 
-__can_inplace(::Number) = false
-__can_inplace(::AbstractArray) = true
-__can_inplace(::StaticArray) = false
+__compatible(::Any, ::Val{:oop}) = true
+__compatible(::Any, ::Val{:oop_cache}) = true
+__compatible(::Number, ::Val{:iip}) = false
+__compatible(::AbstractArray, ::Val{:iip}) = true
+__compatible(::StaticArray, ::Val{:iip}) = false
+__compatible(::Number, ::Val{:iip_cache}) = false
+__compatible(::AbstractArray, ::Val{:iip_cache}) = true
+__compatible(::StaticArray, ::Val{:iip_cache}) = false
 
 __compatible(::Any, ::Number) = true
 __compatible(::Number, ::AbstractArray) = false
@@ -32,37 +47,49 @@ __compatible(u::StaticArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true
 __compatible(u::StaticArray, ::Union{CMINPACK, NLsolveJL}) = false
 __compatible(u, ::Nothing) = true
 
+__compatible(::Any, ::Any) = true
+__compatible(::CMINPACK, ::Val{:iip_cache}) = false
+__compatible(::CMINPACK, ::Val{:oop_cache}) = false
+__compatible(::NLsolveJL, ::Val{:iip_cache}) = false
+__compatible(::NLsolveJL, ::Val{:oop_cache}) = false
+
 @testset "ForwardDiff.jl Integration: $(alg)" for alg in (NewtonRaphson(), TrustRegion(),
     LevenbergMarquardt(), PseudoTransient(; alpha_initial = 10.0), Broyden(), Klement(),
     DFSane(), nothing, NLsolveJL(), CMINPACK())
     us = (2.0, @SVector[1.0, 1.0], [1.0, 1.0], ones(2, 2), @SArray ones(2, 2))
 
     @testset "Scalar AD" begin
-        for p in 1.0:0.1:100.0
-            for u0 in us
-                __compatible(u0, alg) || continue
-                sol = solve(NonlinearProblem(test_f, u0, p), alg)
-                if SciMLBase.successful_retcode(sol)
-                    gs = abs.(ForwardDiff.derivative(solve_with(Val{false}(), u0, alg), p))
-                    gs_true = abs.(jacobian_f(u0, p))
-                    if !(isapprox(gs, gs_true, atol = 1e-5))
-                        @show sol.retcode, sol.u
-                        @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_gradient=gs true_gradient=gs_true
-                    else
-                        @test abs.(gs)≈abs.(gs_true) atol=1e-5
-                    end
+        for p in 1.0:0.1:100.0, u0 in us, mode in (:iip, :oop, :iip_cache, :oop_cache)
+            __compatible(u0, alg) || continue
+            __compatible(u0, Val(mode)) || continue
+            __compatible(alg, Val(mode)) || continue
+
+            sol = solve(NonlinearProblem(test_f, u0, p), alg)
+            if SciMLBase.successful_retcode(sol)
+                gs = abs.(ForwardDiff.derivative(solve_with(Val{mode}(), u0, alg), p))
+                gs_true = abs.(jacobian_f(u0, p))
+                if !(isapprox(gs, gs_true, atol = 1e-5))
+                    @show sol.retcode, sol.u
+                    @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_gradient=gs true_gradient=gs_true
+                else
+                    @test abs.(gs)≈abs.(gs_true) atol=1e-5
                 end
             end
         end
     end
 
     @testset "Jacobian" begin
-        for u0 in us, p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0])
+        for u0 in us, p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0]),
+            mode in (:iip, :oop, :iip_cache, :oop_cache)
+
             __compatible(u0, p) || continue
             __compatible(u0, alg) || continue
+            __compatible(u0, Val(mode)) || continue
+            __compatible(alg, Val(mode)) || continue
+
             sol = solve(NonlinearProblem(test_f, u0, p), alg)
             if SciMLBase.successful_retcode(sol)
-                gs = abs.(ForwardDiff.jacobian(solve_with(Val{false}(), u0, alg), p))
+                gs = abs.(ForwardDiff.jacobian(solve_with(Val{mode}(), u0, alg), p))
                 gs_true = abs.(jacobian_f(u0, p))
                 if !(isapprox(gs, gs_true, atol = 1e-5))
                     @show sol.retcode, sol.u