Skip FunctionWrappersWrapper in maybe_wrap_nonlinear_f under Enzyme AD

[ChrisRackauckas-Claude]

Note

Please ignore until reviewed by @ChrisRackauckas.

Summary

Fixes #939: Enzyme.autodiff(Reverse, ...) over a function that calls NonlinearSolve.solve was failing with EnzymeRuntimeActivityError originating from NonlinearSolveBase.maybe_wrap_nonlinear_f. The FunctionWrappersWrapper it builds (for AutoSpecialize + IIP + array-state + non-Dual eltype problems) emits ptrtoint/store patterns that defeat Enzyme's static activity analysis, and set_runtime_activity is not sufficient to recover correctness.

The wrapper exists so the inner ForwardDiff Jacobian gets precompiled dispatch signatures. It is unnecessary on the outer-AD path — the unwrapped function works fine for Enzyme reverse mode.

Fix

Add an EnzymeCore.within_autodiff() short-circuit at the top of maybe_wrap_nonlinear_f that returns the raw prob.f.f whenever we're being called from inside Enzyme reverse-mode AD. EnzymeCore is already a hard dependency, so no new packages are added.

The same opt-out pattern is consistent with the existing Enzyme-aware infrastructure in autospecialize.jl (_uses_enzyme_ad, maybe_unwrap_prob_for_enzyme) — those handle the inner AD case; this handles the outer AD case.

Effect

• Before (Enzyme v0.13.150, NonlinearSolve v4.19.1, Julia 1.11.9): cryptic EnzymeRuntimeActivityError: Detected potential need for runtime activity raised at autospecialize.jl:134 via a FunctionWrappers / FunctionWrappersWrapper activity-mismatch.
• After, depending on what the user has loaded:
  • With nothing else: clean, actionable error from _concrete_solve_adjoint — "Compatibility with reverse-mode automatic differentiation requires SciMLSensitivity.jl." (the standard SciML message users already know)
  • With using SciMLSensitivity (+ using ChainRulesCore to activate NonlinearSolveBaseEnzymeExt): the MWE from Enzyme reverse-mode broken on NonlinearSolveBase.maybe_wrap_nonlinear_f — autospecialize FunctionWrappers defeats activity analysis (Julia 1.11) #939 returns the correct gradient dp == [1.0, -1.0].

MWE verification

The reproducer from #939:

using ChainRulesCore, SciMLSensitivity, Enzyme, NonlinearSolve

u0 = [0.0]
p = [2.0, 1.0]

function simple_loss(p)
    prob = NonlinearProblem((du, u, p) -> du[1] = u[1] - p[1] + p[2], [0.0], p)
    sol = solve(prob, NewtonRaphson())
    return sum(sol.u)
end

dp = Enzyme.make_zero(p)
Enzyme.autodiff(Enzyme.Reverse, simple_loss, Enzyme.Active, Enzyme.Duplicated(p, dp))
# Before: EnzymeRuntimeActivityError
# After:  dp == [1.0, -1.0]  ✓

Verified locally on Julia 1.11.9 with Enzyme v0.13.150 and dev-checkout of this branch.

Test

Added a regression test in lib/NonlinearSolveBase/test/runtests.jl that probes maybe_wrap_nonlinear_f from inside Enzyme.autodiff(ReverseWithPrimal, ...) and asserts the result is unwrapped. Test deps gain Enzyme. Full NonlinearSolveBase test suite passes (25/25).

Related

• Closes Enzyme reverse-mode broken on NonlinearSolveBase.maybe_wrap_nonlinear_f — autospecialize FunctionWrappers defeats activity analysis (Julia 1.11) #939
• Cross-references LLVM verification failure: reverse-mode mapreduce on mutated array in loop with break (Julia 1.12) EnzymeAD/Enzyme.jl#3010 (now closed; original was a Julia 1.12 segfault that has been fixed in Enzyme, exposing this NonlinearSolveBase-side issue as the wall on 1.11)

🤖 Generated with Claude Code

[ChrisRackauckas-Claude]

Filed upstream Enzyme issue for the macOS LTS GC corruption that the heavier test form triggered: EnzymeAD/Enzyme.jl#3130. The probe-style regression test in this PR avoids the trigger and verifies the NonlinearSolveBase fix directly.