Skip to content

Commit 438fe5f

Browse files
topolaritytopolarity
committed
Support tag = :small in Config(...) constructors
This provides a convenient interface to ask for a SmallTag.
1 parent a515f5a commit 438fe5f

File tree

4 files changed

+148
-43
lines changed

4 files changed

+148
-43
lines changed

‎src/config.jl

Lines changed: 118 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,18 @@ Base.eltype(cfg::AbstractConfig) = eltype(typeof(cfg))
9999

100100
@inline (chunksize(::AbstractConfig{N})::Int) where {N} = N
101101

102+
function maketag(kind::Union{Symbol,Nothing}, f, X)
103+
if kind === :default
104+
return Tag(f, X)
105+
elseif kind === :small
106+
return SmallTag(f, X)
107+
elseif kind === nothing
108+
return nothing
109+
else
110+
throw(ArgumentError("tag may be :default, :small, or nothing"))
111+
end
112+
end
113+
102114
####################
103115
# DerivativeConfig #
104116
####################
@@ -108,7 +120,7 @@ struct DerivativeConfig{T,D} <: AbstractConfig{1}
108120
end
109121

110122
"""
111-
ForwardDiff.DerivativeConfig(f!, y::AbstractArray, x::Real)
123+
ForwardDiff.DerivativeConfig(f!, y::AbstractArray, x::Real; tag::Union{Symbol,Nothing} = :default)
112124
113125
Return a `DerivativeConfig` instance based on the type of `f!`, and the types/shapes of the
114126
output vector `y` and the input value `x`.
@@ -121,12 +133,24 @@ If `f!` is `nothing` instead of the actual target function, then the returned in
121133
be used with any target function. However, this will reduce ForwardDiff's ability to catch
122134
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
123135
136+
If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
137+
with similar accuracy, but is much smaller when printing types.
138+
124139
This constructor does not store/modify `y` or `x`.
125140
"""
141+
@inline function DerivativeConfig(f::F,
142+
y::AbstractArray{Y},
143+
x::X;
144+
tag::Union{Symbol,Nothing} = :default) where {F,X<:Real,Y<:Real}
145+
# @inline ensures that, e.g., DerivativeConfig(...; tag = :small) will be well-inferred
146+
T = @inline maketag(tag, f, X)
147+
return @noinline DerivativeConfig(f,y,x,T)
148+
end
149+
126150
function DerivativeConfig(f::F,
127151
y::AbstractArray{Y},
128152
x::X,
129-
tag::T = Tag(f, X)) where {F,X<:Real,Y<:Real,T}
153+
tag::T) where {F,X<:Real,Y<:Real,T}
130154
duals = similar(y, Dual{T,Y,1})
131155
return DerivativeConfig{T,typeof(duals)}(duals)
132156
end
@@ -144,24 +168,36 @@ struct GradientConfig{T,V,N,D} <: AbstractConfig{N}
144168
end
145169

146170
"""
147-
ForwardDiff.GradientConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x))
171+
ForwardDiff.GradientConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)
148172
149173
Return a `GradientConfig` instance based on the type of `f` and type/shape of the input
150174
vector `x`.
151175
152176
The returned `GradientConfig` instance contains all the work buffers required by
153177
`ForwardDiff.gradient` and `ForwardDiff.gradient!`.
154178
155-
If `f` is `nothing` instead of the actual target function, then the returned instance can
156-
be used with any target function. However, this will reduce ForwardDiff's ability to catch
157-
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
179+
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
180+
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
181+
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
182+
183+
If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
184+
with similar accuracy, but is much smaller when printing types.
158185
159186
This constructor does not store/modify `x`.
160187
"""
188+
@inline function GradientConfig(f::F,
189+
x::AbstractArray{V},
190+
c::Chunk{N} = Chunk(x);
191+
tag::Union{Symbol,Nothing} = :default) where {F,V,N}
192+
# @inline ensures that, e.g., GradientConfig(...; tag = :small) will be well-inferred
193+
T = @inline maketag(tag, f, V)
194+
return @noinline GradientConfig(f,x,c,T)
195+
end
196+
161197
function GradientConfig(f::F,
162198
x::AbstractArray{V},
163-
::Chunk{N} = Chunk(x),
164-
::T = Tag(f, V)) where {F,V,N,T}
199+
::Chunk{N},
200+
::T) where {F,V,N,T}
165201
seeds = construct_seeds(Partials{N,V})
166202
duals = similar(x, Dual{T,V,N})
167203
return GradientConfig{T,V,N,typeof(duals)}(seeds, duals)
@@ -180,7 +216,7 @@ struct JacobianConfig{T,V,N,D} <: AbstractConfig{N}
180216
end
181217

