Enzyme reverse-mode broken on NonlinearSolveBase.maybe_wrap_nonlinear_f — autospecialize FunctionWrappers defeats activity analysis (Julia 1.11) #939

Summary

Enzyme.autodiff(Reverse, ...) on a function that calls NonlinearSolve.solve now fails with EnzymeRuntimeActivityError originating from NonlinearSolveBase.maybe_wrap_nonlinear_f (lib/NonlinearSolveBase/src/autospecialize.jl:134) via the FunctionWrappers wrapping it applies to in-place array-state problems.

This is a regression on Julia 1.11: per EnzymeAD/Enzyme.jl#3010, the same MWE used to work on Julia 1.11.9 + Enzyme v0.13.132 (it was the working baseline that confirmed an unrelated 1.12 segfault in Enzyme). With current packages, 1.12 is now fixed but 1.11 has broken in a different way — and the failure points clearly at NonlinearSolveBase rather than Enzyme.

The Enzyme-suggested set_runtime_activity(Reverse) workaround does not help: it advances past the activity-analysis error but hits a second wall (EnzymeNoDerivativeError in DifferentiationInterfaceForwardDiffExt's dual_type path, also called from NonlinearSolveBase).

MWE

using Pkg
Pkg.activate(; temp=true)
Pkg.add([\"Enzyme\", \"NonlinearSolve\"])

using 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))
# Expected: dp == [1.0, -1.0]
# Actual: EnzymeRuntimeActivityError

Error (default Enzyme.Reverse)

ERROR: LoadError: EnzymeRuntimeActivityError: Detected potential need for runtime activity.

Constant memory is stored (or returned) to a differentiable variable and correctness cannot be guaranteed with static activity analysis.

Failure within method:
maybe_wrap_nonlinear_f(::NonlinearProblem{Vector{Float64}, true, Vector{Float64}, NonlinearFunction{true, SciMLBase.AutoSpecialize, ...}, ...})
     @ NonlinearSolveBase ~/.julia/packages/NonlinearSolveBase/uwOVd/src/autospecialize.jl:134

Mismatched activity for:   store i64 %17, i64 addrspace(11)* %28, ...
const val:   %17 = ptrtoint {}* %16 to i64

Stacktrace:
 [1] FunctionWrapper                  @ FunctionWrappers/src/FunctionWrappers.jl:107
 [2] #2                               @ FunctionWrappersWrappers/src/FunctionWrappersWrappers.jl:134
 [3] map                              @ ./tuple.jl:386
 [4] FunctionWrappersWrapper          @ FunctionWrappersWrappers/src/FunctionWrappersWrappers.jl:128
 [5] wrapfun_iip                      @ NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl:88
 [6] maybe_wrap_nonlinear_f           @ NonlinearSolveBase/src/autospecialize.jl:162
 [7] maybe_wrap_f                     @ NonlinearSolveBase/src/solve.jl:777
 ...
 [14] runtime_generic_augfwd          @ Enzyme/src/rules/jitrules.jl:574
 [15] get_concrete_problem            @ NonlinearSolveBase/src/solve.jl:794
 ...
 [19] solve                           @ NonlinearSolveBase/src/solve.jl:92

Error with set_runtime_activity workaround

Swapping Enzyme.Reverse for Enzyme.set_runtime_activity(Enzyme.Reverse) clears the first error but hits:

ERROR: LoadError: EnzymeNoDerivativeError: 
No augmented forward pass found for jl_get_builtin_fptr
 at context:   %23 = call ... @jl_get_builtin_fptr(... @\"ejl_inserted$_Core___compute_sparams_..._false_...\") ...

Stacktrace:
  [1] dual_type                       @ DifferentiationInterfaceForwardDiffExt/utils.jl:31
  [2] prepare_jacobian_nokwarg        @ DifferentiationInterfaceForwardDiffExt/twoarg.jl:388
  [3] prepare_jacobian                @ DifferentiationInterface/src/first_order/jacobian.jl:18
  [4] construct_jacobian_cache        @ NonlinearSolveBase/src/jacobian.jl:62

So there are two stacked obstacles for this code path under Enzyme:

• (a) the FunctionWrappers/FunctionWrappersWrapper wrapping that maybe_wrap_nonlinear_f applies to AutoSpecialize+IIP+array-state problems defeats Enzyme's static activity analysis, and even runtime activity does not recover correctness;
• (b) once past (a), DifferentiationInterface's ForwardDiff dual_type machinery (used in NonlinearSolveBase's jacobian cache construction) is itself not differentiable by Enzyme.

Suspected fix direction

maybe_wrap_nonlinear_f already has several opt-outs (OOP / non-array state / dual eltype / FullSpecialize). It probably needs one more for "called from inside Enzyme reverse-mode AD" — either via:

• an EnzymeRules.inactive/EnzymeRules.augmented_primal rule on maybe_wrap_nonlinear_f that returns the unwrapped prob.f.f (skipping the FunctionWrappersWrapper construction), or
• detecting Enzyme context another way and skipping the wrap, or
• providing a guidance doc-string note that users hitting Enzyme should construct the problem with SciMLBase.FullSpecialize to bypass.

The cleanest option is probably (1) — a no-op EnzymeRules rule on maybe_wrap_nonlinear_f (and the related wrapfun_iip it calls) so the wrapper is never built under reverse-mode AD.

Even with that, the secondary EnzymeNoDerivativeError in construct_jacobian_cache → DifferentiationInterfaceForwardDiffExt.prepare_jacobian → dual_type is a separate problem that will surface immediately after. May need EnzymeRules coverage on the jacobian-cache construction path too, or to ensure the AD backend used inside solve is selected to be compatible when Enzyme is the outer driver.

Versions

• Julia 1.11.9
• Enzyme v0.13.150
• EnzymeCore v0.8.20
• NonlinearSolve v4.19.1
• NonlinearSolveBase v2.28.0
• DifferentiationInterface v0.7.18

Related

• LLVM verification failure: reverse-mode mapreduce on mutated array in loop with break (Julia 1.12) EnzymeAD/Enzyme.jl#3010 — original was a 1.12 segfault in Enzyme over this same MWE; that segfault is now fixed in Enzyme v0.13.150, exposing this NonlinearSolveBase-side issue as the new wall on 1.11. Closed the upstream Enzyme issue and filed this one to track the NonlinearSolveBase side.
• The MWE and reproducer project (with manifest) live at /home/crackauc/sandbox/tmp_20260527_080116_29711/enzyme_verify/3010_v111/ on my machine if a maintainer wants the exact frozen versions.

🤖 Generated with Claude Code

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