From 0fbf311d750426aff098536a69b49b8ca5460e53 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 15 Sep 2023 17:02:54 -0400 Subject: [PATCH] Fix JacVec for not inplace problems --- src/jacobian.jl | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index 157562752..3c6d2b7c2 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -1,10 +1,13 @@ -@concrete struct JacobianWrapper +@concrete struct JacWrapper{iip} f p end -(uf::JacobianWrapper)(u) = uf.f(u, uf.p) -(uf::JacobianWrapper)(res, u) = uf.f(res, u, uf.p) +# Previous Implementation did not hold onto `iip`, but this causes problems in packages +# where we check for the presence of function signatures to check which dispatch to call +(uf::JacWrapper{false})(u) = uf.f(u, uf.p) +(uf::JacWrapper{false})(res, u) = (vec(res) .= vec(uf.f(u, uf.p))) +(uf::JacWrapper{true})(res, u) = uf.f(res, u, uf.p) sparsity_detection_alg(f, ad) = NoSparsityDetection() function sparsity_detection_alg(f, ad::AbstractSparseADType) @@ -48,7 +51,7 @@ jacobian!!(::Number, cache) = last(value_derivative(cache.uf, cache.u)) # Build Jacobian Caches function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p, ::Val{iip}) where {iip} - uf = JacobianWrapper(f, p) + uf = JacWrapper{iip}(f, p) haslinsolve = hasfield(typeof(alg), :linsolve)