Skip to content

Commit

Permalink
Last fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
@gdalle
gdalle committed Apr 22, 2024
1 parent 6c5811f commit 1365d5e
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 17 deletions.
1 change: 1 addition & 0 deletions ext/SparseDiffToolsPolyesterForwardDiffExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ using ADTypes, SparseDiffTools, PolyesterForwardDiff, UnPack, Random, SparseArra
import ForwardDiff
import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityDetection,
ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache,
sparse_jacobian_cache_aux,
sparse_jacobian!,
sparse_jacobian_static_array, __standard_tag, __chunksize,
polyesterforwarddiff_color_jacobian
Expand Down
4 changes: 2 additions & 2 deletions src/highlevel/common.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -273,8 +273,8 @@ const __init_𝒥 = init_jacobian

# Misc Functions
function __chunksize(
::Union{AutoSparse{<:AutoForwardDiff}{C}, AutoForwardDiff{C},
AutoSparse{<:AutoPolyesterForwardDiff}{C}, AutoPolyesterForwardDiff{C}},
::Union{AutoSparse{<:AutoForwardDiff{C}}, AutoForwardDiff{C},
AutoSparse{<:AutoPolyesterForwardDiff{C}}, AutoPolyesterForwardDiff{C}},
x) where {C}
C isa ForwardDiff.Chunk && return C
return __chunksize(Val(C), x)
Expand Down
8 changes: 4 additions & 4 deletions src/highlevel/forward_or_reverse_mode.jl
View file
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
function sparse_jacobian_cache_aux(::ForwardOrReverseMode, ad::AbstractADType,

Check warning on line 1 in src/highlevel/forward_or_reverse_mode.jl

View check run for this annotation

Codecov / codecov/patch

src/highlevel/forward_or_reverse_mode.jl#L1 

Added line #L1 was not covered by tests
sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F}
if ad isa AutoEnzyme
if ad isa Union{AutoEnzyme, AutoSparse{<:AutoEnzyme}}
return sparse_jacobian_cache_aux(ReverseMode(), ad, sd, f, x; fx)
elseif ad isa AutoDiffractor
elseif ad isa Union{AutoDiffractor, AutoSparse{<:AutoDiffractor}}
return sparse_jacobian_cache_aux(ForwardMode(), ad, sd, f, x; fx)

Check warning on line 6 in src/highlevel/forward_or_reverse_mode.jl

View check run for this annotation

Codecov / codecov/patch

src/highlevel/forward_or_reverse_mode.jl#L3-L6 

Added lines #L3 - L6 were not covered by tests
else
error("Unknown mixed mode AD")

Check warning on line 8 in src/highlevel/forward_or_reverse_mode.jl

View check run for this annotation

Codecov / codecov/patch

src/highlevel/forward_or_reverse_mode.jl#L8 

Added line #L8 was not covered by tests
Expand All @@ -11,9 +11,9 @@ end

function sparse_jacobian_cache_aux(::ForwardOrReverseMode, ad::AbstractADType,

Check warning on line 12 in src/highlevel/forward_or_reverse_mode.jl

View check run for this annotation

Codecov / codecov/patch

src/highlevel/forward_or_reverse_mode.jl#L12 

Added line #L12 was not covered by tests
sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F}
if ad isa AutoEnzyme
if ad isa Union{AutoEnzyme, AutoSparse{<:AutoEnzyme}}
return sparse_jacobian_cache_aux(ReverseMode(), ad, sd, f!, fx, x)
elseif ad isa AutoDiffractor
elseif ad isa Union{AutoDiffractor, AutoSparse{<:AutoDiffractor}}
return sparse_jacobian_cache_aux(ForwardMode(), ad, sd, f!, fx, x)

Check warning on line 17 in src/highlevel/forward_or_reverse_mode.jl

View check run for this annotation

Codecov / codecov/patch

src/highlevel/forward_or_reverse_mode.jl#L14-L17 

Added lines #L14 - L17 were not covered by tests
else
error("Unknown mixed mode AD")

Check warning on line 19 in src/highlevel/forward_or_reverse_mode.jl

View check run for this annotation

Codecov / codecov/patch

src/highlevel/forward_or_reverse_mode.jl#L19 

Added line #L19 was not covered by tests
Expand Down
30 changes: 19 additions & 11 deletions test/test_sparse_jacobian.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,14 @@ using ADTypes, SparseDiffTools,
StaticArrays
using ADTypes: dense_ad

