Fix Jacobian Construction for concrete_jac

SciML · Oct 6, 2023 · a1877cd · a1877cd
1 parent 9d98446
commit a1877cd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "2.0.2"
+version = "2.0.3"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

diff --git a/src/jacobian.jl b/src/jacobian.jl
@@ -1,4 +1,4 @@
-@concrete struct JacobianWrapper{iip}
+@concrete struct JacobianWrapper{iip} <: Function
     f
     p
 end
@@ -73,9 +73,9 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii
         jac_cache = nothing
     end
 
-    J = if !linsolve_needs_jac
+    J = if !(linsolve_needs_jac || alg_wants_jac)
         # We don't need to construct the Jacobian
-        JacVec(uf, u; autodiff = alg.ad)
+        JacVec(uf, u; autodiff = __get_nonsparse_ad(alg.ad))
     else
         if has_analytic_jac
             f.jac_prototype === nothing ? undefmatrix(u) : f.jac_prototype
@@ -98,6 +98,11 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{ii
     return uf, linsolve, J, fu, jac_cache, du
 end
 
+__get_nonsparse_ad(::AutoSparseForwardDiff) = AutoForwardDiff()
+__get_nonsparse_ad(::AutoSparseFiniteDiff) = AutoFiniteDiff()
+__get_nonsparse_ad(::AutoSparseZygote) = AutoZygote()
+__get_nonsparse_ad(ad) = ad
+
 ## Special Handling for Scalars
 function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u::Number, p,
     ::Val{false}; kwargs...)