StatProfilerHTML.jl report
Generated on Thu, 21 Dec 2023 13:06:16
File source code
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