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 (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 (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!
function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) where {T}
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 |
||
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 %) | $(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!
@eval function chunk_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
96 (100 %) (incl.) when called from jacobian! line 80 |
||
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 |