function nice_string(ad::AbstractADType)
if ad isa AutoSparse
return "AutoSparse($(nice_string(dense_ad(ad))))"
else
return nameof(typeof(ad))
end
end

function __chunksize(::Union{
AutoSparse{<:AutoForwardDiff{C}}, AutoForwardDiff{C},
AutoSparse{<:AutoPolyesterForwardDiff{C}}, AutoPolyesterForwardDiff{C}
Expand Down Expand Up @@ -70,7 +78,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
AutoPolyesterForwardDiff(; chunksize = 4)
]

@testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in DIFFTYPES
@testset "sparse_jacobian $(nice_string(difftype)): Out of Place" for difftype in DIFFTYPES
@testset "Cache & Reuse" begin
cache = sparse_jacobian_cache(difftype, sd, fdiff, x)
J = init_jacobian(cache)
Expand All @@ -81,7 +89,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
@inferred sparse_jacobian!(J, difftype, cache, fdiff, x)

t₁ = @elapsed sparse_jacobian!(J, difftype, cache, fdiff, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (only differentiation) time: $(t₁)s"
@info "$(nice_string(difftype))() `sparse_jacobian!` (only differentiation) time: $(t₁)s"

J = sparse_jacobian(difftype, cache, fdiff, x)

Expand All @@ -92,7 +100,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
end

t₂ = @elapsed sparse_jacobian(difftype, cache, fdiff, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian` (with matrix allocation) time: $(t₂)s"
@info "$(nice_string(difftype))() `sparse_jacobian` (with matrix allocation) time: $(t₂)s"
end

@testset "Single Use" begin
Expand All @@ -104,7 +112,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
end

t₁ = @elapsed sparse_jacobian(difftype, sd, fdiff, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian` (complete) time: $(t₁)s"
@info "$(nice_string(difftype))() `sparse_jacobian` (complete) time: $(t₁)s"

cache = sparse_jacobian_cache(difftype, sd, fdiff, x)
J = init_jacobian(cache)
Expand All @@ -115,13 +123,13 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
@inferred sparse_jacobian!(J, difftype, sd, fdiff, x)

t₂ = @elapsed sparse_jacobian!(J, difftype, sd, fdiff, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s"
@info "$(nice_string(difftype))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s"
end
end

@info "Inplace Place Function"

@testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in (
@testset "sparse_jacobian $(nice_string(difftype)): In place" for difftype in (
AutoSparse(AutoForwardDiff()), AutoForwardDiff(),
AutoSparse(AutoForwardDiff(; chunksize = 0)), AutoForwardDiff(; chunksize = 0),
AutoSparse(AutoForwardDiff(; chunksize = 4)), AutoForwardDiff(; chunksize = 4),
Expand All @@ -138,7 +146,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
@inferred sparse_jacobian!(J, difftype, cache, fdiff, y, x)

t₁ = @elapsed sparse_jacobian!(J, difftype, cache, fdiff, y, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (only differentiation) time: $(t₁)s"
@info "$(nice_string(difftype))() `sparse_jacobian!` (only differentiation) time: $(t₁)s"

J = sparse_jacobian(difftype, cache, fdiff, y, x)

Expand All @@ -148,7 +156,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
end

t₂ = @elapsed sparse_jacobian(difftype, cache, fdiff, y, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian` (with jacobian allocation) time: $(t₂)s"
@info "$(nice_string(difftype))() `sparse_jacobian` (with jacobian allocation) time: $(t₂)s"
end

@testset "Single Use" begin
Expand All @@ -160,7 +168,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
end

t₁ = @elapsed sparse_jacobian(difftype, sd, fdiff, y, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian` (complete) time: $(t₁)s"
@info "$(nice_string(difftype))() `sparse_jacobian` (complete) time: $(t₁)s"

J = init_jacobian(cache)

Expand All @@ -170,11 +178,11 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
@inferred sparse_jacobian!(J, difftype, sd, fdiff, y, x)

t₂ = @elapsed sparse_jacobian!(J, difftype, sd, fdiff, y, x)
@info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s"
@info "$(nice_string(difftype))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s"
end
end

@testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in (
@testset "sparse_jacobian $(nice_string(difftype)): In place" for difftype in (
AutoSparse(AutoZygote()),
AutoZygote())
y = similar(x)
Expand Down

0 comments on commit 1365d5e

Please sign in to comment.