fix: auto-set autodiff for ForwardDiff if trying to propagate Duals

SciML · Oct 8, 2024 · 8cf9899 · 8cf9899
1 parent 244d7bb
commit 8cf9899
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 5 deletions.
diff --git a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl
@@ -54,6 +54,9 @@ function CommonSolve.solve(
         alg::AbstractSimpleNonlinearSolveAlgorithm,
         args...;
         kwargs...) where {T, V, P, iip}
+    if hasfield(typeof(alg), :autodiff) && alg.autodiff === nothing
+        @reset alg.autodiff = AutoForwardDiff()
+    end
     prob = convert(ImmutableNonlinearProblem, prob)
     sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...)
     dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p)
@@ -68,6 +71,9 @@ function CommonSolve.solve(
         alg::AbstractSimpleNonlinearSolveAlgorithm,
         args...;
         kwargs...) where {T, V, P, iip}
+    if hasfield(typeof(alg), :autodiff) && alg.autodiff === nothing
+        @reset alg.autodiff = AutoForwardDiff()
+    end
     sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...)
     dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p)
     return SciMLBase.build_solution(

diff --git a/lib/SimpleNonlinearSolve/test/core/forward_diff_tests.jl b/lib/SimpleNonlinearSolve/test/core/forward_diff_tests.jl
@@ -14,7 +14,7 @@
     jacobian_f(u, p::Number) = one.(u) .* (1 / (2 * √p))
     jacobian_f(u, p::AbstractArray) = diagm(vec(@. 1 / (2 * √p)))
 
-    @testset for alg in (
+    @testset "#(nameof(typeof(alg)))" for alg in (
         SimpleNewtonRaphson(),
         SimpleTrustRegion(),
         SimpleTrustRegion(; nlsolve_update_rule = Val(true)),
@@ -118,7 +118,7 @@ end
 
     θ_init = θ_true .+ 0.1
 
-    @testset for alg in (
+    for alg in (
         SimpleGaussNewton(),
         SimpleGaussNewton(; autodiff = AutoForwardDiff()),
         SimpleGaussNewton(; autodiff = AutoFiniteDiff()),

diff --git a/lib/SimpleNonlinearSolve/test/core/qa_tests.jl b/lib/SimpleNonlinearSolve/test/core/qa_tests.jl
@@ -12,8 +12,7 @@ end
     import ReverseDiff, Tracker, StaticArrays, Zygote
     using ExplicitImports, SimpleNonlinearSolve
 
-    @test check_no_implicit_imports(
-        SimpleNonlinearSolve; skip = (Base, Core, SciMLBase)) === nothing
+    @test check_no_implicit_imports(SimpleNonlinearSolve; skip = (Base, Core)) === nothing
     @test check_no_stale_explicit_imports(SimpleNonlinearSolve) === nothing
     @test check_all_qualified_accesses_via_owners(SimpleNonlinearSolve) === nothing
 end
diff --git a/lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl b/lib/SimpleNonlinearSolve/test/core/rootfind_tests.jl
@@ -36,7 +36,7 @@
 end
 
 @testitem "First Order Methods" setup=[RootfindTestSnippet] tags=[:core] begin
-    @testset for alg in (
+    for alg in (
         SimpleNewtonRaphson,
         SimpleTrustRegion,
         (; kwargs...) -> SimpleTrustRegion(; kwargs..., nlsolve_update_rule = Val(true))