From 764af8229f875a95a7c0f7668db8e10871ff5985 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?= Date: Thu, 16 Oct 2025 04:44:38 +0300 Subject: [PATCH 1/4] fix method ambiguities for instantiate_function in the Zygote extension Co-authored-by: Claude --- .../ext/OptimizationZygoteExt.jl | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/lib/OptimizationBase/ext/OptimizationZygoteExt.jl b/lib/OptimizationBase/ext/OptimizationZygoteExt.jl index 461efe699..f2bf61a8e 100644 --- a/lib/OptimizationBase/ext/OptimizationZygoteExt.jl +++ b/lib/OptimizationBase/ext/OptimizationZygoteExt.jl @@ -288,6 +288,18 @@ function OptimizationBase.instantiate_function( f, x, adtype, p, num_cons; kwargs...) end +function OptimizationBase.instantiate_function( + f::OptimizationFunction{true}, cache::OptimizationBase.ReInitCache, + adtype::DifferentiationInterface.SecondOrder{ + <:ADTypes.AbstractADType, <:ADTypes.AutoZygote}, + num_cons = 0; kwargs...) + x = cache.u0 + p = cache.p + + return OptimizationBase.instantiate_function( + f, x, adtype, p, num_cons; kwargs...) +end + function OptimizationBase.instantiate_function( f::OptimizationFunction{true}, x, adtype::ADTypes.AutoSparse{<:Union{ADTypes.AutoZygote, @@ -575,4 +587,15 @@ function OptimizationBase.instantiate_function( return OptimizationBase.instantiate_function(f, x, adtype, p, num_cons; kwargs...) end +function OptimizationBase.instantiate_function( + f::OptimizationFunction{true}, cache::OptimizationBase.ReInitCache, + adtype::ADTypes.AutoSparse{<:DifferentiationInterface.SecondOrder{ + <:ADTypes.AbstractADType, <:ADTypes.AutoZygote}}, + num_cons = 0; kwargs...) + x = cache.u0 + p = cache.p + + return OptimizationBase.instantiate_function(f, x, adtype, p, num_cons; kwargs...) +end + end From 654f65fbb37e77a7f6f96fdb5faa99a1b32f3192 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?= Date: Thu, 16 Oct 2025 04:46:26 +0300 Subject: [PATCH 2/4] Fix Lagrangian Hessian prototype dimensions in OptimizationZygoteExt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Lagrangian Hessian prototype was incorrectly sized as (num_constraints × num_variables) instead of (num_variables × num_variables). This caused a `BoundsError` when computing the Lagrangian Hessian with more variables than constraints, as the prototype was used as a buffer for the n×n Hessian matrix. Changes: - Fix lag_hess_prototype initialization to use zeros(Bool, length(x), length(x)) - Add comprehensive tests to verify prototype dimensions and usability as a buffer This ensures the Lagrangian Hessian is always correctly sized as n×n regardless of the number of constraints, matching the mathematical definition of the Lagrangian Hessian as the second derivative with respect to the decision variables. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- .../ext/OptimizationZygoteExt.jl | 2 +- lib/OptimizationBase/test/adtests.jl | 22 +++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/lib/OptimizationBase/ext/OptimizationZygoteExt.jl b/lib/OptimizationBase/ext/OptimizationZygoteExt.jl index f2bf61a8e..1e26909af 100644 --- a/lib/OptimizationBase/ext/OptimizationZygoteExt.jl +++ b/lib/OptimizationBase/ext/OptimizationZygoteExt.jl @@ -208,7 +208,7 @@ function OptimizationBase.instantiate_function( lag_extras = prepare_hessian( lagrangian, soadtype, x, Constant(one(eltype(x))), Constant(ones(eltype(x), num_cons)), Constant(p), strict = Val(false)) - lag_hess_prototype = zeros(Bool, num_cons, length(x)) + lag_hess_prototype = zeros(Bool, length(x), length(x)) function lag_h!(H::AbstractMatrix, θ, σ, λ) if σ == zero(eltype(θ)) diff --git a/lib/OptimizationBase/test/adtests.jl b/lib/OptimizationBase/test/adtests.jl index 6fe4eea05..f1c12dd9b 100644 --- a/lib/OptimizationBase/test/adtests.jl +++ b/lib/OptimizationBase/test/adtests.jl @@ -257,6 +257,17 @@ optprob.cons_h(H3, x0) optprob.lag_h(H4, x0, σ, μ) @test H4≈σ * H2 + μ[1] * H3[1] rtol=1e-6 + # Test that the AD-generated lag_hess_prototype has correct dimensions + @test !isnothing(optprob.lag_hess_prototype) + @test size(optprob.lag_hess_prototype) == (length(x0), length(x0)) # Should be n×n, not num_cons×n + + # Test that we can actually use it as a buffer (this would fail with the bug) + if !isnothing(optprob.lag_hess_prototype) + H_proto = similar(optprob.lag_hess_prototype, Float64) + optprob.lag_h(H_proto, x0, σ, μ) + @test H_proto ≈ σ * H2 + μ[1] * H3[1] rtol=1e-6 + end + G2 = Array{Float64}(undef, 2) H2 = Array{Float64}(undef, 2, 2) @@ -490,6 +501,17 @@ end optprob.lag_h(H4, x0, σ, μ) @test H4≈σ * H1 + sum(μ .* H3) rtol=1e-6 + # Test that the AD-generated lag_hess_prototype has correct dimensions + @test !isnothing(optprob.lag_hess_prototype) + @test size(optprob.lag_hess_prototype) == (length(x0), length(x0)) # Should be n×n, not num_cons×n + + # Test that we can actually use it as a buffer (this would fail with the bug) + if !isnothing(optprob.lag_hess_prototype) + H_proto = similar(optprob.lag_hess_prototype, Float64) + optprob.lag_h(H_proto, x0, σ, μ) + @test H_proto ≈ σ * H1 + sum(μ .* H3) rtol=1e-6 + end + G2 = Array{Float64}(undef, 2) H2 = Array{Float64}(undef, 2, 2) From 00925e30ba096a0e023d83f35957376d0c1c1db4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?= Date: Thu, 16 Oct 2025 04:48:31 +0300 Subject: [PATCH 3/4] Fix Lagrangian Hessian prototype dimensions in OptimizationDIExt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Lagrangian Hessian prototype was incorrectly sized as (num_constraints × num_variables) instead of (num_variables × num_variables). This caused a `BoundsError` when computing the Lagrangian Hessian with more variables than constraints, as the prototype was used as a buffer for the n×n Hessian matrix. Changes: - Fix lag_hess_prototype initialization to use zeros(Bool, length(x), length(x)) in both in-place and out-of-place function instantiation - Add tests to verify prototype dimensions This ensures the Lagrangian Hessian is always correctly sized as n×n regardless of the number of constraints, matching the mathematical definition of the Lagrangian Hessian as the second derivative with respect to the decision variables. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- lib/OptimizationBase/src/OptimizationDIExt.jl | 4 ++-- lib/OptimizationBase/test/adtests.jl | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/lib/OptimizationBase/src/OptimizationDIExt.jl b/lib/OptimizationBase/src/OptimizationDIExt.jl index adb84b55f..017506a03 100644 --- a/lib/OptimizationBase/src/OptimizationDIExt.jl +++ b/lib/OptimizationBase/src/OptimizationDIExt.jl @@ -198,7 +198,7 @@ function instantiate_function( lag_prep = prepare_hessian( lagrangian, soadtype, x, Constant(one(eltype(x))), Constant(ones(eltype(x), num_cons)), Constant(p)) - lag_hess_prototype = zeros(Bool, num_cons, length(x)) + lag_hess_prototype = zeros(Bool, length(x), length(x)) function lag_h!(H::AbstractMatrix, θ, σ, λ) if σ == zero(eltype(θ)) @@ -457,7 +457,7 @@ function instantiate_function( lag_prep = prepare_hessian( lagrangian, soadtype, x, Constant(one(eltype(x))), Constant(ones(eltype(x), num_cons)), Constant(p)) - lag_hess_prototype = zeros(Bool, num_cons, length(x)) + lag_hess_prototype = zeros(Bool, length(x), length(x)) function lag_h!(θ, σ, λ) if σ == zero(eltype(θ)) diff --git a/lib/OptimizationBase/test/adtests.jl b/lib/OptimizationBase/test/adtests.jl index f1c12dd9b..1a9483deb 100644 --- a/lib/OptimizationBase/test/adtests.jl +++ b/lib/OptimizationBase/test/adtests.jl @@ -144,6 +144,17 @@ optprob.cons_h(H3, x0) optprob.lag_h(H4, x0, σ, μ) @test H4≈σ * H2 + μ[1] * H3[1] rtol=1e-6 + # Test that the AD-generated lag_hess_prototype has correct dimensions + @test !isnothing(optprob.lag_hess_prototype) + @test size(optprob.lag_hess_prototype) == (length(x0), length(x0)) # Should be n×n, not num_cons×n + + # Test that we can actually use it as a buffer + if !isnothing(optprob.lag_hess_prototype) + H_proto = similar(optprob.lag_hess_prototype, Float64) + optprob.lag_h(H_proto, x0, σ, μ) + @test H_proto ≈ σ * H2 + μ[1] * H3[1] rtol=1e-6 + end + G2 = Array{Float64}(undef, 2) H2 = Array{Float64}(undef, 2, 2) From 19a3806143449ab9d7bda7a356c154245309e549 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?= Date: Thu, 16 Oct 2025 04:52:07 +0300 Subject: [PATCH 4/4] Fix typo: allowsconsjvp -> allowsconsvjp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Corrected function name to properly spell VJP (Vector-Jacobian Product). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- lib/OptimizationBase/src/cache.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OptimizationBase/src/cache.jl b/lib/OptimizationBase/src/cache.jl index 375d5aaa1..9454e4f03 100644 --- a/lib/OptimizationBase/src/cache.jl +++ b/lib/OptimizationBase/src/cache.jl @@ -62,7 +62,7 @@ function OptimizationCache(prob::SciMLBase.OptimizationProblem, opt; g = SciMLBase.requiresgradient(opt), h = SciMLBase.requireshessian(opt), hv = SciMLBase.requireshessian(opt), fg = SciMLBase.allowsfg(opt), fgh = SciMLBase.allowsfgh(opt), cons_j = SciMLBase.requiresconsjac(opt), cons_h = SciMLBase.requiresconshess(opt), - cons_vjp = SciMLBase.allowsconsjvp(opt), cons_jvp = SciMLBase.allowsconsjvp(opt), lag_h = SciMLBase.requireslagh(opt)) + cons_vjp = SciMLBase.allowsconsvjp(opt), cons_jvp = SciMLBase.allowsconsjvp(opt), lag_h = SciMLBase.requireslagh(opt)) if structural_analysis obj_res, cons_res = symify_cache(f, prob, num_cons, manifold)