Don't use custom tags

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <[email protected]>"]
-version = "2.17.0"
+version = "2.17.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

ext/SparseDiffToolsPolyesterForwardDiffExt.jl

@@ -6,8 +6,7 @@ import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityD
                         ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache,
                         sparse_jacobian!,
                         sparse_jacobian_static_array, __standard_tag, __chunksize,
-                        polyesterforwarddiff_color_jacobian,
-                        polyesterforwarddiff_color_jacobian!
+                        polyesterforwarddiff_color_jacobian
 
 struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <:
        AbstractMaybeSparseJacobianCache
@@ -27,7 +26,7 @@ function sparse_jacobian_cache(
         cache = __chunksize(ad, x)
         jac_prototype = nothing
     else
-        tag = __standard_tag(nothing, x)
+        tag = __standard_tag(nothing, f, 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))
@@ -47,7 +46,7 @@ function sparse_jacobian_cache(
         @warn """Currently PolyesterForwardDiff does not support sparsity detection
                  natively for inplace functions. Falling back to using
                 ForwardDiff.jl""" maxlog=1
-        tag = __standard_tag(nothing, x)
+        tag = __standard_tag(nothing, f!, 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))

src/highlevel/forward_mode.jl

@@ -8,22 +8,15 @@ end
 
 __getfield(c::ForwardDiffJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype
 
-struct SparseDiffToolsTag end
-
-function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:SparseDiffToolsTag, <:T}}, f::F,
-        x::AbstractArray{T}) where {T, F}
-    return true
-end
-
-__standard_tag(::Nothing, x) = ForwardDiff.Tag(SparseDiffToolsTag(), eltype(x))
-__standard_tag(tag::ForwardDiff.Tag, _) = tag
-__standard_tag(tag, x) = ForwardDiff.Tag(tag, eltype(x))
+__standard_tag(::Nothing, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x))
+__standard_tag(tag::ForwardDiff.Tag, ::F, _) where {F} = tag
+__standard_tag(tag, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x))
 
 function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff},
         sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F}
     coloring_result = sd(ad, f, x)
     fx = fx === nothing ? similar(f(x)) : fx
-    tag = __standard_tag(ad.tag, x)
+    tag = __standard_tag(ad.tag, f, x)
     if coloring_result isa NoMatrixColoring
         cache = ForwardDiff.JacobianConfig(f, x, __chunksize(ad, x), tag)
         jac_prototype = nothing
@@ -39,7 +32,7 @@ end
 function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff},
         sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F}
     coloring_result = sd(ad, f!, fx, x)
-    tag = __standard_tag(ad.tag, x)
+    tag = __standard_tag(ad.tag, f!, x)
     if coloring_result isa NoMatrixColoring
         cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x), tag)
         jac_prototype = nothing