Unwrap AutoSpecializeCallable for Enzyme at __init entry points

Instead of fixing individual call sites (trust_region.jl VecJac/JacVec,
jacobian.jl construct_jacobian_cache), create _ad_prob with unwrapped
function early in __init for both FirstOrder and QuasiNewton. This
ensures ALL downstream AD consumers (Jacobian cache, trust region,
linesearch, forcing) receive the unwrapped problem when Enzyme is used.

- Add maybe_unwrap_prob_for_enzyme helper in NonlinearSolveBase
- FirstOrder: create _ad_prob from alg.autodiff/jvp_autodiff/vjp_autodiff
- QuasiNewton: detect Enzyme from kwargs and alg.linesearch/trustregion
- Revert trust_region.jl inline fix (now handled upstream in solve.jl)

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

Author: ChrisRackauckas

lib/NonlinearSolveBase/src/autospecialize.jl

@@ -92,6 +92,24 @@ _uses_enzyme_ad(::ADTypes.AutoEnzyme) = true
 _uses_enzyme_ad(ad::AutoSparse) = _uses_enzyme_ad(ADTypes.dense_ad(ad))
 _uses_enzyme_ad(_) = false
 
+"""
+    maybe_unwrap_prob_for_enzyme(prob, autodiffs...)
+
+If the problem function is wrapped by AutoSpecialize and any of the given AD backends
+use Enzyme, return a copy of `prob` with the unwrapped raw function. Otherwise return
+`prob` unchanged.
+
+This should be called early in `__init` so that all downstream AD-related constructions
+(Jacobian cache, trust region operators, linesearch, forcing) receive the unwrapped problem.
+"""
+function maybe_unwrap_prob_for_enzyme(prob, autodiffs...)
+    is_fw_wrapped(prob.f.f) || return prob
+    for ad in autodiffs
+        _uses_enzyme_ad(ad) && return @set prob.f.f = get_raw_f(prob.f.f)
+    end
+    return prob
+end
+
 # Default dispatch assumes no ForwardDiff loaded.
 # The ForwardDiff extension overrides these with dual-aware versions.
 

lib/NonlinearSolveFirstOrder/src/solve.jl

