From 627718755e8ada79139d256dd9a1e183925eca9d Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 26 Sep 2024 10:32:33 -0400
Subject: [PATCH 01/11] feat: use DI for dense jacobians

fix: stop storing extra stuff in JacobianCache

feat: use DI for dense jacobians

refactor: remove alg from the cache

fix: don't ignore sparsity for dense_ad
---
 Project.toml             |   4 +-
 src/NonlinearSolve.jl    |   9 ++-
 src/internal/jacobian.jl | 160 ++++++++++++++++++++-------------------
 3 files changed, 93 insertions(+), 80 deletions(-)

diff --git a/Project.toml b/Project.toml
index ab35df1e2..a8028527e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -65,8 +66,9 @@ BandedMatrices = "1.5"
 BenchmarkTools = "1.4"
 CUDA = "5.2"
 ConcreteStructs = "0.2.3"
+DifferentiationInterface = "0.6.1"
 DiffEqBase = "6.149.0"
-Enzyme = "0.12"
+Enzyme = "0.12, 0.13"
 ExplicitImports = "1.5"
 FastBroadcast = "0.2.8, 0.3"
 FastClosures = "0.3.2"
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
index 2e420a761..e68fa5472 100644
--- a/src/NonlinearSolve.jl
+++ b/src/NonlinearSolve.jl
@@ -7,8 +7,8 @@ end
 using Reexport: @reexport
 using PrecompileTools: @compile_workload, @setup_workload
 
-using ADTypes: ADTypes, AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff,
-               AutoZygote, AutoEnzyme, AutoSparse
+using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff,
+               AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse
 # FIXME: deprecated, remove in future
 using ADTypes: AutoSparseFiniteDiff, AutoSparseForwardDiff, AutoSparsePolyesterForwardDiff,
                AutoSparseZygote
@@ -22,6 +22,7 @@ using DiffEqBase: DiffEqBase, AbstractNonlinearTerminationMode,
                   NormTerminationMode, RelNormTerminationMode, RelSafeBestTerminationMode,
                   RelSafeTerminationMode, RelTerminationMode,
                   SimpleNonlinearSolveTerminationMode, SteadyStateDiffEqTerminationMode
+using DifferentiationInterface: DifferentiationInterface, Constant
 using FastBroadcast: @..
 using FastClosures: @closure
 using FiniteDiff: FiniteDiff
@@ -39,7 +40,7 @@ using Printf: @printf
 using Preferences: Preferences, @load_preference, @set_preferences!
 using RecursiveArrayTools: recursivecopy!, recursivefill!
 using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
-                 AbstractSciMLOperator, _unwrap_val, has_jac, isinplace, NLStats
+                 AbstractSciMLOperator, _unwrap_val, isinplace, NLStats
 using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJacOperator,
                               JacVecOperator, StatefulJacobianOperator
 using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC
@@ -55,6 +56,8 @@ using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingPro
 
 @reexport using SciMLBase, SimpleNonlinearSolve
 
+const DI = DifferentiationInterface
+
 # Type-Inference Friendly Check for Extension Loading
 is_extension_loaded(::Val) = false
 
diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
index 0c6080112..fc8e1d0a6 100644
--- a/src/internal/jacobian.jl
+++ b/src/internal/jacobian.jl
@@ -31,35 +31,24 @@ Construct a cache for the Jacobian of `f` w.r.t. `u`.
 @concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip}
     J
     f
-    uf
     fu
     u
     p
-    jac_cache
-    alg
     stats::NLStats
     autodiff
-    vjp_autodiff
-    jvp_autodiff
+    di_extras
+    sdifft_extras
 end
 
 function reinit_cache!(cache::JacobianCache{iip}, args...; p = cache.p,
         u0 = cache.u, kwargs...) where {iip}
     cache.u = u0
     cache.p = p
-    cache.uf = JacobianWrapper{iip}(cache.f, p)
 end
 
 function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
         vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F}
     iip = isinplace(prob)
