Dont break code for non-functions

Author: avik-pal

src/highlevel/coloring.jl

@@ -33,7 +33,7 @@ end
 # Approximate Jacobian Sparsity Detection
 ## Right now we hardcode it to use `ForwardDiff`
 function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, x; fx = nothing,
-    kwargs...) where {F <: Function}
+    kwargs...) where {F}
     @unpack ntrials, rng = alg
     fx = fx === nothing ? f(x) : fx
     J = fill!(similar(fx, length(fx), length(x)), 0)
@@ -48,7 +48,7 @@ function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, x; f
 end
 
 function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f!::F, fx, x;
-    kwargs...) where {F <: Function}
+    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

@@ -252,7 +252,11 @@ function __chunksize(::Val{C}, x) where {C}
         error("$(C)::$(typeof(C)) is not a valid chunksize!")
     end
 end
-__chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C} = C
+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
 
 __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::F, x; fx = nothing) where {F <: Function}
+    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!::F, fx, x) where {F <: Function}
+    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)
@@ -36,13 +36,13 @@ function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff},
 end
 
 function sparse_jacobian!(J::AbstractMatrix, fd, cache::FiniteDiffJacobianCache, f::F,
-    x) where {F <: Function}
+    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!::F, _,
-    x) where {F <: Function}
+    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::F, x; fx = nothing) where {F <: Function}
+    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!::F, fx, x) where {F <: Function}
+    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),
@@ -40,7 +40,7 @@ function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}
 end
 
 function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f::F,
-    x) where {F <: Function}
+    x) where {F}
     if cache.cache isa ForwardColorJacCache
         forwarddiff_color_jacobian(J, f, x, cache.cache) # Use Sparse ForwardDiff
     else
@@ -50,7 +50,7 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache,
 end
 
 function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f!::F, fx,
-    x) where {F <: Function}
+    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::F, x; fx = nothing) where {F <: Function}
+    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!::F, fx, x) where {F <: Function}
+    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::F, x) where {F <: Function}
+    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::F, fx, x) where {F <: Function}
+    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

@@ -41,7 +41,7 @@ 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),
+            AutoSparseForwardDiff(; chunksize = 4), AutoForwardDiff(; chunksize = 4),
             AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme())
             @testset "Cache & Reuse" begin
                 cache = sparse_jacobian_cache(difftype, sd, fdiff, x)
@@ -95,8 +95,8 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
 
         @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in (AutoSparseForwardDiff(),
             AutoForwardDiff(), AutoSparseForwardDiff(; chunksize = 0),
-            AutoForwardDiff(; chunksize = 0), AutoSparseForwardDiff(; chunksize = 8),
-            AutoForwardDiff(; chunksize = 8), AutoSparseFiniteDiff(), AutoFiniteDiff(),
+            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)