From 39664b17135af68bf76cd2a7e3d1c624a88aed16 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Tue, 16 Jun 2026 10:41:45 -0400 Subject: [PATCH 01/10] allow Duals to pass through grad --- lib/OptimizationBase/src/OptimizationDIExt.jl | 70 ++++++++++++++++--- 1 file changed, 60 insertions(+), 10 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index ab3be02e9..15bb132b6 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -15,6 +15,30 @@ import DifferentiationInterface: prepare_gradient, prepare_hessian, prepare_hvp, using ADTypes, SciMLBase using OptimizationBase.FastClosures +# --- Dual-tolerant gradient dispatch ------------------------------------------------------ +# A DI preparation is monomorphic: built once at the construction types (`x`, `p`), its +# internal buffers are concrete (e.g. `Float64`) and reject inputs of any other type. That is +# correct and fast for the optimization solve, but a downstream sensitivity layer (e.g. +# SciMLSensitivity's `OptimizationAdjoint`) differentiates the KKT stationarity conditions +# w.r.t. the parameters by pushing `ForwardDiff.Dual`s through the gradient: a dual `p` (and a +# dual output buffer) evaluated at a real `θ = x*`. Those duals hit the prepared buffers and +# throw `DifferentiationInterface.PreparationMismatchError`. +# +# To support that without slowing the solve, the gradient closures keep the prepared fast path +# for the prepared type and fall back to a prep-free DI call (which prepares at the actual +# argument types each call) for anything else. The fallback cost is irrelevant off the +# optimization hot loop — sensitivity does one such call per solve, not per iteration. +_grad_param_eltype(p) = p isa Union{SciMLBase.NullParameters, Nothing} ? nothing : eltype(p) + +# True when every float-bearing input matches the prepared element type `T0`, so the prepared +# path is valid. A dual `θ` or dual `p` (the sensitivity case) flunks this and routes to the +# prep-free fallback. `NullParameters`/`nothing` carry no differentiable parameters, so they +# never veto the fast path. +@inline function _grad_use_prep(::Type{T0}, θ, p) where {T0} + pe = _grad_param_eltype(p) + eltype(θ) === T0 && (pe === nothing || pe === T0) +end + function instantiate_function( f::OptimizationFunction{true}, x, ::ADTypes.AutoSparse{<:ADTypes.AutoSymbolics}, args...; kwargs... @@ -39,16 +63,31 @@ function instantiate_function( ) adtype, soadtype = generate_adtype(adtype) - # Create gradient closures with proper type stability using let blocks + # Create gradient closures with proper type stability using let blocks. + # `T0 = eltype(x)` is the prepared element type; `_grad_use_prep` gates the prepared fast + # path vs the prep-free fallback (see the note above the imports). grad = if g == true && f.grad === nothing _prep_grad = prepare_gradient(f.f, adtype, x, Constant(p)) + T0 = eltype(x) if p !== SciMLBase.NullParameters() && p !== nothing - let _prep_grad = _prep_grad, f = f, adtype = adtype - (res, θ, p = p) -> gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) + let _prep_grad = _prep_grad, f = f, adtype = adtype, T0 = T0 + function (res, θ, p = p) + if _grad_use_prep(T0, θ, p) && eltype(res) === T0 + gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) + else + gradient!(f.f, res, adtype, θ, Constant(p)) + end + end end else - let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p - (res, θ, p = p) -> gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) + let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, T0 = T0 + function (res, θ, p = p) + if _grad_use_prep(T0, θ, p) && eltype(res) === T0 + gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) + else + gradient!(f.f, res, adtype, θ, Constant(p)) + end + end end end elseif g == true @@ -428,16 +467,27 @@ function instantiate_function( ) adtype, soadtype = generate_adtype(adtype) - # Create gradient closures with proper type stability using let blocks + # Create gradient closures with proper type stability using let blocks. + # `T0 = eltype(x)` is the prepared element type; `_grad_use_prep` gates the prepared fast + # path vs the prep-free fallback (see the note above the imports). grad = if g == true && f.grad === nothing _prep_grad = prepare_gradient(f.f, adtype, x, Constant(p)) + T0 = eltype(x) if p !== SciMLBase.NullParameters() && p !== nothing - let _prep_grad = _prep_grad, f = f, adtype = adtype - (θ, p = p) -> gradient(f.f, _prep_grad, adtype, θ, Constant(p)) + let _prep_grad = _prep_grad, f = f, adtype = adtype, T0 = T0 + function (θ, p = p) + _grad_use_prep(T0, θ, p) ? + gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : + gradient(f.f, adtype, θ, Constant(p)) + end end else - let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p - (θ, p = p) -> gradient(f.f, _prep_grad, adtype, θ, Constant(p)) + let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, T0 = T0 + function (θ, p = p) + _grad_use_prep(T0, θ, p) ? + gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : + gradient(f.f, adtype, θ, Constant(p)) + end end end elseif g == true From 51b5e7d7f4defc45d85e5a3ed099525fa0bf3a6e Mon Sep 17 00:00:00 2001 From: jClugstor Date: Tue, 16 Jun 2026 14:53:43 -0400 Subject: [PATCH 02/10] allow cons_j to take parameter p --- lib/OptimizationBase/src/OptimizationDIExt.jl | 48 +++++++++++++++---- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index 15bb132b6..986f68751 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -233,10 +233,33 @@ function instantiate_function( cons_jac_colorvec = f.cons_jac_colorvec cons_j! = if f.cons !== nothing && cons_j == true && f.cons_j === nothing - _prep_jac = prepare_jacobian(cons_oop, adtype, x) - let cons_oop = cons_oop, _prep_jac = _prep_jac, adtype = adtype - function (J, θ) - jacobian!(cons_oop, J, _prep_jac, adtype, θ) + # A `p`-accepting out-of-place constraint wrapper, so the Jacobian can be evaluated + # at parameters other than the construction `p` — including duals pushed in by a + # sensitivity layer differentiating the constraint Jacobian w.r.t. `p` (the mixed + # ∂²cᵢ/∂x∂p term of the KKT residual). The prepared `cons_oop` bakes `p` in and + # exposes no parameter slot, so we cannot reuse it here. Output eltype promotes + # against `q` so duals propagate (Base.promote_op, not promote_type: ForwardDiff's + # dual ops bypass promote_type). + _cons_oop_p = let f = f, num_cons = num_cons + function (x, p) + T = p isa Union{SciMLBase.NullParameters, Nothing} ? eltype(x) : + Base.promote_op(+, eltype(x), eltype(p)) + res = Vector{T}(undef, num_cons) + f.cons(res, x, p) + return res + end + end + _prep_jac = prepare_jacobian(_cons_oop_p, adtype, x, Constant(p)) + T0 = eltype(x) + let _cons_oop_p = _cons_oop_p, _prep_jac = _prep_jac, adtype = adtype, p = p, T0 = T0 + function (J, θ, p = p) + # Prepared fast path when types match the prep; prep-free fallback for duals + # (see the `_grad_use_prep` note above the imports). + if _grad_use_prep(T0, θ, p) && eltype(J) === T0 + jacobian!(_cons_oop_p, J, _prep_jac, adtype, θ, Constant(p)) + else + jacobian!(_cons_oop_p, J, adtype, θ, Constant(p)) + end return if size(J, 1) == 1 J = vec(J) end @@ -244,7 +267,7 @@ function instantiate_function( end elseif cons_j == true && f.cons !== nothing let f = f, p = p - (J, θ) -> f.cons_j(J, θ, p) + (J, θ, p = p) -> f.cons_j(J, θ, p) end else nothing @@ -610,10 +633,17 @@ function instantiate_function( cons_jac_colorvec = f.cons_jac_colorvec cons_j! = if f.cons !== nothing && cons_j == true && f.cons_j === nothing + # `f.cons` is out-of-place here and the prep already takes `Constant(p)`, so this + # only needs to expose the parameter argument and add the dual-tolerant fallback + # (see the `_grad_use_prep` note above the imports) — unlike the in-place method, + # whose prepared wrapper bakes `p` in. _prep_jac = prepare_jacobian(f.cons, adtype, x, Constant(p)) - let f = f, _prep_jac = _prep_jac, adtype = adtype, p = p - function (θ) - J = jacobian(f.cons, _prep_jac, adtype, θ, Constant(p)) + T0 = eltype(x) + let f = f, _prep_jac = _prep_jac, adtype = adtype, p = p, T0 = T0 + function (θ, p = p) + J = _grad_use_prep(T0, θ, p) ? + jacobian(f.cons, _prep_jac, adtype, θ, Constant(p)) : + jacobian(f.cons, adtype, θ, Constant(p)) if size(J, 1) == 1 J = vec(J) end @@ -622,7 +652,7 @@ function instantiate_function( end elseif cons_j == true && f.cons !== nothing let f = f, p = p - (θ) -> f.cons_j(θ, p) + (θ, p = p) -> f.cons_j(θ, p) end else nothing From 05837efce60f80db08d81dd199fd76ebd71fd56a Mon Sep 17 00:00:00 2001 From: jClugstor Date: Wed, 17 Jun 2026 13:08:46 -0400 Subject: [PATCH 03/10] allow Enzyme cons_j to take parameters, prevent Boxing --- .../ext/OptimizationEnzymeExt.jl | 69 ++++++++++--------- 1 file changed, 35 insertions(+), 34 deletions(-) diff --git a/lib/OptimizationBase/ext/OptimizationEnzymeExt.jl b/lib/OptimizationBase/ext/OptimizationEnzymeExt.jl index d6493e147..1bb7d68ef 100644 --- a/lib/OptimizationBase/ext/OptimizationEnzymeExt.jl +++ b/lib/OptimizationBase/ext/OptimizationEnzymeExt.jl @@ -238,7 +238,7 @@ function OptimizationBase.instantiate_function( if f.cons === nothing cons = nothing else - cons = (res, θ) -> f.cons(res, θ, p) + cons = (res, θ, p = p) -> f.cons(res, θ, p) end if cons !== nothing && cons_j == true && f.cons_j === nothing @@ -271,44 +271,45 @@ function OptimizationBase.instantiate_function( y = zeros(eltype(x), num_cons) - function cons_j!(J, θ) - for jc in Jaccache - Enzyme.make_zero!(jc) - end - Enzyme.make_zero!(y) - if func_annot <: Enzyme.Duplicated || func_annot <: Enzyme.BatchDuplicated || - func_annot <: Enzyme.DuplicatedNoNeed || - func_annot <: Enzyme.BatchDuplicatedNoNeed - for bf in basefunc.dval - Enzyme.make_zero!(bf) + # Precompute the duplicated-annotation check as a Bool. Capturing this (rather than the + # type-valued `func_annot::DataType`) keeps the closure free of type-valued fields, which + # Enzyme's shadow-layout pass cannot reconcile when an outer Enzyme pass differentiates + # this closure. + dup_annot = func_annot <: Enzyme.Duplicated || func_annot <: Enzyme.BatchDuplicated || + func_annot <: Enzyme.DuplicatedNoNeed || func_annot <: Enzyme.BatchDuplicatedNoNeed + + # `let` so the closure captures stable bindings rather than `Core.Box`es: `basefunc` + # is reassigned in the branch above and `y` is captured alongside it, so without this + # both get boxed — which defeats specialization in the solve hot loop. The same + # `let`-capture idiom is used for the DI-built closures. + cons_j! = let basefunc = basefunc, y = y, Jaccache = Jaccache, seeds = seeds, + dup_annot = dup_annot, fmode = fmode, p = p + function (J, θ, p = p) + for jc in Jaccache + Enzyme.make_zero!(jc) end - end - Enzyme.autodiff( - fmode, basefunc, BatchDuplicated(y, Jaccache), - BatchDuplicated(θ, seeds), Const(p) - ) - for i in eachindex(θ) - if J isa Vector - J[i] = Jaccache[i][1] - else - copyto!(@view(J[:, i]), Jaccache[i]) + Enzyme.make_zero!(y) + if dup_annot + for bf in basefunc.dval + Enzyme.make_zero!(bf) + end end + Enzyme.autodiff( + fmode, basefunc, BatchDuplicated(y, Jaccache), + BatchDuplicated(θ, seeds), Const(p) + ) + for i in eachindex(θ) + if J isa Vector + J[i] = Jaccache[i][1] + else + copyto!(@view(J[:, i]), Jaccache[i]) + end + end + return end - # else - # Enzyme.autodiff(Enzyme.Reverse, f.cons, BatchDuplicated(y, seeds), - # BatchDuplicated(θ, Jaccache), Const(p)) - # for i in 1:num_cons - # if J isa Vector - # J .= Jaccache[1] - # else - # J[i, :] = Jaccache[i] - # end - # end - # end - return end elseif cons_j == true && cons !== nothing - cons_j! = (J, θ) -> f.cons_j(J, θ, p) + cons_j! = (J, θ, p = p) -> f.cons_j(J, θ, p) else cons_j! = nothing end From 5b5446503501abc5df8257229101120d5f03d2d7 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Thu, 18 Jun 2026 11:17:44 -0400 Subject: [PATCH 04/10] format --- lib/OptimizationBase/src/OptimizationDIExt.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index 986f68751..ca9e3ec7e 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -36,7 +36,7 @@ _grad_param_eltype(p) = p isa Union{SciMLBase.NullParameters, Nothing} ? nothing # never veto the fast path. @inline function _grad_use_prep(::Type{T0}, θ, p) where {T0} pe = _grad_param_eltype(p) - eltype(θ) === T0 && (pe === nothing || pe === T0) + return eltype(θ) === T0 && (pe === nothing || pe === T0) end function instantiate_function( @@ -72,7 +72,7 @@ function instantiate_function( if p !== SciMLBase.NullParameters() && p !== nothing let _prep_grad = _prep_grad, f = f, adtype = adtype, T0 = T0 function (res, θ, p = p) - if _grad_use_prep(T0, θ, p) && eltype(res) === T0 + return if _grad_use_prep(T0, θ, p) && eltype(res) === T0 gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) else gradient!(f.f, res, adtype, θ, Constant(p)) @@ -82,7 +82,7 @@ function instantiate_function( else let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, T0 = T0 function (res, θ, p = p) - if _grad_use_prep(T0, θ, p) && eltype(res) === T0 + return if _grad_use_prep(T0, θ, p) && eltype(res) === T0 gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) else gradient!(f.f, res, adtype, θ, Constant(p)) @@ -499,7 +499,7 @@ function instantiate_function( if p !== SciMLBase.NullParameters() && p !== nothing let _prep_grad = _prep_grad, f = f, adtype = adtype, T0 = T0 function (θ, p = p) - _grad_use_prep(T0, θ, p) ? + return _grad_use_prep(T0, θ, p) ? gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : gradient(f.f, adtype, θ, Constant(p)) end @@ -507,7 +507,7 @@ function instantiate_function( else let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, T0 = T0 function (θ, p = p) - _grad_use_prep(T0, θ, p) ? + return _grad_use_prep(T0, θ, p) ? gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : gradient(f.f, adtype, θ, Constant(p)) end From 979184fd9a0081f1b9b6d2dfc49f66b3f7b5beb9 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Thu, 2 Jul 2026 11:37:10 -0400 Subject: [PATCH 05/10] add tests --- lib/OptimizationBase/test/core_tests.jl | 1 + .../test/dual_tolerant_tests.jl | 124 ++++++++++++++++++ 2 files changed, 125 insertions(+) create mode 100644 lib/OptimizationBase/test/dual_tolerant_tests.jl diff --git a/lib/OptimizationBase/test/core_tests.jl b/lib/OptimizationBase/test/core_tests.jl index 0322c3672..fbb6fba03 100644 --- a/lib/OptimizationBase/test/core_tests.jl +++ b/lib/OptimizationBase/test/core_tests.jl @@ -3,6 +3,7 @@ using Test @testset "OptimizationBase.jl" begin include("adtests.jl") + include("dual_tolerant_tests.jl") include("cvxtest.jl") include("matrixvalued.jl") include("solver_missing_error_messages.jl") diff --git a/lib/OptimizationBase/test/dual_tolerant_tests.jl b/lib/OptimizationBase/test/dual_tolerant_tests.jl new file mode 100644 index 000000000..744956b17 --- /dev/null +++ b/lib/OptimizationBase/test/dual_tolerant_tests.jl @@ -0,0 +1,124 @@ +# Tests for the dual-tolerant gradient/Jacobian path and the `p`-accepting +# derivative closures added on the `dual-tolerant-grad` branch. +# +# The behaviors under test (all previously uncovered): +# 1. `_grad_use_prep` gating logic in isolation. +# 2. `grad`/`cons_j` still hit the prepared fast path for the construction types +# and stay numerically correct. +# 3. `grad`/`cons_j` accept an explicit `p` different from the construction `p`. +# 4. Pushing `ForwardDiff.Dual`s (dual `p`, real `θ`) through `grad`/`cons_j` +# does not throw `PreparationMismatchError` and yields the correct +# sensitivity (∂/∂p of the derivative) — the SciMLSensitivity use case. + +using OptimizationBase, Test, ForwardDiff, FiniteDiff +using ADTypes, Enzyme +import SciMLBase +using OptimizationBase: _grad_use_prep + +# Parametrized objective and constraint whose derivatives genuinely depend on `p`, +# so a dual `p` produces a nonzero, checkable sensitivity. +objp(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2 +consp!(res, x, p) = (res[1] = p[1] * x[1]^2 + x[2]^2; return nothing) +consp(x, p) = [p[1] * x[1]^2 + x[2]^2] + +x0 = zeros(2) +xt = [0.7, -0.3] # evaluation point, distinct from x0 +p0 = [2.0, 3.0] # construction parameters +p1 = [5.0, 7.0] # a different parameter value + +∇xf(x, p) = ForwardDiff.gradient(xx -> objp(xx, p), x) +consjac(x, p) = ForwardDiff.jacobian(xx -> consp(xx, p), x) + +@testset "_grad_use_prep gating" begin + d = ForwardDiff.Dual{Nothing}(1.0, 1.0) + # Matching real inputs -> prepared fast path. + @test _grad_use_prep(Float64, [1.0, 2.0], [3.0, 4.0]) + # NullParameters / nothing carry no differentiable params, never veto. + @test _grad_use_prep(Float64, [1.0, 2.0], SciMLBase.NullParameters()) + @test _grad_use_prep(Float64, [1.0, 2.0], nothing) + # Dual θ -> fallback. + @test !_grad_use_prep(Float64, [d, d], [3.0, 4.0]) + # Real θ but dual p (the sensitivity case) -> fallback. + @test !_grad_use_prep(Float64, [1.0, 2.0], [d, d]) + # Mismatched θ element type -> fallback. + @test !_grad_use_prep(Float64, Float32[1.0, 2.0], [3.0, 4.0]) +end + +# Runs the full matrix of assertions against an already-instantiated problem. +# `inplace` selects whether grad/cons_j are the mutating (res, θ[, p]) form. +function check_dual_tolerant(optprob; inplace::Bool, rtol = 1.0e-6) + # --- gradient --------------------------------------------------------- + gref = ∇xf(xt, p0) + if inplace + g = zeros(2) + optprob.grad(g, xt) # fast path, default p + @test g ≈ gref rtol=rtol + optprob.grad(g, xt, p1) # explicit different p + @test g ≈ ∇xf(xt, p1) rtol=rtol + else + @test optprob.grad(xt) ≈ gref rtol=rtol + @test optprob.grad(xt, p1) ≈ ∇xf(xt, p1) rtol=rtol + end + + # Dual p, real θ: sensitivity ∂/∂p ∇ₓf. Must not throw, must match FD. + Jsens_ref = ForwardDiff.jacobian(pp -> ∇xf(xt, pp), p0) + if inplace + gof_p = pp -> (buf = zeros(eltype(pp), 2); optprob.grad(buf, xt, pp); buf) + else + gof_p = pp -> optprob.grad(xt, pp) + end + @test ForwardDiff.jacobian(gof_p, p0) ≈ Jsens_ref rtol=rtol + + # --- constraint Jacobian --------------------------------------------- + Jref = vec(consjac(xt, p0)) + if inplace + J = zeros(2) + optprob.cons_j(J, xt) # fast path, default p + @test J ≈ Jref rtol=rtol + optprob.cons_j(J, xt, p1) # explicit different p + @test J ≈ vec(consjac(xt, p1)) rtol=rtol + else + @test optprob.cons_j(xt) ≈ Jref rtol=rtol + @test optprob.cons_j(xt, p1) ≈ vec(consjac(xt, p1)) rtol=rtol + end + + # Dual p through cons_j: sensitivity ∂/∂p of the constraint Jacobian. + Jcons_sens_ref = ForwardDiff.jacobian(pp -> vec(consjac(xt, pp)), p0) + if inplace + cjof_p = pp -> (J = zeros(eltype(pp), 2); optprob.cons_j(J, xt, pp); J) + else + cjof_p = pp -> optprob.cons_j(xt, pp) + end + @test ForwardDiff.jacobian(cjof_p, p0) ≈ Jcons_sens_ref rtol=rtol +end + +@testset "dual-tolerant grad / parametrized cons_j (DI)" begin + @testset "AutoForwardDiff in-place" begin + optf = OptimizationFunction(objp, ADTypes.AutoForwardDiff(); cons = consp!) + optprob = OptimizationBase.instantiate_function( + optf, x0, ADTypes.AutoForwardDiff(), p0, 1; g = true, cons_j = true) + check_dual_tolerant(optprob; inplace = true) + end + + @testset "AutoForwardDiff out-of-place" begin + optf = OptimizationFunction{false}(objp, ADTypes.AutoForwardDiff(); cons = consp) + optprob = OptimizationBase.instantiate_function( + optf, x0, ADTypes.AutoForwardDiff(), p0, 1; g = true, cons_j = true) + check_dual_tolerant(optprob; inplace = false) + end +end + +@testset "parametrized cons_j (Enzyme)" begin + # The dual-through path is not exercised for Enzyme (ForwardDiff-over-Enzyme + # nesting is out of scope); this pins the `p`-accepting closure and the + # de-boxed fast path stay numerically correct at the default and explicit p. + optf = OptimizationFunction(objp, ADTypes.AutoEnzyme(); cons = consp!) + optprob = OptimizationBase.instantiate_function( + optf, x0, ADTypes.AutoEnzyme(), p0, 1; g = true, cons_j = true) + + J = zeros(2) + optprob.cons_j(J, xt) # default p + @test J ≈ vec(consjac(xt, p0)) rtol=1.0e-6 + optprob.cons_j(J, xt, p1) # explicit different p + @test J ≈ vec(consjac(xt, p1)) rtol=1.0e-6 +end From 659da2094b64d8e4a35b5f651fd14013bb3ad94c Mon Sep 17 00:00:00 2001 From: jClugstor Date: Thu, 2 Jul 2026 11:56:14 -0400 Subject: [PATCH 06/10] format test --- .../test/dual_tolerant_tests.jl | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/lib/OptimizationBase/test/dual_tolerant_tests.jl b/lib/OptimizationBase/test/dual_tolerant_tests.jl index 744956b17..f1946e73d 100644 --- a/lib/OptimizationBase/test/dual_tolerant_tests.jl +++ b/lib/OptimizationBase/test/dual_tolerant_tests.jl @@ -52,12 +52,12 @@ function check_dual_tolerant(optprob; inplace::Bool, rtol = 1.0e-6) if inplace g = zeros(2) optprob.grad(g, xt) # fast path, default p - @test g ≈ gref rtol=rtol + @test g ≈ gref rtol = rtol optprob.grad(g, xt, p1) # explicit different p - @test g ≈ ∇xf(xt, p1) rtol=rtol + @test g ≈ ∇xf(xt, p1) rtol = rtol else - @test optprob.grad(xt) ≈ gref rtol=rtol - @test optprob.grad(xt, p1) ≈ ∇xf(xt, p1) rtol=rtol + @test optprob.grad(xt) ≈ gref rtol = rtol + @test optprob.grad(xt, p1) ≈ ∇xf(xt, p1) rtol = rtol end # Dual p, real θ: sensitivity ∂/∂p ∇ₓf. Must not throw, must match FD. @@ -67,19 +67,19 @@ function check_dual_tolerant(optprob; inplace::Bool, rtol = 1.0e-6) else gof_p = pp -> optprob.grad(xt, pp) end - @test ForwardDiff.jacobian(gof_p, p0) ≈ Jsens_ref rtol=rtol + @test ForwardDiff.jacobian(gof_p, p0) ≈ Jsens_ref rtol = rtol # --- constraint Jacobian --------------------------------------------- Jref = vec(consjac(xt, p0)) if inplace J = zeros(2) optprob.cons_j(J, xt) # fast path, default p - @test J ≈ Jref rtol=rtol + @test J ≈ Jref rtol = rtol optprob.cons_j(J, xt, p1) # explicit different p - @test J ≈ vec(consjac(xt, p1)) rtol=rtol + @test J ≈ vec(consjac(xt, p1)) rtol = rtol else - @test optprob.cons_j(xt) ≈ Jref rtol=rtol - @test optprob.cons_j(xt, p1) ≈ vec(consjac(xt, p1)) rtol=rtol + @test optprob.cons_j(xt) ≈ Jref rtol = rtol + @test optprob.cons_j(xt, p1) ≈ vec(consjac(xt, p1)) rtol = rtol end # Dual p through cons_j: sensitivity ∂/∂p of the constraint Jacobian. @@ -89,21 +89,23 @@ function check_dual_tolerant(optprob; inplace::Bool, rtol = 1.0e-6) else cjof_p = pp -> optprob.cons_j(xt, pp) end - @test ForwardDiff.jacobian(cjof_p, p0) ≈ Jcons_sens_ref rtol=rtol + return @test ForwardDiff.jacobian(cjof_p, p0) ≈ Jcons_sens_ref rtol = rtol end @testset "dual-tolerant grad / parametrized cons_j (DI)" begin @testset "AutoForwardDiff in-place" begin optf = OptimizationFunction(objp, ADTypes.AutoForwardDiff(); cons = consp!) optprob = OptimizationBase.instantiate_function( - optf, x0, ADTypes.AutoForwardDiff(), p0, 1; g = true, cons_j = true) + optf, x0, ADTypes.AutoForwardDiff(), p0, 1; g = true, cons_j = true + ) check_dual_tolerant(optprob; inplace = true) end @testset "AutoForwardDiff out-of-place" begin optf = OptimizationFunction{false}(objp, ADTypes.AutoForwardDiff(); cons = consp) optprob = OptimizationBase.instantiate_function( - optf, x0, ADTypes.AutoForwardDiff(), p0, 1; g = true, cons_j = true) + optf, x0, ADTypes.AutoForwardDiff(), p0, 1; g = true, cons_j = true + ) check_dual_tolerant(optprob; inplace = false) end end @@ -114,11 +116,12 @@ end # de-boxed fast path stay numerically correct at the default and explicit p. optf = OptimizationFunction(objp, ADTypes.AutoEnzyme(); cons = consp!) optprob = OptimizationBase.instantiate_function( - optf, x0, ADTypes.AutoEnzyme(), p0, 1; g = true, cons_j = true) + optf, x0, ADTypes.AutoEnzyme(), p0, 1; g = true, cons_j = true + ) J = zeros(2) optprob.cons_j(J, xt) # default p - @test J ≈ vec(consjac(xt, p0)) rtol=1.0e-6 + @test J ≈ vec(consjac(xt, p0)) rtol = 1.0e-6 optprob.cons_j(J, xt, p1) # explicit different p - @test J ≈ vec(consjac(xt, p1)) rtol=1.0e-6 + @test J ≈ vec(consjac(xt, p1)) rtol = 1.0e-6 end From 39ee61f87ac9897e8bb50cf3bed1f26994bb2a75 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Thu, 2 Jul 2026 13:22:33 -0400 Subject: [PATCH 07/10] fix path for p isa Tuple, add test --- lib/OptimizationBase/src/OptimizationDIExt.jl | 33 +++++++++++++------ .../test/dual_tolerant_tests.jl | 23 +++++++++++++ 2 files changed, 46 insertions(+), 10 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index ca9e3ec7e..fb2078020 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -28,17 +28,33 @@ using OptimizationBase.FastClosures # for the prepared type and fall back to a prep-free DI call (which prepares at the actual # argument types each call) for anything else. The fallback cost is irrelevant off the # optimization hot loop — sensitivity does one such call per solve, not per iteration. -_grad_param_eltype(p) = p isa Union{SciMLBase.NullParameters, Nothing} ? nothing : eltype(p) +# Scalar eltype of `p`, or `nothing` when `p` has no scalar to match: NullParameters/nothing, +# and structured `p` (tuple/array-of-arrays) whose `eltype` is not a `Number`. Only a numeric +# `p` returns a type that can veto the fast path. +function _grad_param_eltype(p) + p isa Union{SciMLBase.NullParameters, Nothing} && return nothing + pe = eltype(p) + return pe <: Number ? pe : nothing +end # True when every float-bearing input matches the prepared element type `T0`, so the prepared -# path is valid. A dual `θ` or dual `p` (the sensitivity case) flunks this and routes to the -# prep-free fallback. `NullParameters`/`nothing` carry no differentiable parameters, so they -# never veto the fast path. +# path is valid. A dual `θ` or dual numeric `p` (the sensitivity case) flunks this and routes to +# the prep-free fallback; `nothing`-eltype `p` never vetoes. @inline function _grad_use_prep(::Type{T0}, θ, p) where {T0} pe = _grad_param_eltype(p) return eltype(θ) === T0 && (pe === nothing || pe === T0) end +# Output-buffer eltype for the `p`-accepting constraint wrapper. Promote against `eltype(p)` only +# when it is a `Number` (so a dual `p` propagates in the sensitivity case); a structured `p` has +# a non-`Number` eltype that would promote to `Union{}`, so fall back to `eltype(x)` there. +@inline function _cons_out_eltype(x, p) + p isa Union{SciMLBase.NullParameters, Nothing} && return eltype(x) + pe = eltype(p) + T = pe <: Number ? Base.promote_op(+, eltype(x), pe) : eltype(x) + return T === Union{} ? eltype(x) : T +end + function instantiate_function( f::OptimizationFunction{true}, x, ::ADTypes.AutoSparse{<:ADTypes.AutoSymbolics}, args...; kwargs... @@ -237,14 +253,11 @@ function instantiate_function( # at parameters other than the construction `p` — including duals pushed in by a # sensitivity layer differentiating the constraint Jacobian w.r.t. `p` (the mixed # ∂²cᵢ/∂x∂p term of the KKT residual). The prepared `cons_oop` bakes `p` in and - # exposes no parameter slot, so we cannot reuse it here. Output eltype promotes - # against `q` so duals propagate (Base.promote_op, not promote_type: ForwardDiff's - # dual ops bypass promote_type). + # exposes no parameter slot, so we cannot reuse it here. `_cons_out_eltype` picks the + # output eltype so duals propagate without poisoning it to `Union{}`. _cons_oop_p = let f = f, num_cons = num_cons function (x, p) - T = p isa Union{SciMLBase.NullParameters, Nothing} ? eltype(x) : - Base.promote_op(+, eltype(x), eltype(p)) - res = Vector{T}(undef, num_cons) + res = Vector{_cons_out_eltype(x, p)}(undef, num_cons) f.cons(res, x, p) return res end diff --git a/lib/OptimizationBase/test/dual_tolerant_tests.jl b/lib/OptimizationBase/test/dual_tolerant_tests.jl index f1946e73d..1802046c8 100644 --- a/lib/OptimizationBase/test/dual_tolerant_tests.jl +++ b/lib/OptimizationBase/test/dual_tolerant_tests.jl @@ -110,6 +110,29 @@ end end end +@testset "structured (tuple) parameters" begin + # Regression: a tuple-valued `p` has a non-`Number` eltype. The output-buffer eltype + # must not promote to `Union{}` (which crashed the constraint Jacobian), and such a `p` + # must not veto the prepared fast path — it carries no scalar for the prep to match. + losst(x, p) = sum(abs2, p[1] * x .- p[2]) + tcons!(res, x, p) = (res[1] = sum(abs2, x) - 1.0; return nothing) + pt = ([1.0 0.5; 0.5 1.0; 0.2 0.3], [0.1, 0.2, 0.3]) # (Matrix, Vector) tuple + + optf = OptimizationFunction(losst, ADTypes.AutoForwardDiff(); cons = tcons!) + optprob = OptimizationBase.instantiate_function( + optf, x0, ADTypes.AutoForwardDiff(), pt, 1; g = true, cons_j = true + ) + + J = zeros(2) + optprob.cons_j(J, xt) + @test J ≈ [2xt[1], 2xt[2]] rtol = 1.0e-6 # ∂(‖x‖²-1)/∂x + g = zeros(2) + optprob.grad(g, xt) + @test g ≈ ForwardDiff.gradient(xx -> losst(xx, pt), xt) rtol = 1.0e-6 + # A structured `p` must still route through the prepared fast path. + @test OptimizationBase._grad_use_prep(Float64, xt, pt) +end + @testset "parametrized cons_j (Enzyme)" begin # The dual-through path is not exercised for Enzyme (ForwardDiff-over-Enzyme # nesting is out of scope); this pins the `p`-accepting closure and the From 181f01e4ca69b41a6e320f9d8c5eae27abcc4976 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Thu, 2 Jul 2026 15:02:43 -0400 Subject: [PATCH 08/10] use anyeltypedual and make more robust --- lib/OptimizationBase/src/OptimizationDIExt.jl | 101 ++++++++---------- .../test/dual_tolerant_tests.jl | 29 ++--- 2 files changed, 62 insertions(+), 68 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index fb2078020..9399d08c4 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -17,42 +17,37 @@ using OptimizationBase.FastClosures # --- Dual-tolerant gradient dispatch ------------------------------------------------------ # A DI preparation is monomorphic: built once at the construction types (`x`, `p`), its -# internal buffers are concrete (e.g. `Float64`) and reject inputs of any other type. That is -# correct and fast for the optimization solve, but a downstream sensitivity layer (e.g. -# SciMLSensitivity's `OptimizationAdjoint`) differentiates the KKT stationarity conditions -# w.r.t. the parameters by pushing `ForwardDiff.Dual`s through the gradient: a dual `p` (and a -# dual output buffer) evaluated at a real `θ = x*`. Those duals hit the prepared buffers and -# throw `DifferentiationInterface.PreparationMismatchError`. +# internal buffers are concrete (e.g. `Float64`) and reject `ForwardDiff.Dual`s. That is correct +# and fast for the solve, but a downstream sensitivity layer (e.g. SciMLSensitivity's +# `OptimizationAdjoint`) differentiates the KKT conditions w.r.t. the parameters by pushing duals +# through `grad`/`cons_j` — a dual `p` at a real `θ = x*`. Those duals hit the prepared buffers +# and throw `DifferentiationInterface.PreparationMismatchError`. # -# To support that without slowing the solve, the gradient closures keep the prepared fast path -# for the prepared type and fall back to a prep-free DI call (which prepares at the actual -# argument types each call) for anything else. The fallback cost is irrelevant off the -# optimization hot loop — sensitivity does one such call per solve, not per iteration. -# Scalar eltype of `p`, or `nothing` when `p` has no scalar to match: NullParameters/nothing, -# and structured `p` (tuple/array-of-arrays) whose `eltype` is not a `Number`. Only a numeric -# `p` returns a type that can veto the fast path. -function _grad_param_eltype(p) - p isa Union{SciMLBase.NullParameters, Nothing} && return nothing - pe = eltype(p) - return pe <: Number ? pe : nothing +# So the closures keep the prepared fast path whenever the call carries no dual, and fall back to +# a prep-free DI call (which prepares at the actual argument types) when one does. Dual-detection +# reuses SciMLBase's `anyeltypedual` (`== Any` ⇒ no dual, the same convention `promote_u0` uses), +# which recurses into a structured `p`, so a dual nested in a tuple/ComponentArray is caught too. +# A *real* `p` of any eltype (Int, Float32, …) keeps the fast path: `p` enters as a Constant, so +# its eltype never invalidates the prepared gradient. +@inline function _use_prep(θ, p) + return SciMLBase.anyeltypedual(θ) == Any && SciMLBase.anyeltypedual(p) == Any end -# True when every float-bearing input matches the prepared element type `T0`, so the prepared -# path is valid. A dual `θ` or dual numeric `p` (the sensitivity case) flunks this and routes to -# the prep-free fallback; `nothing`-eltype `p` never vetoes. -@inline function _grad_use_prep(::Type{T0}, θ, p) where {T0} - pe = _grad_param_eltype(p) - return eltype(θ) === T0 && (pe === nothing || pe === T0) -end - -# Output-buffer eltype for the `p`-accepting constraint wrapper. Promote against `eltype(p)` only -# when it is a `Number` (so a dual `p` propagates in the sensitivity case); a structured `p` has -# a non-`Number` eltype that would promote to `Union{}`, so fall back to `eltype(x)` there. +# Output-buffer eltype for the `p`-accepting constraint wrapper: the type `f.cons` produces, +# including the *nested* dual when both `x` (DI's seeds) and `p` (the sensitivity layer) carry +# duals with different tags (`promote_op(+, …)` nests them). Deliberately uses plain `eltype(p)` +# rather than `SciMLBase.anyeltypedual`: this runs *inside* the DI-differentiated wrapper, and +# Enzyme's forward mode corrupts the derivative shadow (DataType-valued entries) when the +# allocation type flows through `anyeltypedual` — in either its value or its type-level form, +# and even with an `EnzymeRules.inactive` mark on this helper. A structured `p` (non-`Number` +# eltype) therefore falls back to `eltype(x)`; duals nested inside such a `p` are the one +# unsupported case. (`_use_prep` above is free to use `anyeltypedual`: it runs in the outer +# closure, outside anything a backend differentiates.) @inline function _cons_out_eltype(x, p) - p isa Union{SciMLBase.NullParameters, Nothing} && return eltype(x) - pe = eltype(p) - T = pe <: Number ? Base.promote_op(+, eltype(x), pe) : eltype(x) - return T === Union{} ? eltype(x) : T + Tu = eltype(x) + p isa Union{SciMLBase.NullParameters, Nothing} && return Tu + Tp = eltype(p) + return Tp <: Number ? Base.promote_op(+, Tu, Tp) : Tu end function instantiate_function( @@ -80,15 +75,14 @@ function instantiate_function( adtype, soadtype = generate_adtype(adtype) # Create gradient closures with proper type stability using let blocks. - # `T0 = eltype(x)` is the prepared element type; `_grad_use_prep` gates the prepared fast - # path vs the prep-free fallback (see the note above the imports). + # `_use_prep` gates the prepared fast path vs the prep-free dual fallback (see the note + # above the imports). grad = if g == true && f.grad === nothing _prep_grad = prepare_gradient(f.f, adtype, x, Constant(p)) - T0 = eltype(x) if p !== SciMLBase.NullParameters() && p !== nothing - let _prep_grad = _prep_grad, f = f, adtype = adtype, T0 = T0 + let _prep_grad = _prep_grad, f = f, adtype = adtype function (res, θ, p = p) - return if _grad_use_prep(T0, θ, p) && eltype(res) === T0 + return if _use_prep(θ, p) gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) else gradient!(f.f, res, adtype, θ, Constant(p)) @@ -96,9 +90,9 @@ function instantiate_function( end end else - let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, T0 = T0 + let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p function (res, θ, p = p) - return if _grad_use_prep(T0, θ, p) && eltype(res) === T0 + return if _use_prep(θ, p) gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) else gradient!(f.f, res, adtype, θ, Constant(p)) @@ -263,12 +257,11 @@ function instantiate_function( end end _prep_jac = prepare_jacobian(_cons_oop_p, adtype, x, Constant(p)) - T0 = eltype(x) - let _cons_oop_p = _cons_oop_p, _prep_jac = _prep_jac, adtype = adtype, p = p, T0 = T0 + let _cons_oop_p = _cons_oop_p, _prep_jac = _prep_jac, adtype = adtype, p = p function (J, θ, p = p) - # Prepared fast path when types match the prep; prep-free fallback for duals - # (see the `_grad_use_prep` note above the imports). - if _grad_use_prep(T0, θ, p) && eltype(J) === T0 + # Prepared fast path unless a dual is present; prep-free fallback for duals + # (see the `_use_prep` note above the imports). + if _use_prep(θ, p) jacobian!(_cons_oop_p, J, _prep_jac, adtype, θ, Constant(p)) else jacobian!(_cons_oop_p, J, adtype, θ, Constant(p)) @@ -504,23 +497,22 @@ function instantiate_function( adtype, soadtype = generate_adtype(adtype) # Create gradient closures with proper type stability using let blocks. - # `T0 = eltype(x)` is the prepared element type; `_grad_use_prep` gates the prepared fast - # path vs the prep-free fallback (see the note above the imports). + # `_use_prep` gates the prepared fast path vs the prep-free dual fallback (see the note + # above the imports). grad = if g == true && f.grad === nothing _prep_grad = prepare_gradient(f.f, adtype, x, Constant(p)) - T0 = eltype(x) if p !== SciMLBase.NullParameters() && p !== nothing - let _prep_grad = _prep_grad, f = f, adtype = adtype, T0 = T0 + let _prep_grad = _prep_grad, f = f, adtype = adtype function (θ, p = p) - return _grad_use_prep(T0, θ, p) ? + return _use_prep(θ, p) ? gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : gradient(f.f, adtype, θ, Constant(p)) end end else - let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, T0 = T0 + let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p function (θ, p = p) - return _grad_use_prep(T0, θ, p) ? + return _use_prep(θ, p) ? gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : gradient(f.f, adtype, θ, Constant(p)) end @@ -648,13 +640,12 @@ function instantiate_function( cons_j! = if f.cons !== nothing && cons_j == true && f.cons_j === nothing # `f.cons` is out-of-place here and the prep already takes `Constant(p)`, so this # only needs to expose the parameter argument and add the dual-tolerant fallback - # (see the `_grad_use_prep` note above the imports) — unlike the in-place method, + # (see the `_use_prep` note above the imports) — unlike the in-place method, # whose prepared wrapper bakes `p` in. _prep_jac = prepare_jacobian(f.cons, adtype, x, Constant(p)) - T0 = eltype(x) - let f = f, _prep_jac = _prep_jac, adtype = adtype, p = p, T0 = T0 + let f = f, _prep_jac = _prep_jac, adtype = adtype, p = p function (θ, p = p) - J = _grad_use_prep(T0, θ, p) ? + J = _use_prep(θ, p) ? jacobian(f.cons, _prep_jac, adtype, θ, Constant(p)) : jacobian(f.cons, adtype, θ, Constant(p)) if size(J, 1) == 1 diff --git a/lib/OptimizationBase/test/dual_tolerant_tests.jl b/lib/OptimizationBase/test/dual_tolerant_tests.jl index 1802046c8..1a01b6dde 100644 --- a/lib/OptimizationBase/test/dual_tolerant_tests.jl +++ b/lib/OptimizationBase/test/dual_tolerant_tests.jl @@ -2,7 +2,7 @@ # derivative closures added on the `dual-tolerant-grad` branch. # # The behaviors under test (all previously uncovered): -# 1. `_grad_use_prep` gating logic in isolation. +# 1. `_use_prep` dual-detection gating logic in isolation. # 2. `grad`/`cons_j` still hit the prepared fast path for the construction types # and stay numerically correct. # 3. `grad`/`cons_j` accept an explicit `p` different from the construction `p`. @@ -13,7 +13,7 @@ using OptimizationBase, Test, ForwardDiff, FiniteDiff using ADTypes, Enzyme import SciMLBase -using OptimizationBase: _grad_use_prep +using OptimizationBase: _use_prep # Parametrized objective and constraint whose derivatives genuinely depend on `p`, # so a dual `p` produces a nonzero, checkable sensitivity. @@ -29,19 +29,22 @@ p1 = [5.0, 7.0] # a different parameter value ∇xf(x, p) = ForwardDiff.gradient(xx -> objp(xx, p), x) consjac(x, p) = ForwardDiff.jacobian(xx -> consp(xx, p), x) -@testset "_grad_use_prep gating" begin +@testset "_use_prep dual gating" begin d = ForwardDiff.Dual{Nothing}(1.0, 1.0) - # Matching real inputs -> prepared fast path. - @test _grad_use_prep(Float64, [1.0, 2.0], [3.0, 4.0]) - # NullParameters / nothing carry no differentiable params, never veto. - @test _grad_use_prep(Float64, [1.0, 2.0], SciMLBase.NullParameters()) - @test _grad_use_prep(Float64, [1.0, 2.0], nothing) + # No duals -> prepared fast path, regardless of real parameter eltype. + @test _use_prep([1.0, 2.0], [3.0, 4.0]) + @test _use_prep([1.0, 2.0], SciMLBase.NullParameters()) + @test _use_prep([1.0, 2.0], nothing) + @test _use_prep([1.0, 2.0], [1, 100]) # Int p keeps the fast path + @test _use_prep([1.0, 2.0], Float32[1, 2]) # mixed-precision p keeps it too + @test _use_prep([1.0, 2.0], (rand(2, 2), rand(2))) # structured (tuple) p keeps it + @test _use_prep(Float32[1.0, 2.0], [3.0, 4.0]) # off-construction real eltype: still fast # Dual θ -> fallback. - @test !_grad_use_prep(Float64, [d, d], [3.0, 4.0]) + @test !_use_prep([d, d], [3.0, 4.0]) # Real θ but dual p (the sensitivity case) -> fallback. - @test !_grad_use_prep(Float64, [1.0, 2.0], [d, d]) - # Mismatched θ element type -> fallback. - @test !_grad_use_prep(Float64, Float32[1.0, 2.0], [3.0, 4.0]) + @test !_use_prep([1.0, 2.0], [d, d]) + # Dual nested inside a structured p -> fallback (anyeltypedual recurses). + @test !_use_prep([1.0, 2.0], ([d, d], [1.0])) end # Runs the full matrix of assertions against an already-instantiated problem. @@ -130,7 +133,7 @@ end optprob.grad(g, xt) @test g ≈ ForwardDiff.gradient(xx -> losst(xx, pt), xt) rtol = 1.0e-6 # A structured `p` must still route through the prepared fast path. - @test OptimizationBase._grad_use_prep(Float64, xt, pt) + @test _use_prep(xt, pt) end @testset "parametrized cons_j (Enzyme)" begin From 37f5d9d82bc3b2da278e573cce86306b3179ef9a Mon Sep 17 00:00:00 2001 From: jClugstor Date: Wed, 8 Jul 2026 11:31:56 -0400 Subject: [PATCH 09/10] check full types instead of anyeltypedual for prep gating --- lib/OptimizationBase/src/OptimizationDIExt.jl | 74 +++++++++---------- .../test/dual_tolerant_tests.jl | 69 ++++++++++++----- 2 files changed, 85 insertions(+), 58 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index 9399d08c4..c762d2c65 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -15,23 +15,13 @@ import DifferentiationInterface: prepare_gradient, prepare_hessian, prepare_hvp, using ADTypes, SciMLBase using OptimizationBase.FastClosures -# --- Dual-tolerant gradient dispatch ------------------------------------------------------ -# A DI preparation is monomorphic: built once at the construction types (`x`, `p`), its -# internal buffers are concrete (e.g. `Float64`) and reject `ForwardDiff.Dual`s. That is correct -# and fast for the solve, but a downstream sensitivity layer (e.g. SciMLSensitivity's -# `OptimizationAdjoint`) differentiates the KKT conditions w.r.t. the parameters by pushing duals -# through `grad`/`cons_j` — a dual `p` at a real `θ = x*`. Those duals hit the prepared buffers -# and throw `DifferentiationInterface.PreparationMismatchError`. -# -# So the closures keep the prepared fast path whenever the call carries no dual, and fall back to -# a prep-free DI call (which prepares at the actual argument types) when one does. Dual-detection -# reuses SciMLBase's `anyeltypedual` (`== Any` ⇒ no dual, the same convention `promote_u0` uses), -# which recurses into a structured `p`, so a dual nested in a tuple/ComponentArray is caught too. -# A *real* `p` of any eltype (Int, Float32, …) keeps the fast path: `p` enters as a Constant, so -# its eltype never invalidates the prepared gradient. -@inline function _use_prep(θ, p) - return SciMLBase.anyeltypedual(θ) == Any && SciMLBase.anyeltypedual(p) == Any -end +# A DI preparation is built for the exact construction types (`x` and `Constant(p)`) and only +# works for those — otherwise DI throws `PreparationMismatchError`. Reuse it while `θ`/`p` keep +# those types; fall back to a prep-free call otherwise. We compare types exactly rather than just +# looking for duals, so any off-type is caught (a dual `p` from a sensitivity layer, but equally a +# `Float32`/tracked/`BigFloat` `θ` that a prior `anyeltypedual` gate let slip through and error). +# `T` is a constant, so on the solve path this folds away. +@inline _prep_valid(::Type{T}, v) where {T} = typeof(v) === T # Output-buffer eltype for the `p`-accepting constraint wrapper: the type `f.cons` produces, # including the *nested* dual when both `x` (DI's seeds) and `p` (the sensitivity layer) carry @@ -41,8 +31,8 @@ end # allocation type flows through `anyeltypedual` — in either its value or its type-level form, # and even with an `EnzymeRules.inactive` mark on this helper. A structured `p` (non-`Number` # eltype) therefore falls back to `eltype(x)`; duals nested inside such a `p` are the one -# unsupported case. (`_use_prep` above is free to use `anyeltypedual`: it runs in the outer -# closure, outside anything a backend differentiates.) +# unsupported case. (`_prep_valid` above runs in the outer closure, outside anything a backend +# differentiates, so its type comparison is safe there.) @inline function _cons_out_eltype(x, p) Tu = eltype(x) p isa Union{SciMLBase.NullParameters, Nothing} && return Tu @@ -74,15 +64,18 @@ function instantiate_function( ) adtype, soadtype = generate_adtype(adtype) + # Construction types the DI preps are built at; `_prep_valid` gates the prepared fast path + # vs the prep-free fallback (see the note above the imports). + Tx0 = typeof(x) + Tp0 = typeof(p) + # Create gradient closures with proper type stability using let blocks. - # `_use_prep` gates the prepared fast path vs the prep-free dual fallback (see the note - # above the imports). grad = if g == true && f.grad === nothing _prep_grad = prepare_gradient(f.f, adtype, x, Constant(p)) if p !== SciMLBase.NullParameters() && p !== nothing - let _prep_grad = _prep_grad, f = f, adtype = adtype + let _prep_grad = _prep_grad, f = f, adtype = adtype, Tx0 = Tx0, Tp0 = Tp0 function (res, θ, p = p) - return if _use_prep(θ, p) + return if _prep_valid(Tx0, θ) && _prep_valid(Tp0, p) gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) else gradient!(f.f, res, adtype, θ, Constant(p)) @@ -90,9 +83,9 @@ function instantiate_function( end end else - let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p + let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, Tx0 = Tx0, Tp0 = Tp0 function (res, θ, p = p) - return if _use_prep(θ, p) + return if _prep_valid(Tx0, θ) && _prep_valid(Tp0, p) gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) else gradient!(f.f, res, adtype, θ, Constant(p)) @@ -257,11 +250,11 @@ function instantiate_function( end end _prep_jac = prepare_jacobian(_cons_oop_p, adtype, x, Constant(p)) - let _cons_oop_p = _cons_oop_p, _prep_jac = _prep_jac, adtype = adtype, p = p + let _cons_oop_p = _cons_oop_p, _prep_jac = _prep_jac, adtype = adtype, p = p, Tx0 = Tx0, Tp0 = Tp0 function (J, θ, p = p) - # Prepared fast path unless a dual is present; prep-free fallback for duals - # (see the `_use_prep` note above the imports). - if _use_prep(θ, p) + # Prepared fast path when the call types match construction; prep-free fallback + # otherwise (see the `_prep_valid` note above the imports). + if _prep_valid(Tx0, θ) && _prep_valid(Tp0, p) jacobian!(_cons_oop_p, J, _prep_jac, adtype, θ, Constant(p)) else jacobian!(_cons_oop_p, J, adtype, θ, Constant(p)) @@ -496,23 +489,26 @@ function instantiate_function( ) adtype, soadtype = generate_adtype(adtype) + # Construction types the DI preps are built at; `_prep_valid` gates the prepared fast path + # vs the prep-free fallback (see the note above the imports). + Tx0 = typeof(x) + Tp0 = typeof(p) + # Create gradient closures with proper type stability using let blocks. - # `_use_prep` gates the prepared fast path vs the prep-free dual fallback (see the note - # above the imports). grad = if g == true && f.grad === nothing _prep_grad = prepare_gradient(f.f, adtype, x, Constant(p)) if p !== SciMLBase.NullParameters() && p !== nothing - let _prep_grad = _prep_grad, f = f, adtype = adtype + let _prep_grad = _prep_grad, f = f, adtype = adtype, Tx0 = Tx0, Tp0 = Tp0 function (θ, p = p) - return _use_prep(θ, p) ? + return _prep_valid(Tx0, θ) && _prep_valid(Tp0, p) ? gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : gradient(f.f, adtype, θ, Constant(p)) end end else - let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p + let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, Tx0 = Tx0, Tp0 = Tp0 function (θ, p = p) - return _use_prep(θ, p) ? + return _prep_valid(Tx0, θ) && _prep_valid(Tp0, p) ? gradient(f.f, _prep_grad, adtype, θ, Constant(p)) : gradient(f.f, adtype, θ, Constant(p)) end @@ -639,13 +635,13 @@ function instantiate_function( cons_j! = if f.cons !== nothing && cons_j == true && f.cons_j === nothing # `f.cons` is out-of-place here and the prep already takes `Constant(p)`, so this - # only needs to expose the parameter argument and add the dual-tolerant fallback - # (see the `_use_prep` note above the imports) — unlike the in-place method, + # only needs to expose the parameter argument and add the prep-validity fallback + # (see the `_prep_valid` note above the imports) — unlike the in-place method, # whose prepared wrapper bakes `p` in. _prep_jac = prepare_jacobian(f.cons, adtype, x, Constant(p)) - let f = f, _prep_jac = _prep_jac, adtype = adtype, p = p + let f = f, _prep_jac = _prep_jac, adtype = adtype, p = p, Tx0 = Tx0, Tp0 = Tp0 function (θ, p = p) - J = _use_prep(θ, p) ? + J = _prep_valid(Tx0, θ) && _prep_valid(Tp0, p) ? jacobian(f.cons, _prep_jac, adtype, θ, Constant(p)) : jacobian(f.cons, adtype, θ, Constant(p)) if size(J, 1) == 1 diff --git a/lib/OptimizationBase/test/dual_tolerant_tests.jl b/lib/OptimizationBase/test/dual_tolerant_tests.jl index 1a01b6dde..6bda15ecc 100644 --- a/lib/OptimizationBase/test/dual_tolerant_tests.jl +++ b/lib/OptimizationBase/test/dual_tolerant_tests.jl @@ -2,18 +2,20 @@ # derivative closures added on the `dual-tolerant-grad` branch. # # The behaviors under test (all previously uncovered): -# 1. `_use_prep` dual-detection gating logic in isolation. +# 1. `_prep_valid` type-match gating logic in isolation. # 2. `grad`/`cons_j` still hit the prepared fast path for the construction types # and stay numerically correct. # 3. `grad`/`cons_j` accept an explicit `p` different from the construction `p`. # 4. Pushing `ForwardDiff.Dual`s (dual `p`, real `θ`) through `grad`/`cons_j` # does not throw `PreparationMismatchError` and yields the correct # sensitivity (∂/∂p of the derivative) — the SciMLSensitivity use case. +# 5. Any other off-construction eltype (`Float32`, `BigFloat`, …) routes through +# the prep-free fallback instead of erroring on the prep built at the construction types. using OptimizationBase, Test, ForwardDiff, FiniteDiff using ADTypes, Enzyme import SciMLBase -using OptimizationBase: _use_prep +using OptimizationBase: _prep_valid # Parametrized objective and constraint whose derivatives genuinely depend on `p`, # so a dual `p` produces a nonzero, checkable sensitivity. @@ -29,22 +31,22 @@ p1 = [5.0, 7.0] # a different parameter value ∇xf(x, p) = ForwardDiff.gradient(xx -> objp(xx, p), x) consjac(x, p) = ForwardDiff.jacobian(xx -> consp(xx, p), x) -@testset "_use_prep dual gating" begin +@testset "_prep_valid type-match gating" begin d = ForwardDiff.Dual{Nothing}(1.0, 1.0) - # No duals -> prepared fast path, regardless of real parameter eltype. - @test _use_prep([1.0, 2.0], [3.0, 4.0]) - @test _use_prep([1.0, 2.0], SciMLBase.NullParameters()) - @test _use_prep([1.0, 2.0], nothing) - @test _use_prep([1.0, 2.0], [1, 100]) # Int p keeps the fast path - @test _use_prep([1.0, 2.0], Float32[1, 2]) # mixed-precision p keeps it too - @test _use_prep([1.0, 2.0], (rand(2, 2), rand(2))) # structured (tuple) p keeps it - @test _use_prep(Float32[1.0, 2.0], [3.0, 4.0]) # off-construction real eltype: still fast - # Dual θ -> fallback. - @test !_use_prep([d, d], [3.0, 4.0]) - # Real θ but dual p (the sensitivity case) -> fallback. - @test !_use_prep([1.0, 2.0], [d, d]) - # Dual nested inside a structured p -> fallback (anyeltypedual recurses). - @test !_use_prep([1.0, 2.0], ([d, d], [1.0])) + Tx0 = Vector{Float64} # what the closures capture as the construction type + + # Exact construction type -> prepared fast path. + @test _prep_valid(Tx0, [1.0, 2.0]) + @test _prep_valid(typeof(SciMLBase.NullParameters()), SciMLBase.NullParameters()) + @test _prep_valid(Nothing, nothing) + @test _prep_valid(typeof((rand(2, 2), rand(2))), (rand(2, 2), rand(2))) # structured p + + # Any deviating type -> fallback. Unlike `anyeltypedual`, non-dual types are caught too. + @test !_prep_valid(Tx0, [d, d]) # dual + @test !_prep_valid(Tx0, Float32[1.0, 2.0]) # Float32 (the bug fix) + @test !_prep_valid(Tx0, big.([1.0, 2.0])) # BigFloat + @test !_prep_valid(Tx0, [1, 100]) # Int + @test !_prep_valid(typeof((rand(2), [1.0])), ([d, d], [1.0])) # dual nested in a tuple end # Runs the full matrix of assertions against an already-instantiated problem. @@ -132,8 +134,37 @@ end g = zeros(2) optprob.grad(g, xt) @test g ≈ ForwardDiff.gradient(xx -> losst(xx, pt), xt) rtol = 1.0e-6 - # A structured `p` must still route through the prepared fast path. - @test _use_prep(xt, pt) + # A structured `p`, at its construction type, must still route through the fast path. + @test _prep_valid(typeof(x0), xt) && _prep_valid(typeof(pt), pt) +end + +@testset "foreign θ eltype routes through the fallback (no PreparationMismatchError)" begin + # A non-dual off-construction eltype (Float32, BigFloat) used to hit the Float64 prep and + # throw; the type-match gate routes it to the prep-free fallback instead. + for (inplace, cons) in ((true, consp!), (false, consp)) + optf = inplace ? + OptimizationFunction(objp, ADTypes.AutoForwardDiff(); cons = cons) : + OptimizationFunction{false}(objp, ADTypes.AutoForwardDiff(); cons = cons) + optprob = OptimizationBase.instantiate_function( + optf, x0, ADTypes.AutoForwardDiff(), p0, 1; g = true, cons_j = true + ) + + for T in (Float32, BigFloat) + xT = T.(xt) + gref = ∇xf(xt, p0) # reference in Float64; compare at loose tol + if inplace + gT = zeros(T, 2) + @test_nowarn optprob.grad(gT, xT) + @test Float64.(gT) ≈ gref rtol = 1.0e-3 + JT = zeros(T, 2) + @test_nowarn optprob.cons_j(JT, xT) + @test Float64.(JT) ≈ vec(consjac(xt, p0)) rtol = 1.0e-3 + else + @test Float64.(optprob.grad(xT)) ≈ gref rtol = 1.0e-3 + @test Float64.(optprob.cons_j(xT)) ≈ vec(consjac(xt, p0)) rtol = 1.0e-3 + end + end + end end @testset "parametrized cons_j (Enzyme)" begin From a1e8fc13c9a82d19a605d9586d63ecdf0d4fc742 Mon Sep 17 00:00:00 2001 From: jClugstor Date: Wed, 8 Jul 2026 11:36:57 -0400 Subject: [PATCH 10/10] clean up comment --- lib/OptimizationBase/src/OptimizationDIExt.jl | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index c762d2c65..d3c41157f 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -17,10 +17,9 @@ using OptimizationBase.FastClosures # A DI preparation is built for the exact construction types (`x` and `Constant(p)`) and only # works for those — otherwise DI throws `PreparationMismatchError`. Reuse it while `θ`/`p` keep -# those types; fall back to a prep-free call otherwise. We compare types exactly rather than just -# looking for duals, so any off-type is caught (a dual `p` from a sensitivity layer, but equally a -# `Float32`/tracked/`BigFloat` `θ` that a prior `anyeltypedual` gate let slip through and error). -# `T` is a constant, so on the solve path this folds away. +# those types (e.g. a dual `p` from a sensitivity layer, or a `Float32`/`BigFloat` `θ`, does not), +# and fall back to a prep-free call otherwise. `T` is a constant, so on the solve path this folds +# away. @inline _prep_valid(::Type{T}, v) where {T} = typeof(v) === T # Output-buffer eltype for the `p`-accepting constraint wrapper: the type `f.cons` produces,