-    uf = JacobianWrapper{iip}(f, p)
-
-    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
-    jvp_autodiff = get_concrete_forward_ad(
-        jvp_autodiff, prob, Val(false); check_forward_mode = true)
-    vjp_autodiff = get_concrete_reverse_ad(
-        vjp_autodiff, prob, Val(false); check_reverse_mode = false)
 
     has_analytic_jac = SciMLBase.has_jac(f)
     linsolve_needs_jac = concrete_jac(alg) === nothing && (linsolve === missing ||
@@ -70,90 +59,128 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
     @bb fu = similar(fu_)
 
     if !has_analytic_jac && needs_jac
-        sd = __sparsity_detection_alg(f, autodiff)
-        jac_cache = iip ? sparse_jacobian_cache(autodiff, sd, uf, fu, u) :
-                    sparse_jacobian_cache(
-            autodiff, sd, uf, __maybe_mutable(u, autodiff); fx = fu)
+        autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
+        sd = sparsity_detection_alg(f, autodiff)
+        sparse_jac = !(sd isa NoSparsityDetection)
+        # Eventually we want to do everything via DI. But for now, we just do the dense via DI
+        if sparse_jac
+            di_extras = nothing
+            uf = JacobianWrapper{iip}(f, p)
+            sdifft_extras = if iip
+                sparse_jacobian_cache(autodiff, sd, uf, fu, u)
+            else
+                sparse_jacobian_cache(
+                    autodiff, sd, uf, __maybe_mutable(u, autodiff); fx = fu)
+            end
+        else
+            sdifft_extras = nothing
+            di_extras = if iip
+                DI.prepare_jacobian(f, fu, autodiff, u, Constant(p))
+            else
+                DI.prepare_jacobian(f, autodiff, u, Constant(p))
+            end
+        end
     else
-        jac_cache = nothing
+        sparse_jac = false
+        di_extras = nothing
+        sdifft_extras = nothing
     end
 
     J = if !needs_jac
+        jvp_autodiff = get_concrete_forward_ad(
+            jvp_autodiff, prob, Val(false); check_forward_mode = true)
+        vjp_autodiff = get_concrete_reverse_ad(
+            vjp_autodiff, prob, Val(false); check_reverse_mode = false)
         JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff)
     else
-        if has_analytic_jac
-            f.jac_prototype === nothing ?
-            __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
-            copy(f.jac_prototype)
-        elseif f.jac_prototype === nothing
-            zero(init_jacobian(jac_cache; preserve_immutable = Val(true)))
+        if f.jac_prototype === nothing
+            if !sparse_jac
+                __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
+            else
+                zero(init_jacobian(sdifft_extras; preserve_immutable = Val(true)))
+            end
         else
-            f.jac_prototype
+            similar(f.jac_prototype)
         end
     end
 
     return JacobianCache{iip}(
-        J, f, uf, fu, u, p, jac_cache, alg, stats, autodiff, vjp_autodiff, jvp_autodiff)
+        J, f, fu, u, p, stats, autodiff, di_extras, sdifft_extras)
 end
 
 function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
         autodiff = nothing, kwargs...) where {F}
-    uf = JacobianWrapper{false}(f, p)
-    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
-    if !(autodiff isa AutoForwardDiff ||
-         autodiff isa AutoPolyesterForwardDiff ||
-         autodiff isa AutoFiniteDiff)
-        # Other cases are not properly supported so we fallback to finite differencing
-        @warn "Scalar AD is supported only for AutoForwardDiff and AutoFiniteDiff. \
-               Detected $(autodiff). Falling back to AutoFiniteDiff."
-        autodiff = AutoFiniteDiff()
+    fu = f(u, p)
+    if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
+        return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing)
     end
-    return JacobianCache{false}(
-        u, f, uf, u, u, p, nothing, alg, stats, autodiff, nothing, nothing)
+    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
+    di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p))
+    return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras, nothing)
 end
 
-@inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p)
-@inline function (cache::JacobianCache)(::Nothing)
+(cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p)
+function (cache::JacobianCache)(::Nothing)
     cache.J isa JacobianOperator &&
         return StatefulJacobianOperator(cache.J, cache.u, cache.p)
     return cache.J
 end
 
+# Operator
 function (cache::JacobianCache)(J::JacobianOperator, u, p = cache.p)
     return StatefulJacobianOperator(J, u, p)
 end
