Add sparsity detector field to sparse backends (prototype)

Project.toml

@@ -1,7 +1,7 @@
 name = "ADTypes"
 uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 authors = ["Vaibhav Dixit <[email protected]> and contributors"]
-version = "0.2.7"
+version = "0.2.8"
 
 [compat]
 julia = "1.6"

src/ADTypes.jl

@@ -21,6 +21,8 @@ abstract type AbstractSparseFiniteDifferences <: AbstractFiniteDifferencesMode e
 abstract type AbstractSparseSymbolicDifferentiationMode <:
               AbstractSymbolicDifferentiationMode end
 
+## Dense
+
 """
     AutoChainRules{RC}
 
@@ -34,6 +36,33 @@ Base.@kwdef struct AutoChainRules{RC} <: AbstractADType
     ruleconfig::RC
 end
 
+"""
+    AutoDiffractor
+
+Chooses [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl).
+"""
+struct AutoDiffractor <: AbstractADType end
+
+"""
+    AutoEnzyme{M}
+
+Chooses [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
+
+# Fields
+
+- `mode::M = nothing`
+"""
+Base.@kwdef struct AutoEnzyme{M} <: AbstractADType
+    mode::M = nothing
+end
+
+"""
+    AutoFastDifferentiation
+
+Chooses [FastDifferentiation.jl](https://github.com/brianguenter/FastDifferentiation.jl).
+"""
+struct AutoFastDifferentiation <: AbstractSymbolicDifferentiationMode end
+
 """
     AutoFiniteDiff{T1,T2,T3}
 
@@ -86,6 +115,21 @@ function AutoForwardDiff(; chunksize = nothing, tag = nothing)
     AutoForwardDiff{chunksize, typeof(tag)}(tag)
 end
 
+"""
+    AutoModelingToolkit
+
+Chooses [ModelingToolkit.jl](https://github.com/SciML/ModelingToolkit.jl).
+
+# Fields
+
+- `obj_sparse::Bool = false`
+- `cons_sparse::Bool = false`
+"""
+Base.@kwdef struct AutoModelingToolkit <: AbstractSymbolicDifferentiationMode
+    obj_sparse::Bool = false
+    cons_sparse::Bool = false
+end
+
 """
     AutoPolyesterForwardDiff{chunksize}
 
@@ -116,33 +160,6 @@ Base.@kwdef struct AutoReverseDiff <: AbstractReverseMode
     compile::Bool = false
 end
 
-"""
-    AutoZygote
-
-Chooses [Zygote.jl](https://github.com/FluxML/Zygote.jl).
-"""
-struct AutoZygote <: AbstractReverseMode end
-
-"""
-    AutoSparseZygote
-
-Chooses [Zygote.jl](https://github.com/FluxML/Zygote.jl) while exploiting sparsity.
-"""
-struct AutoSparseZygote <: AbstractSparseReverseMode end
-
-"""
-    AutoEnzyme{M}
-
-Chooses [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl).
-
-# Fields
-
-- `mode::M = nothing`
-"""
-Base.@kwdef struct AutoEnzyme{M} <: AbstractADType
-    mode::M = nothing
-end
-
 """
     AutoTracker
 
@@ -151,29 +168,30 @@ Chooses [Tracker.jl](https://github.com/FluxML/Tracker.jl).
 struct AutoTracker <: AbstractReverseMode end
 
 """
-    AutoModelingToolkit
-
-Chooses [ModelingToolkit.jl](https://github.com/SciML/ModelingToolkit.jl).
-
-# Fields
+    AutoZygote
 
-- `obj_sparse::Bool = false`
-- `cons_sparse::Bool = false`
+Chooses [Zygote.jl](https://github.com/FluxML/Zygote.jl).
 """
-Base.@kwdef struct AutoModelingToolkit <: AbstractSymbolicDifferentiationMode
-    obj_sparse::Bool = false
-    cons_sparse::Bool = false
-end
+struct AutoZygote <: AbstractReverseMode end
+
+## Sparse
 
 """
-    AutoSparseFiniteDiff
+    AutoSparseFiniteDiff{T1,T2,T3}
 
 Chooses [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) while exploiting sparsity.
+
+# Fields
+
+- `fdtype::T1 = Val(:forward)`
+- `fdjtype::T2 = fdtype`
+- `fdhtype::T3 = Val(:hcentral)`
 """
-Base.@kwdef struct AutoSparseFiniteDiff{T1, T2, T3} <: AbstractSparseFiniteDifferences
+Base.@kwdef struct AutoSparseFiniteDiff{T1, T2, T3, S} <: AbstractSparseFiniteDifferences
     fdtype::T1 = Val(:forward)
     fdjtype::T2 = fdtype
     fdhtype::T3 = Val(:hcentral)
+    sparsity_detector::S = nothing
 end
 
 """
@@ -185,34 +203,38 @@ Chooses [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) while expl
 
 - `tag::T`
 """
-struct AutoSparseForwardDiff{chunksize, T} <: AbstractSparseForwardMode
+struct AutoSparseForwardDiff{chunksize, T, S} <: AbstractSparseForwardMode
     tag::T
+    sparsity_detector::S
 end
 
 """
     AutoSparseForwardDiff(; chunksize = nothing, tag = nothing)
 
 Constructor.
 """
-function AutoSparseForwardDiff(; chunksize = nothing, tag = nothing)
-    AutoSparseForwardDiff{chunksize, typeof(tag)}(tag)
+function AutoSparseForwardDiff(;
+        chunksize = nothing, tag = nothing, sparsity_detector = nothing)
+    AutoSparseForwardDiff{chunksize, typeof(tag), typeof(sparsity_detector)}(
+        tag, sparsity_detector)
 end
 
 """
     AutoSparsePolyesterForwardDiff{chunksize}
 
 Chooses [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) while exploiting sparsity.
 """
-struct AutoSparsePolyesterForwardDiff{chunksize} <: AbstractSparseForwardMode
+struct AutoSparsePolyesterForwardDiff{chunksize, S} <: AbstractSparseForwardMode
+    sparsity_detector::S
 end
 
 """
     AutoSparsePolyesterForwardDiff(; chunksize = nothing)
 
 Constructor.
 """
-function AutoSparsePolyesterForwardDiff(; chunksize = nothing)
-    AutoSparsePolyesterForwardDiff{chunksize}()
+function AutoSparsePolyesterForwardDiff(; chunksize = nothing, sparsity_detector = nothing)
+    AutoSparsePolyesterForwardDiff{chunksize, typeof(sparsity_detector)}(sparsity_detector)
 end
 
 """
@@ -224,23 +246,19 @@ Chooses [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl) while expl
 
 - `compile::Bool = false`
 """
-Base.@kwdef struct AutoSparseReverseDiff <: AbstractSparseReverseMode
+Base.@kwdef struct AutoSparseReverseDiff{S} <: AbstractSparseReverseMode
     compile::Bool = false
+    sparsity_detector::S = nothing
 end
 
 """
-    AutoDiffractor
-
-Chooses [Diffractor.jl](https://github.com/JuliaDiff/Diffractor.jl).
-"""
-struct AutoDiffractor <: AbstractADType end
-
-"""
-    AutoFastDifferentiation
+    AutoSparseZygote
 
-Chooses [FastDifferentiation.jl](https://github.com/brianguenter/FastDifferentiation.jl).
+Chooses [Zygote.jl](https://github.com/FluxML/Zygote.jl) while exploiting sparsity.
 """
-struct AutoFastDifferentiation <: AbstractSymbolicDifferentiationMode end
+Base.@kwdef struct AutoSparseZygote{S} <: AbstractSparseReverseMode
+    sparsity_detector::S = nothing
+end
 
 """
     AutoSparseFastDifferentiation