fix: extension for forward AD support

SciML · Sep 17, 2024 · 8fadcc4 · 8fadcc4
1 parent 0b70d62
commit 8fadcc4
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 5 deletions.
diff --git a/lib/BracketingNonlinearSolve/Project.toml b/lib/BracketingNonlinearSolve/Project.toml
@@ -14,7 +14,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [extensions]
-NonlinearSolveBaseForwardDiffExt = "ForwardDiff"
+BracketingNonlinearSolveForwardDiffExt = "ForwardDiff"
 
 [compat]
 CommonSolve = "0.2.4"

diff --git a/lib/BracketingNonlinearSolve/test/rootfind_tests.jl b/lib/BracketingNonlinearSolve/test/rootfind_tests.jl
@@ -1,6 +1,4 @@
 @testsnippet RootfindingTestSnippet begin
-    using NonlinearSolveBase, BracketingNonlinearSolve
-
     quadratic_f(u, p) = u .* u .- p
     quadratic_f!(du, u, p) = (du .= u .* u .- p)
     quadratic_f2(u, p) = @. p[1] * u * u - p[2]

diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml
@@ -5,6 +5,7 @@ version = "1.0.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
@@ -24,6 +25,7 @@ NonlinearSolveBaseSparseArraysExt = "SparseArrays"
 
 [compat]
 ArrayInterface = "7.9"
+CommonSolve = "0.2.4"
 Compat = "4.15"
 ConcreteStructs = "0.2.3"
 FastClosures = "0.3"

diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl
@@ -43,7 +43,7 @@ function NonlinearSolveBase.nonlinearsolve_forwarddiff_solve(
 end
 
 function nonlinearsolve_∂f_∂p(prob, f::F, u, p) where {F}
-    if isinplace(prob)
+    if SciMLBase.isinplace(prob)
         f = @closure p -> begin
             du = Utils.safe_similar(u, promote_type(eltype(u), eltype(p)))
             f(du, u, p)
@@ -62,10 +62,11 @@ function nonlinearsolve_∂f_∂p(prob, f::F, u, p) where {F}
 end
 
 function nonlinearsolve_∂f_∂u(prob, f::F, u, p) where {F}
-    if isinplace(prob)
+    if SciMLBase.isinplace(prob)
         return ForwardDiff.jacobian(
             @closure((du, u)->f(du, u, p)), Utils.safe_similar(u), u)
     end
+    u isa Number && return ForwardDiff.derivative(Base.Fix2(f, p), u)
     return ForwardDiff.jacobian(Base.Fix2(f, p), u)
 end