+# Numbers
 function (cache::JacobianCache)(::Number, u, p = cache.p) # Scalar
     cache.stats.njacs += 1
-    J = last(__value_derivative(cache.autodiff, cache.uf, u))
-    return J
+    if SciMLBase.has_jac(cache.f)
+        return cache.f.jac(u, p)
+    elseif SciMLBase.has_vjp(cache.f)
+        return cache.f.vjp(one(u), u, p)
+    elseif SciMLBase.has_jvp(cache.f)
+        return cache.f.jvp(one(u), u, p)
+    end
+    return DI.derivative(cache.f, cache.di_extras, cache.autodiff, u, Constant(p))
 end
-# Compute the Jacobian
+# Actually Compute the Jacobian
 function (cache::JacobianCache{iip})(
         J::Union{AbstractMatrix, Nothing}, u, p = cache.p) where {iip}
     cache.stats.njacs += 1
     if iip
-        if has_jac(cache.f)
+        if SciMLBase.has_jac(cache.f)
             cache.f.jac(J, u, p)
+        elseif cache.di_extras !== nothing
+            DI.jacobian!(
+                cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p))
         else
-            sparse_jacobian!(J, cache.autodiff, cache.jac_cache, cache.uf, cache.fu, u)
+            uf = JacobianWrapper{iip}(cache.f, p)
+            sparse_jacobian!(J, cache.autodiff, cache.jac_cache, uf, cache.fu, u)
         end
-        J_ = J
+        return J
     else
-        J_ = if has_jac(cache.f)
-            cache.f.jac(u, p)
-        elseif __can_setindex(typeof(J))
-            sparse_jacobian!(J, cache.autodiff, cache.jac_cache, cache.uf, u)
-            J
+        if SciMLBase.has_jac(cache.f)
+            return cache.f.jac(u, p)
+        elseif cache.di_extras !== nothing
+            return DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p))
         else
-            sparse_jacobian(cache.autodiff, cache.jac_cache, cache.uf, u)
+            uf = JacobianWrapper{iip}(cache.f, p)
+            if __can_setindex(typeof(J))
+                sparse_jacobian!(J, cache.autodiff, cache.sdifft_extras, uf, u)
+                return J
+            else
+                return sparse_jacobian(cache.autodiff, cache.sdifft_extras, uf, u)
+            end
         end
     end
-    return J_
 end
 
-# Sparsity Detection Choices
-@inline __sparsity_detection_alg(_, _) = NoSparsityDetection()
-@inline function __sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
+function sparsity_detection_alg(f::NonlinearFunction, ad::AbstractADType)
+    # TODO: Also handle case where colorvec is provided
+    f.sparsity === nothing && return NoSparsityDetection()
+    return sparsity_detection_alg(f, AutoSparse(ad; sparsity_detector = f.sparsity))
+end
+
+function sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
     if f.sparsity === nothing
         if f.jac_prototype === nothing
             is_extension_loaded(Val(:Symbolics)) && return SymbolicsSparsityDetection()
@@ -177,28 +204,9 @@ end
     end
 
     if SciMLBase.has_colorvec(f)
-        return PrecomputedJacobianColorvec(; jac_prototype,
-            f.colorvec,
-            partition_by_rows = (ad isa AutoSparse &&
-                                 ADTypes.mode(ad) isa ADTypes.ReverseMode))
+        return PrecomputedJacobianColorvec(; jac_prototype, f.colorvec,
+            partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode)
     else
         return JacPrototypeSparsityDetection(; jac_prototype)
     end
 end
-
-@inline function __value_derivative(
-        ::Union{AutoForwardDiff, AutoPolyesterForwardDiff}, f::F, x::R) where {F, R}
-    T = typeof(ForwardDiff.Tag(f, R))
-    out = f(ForwardDiff.Dual{T}(x, one(x)))
-    return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
-end
-
-@inline function __value_derivative(ad::AutoFiniteDiff, f::F, x::R) where {F, R}
-    return f(x), FiniteDiff.finite_difference_derivative(f, x, ad.fdtype)
-end
-
-@inline function __scalar_jacvec(f::F, x::R, v::V) where {F, R, V}
-    T = typeof(ForwardDiff.Tag(f, R))
-    out = f(ForwardDiff.Dual{T}(x, v))
-    return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
-end

