StatProfilerHTML.jl report
Generated on Thu, 21 Dec 2023 13:06:16
File source code
Line Exclusive Inclusive Code
1 const AbstractSparseADType = Union{AbstractSparseForwardMode, AbstractSparseReverseMode,
2 AbstractSparseFiniteDifferences}
3
4 struct AutoSparseEnzyme <: AbstractSparseReverseMode end
5
6 # Sparsity Detection
7 abstract type AbstractMaybeSparsityDetection end
8 abstract type AbstractSparsityDetection <: AbstractMaybeSparsityDetection end
9
10 struct NoSparsityDetection <: AbstractMaybeSparsityDetection end
11
12 """
13 SymbolicsSparsityDetection(; alg = GreedyD1Color())
14
15 Use Symbolics to compute the sparsity pattern of the Jacobian. This requires `Symbolics.jl`
16 to be explicitly loaded.
17
18 ## Keyword Arguments
19
20 - `alg`: The algorithm used for computing the matrix colors
21
22 See Also: [JacPrototypeSparsityDetection](@ref), [PrecomputedJacobianColorvec](@ref)
23 """
24 Base.@kwdef struct SymbolicsSparsityDetection{A <: ArrayInterface.ColoringAlgorithm} <:
25 AbstractSparsityDetection
26 alg::A = GreedyD1Color()
27 end
28
29 """
30 JacPrototypeSparsityDetection(; jac_prototype, alg = GreedyD1Color())
31
32 Use a pre-specified `jac_prototype` to compute the matrix colors of the Jacobian.
33
34 ## Keyword Arguments
35
36 - `jac_prototype`: The prototype Jacobian used for computing the matrix colors
37 - `alg`: The algorithm used for computing the matrix colors
38
39 See Also: [SymbolicsSparsityDetection](@ref), [PrecomputedJacobianColorvec](@ref)
40 """
41 Base.@kwdef struct JacPrototypeSparsityDetection{
42 J, A <: ArrayInterface.ColoringAlgorithm} <: AbstractSparsityDetection
43 jac_prototype::J
44 alg::A = GreedyD1Color()
45 end
46
47 """
48 PrecomputedJacobianColorvec(jac_prototype, row_colorvec, col_colorvec)
49
50 Use a pre-specified `colorvec` which can be directly used for sparse differentiation. Based
51 on whether a reverse mode or forward mode or finite differences is used, the corresponding
52 `row_colorvec` or `col_colorvec` is used. Atmost one of them can be set to `nothing`.
53
54 ## Arguments
55
56 - `jac_prototype`: The prototype Jacobian used for computing structural nonzeros
57 - `row_colorvec`: The row colorvec of the Jacobian
58 - `col_colorvec`: The column colorvec of the Jacobian
59
60 See Also: [SymbolicsSparsityDetection](@ref), [JacPrototypeSparsityDetection](@ref)
61 """
62 struct PrecomputedJacobianColorvec{J, RC, CC} <: AbstractSparsityDetection
63 jac_prototype::J
64 row_colorvec::RC
65 col_colorvec::CC
66 end
67
68 """
69 PrecomputedJacobianColorvec(; jac_prototype, partition_by_rows::Bool = false,
70 colorvec = missing, row_colorvec = missing, col_colorvec = missing)
71
72 Use a pre-specified `colorvec` which can be directly used for sparse differentiation. Based
73 on whether a reverse mode or forward mode or finite differences is used, the corresponding
74 `row_colorvec` or `col_colorvec` is used. Atmost one of them can be set to `nothing`.
75
76 ## Keyword Arguments
77
78 - `jac_prototype`: The prototype Jacobian used for computing structural nonzeros
79 - `partition_by_rows`: Whether to partition the Jacobian by rows or columns (row
80 partitioning is used for reverse mode AD)
81 - `colorvec`: The colorvec of the Jacobian. If `partition_by_rows` is `true` then this
82 is the row colorvec, otherwise it is the column colorvec
83 - `row_colorvec`: The row colorvec of the Jacobian
84 - `col_colorvec`: The column colorvec of the Jacobian
85
86 See Also: [SymbolicsSparsityDetection](@ref), [JacPrototypeSparsityDetection](@ref)
87 """
88 function PrecomputedJacobianColorvec(; jac_prototype, partition_by_rows::Bool = false,
89 colorvec = missing, row_colorvec = missing, col_colorvec = missing)
90 if colorvec === missing
91 @assert row_colorvec !== missing||col_colorvec !== missing "Either `colorvec` or `row_colorvec` and `col_colorvec` must be specified!"
92 row_colorvec = row_colorvec === missing ? nothing : row_colorvec
93 col_colorvec = col_colorvec === missing ? nothing : col_colorvec
94 return PrecomputedJacobianColorvec(jac_prototype, row_colorvec, col_colorvec)
95 else
96 @assert row_colorvec === missing&&col_colorvec === missing "Specifying `colorvec` is incompatible with specifying `row_colorvec` or `col_colorvec`!"
97 row_colorvec = partition_by_rows ? colorvec : nothing
98 col_colorvec = partition_by_rows ? nothing : colorvec
99 return PrecomputedJacobianColorvec(jac_prototype, row_colorvec, col_colorvec)
100 end
101 end
102
103 function _get_colorvec(alg::PrecomputedJacobianColorvec, ad)
104 cvec = alg.col_colorvec
105 @assert cvec!==nothing "`col_colorvec` is nothing, but Forward Mode AD or Finite Differences is being used!"
106 return cvec
107 end
108
109 function _get_colorvec(alg::PrecomputedJacobianColorvec, ::AbstractReverseMode)
110 cvec = alg.row_colorvec
111 @assert cvec!==nothing "`row_colorvec` is nothing, but Reverse Mode AD is being used!"
112 return cvec
113 end
114
115 """
116 ApproximateJacobianSparsity(; ntrials = 5, rng = Random.default_rng(),
117 epsilon = nothing, alg = GreedyD1Color())
118
119 Use `ntrials` random vectors to compute the sparsity pattern of the Jacobian. This is an
120 approximate method and the sparsity pattern may not be exact.
121
122 ## Keyword Arguments
123
124 - `ntrials`: The number of random vectors to use for computing the sparsity pattern
125 - `rng`: The random number generator used for generating the random vectors
126 - `alg`: The algorithm used for computing the matrix colors
127 - `epsilon`: For Finite Differencing based Jacobian Approximations, any number smaller
128 than `epsilon` is considered to be zero. If `nothing` is specified, then this value
129 is calculated as `100 * eps(eltype(x))`
130 """
131 struct ApproximateJacobianSparsity{R <: AbstractRNG, A <: ArrayInterface.ColoringAlgorithm,
132 E} <: AbstractSparsityDetection
133 ntrials::Int
134 rng::R
135 alg::A
136 epsilon::E
137 end
138
139 function ApproximateJacobianSparsity(; ntrials::Int = 3, epsilon = nothing,
140 rng::AbstractRNG = Random.default_rng(), alg = GreedyD1Color())
141 return ApproximateJacobianSparsity(ntrials, rng, alg, epsilon)
142 end
143
144 # No one should be using this currently
145 Base.@kwdef struct AutoSparsityDetection{A <: ArrayInterface.ColoringAlgorithm} <:
146 AbstractSparsityDetection
147 alg::A = GreedyD1Color()
148 end
149
150 # Function Specifications
151 abstract type AbstractMaybeSparseJacobianCache end
152
153 """
154 sparse_jacobian!(J::AbstractMatrix, ad, cache::AbstractMaybeSparseJacobianCache, f, x)
155 sparse_jacobian!(J::AbstractMatrix, ad, cache::AbstractMaybeSparseJacobianCache, f!, fx,
156 x)
157
158 Inplace update the matrix `J` with the Jacobian of `f` at `x` using the AD backend `ad`.
159
160 `cache` is the cache object returned by `sparse_jacobian_cache`.
161 """
162 function sparse_jacobian! end
163
164 """
165 sparse_jacobian_cache(ad::AbstractADType, sd::AbstractSparsityDetection, f, x;
166 fx=nothing)
167 sparse_jacobian_cache(ad::AbstractADType, sd::AbstractSparsityDetection, f!, fx, x)
168
169 Takes the underlying AD backend `ad`, sparsity detection algorithm `sd`, function `f`,
170 and input `x` and returns a cache object that can be used to compute the Jacobian.
171
172 If `fx` is not specified, it will be computed by calling `f(x)`.
173
174 ## Returns
175
176 A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`.
177 """
178 function sparse_jacobian_cache end
179
180 function sparse_jacobian_static_array(ad, cache, f, x::SArray)
181 # Not the most performant fallback
182 J = init_jacobian(cache)
183 sparse_jacobian!(J, ad, cache, f, MArray(x))
184 return J
185 end
186
187 """
188 sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f, x; fx=nothing)
189 sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f!, fx, x)
190
191 Sequentially calls `sparse_jacobian_cache` and `sparse_jacobian!` to compute the Jacobian of
192 `f` at `x`. Use this if the jacobian for `f` is computed exactly once. In all other
193 cases, use `sparse_jacobian_cache` once to generate the cache and use `sparse_jacobian!`
194 with the same cache to compute the jacobian.
195
196 If `x` is a StaticArray, then this function tries to use a non-allocating implementation for
197 the jacobian computation. This is possible only for a limited backends currently.
198 """
199 function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, args...;
200 kwargs...)
201 cache = sparse_jacobian_cache(ad, sd, args...; kwargs...)
202 J = init_jacobian(cache)
203 sparse_jacobian!(J, ad, cache, args...)
204 return J
205 end
206 function sparse_jacobian(ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f,
207 x::SArray; kwargs...)
208 cache = sparse_jacobian_cache(ad, sd, f, x; kwargs...)
209 return sparse_jacobian_static_array(ad, cache, f, x)
210 end
211
212 """
213 sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f, x)
214 sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f!, fx, x)
215
216 Use the sparsity detection `cache` for computing the sparse Jacobian. This allocates a new
217 Jacobian at every function call.
218
219 If `x` is a StaticArray, then this function tries to use a non-allocating implementation for
220 the jacobian computation. This is possible only for a limited backends currently.
221 """
222 function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache,
223 args...)
224 J = init_jacobian(cache)
225 sparse_jacobian!(J, ad, cache, args...)
226 return J
227 end
228 function sparse_jacobian(ad::AbstractADType, cache::AbstractMaybeSparseJacobianCache, f,
229 x::SArray)
230 return sparse_jacobian_static_array(ad, cache, f, x)
231 end
232
233 """
234 sparse_jacobian!(J::AbstractMatrix, ad::AbstractADType, sd::AbstractSparsityDetection,
235 f, x; fx=nothing)
236 sparse_jacobian!(J::AbstractMatrix, ad::AbstractADType, sd::AbstractSparsityDetection,
237 f!, fx, x)
238
239 Sequentially calls `sparse_jacobian_cache` and `sparse_jacobian!` to compute the Jacobian of
240 `f` at `x`. Use this if the jacobian for `f` is computed exactly once. In all other
241 cases, use `sparse_jacobian_cache` once to generate the cache and use `sparse_jacobian!`
242 with the same cache to compute the jacobian.
243 """
244 function sparse_jacobian!(J::AbstractMatrix, ad::AbstractADType,
245 sd::AbstractMaybeSparsityDetection, args...; kwargs...)
246 cache = sparse_jacobian_cache(ad, sd, args...; kwargs...)
247 sparse_jacobian!(J, ad, cache, args...)
248 return J
249 end
250
251 ## Internal
252 function __gradient end
253 function __gradient! end
254 function __jacobian! end
255
256 """
257 init_jacobian(cache::AbstractMaybeSparseJacobianCache)
258
259 Initialize the Jacobian based on the cache. Uses sparse jacobians if possible.
260
261 !!! note
262 This function doesn't alias the provided jacobian prototype. It always initializes a
263 fresh jacobian that can be mutated without any side effects.
264 """
265 function init_jacobian end
266
267 # Never thought this was a useful function externally, but I ended up using it in quite a
268 # few places. Keeping this till I remove uses of those.
269 const __init_𝒥 = init_jacobian
270
271 # Misc Functions
272 function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}, x) where {C}
273 C isa ForwardDiff.Chunk && return C
274 return __chunksize(Val(C), x)
275 end
276 __chunksize(::Val{nothing}, x) = __chunksize(x)
277 function __chunksize(::Val{C}, x) where {C}
278 if C isa Integer && !(C isa Bool)
279 return C ≤ 0 ? __chunksize(x) : ForwardDiff.Chunk{C}()
280 else
281 error("$(C)::$(typeof(C)) is not a valid chunksize!")
282 end
283 end
284
285 __chunksize(x) = ForwardDiff.Chunk(x)
286 __chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}()
287
288 function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C}
289 C === nothing && return nothing
290 C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C)
291 return nothing
292 end
293
294 __f̂(f, x, idxs) = dot(vec(f(x)), idxs)
295
296 function __f̂(f!, fx, x, idxs)
297 f!(fx, x)
298 return dot(vec(fx), idxs)
299 end
300
301 @generated function __getfield(c::T, ::Val{S}) where {T, S}
302 hasfield(T, S) && return :(c.$(S))
303 return :(nothing)
304 end
305
306 """
307 init_jacobian(cache::AbstractMaybeSparseJacobianCache;
308 preserve_immutable::Val = Val(false))
309
310 Initialize the Jacobian based on the cache. Uses sparse jacobians if possible.
311
312 If `preserve_immutable` is `true`, then the Jacobian returned might be immutable, this is
313 relevant if the inputs are immutable like `StaticArrays`.
314 """
315 7 (3 %)
7 (3 %) samples spent in init_jacobian
7 (100 %) (incl.) when called from #jacobian_caches#82 line 126
7 (100 %) samples spent calling #init_jacobian#86
function init_jacobian(c::AbstractMaybeSparseJacobianCache;
316 preserve_immutable::Val = Val(false))
317 T = promote_type(eltype(c.fx), eltype(c.x))
318 7 (3 %)
7 (3 %) samples spent in #init_jacobian#86
7 (100 %) (incl.) when called from init_jacobian line 315
7 (100 %) samples spent calling init_jacobian
return init_jacobian(__getfield(c, Val(:jac_prototype)), T, c.fx, c.x;
319 preserve_immutable)
320 end
321 7 (3 %)
7 (3 %) samples spent in init_jacobian
7 (100 %) (incl.) when called from #init_jacobian#86 line 318
7 (100 %) samples spent calling #init_jacobian#87
function init_jacobian(::Nothing, ::Type{T}, fx, x; kwargs...) where {T}
322 7 (3 %)
7 (3 %) samples spent in #init_jacobian#87
7 (100 %) (incl.) when called from init_jacobian line 321
7 (100 %) samples spent calling similar
return similar(fx, T, length(fx), length(x))
323 end
324 function init_jacobian(::Nothing, ::Type{T}, fx::StaticArray, x::StaticArray;
325 preserve_immutable::Val{PI} = Val(true)) where {T, PI}
326 if PI
327 return StaticArrays.SArray{Tuple{length(fx), length(x)}, T}(I)
328 else
329 return StaticArrays.MArray{Tuple{length(fx), length(x)}, T}(undef)
330 end
331 end
332 function init_jacobian(J, ::Type{T}, fx, x; kwargs...) where {T}
333 return similar(J, T, size(J, 1), size(J, 2))
334 end
335 init_jacobian(J::SparseMatrixCSC, ::Type{T}, fx, x; kwargs...) where {T} = T.(J)
336
337 __maybe_copy_x(_, x) = x
338 __maybe_copy_x(_, ::Nothing) = nothing
339
340 # Check Backend has been loaded
341 ## We pay a small compile time cost for this, but it avoids cryptic error messages
342 @inline function __test_backend_loaded(ad::AbstractADType)
343 error("$(ad) requires $(__backend(ad)).jl to be loaded. Please load it.")
344 end
345
346 @inline __backend(ad) = nothing
347 @inline __backend(::Union{AutoEnzyme, AutoSparseEnzyme}) = :Enzyme
348 @inline __backend(::Union{AutoZygote, AutoSparseZygote}) = :Zygote
349 @inline __backend(::Union{AutoForwardDiff, AutoSparseForwardDiff}) = :ForwardDiff
350 @inline __backend(::Union{AutoFiniteDiff, AutoSparseFiniteDiff}) = :FiniteDiff