|
8 | 8 |
|
9 | 9 | __getfield(c::ForwardDiffJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype |
10 | 10 |
|
11 | | -struct SparseDiffToolsTag end |
12 | | - |
13 | | -function ForwardDiff.checktag(::Type{<:ForwardDiff.Tag{<:SparseDiffToolsTag, <:T}}, f::F, |
14 | | - x::AbstractArray{T}) where {T, F} |
15 | | - return true |
16 | | -end |
17 | | - |
18 | | -__standard_tag(::Nothing, x) = ForwardDiff.Tag(SparseDiffToolsTag(), eltype(x)) |
19 | | -__standard_tag(tag::ForwardDiff.Tag, _) = tag |
20 | | -__standard_tag(tag, x) = ForwardDiff.Tag(tag, eltype(x)) |
| 11 | +__standard_tag(::Nothing, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x)) |
| 12 | +__standard_tag(tag::ForwardDiff.Tag, ::F, _) where {F} = tag |
| 13 | +__standard_tag(tag, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x)) |
21 | 14 |
|
22 | 15 | function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, |
23 | 16 | sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} |
24 | 17 | coloring_result = sd(ad, f, x) |
25 | 18 | fx = fx === nothing ? similar(f(x)) : fx |
26 | | - tag = __standard_tag(ad.tag, x) |
| 19 | + tag = __standard_tag(ad.tag, f, x) |
27 | 20 | if coloring_result isa NoMatrixColoring |
28 | 21 | cache = ForwardDiff.JacobianConfig(f, x, __chunksize(ad, x), tag) |
29 | 22 | jac_prototype = nothing |
|
39 | 32 | function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, |
40 | 33 | sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} |
41 | 34 | coloring_result = sd(ad, f!, fx, x) |
42 | | - tag = __standard_tag(ad.tag, x) |
| 35 | + tag = __standard_tag(ad.tag, f!, x) |
43 | 36 | if coloring_result isa NoMatrixColoring |
44 | 37 | cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x), tag) |
45 | 38 | jac_prototype = nothing |
|
0 commit comments