From 1f2cc3afa070f459e8bf46662b34a1c0276286c8 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 8 Oct 2025 12:56:31 +0200 Subject: [PATCH 1/2] fix: handle empty row or column colors in mixed mode sparse Jacobian --- .../jacobian_mixed.jl | 32 +++++++++++++++---- .../test/Core/SimpleFiniteDiff/test.jl | 5 +++ DifferentiationInterface/test/testutils.jl | 2 +- 3 files changed, 32 insertions(+), 7 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl index 463d5e173..b17345fd9 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl @@ -7,6 +7,8 @@ struct SMCMixedModeSparseJacobianPrep{ P<:AbstractMatrix, C<:AbstractColoringResult{:nonsymmetric,:bidirectional}, M<:AbstractMatrix{<:Number}, + Sfp<:NTuple, + Srp<:NTuple, Sf<:Vector{<:NTuple}, Sr<:Vector{<:NTuple}, Rf<:Vector{<:NTuple}, @@ -21,6 +23,8 @@ struct SMCMixedModeSparseJacobianPrep{ coloring_result::C compressed_matrix_forward::M compressed_matrix_reverse::M + batched_seed_forward_prep::Sfp + batched_seed_reverse_prep::Srp batched_seeds_forward::Sf batched_seeds_reverse::Sr batched_results_forward::Rf @@ -111,12 +115,24 @@ function _prepare_mixed_sparse_jacobian_aux_aux( groups_forward = column_groups(coloring_result) groups_reverse = row_groups(coloring_result) + seed_forward_prep = DI.multibasis(x, eachindex(x)) + seed_reverse_prep = DI.multibasis(y, eachindex(y)) seeds_forward = [DI.multibasis(x, eachindex(x)[group]) for group in groups_forward] seeds_reverse = [DI.multibasis(y, eachindex(y)[group]) for group in groups_reverse] - compressed_matrix_forward = stack(_ -> vec(similar(y)), groups_forward; dims=2) - compressed_matrix_reverse = stack(_ -> vec(similar(x)), groups_reverse; dims=1) + compressed_matrix_forward = if isempty(groups_forward) + similar(vec(y), length(y), 0) + else + stack(_ -> vec(similar(y)), groups_forward; dims=2) + end + compressed_matrix_reverse = if isempty(groups_reverse) + similar(vec(x), 0, length(x)) + else + stack(_ -> vec(similar(x)), groups_reverse; dims=1) + end + batched_seed_forward_prep = ntuple(b -> copy(seed_forward_prep), Val(Bf)) + batched_seed_reverse_prep = ntuple(b -> copy(seed_reverse_prep), Val(Br)) batched_seeds_forward = [ ntuple(b -> seeds_forward[1 + ((a - 1) * Bf + (b - 1)) % Nf], Val(Bf)) for a in 1:Af ] @@ -136,7 +152,7 @@ function _prepare_mixed_sparse_jacobian_aux_aux( f_or_f!y..., DI.forward_backend(dense_backend), x, - batched_seeds_forward[1], + batched_seed_forward_prep, contexts...; ) pullback_prep = DI.prepare_pullback_nokwarg( @@ -144,7 +160,7 @@ function _prepare_mixed_sparse_jacobian_aux_aux( f_or_f!y..., DI.reverse_backend(dense_backend), x, - batched_seeds_reverse[1], + batched_seed_reverse_prep, contexts...; ) @@ -156,6 +172,8 @@ function _prepare_mixed_sparse_jacobian_aux_aux( coloring_result, compressed_matrix_forward, compressed_matrix_reverse, + batched_seed_forward_prep, + batched_seed_reverse_prep, batched_seeds_forward, batched_seeds_reverse, batched_results_forward, @@ -183,6 +201,8 @@ function _sparse_jacobian_aux!( coloring_result, compressed_matrix_forward, compressed_matrix_reverse, + batched_seed_forward_prep, + batched_seed_reverse_prep, batched_seeds_forward, batched_seeds_reverse, batched_results_forward, @@ -200,7 +220,7 @@ function _sparse_jacobian_aux!( pushforward_prep, DI.forward_backend(dense_backend), x, - batched_seeds_forward[1], + batched_seed_forward_prep, contexts..., ) pullback_prep_same = DI.prepare_pullback_same_point( @@ -208,7 +228,7 @@ function _sparse_jacobian_aux!( pullback_prep, DI.reverse_backend(dense_backend), x, - batched_seeds_reverse[1], + batched_seed_reverse_prep, contexts..., ) diff --git a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl index 573aae0af..6c5595475 100644 --- a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl +++ b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl @@ -128,6 +128,11 @@ end @test only(row_groups(jac_rev_prep)) == 1:10 @test only(column_groups(hess_prep)) == 1:10 end + + @testset "Empty colors for mixed mode" begin # issue 857 + backend = MyAutoSparse(MixedMode(adaptive_backends[1], adaptive_backends[2])) + @test jacobian(copyto!, zeros(10), backend, ones(10)) isa AbstractMatrix + end end @testset "Misc" begin diff --git a/DifferentiationInterface/test/testutils.jl b/DifferentiationInterface/test/testutils.jl index 657a1834d..8dec57bfa 100644 --- a/DifferentiationInterface/test/testutils.jl +++ b/DifferentiationInterface/test/testutils.jl @@ -17,7 +17,7 @@ function MyAutoSparse(backend::AbstractADType) return AutoSparse( backend; sparsity_detector=TracerSparsityDetector(), - coloring_algorithm=GreedyColoringAlgorithm(), + coloring_algorithm=GreedyColoringAlgorithm(; postprocessing=true), ) end From ad1c9c88ae7905fd7789b7d7d8761f9b631bf314 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 8 Oct 2025 19:43:16 +0200 Subject: [PATCH 2/2] Fix matrix types --- .../jacobian_mixed.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl index b17345fd9..cd611d5cd 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl @@ -6,7 +6,8 @@ struct SMCMixedModeSparseJacobianPrep{ BSr<:DI.BatchSizeSettings, P<:AbstractMatrix, C<:AbstractColoringResult{:nonsymmetric,:bidirectional}, - M<:AbstractMatrix{<:Number}, + Mf<:AbstractMatrix{<:Number}, + Mr<:AbstractMatrix{<:Number}, Sfp<:NTuple, Srp<:NTuple, Sf<:Vector{<:NTuple}, @@ -21,8 +22,8 @@ struct SMCMixedModeSparseJacobianPrep{ batch_size_settings_reverse::BSr sparsity::P coloring_result::C - compressed_matrix_forward::M - compressed_matrix_reverse::M + compressed_matrix_forward::Mf + compressed_matrix_reverse::Mr batched_seed_forward_prep::Sfp batched_seed_reverse_prep::Srp batched_seeds_forward::Sf