|
1 |
| -@concrete struct JacobianWrapper |
| 1 | +@concrete struct JacobianWrapper{iip} |
2 | 2 | f
|
3 | 3 | p
|
4 | 4 | end
|
5 | 5 |
|
6 |
| -(uf::JacobianWrapper)(u) = uf.f(u, uf.p) |
7 |
| -(uf::JacobianWrapper)(res, u) = uf.f(res, u, uf.p) |
| 6 | +# Previous Implementation did not hold onto `iip`, but this causes problems in packages |
| 7 | +# where we check for the presence of function signatures to check which dispatch to call |
| 8 | +(uf::JacobianWrapper{false})(u) = uf.f(u, uf.p) |
| 9 | +(uf::JacobianWrapper{false})(res, u) = (vec(res) .= vec(uf.f(u, uf.p))) |
| 10 | +(uf::JacobianWrapper{true})(res, u) = uf.f(res, u, uf.p) |
8 | 11 |
|
9 | 12 | sparsity_detection_alg(f, ad) = NoSparsityDetection()
|
10 | 13 | function sparsity_detection_alg(f, ad::AbstractSparseADType)
|
@@ -48,7 +51,7 @@ jacobian!!(::Number, cache) = last(value_derivative(cache.uf, cache.u))
|
48 | 51 | # Build Jacobian Caches
|
49 | 52 | function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p,
|
50 | 53 | ::Val{iip}) where {iip}
|
51 |
| - uf = JacobianWrapper(f, p) |
| 54 | + uf = JacobianWrapper{iip}(f, p) |
52 | 55 |
|
53 | 56 | haslinsolve = hasfield(typeof(alg), :linsolve)
|
54 | 57 |
|
|
98 | 101 | function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u::Number, p,
|
99 | 102 | ::Val{false})
|
100 | 103 | # NOTE: Scalar `u` assumes scalar output from `f`
|
101 |
| - uf = JacobianWrapper(f, p) |
| 104 | + uf = JacobianWrapper{false}(f, p) |
102 | 105 | return uf, nothing, u, nothing, nothing, u
|
103 | 106 | end
|
0 commit comments