Merge pull request #268 from avik-pal/patch-1

Fix Chunking

Author: ChrisRackauckas

.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

.github/workflows/Downstream.yml

@@ -12,11 +12,14 @@ jobs:
     env:
       GROUP: ${{ matrix.package.group }}
     strategy:
+      fail-fast: false
       matrix:
         julia-version: [1,1.6]
         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

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>", "Chris Rackauckas <contact@chrisrackauckas.com>"]
-version = "2.9.0"
+version = "2.9.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

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}
     @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}
     @unpack ntrials, rng = alg
     cfg = ForwardDiff.JacobianConfig(f!, fx, x)
     J = fill!(similar(fx, length(fx), length(x)), 0)

src/highlevel/common.jl

@@ -240,12 +240,23 @@ function init_jacobian end
 const __init_𝒥 = init_jacobian
 
 # Misc Functions
-function __chunksize(::AutoSparseForwardDiff{C}, x) where {C}
-    return C === nothing ? ForwardDiff.Chunk(x) : C
+function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) where {C}
+    C isa ForwardDiff.Chunk && return C
+    return __chunksize(Val(C), x)
+end
+__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
+function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C}
+    C === nothing && return nothing
+    C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C)
+    return nothing
 end
-__chunksize(::AutoSparseForwardDiff{C}) where {C} = C
-__chunksize(::AutoForwardDiff{C}, x) where {C} = C === nothing ? ForwardDiff.Chunk(x) : C
-__chunksize(::AutoForwardDiff{C}) where {C} = C
 
 __f̂(f, x, idxs) = dot(vec(f(x)), idxs)
 

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}
     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}
     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}
     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}
     FiniteDiff.finite_difference_jacobian!(J, f!, x, cache.cache)
     return J
 end

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}
     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}
     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}
     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}
     if cache.cache isa ForwardColorJacCache
         forwarddiff_color_jacobian!(J, f!, x, cache.cache) # Use Sparse ForwardDiff
     else

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}
     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}
     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}
     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}
     # If `fx` is `nothing` then assume `f` is not in-place
     x_ = __maybe_copy_x(ad, x)
     fx_ = __maybe_copy_x(ad, fx)

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 = 4), AutoForwardDiff(; chunksize = 4),
             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 = 4),
+            AutoForwardDiff(; chunksize = 4), AutoSparseFiniteDiff(), AutoFiniteDiff(),
+            AutoEnzyme(), AutoSparseEnzyme())
             y = similar(x)
             cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x)