Unwrap AutoSpecialize in AD-over-solve paths and external wrapper extensions

ChrisRackauckas · claude · ChrisRackauckas · commit d190b980d3fb · 2026-03-30T16:54:59.000-04:00
Instead of relying on FunctionWrappersWrapper fallback mode, explicitly
unwrap AutoSpecializeCallable in all code paths where ForwardDiff uses
non-standard tags:

1. ForwardDiff sensitivity path (nonlinearsolve_forwarddiff_solve):
   unwrap prob.f.f before computing ∂f/∂p and ∂f/∂u Jacobians, which
   use closure-based ForwardDiff tags incompatible with the wrappers.
   Removes the now-unnecessary _is_wrapped_nlf checks and _nls_tag
   config overrides.

2. External wrapper extensions (LeastSquaresOptim, NLsolve, NLSolvers,
   SIAM, FastLevenbergMarquardt, MINPACK, PETSc, FixedPointAcceleration,
   SpeedMapping, Sundials): unwrap at the top of __solve since these
   packages do their own AD with arbitrary tags.

3. Bounds transform: unwrap before wrapping in BoundedWrapper since the
   transform changes argument types.

4. Restrict wrapping to supported param types only (Vector{Float64},
   NullParameters) — scalar params skip wrapping entirely.

Co-Authored-By: Chris Rackauckas &lt;accounts@chrisrackauckas.com&gt;
Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/Project.toml b/Project.toml
@@ -24,6 +24,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SciMLLogging = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
@@ -119,6 +120,7 @@ Reexport = "1.2.2"
 ReverseDiff = "1.15"
 SIAMFANLEquations = "1.0.1"
 SciMLBase = "2.153"
+Setfield = "1.1.2"
 SciMLLogging = "1.7"
 SimpleNonlinearSolve = "2.10"
 SparseArrays = "1.10"
diff --git a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl
@@ -6,9 +6,10 @@ using ArrayInterface: ArrayInterface
 using FastLevenbergMarquardt: FastLevenbergMarquardt
 using StaticArraysCore: SArray
 
-using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, FastLevenbergMarquardtJL
-using SciMLBase: SciMLBase, AbstractNonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, AbstractNonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 const FastLM = FastLevenbergMarquardt
 
