feat: AutoSpecialize norecompile infrastructure for NonlinearSolveBase

[ChrisRackauckas-Claude]

Summary

Ports the FunctionWrappersWrappers-based norecompile infrastructure from DiffEqBase to NonlinearSolveBase, following the approach described in https://sciml.ai/news/2022/09/21/compile_time/.

This PR adds the infrastructure only — the wrappers, tag types, and extension methods needed for the norecompile/AutoSpecialize pattern. The automatic wrapping at solve-time is not yet activated because NonlinearSolve has several direct ForwardDiff.jacobian call sites (sensitivity analysis df/dp, df/du, bounds transforms) that bypass the DI-based Jacobian path and would receive duals with mismatched chunk sizes.

What's included

• src/autospecialize.jl (new): NonlinearSolveTag, wrapfun_iip/wrapfun_oop stub methods, maybe_wrap_nonlinear_f, standardize_forwarddiff_tag fallback
• ForwardDiff extension (NonlinearSolveBaseForwardDiffExt.jl): Dual-aware wrapfun_iip/wrapfun_oop dispatches with 6 type combinations each (Float64, Dual{NonlinearSolveTag}, NullParameters). Tag standardization that stamps NonlinearSolveTag on AutoForwardDiff and forces chunksize=1 when the function is wrapped via EvalFunc.
• NonlinearSolveBase.jl: FunctionWrappers/FunctionWrappersWrappers imports, exports for the new public API
• Project.toml: FunctionWrappers and FunctionWrappersWrappers dependencies

What's NOT included (deferred)

• Automatic wrapping in get_concrete_problem / solve path — requires standardizing all direct ForwardDiff call sites first (nonlinearsolve_∂f_∂p, nonlinearsolve_∂f_∂u, bounds transform code)
• Tag standardization at Jacobian construction or NLLS VJP generation — same reason

Design notes

The infrastructure mirrors DiffEqBase's pattern:

• For standard problem types (Vector{Float64} state, Vector{Float64} or NullParameters parameters), maybe_wrap_nonlinear_f wraps the function in a FunctionWrappersWrapper with precompiled dual type signatures
• standardize_forwarddiff_tag coordinates the ForwardDiff tag (NonlinearSolveTag) and chunk size (N=1) so duals match the wrapper signatures
• Non-standard types (Float32, StaticArray, scalar) pass through unchanged

Next steps to activate

1. Route all direct ForwardDiff.jacobian/derivative/gradient calls through DI or standardize their tag/chunksize
2. Wire maybe_wrap_f into get_concrete_problem
3. Add standardize_forwarddiff_tag before construct_concrete_adtype in jacobian.jl

[ChrisRackauckas-Claude]

CI Fix: ImmutableNonlinearProblem wrapping

The adjoint test in SimpleNonlinearSolve failed because @set! (Setfield) couldn't reconstruct ImmutableNonlinearProblem with the wrapped function type. The constructor doesn't accept EvalFunc{FunctionWrappersWrapper{...}}.

Fix: Skip FunctionWrapper wrapping for ImmutableNonlinearProblem. SimpleNonlinearSolve's lighter solvers don't benefit from the norecompile pathway anyway.

The other CI failures (runic, alloc_check, wrappers) are pre-existing/infrastructure issues:

• runic: Formatting issue in lib/SCCNonlinearSolve/ (not our files)
• alloc_check, wrappers: Unable to locate executable file: julia (CI runner issue)

[ChrisRackauckas-Claude]

Reverted automatic wrapping - infrastructure only for now

The core and wrappers CI tests failed with No matching function wrapper was found! because the automatic wrapping at get_concrete_problem time applies the N=1 FunctionWrappers to ALL code paths, but several paths call ForwardDiff directly with default chunk sizes:

1. nonlinearsolve_∂f_∂p → ForwardDiff.jacobian(f2, p) (direct call)
2. nonlinearsolve_∂f_∂u → ForwardDiff.jacobian(...) (direct call)
3. Bounds transform code wraps and applies ForwardDiff

