Add ChainRulesCore support for SCCNonlinearProblem (#757)

This enables automatic differentiation through SCCNonlinearProblem solves
by adding a ChainRulesCore extension that hooks into SciMLSensitivity's
adjoint machinery.

Changes:
- Add `scc_solve_up` internal function that can be hooked by ChainRulesCore
- Route public `solve` through `scc_solve_up` → `_scc_solve`
- Add ChainRulesCore weak dependency and extension
- Extension defines rrule for `scc_solve_up` that calls `_concrete_solve_adjoint`

This is needed to fix zero gradients when differentiating through MTK v10
initialization which uses SCCNonlinearProblem.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: ChrisRackauckas <[email protected]>
Co-authored-by: Claude <[email protected]>

Author: 3 people

lib/SCCNonlinearSolve/Project.toml

@@ -10,9 +10,16 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+
+[extensions]
+SCCNonlinearSolveChainRulesCoreExt = "ChainRulesCore"
+
 [compat]
 Aqua = "0.8"
 BenchmarkTools = "1.5.0"
+ChainRulesCore = "1"
 CommonSolve = "0.2.4"
 ExplicitImports = "1.5"
 Hwloc = "3"

lib/SCCNonlinearSolve/ext/SCCNonlinearSolveChainRulesCoreExt.jl

@@ -0,0 +1,17 @@
+module SCCNonlinearSolveChainRulesCoreExt
+
+using SCCNonlinearSolve
+using SCCNonlinearSolve: SCCAlg, scc_solve_up
+using SciMLBase: SCCNonlinearProblem, AbstractSensitivityAlgorithm, ChainRulesOriginator,
+    _concrete_solve_adjoint
+
+import ChainRulesCore
+
+function ChainRulesCore.rrule(
+        ::typeof(scc_solve_up), prob::SCCNonlinearProblem,
+        sensealg::Union{Nothing, AbstractSensitivityAlgorithm},
+        u0, p, alg::SCCAlg; kwargs...)
+    _concrete_solve_adjoint(prob, alg, sensealg, u0, p, ChainRulesOriginator(); kwargs...)
+end
+
+end

lib/SCCNonlinearSolve/src/SCCNonlinearSolve.jl

@@ -23,13 +23,30 @@ end
 
 SCCAlg(; nlalg = nothing, linalg = nothing) = SCCAlg(nlalg, linalg)
 
-function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem; kwargs...)
-    CommonSolve.solve(prob, SCCAlg(nothing, nothing); kwargs...)
+function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem;
+        sensealg = nothing, u0 = nothing, p = nothing, kwargs...)
+    CommonSolve.solve(prob, SCCAlg(nothing, nothing); sensealg, u0, p, kwargs...)
 end
 
 function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem,
-        alg::SciMLBase.AbstractNonlinearAlgorithm; kwargs...)
-    CommonSolve.solve(prob, SCCAlg(alg, nothing); kwargs...)
+        alg::SciMLBase.AbstractNonlinearAlgorithm;
+        sensealg = nothing, u0 = nothing, p = nothing, kwargs...)
+    CommonSolve.solve(prob, SCCAlg(alg, nothing); sensealg, u0, p, kwargs...)
+end
+
+function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg::SCCAlg;
+        sensealg = nothing, u0 = nothing, p = nothing, kwargs...)
+    u0 = u0 !== nothing ? u0 : prob.u0
+    p = p !== nothing ? p : prob.p
+    scc_solve_up(prob, sensealg, u0, p, alg; kwargs...)
+end
+
+"""
+Internal solve function that can be hooked by ChainRulesCore for AD.
+"""
+function scc_solve_up(prob::SciMLBase.SCCNonlinearProblem, sensealg, u0, p, alg::SCCAlg;
+        kwargs...)
+    _scc_solve(prob, alg; kwargs...)
 end
 
 probvec(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}) = prob.u0
@@ -60,7 +77,11 @@ function iteratively_build_sols(alg, sols, (prob, explicitfun), args...; kwargs.
     iteratively_build_sols(alg, (sols..., _sol), args...; kwargs...)
 end
 
-function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg::SCCAlg; kwargs...)
+"""
+Internal solve implementation for SCCNonlinearProblem.
+This is called by scc_solve_up and should NOT be hooked by ChainRulesCore.
+"""
+function _scc_solve(prob::SciMLBase.SCCNonlinearProblem, alg::SCCAlg; kwargs...)
     numscc = length(prob.probs)
     sols = iteratively_build_sols(
         alg, (), zip(prob.probs, prob.explicitfuns!)...; kwargs...)
@@ -79,4 +100,6 @@ function CommonSolve.solve(prob::SciMLBase.SCCNonlinearProblem, alg::SCCAlg; kwa
     SciMLBase.build_solution(prob, alg, u, resid; retcode, original = sols)
 end
 
+export scc_solve_up
+
 end