StatProfilerHTML.jl report
Generated on Thu, 21 Dec 2023 12:59:22
File source code
Line Exclusive Inclusive Code
1 ###############
2 # API methods #
3 ###############
4
5 """
6 ForwardDiff.jacobian(f, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f, x), check=Val{true}())
7
8 Return `J(f)` evaluated at `x`, assuming `f` is called as `f(x)`.
9 Multidimensional arrays are flattened in iteration order: the array
10 `J(f)` has shape `length(f(x)) × length(x)`, and its elements are
11 `J(f)[j,k] = ∂f(x)[j]/∂x[k]`. When `x` is a vector, this means
12 that `jacobian(x->[f(x)], x)` is the transpose of `gradient(f, x)`.
13
14 This method assumes that `isa(f(x), AbstractArray)`.
15
16 Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
17 """
18 function jacobian(f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {T,CHK}
19 CHK && checktag(T, f, x)
20 if chunksize(cfg) == length(x)
21 return vector_mode_jacobian(f, x, cfg)
22 else
23 return chunk_mode_jacobian(f, x, cfg)
24 end
25 end
26
27 """
28 ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())
29
30 Return `J(f!)` evaluated at `x`, assuming `f!` is called as `f!(y, x)` where the result is
31 stored in `y`.
32
33 Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
34 """
35 function jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK}
36 CHK && checktag(T, f!, x)
37 if chunksize(cfg) == length(x)
38 return vector_mode_jacobian(f!, y, x, cfg)
39 else
40 return chunk_mode_jacobian(f!, y, x, cfg)
41 end
42 end
43
44
45 """
46 ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f, x), check=Val{true}())
47
48 Compute `J(f)` evaluated at `x` and store the result(s) in `result`, assuming `f` is called
49 as `f(x)`.
50
51 This method assumes that `isa(f(x), AbstractArray)`.
52
53 Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
54 """
55 function jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK}
56 CHK && checktag(T, f, x)
57 if chunksize(cfg) == length(x)
58 vector_mode_jacobian!(result, f, x, cfg)
59 else
60 chunk_mode_jacobian!(result, f, x, cfg)
61 end
62 return result
63 end
64
65 """
66 ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())
67
68 Compute `J(f!)` evaluated at `x` and store the result(s) in `result`, assuming `f!` is
69 called as `f!(y, x)` where the result is stored in `y`.
70
71 This method assumes that `isa(f(x), AbstractArray)`.
72
73 Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
74 """
75 function jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T,CHK}
76 96 (33 %)
96 (33 %) samples spent in jacobian!
96 (100 %) (incl.) when called from sparse_jacobian! line 70
96 (100 %) samples spent calling jacobian!
CHK && checktag(T, f!, x)
77 if chunksize(cfg) == length(x)
78 vector_mode_jacobian!(result, f!, y, x, cfg)
79 else
80 96 (33 %)
96 (33 %) samples spent in jacobian!
96 (100 %) (incl.) when called from jacobian! line 76
96 (100 %) samples spent calling chunk_mode_jacobian!
chunk_mode_jacobian!(result, f!, y, x, cfg)
81 end
82 return result
83 end
84
85 jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is an array. Perhaps you meant derivative(f, x)?"))
86
87 #####################
88 # result extraction #
89 #####################
90
91 function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T}
92 out_reshaped = reshape(result, length(ydual), n)
93 ydual_reshaped = vec(ydual)
94 # Use closure to avoid GPU broadcasting with Type
95 partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
96 out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n))
97 return result
98 end
99
100 function extract_jacobian!(::Type{T}, result::MutableDiffResult, ydual::AbstractArray, n) where {T}
101 extract_jacobian!(T, DiffResults.jacobian(result), ydual, n)
102 return result
103 end
104
105
69 (24 %) samples spent in extract_jacobian_chunk!
67 (97 %) (incl.) when called from chunk_mode_jacobian! line 195
1 (1 %) (incl.) when called from chunk_mode_jacobian! line 202
1 (1 %) (incl.) when called from chunk_mode_jacobian! line 187
function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) where {T}
106 ydual_reshaped = vec(ydual)
107 offset = index - 1
108 irange = 1:chunksize
109 col = irange .+ offset
110 # Use closure to avoid GPU broadcasting with Type
111 partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
112 69 (24 %)
69 (100 %) samples spent calling materialize!
result[:, col] .= partials_wrap.(ydual_reshaped, transpose(irange))
113 return result
114 end
115
116 reshape_jacobian(result, ydual, xdual) = reshape(result, length(ydual), length(xdual))
117 reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResults.jacobian(result), ydual, xdual)
118
119 ###############
120 # vector mode #
121 ###############
122
123 function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T}) where {F,T}
124 N = chunksize(cfg)
125 ydual = vector_mode_dual_eval!(f, cfg, x)
126 ydual isa AbstractArray || throw(JACOBIAN_ERROR)
127 result = similar(ydual, valtype(eltype(ydual)), length(ydual), N)
128 extract_jacobian!(T, result, ydual, N)
129 extract_value!(T, result, ydual)
130 return result
131 end
132
133 function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
134 N = chunksize(cfg)
135 ydual = vector_mode_dual_eval!(f!, cfg, y, x)
136 map!(d -> value(T,d), y, ydual)
137 result = similar(y, length(y), N)
138 extract_jacobian!(T, result, ydual, N)
139 map!(d -> value(T,d), y, ydual)
140 return result
141 end
142
143 function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T}) where {F,T}
144 N = chunksize(cfg)
145 ydual = vector_mode_dual_eval!(f, cfg, x)
146 extract_jacobian!(T, result, ydual, N)
147 extract_value!(T, result, ydual)
148 return result
149 end
150
151 function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
152 N = chunksize(cfg)
153 ydual = vector_mode_dual_eval!(f!, cfg, y, x)
154 map!(d -> value(T,d), y, ydual)
155 extract_jacobian!(T, result, ydual, N)
156 extract_value!(T, result, y, ydual)
157 return result
158 end
159
160 const JACOBIAN_ERROR = DimensionMismatch("jacobian(f, x) expects that f(x) is an array. Perhaps you meant gradient(f, x)?")
161
162 # chunk mode #
163 #------------#
164
165 function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Expr,
166 result_definition::Expr, y_definition::Expr)
167 return quote
168 @assert length(x) >= N "chunk size cannot be greater than length(x) ($(N) > $(length(x)))"
169
170 # precalculate loop bounds
171 xlen = length(x)
172 remainder = xlen % N
173 lastchunksize = ifelse(remainder == 0, N, remainder)
174 lastchunkindex = xlen - lastchunksize + 1
175 middlechunks = 2:div(xlen - lastchunksize, N)
176
177 # seed work arrays
178 $(work_array_definition)
179 seeds = cfg.seeds
180
181 # do first chunk manually to calculate output type
182 seed!(xdual, x, 1, seeds)
183 $(compute_ydual)
184 ydual isa AbstractArray || throw(JACOBIAN_ERROR)
185 $(result_definition)
186 out_reshaped = reshape_jacobian(result, ydual, xdual)
187 1 (0 %)
1 (100 %) samples spent calling extract_jacobian_chunk!
extract_jacobian_chunk!(T, out_reshaped, ydual, 1, N)
188 seed!(xdual, x, 1)
189
190 # do middle chunks
191 for c in middlechunks
192 i = ((c - 1) * N + 1)
193 seed!(xdual, x, i, seeds)
194 27 (9 %)
26 (96 %) samples spent calling JacobianWrapper
1 (4 %) samples spent calling seed!
$(compute_ydual)
195 67 (23 %)
67 (100 %) samples spent calling extract_jacobian_chunk!
extract_jacobian_chunk!(T, out_reshaped, ydual, i, N)
196 seed!(xdual, x, i)
197 end
198
199 # do final chunk
200 seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
201 $(compute_ydual)
202 1 (0 %)
1 (100 %) samples spent calling extract_jacobian_chunk!
extract_jacobian_chunk!(T, out_reshaped, ydual, lastchunkindex, lastchunksize)
203
204 $(y_definition)
205
206 return result
207 end
208 end
209
210 @eval function chunk_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
211 $(jacobian_chunk_mode_expr(quote
212 xdual = cfg.duals
213 seed!(xdual, x)
214 end,
215 :(ydual = f(xdual)),
216 :(result = similar(ydual, valtype(eltype(ydual)), length(ydual), xlen)),
217 :()))
218 end
219
220 @eval function chunk_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
221 $(jacobian_chunk_mode_expr(quote
222 ydual, xdual = cfg.duals
223 seed!(xdual, x)
224 end,
225 :(f!(seed!(ydual, y), xdual)),
226 :(result = similar(y, length(y), xlen)),
227 :(map!(d -> value(T,d), y, ydual))))
228 end
229
230 @eval function chunk_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
231 $(jacobian_chunk_mode_expr(quote
232 xdual = cfg.duals
233 seed!(xdual, x)
234 end,
235 :(ydual = f(xdual)),
236 :(),
237 :(extract_value!(T, result, ydual))))
238 end
239
240
96 (33 %) samples spent in chunk_mode_jacobian!
96 (100 %) (incl.) when called from jacobian! line 80
@eval function chunk_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
241 $(jacobian_chunk_mode_expr(quote
242 ydual, xdual = cfg.duals
243 seed!(xdual, x)
244 end,
245 :(f!(seed!(ydual, y), xdual)),
246 :(),
247 :(extract_value!(T, result, y, ydual))))
248 end