These bypass standardize_forwarddiff_tag, so they use chunksize based on problem dimension (N=2 for 2-element vectors), which doesn't match the N=1 wrappers.

Current state: The AutoSpecialize infrastructure is fully implemented and ready:

• autospecialize.jl: NonlinearSolveTag, maybe_wrap_nonlinear_f, wrapfun_iip/oop
• ForwardDiff extension: Dual-aware wrappers, standardize_forwarddiff_tag
• Jacobian cache: Tag standardization integrated

What's needed for automatic wrapping: All ForwardDiff call sites need to be coordinated to use chunksize=1 when the function is wrapped. This is analogous to how DiffEqBase/OrdinaryDiffEq controls all internal ForwardDiff calls.

[ChrisRackauckas-Claude]

Update: AutoSpecialize wrapping now fully activated

This commit activates the wrapping infrastructure end-to-end. Key design decisions:

Architecture

• AutoSpecializeCallable{FW} holds both the FunctionWrappersWrapper (precompiled dispatch) and the original function (orig::Any, type-erased for precompilation)
• No EvalFunc/invokelatest — wrapping sets prob.f.f directly, avoiding the world-age indirection that would break IIP @inferred
• Try-catch fallback in AutoSpecializeCallable.__call__: when FunctionWrappersWrapper throws NoFunctionWrapperFoundError (mismatched dual tags), falls back to the original function

What gets wrapped

• IIP NonlinearProblem: wrapped (both Vector{Float64} state + params)
• OOP NonlinearProblem: wrapped
• IIP NonlinearLeastSquaresProblem: wrapped
• OOP NonlinearLeastSquaresProblem: NOT wrapped (return type may differ from u0)
• ImmutableNonlinearProblem: NOT wrapped

Tag standardization

• Main autodiff → standardize_forwarddiff_tag stamps NonlinearSolveTag + forces chunksize=1 when wrapped
• jvp_autodiff/vjp_autodiff → also standardized in construct_jacobian_cache before passing to JacobianOperator
• AutoPolyesterForwardDiff → replaced with AutoForwardDiff{1, tag} when wrapped (no custom tag support)
• Sensitivity analysis (∂f/∂u, ∂f/∂p) → use ForwardDiff.JacobianConfig with chunksize=1 + NonlinearSolveTag when wrapped

Known trade-off

• OOP @inferred regresses — try-catch with orig::Any makes return type Any. Same trade-off as DiffEqBase. Updated test to @test_broken.

Local test results (all clean)

Suite	Pass	Broken	Fail	Error
Core	727	88	0	0
Wrapper	195	7	0	0
ForwardDiff	135,636	0	0	0

[ChrisRackauckas-Claude]

Adjoint/Reverse-Mode AD Fix (commit 7abe63b)

The nopre CI jobs were failing with llvmcall must be compiled to be called because reverse-mode AD backends (Zygote, Mooncake, Enzyme) cannot differentiate through FunctionWrapper internals.

Root Cause

When solve_up is called inside an rrule, SciMLSensitivity._concrete_solve_adjoint internally calls solve, which goes through get_concrete_problem → maybe_wrap_f → wraps with FunctionWrapper again. Mooncake in particular compiles tangent rules for ALL types during the forward pass, so the FunctionWrapper type must never appear in the computation graph.

Fix

Two-pronged approach:

1. ChainRulesCore.rrule for AutoSpecializeCallable: Redirects reverse-mode AD through f.orig (the unwrapped callable) instead of f.fw (the FunctionWrapper)
2. _DISABLE_AUTOSPECIALIZE flag: Set to true in the solve_up rrule before calling _solve_adjoint. This prevents maybe_wrap_nonlinear_f from wrapping during the entire adjoint code path, ensuring FunctionWrapper types never enter the AD computation graph.

Test Results (local, Julia 1.10)

Suite	Pass	Broken	Fail	Error
Core	727	89	0	0
Wrapper	195	7	0	0
Adjoint	3	0	0	0
ForwardDiff	135,636	0	0	0

