From bd2f0fec9738ed1ecb487fa8b41a2626925e85c7 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 28 Oct 2023 11:19:06 -0400 Subject: [PATCH] Add tests and specialize on functions --- .github/workflows/CI.yml | 3 +-- .github/workflows/Downstream.yml | 2 ++ Project.toml | 2 +- src/highlevel/coloring.jl | 7 ++++--- src/highlevel/common.jl | 17 ++++++++++++----- src/highlevel/finite_diff.jl | 10 ++++++---- src/highlevel/forward_mode.jl | 10 ++++++---- src/highlevel/reverse_mode.jl | 8 ++++---- test/test_sparse_jacobian.jl | 8 ++++++-- 9 files changed, 42 insertions(+), 25 deletions(-) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 5dadef0d..c77dcf45 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -17,8 +17,7 @@ jobs: - InterfaceI version: - '1' # Latest Release - - '~1.6' # Current LTS - - '1.8' + - '1.6' # Current LTS steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v1 diff --git a/.github/workflows/Downstream.yml b/.github/workflows/Downstream.yml index 5b1be058..63167e87 100644 --- a/.github/workflows/Downstream.yml +++ b/.github/workflows/Downstream.yml @@ -17,6 +17,8 @@ jobs: os: [ubuntu-latest] package: - {user: SciML, repo: OrdinaryDiffEq.jl, group: InterfaceII} + - {user: SciML, repo: NonlinearSolve.jl, group: All} + - {user: SciML, repo: BoundaryValueDiffEq.jl, group: All} steps: - uses: actions/checkout@v4 diff --git a/Project.toml b/Project.toml index 4fca1f27..7d1575a5 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseDiffTools" uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" authors = ["Pankaj Mishra ", "Chris Rackauckas "] -version = "2.9.0" +version = "2.9.1" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/highlevel/coloring.jl b/src/highlevel/coloring.jl index 774b8010..a22ae6de 100644 --- a/src/highlevel/coloring.jl +++ b/src/highlevel/coloring.jl @@ -32,8 +32,8 @@ end # Approximate Jacobian Sparsity Detection ## Right now we hardcode it to use `ForwardDiff` -function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f, x; fx = nothing, - kwargs...) +function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, x; fx = nothing, + kwargs...) where {F <: Function} @unpack ntrials, rng = alg fx = fx === nothing ? f(x) : fx J = fill!(similar(fx, length(fx), length(x)), 0) @@ -47,7 +47,8 @@ function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f, x; fx = fx, kwargs...) end -function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f!, fx, x; kwargs...) +function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f!::F, fx, x; + kwargs...) where {F <: Function} @unpack ntrials, rng = alg cfg = ForwardDiff.JacobianConfig(f!, fx, x) J = fill!(similar(fx, length(fx), length(x)), 0) diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index b8ef2921..b0037435 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -240,12 +240,19 @@ function init_jacobian end const __init_𝒥 = init_jacobian # Misc Functions -function __chunksize(::AutoSparseForwardDiff{C}, x) where {C} - return C === nothing ? ForwardDiff.Chunk(x) : (C isa Number ? ForwardDiff.Chunk{C}() : C) +function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) where {C} + C isa ForwardDiff.Chunk && return C + return __chunksize(Val(C), x) end -__chunksize(::AutoSparseForwardDiff{C}) where {C} = C -__chunksize(::AutoForwardDiff{C}, x) where {C} = C === nothing ? ForwardDiff.Chunk(x) : (C isa Number ? ForwardDiff.Chunk{C}() : C) -__chunksize(::AutoForwardDiff{C}) where {C} = C +__chunksize(::Val{nothing}, x) = ForwardDiff.Chunk(x) +function __chunksize(::Val{C}, x) where {C} + if C isa Integer && !(C isa Bool) + return C ≤ 0 ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{C}() + else + error("$(C)::$(typeof(C)) is not a valid chunksize!") + end +end +__chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C} = C __f̂(f, x, idxs) = dot(vec(f(x)), idxs) diff --git a/src/highlevel/finite_diff.jl b/src/highlevel/finite_diff.jl index c6d2d5e7..5f0e6133 100644 --- a/src/highlevel/finite_diff.jl +++ b/src/highlevel/finite_diff.jl @@ -7,7 +7,7 @@ struct FiniteDiffJacobianCache{CO, CA, J, FX, X} <: AbstractMaybeSparseJacobianC end function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff}, - sd::AbstractMaybeSparsityDetection, f, x; fx = nothing) + sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F <: Function} coloring_result = sd(fd, f, x) fx = fx === nothing ? similar(f(x)) : fx if coloring_result isa NoMatrixColoring @@ -22,7 +22,7 @@ function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff}, end function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff}, - sd::AbstractMaybeSparsityDetection, f!, fx, x) + sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F <: Function} coloring_result = sd(fd, f!, fx, x) if coloring_result isa NoMatrixColoring cache = FiniteDiff.JacobianCache(x, fx) @@ -35,12 +35,14 @@ function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff}, return FiniteDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian!(J::AbstractMatrix, fd, cache::FiniteDiffJacobianCache, f, x) +function sparse_jacobian!(J::AbstractMatrix, fd, cache::FiniteDiffJacobianCache, f::F, + x) where {F <: Function} f!(y, x) = (y .= f(x)) return sparse_jacobian!(J, fd, cache, f!, cache.fx, x) end -function sparse_jacobian!(J::AbstractMatrix, _, cache::FiniteDiffJacobianCache, f!, _, x) +function sparse_jacobian!(J::AbstractMatrix, _, cache::FiniteDiffJacobianCache, f!::F, _, + x) where {F <: Function} FiniteDiff.finite_difference_jacobian!(J, f!, x, cache.cache) return J end diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl index a563bc68..c027e19a 100644 --- a/src/highlevel/forward_mode.jl +++ b/src/highlevel/forward_mode.jl @@ -9,7 +9,7 @@ end struct SparseDiffToolsTag end function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, - sd::AbstractMaybeSparsityDetection, f, x; fx = nothing) + sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F <: Function} coloring_result = sd(ad, f, x) fx = fx === nothing ? similar(f(x)) : fx if coloring_result isa NoMatrixColoring @@ -25,7 +25,7 @@ function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff} end function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, - sd::AbstractMaybeSparsityDetection, f!, fx, x) + sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F <: Function} coloring_result = sd(ad, f!, fx, x) if coloring_result isa NoMatrixColoring cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x), @@ -39,7 +39,8 @@ function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff} return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f, x) +function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f::F, + x) where {F <: Function} if cache.cache isa ForwardColorJacCache forwarddiff_color_jacobian(J, f, x, cache.cache) # Use Sparse ForwardDiff else @@ -48,7 +49,8 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, return J end -function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f!, fx, x) +function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f!::F, fx, + x) where {F <: Function} if cache.cache isa ForwardColorJacCache forwarddiff_color_jacobian!(J, f!, x, cache.cache) # Use Sparse ForwardDiff else diff --git a/src/highlevel/reverse_mode.jl b/src/highlevel/reverse_mode.jl index dc34ab1b..0fd70284 100644 --- a/src/highlevel/reverse_mode.jl +++ b/src/highlevel/reverse_mode.jl @@ -8,7 +8,7 @@ struct ReverseModeJacobianCache{CO, CA, J, FX, X, I} <: AbstractMaybeSparseJacob end function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode}, - sd::AbstractMaybeSparsityDetection, f, x; fx = nothing) + sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F <: Function} fx = fx === nothing ? similar(f(x)) : fx coloring_result = sd(ad, f, x) jac_prototype = __getfield(coloring_result, Val(:jacobian_sparsity)) @@ -17,7 +17,7 @@ function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode}, end function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode}, - sd::AbstractMaybeSparsityDetection, f!, fx, x) + sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F <: Function} coloring_result = sd(ad, f!, fx, x) jac_prototype = __getfield(coloring_result, Val(:jacobian_sparsity)) return ReverseModeJacobianCache(coloring_result, nothing, jac_prototype, fx, x, @@ -34,12 +34,12 @@ function sparse_jacobian!(J::AbstractMatrix, ad, cache::ReverseModeJacobianCache end function __sparse_jacobian_reverse_impl!(J::AbstractMatrix, ad, idx_vec, - cache::MatrixColoringResult, f, x) + cache::MatrixColoringResult, f::F, x) where {F <: Function} return __sparse_jacobian_reverse_impl!(J, ad, idx_vec, cache, f, nothing, x) end function __sparse_jacobian_reverse_impl!(J::AbstractMatrix, ad, idx_vec, - cache::MatrixColoringResult, f, fx, x) + cache::MatrixColoringResult, f::F, fx, x) where {F <: Function} # If `fx` is `nothing` then assume `f` is not in-place x_ = __maybe_copy_x(ad, x) fx_ = __maybe_copy_x(ad, fx) diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index 290d350d..5bfdd8d1 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -40,6 +40,8 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in (AutoSparseZygote(), AutoZygote(), AutoSparseForwardDiff(), AutoForwardDiff(), + AutoSparseForwardDiff(; chunksize = 0), AutoForwardDiff(; chunksize = 0), + AutoSparseForwardDiff(; chunksize = 8), AutoForwardDiff(; chunksize = 8), AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme()) @testset "Cache & Reuse" begin cache = sparse_jacobian_cache(difftype, sd, fdiff, x) @@ -92,8 +94,10 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Inplace Place Function" @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in (AutoSparseForwardDiff(), - AutoForwardDiff(), AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), - AutoSparseEnzyme()) + AutoForwardDiff(), AutoSparseForwardDiff(; chunksize = 0), + AutoForwardDiff(; chunksize = 0), AutoSparseForwardDiff(; chunksize = 8), + AutoForwardDiff(; chunksize = 8), AutoSparseFiniteDiff(), AutoFiniteDiff(), + AutoEnzyme(), AutoSparseEnzyme()) y = similar(x) cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x)