@@ -17,6 +18,11 @@ function SciMLBase.__solve(
         alias_u0 = false, abstol = nothing, reltol = nothing, maxiters = 1000,
         termination_condition = nothing, kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     NonlinearSolveBase.assert_extension_supported_termination_condition(
         termination_condition, alg
     )
diff --git a/ext/NonlinearSolveFixedPointAccelerationExt.jl b/ext/NonlinearSolveFixedPointAccelerationExt.jl
@@ -2,15 +2,21 @@ module NonlinearSolveFixedPointAccelerationExt
 
 using FixedPointAcceleration: FixedPointAcceleration, fixed_point
 
-using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, FixedPointAccelerationJL
-using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 function SciMLBase.__solve(
         prob::NonlinearProblem, alg::FixedPointAccelerationJL, args...;
         abstol = nothing, maxiters = 1000, alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false),
         show_trace::Val = Val(false), termination_condition = nothing, kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     if haskey(kwargs, :alias_u0)
         alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
     end
diff --git a/ext/NonlinearSolveLeastSquaresOptimExt.jl b/ext/NonlinearSolveLeastSquaresOptimExt.jl
@@ -2,9 +2,10 @@ module NonlinearSolveLeastSquaresOptimExt
 
 using LeastSquaresOptim: LeastSquaresOptim
 
-using NonlinearSolveBase: NonlinearSolveBase, TraceMinimal
+using NonlinearSolveBase: NonlinearSolveBase, TraceMinimal, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, LeastSquaresOptimJL
-using SciMLBase: SciMLBase, AbstractNonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, AbstractNonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 const LSO = LeastSquaresOptim
 
@@ -14,6 +15,11 @@ function SciMLBase.__solve(
         trace_level = TraceMinimal(), termination_condition = nothing,
         show_trace::Val = Val(false), store_trace::Val = Val(false), kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     NonlinearSolveBase.assert_extension_supported_termination_condition(
         termination_condition, alg
     )
diff --git a/ext/NonlinearSolveMINPACKExt.jl b/ext/NonlinearSolveMINPACKExt.jl
@@ -3,16 +3,22 @@ module NonlinearSolveMINPACKExt
 using FastClosures: @closure
 using MINPACK: MINPACK
 
-using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, CMINPACK
-using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 function SciMLBase.__solve(
         prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, alg::CMINPACK, args...;
         abstol = nothing, maxiters = 1000, alias_u0::Bool = false,
         show_trace::Val = Val(false), store_trace::Val = Val(false),
         termination_condition = nothing, kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     NonlinearSolveBase.assert_extension_supported_termination_condition(
         termination_condition, alg
     )
diff --git a/ext/NonlinearSolveNLSolversExt.jl b/ext/NonlinearSolveNLSolversExt.jl
@@ -5,9 +5,10 @@ using FastClosures: @closure
 
 using NLSolvers: NLSolvers, NEqOptions, NEqProblem
 
-using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, NLSolversJL
-using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 const DI = DifferentiationInterface
 
@@ -16,6 +17,11 @@ function SciMLBase.__solve(
         abstol = nothing, reltol = nothing, maxiters = 1000, alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false),
         termination_condition = nothing, kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     if haskey(kwargs, :alias_u0)
         alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
     end
diff --git a/ext/NonlinearSolveNLsolveExt.jl b/ext/NonlinearSolveNLsolveExt.jl
@@ -3,16 +3,22 @@ module NonlinearSolveNLsolveExt
 using LineSearches: Static
 using NLsolve: NLsolve, OnceDifferentiable, nlsolve
 
-using NonlinearSolveBase: NonlinearSolveBase, Utils, TraceMinimal
+using NonlinearSolveBase: NonlinearSolveBase, Utils, TraceMinimal, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, NLsolveJL
-using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 function SciMLBase.__solve(
         prob::NonlinearProblem, alg::NLsolveJL, args...;
         abstol = nothing, maxiters = 1000, alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false),
         termination_condition = nothing, trace_level = TraceMinimal(),
         store_trace::Val = Val(false), show_trace::Val = Val(false), kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     if haskey(kwargs, :alias_u0)
         alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
     end
diff --git a/ext/NonlinearSolvePETScExt.jl b/ext/NonlinearSolvePETScExt.jl
@@ -3,9 +3,10 @@ module NonlinearSolvePETScExt
 using MPI: MPI
 using PETSc: PETSc
 
-using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, PETScSNES
-using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC, sparse, rowvals, nonzeros, nzrange
 
@@ -40,6 +41,11 @@ function SciMLBase.__solve(
         maxiters = 1000, alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false), termination_condition = nothing,
         show_trace::Val = Val(false), kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     if haskey(kwargs, :alias_u0)
         alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
     end
diff --git a/ext/NonlinearSolveSIAMFANLEquationsExt.jl b/ext/NonlinearSolveSIAMFANLEquationsExt.jl
@@ -4,9 +4,10 @@ using FastClosures: @closure
 using SIAMFANLEquations: SIAMFANLEquations, aasol, nsol, nsoli, nsolsc, ptcsol, ptcsoli,
     ptcsolsc, secant
 
-using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, SIAMFANLEquationsJL
-using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 function siamfanlequations_retcode_mapping(sol)
     if sol.errcode == 0
@@ -42,6 +43,11 @@ function SciMLBase.__solve(
         abstol = nothing, reltol = nothing, alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false), maxiters = 1000,
         termination_condition = nothing, show_trace = Val(false), kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     if haskey(kwargs, :alias_u0)
         alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
     end
diff --git a/ext/NonlinearSolveSpeedMappingExt.jl b/ext/NonlinearSolveSpeedMappingExt.jl
@@ -2,16 +2,22 @@ module NonlinearSolveSpeedMappingExt
 
 using SpeedMapping: speedmapping
 
-using NonlinearSolveBase: NonlinearSolveBase
+using NonlinearSolveBase: NonlinearSolveBase, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, SpeedMappingJL
-using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode
+using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode, remake
+using Setfield: @set
 
 function SciMLBase.__solve(
         prob::NonlinearProblem, alg::SpeedMappingJL, args...;
         abstol = nothing, maxiters = 1000, alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = false),
         maxtime = nothing, store_trace::Val = Val(false),
         termination_condition = nothing, kwargs...
     )
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     if haskey(kwargs, :alias_u0)
         alias = SciMLBase.NonlinearAliasSpecifier(alias_u0 = kwargs[:alias_u0])
     end
diff --git a/ext/NonlinearSolveSundialsExt.jl b/ext/NonlinearSolveSundialsExt.jl
@@ -4,11 +4,17 @@ using Sundials: KINSOL
 
 using CommonSolve: CommonSolve
 using NonlinearSolveBase: NonlinearSolveBase, nonlinearsolve_forwarddiff_solve,
-    nonlinearsolve_dual_solution
+    nonlinearsolve_dual_solution, is_fw_wrapped, get_raw_f
 using NonlinearSolve: NonlinearSolve, DualNonlinearProblem
-using SciMLBase: SciMLBase
+using SciMLBase: SciMLBase, remake
+using Setfield: @set
 
 function SciMLBase.__solve(prob::DualNonlinearProblem, alg::KINSOL, args...; kwargs...)
+    # Unwrap AutoSpecialize — external packages do their own AD
+    if is_fw_wrapped(prob.f.f)
+        prob = @set prob.f.f = get_raw_f(prob.f.f)
+    end
+
     sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...)
     dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p)
     return SciMLBase.build_solution(
diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl
@@ -10,6 +10,7 @@ using FunctionWrappers: FunctionWrappers
 import FunctionWrappersWrappers
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, IntervalNonlinearProblem,
     NonlinearProblem, NonlinearLeastSquaresProblem, remake
+using Setfield: @set
 
 using LinearAlgebra: LinearAlgebra, dot, norm
 using NonlinearSolveBase: NonlinearSolveBase, ImmutableNonlinearProblem, Utils, InternalAPI,
@@ -76,7 +77,7 @@ end
     )
     iip_returnlists = (Nothing, Nothing, Nothing, Nothing)
     return FunctionWrappersWrappers.FunctionWrappersWrapper(
-        SciMLBase.Void(ff), iip_arglists, iip_returnlists, Val{true}()
+        SciMLBase.Void(ff), iip_arglists, iip_returnlists
     )
 end
 
