From 819a8088ba19d1e4e348576a86c023d87f9ea536 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 28 Jul 2026 04:30:01 -0400 Subject: [PATCH 1/2] Reach the AutoDePSpecialize wrappers through public APIs The NonlinearSolveBase QA group fails on `master` with two `ExplicitImports.NonPublicQualifiedAccessException`s in the ForwardDiff extension, and two more in NonlinearSolveBase itself that the first pair was masking. Rather than reconstruct the same objects through indirect public surface, promote each name at its owning package and use it. - `FunctionWrappersWrappers.SingleCacheStorage` plus the raw `FunctionWrappersWrapper{FW, P, CS}(fwt, cs)` inner constructor become the public `FunctionWrappersWrapper(fwt; policy)` tuple constructor, which exists for exactly this inferrability case. It yields the identical concrete wrapper type, so the norecompile path is unchanged. Requires FunctionWrappersWrappers 1.11. - `SciMLBase.has_sys` and `SciMLBase.ImmutableNonlinearProblem` are public as of SciMLBase 3.41, so the guards in `maybe_opaque_wrap` keep their exact semantics. `ImmutableNonlinearProblem` comes off the ExplicitImports ignore list. - `wrapfun_iip_opaque` is NonlinearSolveBase's own extension point, overloaded in its own extension exactly like `wrapfun_iip`. It joins `wrapfun_iip` as an explicit import and in the same ignore entry, which already documents this category. `RespecializeParams.wrap_void_opaque` needed nothing: it is already exported and documented, and ExplicitImports never flagged it. Adds a direct dual-call regression test for the opaque wrapper's Jacobian signature, the shape AD actually reaches. Co-Authored-By: Chris Rackauckas --- lib/NonlinearSolveBase/Project.toml | 6 ++-- .../ext/NonlinearSolveBaseForwardDiffExt.jl | 34 +++++++++---------- lib/NonlinearSolveBase/src/autospecialize.jl | 3 +- .../test/autodepspecialize.jl | 14 ++++++++ lib/NonlinearSolveBase/test/qa/qa.jl | 10 +++--- 5 files changed, 40 insertions(+), 27 deletions(-) diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index 3a06ba772..21698ab45 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -1,6 +1,6 @@ name = "NonlinearSolveBase" uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" -version = "2.38.0" +version = "2.38.1" authors = ["Avik Pal and contributors"] [deps] @@ -79,7 +79,7 @@ EnzymeCore = "0.8.16" FastClosures = "0.3.2" ForwardDiff = "0.10.36, 1" FunctionWrappers = "1.1.2" -FunctionWrappersWrappers = "1" +FunctionWrappersWrappers = "1.11" InteractiveUtils = "<0.0.1, 1" LineSearch = "0.1.4" LinearAlgebra = "1.10" @@ -95,7 +95,7 @@ Printf = "1.10" RecursiveArrayTools = "3, 4" ReverseDiff = "1.15" RespecializeParams = "1" -SciMLBase = "3.37" +SciMLBase = "3.41" SciMLJacobianOperators = "0.1.1" SciMLLogging = "1.10.1, 2" SciMLOperators = "1.24" diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl index d5fc68995..5d9e9c885 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl @@ -18,7 +18,7 @@ using NonlinearSolveBase: NonlinearSolveBase, Utils, InternalAPI, NonlinearSolvePolyAlgorithm, NonlinearSolveForwardDiffCache, NonlinearSolveTag, is_fw_wrapped -import NonlinearSolveBase: wrapfun_iip, standardize_forwarddiff_tag +import NonlinearSolveBase: wrapfun_iip, wrapfun_iip_opaque, standardize_forwarddiff_tag const DI = DifferentiationInterface @@ -64,17 +64,18 @@ function standardize_forwarddiff_tag( return _wrapped_forwarddiff_ad(eltype(prob.u0)) end -# Construct the `FunctionWrappersWrapper` bypassing the convenience constructor, mirroring -# DiffEqBase's `_make_fww` (ext/DiffEqBaseForwardDiffExt.jl). The convenience constructor -# builds its wrapper tuple with `map(argtypes, rettypes) do A, R; …; end`, whose closure -# leaves `A`/`R` inferred as the join of the heterogeneous `argtypes` tuple, so the wrapper -# type — and every cache built from the wrapped problem — widens to abstract whenever the -# wrapped callable is one inference cannot see through: a functor such as the homotopy -# drivers' `FixLambda` / `AugmentedHomotopy` / `HomotopyResidual` residuals (the same reason -# DiffEqBase hit this with many-parameter `ODEFunction`s). Binding each arglist as a -# `::Type{A}` type parameter makes `FunctionWrapper{Nothing, A}(vff)` fully inferrable, so -# `typeof(fwt)` — and the wrapper — stays concrete. `vff` is `@nospecialize`d because it is -# already type-erased through `Void` (that is what makes the norecompile path precompile). +# Build the `FunctionWrapper` tuple here rather than handing `argtypes`/`rettypes` to the +# `FunctionWrappersWrapper(f, argtypes, rettypes)` constructor, mirroring DiffEqBase's +# `_make_fww` (ext/DiffEqBaseForwardDiffExt.jl). That constructor builds its wrapper tuple +# with `map(argtypes, rettypes) do A, R; …; end`, whose closure leaves `A`/`R` inferred as +# the join of the heterogeneous `argtypes` tuple, so the wrapper type — and every cache built +# from the wrapped problem — widens to abstract whenever the wrapped callable is one +# inference cannot see through: a functor such as the homotopy drivers' `FixLambda` / +# `AugmentedHomotopy` / `HomotopyResidual` residuals (the same reason DiffEqBase hit this +# with many-parameter `ODEFunction`s). Binding each arglist as a `::Type{A}` type parameter +# makes `FunctionWrapper{Nothing, A}(vff)` fully inferrable, so `typeof(fwt)` — and the +# wrapper — stays concrete. `vff` is `@nospecialize`d because it is already type-erased +# through `Void` (that is what makes the norecompile path precompile). function _make_fww_iip( @nospecialize(vff), ::Type{A1}, ::Type{A2}, ::Type{A3}, ::Type{A4} ) where {A1, A2, A3, A4} @@ -83,10 +84,9 @@ function _make_fww_iip( FW{Nothing, A1}(vff), FW{Nothing, A2}(vff), FW{Nothing, A3}(vff), FW{Nothing, A4}(vff), ) - cs = FunctionWrappersWrappers.SingleCacheStorage() - return FunctionWrappersWrappers.FunctionWrappersWrapper{ - typeof(fwt), FunctionWrappersWrappers.AllowNonIsBits, typeof(cs), - }(fwt, cs) + return FunctionWrappersWrappers.FunctionWrappersWrapper( + fwt; policy = FunctionWrappersWrappers.AllowNonIsBits() + ) end # IIP wrapfun: wraps f(du, u, p) with dual-aware type combinations. @@ -119,7 +119,7 @@ end # the third slot). `p` is never a `Dual` on this path (opaque-ification is # skipped for dual state/params), so only the plain and Jacobian-`Dual`-`u` # signatures are needed. -@inline function NonlinearSolveBase.wrapfun_iip_opaque( +@inline function wrapfun_iip_opaque( ff, ::Type{P}, inputs::Tuple{T1, T2, T3} ) where {P, T1 <: AbstractArray, T2 <: AbstractArray, T3} T = eltype(T1) diff --git a/lib/NonlinearSolveBase/src/autospecialize.jl b/lib/NonlinearSolveBase/src/autospecialize.jl index c840a4f54..8fbd996dc 100644 --- a/lib/NonlinearSolveBase/src/autospecialize.jl +++ b/lib/NonlinearSolveBase/src/autospecialize.jl @@ -292,8 +292,7 @@ function maybe_opaque_wrap(prob::AbstractNonlinearProblem) SciMLBase.specialization(prob.f) === SciMLBase.AutoDePSpecialize || return nothing EnzymeCore.within_autodiff() && return nothing is_fw_wrapped(prob.f.f) && return nothing - (prob isa NonlinearProblem || prob isa SciMLBase.ImmutableNonlinearProblem) || - return nothing + (prob isa NonlinearProblem || prob isa ImmutableNonlinearProblem) || return nothing SciMLBase.isinplace(prob) || return nothing u0 = prob.u0 diff --git a/lib/NonlinearSolveBase/test/autodepspecialize.jl b/lib/NonlinearSolveBase/test/autodepspecialize.jl index 55cfad6a9..a3172a131 100644 --- a/lib/NonlinearSolveBase/test/autodepspecialize.jl +++ b/lib/NonlinearSolveBase/test/autodepspecialize.jl @@ -38,6 +38,20 @@ u0 = [1.0, 1.0] res = [0.0, 0.0] cp.f.f(res, [2.0, 3.0], cp.p) @test res ≈ [4.0 - 2.0, 9.0 - 3.0] + + # The Jacobian shape — `Dual` state against the still-opaque `p` — is the + # second signature the opaque wrapper carries, and the only one AD reaches. + DualT = ForwardDiff.Dual{ + ForwardDiff.Tag{NonlinearSolveBase.NonlinearSolveTag, Float64}, Float64, 1, + } + dual_u = DualT[ + DualT(2.0, ForwardDiff.Partials((1.0,))), + DualT(3.0, ForwardDiff.Partials((0.0,))), + ] + dual_res = similar(dual_u) + cp.f.f(dual_res, dual_u, cp.p) + @test ForwardDiff.value.(dual_res) ≈ res + @test first.(ForwardDiff.partials.(dual_res)) ≈ [4.0, 0.0] end @testset "AutoSpecialize / FullSpecialize leave p untouched" begin diff --git a/lib/NonlinearSolveBase/test/qa/qa.jl b/lib/NonlinearSolveBase/test/qa/qa.jl index 89cc3fe47..8f991d71c 100644 --- a/lib/NonlinearSolveBase/test/qa/qa.jl +++ b/lib/NonlinearSolveBase/test/qa/qa.jl @@ -64,8 +64,8 @@ run_qa( ), ), # Still non-public in their owning packages. AbstractODEIntegrator / __init / - # __solve dropped here: now public in SciMLBase. - # SciMLBase: ImmutableNonlinearProblem, KeywordArgError, NoDefaultAlgorithmError, + # __solve / ImmutableNonlinearProblem dropped here: now public in SciMLBase. + # SciMLBase: KeywordArgError, NoDefaultAlgorithmError, # NonSolverError, _concrete_solve_adjoint, _concrete_solve_forward, checkkwargs, # extract_alg, get_concrete_p, get_concrete_u0, get_root_indp, has_kwargs, # promote_u0, wrap_sol @@ -73,17 +73,17 @@ run_qa( # StaticArraysCore: StaticArray # NonlinearSolveBase (own internal): NonlinearSolveForwardDiffCache, # NonlinearSolveTag, Utils, is_fw_wrapped, standardize_forwarddiff_tag, - # wrapfun_iip + # wrapfun_iip, wrapfun_iip_opaque all_explicit_imports_are_public = (; ignore = ( - :ImmutableNonlinearProblem, :KeywordArgError, + :KeywordArgError, :NoDefaultAlgorithmError, :NonSolverError, :_concrete_solve_adjoint, :_concrete_solve_forward, :checkkwargs, :extract_alg, :get_concrete_p, :get_concrete_u0, :get_root_indp, :has_kwargs, :promote_u0, :wrap_sol, :Dual, :pickchunksize, :AbstractSparseMatrixCSC, :StaticArray, :NonlinearSolveForwardDiffCache, :NonlinearSolveTag, :Utils, :is_fw_wrapped, :standardize_forwarddiff_tag, - :wrapfun_iip, + :wrapfun_iip, :wrapfun_iip_opaque, ), ), ), From fb187093d941ce503ced803220e4c797808960e4 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 28 Jul 2026 04:47:02 -0400 Subject: [PATCH 2/2] Keep master's explicit FunctionWrappersWrapper construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reverts `_make_fww_iip` to its `master` form. With FunctionWrappersWrappers 1.11.0 registered, `SingleCacheStorage` is public, so the original code passes the non-public-access check unchanged — rebuilding the wrapper through the keyword constructor was never necessary. It was also the wrong trade. The explicit `{FW, P, CS}(fwt, cs)` form fixes the wrapper type from type parameters alone, while the keyword form derives `CS` from `_make_cache_storage(cache)` and so leans on constant propagation through a keyword-argument frame to stay concrete. Nothing on the norecompile path should depend on that holding. Co-Authored-By: Chris Rackauckas --- .../ext/NonlinearSolveBaseForwardDiffExt.jl | 30 +++++++++---------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl index 5d9e9c885..729f80fe8 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl @@ -64,18 +64,17 @@ function standardize_forwarddiff_tag( return _wrapped_forwarddiff_ad(eltype(prob.u0)) end -# Build the `FunctionWrapper` tuple here rather than handing `argtypes`/`rettypes` to the -# `FunctionWrappersWrapper(f, argtypes, rettypes)` constructor, mirroring DiffEqBase's -# `_make_fww` (ext/DiffEqBaseForwardDiffExt.jl). That constructor builds its wrapper tuple -# with `map(argtypes, rettypes) do A, R; …; end`, whose closure leaves `A`/`R` inferred as -# the join of the heterogeneous `argtypes` tuple, so the wrapper type — and every cache built -# from the wrapped problem — widens to abstract whenever the wrapped callable is one -# inference cannot see through: a functor such as the homotopy drivers' `FixLambda` / -# `AugmentedHomotopy` / `HomotopyResidual` residuals (the same reason DiffEqBase hit this -# with many-parameter `ODEFunction`s). Binding each arglist as a `::Type{A}` type parameter -# makes `FunctionWrapper{Nothing, A}(vff)` fully inferrable, so `typeof(fwt)` — and the -# wrapper — stays concrete. `vff` is `@nospecialize`d because it is already type-erased -# through `Void` (that is what makes the norecompile path precompile). +# Construct the `FunctionWrappersWrapper` bypassing the convenience constructor, mirroring +# DiffEqBase's `_make_fww` (ext/DiffEqBaseForwardDiffExt.jl). The convenience constructor +# builds its wrapper tuple with `map(argtypes, rettypes) do A, R; …; end`, whose closure +# leaves `A`/`R` inferred as the join of the heterogeneous `argtypes` tuple, so the wrapper +# type — and every cache built from the wrapped problem — widens to abstract whenever the +# wrapped callable is one inference cannot see through: a functor such as the homotopy +# drivers' `FixLambda` / `AugmentedHomotopy` / `HomotopyResidual` residuals (the same reason +# DiffEqBase hit this with many-parameter `ODEFunction`s). Binding each arglist as a +# `::Type{A}` type parameter makes `FunctionWrapper{Nothing, A}(vff)` fully inferrable, so +# `typeof(fwt)` — and the wrapper — stays concrete. `vff` is `@nospecialize`d because it is +# already type-erased through `Void` (that is what makes the norecompile path precompile). function _make_fww_iip( @nospecialize(vff), ::Type{A1}, ::Type{A2}, ::Type{A3}, ::Type{A4} ) where {A1, A2, A3, A4} @@ -84,9 +83,10 @@ function _make_fww_iip( FW{Nothing, A1}(vff), FW{Nothing, A2}(vff), FW{Nothing, A3}(vff), FW{Nothing, A4}(vff), ) - return FunctionWrappersWrappers.FunctionWrappersWrapper( - fwt; policy = FunctionWrappersWrappers.AllowNonIsBits() - ) + cs = FunctionWrappersWrappers.SingleCacheStorage() + return FunctionWrappersWrappers.FunctionWrappersWrapper{ + typeof(fwt), FunctionWrappersWrappers.AllowNonIsBits, typeof(cs), + }(fwt, cs) end # IIP wrapfun: wraps f(du, u, p) with dual-aware type combinations.