Skip to content

Commit 54ce5d5

Browse files
Avoid mutation on user x
Fixes #265
1 parent dc28eb2 commit 54ce5d5

File tree

5 files changed

+58
-53
lines changed

5 files changed

+58
-53
lines changed

README.md

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,7 @@ H = numauto_color_hessian(f, x, colorvec, sparsity)
282282
numauto_color_hessian!(H, f, x, colorvec, sparsity)
283283
```
284284

285-
To avoid unnecessary allocations every time the Hessian is computed,
285+
To avoid unnecessary allocations every time the Hessian is computed,
286286
construct a `ForwardColorHesCache` beforehand:
287287

288288
```julia
@@ -292,7 +292,7 @@ numauto_color_hessian!(H, f, x, hescache)
292292

293293
By default, these methods use a mix of numerical and automatic differentiation,
294294
namely by taking finite differences of gradients calculated via ForwardDiff.jl.
295-
Alternatively, if you have your own custom gradient function `g!`, you can specify
295+
Alternatively, if you have your own custom gradient function `g!`, you can specify
296296
it as an argument to `ForwardColorHesCache`:
297297

298298
```julia
@@ -301,7 +301,6 @@ hescache = ForwardColorHesCache(f, x, colorvec, sparsity, g!)
301301
Note that any user-defined gradient needs to have the signature `g!(G, x)`,
302302
i.e. updating the gradient `G` in place.
303303

304-
305304
### Jacobian-Vector and Hessian-Vector Products
306305

307306
Matrix-free implementations of Jacobian-Vector and Hessian-Vector products is
@@ -322,7 +321,8 @@ auto_jacvec(f, x, v)
322321

323322
# If compute_f0 is false, then `f(cache1,x)` will be computed
324323
num_jacvec!(dy,f,x,v,cache1 = similar(v),
325-
cache2 = similar(v);
324+
cache2 = similar(v),
325+
cache3 = similar(v);
326326
compute_f0 = true)
327327
num_jacvec(f,x,v,f0=nothing)
328328
```
@@ -333,14 +333,16 @@ For Hessians, the following are provided:
333333
num_hesvec!(dy,f,x,v,
334334
cache1 = similar(v),
335335
cache2 = similar(v),
336-
cache3 = similar(v))
336+
cache3 = similar(v),
337+
cache4 = similar(v))
337338

338339
num_hesvec(f,x,v)
339340

340341
numauto_hesvec!(dy,f,x,v,
341342
cache = ForwardDiff.GradientConfig(f,v),
342343
cache1 = similar(v),
343-
cache2 = similar(v))
344+
cache2 = similar(v),
345+
cache3 = similar(v))
344346

345347
numauto_hesvec(f,x,v)
346348

@@ -358,6 +360,7 @@ respectively:
358360

359361
```julia
360362
num_hesvecgrad!(dy,g,x,v,
363+
cache1 = similar(v),
361364
cache2 = similar(v),
362365
cache3 = similar(v))
363366

@@ -384,7 +387,8 @@ using Zygote # Required
384387

385388
numback_hesvec!(dy,f,x,v,
386389
cache1 = similar(v),
387-
cache2 = similar(v))
390+
cache2 = similar(v),
391+
cache3 = similar(v))
388392

389393
numback_hesvec(f,x,v)
390394

@@ -396,7 +400,7 @@ autoback_hesvec!(dy,f,x,v,
396400
autoback_hesvec(f,x,v)
397401
```
398402

399-
#### J*v and H*v Operators
403+
#### `J*v` and `H*v` Operators
400404

401405
The following produce matrix-free operators which are used for calculating
402406
Jacobian-vector and Hessian-vector products where the differentiation takes

docs/src/index.md

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
244244
`dx` is a pre-allocated output vector which is used to declare the output size,
245245
and thus allows for specifying a non-square Jacobian.
246246

247-
Also, it is possible retrieve the function value via `value(jac_cache)` or
247+
Also, it is possible retrieve the function value via `value(jac_cache)` or
248248
`value!(result, jac_cache)`
249249

250250

@@ -276,7 +276,7 @@ H = numauto_color_hessian(f, x, colorvec, sparsity)
276276
numauto_color_hessian!(H, f, x, colorvec, sparsity)
277277
```
278278

279-
To avoid unnecessary allocations every time the Hessian is computed,
279+
To avoid unnecessary allocations every time the Hessian is computed,
280280
construct a `ForwardColorHesCache` beforehand:
281281

282282
```julia
@@ -286,7 +286,7 @@ numauto_color_hessian!(H, f, x, hescache)
286286

