Skip to content

Commit 2a2a50d

Browse files
topolaritytopolarity
committed
Support tag = :small in Config(...) constructors
This provides a convenient interface to ask for a SmallTag.
1 parent 48b4d69 commit 2a2a50d

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
@@ -98,6 +98,18 @@ Base.eltype(cfg::AbstractConfig) = eltype(typeof(cfg))
9898

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

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

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

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

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

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

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

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