182218
"""
183-
ForwardDiff.JacobianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x))
219+
ForwardDiff.JacobianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)
184220
185221
Return a `JacobianConfig` instance based on the type of `f` and type/shape of the input
186222
vector `x`.
@@ -189,23 +225,35 @@ The returned `JacobianConfig` instance contains all the work buffers required by
189225
`ForwardDiff.jacobian` and `ForwardDiff.jacobian!` when the target function takes the form
190226
`f(x)`.
191227
192-
If `f` is `nothing` instead of the actual target function, then the returned instance can
193-
be used with any target function. However, this will reduce ForwardDiff's ability to catch
194-
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
228+
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
229+
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
230+
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
231+
232+
If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
233+
with similar accuracy, but is much smaller when printing types.
195234
196235
This constructor does not store/modify `x`.
197236
"""
237+
@inline function JacobianConfig(f::F,
238+
x::AbstractArray{V},
239+
c::Chunk{N} = Chunk(x);
240+
tag::Union{Symbol,Nothing} = :default) where {F,V,N}
241+
# @inline ensures that, e.g., JacobianConfig(...; tag = :small) will be well-inferred
242+
T = @inline maketag(tag, f, V)
243+
return @noinline JacobianConfig(f,x,c,T)
244+
end
245+
198246
function JacobianConfig(f::F,
199247
x::AbstractArray{V},
200-
::Chunk{N} = Chunk(x),
201-
::T = Tag(f, V)) where {F,V,N,T}
248+
::Chunk{N},
249+
::T) where {F,V,N,T}
202250
seeds = construct_seeds(Partials{N,V})
203251
duals = similar(x, Dual{T,V,N})
204252
return JacobianConfig{T,V,N,typeof(duals)}(seeds, duals)
205253
end
206254

207255
"""
208-
ForwardDiff.JacobianConfig(f!, y::AbstractArray, x::AbstractArray, chunk::Chunk = Chunk(x))
256+
ForwardDiff.JacobianConfig(f!, y::AbstractArray, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)
209257
210258
Return a `JacobianConfig` instance based on the type of `f!`, and the types/shapes of the
211259
output vector `y` and the input vector `x`.
@@ -214,17 +262,30 @@ The returned `JacobianConfig` instance contains all the work buffers required by
214262
`ForwardDiff.jacobian` and `ForwardDiff.jacobian!` when the target function takes the form
215263
`f!(y, x)`.
216264
217-
If `f!` is `nothing` instead of the actual target function, then the returned instance can
218-
be used with any target function. However, this will reduce ForwardDiff's ability to catch
219-
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
265+
If `f!` or `tag` is `nothing`, then the returned instance can be used with any target function.
266+
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
267+
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
268+
269+
If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
270+
with similar accuracy, but is much smaller when printing types.
220271
221272
This constructor does not store/modify `y` or `x`.
222273
"""
274+
@inline function JacobianConfig(f::F,
275+
y::AbstractArray{Y},
276+
x::AbstractArray{X},
277+
c::Chunk{N} = Chunk(x);
278+
tag::Union{Symbol,Nothing} = :default) where {F,Y,X,N}
279+
# @inline ensures that, e.g., JacobianConfig(...; tag = :small) will be well-inferred
280+
T = @inline maketag(tag, f, X)
281+
return @noinline JacobianConfig(f,y,x,c,T)
282+
end
283+
223284
function JacobianConfig(f::F,
224285
y::AbstractArray{Y},
225286
x::AbstractArray{X},
226-
::Chunk{N} = Chunk(x),
227-
::T = Tag(f, X)) where {F,Y,X,N,T}
287+
::Chunk{N},
288+
::T) where {F,Y,X,N,T}
228289
seeds = construct_seeds(Partials{N,X})
229290
yduals = similar(y, Dual{T,Y,N})
230291
xduals = similar(x, Dual{T,X,N})
@@ -245,7 +306,7 @@ struct HessianConfig{T,V,N,DG,DJ} <: AbstractConfig{N}
245306
end
246307