287287
By default, these methods use a mix of numerical and automatic differentiation,
288288
namely by taking finite differences of gradients calculated via ForwardDiff.jl.
289-
Alternatively, if you have your own custom gradient function `g!`, you can specify
289+
Alternatively, if you have your own custom gradient function `g!`, you can specify
290290
it as an argument to `ForwardColorHesCache`:
291291

292292
```julia
@@ -295,7 +295,6 @@ hescache = ForwardColorHesCache(f, x, colorvec, sparsity, g!)
295295
Note that any user-defined gradient needs to have the signature `g!(G, x)`,
296296
i.e. updating the gradient `G` in place.
297297

298-
299298
### Jacobian-Vector and Hessian-Vector Products
300299

301300
Matrix-free implementations of Jacobian-Vector and Hessian-Vector products is
@@ -316,7 +315,8 @@ auto_jacvec(f, x, v)
316315

317316
# If compute_f0 is false, then `f(cache1,x)` will be computed
318317
num_jacvec!(dy,f,x,v,cache1 = similar(v),
319-
cache2 = similar(v);
318+
cache2 = similar(v),
319+
cache3 = similar(v);
320320
compute_f0 = true)
321321
num_jacvec(f,x,v,f0=nothing)
322322
```
@@ -327,14 +327,16 @@ For Hessians, the following are provided:
327327
num_hesvec!(dy,f,x,v,
328328
cache1 = similar(v),
329329
cache2 = similar(v),
330-
cache3 = similar(v))
330+
cache3 = similar(v),
331+
cache4 = similar(v))
331332

332333
num_hesvec(f,x,v)
333334

334335
numauto_hesvec!(dy,f,x,v,
335336
cache = ForwardDiff.GradientConfig(f,v),
336337
cache1 = similar(v),
337-
cache2 = similar(v))
338+
cache2 = similar(v),
339+
cache3 = similar(v))
338340

339341
numauto_hesvec(f,x,v)
340342

@@ -352,6 +354,7 @@ respectively:
352354

353355
```julia
354356
num_hesvecgrad!(dy,g,x,v,
357+
cache1 = similar(v),
355358
cache2 = similar(v),
356359
cache3 = similar(v))
357360

@@ -378,7 +381,8 @@ using Zygote # Required
378381

379382
numback_hesvec!(dy,f,x,v,
380383
cache1 = similar(v),
381-
cache2 = similar(v))
384+
cache2 = similar(v),
385+
cache3 = similar(v))
382386

383387
numback_hesvec(f,x,v)
384388

@@ -390,7 +394,7 @@ autoback_hesvec!(dy,f,x,v,
390394
autoback_hesvec(f,x,v)
391395
```
392396

393-
#### J*v and H*v Operators
397+
#### `J*v` and `H*v` Operators
394398

395399
The following produce matrix-free operators which are used for calculating
396400
Jacobian-vector and Hessian-vector products where the differentiation takes

ext/SparseDiffToolsZygoteExt.jl

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,18 +45,17 @@ end
4545

4646
### Jac, Hes products
4747

48-
function numback_hesvec!(dy, f::F, x, v, cache1 = similar(v), cache2 = similar(v)) where {F}
48+
function numback_hesvec!(dy, f::F, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v)) where {F}
4949
g = let f = f
5050
(dx, x) -> dx .= first(Zygote.gradient(f, x))
5151
end
5252
T = eltype(x)
5353
# Should it be min? max? mean?
5454
ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
55-
@. x += ϵ * v
56-
g(cache1, x)
57-
@. x -= 2ϵ * v
58-
g(cache2, x)
59-
@. x += ϵ * v
55+
@. cache3 = x + ϵ * v
56+
g(cache1, cache3)
57+
@. cache3 = x - ϵ * v
58+
g(cache2, cache3)
6059
@. dy = (cache1 - cache2) / (2ϵ)
6160
end
6261

src/differentiation/jaches_products.jl

Lines changed: 19 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -32,16 +32,15 @@ function auto_jacvec(f, x, v)
3232
vec(partials.(vec(f(y)), 1))
3333
end
3434

35-
function num_jacvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v);
35+
function num_jacvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v);
3636
compute_f0 = true)
3737
vv = reshape(v, axes(x))
3838
compute_f0 && (f(cache1, x))
3939
T = eltype(x)
4040
# Should it be min? max? mean?
4141
ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
42-
@. x += ϵ * vv
43-
f(cache2, x)
44-
@. x -= ϵ * vv
42+
@. cache3 = x + ϵ * vv
43+
f(cache2, cache3)
4544
vecdy = _vec(dy)
4645
veccache1 = _vec(cache1)
4746
veccache2 = _vec(cache2)
@@ -58,19 +57,18 @@ function num_jacvec(f, x, v, f0 = nothing)
5857
end
5958

