Add diff_tunables flag to SteadyStateAdjoint for full-parameter VJP

Add `diff_tunables::Val` field to SteadyStateAdjoint (default Val(true)).
When Val(false), the parameter VJP via vecjacobian! is computed w.r.t.
the full parameter object (including caches) instead of just tunables.

This is needed for SCCNonlinearProblem where explicitfuns! write active
data into non-tunable parameter components (caches). The automatic
sensealg choice detects non-empty caches and sets diff_tunables=Val(false)
with a structured-VJP-compatible backend.

Changes:
- sensitivity_algorithms.jl: Add diff_tunables field to SteadyStateAdjoint
- adjoint_common.jl: Add use_full_p kwarg to adjointdiffcache
- steadystate_adjoint.jl: Gate use_full_p on diff_tunables flag
- concrete_solve.jl: Pass original_p to automatic_sensealg_choice for
  cache detection; return full gradient for EnzymeOriginator when
  diff_tunables=Val(false)
- test/scc_enzyme.jl: Direct SCC differentiation test with Enzyme

Companion PR: NonlinearSolve.jl#884

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

src/adjoint_common.jl

@@ -33,12 +33,17 @@ return (AdjointDiffCache, y)
 function adjointdiffcache(
         g::G, sensealg, discrete, sol, dgdu::DG1, dgdp::DG2, f, alg;
         quad = false,
-        noiseterm = false, needs_jac = false
+        noiseterm = false, needs_jac = false, use_full_p = false
     ) where {G, DG1, DG2}
     prob = sol.prob
     u0 = state_values(prob)
     p = parameter_values(prob)
-    if p === nothing || p isa SciMLBase.NullParameters
+    if use_full_p && p !== nothing && !(p isa SciMLBase.NullParameters)
+        # Use full parameter object (including caches) for VJP computation.
+        # Required for SCCNonlinearProblem where explicitfuns! write active
+        # data into non-tunable parameter components.
+        tunables, repack = p, identity
+    elseif p === nothing || p isa SciMLBase.NullParameters
         tunables, repack = p, identity
     elseif isscimlstructure(p)
         tunables, repack, _ = canonicalize(Tunable(), p)

src/concrete_solve.jl

