fix: reinit! on forwarddiff cache (#491)

SciML · Nov 1, 2024 · 3ed48f8 · 3ed48f8 · avik-pal · Nov 3, 2024
1 parent 5c722c0
commit 3ed48f8
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/src/forward_diff.jl b/src/forward_diff.jl
@@ -37,10 +37,9 @@ function InternalAPI.reinit!(
         cache::NonlinearSolveForwardDiffCache, args...;
         p = cache.p, u0 = NonlinearSolveBase.get_u(cache.cache), kwargs...
 )
-    inner_cache = InternalAPI.reinit!(
+    InternalAPI.reinit!(
         cache.cache; p = nodual_value(p), u0 = nodual_value(u0), kwargs...
     )
-    cache.cache = inner_cache
     cache.p = p
     cache.values_p = nodual_value(p)
     cache.partials_p = ForwardDiff.partials(p)

diff --git a/test/forward_ad_tests.jl b/test/forward_ad_tests.jl
@@ -218,3 +218,34 @@ end
 
     @test hess1≈hess2 atol=1e-3
 end
+
+@testitem "reinit! on ForwardDiff cache SciML/NonlinearSolve.jl#391" tags=[:core] begin
+    using ForwardDiff
+
+    function multiple_solves(ps::Vector)
+        res = similar(ps, 4, length(ps))
+        for (i, p) in enumerate(ps)
+            prob = NonlinearProblem{false}((u, p) -> u .* u .- p, rand(4), ps[i])
+            sol = solve(prob)
+            res[:, i] .= sol.u
+        end
+        return sum(abs2, res)
+    end
+
+    function multiple_solves_cached(ps::Vector)
+        res = similar(ps, 4, length(ps))
+        prob = NonlinearProblem{false}((u, p) -> u .* u .- p, rand(4), ps[1])
+        cache = init(prob, NewtonRaphson())
+        for (i, p) in enumerate(ps)
+            reinit!(cache; p)
+            sol = solve!(cache)
+            res[:, i] .= sol.u
+        end
+        return sum(abs2, res)
+    end
+
+    ps = collect(1.0:5.0)
+
+    @test ForwardDiff.gradient(multiple_solves, ps) ≈
+          ForwardDiff.gradient(multiple_solves_cached, ps)
+end