6059
function num_hesvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v),
61-
cache3 = similar(v))
60+
cache3 = similar(v), cache4 = similar(v))
6261
cache = FiniteDiff.GradientCache(v[1], cache1, Val{:central})
6362
g = let f = f, cache = cache
6463
(dx, x) -> FiniteDiff.finite_difference_gradient!(dx, f, x, cache)
6564
end
6665
T = eltype(x)
6766
# Should it be min? max? mean?
6867
ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
69-
@. x += ϵ * v
70-
g(cache2, x)
71-
@. x -= 2ϵ * v
72-
g(cache3, x)
73-
@. x += ϵ * v
68+
@. cache4 = x + ϵ * v
69+
g(cache2, cache4)
70+
@. cache4 = x - ϵ * v
71+
g(cache3, cache4)
7472
@. dy = (cache2 - cache3) / (2ϵ)
7573
end
7674

@@ -87,18 +85,17 @@ function num_hesvec(f, x, v)
8785
end
8886

8987
function numauto_hesvec!(dy, f, x, v, cache = ForwardDiff.GradientConfig(f, v),
90-
cache1 = similar(v), cache2 = similar(v))
88+
cache1 = similar(v), cache2 = similar(v), cache3 = similar(v))
9189
g = let f = f, x = x, cache = cache
9290
g = (dx, x) -> ForwardDiff.gradient!(dx, f, x, cache)
9391
end
9492
T = eltype(x)
9593
# Should it be min? max? mean?
9694
ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
97-
@. x += ϵ * v
98-
g(cache1, x)
99-
@. x -= 2ϵ * v
100-
g(cache2, x)
101-
@. x += ϵ * v
95+
@. cache3 = x + ϵ * v
96+
g(cache1, cache3)
97+
@. cache3 = x - ϵ * v
98+
g(cache2, cache3)
10299
@. dy = (cache1 - cache2) / (2ϵ)
103100
end
104101

@@ -137,16 +134,15 @@ function autonum_hesvec(f, x, v)
137134
partials.(g(Dual{DeivVecTag}.(x, v)), 1)
138135
end
139136

140-
function num_hesvecgrad!(dy, g, x, v, cache2 = similar(v), cache3 = similar(v))
137+
function num_hesvecgrad!(dy, g, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v))
141138
T = eltype(x)
142139
# Should it be min? max? mean?
143140
ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
144-
@. x += ϵ * v
145-
g(cache2, x)
146-
@. x -= 2ϵ * v
147-
g(cache3, x)
148-
@. x += ϵ * v
149-
@. dy = (cache2 - cache3) / (2ϵ)
141+
@. cache3 = x + ϵ * v
142+
g(cache1, cache3)
143+
@. cache3 = x - ϵ * v
144+
g(cache2, cache3)
145+
@. dy = (cache1 - cache2) / (2ϵ)
150146
end
151147

152148
function num_hesvecgrad(g, x, v)

src/differentiation/vecjac_products.jl

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
1-
function num_vecjac!(du, f::F, x, v, cache1 = similar(v), cache2 = similar(v);
1+
function num_vecjac!(du, f::F, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v);
22
compute_f0 = true) where {F}
33
compute_f0 && (f(cache1, x))
44
T = eltype(x)
55
# Should it be min? max? mean?
66
ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
77
vv = reshape(v, size(cache1))
8+
cache3 .= x
89
for i in 1:length(x)
9-
x[i] += ϵ
10-
f(cache2, x)
11-
x[i] -= ϵ
10+
cache3[i] += ϵ
11+
f(cache2, cache3)
12+
cache3[i] = x[i]
1213
du[i] = (((cache2 .- cache1) ./ ϵ)' * vv)[1]
1314
end
1415
return du
@@ -21,10 +22,11 @@ function num_vecjac(f::F, x, v, f0 = nothing) where {F}
2122
# Should it be min? max? mean?
2223
ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
2324
du = similar(x)
25+
cache = copy(x)
2426
for i in 1:length(x)
25-
x[i] += ϵ
27+
cache[i] += ϵ
2628
f0 = f(x)
27-
x[i] -= ϵ
29+
cache[i] = x[i]
2830
du[i] = (((f0 .- _f0) ./ ϵ)' * vv)[1]
2931
end
3032
return vec(du)
@@ -91,7 +93,7 @@ function VecJac(f, u::AbstractArray, p = nothing, t = nothing; fu = nothing,
9193
end
9294

9395
function _vecjac(f::F, fu, u, autodiff::AutoFiniteDiff) where {F}
94-
cache = (similar(fu), similar(fu))
96+
cache = (similar(fu), similar(fu), similar(fu))
9597
pullback = nothing
9698
return AutoDiffVJP(f, u, cache, autodiff, pullback)
9799
end

0 commit comments

Comments
 (0)