Enable tag checking

JuliaDiff · Oct 31, 2023 · d6fae4e · d6fae4e
1 parent a7666fa
commit d6fae4e
Showing 1 changed file with 16 additions and 11 deletions.
diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl
@@ -10,20 +10,26 @@ __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, _) = tag
+__standard_tag(tag, x) = ForwardDiff.Tag(tag, 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)
     if coloring_result isa NoMatrixColoring
-        cache = ForwardDiff.JacobianConfig(f, x, __chunksize(ad, x),
-            __standard_tag(ad.tag, x))
+        cache = ForwardDiff.JacobianConfig(f, x, __chunksize(ad, x), tag)
         jac_prototype = nothing
     else
+        # 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, ad.tag)
+            dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag))
         jac_prototype = coloring_result.jacobian_sparsity
     end
     return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x)
@@ -32,13 +38,14 @@ 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)
     if coloring_result isa NoMatrixColoring
-        cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x),
-            __standard_tag(ad.tag, x))
+        cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x), tag)
         jac_prototype = nothing
     else
+        # 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, ad.tag)
+            dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag))
         jac_prototype = coloring_result.jacobian_sparsity
     end
     return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x)
@@ -49,8 +56,7 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache,
     if cache.cache isa ForwardColorJacCache
         forwarddiff_color_jacobian(J, f, x, cache.cache) # Use Sparse ForwardDiff
     else
-        # Disable tag checking since we set the tag to our custom tag
-        ForwardDiff.jacobian!(J, f, x, cache.cache, Val(false)) # Don't try to exploit sparsity
+        ForwardDiff.jacobian!(J, f, x, cache.cache) # Don't try to exploit sparsity
     end
     return J
 end
@@ -60,8 +66,7 @@ function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache,
     if cache.cache isa ForwardColorJacCache
         forwarddiff_color_jacobian!(J, f!, x, cache.cache) # Use Sparse ForwardDiff
     else
-        # Disable tag checking since we set the tag to our custom tag
-        ForwardDiff.jacobian!(J, f!, fx, x, cache.cache, Val(false)) # Don't try to exploit sparsity
+        ForwardDiff.jacobian!(J, f!, fx, x, cache.cache) # Don't try to exploit sparsity
     end
     return J
 end