fix type stability of SimpleNonlinearSolve

lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl

@@ -39,6 +39,7 @@ const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{
 } where {iip, T, V, P}
 
 abstract type AbstractSimpleNonlinearSolveAlgorithm <: AbstractNonlinearSolveAlgorithm end
+configure_autodiff(prob, alg::AbstractSimpleNonlinearSolveAlgorithm) = alg
 
 const NLBUtils = NonlinearSolveBase.Utils
 
@@ -59,12 +60,6 @@ function CommonSolve.solve(
         prob::NonlinearProblem, alg::AbstractSimpleNonlinearSolveAlgorithm, args...;
         kwargs...
 )
-    cache = SciMLBase.__init(prob, alg, args...; kwargs...)
-    prob = cache.prob
-    if cache.retcode == ReturnCode.InitialFailure
-        return SciMLBase.build_solution(prob, alg, prob.u0,
-            NonlinearSolveBase.Utils.evaluate_f(prob, prob.u0); cache.retcode)
-    end
     prob = convert(ImmutableNonlinearProblem, prob)
     return solve(prob, alg, args...; kwargs...)
 end
@@ -73,9 +68,7 @@ function CommonSolve.solve(
         prob::DualNonlinearProblem, alg::AbstractSimpleNonlinearSolveAlgorithm,
         args...; kwargs...
 )
-    if hasfield(typeof(alg), :autodiff) && alg.autodiff === nothing
-        @set! alg.autodiff = AutoForwardDiff()
-    end
+    alg = configure_autodiff(prob, alg)
     prob = convert(ImmutableNonlinearProblem, prob)
     sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...)
     dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p)
@@ -88,9 +81,7 @@ function CommonSolve.solve(
         prob::DualNonlinearLeastSquaresProblem, alg::AbstractSimpleNonlinearSolveAlgorithm,
         args...; kwargs...
 )
-    if hasfield(typeof(alg), :autodiff) && alg.autodiff === nothing
-        @set! alg.autodiff = AutoForwardDiff()
-    end
+    alg = configure_autodiff(prob, alg)
     sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...)
     dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p)
     return SciMLBase.build_solution(
@@ -103,6 +94,7 @@ function CommonSolve.solve(
         alg::AbstractSimpleNonlinearSolveAlgorithm,
         args...; sensealg = nothing, u0 = nothing, p = nothing, kwargs...
 )
+    alg = configure_autodiff(prob, alg)
     cache = SciMLBase.__init(prob, alg, args...; kwargs...)
     prob = cache.prob
     if cache.retcode == ReturnCode.InitialFailure

lib/SimpleNonlinearSolve/src/halley.jl

@@ -20,11 +20,18 @@ A low-overhead implementation of Halley's Method.
     autodiff = nothing
 end
 
+function configure_autodiff(prob, alg::SimpleHalley)
+    # The way we write the 2nd order derivatives, we know Enzyme won't work there
+    @set! alg.autodiff = something(alg.autodiff, AutoForwardDiff())
+    alg
+end
+
 function SciMLBase.__solve(
         prob::ImmutableNonlinearProblem, alg::SimpleHalley, args...;
         abstol = nothing, reltol = nothing, maxiters = 1000,
         alias_u0 = false, termination_condition = nothing, kwargs...
 )
+    autodiff = alg.autodiff
     x = NLBUtils.maybe_unaliased(prob.u0, alias_u0)
     fx = NLBUtils.evaluate_f(prob, x)
     T = promote_type(eltype(fx), eltype(x))
@@ -36,10 +43,6 @@ function SciMLBase.__solve(
         prob, abstol, reltol, fx, x, termination_condition, Val(:simple)
     )
 
-    # The way we write the 2nd order derivatives, we know Enzyme won't work there
-    autodiff = alg.autodiff === nothing ? AutoForwardDiff() : alg.autodiff
-    @set! alg.autodiff = autodiff
-
     @bb xo = copy(x)
 
     if NLBUtils.can_setindex(x)

lib/SimpleNonlinearSolve/src/raphson.jl

@@ -23,12 +23,21 @@ end
 
 const SimpleGaussNewton = SimpleNewtonRaphson
 
+function configure_autodiff(prob, alg::SimpleNewtonRaphson)
+    autodiff = something(alg.autodiff, AutoForwardDiff())
+    autodiff = SciMLBase.has_jac(prob.f) ? autodiff :
+               NonlinearSolveBase.select_jacobian_autodiff(prob, autodiff)
+    @set! alg.autodiff = autodiff
+    alg
+end
+
 function SciMLBase.__solve(
         prob::Union{ImmutableNonlinearProblem, NonlinearLeastSquaresProblem},
         alg::SimpleNewtonRaphson, args...;
         abstol = nothing, reltol = nothing, maxiters = 1000,
         alias_u0 = false, termination_condition = nothing, kwargs...
 )
+    autodiff = alg.autodiff
     x = NLBUtils.maybe_unaliased(prob.u0, alias_u0)
     fx = NLBUtils.evaluate_f(prob, x)
 
@@ -39,10 +48,6 @@ function SciMLBase.__solve(
         prob, abstol, reltol, fx, x, termination_condition, Val(:simple)
     )
 
-    autodiff = SciMLBase.has_jac(prob.f) ? alg.autodiff :
-               NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff)
-    @set! alg.autodiff = autodiff
-
     @bb xo = similar(x)
     fx_cache = (SciMLBase.isinplace(prob) && !SciMLBase.has_jac(prob.f)) ?
                NLBUtils.safe_similar(fx) : fx

lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl

@@ -28,10 +28,10 @@
     ]
 
     function run_nlsolve_oop(f::F, u0, p = 2.0; solver) where {F}
-        return solve(NonlinearProblem{false}(f, u0, p), solver; abstol = 1e-9)
+        return @inferred solve(NonlinearProblem{false}(f, u0, p), solver; abstol = 1e-9)
     end
     function run_nlsolve_iip(f!::F, u0, p = 2.0; solver) where {F}
-        return solve(NonlinearProblem{true}(f!, u0, p), solver; abstol = 1e-9)
+        return @inferred solve(NonlinearProblem{true}(f!, u0, p), solver; abstol = 1e-9)
     end
 end