From afd083868e0c0ffbd98dcdb977d3db161255c809 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 16 Oct 2026 12:10:45 +0200 Subject: [PATCH 1/2] Fix allocation-free Hessian recovery with star coloring --- src/coloring.jl | 34 ++++++++++++++++- src/forward_over_reverse.jl | 3 +- test/test_ReverseAD.jl | 73 +++++++++++++++++++++++++++++++------ 3 files changed, 95 insertions(+), 15 deletions(-) diff --git a/src/coloring.jl b/src/coloring.jl index 10a1dc3..3914ee0 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -31,7 +31,8 @@ diagonal. Diagonal entries absent from `edgelist` are not added to the pattern. Returns `(colptr, I, J, result)` where `colptr` contains the local lower- triangular CSC column pointers and `I`, `J` give the corresponding entries -in global variable indices. `result` wraps an SMC coloring result. +in global variable indices. `result` wraps an SMC coloring result, or is +`nothing` if the pattern is empty. """ function _hessian_color_preprocess( edgelist, @@ -49,6 +50,10 @@ function _hessian_color_preprocess( end local_indices = sort!(collect(seen_idx)) empty!(seen_idx) + # SMC's star coloring does not support an empty pattern. + if isempty(local_indices) + return [1], Int[], Int[], nothing + end global_to_local_idx = seen_idx.nzidx # steal for storage for k in eachindex(local_indices) @@ -86,6 +91,10 @@ function _seed_matrix(result::ColoringResult) return Matrix{Float64}(undef, n, ncolors) end +function _seed_matrix(::Nothing) + return Matrix{Float64}(undef, 0, 0) +end + """ _prepare_seed_matrix!(R, result::ColoringResult) @@ -125,3 +134,26 @@ function _recover_from_matmat!( SMC.decompress_csc!(V, colptr, R, result.result, :L) return end + +function _recover_from_matmat!( + colptr::AbstractVector, + V::AbstractVector{T}, + R::AbstractMatrix{T}, + result::ColoringResult{<:SMC.StarSetColoringResult}, +) where {T} + # SMC.decompress_csc! only supports substitution decompression. Mirror + # SMC.decompress! for a lower-triangular CSC matrix without constructing it. + S = SMC.sparsity_pattern(result.result) + rows = SparseArrays.rowvals(S) + compressed_indices = result.result.compressed_indices + l = 0 + for j in axes(S, 2) + for k in SparseArrays.nzrange(S, j) + if rows[k] >= j + l += 1 + V[l] = R[compressed_indices[k]] + end + end + end + return +end diff --git a/src/forward_over_reverse.jl b/src/forward_over_reverse.jl index c103ebd..cc4328b 100644 --- a/src/forward_over_reverse.jl +++ b/src/forward_over_reverse.jl @@ -38,8 +38,7 @@ function _eval_hessian( scale::Float64, nzcount::Int, )::Int - if ex.expr.linearity == LINEAR - @assert length(ex.hess_I) == 0 + if isempty(ex.hess_I) return 0 end chunk = min(size(ex.seed_matrix, 2), d.max_chunk) diff --git a/test/test_ReverseAD.jl b/test/test_ReverseAD.jl index 99bbc17..345fc03 100644 --- a/test/test_ReverseAD.jl +++ b/test/test_ReverseAD.jl @@ -450,17 +450,32 @@ end function test_hessian_structurally_zero_diagonal() variables = MOI.VariableIndex.(1:4) x, y = variables[2], variables[4] - for (expr, structure) in - ((:($x * $y), [(4, 2)]), (:($x * $y + $y^3), [(4, 2), (4, 4)])) + for decompression in (:direct, :substitution), + (expr, structure) in ( + (:($x * $y), [(4, 2)]), + (:($x * $y + $y^3), [(4, 2), (4, 4)]), + (:($x + $y), Tuple{Int,Int}[]), + (:(3.0), Tuple{Int,Int}[]), + (:(ifelse($x < 0, $x, $y)), Tuple{Int,Int}[]), + ) + model = ArrayDiff.Model() ArrayDiff.set_objective(model, expr) - evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(), variables) + algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(; decompression) + evaluator = + ArrayDiff.Evaluator(model, ArrayDiff.Mode(algorithm), variables) MOI.initialize(evaluator, [:Hess]) @test MOI.hessian_objective_structure(evaluator) == structure @test MOI.hessian_lagrangian_structure(evaluator) == structure for y_value in (0.0, 2.0) values = [0.0, 3.0, 0.0, y_value] - expected = length(structure) == 1 ? [1.0] : [1.0, 6 * y_value] + expected = if isempty(structure) + Float64[] + elseif length(structure) == 1 + [1.0] + else + [1.0, 6 * y_value] + end H = fill(NaN, length(structure)) MOI.eval_hessian_objective(evaluator, H, values) @test H == expected @@ -472,11 +487,11 @@ function test_hessian_structurally_zero_diagonal() end function test_coloring_lower_triangular_recovery() - for order in (ArrayDiff.SMC.NaturalOrder(), ArrayDiff.SMC.LargestFirst()) - algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm( - order; - decompression = :substitution, - ) + for (T, decompression) in + ((Float32, :direct), (Float64, :direct), (Float64, :substitution)), + order in (ArrayDiff.SMC.NaturalOrder(), ArrayDiff.SMC.LargestFirst()) + + algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(order; decompression) for edges in ( Set([(2, 4), (4, 7), (7, 9)]), Set([(4, 2), (2, 7), (9, 2), (7, 7)]), @@ -485,7 +500,7 @@ function test_coloring_lower_triangular_recovery() Set([(2, 2), (9, 9)]), Set{Tuple{Int,Int}}(), ) - H = zeros(9, 9) + H = zeros(T, 9, 9) for (i, j) in edges H[i, j] = H[j, i] = 10 * min(i, j) + max(i, j) end @@ -496,11 +511,17 @@ function test_coloring_lower_triangular_recovery() by = reverse, ) @test collect(zip(I, J)) == expected_structure - R = ArrayDiff._seed_matrix(result) + R = Matrix{T}(ArrayDiff._seed_matrix(result)) + if isempty(expected_structure) + @test result === nothing + @test colptr == [1] + @test size(R) == (0, 0) + continue + end ArrayDiff._prepare_seed_matrix!(R, result) local_indices = result.local_indices compressed = H[local_indices, local_indices] * R - V = fill(NaN, length(I)) + V = fill(T(NaN), length(I)) ArrayDiff._recover_from_matmat!(colptr, V, compressed, result) @test V == [H[i, j] for (i, j) in expected_structure] @test 0 == @allocated ArrayDiff._recover_from_matmat!( @@ -514,6 +535,34 @@ function test_coloring_lower_triangular_recovery() return end +function test_hessian_star_coloring() + x, y = MOI.VariableIndex.(1:2) + model = ArrayDiff.Model() + ArrayDiff.set_objective(model, :((1 - $x)^2 + 100 * ($y - $x^2)^2)) + ArrayDiff.add_constraint(model, :($x * $y), MOI.LessThan(2.0)) + ArrayDiff.add_constraint(model, :($y^2), MOI.LessThan(2.0)) + algorithm = ArrayDiff.SMC.GreedyColoringAlgorithm(; decompression = :direct) + evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(algorithm), [x, y]) + MOI.initialize(evaluator, [:Hess]) + @test MOI.hessian_lagrangian_structure(evaluator) == + [(1, 1), (2, 1), (2, 2), (2, 1), (2, 2)] + H = zeros(5) + μ = [2.0, 3.0] + for xv in ([1.0, 1.0], [-1.2, 1.0]) + MOI.eval_hessian_lagrangian(evaluator, H, xv, 1.5, μ) + @test H ≈ [ + 1.5 * (2 + 1200 * xv[1]^2 - 400 * xv[2]), + 1.5 * (-400 * xv[1]), + 300.0, + 2.0, + 6.0, + ] + @test 0 == + @allocated MOI.eval_hessian_lagrangian(evaluator, H, xv, 1.5, μ) + end + return +end + function test_derivatives() a = MOI.VariableIndex(1) b = MOI.VariableIndex(2) From 7b725614059993bcbe4ab4b9daa0e8d4f808853f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 16 Oct 2026 12:26:43 +0200 Subject: [PATCH 2/2] Add TODO --- src/coloring.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/coloring.jl b/src/coloring.jl index 3914ee0..605e5d6 100644 --- a/src/coloring.jl +++ b/src/coloring.jl @@ -51,6 +51,7 @@ function _hessian_color_preprocess( local_indices = sort!(collect(seen_idx)) empty!(seen_idx) # SMC's star coloring does not support an empty pattern. + # TODO: remove when https://github.com/JuliaDiff/SparseMatrixColorings.jl/pull/352 lands if isempty(local_indices) return [1], Int[], Int[], nothing end