Use SparseDiffTools v2 for steadystateadjoint

Project.toml

@@ -4,6 +4,7 @@ authors = ["Christopher Rackauckas <[email protected]>", "Yingbo Ma <
 version = "7.37.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
@@ -71,7 +72,7 @@ ReverseDiff = "1.9"
 SciMLBase = "1.66.0"
 SciMLOperators = "0.1, 0.2, 0.3"
 SimpleNonlinearSolve = "0.1.8"
-SparseDiffTools = "2"
+SparseDiffTools = "2.4"
 StaticArraysCore = "1.4"
 StochasticDiffEq = "6.20"
 Tracker = "0.2"

src/SciMLSensitivity.jl

@@ -14,6 +14,7 @@ import ArrayInterface
 import Enzyme
 import GPUArraysCore
 using StaticArraysCore
+using ADTypes
 using SparseDiffTools
 using SciMLOperators
 import TruncatedStacktraces

src/sensitivity_algorithms.jl

@@ -1313,3 +1313,8 @@ struct ForwardDiffOverAdjoint{A} <:
        AbstractSecondOrderSensitivityAlgorithm{nothing, true, nothing}
     adjalg::A
 end
+
+get_autodiff_from_vjp(vjp::ReverseDiffVJP{compile}) where{compile} = AutoReverseDiff(; compile = compile)
+get_autodiff_from_vjp(::ZygoteVJP) = AutoZygote()
+get_autodiff_from_vjp(::EnzymeVJP) = AutoEnzyme()
+get_autodiff_from_vjp(::TrackerVJP) = AutoTracker()

src/steadystate_adjoint.jl

@@ -49,9 +49,11 @@ end
     dgdu === nothing && dgdp === nothing && g === nothing &&
         error("Either `dgdu`, `dgdp`, or `g` must be specified.")
 
+    needs_jac = if has_adjoint(f)
+        false
     # TODO: What is the correct heuristic? Can we afford to compute Jacobian for
     #       cases where the length(u0) > 50 and if yes till what threshold
-    needs_jac = if sensealg.linsolve === nothing
+    elseif sensealg.linsolve === nothing
         length(u0) <= 50
     else
         LinearSolve.needs_concrete_A(sensealg.linsolve)
@@ -98,8 +100,8 @@ end
     end
 
     if !needs_jac
-        # TODO: FixedVecJacOperator should respect the `autojacvec` of the algorithm
-        operator = FixedVecJacOperator(f, y, p, Val(DiffEqBase.isinplace(sol.prob)))
+        # operator = VecJac(f, y, p; Val(DiffEqBase.isinplace(sol.prob)))
+        operator = VecJac(f, y, p; autodiff = get_autodiff_from_vjp(vjp))
         linear_problem = LinearProblem(operator, vec(dgdu_val); u0 = vec(λ))
     else
         linear_problem = LinearProblem(diffcache.J', vec(dgdu_val'); u0 = vec(λ))
@@ -136,24 +138,3 @@ end
         return vjp
     end
 end
-
-function FixedVecJacOperator(f_in, y, p, ::Val{false})
-    # NOTE: Zygote doesn't support inplace
-    input, f = Zygote.pullback(x -> f_in(x, p, nothing), y)
-    output = f(input)[1]
-    function f_operator!(du, u, p, t)
-        du .= vec(f(reshape(u, size(input)))[1])
-        return du
-    end
-    op = FunctionOperator(f_operator!, vec(input), vec(output))
-    return op
-end
-
-function FixedVecJacOperator(f, y, p, ::Val{true})
-    function f_operator!(du, u, p, t)
-        num_vecjac!(du, (_du, _u) -> f(reshape(_du, size(y)), _u, p, t),
-            y, reshape(u, size(y)))
-        return du
-    end
-    return FunctionOperator(f_operator!, vec(y); p)
-end