247308
"""
248-
ForwardDiff.HessianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x))
309+
ForwardDiff.HessianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)
249310
250311
Return a `HessianConfig` instance based on the type of `f` and type/shape of the input
251312
vector `x`.
@@ -256,41 +317,66 @@ configured for the case where the `result` argument is an `AbstractArray`. If
256317
it is a `DiffResult`, the `HessianConfig` should instead be constructed via
257318
`ForwardDiff.HessianConfig(f, result, x, chunk)`.
258319
259-
If `f` is `nothing` instead of the actual target function, then the returned instance can
260-
be used with any target function. However, this will reduce ForwardDiff's ability to catch
261-
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
320+
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
321+
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
322+
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
323+
324+
If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
325+
with similar accuracy, but is much smaller when printing types.
262326
263327
This constructor does not store/modify `x`.
264328
"""
329+
@inline function HessianConfig(f::F,
330+
x::AbstractArray{V},
331+
chunk::Chunk = Chunk(x);
332+
tag::Union{Symbol,Nothing} = :default) where {F,V}
333+
# @inline ensures that, e.g., HessianConfig(...; tag = :small) will be well-inferred
334+
T = @inline maketag(tag, f, V)
335+
return @noinline HessianConfig(f, x, chunk, T)
336+
end
337+
265338
function HessianConfig(f::F,
266339
x::AbstractArray{V},
267-
chunk::Chunk = Chunk(x),
268-
tag = Tag(f, V)) where {F,V}
340+
chunk::Chunk,
341+
tag) where {F,V}
269342
jacobian_config = JacobianConfig(f, x, chunk, tag)
270343
gradient_config = GradientConfig(f, jacobian_config.duals, chunk, tag)
271344
return HessianConfig(jacobian_config, gradient_config)
272345
end
273346

274347
"""
275-
ForwardDiff.HessianConfig(f, result::DiffResult, x::AbstractArray, chunk::Chunk = Chunk(x))
348+
ForwardDiff.HessianConfig(f, result::DiffResult, x::AbstractArray, chunk::Chunk = Chunk(x); tag::Union{Symbol,Nothing} = :default)
276349
277350
Return a `HessianConfig` instance based on the type of `f`, types/storage in `result`, and
278351
type/shape of the input vector `x`.
279352
280353
The returned `HessianConfig` instance contains all the work buffers required by
281354
`ForwardDiff.hessian!` for the case where the `result` argument is an `DiffResult`.
282355
283-
If `f` is `nothing` instead of the actual target function, then the returned instance can
284-
be used with any target function. However, this will reduce ForwardDiff's ability to catch
285-
and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
356+
If `f` or `tag` is `nothing`, then the returned instance can be used with any target function.
357+
However, this will reduce ForwardDiff's ability to catch and prevent perturbation confusion
358+
(see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
359+
360+
If `tag` is `:small`, a small hash-based tag is provided. This tracks perturbation confusion
361+
with similar accuracy, but is much smaller when printing types.
286362
287363
This constructor does not store/modify `x`.
288364
"""
365+
@inline function HessianConfig(f::F,
366+
result::DiffResult,
367+
x::AbstractArray{V},
368+
chunk::Chunk = Chunk(x);
369+
tag::Union{Symbol,Nothing} = :default) where {F,V}
370+
# @inline ensures that, e.g., HessianConfig(...; tag = :small) will be well-inferred
371+
T = @inline maketag(tag, f, V)
372+
return @noinline HessianConfig(f, result, x, chunk, T)
373+
end
374+
289375
function HessianConfig(f::F,
290376
result::DiffResult,
291377
x::AbstractArray{V},
292-
chunk::Chunk = Chunk(x),
293-
tag = Tag(f, V)) where {F,V}
378+
chunk::Chunk,
379+
tag) where {F,V}
294380
jacobian_config = JacobianConfig((f,gradient), DiffResults.gradient(result), x, chunk, tag)
295381
gradient_config = GradientConfig(f, jacobian_config.duals[2], chunk, tag)
296382
return HessianConfig(jacobian_config, gradient_config)

‎test/GradientTest.jl

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ import NaNMath
66
using Test
77
using LinearAlgebra
88
using ForwardDiff
9-
using ForwardDiff: Dual, Tag
9+
using ForwardDiff: Dual, maketag
1010
using StaticArrays
1111
using DiffTests
1212