Also updated the IIP @inferred test to @test_broken (same wrapping-induced type inference regression as OOP).

[ChrisRackauckas-Claude]

Fix: Extend Enzyme unwrap to JacobianOperator path

The previous push fixed Enzyme compatibility for the concrete Jacobian path (DI.prepare_jacobian/DI.jacobian), but CI still had 112 errors in NonlinearSolveFirstOrder and NonlinearSolveQuasiNewton (ubuntu 1.10/1.11) from the JacobianOperator path.

Root cause analysis from CI logs:
The EnzymeMutabilityException stacktrace showed:

FunctionWrappers.jl:137 → AutoSpecializeCallable → NonlinearFunction → 
SciMLJacobianOperators.jl:404 (DI.pushforward!) → JacobianOperator → 
mul! → Krylov.gmres!

When concrete_jac=false (used with Krylov solvers like KrylovJL_GMRES() and \ linsolve), construct_jacobian_cache creates a JacobianOperator by passing prob directly. The SciMLJacobianOperators package then calls DI.pushforward!/DI.pullback! with prob.f (still containing AutoSpecializeCallable), triggering Enzyme's mutability exception.

Fix: In construct_jacobian_cache's !needs_jac branch, check if either jvp_autodiff or vjp_autodiff uses Enzyme, and if so, create a modified prob with the unwrapped raw function via @set prob.f.f = get_raw_f(f.f) before passing to JacobianOperator.

Verified locally:

• NewtonRaphson(linsolve=\, autodiff=AutoEnzyme(), concrete_jac=Val(false)) → Success ✓
• NewtonRaphson(linsolve=KrylovJL_GMRES(), autodiff=AutoEnzyme(), concrete_jac=Val(false)) → Success ✓
• TrustRegion(linsolve=KrylovJL_GMRES(), autodiff=AutoEnzyme()) → Success ✓

[ChrisRackauckas-Claude]

Additional Enzyme fix: TrustRegion VecJac/JacVec operators

The previous commit fixed the JacobianOperator path in construct_jacobian_cache, but CI still showed Enzyme failures because the TrustRegion scheme creates its own VecJacOperator and JacVecOperator directly from the problem at trust_region.jl:227-230, bypassing construct_jacobian_cache entirely.

Root cause from CI log:

SciMLJacobianOperators.jl:335  # VJP closure with AutoSpecializeCallable
  ← trust_region.jl:233        # StatefulJacobianOperator * fu
  ← trust_region.jl:227        # VecJacOperator(prob, fu, u; autodiff = vjp_autodiff)

Fix: Added Enzyme unwrap check before creating VecJac/JacVec operators in trust_region.jl, using the same pattern as in jacobian.jl.

[ChrisRackauckas-Claude]

Note: Enzyme workaround simplification

The Enzyme-specific workarounds in this PR (_uses_enzyme_ad, maybe_unwrap_prob_for_enzyme, and the _ad_prob unwrapping in __init methods) exist because Enzyme cannot differentiate through FunctionWrappers' llvmcall/ccall.

A companion PR has been opened at EnzymeAD/Enzyme.jl#2980 that adds an EnzymeFunctionWrappersExt extension. This extension defines EnzymeRules for FunctionWrapper that automatically extract the wrapped function and delegate to autodiff_deferred, making Enzyme transparently differentiate through FunctionWrappers.

Once that Enzyme PR is merged and released, the following code can be removed from this PR (~73 lines):

• _uses_enzyme_ad() and maybe_unwrap_prob_for_enzyme() in autospecialize.jl
• Two Enzyme unwrap blocks in jacobian.jl
• _ad_prob construction in NonlinearSolveFirstOrder/src/solve.jl
• _ad_autodiffs array + _ad_prob construction in NonlinearSolveQuasiNewton/src/solve.jl

All references to _ad_prob would revert back to prob.

[ChrisRackauckas]

Waiting on the Enzyme PR

[ChrisRackauckas-Claude]

Session State Summary (for continuation)

What was done

