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 diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index ab3be02e9..d3c41157f 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 +# 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 (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, +# 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. (`_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 + Tp = eltype(p) + return Tp <: Number ? Base.promote_op(+, Tu, Tp) : Tu +end + function instantiate_function( f::OptimizationFunction{true}, x, ::ADTypes.AutoSparse{<:ADTypes.AutoSymbolics}, args...; kwargs... @@ -39,16 +63,33 @@ function instantiate_function( ) adtype, soadtype = generate_adtype(adtype) - # Create gradient closures with proper type stability using let blocks + # 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. 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 - (res, θ, p = p) -> gradient!(f.f, res, _prep_grad, adtype, θ, Constant(p)) + let _prep_grad = _prep_grad, f = f, adtype = adtype, Tx0 = Tx0, Tp0 = Tp0 + function (res, θ, p = 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)) + 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, Tx0 = Tx0, Tp0 = Tp0 + function (res, θ, p = 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)) + end + end end end elseif g == true @@ -194,10 +235,29 @@ 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. `_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) + res = Vector{_cons_out_eltype(x, p)}(undef, num_cons) + f.cons(res, x, p) + return res + 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, Tx0 = Tx0, Tp0 = Tp0 + function (J, θ, p = 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)) + end return if size(J, 1) == 1 J = vec(J) end @@ -205,7 +265,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 @@ -428,16 +488,29 @@ function instantiate_function( ) adtype, soadtype = generate_adtype(adtype) - # Create gradient closures with proper type stability using let blocks + # 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. 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 - (θ, p = p) -> gradient(f.f, _prep_grad, adtype, θ, Constant(p)) + let _prep_grad = _prep_grad, f = f, adtype = adtype, Tx0 = Tx0, Tp0 = Tp0 + function (θ, p = 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 - (θ, p = p) -> gradient(f.f, _prep_grad, adtype, θ, Constant(p)) + let _prep_grad = _prep_grad, f = f, adtype = adtype, p = p, Tx0 = Tx0, Tp0 = Tp0 + function (θ, p = 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 end elseif g == true @@ -560,10 +633,16 @@ 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 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 - function (θ) - J = jacobian(f.cons, _prep_jac, adtype, θ, Constant(p)) + let f = f, _prep_jac = _prep_jac, adtype = adtype, p = p, Tx0 = Tx0, Tp0 = Tp0 + function (θ, p = 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 J = vec(J) end @@ -572,7 +651,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 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..6bda15ecc --- /dev/null +++ b/lib/OptimizationBase/test/dual_tolerant_tests.jl @@ -0,0 +1,184 @@ +# 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. `_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: _prep_valid + +# 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 "_prep_valid type-match gating" begin + d = ForwardDiff.Dual{Nothing}(1.0, 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. +# `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 + 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 + ) + 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 "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`, 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 + # 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