@@ -176,6 +176,12 @@ function SciMLBase.__init(
         verbose = NonlinearVerbosity(verbose)
     end
 
+    # Enzyme cannot differentiate through FunctionWrappers' llvmcall.
+    # Create unwrapped prob for all AD-related constructions when using Enzyme.
+    _ad_prob = NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(
+        prob, alg.autodiff, alg.jvp_autodiff, alg.vjp_autodiff
+    )
+
     timer = get_timer_output()
     @static_timeit timer "cache construction" begin
         u = Utils.maybe_unaliased(prob.u0, alias_u0)
@@ -191,7 +197,7 @@ function SciMLBase.__init(
         linsolve_kwargs = merge((; verbose = verbose.linear_verbosity, abstol, reltol), linsolve_kwargs)
 
         jac_cache = NonlinearSolveBase.construct_jacobian_cache(
-            prob, alg, prob.f, fu, u, prob.p;
+            _ad_prob, alg, _ad_prob.f, fu, u, _ad_prob.p;
             stats, alg.autodiff, linsolve, alg.jvp_autodiff, alg.vjp_autodiff
         )
         J = reused_jacobian(jac_cache, u)
@@ -219,7 +225,7 @@ function SciMLBase.__init(
             NonlinearSolveBase.supports_trust_region(alg.descent) ||
                 error("Trust Region not supported by $(alg.descent).")
             trustregion_cache = InternalAPI.init(
-                prob, alg.trustregion, prob.f, fu, u, prob.p;
+                _ad_prob, alg.trustregion, _ad_prob.f, fu, u, _ad_prob.p;
                 alg.vjp_autodiff, alg.jvp_autodiff, stats, internalnorm, kwargs...
             )
             globalization = Val(:TrustRegion)
@@ -229,7 +235,7 @@ function SciMLBase.__init(
             NonlinearSolveBase.supports_line_search(alg.descent) ||
                 error("Line Search not supported by $(alg.descent).")
             linesearch_cache = CommonSolve.init(
-                prob, alg.linesearch, fu, u; stats, internalnorm,
+                _ad_prob, alg.linesearch, fu, u; stats, internalnorm,
                 autodiff = ifelse(
                     provided_jvp_autodiff, alg.jvp_autodiff, alg.vjp_autodiff
                 ),
@@ -240,7 +246,7 @@ function SciMLBase.__init(
 
         if has_forcing
             forcing_cache = InternalAPI.init(
-                prob, alg.forcing, fu, u, u, prob.p; stats, internalnorm,
+                _ad_prob, alg.forcing, fu, u, u, _ad_prob.p; stats, internalnorm,
                 autodiff = ifelse(
                     provided_jvp_autodiff, alg.jvp_autodiff, alg.vjp_autodiff
                 ),

lib/NonlinearSolveFirstOrder/src/trust_region.jl

@@ -223,20 +223,11 @@ function InternalAPI.init(
     p1, p2, p3, p4 = get_parameters(T, alg.method)
     ϵ = T(1.0e-8)
 
-    # Enzyme cannot differentiate through FunctionWrappers' llvmcall.
-    # Unwrap AutoSpecializeCallable so DI sees the raw user function.
-    _ad_prob = prob
-    if NonlinearSolveBase.is_fw_wrapped(prob.f.f) &&
-            (NonlinearSolveBase._uses_enzyme_ad(vjp_autodiff) ||
-                NonlinearSolveBase._uses_enzyme_ad(jvp_autodiff))
-        @set! _ad_prob.f.f = NonlinearSolveBase.get_raw_f(prob.f.f)
-    end
-
     vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ?
-        VecJacOperator(_ad_prob, fu, u; autodiff = vjp_autodiff) : nothing
+        VecJacOperator(prob, fu, u; autodiff = vjp_autodiff) : nothing
 
     jvp_operator = alg.method isa RUS.__Bastin ?
-        JacVecOperator(_ad_prob, fu, u; autodiff = jvp_autodiff) : nothing
+        JacVecOperator(prob, fu, u; autodiff = jvp_autodiff) : nothing
 
     if alg.method isa RUS.__Yuan
         Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * Utils.safe_vec(fu)

lib/NonlinearSolveQuasiNewton/src/solve.jl

@@ -160,6 +160,24 @@ function SciMLBase.__init(
         alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
     end
     alias_u0 = alias.alias_u0
+    # Enzyme cannot differentiate through FunctionWrappers' llvmcall.
+    # QuasiNewton doesn't have alg.autodiff fields; autodiff may come through kwargs
+    # or from the linesearch/trustregion algorithm's own autodiff field.
+    _ad_autodiffs = Any[
+        get(kwargs, :autodiff, nothing),
+        get(kwargs, :jvp_autodiff, nothing),
+        get(kwargs, :vjp_autodiff, nothing),
+    ]
+    if alg.linesearch !== missing && alg.linesearch !== nothing &&
+            hasfield(typeof(alg.linesearch), :autodiff)
+        push!(_ad_autodiffs, alg.linesearch.autodiff)
+    end
+    if alg.trustregion !== missing && alg.trustregion !== nothing &&
+            hasfield(typeof(alg.trustregion), :autodiff)
+        push!(_ad_autodiffs, alg.trustregion.autodiff)
+    end
+    _ad_prob = NonlinearSolveBase.maybe_unwrap_prob_for_enzyme(prob, _ad_autodiffs...)
+
     timer = get_timer_output()
     @static_timeit timer "cache construction" begin
 
@@ -222,7 +240,7 @@ function SciMLBase.__init(
             NonlinearSolveBase.supports_trust_region(alg.descent) ||
                 error("Trust Region not supported by $(alg.descent).")
             trustregion_cache = InternalAPI.init(
-                prob, alg.trustregion, fu, u, p; stats, internalnorm, kwargs...
+                _ad_prob, alg.trustregion, fu, u, _ad_prob.p; stats, internalnorm, kwargs...
             )
             globalization = Val(:TrustRegion)
         end
@@ -231,7 +249,7 @@ function SciMLBase.__init(
             NonlinearSolveBase.supports_line_search(alg.descent) ||
                 error("Line Search not supported by $(alg.descent).")
             linesearch_cache = CommonSolve.init(
-                prob, alg.linesearch, fu, u; stats, internalnorm, kwargs...
+                _ad_prob, alg.linesearch, fu, u; stats, internalnorm, kwargs...
             )
             globalization = Val(:LineSearch)
         end