From 1cc6c34fb68d9e2ece1869d38eed375b630fa0c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 21 May 2026 15:50:49 +0200 Subject: [PATCH 1/2] Refactor transformer to use _eval in tests --- test/JuMP.jl | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/test/JuMP.jl b/test/JuMP.jl index 251323a..c780271 100644 --- a/test/JuMP.jl +++ b/test/JuMP.jl @@ -166,7 +166,13 @@ function test_parse_moi() return end -function _eval(model::JuMP.GenericModel{T}, func, x) where {T} +function _eval( + model::JuMP.GenericModel{T}, + func, + x; + x_grad = T.(collect(1:length(x))), + check_alloc::Bool = true, +) where {T} mode = ArrayDiff.Mode{Vector{T}}() ad = ArrayDiff.model(mode) MOI.Nonlinear.set_objective(ad, JuMP.moi_function(func)) @@ -178,20 +184,19 @@ function _eval(model::JuMP.GenericModel{T}, func, x) where {T} MOI.initialize(evaluator, [:Grad]) sizes = evaluator.backend.objective.expr.sizes val = MOI.eval_objective(evaluator, x) - if VERSION >= v"1.12" + if check_alloc && VERSION >= v"1.12" fill!(evaluator.backend.last_x, NaN) @test 0 == @allocated MOI.eval_objective(evaluator, x) end - x_grad = T.(collect(1:8)) g = zero(x) MOI.eval_objective_gradient(evaluator, g, x_grad) - if VERSION >= v"1.12" + if check_alloc && VERSION >= v"1.12" fill!(evaluator.backend.last_x, NaN) @test 0 == @allocated MOI.eval_objective_gradient(evaluator, g, x_grad) end MOI.Nonlinear.set_objective(ad, nothing) @test isnothing(ad.objective) - return sizes, val, g + return sizes, val, g, evaluator end function _test_neural( @@ -646,21 +651,11 @@ function _check_transformer_loss(build_loss; seq = 2, d_emb = 2) model = Model() @variable(model, x[1:seq, 1:d_emb], container = ArrayDiff.ArrayOfVariables) loss = build_loss(x) - mode = ArrayDiff.Mode() - ad = ArrayDiff.model(mode) - MOI.Nonlinear.set_objective(ad, JuMP.moi_function(loss)) - evaluator = MOI.Nonlinear.Evaluator( - ad, - mode, - JuMP.index.(JuMP.all_variables(model)), - ) - MOI.initialize(evaluator, [:Grad]) nvar = JuMP.num_variables(model) x_pt = randn(nvar) - val = MOI.eval_objective(evaluator, x_pt) + _, val, g, evaluator = + _eval(model, loss, x_pt; x_grad = x_pt, check_alloc = false) @test isfinite(val) - g = zeros(nvar) - MOI.eval_objective_gradient(evaluator, g, x_pt) @test all(isfinite, g) h = 1e-6 g_fd = zeros(nvar) From c5e0da57f7fd5132638502fa00cff0ca2b209413 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Thu, 21 May 2026 15:56:39 +0200 Subject: [PATCH 2/2] enable alloc --- test/JuMP.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/test/JuMP.jl b/test/JuMP.jl index c780271..4d266ec 100644 --- a/test/JuMP.jl +++ b/test/JuMP.jl @@ -171,7 +171,6 @@ function _eval( func, x; x_grad = T.(collect(1:length(x))), - check_alloc::Bool = true, ) where {T} mode = ArrayDiff.Mode{Vector{T}}() ad = ArrayDiff.model(mode) @@ -184,13 +183,13 @@ function _eval( MOI.initialize(evaluator, [:Grad]) sizes = evaluator.backend.objective.expr.sizes val = MOI.eval_objective(evaluator, x) - if check_alloc && VERSION >= v"1.12" + if VERSION >= v"1.12" fill!(evaluator.backend.last_x, NaN) @test 0 == @allocated MOI.eval_objective(evaluator, x) end g = zero(x) MOI.eval_objective_gradient(evaluator, g, x_grad) - if check_alloc && VERSION >= v"1.12" + if VERSION >= v"1.12" fill!(evaluator.backend.last_x, NaN) @test 0 == @allocated MOI.eval_objective_gradient(evaluator, g, x_grad) end @@ -653,8 +652,7 @@ function _check_transformer_loss(build_loss; seq = 2, d_emb = 2) loss = build_loss(x) nvar = JuMP.num_variables(model) x_pt = randn(nvar) - _, val, g, evaluator = - _eval(model, loss, x_pt; x_grad = x_pt, check_alloc = false) + _, val, g, evaluator = _eval(model, loss, x_pt; x_grad = x_pt) @test isfinite(val) @test all(isfinite, g) h = 1e-6