Add PolyesterForwardDiff Support

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.15.1"
+version = "2.16.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -27,17 +27,19 @@ VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [weakdeps]
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 SparseDiffToolsEnzymeExt = "Enzyme"
+SparseDiffToolsPolyesterForwardDiffExt = "PolyesterForwardDiff"
 SparseDiffToolsSymbolicsExt = "Symbolics"
 SparseDiffToolsZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "0.2.1"
-Adapt = "3.0, 4"
+ADTypes = "0.2.6"
+Adapt = "3, 4"
 ArrayInterface = "7.4.2"
 Compat = "4"
 DataStructures = "0.18"
@@ -47,7 +49,8 @@ ForwardDiff = "0.10"
 Graphs = "1"
 LinearAlgebra = "<0.0.1, 1"
 PackageExtensionCompat = "1"
-Random = "<0.0.1, 1"
+PolyesterForwardDiff = "0.1.1"
+Random = "1.6"
 Reexport = "1"
 SciMLOperators = "0.3.7"
 Setfield = "1"
@@ -67,6 +70,7 @@ BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -75,4 +79,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]
+test = ["Test", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays", "PolyesterForwardDiff"]

ext/SparseDiffToolsPolyesterForwardDiffExt.jl

@@ -0,0 +1,77 @@
+module SparseDiffToolsPolyesterForwardDiffExt
+
+using ADTypes, SparseDiffTools, PolyesterForwardDiff
+import ForwardDiff
+import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityDetection,
+    ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache, sparse_jacobian!,
+    sparse_jacobian_static_array, __standard_tag, __chunksize
+
+struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <:
+       AbstractMaybeSparseJacobianCache
+    coloring::CO
+    cache::CA
+    jac_prototype::J
+    fx::FX
+    x::X
+end
+
+function sparse_jacobian_cache(ad::Union{AutoSparsePolyesterForwardDiff,
+            AutoPolyesterForwardDiff}, 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
+        cache = __chunksize(ad, x)
+        jac_prototype = nothing
+    else
+        @warn """Currently PolyesterForwardDiff does not support sparsity detection
+                 natively. Falling back to using ForwardDiff.jl""" maxlog=1
+        tag = __standard_tag(nothing, x)
+        # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof`
+        cache = ForwardColorJacCache(f, x, __chunksize(ad); coloring_result.colorvec,
+            dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag))
+        jac_prototype = coloring_result.jacobian_sparsity
+    end
+    return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x)
+end
+
+function sparse_jacobian_cache(ad::Union{AutoSparsePolyesterForwardDiff,
+            AutoPolyesterForwardDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx,
+        x) where {F}
+    coloring_result = sd(ad, f!, fx, x)
+    if coloring_result isa NoMatrixColoring
+        cache = __chunksize(ad, x)
+        jac_prototype = nothing
+    else
+        @warn """Currently PolyesterForwardDiff does not support sparsity detection
+                 natively. Falling back to using ForwardDiff.jl""" maxlog=1
+        tag = __standard_tag(nothing, x)
+        # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof`
+        cache = ForwardColorJacCache(f!, x, __chunksize(ad); coloring_result.colorvec,
+            dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag))
+        jac_prototype = coloring_result.jacobian_sparsity
+    end
+    return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x)
+end
+
+function sparse_jacobian!(J::AbstractMatrix, _, cache::PolyesterForwardDiffJacobianCache,
+        f::F, x) where {F}
+    if cache.cache isa ForwardColorJacCache
+        forwarddiff_color_jacobian(J, f, x, cache.cache) # Use Sparse ForwardDiff
+    else
+        PolyesterForwardDiff.threaded_jacobian!(f, J, x, cache.cache) # Don't try to exploit sparsity
+    end
+    return J
+end
+
+function sparse_jacobian!(J::AbstractMatrix, _, cache::PolyesterForwardDiffJacobianCache,
+        f!::F, fx, x) where {F}
+    if cache.cache isa ForwardColorJacCache
+        forwarddiff_color_jacobian!(J, f!, x, cache.cache) # Use Sparse ForwardDiff
+    else
+        PolyesterForwardDiff.threaded_jacobian!(f!, fx, J, x, cache.cache) # Don't try to exploit sparsity
+    end
+    return J
+end
+
+end

