diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl index 463d5e173..cd611d5cd 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl @@ -6,7 +6,10 @@ 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}, Sr<:Vector{<:NTuple}, Rf<:Vector{<:NTuple}, @@ -19,8 +22,10 @@ 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 batched_seeds_reverse::Sr batched_results_forward::Rf @@ -111,12 +116,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 +153,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 +161,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 +173,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 +202,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 +221,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 +229,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