Trace the nonlinear solver loops under Reactant

[ChrisRackauckas-Claude]

Draft. Ignore this PR until reviewed by @ChrisRackauckas.

Stack: SciML/SciMLBase.jl#1563 (retcode type parameter) → this PR → the SimpleNonlinearSolve PR (single branchless loops). Requires EnzymeAD/Reactant.jl#3232 (traced enums) — the Reactant test group pins the branch until it is released.

What changed and why

This replaces the previous content of this PR. Every nonlinear algorithm's solve dispatch now runs under Reactant.@compile/@jit by tracing the ordinary code path: there is one implementation per solver, and the Reactant-specific machinery is three helpers in NonlinearSolveBase/src/reactant.jl plus a few one-line ReactantCore.within_compile() gates where the host semantics genuinely cannot be traced.

• Return codes are traced enums. With Reactant tracing ReturnCode.T (Trace Base.Enum values as an enum wrapper around the traced base integer EnzymeAD/Reactant.jl#3232) and NonlinearSolution.retcode parametric (SciMLBase#1563), cache.retcode = tc_cache.retcode inside a @trace if and ifelse(converged, ReturnCode.Success, cache.retcode) work as written. The shadow _ReactantNonlinearSolution, its fabricated NLStats, InternalAPI.build_reactant_solution, InternalAPI.prepare_reactant_loop_state! and the converged/finite reconstruction are gone; build_nonlinear_solution (the positional NonlinearSolution constructor) is used on every path. Under compilation stats is nothing (NLStats holds Ints and would only count trace-time evaluations) and prob is nothing (a problem's keyword arguments are a Base.Pairs, which Reactant's result codegen cannot rebuild — it treats it as an AbstractDict and calls D()).
• Loop state. maybe_traced(x) promotes a scalar only during compilation; caches are @concrete, so nsteps/force_stop/retcode are promoted where the cache is constructed (FirstOrder, QuasiNewton, SpectralMethods) and in the termination cache. dealias_traced!(cache) is one reflective walk that gives every traced leaf reachable from the cache a fresh copy before the loop and after each step, replacing eleven hand-written per-cache field lists (Newton, steepest, damped Newton, geodesic, dogleg, first-order, trust region, LM damping/trust region, pseudo-transient). Reactant records one path per traced object among loop-carried values and requires the same set after the body, so aliases such as u_cache === u (a @bb copy on a traced array rebinds), the shared p, and the problem stored in the trace are exactly what breaks; the trace now stores nothing for the problem under compilation since it only inspects its type.
• Shared loop and termination. _run_cache_to_completion! is a single ReactantCore.@trace while over CommonSolve.step!; the post-loop MaxIters/Success decision is an ifelse, so it is one path. Termination and check_and_update! set retcode/force_stop with ifelse; no @trace if captures the mutable cache. Two Reactant facts shaped this and are recorded as comments: inside a @trace body ReactantCore.within_compile() is false (the macro captures the module as a loop variable), so loop bodies only call self-gating helpers; and a @trace if binds its captured variables through a parameter named args, so nothing in scope may use that name.
• First-order step and globalization. Accept/reject in the trust-region step, the __Simple radius update, Levenberg–Marquardt acceptance and damping, geodesic acceleration and the descent selection in dogleg are written with ifelse/select once, valid on the host and traced. Where the host path has a real early-out that saves work (dogleg's Newton/Cauchy returns, the isnan shortcut in the trust-region ratio) it is kept behind !within_compile() with the branchless form as the fallthrough; dogleg's Cauchy step is factored into dogleg_cauchy_step! so it is not duplicated. make_new_jacobian stays a host Bool (within_compile() || accepted): under compilation the Jacobian is recomputed every step.
• Polyalgorithms are traced as a chain: every member is compiled, a later one runs only if the earlier ones did not succeed, and the smallest residual is kept as the fallback, mirroring __generated_polysolve. The previous branch silently ran only the first member.
• Jacobian reuse (Add adaptive Jacobian reuse to first-order solvers #1072). Rebased onto Add adaptive Jacobian reuse to first-order solvers #1072: its reuse policy decides make_new_jacobian from residual-norm ratios, which are traced under Reactant, so __init switches the policy off (without_jacobian_reuse) during compilation and every step recomputes the Jacobian there; on the host the policy is untouched, including the stale-Jacobian retries.
• AD selection prefers forward-mode AutoEnzyme during compilation (unchanged from the previous branch); maybe_wrap_f and the ReactantLinearSolveCache are also kept.
• SciMLJacobianOperators.JacobianOperator{iip, T} accepts T <: Number (was Real) so that traced element types can build the trust-region operators; its tests pass unchanged.
• maybe_traced, dealias_traced!, select and build_nonlinear_solution are public and documented (NonlinearSolveBase 2.50.0 — 2.49.0 was released from master meanwhile; the sublibraries' compat and patch versions are bumped accordingly). ReactantCore.is_traced is allow-listed in the QA check with the reason: it is the predicate behind @trace, not declared public upstream, and cannot be reproduced locally because Reactant defines its methods.
• A root AGENTS.md records the Reactant conventions and the repository pitfalls hit while doing this (sublibrary [sources] paths, Pkg.develop into a sublibrary project, @trace body rules), per the agent-memory rule.
• build_solution_less_specialize and SCCNonlinearSolve spelled out all of NonlinearSolution's type parameters; they now pick the layout at load time (fieldtype(NonlinearSolution, :retcode)), so this works with both released SciMLBase and #1563.

Deliberately unchanged semantics under compilation, still documented in nonlinear_solve_gpus.md: the default termination mode is AbsNormTerminationMode rather than the SafeBest variant (its stall bookkeeping is host control flow), initialization failure is a host event and cannot be returned from a compiled solve, and LineSearch-globalized and DFSane paths are not traced.

Against master (after #1072): 35 files, +876/−295 (the previous branch content was +1254/−212 against master; this replaces it with a smaller and more uniform change).

Verification

Local, Julia 1.12.4 unless noted, with EnzymeAD/Reactant.jl#3232 (~/sandbox/Reactant.jl, branch traced-enums) and SciML/SciMLBase.jl#1563 developed into the environments and DifferentiationInterface pinned to 571fc178 as in test/runtests.jl.

Reactant group on this branch (test/Reactant/reactant_tests.jl as in this PR, i.e. without the Simple solvers):

Test Summary:               | Pass  Total     Time
Reactant Integration (core) |  102    102  6m12.5s

The same environment with the stacked SimpleNonlinearSolve PR applied (full matrix):

Test Summary:        | Pass  Total     Time
Reactant Integration |  147    147  8m42.9s

Host groups, run in each sublibrary's own project:

GROUP=Core  NonlinearSolveBase          Testing NonlinearSolveBase tests passed
GROUP=QA    NonlinearSolveBase          Testing NonlinearSolveBase tests passed
GROUP=Core  NonlinearSolveFirstOrder    Testing NonlinearSolveFirstOrder tests passed
GROUP=QA    NonlinearSolveFirstOrder    Testing NonlinearSolveFirstOrder tests passed
GROUP=Core  NonlinearSolveQuasiNewton   Testing NonlinearSolveQuasiNewton tests passed
GROUP=Core  NonlinearSolveSpectralMethods  Testing NonlinearSolveSpectralMethods tests passed
GROUP=Core  SCCNonlinearSolve           Testing SCCNonlinearSolve tests passed
GROUP=Core  SciMLJacobianOperators      Testing SciMLJacobianOperators tests passed
GROUP=Core  SimpleNonlinearSolve (master's sources, against this branch's NonlinearSolveBase)  Testing SimpleNonlinearSolve tests passed

Formatted with Runic 1.10.0; typos clean on the changed files. The final commit differs from the tree the runs above used only by public declarations, version/compat bumps, docs entries and test comments.

Not traced, and therefore not in the Reactant matrix, with the reason next to the matrix in the test file: the least-squares polyalgorithms and default least-squares solve (a line-search member calls norm(x, Inf), whose Reactant overload scalar-indexes — a Reactant bug, one line in src/stdlibs/LinearAlgebra.jl), and RobustMultiNewton (its Yuan/Bastin members need a reverse-mode vector-Jacobian product, which DifferentiationInterface's Enzyme backend does not route through Reactant). Both compiled under the previous branch only because the polyalgorithm collapsed to its first member.

Not verified

• GROUP=Everything; CUDA/GPU groups; the pre job.
• Downstream packages consuming NonlinearSolution's explicit type parameters other than the two sites in this repo.

Reviewer judgment points

• sol.stats === nothing under compilation rather than trace-time counts.
• Polyalgorithms compile every member (compile time grows with the number of members).
• The 1.10 @inferred situation in SimpleNonlinearSolve is handled in the stacked PR.

---

🤖 Generated with Claude Code (model: claude-fable-5) on behalf of Chris Rackauckas. Session: https://claude.ai/code/session_016LsC6pp9z6s5EABX9DnVjE

[ChrisRackauckas-Claude]

Rework plan for this draft:

• Remove the solver allowlist and bespoke Reactant solve interception so algorithms use their normal solve dispatch.
• Remove the local differentiation overloads; derivative-backed algorithms should rely on upstream backend support and surface upstream errors normally.
• Make the derivative-free SimpleBroyden and SimpleKlement loops traceable and test those algorithms as complete compiled solves.
• Record failing-before and passing-after evidence, then run Reactant, Core, Runic, typos, QA, and docs locally before the next push.

QA currently fails during environment resolution on unchanged upstream master because the QA environment requires NLsolve 4.x while NonlinearSolve requires NLsolve 5.x; that baseline regression is being bisected separately.

[ChrisRackauckas-Claude]

CI triage: the only red lanes so far are the three Wrappers jobs, which also fail on master's latest CI_NonlinearSolve run (6 Wrappers failures there; the released-NLsolve regression addressed by #1209). Everything else is pending or green.

🤖 Posted by an AI agent (Claude Code, model claude-fable-5) on behalf of Chris Rackauckas. Session: https://claude.ai/code/session_016LsC6pp9z6s5EABX9DnVjE

[ChrisRackauckas-Claude]

Rebuilt on current master (head 90ccb4245): the previous push had been generated from a tree one commit behind and silently reverted #1211's SCC line; that is fixed and SCC's Core group passes again locally. The layout-agnostic NonlinearSolution constructors are also carved out as #1214 so they can be released ahead of SciMLBase#1563; this PR will be rebased on it once merged.

🤖 Posted by an AI agent (Claude Code, model claude-fable-5) on behalf of Chris Rackauckas. Session: https://claude.ai/code/session_016LsC6pp9z6s5EABX9DnVjE

[ChrisRackauckas-Claude]

Rebased onto master after #1072 (adaptive Jacobian reuse). The reuse policy decides make_new_jacobian from residual-norm ratios, which are traced under Reactant, so __init switches it off (without_jacobian_reuse) during compilation and every traced step recomputes the Jacobian; on the host the policy, including the stale-Jacobian retries, is untouched. Versions: NonlinearSolveBase 2.50.0 (2.49.0 was released meanwhile), FirstOrder 2.5.1, SCC 1.15.3, QuasiNewton 1.15.3. Re-validated locally: Base Core+QA, FirstOrder Core+QA, Reactant group 102/102 (this PR) and 147/147 (with #1213), Simple Core 1.10/1.12 + QA.

🤖 Posted by an AI agent (Claude Code, model claude-fable-5) on behalf of Chris Rackauckas. Session: https://claude.ai/code/session_016LsC6pp9z6s5EABX9DnVjE

[ChrisRackauckas-Claude]

CI triage: the red OrdinaryDiffEq.jl/* lanes (Downstream workflow) fail during environment resolution — Unsatisfiable requirements detected for package AlgebraicMultigrid against NonlinearSolveBase's TimerOutputs = "1" compat — and they fail identically on master's own push run 33303882132 (and on #1072's last run), after having passed earlier the same day, so it is a registry-side change unrelated to this PR. It is being investigated separately.

🤖 Posted by an AI agent (Claude Code, model claude-fable-5) on behalf of Chris Rackauckas. Session: https://claude.ai/code/session_016LsC6pp9z6s5EABX9DnVjE