From 1c3f387984d11e3d580cad11656187b552f28cc1 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 13:47:47 -0400
Subject: [PATCH 02/11] fix: update minimum compats

---
 Project.toml | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index a8028527e..55283a7d4 100644
--- a/Project.toml
+++ b/Project.toml
@@ -67,8 +67,8 @@ BenchmarkTools = "1.4"
 CUDA = "5.2"
 ConcreteStructs = "0.2.3"
 DifferentiationInterface = "0.6.1"
-DiffEqBase = "6.149.0"
-Enzyme = "0.12, 0.13"
+DiffEqBase = "6.155.1"
+Enzyme = "0.13"
 ExplicitImports = "1.5"
 FastBroadcast = "0.2.8, 0.3"
 FastClosures = "0.3.2"
@@ -82,7 +82,7 @@ LazyArrays = "1.8.2, 2"
 LeastSquaresOptim = "0.8.5"
 LineSearches = "7.2"
 LinearAlgebra = "1.10"
-LinearSolve = "2.30"
+LinearSolve = "2.34"
 MINPACK = "1.2"
 MaybeInplace = "0.1.3"
 ModelingToolkit = "9.15.0"
@@ -100,11 +100,11 @@ ReTestItems = "1.24"
 RecursiveArrayTools = "3.8"
 Reexport = "1.2"
 SIAMFANLEquations = "1.0.1"
-SciMLBase = "2.34.0"
+SciMLBase = "2.54.0"
 SciMLJacobianOperators = "0.1"
-SimpleNonlinearSolve = "1.8"
+SimpleNonlinearSolve = "1.9.3"
 SparseArrays = "1.10"
-SparseDiffTools = "2.19"
+SparseDiffTools = "2.22"
 SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.9"

From f9ea92c74863a535e7f38655b053bf0b4cfd816f Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 14:17:53 -0400
Subject: [PATCH 03/11] fix: sparse jacobian

---
 docs/src/tutorials/code_optimization.md | 22 +++++++++-------------
 src/internal/jacobian.jl                | 11 +++++++++--
 2 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/docs/src/tutorials/code_optimization.md b/docs/src/tutorials/code_optimization.md
index fa0f61657..ce114961b 100644
--- a/docs/src/tutorials/code_optimization.md
+++ b/docs/src/tutorials/code_optimization.md
@@ -33,9 +33,7 @@ Take for example a prototypical small nonlinear solver code in its out-of-place
 ```@example small_opt
 using NonlinearSolve
 
-function f(u, p)
-    u .* u .- p
-end
+f(u, p) = u .* u .- p
 u0 = [1.0, 1.0]
 p = 2.0
 prob = NonlinearProblem(f, u0, p)
@@ -53,9 +51,7 @@ using BenchmarkTools
 Note that this way of writing the function is a shorthand for:
 
 ```@example small_opt
-function f(u, p)
-    [u[1] * u[1] - p, u[2] * u[2] - p]
-end
+f(u, p) = [u[1] * u[1] - p, u[2] * u[2] - p]
 ```
 
 where the function `f` returns an array. This is a common pattern from things like MATLAB's
@@ -71,7 +67,7 @@ by hand, this looks like:
 function f(du, u, p)
     du[1] = u[1] * u[1] - p
     du[2] = u[2] * u[2] - p
-    nothing
+    return nothing
 end
 
 prob = NonlinearProblem(f, u0, p)
@@ -84,6 +80,7 @@ the `.=` in-place broadcasting.
 ```@example small_opt
 function f(du, u, p)
     du .= u .* u .- p
+    return nothing
 end
 
 @benchmark sol = solve(prob, NewtonRaphson())
@@ -114,6 +111,7 @@ to normal array expressions, for example:
 
 ```@example small_opt
 using StaticArrays
+
 A = SA[2.0, 3.0, 5.0]
 typeof(A)
 ```
