Use LinearSolve for ForwardDiff implicit sensitivities (#1055)

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>

Author: ChrisRackauckas-Claude

lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl

@@ -178,7 +178,7 @@ function NonlinearSolveBase.nonlinearsolve_forwarddiff_solve(
 
     Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(ad_prob, fn, uu, p)
     Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(ad_prob, fn, uu, p)
-    z = -Jᵤ \ Jₚ
+    z = -_forwarddiff_implicit_solve(Jᵤ, Jₚ)
     pp = prob.p
     sumfun = ((z, p),) -> map(Base.Fix2(*, ForwardDiff.partials(p)), z)
 
@@ -222,6 +222,22 @@ function NonlinearSolveBase.nonlinearsolve_∂f_∂u(prob, f::F, u, p) where {F}
     return ForwardDiff.jacobian(Base.Fix2(f, p), u)
 end
 
+_forwarddiff_implicit_solve(A::Number, B) = A \ B
+
+function _forwarddiff_implicit_solve(A, B)
+    Utils.is_extension_loaded(Val(:LinearSolve)) || return A \ B
+
+    T = promote_type(typeof(oneunit(eltype(A)) / oneunit(eltype(A))), eltype(B))
+    u = similar(B, T, size(B))
+    lincache = NonlinearSolveBase.construct_linear_solver(
+        nothing, nothing, A, B, u, nothing;
+        stats = SciMLBase.NLStats(0, 0, 0, 0, 0), verbose = false
+    )
+    linres = lincache()
+    linres.success || error("Linear solve failed while differentiating a nonlinear solve.")
+    return linres.u
+end
+
 function NonlinearSolveBase.nonlinearsolve_dual_solution(
         u::Number, partials,
         ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}
@@ -302,7 +318,7 @@ function CommonSolve.solve!(cache::NonlinearSolveForwardDiffCache)
     Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(ad_prob, fn, uu, cache.values_p)
     Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(ad_prob, fn, uu, cache.values_p)
 
-    z_arr = -Jᵤ \ Jₚ
+    z_arr = -_forwarddiff_implicit_solve(Jᵤ, Jₚ)
 
     sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z)
     if cache.p isa Number

test/Core/forward_ad_tests__item2.jl

@@ -1,5 +1,6 @@
 using NonlinearSolve
 
+using SciMLBase
 using ForwardDiff, FiniteDiff
 
 function objfn(F, init, params)
@@ -107,3 +108,45 @@ hess1 = ForwardDiff.hessian(solve_nlprob_with_cache, [34.0, 87.0])
 hess2 = FiniteDiff.finite_difference_hessian(solve_nlprob_with_cache, [34.0, 87.0])
 
 @test hess1 ≈ hess2 atol = 1.0e-3
+
+function singular_nlls_resid!(resid, u, p)
+    resid[1] = u[1] - p[1]
+    resid[2] = 0
+    return nothing
+end
+
+function singular_nlls_solution(p)
+    prob = NonlinearLeastSquaresProblem(
+        NonlinearFunction(singular_nlls_resid!, resid_prototype = zeros(2)),
+        [1.0, 2.0], p
+    )
+    return solve(prob, NewtonRaphson(); abstol = 1.0e-12, reltol = 1.0e-12)
+end
+
+singular_nlls_sum(p) = sum(singular_nlls_solution(p).u)
+
+function singular_nlls_cached_solution(p)
+    prob = NonlinearLeastSquaresProblem(
+        NonlinearFunction(singular_nlls_resid!, resid_prototype = zeros(2)),
+        [1.0, 2.0], p
+    )
+    cache = init(prob, NewtonRaphson(); abstol = 1.0e-12, reltol = 1.0e-12)
+    return solve!(cache)
+end
+
+singular_nlls_cached_sum(p) = sum(singular_nlls_cached_solution(p).u)
+
+@test SciMLBase.successful_retcode(singular_nlls_solution([1.0]).retcode)
+@test SciMLBase.successful_retcode(singular_nlls_cached_solution([1.0]).retcode)
+@test ForwardDiff.gradient(singular_nlls_sum, [1.0]) ≈ [1.0]
+@test ForwardDiff.gradient(singular_nlls_cached_sum, [1.0]) ≈ [1.0]
+
+function scalar_nl_solution(p)
+    prob = NonlinearProblem((u, p) -> u^2 - p[1] - p[2], 2.0, p)
+    return solve(prob, NewtonRaphson(); abstol = 1.0e-12, reltol = 1.0e-12)
+end
+
+scalar_nl_sum(p) = scalar_nl_solution(p).u
+
+@test SciMLBase.successful_retcode(scalar_nl_solution([4.0, 0.0]).retcode)
+@test ForwardDiff.gradient(scalar_nl_sum, [4.0, 0.0]) ≈ [0.25, 0.25]