From ad2b1506124a79fa844437e38c642e02c696c2b2 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Fri, 27 Oct 2023 17:34:21 -0400 Subject: [PATCH] Pass Tag and Chunksize from the types --- Project.toml | 2 +- src/highlevel/common.jl | 5 +++++ src/highlevel/forward_mode.jl | 13 ++++++++----- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/Project.toml b/Project.toml index bc720e6a..4fca1f27 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseDiffTools" uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" authors = ["Pankaj Mishra ", "Chris Rackauckas "] -version = "2.8.0" +version = "2.9.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index e8582058..318ee251 100644 --- a/src/highlevel/common.jl +++ b/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) diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl index d073e323..a563bc68 100644 --- a/src/highlevel/forward_mode.jl +++ b/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)