@@ -124,11 +125,19 @@ function NonlinearSolveBase.nonlinearsolve_forwarddiff_solve(
     sol = solve(newprob, alg, args...; kwargs...)
     uu = sol.u
 
-    fn = prob isa NonlinearLeastSquaresProblem ?
-        NonlinearSolveBase.nlls_generate_vjp_function(prob, sol, uu) : prob.f
+    # Unwrap AutoSpecializeCallable for the AD-over-solve Jacobian computations.
+    # These use ForwardDiff with closure-based tags that don't match the wrapper signatures.
+    ad_prob = if is_fw_wrapped(prob.f.f)
+        @set prob.f.f = NonlinearSolveBase.get_raw_f(prob.f.f)
+    else
+        prob
+    end
 
-    Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(prob, fn, uu, p)
-    Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(prob, fn, uu, p)
+    fn = ad_prob isa NonlinearLeastSquaresProblem ?
+        NonlinearSolveBase.nlls_generate_vjp_function(ad_prob, sol, uu) : ad_prob.f
+
+    Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(ad_prob, fn, uu, p)
+    Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(ad_prob, fn, uu, p)
     z = -Jᵤ \ Jₚ
     pp = prob.p
     sumfun = ((z, p),) -> map(Base.Fix2(*, ForwardDiff.partials(p)), z)
@@ -144,16 +153,6 @@ function NonlinearSolveBase.nonlinearsolve_forwarddiff_solve(
     return sol, partials
 end
 
-# Check if a NonlinearFunction has an AutoSpecialize-wrapped callable.
-# Wrapping is only active for Vector{Float64} state/params, so the Number code paths
-# (derivative, gradient) never encounter it — only the jacobian paths need configs.
-_is_wrapped_nlf(f) = false
-function _is_wrapped_nlf(f::SciMLBase.NonlinearFunction)
-    return is_fw_wrapped(f.f)
-end
-
-const _nls_tag = ForwardDiff.Tag(NonlinearSolveTag(), Float64)
-
 function NonlinearSolveBase.nonlinearsolve_∂f_∂p(prob, f::F, u, p) where {F}
     if SciMLBase.isinplace(prob)
         f2 = @closure p -> begin
@@ -169,10 +168,6 @@ function NonlinearSolveBase.nonlinearsolve_∂f_∂p(prob, f::F, u, p) where {F}
     elseif u isa Number
         return Utils.safe_reshape(ForwardDiff.gradient(f2, p), 1, :)
     else
-        if _is_wrapped_nlf(f)
-            cfg = ForwardDiff.JacobianConfig(f2, p, ForwardDiff.Chunk{1}(), _nls_tag)
-            return ForwardDiff.jacobian(f2, p, cfg)
-        end
         return ForwardDiff.jacobian(f2, p)
     end
 end
@@ -181,21 +176,9 @@ function NonlinearSolveBase.nonlinearsolve_∂f_∂u(prob, f::F, u, p) where {F}
     if SciMLBase.isinplace(prob)
         jac_f = @closure((du, u) -> f(du, u, p))
         du_cache = Utils.safe_similar(u)
-        if _is_wrapped_nlf(f)
-            cfg = ForwardDiff.JacobianConfig(
-                jac_f, du_cache, u, ForwardDiff.Chunk{1}(), _nls_tag
-            )
-            return ForwardDiff.jacobian(jac_f, du_cache, u, cfg)
-        end
         return ForwardDiff.jacobian(jac_f, du_cache, u)
     end
     u isa Number && return ForwardDiff.derivative(Base.Fix2(f, p), u)
-    if _is_wrapped_nlf(f)
-        cfg = ForwardDiff.JacobianConfig(
-            Base.Fix2(f, p), u, ForwardDiff.Chunk{1}(), _nls_tag
-        )
-        return ForwardDiff.jacobian(Base.Fix2(f, p), u, cfg)
-    end
     return ForwardDiff.jacobian(Base.Fix2(f, p), u)
 end
 
@@ -260,11 +243,18 @@ function CommonSolve.solve!(cache::NonlinearSolveForwardDiffCache)
     prob = cache.prob
     uu = sol.u
 
-    fn = prob isa NonlinearLeastSquaresProblem ?
-        NonlinearSolveBase.nlls_generate_vjp_function(prob, sol, uu) : prob.f
+    # Unwrap AutoSpecializeCallable for the AD-over-solve Jacobian computations.
+    ad_prob = if is_fw_wrapped(prob.f.f)
+        @set prob.f.f = NonlinearSolveBase.get_raw_f(prob.f.f)
+    else
+        prob
+    end
+
+    fn = ad_prob isa NonlinearLeastSquaresProblem ?
+        NonlinearSolveBase.nlls_generate_vjp_function(ad_prob, sol, uu) : ad_prob.f
 
-    Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(prob, fn, uu, cache.values_p)
-    Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(prob, fn, uu, cache.values_p)
+    Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(ad_prob, fn, uu, cache.values_p)
+    Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(ad_prob, fn, uu, cache.values_p)
 
     z_arr = -Jᵤ \ Jₚ
 
diff --git a/lib/NonlinearSolveBase/src/autospecialize.jl b/lib/NonlinearSolveBase/src/autospecialize.jl
diff --git a/lib/NonlinearSolveBase/src/bounds_transform.jl b/lib/NonlinearSolveBase/src/bounds_transform.jl
diff --git a/test/bounds_tests.jl b/test/bounds_tests.jl
