Path to use FiniteDiff for ApproximateSparsityDetection

Author: avik-pal

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.14.0"
+version = "2.15.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

ext/SparseDiffToolsEnzymeExt.jl

@@ -1,14 +1,16 @@
 module SparseDiffToolsEnzymeExt
 
 import ArrayInterface: fast_scalar_indexing
-import SparseDiffTools: __f̂,
-    __maybe_copy_x, __jacobian!, __gradient, __gradient!, AutoSparseEnzyme
+import SparseDiffTools: __f̂, __maybe_copy_x, __jacobian!, __gradient, __gradient!,
+    AutoSparseEnzyme, __test_backend_loaded
 # FIXME: For Enzyme we currently assume reverse mode
 import ADTypes: AutoEnzyme
 using Enzyme
 
 using ForwardDiff
 
+@inline __test_backend_loaded(::Union{AutoSparseEnzyme, AutoEnzyme}) = nothing
+
 ## Satisfying High-Level Interface for Sparse Jacobians
 function __gradient(::Union{AutoSparseEnzyme, AutoEnzyme}, f, x, cols)
     dx = zero(x)

ext/SparseDiffToolsZygoteExt.jl

@@ -1,7 +1,7 @@
 module SparseDiffToolsZygoteExt
 
 using ADTypes, LinearAlgebra, Zygote
-import SparseDiffTools: SparseDiffTools, DeivVecTag, AutoDiffVJP
+import SparseDiffTools: SparseDiffTools, DeivVecTag, AutoDiffVJP, __test_backend_loaded
 import ForwardDiff: ForwardDiff, Dual, partials
 import SciMLOperators: update_coefficients, update_coefficients!
 import Setfield: @set!
@@ -12,6 +12,8 @@ import SparseDiffTools: numback_hesvec!,
 import SparseDiffTools: __f̂, __jacobian!, __gradient, __gradient!
 import ADTypes: AutoZygote, AutoSparseZygote
 
+@inline __test_backend_loaded(::Union{AutoSparseZygote, AutoZygote}) = nothing
+
 ## Satisfying High-Level Interface for Sparse Jacobians
 function __gradient(::Union{AutoSparseZygote, AutoZygote}, f::F, x, cols) where {F}
     _, ∂x, _ = Zygote.gradient(__f̂, f, x, cols)

src/highlevel/coloring.jl

@@ -34,28 +34,70 @@ end
 ## Right now we hardcode it to use `ForwardDiff`
 function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, x; fx = nothing,
         kwargs...) where {F}
+    if !(ad isa AutoSparseForwardDiff)
+        @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
+    end
     @unpack ntrials, rng = alg
     fx = fx === nothing ? f(x) : fx
-    J = fill!(similar(fx, length(fx), length(x)), 0)
     cfg = ForwardDiff.JacobianConfig(f, x)
+    J = fill!(similar(fx, length(fx), length(x)), 0)
+    J_cache = similar(J)
+    x_ = similar(x)
+    for _ in 1:ntrials
+        randn!(rng, x_)
+        ForwardDiff.jacobian!(J_cache, f, x_, cfg)
+        @. J += abs(J_cache)
+    end
+    return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))(ad, f, x;
+        fx, kwargs...)
+end
+
+function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, fx, x;
+        kwargs...) where {F}
+    if !(ad isa AutoSparseForwardDiff)
+        @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
+    end
+    @unpack ntrials, rng = alg
+    cfg = ForwardDiff.JacobianConfig(f, fx, x)
+    J = fill!(similar(fx, length(fx), length(x)), 0)
+    J_cache = similar(J)
+    x_ = similar(x)
+    for _ in 1:ntrials
+        randn!(rng, x_)
+        ForwardDiff.jacobian!(J_cache, f, fx, x_, cfg)
+        @. J += abs(J_cache)
+    end
+    return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))(ad, f, x;
+        fx, kwargs...)
+end
+
+function (alg::ApproximateJacobianSparsity)(ad::AutoSparseFiniteDiff, f::F, x; fx = nothing,
+        kwargs...) where {F}
+    @unpack ntrials, rng = alg
+    fx = fx === nothing ? f(x) : fx
+    cache = FiniteDiff.JacobianCache(x, fx)
+    J = fill!(similar(fx, length(fx), length(x)), 0)
+    x_ = similar(x)
     for _ in 1:ntrials
-        x_ = similar(x)
         randn!(rng, x_)
-        J .+= abs.(ForwardDiff.jacobian(f, x_, cfg))
+        J_cache = FiniteDiff.finite_difference_jacobian(f, x, cache)
+        @. J += abs(J_cache)
     end
     return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))(ad, f, x;
         fx, kwargs...)
 end
 
-function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f!::F, fx, x;
+function (alg::ApproximateJacobianSparsity)(ad::AutoSparseFiniteDiff, f!::F, fx, x;
         kwargs...) where {F}
     @unpack ntrials, rng = alg
-    cfg = ForwardDiff.JacobianConfig(f!, fx, x)
+    cache = FiniteDiff.JacobianCache(x, fx)
     J = fill!(similar(fx, length(fx), length(x)), 0)
+    J_cache = similar(J)
+    x_ = similar(x)
     for _ in 1:ntrials
-        x_ = similar(x)
         randn!(rng, x_)
-        J .+= abs.(ForwardDiff.jacobian(f!, fx, x_, cfg))
+        FiniteDiff.finite_difference_jacobian!(J_cache, f!, x_, cache)
+        @. J += abs(J_cache)
     end
     return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))(ad, f!, fx,
         x; kwargs...)

src/highlevel/common.jl

@@ -332,3 +332,15 @@ init_jacobian(J::SparseMatrixCSC, ::Type{T}, fx, x; kwargs...) where {T} = T.(J)
 
 __maybe_copy_x(_, x) = x
 __maybe_copy_x(_, ::Nothing) = nothing
+
+# Check Backend has been loaded
+## We pay a small compile time cost for this, but it avoids cryptic error messages
+@inline function __test_backend_loaded(ad::AbstractADType)
+    error("$(ad) requires $(__backend(ad)).jl to be loaded. Please load it.")
+end
+
+@inline __backend(ad) = nothing
+@inline __backend(::Union{AutoEnzyme, AutoSparseEnzyme}) = :Enzyme
+@inline __backend(::Union{AutoZygote, AutoSparseZygote}) = :Zygote
+@inline __backend(::Union{AutoForwardDiff, AutoSparseForwardDiff}) = :ForwardDiff
+@inline __backend(::Union{AutoFiniteDiff, AutoSparseFiniteDiff}) = :FiniteDiff

src/highlevel/reverse_mode.jl

@@ -28,6 +28,7 @@ end
 
 function sparse_jacobian!(J::AbstractMatrix, ad, cache::ReverseModeJacobianCache, args...)
     if cache.coloring isa NoMatrixColoring
+        __test_backend_loaded(ad)
         return __jacobian!(J, ad, args...)
     else
         return __sparse_jacobian_reverse_impl!(J, ad, cache.idx_vec, cache.coloring,
@@ -43,6 +44,7 @@ end
 function __sparse_jacobian_reverse_impl!(J::AbstractMatrix, ad, idx_vec,
         cache::MatrixColoringResult, f::F, fx, x) where {F}
     # If `fx` is `nothing` then assume `f` is not in-place
+    __test_backend_loaded(ad)
     x_ = __maybe_copy_x(ad, x)
     fx_ = __maybe_copy_x(ad, fx)
     @unpack colorvec, nz_rows, nz_cols = cache