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 |