@@ -324,17 +324,27 @@ function automatic_sensealg_choice(
 end
 
 function automatic_sensealg_choice(
-        prob::ConcreteNonlinearProblem, u0, p,
-        verbose, repack
+        prob::ConcreteNonlinearProblem, u0, tunables,
+        verbose, repack, original_p = tunables
     )
+    # Check if the original parameter has non-tunable active components
+    # (e.g. caches from SCCNonlinearProblem explicitfuns!).
+    _has_caches = isscimlstructure(original_p) && !(original_p isa AbstractArray) &&
+        hasfield(typeof(original_p), :caches) && !isempty(original_p.caches)
+    _diff_tunables = _has_caches ? Val(false) : Val(true)
+
     default_sensealg = if u0 isa GPUArraysCore.AbstractGPUArray ||
             !DiffEqBase.isinplace(prob)
-        # autodiff = false because forwarddiff fails on many GPU kernels
-        # this only effects the Jacobian calculation and is same computation order
-        SteadyStateAdjoint(autodiff = false, autojacvec = ZygoteVJP())
+        SteadyStateAdjoint(
+            autodiff = false, autojacvec = ZygoteVJP(),
+            diff_tunables = _diff_tunables,
+        )
     else
-        vjp = inplace_vjp(prob, u0, p, verbose, repack)
-        SteadyStateAdjoint(autojacvec = vjp)
+        vjp = inplace_vjp(prob, u0, tunables, verbose, repack)
+        if _diff_tunables isa Val{false} && !supports_structured_vjp(vjp)
+            vjp = ZygoteVJP()
+        end
+        SteadyStateAdjoint(autojacvec = vjp, diff_tunables = _diff_tunables)
     end
     return default_sensealg
 end
@@ -371,7 +381,7 @@ function SciMLBase._concrete_solve_adjoint(
         throw(SciMLStructuresCompatibilityError())
     end
 
-    default_sensealg = automatic_sensealg_choice(prob, u0, tunables, verbose, repack)
+    default_sensealg = automatic_sensealg_choice(prob, u0, tunables, verbose, repack, p)
     if has_cb && default_sensealg isa AbstractAdjointSensitivityAlgorithm &&
             !(typeof(default_sensealg.autojacvec) <: Union{EnzymeVJP, ReverseDiffVJP, ReactantVJP})
         default_sensealg = setvjp(default_sensealg, ReverseDiffVJP())
@@ -404,7 +414,7 @@ function SciMLBase._concrete_solve_adjoint(
     end
 
     u0 = state_values(prob) === nothing ? Float64[] : u0
-    default_sensealg = automatic_sensealg_choice(prob, u0, tunables, verbose, repack)
+    default_sensealg = automatic_sensealg_choice(prob, u0, tunables, verbose, repack, p)
     return SciMLBase._concrete_solve_adjoint(
         prob, alg, default_sensealg, u0, p,
         originator::SciMLBase.ADOriginator, args...; verbose,
@@ -2376,56 +2386,49 @@ function SciMLBase._concrete_solve_adjoint(
             end
         end
 
-        dp = adjoint_sensitivities(sol, alg; sensealg, dgdu = df)
+        dp_full = adjoint_sensitivities(sol, alg; sensealg, dgdu = df)
 
-        dp,
-            Δtunables = if Δ isa AbstractArray || Δ isa Number
-            # if Δ isa AbstractArray, the gradients correspond to `u`
-            # this is something that needs changing in the future, but
-            # this is the applicable till the movement to structuaral
-            # tangents is completed
-            dp, Δtunables = if isscimlstructure(dp)
-                dp, _, _ = canonicalize(Tunable(), dp)
-                dp, nothing
-            elseif isfunctor(dp)
-                dp, _ = Functors.functor(dp)
-                dp, nothing
+        # When diff_tunables=Val(false), dp_full is the full parameter
+        # gradient (SciMLStructure). For Enzyme, return it directly so
+        # the reverse rule can accumulate into all shadow components
+        # (including caches for SCCNonlinearProblem).
+        dp_tangent = if originator isa SciMLBase.EnzymeOriginator &&
+                sensealg.diff_tunables isa Val{false} &&
+                isscimlstructure(dp_full)
+            dp_full
+        else
+            dp = if isscimlstructure(dp_full)
+                canonicalize(Tunable(), dp_full)[1]
+            elseif isfunctor(dp_full)
+                Functors.functor(dp_full)[1]
             else
-                dp, nothing
+                dp_full
             end
-        else
-            dp, Δtunables = if isscimlstructure(p)
-                if (Δ.prob.p == ZeroTangent() || Δ.prob.p == NoTangent())
-                    dp, _, _ = canonicalize(Tunable(), dp)
-                    dp, nothing
+
+            Δtunables = if !(Δ isa AbstractArray || Δ isa Number)
+                if isscimlstructure(p) &&
+                        !(Δ.prob.p == ZeroTangent() || Δ.prob.p == NoTangent())
+                    Δp = setproperties(dp_full, to_nt(Δ.prob.p))
+                    canonicalize(Tunable(), Δp)[1]
+                elseif isfunctor(p)
+                    Functors.functor(Δ.prob.p)[1]
                 else
-                    Δp = setproperties(dp, to_nt(Δ.prob.p))
-                    Δtunables, _, _ = canonicalize(Tunable(), Δp)
-                    dp, _, _ = canonicalize(Tunable(), dp)
-                    dp, Δtunables
+                    nothing
                 end
-            elseif isfunctor(p)
-                dp, _ = Functors.functor(dp)
-                Δtunables, _ = Functors.functor(Δ.prob.p)
-                dp, Δtunables
             else
-                dp, Δ.prob.p
+                nothing
             end
-        end
 
-        dp = Zygote.accum(
-            dp, (isnothing(Δtunables) || isempty(Δtunables)) ? nothing :
-                Δtunables
-        )
+            dp = Zygote.accum(
+                dp, (isnothing(Δtunables) || isempty(Δtunables)) ? nothing :
+                    Δtunables
+            )
 
-        # For Enzyme with SciMLStructure parameters, return the tunable gradient
-        # vector directly instead of the Zygote-repacked NamedTuple. The Enzyme
-        # reverse rule in NonlinearSolveBaseEnzymeExt uses
-        # SciMLStructures.replace! to accumulate it into the parameter shadow.
-        dp_tangent = if originator isa SciMLBase.EnzymeOriginator && isscimlstructure(p)
-            dp
-        else
-            repack_adjoint(dp)[1]
+            if originator isa SciMLBase.EnzymeOriginator && isscimlstructure(p)
+                dp
+            else
+                repack_adjoint(dp)[1]
+            end
         end
 
         return if originator isa SciMLBase.TrackerOriginator ||

src/sensitivity_algorithms.jl

@@ -1293,30 +1293,42 @@ documentation page or the docstrings of the vjp types.
 Johnson, S. G., Notes on Adjoint Methods for 18.336, Online at
 http://math.mit.edu/stevenj/18.336/adjoint.pdf (2007)
 """
-struct SteadyStateAdjoint{CS, AD, FDT, VJP, LS, LK} <:
+struct SteadyStateAdjoint{CS, AD, FDT, VJP, LS, LK, DT <: Val} <:
     AbstractAdjointSensitivityAlgorithm{CS, AD, FDT}
     autojacvec::VJP
     linsolve::LS
     linsolve_kwargs::LK
+    diff_tunables::DT
 end
 
+"""
+    SteadyStateAdjoint(; autojacvec=nothing, linsolve=nothing, diff_tunables=Val(true), ...)
+
+When `diff_tunables = Val(true)` (default), the parameter VJP is computed
+w.r.t. the tunable portion of `p` only. When `diff_tunables = Val(false)`,
+the VJP is computed w.r.t. the full parameter object (including caches,
+initials, etc.). This is needed for SCCNonlinearProblem where `explicitfuns!`
+write active data into non-tunable components. Requires an `autojacvec`
+backend that supports structured parameters (ZygoteVJP, EnzymeVJP,
+MooncakeVJP, ReactantVJP).
+"""
 Base.@pure function SteadyStateAdjoint(;
         chunk_size = 0, autodiff = true,
         diff_type = Val{:central}, autojacvec = nothing, linsolve = nothing,
-        linsolve_kwargs = (;)
+        linsolve_kwargs = (;), diff_tunables = Val(true)
     )
     return SteadyStateAdjoint{
         chunk_size, autodiff, diff_type, typeof(autojacvec),
-        typeof(linsolve), typeof(linsolve_kwargs),
-    }(autojacvec, linsolve, linsolve_kwargs)
+        typeof(linsolve), typeof(linsolve_kwargs), typeof(diff_tunables),
+    }(autojacvec, linsolve, linsolve_kwargs, diff_tunables)
 end
 function setvjp(
-        sensealg::SteadyStateAdjoint{CS, AD, FDT, VJP, LS, LK},
+        sensealg::SteadyStateAdjoint{CS, AD, FDT, VJP, LS, LK, DT},
         vjp
-    ) where {CS, AD, FDT, VJP, LS, LK}
-    return SteadyStateAdjoint{CS, AD, FDT, typeof(vjp), LS, LK}(
+    ) where {CS, AD, FDT, VJP, LS, LK, DT}
+    return SteadyStateAdjoint{CS, AD, FDT, typeof(vjp), LS, LK, DT}(
         vjp, sensealg.linsolve,
-        sensealg.linsolve_kwargs
+        sensealg.linsolve_kwargs, sensealg.diff_tunables
     )
 end
 

src/steadystate_adjoint.jl

@@ -20,10 +20,14 @@ function SteadyStateAdjointSensitivityFunction(
     )
     (; p, u0) = sol.prob
 
+    # When diff_tunables = Val(false), use the full parameter object so
+    # the VJP includes gradients w.r.t. all parameter components.
+    _use_full_p = sensealg.diff_tunables isa Val{false} &&
+        isscimlstructure(p) && !(p isa AbstractArray)
     diffcache,
         y = adjointdiffcache(
         g, sensealg, false, sol, dgdu, dgdp, f, alg;
-        quad = false, needs_jac
+        quad = false, needs_jac, use_full_p = _use_full_p
     )
 
     λ = zero(y)
@@ -161,6 +165,7 @@ end
         end
     end
 
+
     if g !== nothing || dgdp !== nothing
         # compute del g/del p
         if dgdp !== nothing

test/scc_enzyme.jl

@@ -0,0 +1,83 @@
+using Test
+using NonlinearSolve, SCCNonlinearSolve
+using SciMLSensitivity
+using Enzyme
+using FiniteDiff
+import SciMLStructures as SS
+
+# Two-component SCC problem with parameter coupling through caches.
+# Component 1: u1^2 - p[1] = 0 (root: u1 = sqrt(p[1]))
+# Component 2: u2 - p[2]*u1 = 0 (root: u2 = p[2]*u1)
+#   where u1 from component 1 is passed via explicitfun into component 2's
+#   parameter cache.
+
+@testset "SCCNonlinearProblem Enzyme differentiation" begin
+    # Sub-problem 1: u^2 - p = 0
+    function f1(du, u, p)
+        du[1] = u[1]^2 - p[1]
+    end
+    explicitfun1!(p, sols) = nothing
+
+    # Sub-problem 2: u - cache[1] * p[2] = 0
+    # cache[1] will be set to sol1[1] (= sqrt(p[1])) by explicitfun2
+    function f2(du, u, p)
+        du[1] = u[1] - p[1] * p[2]  # p[1] is cache, p[2] is tunable
+    end
+    function explicitfun2!(p, sols)
+        p[1] = sols[1].u[1]  # transfer u1 from component 1 into cache
+        return nothing
+    end
+
+    p_shared = [0.0, 2.0]  # p[1] = cache (written by explicitfun2), p[2] = tunable
+    prob1 = NonlinearProblem(
+        NonlinearFunction{true, SciMLBase.NoSpecialize}(f1), [1.0], p_shared,
+    )
+    prob2 = NonlinearProblem(
+        NonlinearFunction{true, SciMLBase.NoSpecialize}(f2), [1.0], p_shared,
+    )
+
+    sccprob = SciMLBase.SCCNonlinearProblem(
+        [prob1, prob2],
+        SciMLBase.Void{Any}.([explicitfun1!, explicitfun2!]),
+    )
+
+    alg = SCCNonlinearSolve.SCCAlg(nlalg = NewtonRaphson())
+
+    # Forward solve works
+    p_test = [4.0, 3.0]
+    p_shared .= p_test
+    sol = solve(sccprob, alg)
+    @test SciMLBase.successful_retcode(sol)
+    @test sol.u[1] ≈ 2.0 atol = 1.0e-10  # sqrt(4)
+    @test sol.u[2] ≈ 6.0 atol = 1.0e-10  # 3 * 2
+
+    # FiniteDiff ground truth
+    function loss(p_val)
+        p_shared .= p_val
+        sol = solve(sccprob, alg)
+        sum(sol.u)
+    end
+    fd = FiniteDiff.finite_difference_gradient(loss, p_test)
+    @test any(!iszero, fd)
+
+    # Enzyme gradient
+    @testset "Enzyme through SCC" begin
+        loss_enzyme = let sccprob = sccprob, alg = alg, p_shared = p_shared
+            p_val -> begin
+                p_shared .= p_val
+                sol = solve(sccprob, alg)
+                sum(sol.u)
+            end
+        end
+
+        dloss = Enzyme.make_zero(loss_enzyme)
+        dp = zeros(length(p_test))
+        Enzyme.autodiff(
+            Enzyme.set_runtime_activity(Enzyme.Reverse),
+            Enzyme.Duplicated(loss_enzyme, dloss),
+            Enzyme.Active,
+            Enzyme.Duplicated(copy(p_test), dp),
+        )
+        @test isapprox(dp, fd, rtol = 0.05)
+    end
+end