@@ -135,12 +133,12 @@ want to use the out-of-place allocating form, but this time we want to output a
 array. Doing it with broadcasting looks like:
 
 ```@example small_opt
-function f_SA(u, p)
-    u .* u .- p
-end
+f_SA(u, p) = u .* u .- p
+
 u0 = SA[1.0, 1.0]
 p = 2.0
 prob = NonlinearProblem(f_SA, u0, p)
+
 @benchmark solve(prob, NewtonRaphson())
 ```
 
@@ -148,9 +146,7 @@ Note that only change here is that `u0` is made into a StaticArray! If we needed
 `f` out for a more complex nonlinear case, then we'd simply do the following:
 
 ```@example small_opt
-function f_SA(u, p)
-    SA[u[1] * u[1] - p, u[2] * u[2] - p]
-end
+f_SA(u, p) = SA[u[1] * u[1] - p, u[2] * u[2] - p]
 
 @benchmark solve(prob, NewtonRaphson())
 ```
diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
index fc8e1d0a6..bc64bbd06 100644
--- a/src/internal/jacobian.jl
+++ b/src/internal/jacobian.jl
@@ -95,7 +95,14 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
     else
         if f.jac_prototype === nothing
             if !sparse_jac
-                __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
+                # While this is technically wasteful, it gives out the type of the Jacobian
+                # which is needed to create the linear solver cache
+                stats.njacs += 1
+                if iip
+                    DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p))
+                else
+                    DI.jacobian(f, autodiff, u, Constant(p))
+                end
             else
                 zero(init_jacobian(sdifft_extras; preserve_immutable = Val(true)))
             end
@@ -154,7 +161,7 @@ function (cache::JacobianCache{iip})(
                 cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p))
         else
             uf = JacobianWrapper{iip}(cache.f, p)
-            sparse_jacobian!(J, cache.autodiff, cache.jac_cache, uf, cache.fu, u)
+            sparse_jacobian!(J, cache.autodiff, cache.sdifft_extras, uf, cache.fu, u)
         end
         return J
     else

From 3134e5d2891536d9e61e94df3944aefb874ca24b Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 14:18:39 -0400
Subject: [PATCH 04/11] chore: run formatter

---
 lib/SciMLJacobianOperators/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/SciMLJacobianOperators/README.md b/lib/SciMLJacobianOperators/README.md
index c9d14ff50..0bce0f919 100644
--- a/lib/SciMLJacobianOperators/README.md
+++ b/lib/SciMLJacobianOperators/README.md
@@ -56,4 +56,4 @@ snormal_form * v  # Computes JᵀJ * v
 #  8.0
 #  8.0
 #  8.0
-```
\ No newline at end of file
+```

From 6d2b026c8d1fff91882fbc87f01d69689fbb44dd Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 14:31:35 -0400
Subject: [PATCH 05/11] fix: unwrap sparse AD for derivative call

---
 src/internal/jacobian.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
index bc64bbd06..010b9f6c5 100644
--- a/src/internal/jacobian.jl
+++ b/src/internal/jacobian.jl
@@ -122,7 +122,7 @@ function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
         return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing)
     end
     autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
-    di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p))
+    di_extras = DI.prepare_derivative(f, get_dense_ad(autodiff), u, Constant(prob.p))
     return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras, nothing)
 end
 
@@ -217,3 +217,6 @@ function sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
         return JacPrototypeSparsityDetection(; jac_prototype)
     end
 end
+
+get_dense_ad(ad) = ad
+get_dense_ad(ad::AutoSparse) = ADTypes.dense_ad(ad)

From 02430ff8624fa79f6c433a8d7e1ea8769a203c81 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 14:47:54 -0400
Subject: [PATCH 06/11] chore: bump minimum versions

---
 Project.toml                                  | 32 +++++++++----------
 docs/Project.toml                             |  5 +--
 .../tutorials/optimizing_parameterized_ode.md |  2 +-
 test/misc/polyalg_tests.jl                    |  2 +-
 4 files changed, 21 insertions(+), 20 deletions(-)

diff --git a/Project.toml b/Project.toml
index 55283a7d4..5bd304bb3 100644
--- a/Project.toml
+++ b/Project.toml
@@ -59,21 +59,21 @@ NonlinearSolveSymbolicsExt = "Symbolics"
 NonlinearSolveZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "1.1.0"
+ADTypes = "1.9"
 Aqua = "0.8"