@@ -21,7 +21,8 @@ x = [0.1, 0.2, 0.3]
2121
v = f(x)
2222
g = [-9.4, 15.6, 52.0]
2323

24-
@testset "Rosenbrock, chunk size = $c and tag = $(repr(tag))" for c in (1, 2, 3), tag in (nothing, Tag(f, eltype(x)))
24+
@testset "Rosenbrock, chunk size = $c and tag = $(repr(maketag(tag, f, eltype(x))))" for c in (1, 2, 3), tag in (nothing, :default, :small)
25+
tag = maketag(tag, f, eltype(x))
2526
cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}(), tag)
2627

2728
@test eltype(cfg) == Dual{typeof(tag), eltype(x), c}
@@ -60,7 +61,8 @@ cfgx = ForwardDiff.GradientConfig(sin, x)
6061
v = f(X)
6162
g = ForwardDiff.gradient(f, X)
6263
@test isapprox(g, Calculus.gradient(f, X), atol=FINITEDIFF_ERROR)
63-
@testset "... with chunk size = $c and tag = $(repr(tag))" for c in CHUNK_SIZES, tag in (nothing, Tag(f, eltype(x)))
64+
@testset "... with chunk size = $c and tag = $(repr(maketag(tag, f, eltype(x))))" for c in CHUNK_SIZES, tag in (nothing, :default, :small)
65+
tag = maketag(tag, f, eltype(x))
6466
cfg = ForwardDiff.GradientConfig(f, X, ForwardDiff.Chunk{c}(), tag)
6567

6668
out = ForwardDiff.gradient(f, X, cfg)
@@ -140,6 +142,12 @@ end
140142

141143
# make sure this is not a source of type instability
142144
@inferred ForwardDiff.GradientConfig(f, sx)
145+
if VERSION ≥ v"1.11"
146+
# make sure that `GradientConfig(...; tag = compile-time-constant)` also
147+
# infers well (requires that Base.hash(::Type) is foldable, which is true
148+
# in Julia ≥ 1.11)
149+
@inferred ((f, sx)->ForwardDiff.GradientConfig(f, sx; tag=:small))(f, sx)
150+
end
143151
end
144152

145153
@testset "exponential function at base zero" begin

‎test/HessianTest.jl

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import Calculus
55
using Test
66
using LinearAlgebra
77
using ForwardDiff
8-
using ForwardDiff: Dual, Tag
8+
using ForwardDiff: Dual, maketag
99
using StaticArrays
1010
using DiffTests
1111

@@ -23,7 +23,8 @@ h = [-66.0 -40.0 0.0;
2323
-40.0 130.0 -80.0;
2424
0.0 -80.0 200.0]
2525

26-
@testset "running hardcoded test with chunk size = $c and tag = $(repr(tag))" for c in HESSIAN_CHUNK_SIZES, tag in (nothing, Tag((f,ForwardDiff.gradient), eltype(x)))
26+
@testset "running hardcoded test with chunk size = $c and tag = $(repr(maketag(tag, (f, ForwardDiff.gradient), eltype(x))))" for c in HESSIAN_CHUNK_SIZES, tag in (nothing, :default, :small)
27+
tag = maketag(tag, (f, ForwardDiff.gradient), eltype(x))
2728
cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}(), tag)
2829
resultcfg = ForwardDiff.HessianConfig(f, DiffResults.HessianResult(x), x, ForwardDiff.Chunk{c}(), tag)
2930

@@ -68,7 +69,8 @@ for f in DiffTests.VECTOR_TO_NUMBER_FUNCS
6869
h = ForwardDiff.hessian(f, X)
6970
# finite difference approximation error is really bad for Hessians...
7071
@test isapprox(h, Calculus.hessian(f, X), atol=0.02)
71-
@testset "$f with chunk size = $c and tag = $(repr(tag))" for c in HESSIAN_CHUNK_SIZES, tag in (nothing, Tag((f,ForwardDiff.gradient), eltype(x)))
72+
@testset "$f with chunk size = $c and tag = $(repr(maketag(tag, (f, ForwardDiff.gradient), eltype(x))))" for c in HESSIAN_CHUNK_SIZES, tag in (nothing, :default, :small)
73+
tag = maketag(tag, (f, ForwardDiff.gradient), eltype(x))
7274
cfg = ForwardDiff.HessianConfig(f, X, ForwardDiff.Chunk{c}(), tag)
7375
resultcfg = ForwardDiff.HessianConfig(f, DiffResults.HessianResult(X), X, ForwardDiff.Chunk{c}(), tag)
7476

