Fix SteadyStateAdjoint Breakage

Project.toml

@@ -1,7 +1,7 @@
 name = "SciMLSensitivity"
 uuid = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 authors = ["Christopher Rackauckas <[email protected]>", "Yingbo Ma <[email protected]>"]
-version = "7.38.0"
+version = "7.38.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -73,7 +73,7 @@ ReverseDiff = "1.9"
 SciMLBase = "1.66.0"
 SciMLOperators = "0.1, 0.2, 0.3"
 SimpleNonlinearSolve = "0.1.8"
-SparseDiffTools = "2.4"
+SparseDiffTools = "2.5"
 StaticArraysCore = "1.4"
 StochasticDiffEq = "6.20"
 Tracker = "0.2"

src/steadystate_adjoint.jl

@@ -101,7 +101,8 @@ end
 
     if !needs_jac
         # operator = VecJac(f, y, p; Val(DiffEqBase.isinplace(sol.prob)))
-        operator = VecJac(f, y, p; autodiff = get_autodiff_from_vjp(vjp))
+        __f = y -> f(y, p, nothing)
+        operator = VecJac(__f, y; autodiff = get_autodiff_from_vjp(sensealg.autojacvec))
         linear_problem = LinearProblem(operator, vec(dgdu_val); u0 = vec(λ))
     else
         linear_problem = LinearProblem(diffcache.J', vec(dgdu_val'); u0 = vec(λ))

test/steady_state.jl

@@ -424,6 +424,27 @@ end
     @test dp1≈dp6 rtol=1e-10
     @test dp1≈dp7 rtol=1e-10
     @test dp1≈dp8 rtol=1e-10
+
+    # Larger Batched Problem: For testing the Iterative Solvers Path
+    u0 = zeros(128)
+    p = [2.0, 1.0]
+
+    prob = NonlinearProblem((u, p) -> u .- p[1] .+ p[2], u0, p)
+    solve1 = solve(remake(prob, p = p), NewtonRaphson())
+
+    function test_loss(p, prob; alg = NewtonRaphson())
+        _prob = remake(prob, p = p)
+        sol = sum(solve(_prob, alg,
+            sensealg = SteadyStateAdjoint(autojacvec = ZygoteVJP())))
+        return sol
+    end
+
+    test_loss(p, prob)
+
+    dp1 = Zygote.gradient(p -> test_loss(p, prob), p)[1]
+
+    @test dp1[1] ≈ 128
+    @test dp1[2] ≈ -128
 end
 
 @testset "Continuous sensitivity tools" begin