Enzyme.jl PR #2980 — EnzymeAD/Enzyme.jl#2980

• Created EnzymeFunctionWrappersExt extension that defines EnzymeRules for FunctionWrapper
• Forward mode rule: extracts fw.obj[] and delegates to Enzyme.autodiff_deferred
• Reverse mode rules: augmented_primal + reverse that extract and differentiate the unwrapped function
• Handles both IIP (FunctionWrapper{Nothing,...}) and OOP (FunctionWrapper{Ret,...}) patterns
• All 8 tests pass locally (IIP/OOP × Forward/Reverse)
• Files: ext/EnzymeFunctionWrappersExt.jl, test/ext/functionwrappers.jl, Project.toml edits
• Branch: functionwrappers-ext on EnzymeAD/Enzyme.jl (pushed from ChrisRackauckas-Claude fork)
• CI status: action_required — needs maintainer to approve workflow run
• Local repo: /home/crackauc/sandbox/tmp_20260218_173108_96644/Enzyme.jl

NonlinearSolve.jl PR #838 — #838

• Branch: autospecialize on SciML/NonlinearSolve.jl (pushed via ChrisRackauckas-Claude fork remote botfork)
• Adds AutoSpecializeCallable / FunctionWrappersWrapper infrastructure for norecompile mode
• Includes Enzyme-specific workarounds: _uses_enzyme_ad(), maybe_unwrap_prob_for_enzyme(), _ad_prob indirection in FirstOrder/QuasiNewton solvers
• 105/105 test jobs passing on CI (3 "failures" are infrastructure: doc deploy, codecov, trim — not code)
• Local repo: /home/crackauc/sandbox/tmp_20260218_173108_96644/NonlinearSolve.jl
• Working tree: clean

Simplification attempt (reverted)

Attempted to remove ~73 lines of Enzyme workaround code from 4 files:

1. lib/NonlinearSolveBase/src/autospecialize.jl — remove _uses_enzyme_ad() (3 methods) and maybe_unwrap_prob_for_enzyme()
2. lib/NonlinearSolveBase/src/jacobian.jl — remove 2 Enzyme unwrap blocks
3. lib/NonlinearSolveFirstOrder/src/solve.jl — remove _ad_prob construction
4. lib/NonlinearSolveQuasiNewton/src/solve.jl — remove _ad_autodiffs/_ad_prob construction

Result: 116 EnzymeMutabilityException failures across ALL AD backends (not just Enzyme) because BackTracking{AutoEnzyme} line search uses Enzyme even when the main Jacobian uses ForwardDiff/FiniteDiff. Reverted all changes with git checkout --.

Dependency chain: Enzyme PR #2980 must be merged and released BEFORE these workarounds can be removed.

Key technical details

• FunctionWrapper stores original function in fw.obj[] (a Ref)
• Enzyme cannot differentiate through ccall/llvmcall used by FunctionWrappers — throws EnzymeMutabilityException
• The workarounds detect Enzyme AD via _uses_enzyme_ad() and unwrap AutoSpecializeCallable back to the raw function
• BackTracking{AutoEnzyme} is the default line search, so ALL solver configurations (even ForwardDiff Jacobians) hit the Enzyme code path
• ForwardDiff-only tests pass fine without workarounds (ForwardDiff doesn't use ccall)

Remotes

NonlinearSolve.jl:

• origin = SciML/NonlinearSolve.jl
• botfork = ChrisRackauckas-Claude/NonlinearSolve.jl

Enzyme.jl:

• origin = ChrisRackauckas-Claude/Enzyme.jl
• upstream = EnzymeAD/Enzyme.jl

Next steps

1. Wait for Enzyme PR #2980 to get CI approved and reviewed/merged
2. Once merged and a new Enzyme.jl version is released, re-apply the simplification to NonlinearSolve PR feat: AutoSpecialize norecompile infrastructure for NonlinearSolveBase #838 (remove the ~73 lines from the 4 files listed above)
3. Both PRs need maintainer review

Co-Authored-By: Chris Rackauckas accounts@chrisrackauckas.com