‎test/JacobianTest.jl

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import Calculus
44

55
using Test
66
using ForwardDiff
7-
using ForwardDiff: Dual, Tag, JacobianConfig
7+
using ForwardDiff: Dual, Tag, SmallTag, JacobianConfig, maketag
88
using StaticArrays
99
using DiffTests
1010
using LinearAlgebra
@@ -31,8 +31,8 @@ j = [0.8242369704835132 0.4121184852417566 -10.933563142616123
3131
0.169076696546684 0.084538348273342 -2.299173530851733
3232
0.0 0.0 1.0]
3333

34-
for c in (1, 2, 3), tags in ((nothing, nothing),
35-
(Tag(f, eltype(x)), Tag(f!, eltype(x))))
34+
for c in (1, 2, 3), tag in (nothing, :small, :default)
35+
tags = (maketag(tag, f, eltype(x)), maketag(tag, f!, eltype(x)))
3636
println(" ...running hardcoded test with chunk size = $c and tag = $(repr(tags))")
3737
cfg = JacobianConfig(f, x, ForwardDiff.Chunk{c}(), tags[1])
3838
ycfg = JacobianConfig(f!, fill(0.0, 4), x, ForwardDiff.Chunk{c}(), tags[2])
@@ -103,9 +103,11 @@ for f in DiffTests.ARRAY_TO_ARRAY_FUNCS
103103
v = f(X)
104104
j = ForwardDiff.jacobian(f, X)
105105
@test isapprox(j, Calculus.jacobian(x -> vec(f(x)), X, :forward), atol=1.3FINITEDIFF_ERROR)
106-
@testset "$f with chunk size = $c and tag = $(repr(tag))" for c in CHUNK_SIZES, tag in (nothing, Tag)
106+
@testset "$f with chunk size = $c and tag = $(repr(tag))" for c in CHUNK_SIZES, tag in (nothing, Tag, SmallTag)
107107
if tag == Tag
108108
tag = Tag(f, eltype(X))
109+
elseif tag == SmallTag
110+
tag = SmallTag(f, eltype(X))
109111
end
110112
cfg = JacobianConfig(f, X, ForwardDiff.Chunk{c}(), tag)
111113

@@ -128,7 +130,8 @@ for f! in DiffTests.INPLACE_ARRAY_TO_ARRAY_FUNCS
128130
f!(v, X)
129131
j = ForwardDiff.jacobian(f!, fill!(similar(Y), 0.0), X)
130132
@test isapprox(j, Calculus.jacobian(x -> (y = fill!(similar(Y), 0.0); f!(y, x); vec(y)), X, :forward), atol=FINITEDIFF_ERROR)
131-
@testset "$(f!) with chunk size = $c and tag = $(repr(tag))" for c in CHUNK_SIZES, tag in (nothing, Tag(f!, eltype(X)))
133+
@testset "$(f!) with chunk size = $c and tag = $(repr(maketag(tag, f!, eltype(X))))" for c in CHUNK_SIZES, tag in (nothing, :small, :default)
134+
tag = maketag(tag, f!, eltype(X))
132135
ycfg = JacobianConfig(f!, fill!(similar(Y), 0.0), X, ForwardDiff.Chunk{c}(), tag)
133136

134137
y = fill!(similar(Y), 0.0)
@@ -225,6 +228,12 @@ for T in (StaticArrays.SArray, StaticArrays.MArray)
225228

226229
# make sure this is not a source of type instability
227230
@inferred ForwardDiff.JacobianConfig(f, sx)
231+
if VERSION ≥ v"1.11"
232+
# make sure that `JacobianConfig(...; tag = compile-time-constant)` also
233+
# infers well (requires that Base.hash(::Type) is foldable, which is true
234+
# in Julia ≥ 1.11)
235+
@inferred ((f, sx)->ForwardDiff.JacobianConfig(f, sx; tag=:small))(f, sx)
236+
end
228237
end
229238

230239
#########

0 commit comments

Comments
 (0)