Merge pull request #266 from avik-pal/ap/tstable

Pass Tag and Chunksize from the types

Author: ChrisRackauckas

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

src/highlevel/common.jl

@@ -240,7 +240,12 @@ function init_jacobian end
 const __init_𝒥 = init_jacobian
 
 # Misc Functions
+function __chunksize(::AutoSparseForwardDiff{C}, x) where {C}
+    return C === nothing ? ForwardDiff.Chunk(x) : C
+end
 __chunksize(::AutoSparseForwardDiff{C}) where {C} = C
+__chunksize(::AutoForwardDiff{C}, x) where {C} = C === nothing ? ForwardDiff.Chunk(x) : C
+__chunksize(::AutoForwardDiff{C}) where {C} = C
 
 __f̂(f, x, idxs) = dot(vec(f(x)), idxs)
 

src/highlevel/forward_mode.jl

@@ -6,17 +6,19 @@ struct ForwardDiffJacobianCache{CO, CA, J, FX, X} <: AbstractMaybeSparseJacobian
     x::X
 end
 
+struct SparseDiffToolsTag end
+
 function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff},
     sd::AbstractMaybeSparsityDetection, f, x; fx = nothing)
     coloring_result = sd(ad, f, x)
     fx = fx === nothing ? similar(f(x)) : fx
     if coloring_result isa NoMatrixColoring
-        cache = ForwardDiff.JacobianConfig(f, x)
+        cache = ForwardDiff.JacobianConfig(f, x, __chunksize(ad, x),
+            ifelse(ad.tag === nothing, SparseDiffToolsTag(), ad.tag))
         jac_prototype = nothing
     else
         cache = ForwardColorJacCache(f, x, __chunksize(ad); coloring_result.colorvec,
-            dx = fx,
-            sparsity = coloring_result.jacobian_sparsity)
+            dx = fx, sparsity = coloring_result.jacobian_sparsity, ad.tag)
         jac_prototype = coloring_result.jacobian_sparsity
     end
     return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x)
@@ -26,11 +28,12 @@ function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}
     sd::AbstractMaybeSparsityDetection, f!, fx, x)
     coloring_result = sd(ad, f!, fx, x)
     if coloring_result isa NoMatrixColoring
-        cache = ForwardDiff.JacobianConfig(f!, fx, x)
+        cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x),
+            ifelse(ad.tag === nothing, SparseDiffToolsTag(), ad.tag))
         jac_prototype = nothing
     else
         cache = ForwardColorJacCache(f!, x, __chunksize(ad); coloring_result.colorvec,
-            dx = fx, sparsity = coloring_result.jacobian_sparsity)
+            dx = fx, sparsity = coloring_result.jacobian_sparsity, ad.tag)
         jac_prototype = coloring_result.jacobian_sparsity
     end
     return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x)