Line | Exclusive | Inclusive | Code |
---|---|---|---|
1 | struct ForwardDiffJacobianCache{CO, CA, J, FX, X} <: AbstractMaybeSparseJacobianCache | ||
2 | coloring::CO | ||
3 | cache::CA | ||
4 | jac_prototype::J | ||
5 | fx::FX | ||
6 | x::X | ||
7 | end | ||
8 | |||
9 | __getfield(c::ForwardDiffJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype | ||
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)) | ||
21 | |||
22 | function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, | ||
23 | sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} | ||
24 | coloring_result = sd(ad, f, x) | ||
25 | fx = fx === nothing ? similar(f(x)) : fx | ||
26 | tag = __standard_tag(ad.tag, x) | ||
27 | if coloring_result isa NoMatrixColoring | ||
28 | cache = ForwardDiff.JacobianConfig(f, x, __chunksize(ad, x), tag) | ||
29 | jac_prototype = nothing | ||
30 | else | ||
31 | # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof` | ||
32 | cache = ForwardColorJacCache(f, x, __chunksize(ad); coloring_result.colorvec, | ||
33 | dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag)) | ||
34 | jac_prototype = coloring_result.jacobian_sparsity | ||
35 | end | ||
36 | return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) | ||
37 | end | ||
38 | |||
39 | function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, | ||
40 | sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} | ||
41 | coloring_result = sd(ad, f!, fx, x) | ||
42 | tag = __standard_tag(ad.tag, x) | ||
43 | if coloring_result isa NoMatrixColoring | ||
44 | cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x), tag) | ||
45 | jac_prototype = nothing | ||
46 | else | ||
47 | # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof` | ||
48 | cache = ForwardColorJacCache(f!, x, __chunksize(ad); coloring_result.colorvec, | ||
49 | dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag)) | ||
50 | jac_prototype = coloring_result.jacobian_sparsity | ||
51 | end | ||
52 | return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) | ||
53 | end | ||
54 | |||
55 | function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f::F, | ||
56 | x) where {F} | ||
57 | if cache.cache isa ForwardColorJacCache | ||
58 | forwarddiff_color_jacobian(J, f, x, cache.cache) # Use Sparse ForwardDiff | ||
59 | else | ||
60 | ForwardDiff.jacobian!(J, f, x, cache.cache) # Don't try to exploit sparsity | ||
61 | end | ||
62 | return J | ||
63 | end | ||
64 | |||
65 | function sparse_jacobian!(J::AbstractMatrix, _, cache::ForwardDiffJacobianCache, f!::F, fx, | ||
66 | x) where {F} | ||
67 | if cache.cache isa ForwardColorJacCache | ||
68 | forwarddiff_color_jacobian!(J, f!, x, cache.cache) # Use Sparse ForwardDiff | ||
69 | else | ||
70 | 113 (42 %) |
113 (42 %)
samples spent in sparse_jacobian!
113 (100 %) (incl.) when called from jacobian!! line 61
113 (100 %)
samples spent calling
jacobian!
ForwardDiff.jacobian!(J, f!, fx, x, cache.cache) # Don't try to exploit sparsity
|
|
71 | end | ||
72 | return J | ||
73 | end | ||
74 | |||
75 | function sparse_jacobian_static_array(_, cache::ForwardDiffJacobianCache, f, x::SArray) | ||
76 | if cache.cache isa ForwardColorJacCache | ||
77 | return forwarddiff_color_jacobian(f, x, cache.cache) | ||
78 | else | ||
79 | return ForwardDiff.jacobian(f, x, cache.cache) | ||
80 | end | ||
81 | end |