src/highlevel/common.jl

@@ -269,7 +269,8 @@ function init_jacobian end
 const __init_𝒥 = init_jacobian
 
 # Misc Functions
-function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) where {C}
+function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C},
+            AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}, x) where {C}
     C isa ForwardDiff.Chunk && return C
     return __chunksize(Val(C), x)
 end
@@ -285,7 +286,8 @@ end
 __chunksize(x) = ForwardDiff.Chunk(x)
 __chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}()
 
-function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C}
+function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C},
+        AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}) where {C}
     C === nothing && return nothing
     C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C)
     return nothing
@@ -347,4 +349,5 @@ end
 @inline __backend(::Union{AutoEnzyme, AutoSparseEnzyme}) = :Enzyme
 @inline __backend(::Union{AutoZygote, AutoSparseZygote}) = :Zygote
 @inline __backend(::Union{AutoForwardDiff, AutoSparseForwardDiff}) = :ForwardDiff
+@inline __backend(::Union{AutoPolyesterForwardDiff, AutoSparsePolyesterForwardDiff}) = :PolyesterForwardDiff
 @inline __backend(::Union{AutoFiniteDiff, AutoSparseFiniteDiff}) = :FiniteDiff

test/test_sparse_jacobian.jl

@@ -1,6 +1,6 @@
 ## Sparse Jacobian tests
-using SparseDiffTools,
-    Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, StaticArrays
+using SparseDiffTools, PolyesterForwardDiff, Symbolics, ForwardDiff, LinearAlgebra,
+    SparseArrays, Zygote, Enzyme, Test, StaticArrays
 
 @views function fdiff(y, x) # in-place
     L = length(x)
@@ -42,7 +42,12 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
             AutoZygote(), AutoSparseForwardDiff(), AutoForwardDiff(),
             AutoSparseForwardDiff(; chunksize = 0), AutoForwardDiff(; chunksize = 0),
             AutoSparseForwardDiff(; chunksize = 4), AutoForwardDiff(; chunksize = 4),
-            AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme())
+            AutoSparsePolyesterForwardDiff(), AutoPolyesterForwardDiff(),
+            AutoSparsePolyesterForwardDiff(; chunksize = 0),
+            AutoPolyesterForwardDiff(; chunksize = 0),
+            AutoSparsePolyesterForwardDiff(; chunksize = 4),
+            AutoPolyesterForwardDiff(; chunksize = 4), AutoSparseFiniteDiff(),
+            AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme())
             @testset "Cache & Reuse" begin
                 cache = sparse_jacobian_cache(difftype, sd, fdiff, x)
                 J = init_jacobian(cache)
@@ -59,7 +64,9 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
 
                 @test J ≈ J_true
 
-                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff)
+                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff ||
+                     difftype isa AutoSparsePolyesterForwardDiff ||
+                     difftype isa AutoPolyesterForwardDiff)
                     @inferred sparse_jacobian(difftype, cache, fdiff, x)
                 end
 
@@ -71,7 +78,9 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
                 J = sparse_jacobian(difftype, sd, fdiff, x)
 
                 @test J ≈ J_true
-                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff)
+                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff ||
+                     difftype isa AutoSparsePolyesterForwardDiff ||
+                     difftype isa AutoPolyesterForwardDiff)
                     @inferred sparse_jacobian(difftype, sd, fdiff, x)
                 end
 
@@ -114,7 +123,9 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
                 J = sparse_jacobian(difftype, cache, fdiff, y, x)
 
                 @test J ≈ J_true
-                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff)
+                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff ||
+                     difftype isa AutoSparsePolyesterForwardDiff ||
+                     difftype isa AutoPolyesterForwardDiff)
                     @inferred sparse_jacobian(difftype, cache, fdiff, y, x)
                 end
 
@@ -126,7 +137,9 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa
                 J = sparse_jacobian(difftype, sd, fdiff, y, x)
 
                 @test J ≈ J_true
-                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff)
+                if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff ||
+                     difftype isa AutoSparsePolyesterForwardDiff ||
+                     difftype isa AutoPolyesterForwardDiff)
                     @inferred sparse_jacobian(difftype, sd, fdiff, y, x)
                 end