-ArrayInterface = "7.9"
+ArrayInterface = "7.16"
 BandedMatrices = "1.5"
 BenchmarkTools = "1.4"
-CUDA = "5.2"
+CUDA = "5.5"
 ConcreteStructs = "0.2.3"
+DiffEqBase = "6.155.3"
 DifferentiationInterface = "0.6.1"
-DiffEqBase = "6.155.1"
-Enzyme = "0.13"
+Enzyme = "0.13.2"
 ExplicitImports = "1.5"
-FastBroadcast = "0.2.8, 0.3"
+FastBroadcast = "0.3.5"
 FastClosures = "0.3.2"
 FastLevenbergMarquardt = "0.1"
-FiniteDiff = "2.22"
+FiniteDiff = "2.24"
 FixedPointAcceleration = "0.3"
 ForwardDiff = "0.10.36"
 Hwloc = "3"
@@ -82,27 +82,27 @@ LazyArrays = "1.8.2, 2"
 LeastSquaresOptim = "0.8.5"
 LineSearches = "7.2"
 LinearAlgebra = "1.10"
-LinearSolve = "2.34"
+LinearSolve = "2.35"
 MINPACK = "1.2"
 MaybeInplace = "0.1.3"
-ModelingToolkit = "9.15.0"
+ModelingToolkit = "9.41.0"
 NLSolvers = "0.5"
 NLsolve = "4.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.2"
-OrdinaryDiffEq = "6.75"
+OrdinaryDiffEqTsit5 = "1.1.0"
 Pkg = "1.10"
 PrecompileTools = "1.2"
 Preferences = "1.4"
 Printf = "1.10"
 Random = "1.91"
 ReTestItems = "1.24"
-RecursiveArrayTools = "3.8"
+RecursiveArrayTools = "3.27"
 Reexport = "1.2"
 SIAMFANLEquations = "1.0.1"
 SciMLBase = "2.54.0"
 SciMLJacobianOperators = "0.1"
-SimpleNonlinearSolve = "1.9.3"
+SimpleNonlinearSolve = "1.12.3"
 SparseArrays = "1.10"
 SparseDiffTools = "2.22"
 SpeedMapping = "0.3"
@@ -110,8 +110,8 @@ StableRNGs = "1"
 StaticArrays = "1.9"
 StaticArraysCore = "1.4"
 Sundials = "4.23.1"
-SymbolicIndexingInterface = "0.3.15"
-Symbolics = "5.26, 6"
+SymbolicIndexingInterface = "0.3.31"
+Symbolics = "6.12"
 Test = "1.10"
 TimerOutputs = "0.5.23"
 Zygote = "0.6.69"
@@ -135,7 +135,7 @@ NLSolvers = "337daf1e-9722-11e9-073e-8b9effe078ba"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
@@ -149,4 +149,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEq", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
diff --git a/docs/Project.toml b/docs/Project.toml
index 63d1803d4..9fbee9e39 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -10,7 +10,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -30,10 +30,11 @@ DiffEqBase = "6.136"
 Documenter = "1"
 DocumenterCitations = "1"
 IncompleteLU = "0.2"
+InteractiveUtils = "<0.0.1, 1"
 LinearSolve = "2"
 ModelingToolkit = "8, 9"
 NonlinearSolve = "3"
-OrdinaryDiffEq = "6"
+OrdinaryDiffEqTsit5 = "1.1.0"
 Plots = "1"
 Random = "<0.0.1, 1"
 SciMLBase = "2.4"
diff --git a/docs/src/tutorials/optimizing_parameterized_ode.md b/docs/src/tutorials/optimizing_parameterized_ode.md
index d3b409eca..69cca060d 100644
--- a/docs/src/tutorials/optimizing_parameterized_ode.md
+++ b/docs/src/tutorials/optimizing_parameterized_ode.md
@@ -4,7 +4,7 @@ Let us fit a parameterized ODE to some data. We will use the Lotka-Volterra mode
 example. We will use Single Shooting to fit the parameters.
 
 ```@example parameterized_ode
-using OrdinaryDiffEq, NonlinearSolve, Plots
+using OrdinaryDiffEqTsit5, NonlinearSolve, Plots
 ```
 
 Let us simulate some real data from the Lotka-Volterra model.
