Reach the AutoDePSpecialize wrappers through public APIs

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 <avikpal@mit.edu> 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"

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
 
@@ -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)

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

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

lib/NonlinearSolveBase/test/qa/qa.jl

@@ -64,26 +64,26 @@ 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
         #   ForwardDiff: Dual, pickchunksize;  SparseArrays: AbstractSparseMatrixCSC
         #   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,
             ),
         ),
     ),