diff --git a/test/misc/polyalg_tests.jl b/test/misc/polyalg_tests.jl
index 88be7e8cc..a106ace44 100644
--- a/test/misc/polyalg_tests.jl
+++ b/test/misc/polyalg_tests.jl
@@ -127,7 +127,7 @@ end
 # Testing for Complex Valued Root Finding. For Complex valued inputs we drop some of the
 # algorithms which dont support those.
 @testitem "Complex Valued Problems: Single-Shooting" tags=[:misc] begin
-    using OrdinaryDiffEq
+    using OrdinaryDiffEqTsit5
 
     function ode_func!(du, u, p, t)
         du[1] = u[2]

From 5ad2fb20c05986dc97084b8df3d581023b8c628a Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 14:49:55 -0400
Subject: [PATCH 07/11] fix: update bruss testing

---
 test/misc/bruss_tests.jl | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/test/misc/bruss_tests.jl b/test/misc/bruss_tests.jl
index 21b7d792b..d18c67e6a 100644
--- a/test/misc/bruss_tests.jl
+++ b/test/misc/bruss_tests.jl
@@ -66,7 +66,6 @@
 
     sol = solve(prob_brusselator_2d, NewtonRaphson(); abstol = 1e-8)
     @test norm(sol.resid, Inf) < 1e-8
-    @test !all(iszero, jac_prototype)
 
     sol = solve(prob_brusselator_2d,
         NewtonRaphson(autodiff = AutoSparse(AutoFiniteDiff())); abstol = 1e-8)
@@ -74,6 +73,6 @@
 
     cache = init(
         prob_brusselator_2d, NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff())))
-    @test maximum(cache.jac_cache.jac_cache.coloring.colorvec) == 12
-    @test cache.jac_cache.autodiff isa AutoSparse{<:AutoForwardDiff}
+    @test maximum(cache.sdifft_extras.jac_cache.coloring.colorvec) == 12
+    @test cache.sdifft_extras.autodiff isa AutoSparse{<:AutoForwardDiff}
 end

From 25efa1679908734876e882fbd2b6493df7ea3667 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 14:50:53 -0400
Subject: [PATCH 08/11] fix: autodiff cannot be nothing

---
 src/internal/jacobian.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
index 010b9f6c5..46f255149 100644
--- a/src/internal/jacobian.jl
+++ b/src/internal/jacobian.jl
@@ -58,8 +58,9 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
 
     @bb fu = similar(fu_)
 
+    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
+
     if !has_analytic_jac && needs_jac
-        autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
         sd = sparsity_detection_alg(f, autodiff)
         sparse_jac = !(sd isa NoSparsityDetection)
         # Eventually we want to do everything via DI. But for now, we just do the dense via DI

From be0f8da997e348a8813e0571d4edf3aa133852be Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 14:52:45 -0400
Subject: [PATCH 09/11] fix: __value_derivative removal from line searches

---
 src/globalization/line_search.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl
index 5de4610b6..d43e97247 100644
--- a/src/globalization/line_search.jl
+++ b/src/globalization/line_search.jl
@@ -110,10 +110,10 @@ function __internal_init(
             @warn "Scalar AD is supported only for AutoForwardDiff and AutoFiniteDiff. \
                    Detected $(autodiff). Falling back to AutoFiniteDiff."
         end
-        deriv_op = @closure (du, u, fu, p) -> last(__value_derivative(
-                                                  autodiff, Base.Fix2(f, p), u)) *
-                                              fu *
-                                              du
+        deriv_op = @closure (du, u, fu, p) -> begin
+            # Temporary solution, we are anyways moving to LineSearch.jl
+            return DI.derivative(f, autodiff, u, Constant(p)) * fu * du
+        end
     else
         # Both forward and reverse AD can be used for line-search.
         # We prefer forward AD for better performance, however, reverse AD is also supported if user explicitly requests it.

From 59347e39f20b86200b5be4cec74a22b2554cdb26 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 15:07:22 -0400
Subject: [PATCH 10/11] fix: minor test fixes

---
 src/globalization/line_search.jl |  3 ++-
 src/internal/jacobian.jl         | 16 +++++++++++-----
 test/misc/bruss_tests.jl         |  4 ++--
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl
index d43e97247..323c49258 100644
--- a/src/globalization/line_search.jl
+++ b/src/globalization/line_search.jl
@@ -101,7 +101,8 @@ function __internal_init(
         args...; stats, internalnorm::IN = DEFAULT_NORM, kwargs...) where {F, IN}
     T = promote_type(eltype(fu), eltype(u))
     if u isa Number
-        autodiff = get_concrete_forward_ad(alg.autodiff, prob; check_forward_mode = true)
+        autodiff = get_dense_ad(get_concrete_forward_ad(
+            alg.autodiff, prob; check_forward_mode = true))
         if !(autodiff isa AutoForwardDiff ||
              autodiff isa AutoPolyesterForwardDiff ||
              autodiff isa AutoFiniteDiff)
diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
index 46f255149..67ab839a2 100644
--- a/src/internal/jacobian.jl
+++ b/src/internal/jacobian.jl
@@ -99,10 +99,15 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
                 # While this is technically wasteful, it gives out the type of the Jacobian
                 # which is needed to create the linear solver cache
                 stats.njacs += 1
-                if iip
-                    DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p))
+                if has_analytic_jac
+                    __similar(
+                        fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
                 else
-                    DI.jacobian(f, autodiff, u, Constant(p))
+                    if iip
+                        DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p))
+                    else
+                        DI.jacobian(f, autodiff, u, Constant(p))
+                    end
                 end
             else
                 zero(init_jacobian(sdifft_extras; preserve_immutable = Val(true)))
@@ -120,9 +125,10 @@ function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
         autodiff = nothing, kwargs...) where {F}
     fu = f(u, p)
     if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
-        return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing)
+        return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing, nothing)
     end
-    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
+    autodiff = get_dense_ad(get_concrete_forward_ad(
+        autodiff, prob; check_forward_mode = false))
     di_extras = DI.prepare_derivative(f, get_dense_ad(autodiff), u, Constant(prob.p))
     return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras, nothing)
 end
diff --git a/test/misc/bruss_tests.jl b/test/misc/bruss_tests.jl
index d18c67e6a..8ea1fdb18 100644
--- a/test/misc/bruss_tests.jl
+++ b/test/misc/bruss_tests.jl
@@ -73,6 +73,6 @@
 
     cache = init(
         prob_brusselator_2d, NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff())))
-    @test maximum(cache.sdifft_extras.jac_cache.coloring.colorvec) == 12
-    @test cache.sdifft_extras.autodiff isa AutoSparse{<:AutoForwardDiff}
+    @test maximum(cache.jac_cache.sdifft_extras.coloring.colorvec) == 12
+    @test cache.jac_cache.autodiff isa AutoSparse{<:AutoForwardDiff}
 end

From 464cfa18daac530deeb42b532d0c7aa05cac0a70 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Fri, 27 Sep 2024 15:29:07 -0400
Subject: [PATCH 11/11] test: issue 451

---
 test/misc/issues_tests.jl | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)
 create mode 100644 test/misc/issues_tests.jl

diff --git a/test/misc/issues_tests.jl b/test/misc/issues_tests.jl
new file mode 100644
index 000000000..24e39943c
--- /dev/null
+++ b/test/misc/issues_tests.jl
@@ -0,0 +1,22 @@
+@testitem "Issue #451" tags=[:misc] begin
+    f(u, p) = u^2 - p
+
+    jac_calls = 0
+    function df(u, p)
+        global jac_calls += 1
+        return 2u
+    end
+
+    fn = NonlinearFunction(f; jac = df)
+    prob = NonlinearProblem(fn, 1.0, 2.0)
+    sol = solve(prob, NewtonRaphson())
+    @test sol.retcode == ReturnCode.Success
+    @test jac_calls ≥ 1
+
+    jac_calls = 0
+    fn2 = NonlinearFunction(f)
+    prob = NonlinearProblem(fn2, 1.0, 2.0)
+    sol = solve(prob, NewtonRaphson())
+    @test sol.retcode == ReturnCode.Success
+    @test jac_calls == 0
+end