tweak tag generation (#279)

* tweak tag generation

This does 3 related things
1. Re-enables void tags by providing `nothing` as function argument
2. Allows opt out of tag checking via `Val{false}()` argument.
3. Simplifies the Hessian to use the same tag (which is useful for my own purposes of computing Riemannian metrics, but also makes the code a bit cleaner).

* add docs

* derivative fixes

* add tests for tag checks

* add tag docs

Author: simonbyrne

docs/src/user/advanced.md

@@ -283,3 +283,36 @@ Note that whether or not SIMD instructions can actually be used will depend on y
 and Julia build. For example, pre-built Julia binaries might not emit vectorized LLVM
 bitcode. To overcome this specific issue, you can [locally rebuild Julia's system
 image](http://docs.julialang.org/en/latest/devdocs/sysimg).
+
+## Custom tags and tag checking
+
+The `Dual` type includes a "tag" parameter indicating the particular function call to
+which it belongs. This is to avoid a problem known as [_perturbation
+confusion_](https://github.com/JuliaDiff/ForwardDiff.jl/issues/83) which can arise when
+there are nested differentiation calls. Tags are automatically generated as part of the
+appropriate config object, and the tag is checked when the config is used as part of a
+differentiation call (`derivative`, `gradient`, etc.): an `InvalidTagException` will be
+thrown if the incorrect config object is used.
+
+This checking can sometimes be inconvenient, and there are certain cases where you may
+want to disable this checking.
+
+!!! warning
+    Disabling tag checking should only be done with caution, especially if the code itself
+    could be used inside another differentiation call.
+
+1. (preferred) Provide an extra `Val{false}()` argument to the differentiation function, e.g.
+   ```julia
+   cfg = ForwarDiff.GradientConfig(g, x)
+   ForwarDiff.gradient(f, x, cfg, Val{false}())
+   ```
+   If using as part of a library, the tag can be checked manually via
+   ```julia
+   ForwardDiff.checktag(cfg, g, x)
+   ```
+
+2. (discouraged) Construct the config object with `nothing` instead of a function argument, e.g.
+   ```julia
+   cfg = GradientConfig(nothing, x)
+   gradient(f, x, cfg)
+   ```

src/config.jl

@@ -18,14 +18,21 @@ function Tag(f::F, ::Type{V}) where {F,V}
     Tag{F,V}()
 end
 
+Tag(::Void, ::Type{V}) where {V} = nothing
+
+
 @inline function ≺(::Type{Tag{F1,V1}}, ::Type{Tag{F2,V2}}) where {F1,V1,F2,V2}
     tagcount(Tag{F1,V1}) < tagcount(Tag{F2,V2})
 end
 
+struct InvalidTagException{E,O} <: Exception
+end
 
+Base.showerror(io::IO, e::InvalidTagException{E,O}) where {E,O} =
+    print(io, "Invalid Tag object:\n  Expected $E,\n  Observed $O.")
 
 checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT,VT,F,V} =
-    error("Invalid Tag object:\n  Expected $(Tag{F,V}),\n  Observed $(Tag{FT,VT}).")
+    throw(InvalidTagException{Tag{F,V},Tag{FT,VT}}())
 
 checktag(::Type{Tag{F,V}}, f::F, x::AbstractArray{V}) where {F,V} = true
 
@@ -80,6 +87,7 @@ function DerivativeConfig(f::F,
     return DerivativeConfig{T,typeof(duals)}(duals)
 end
 
+checktag(::DerivativeConfig{T},f,x) where {T} = checktag(T,f,x)
 Base.eltype(::Type{DerivativeConfig{T,D}}) where {T,D} = eltype(D)
 
 ##################
@@ -115,6 +123,7 @@ function GradientConfig(f::F,
     return GradientConfig{T,V,N,typeof(duals)}(seeds, duals)
 end
 
+checktag(::GradientConfig{T},f,x) where {T} = checktag(T,f,x)
 Base.eltype(::Type{GradientConfig{T,V,N,D}}) where {T,V,N,D} = Dual{T,V,N}
 
 ##################
@@ -179,15 +188,16 @@ function JacobianConfig(f::F,
     return JacobianConfig{T,X,N,typeof(duals)}(seeds, duals)
 end
 
+checktag(::JacobianConfig{T},f,x) where {T} = checktag(T,f,x)
 Base.eltype(::Type{JacobianConfig{T,V,N,D}}) where {T,V,N,D} = Dual{T,V,N}
 
 #################
 # HessianConfig #
 #################
 
-struct HessianConfig{TG,TJ,V,N,DG,DJ} <: AbstractConfig{N}
-    jacobian_config::JacobianConfig{TJ,V,N,DJ}
-    gradient_config::GradientConfig{TG,Dual{TJ,V,N},N,DG}
+struct HessianConfig{T,V,N,DG,DJ} <: AbstractConfig{N}
+    jacobian_config::JacobianConfig{T,V,N,DJ}
+    gradient_config::GradientConfig{T,Dual{T,V,N},N,DG}
 end
 
 """
@@ -211,10 +221,9 @@ This constructor does not store/modify `x`.
 function HessianConfig(f::F,
                        x::AbstractArray{V},
                        chunk::Chunk = Chunk(x),
-                       tagj = Tag((f,gradient), V),
-                       tagg = Tag(f, Dual{typeof(tagj),V,chunksize(chunk)})) where {F,V}
-    jacobian_config = JacobianConfig((f,gradient), x, chunk, tagj)
-    gradient_config = GradientConfig(f, jacobian_config.duals, chunk, tagg)
+                       tag = Tag(f, V)) where {F,V}
+    jacobian_config = JacobianConfig(f, x, chunk, tag)
+    gradient_config = GradientConfig(f, jacobian_config.duals, chunk, tag)
     return HessianConfig(jacobian_config, gradient_config)
 end
 
@@ -237,12 +246,12 @@ function HessianConfig(f::F,
                        result::DiffResult,
                        x::AbstractArray{V},
                        chunk::Chunk = Chunk(x),
-                       tagj = Tag((f,typeof(gradient)), V),
-                       tagg = Tag(f, Dual{typeof(tagj),V,chunksize(chunk)})) where {F,V}
-    jacobian_config = JacobianConfig((f,gradient), DiffResults.gradient(result), x, chunk, tagj)
-    gradient_config = GradientConfig(f, jacobian_config.duals[2], chunk, tagg)
+                       tag = Tag(f, V)) where {F,V}
+    jacobian_config = JacobianConfig((f,gradient), DiffResults.gradient(result), x, chunk, tag)
+    gradient_config = GradientConfig(f, jacobian_config.duals[2], chunk, tag)
     return HessianConfig(jacobian_config, gradient_config)
 end
 
-Base.eltype(::Type{HessianConfig{TG,TJ,V,N,DG,DJ}}) where {TG,TJ,V,N,DG,DJ} =
-    Dual{TG,Dual{TJ,V,N},N}
+checktag(::HessianConfig{T},f,x) where {T} = checktag(T,f,x)
+Base.eltype(::Type{HessianConfig{T,V,N,DG,DJ}}) where {T,V,N,DG,DJ} =
+    Dual{T,Dual{T,V,N},N}

src/derivative.jl

@@ -15,14 +15,16 @@ This method assumes that `isa(f(x), Union{Real,AbstractArray})`.
 end
 
 """
-    ForwardDiff.derivative(f!, y::AbstractArray, x::Real, cfg::DerivativeConfig = DerivativeConfig(f!, y, x))
+    ForwardDiff.derivative(f!, y::AbstractArray, x::Real, cfg::DerivativeConfig = DerivativeConfig(f!, y, x), check=Val{true}())
 
 Return `df!/dx` evaluated at `x`, assuming `f!` is called as `f!(y, x)` where the result is
 stored in `y`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
 @inline function derivative(f!, y::AbstractArray, x::Real,
-                            cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x)) where {T}
-    checktag(T, f!, x)
+                            cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK}
+    CHK && checktag(T, f!, x)
     ydual = cfg.duals
     seed!(ydual, y)
     f!(ydual, Dual{T}(x, one(x)))
@@ -48,15 +50,17 @@ This method assumes that `isa(f(x), Union{Real,AbstractArray})`.
 end
 
 """
-    ForwardDiff.derivative!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::Real, cfg::DerivativeConfig = DerivativeConfig(f!, y, x))
+    ForwardDiff.derivative!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::Real, cfg::DerivativeConfig = DerivativeConfig(f!, y, x), check=Val{true}())
 
 Compute `df!/dx` evaluated at `x` and store the result(s) in `result`, assuming `f!` is
 called as `f!(y, x)` where the result is stored in `y`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
 @inline function derivative!(result::Union{AbstractArray,DiffResult},
                              f!, y::AbstractArray, x::Real,
-                             cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x)) where {T}
-    checktag(T, f!, x)
+                             cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK}
+    CHK && checktag(T, f!, x)
     ydual = cfg.duals
     seed!(ydual, y)
     f!(ydual, Dual{T}(x, one(x)))

src/gradient.jl

@@ -3,14 +3,16 @@
 ###############
 
 """
-    ForwardDiff.gradient(f, x::AbstractArray, cfg::GradientConfig = GradientConfig(f, x))
+    ForwardDiff.gradient(f, x::AbstractArray, cfg::GradientConfig = GradientConfig(f, x), check=Val{true}())
 
 Return `∇f` evaluated at `x`, assuming `f` is called as `f(x)`.
 
 This method assumes that `isa(f(x), Real)`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function gradient(f, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x)) where {T}
-    checktag(T, f, x)
+function gradient(f, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK}
+    CHK && checktag(T, f, x)
     if chunksize(cfg) == length(x)
         return vector_mode_gradient(f, x, cfg)
     else
@@ -19,15 +21,16 @@ function gradient(f, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f
 end
 
 """
-    ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::GradientConfig = GradientConfig(f, x))
+    ForwardDiff.gradient!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::GradientConfig = GradientConfig(f, x), check=Val{true}())
 
 Compute `∇f` evaluated at `x` and store the result(s) in `result`, assuming `f` is called as
 `f(x)`.
 
 This method assumes that `isa(f(x), Real)`.
+
 """
-function gradient!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x)) where {T}
-    checktag(T, f, x)
+function gradient!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK}
+    CHK && checktag(T, f, x)
     if chunksize(cfg) == length(x)
         vector_mode_gradient!(result, f, x, cfg)
     else

src/hessian.jl

@@ -3,46 +3,55 @@
 ###############
 
 """
-    ForwardDiff.hessian(f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x))
+    ForwardDiff.hessian(f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x), check=Val{true}())
 
 Return `H(f)` (i.e. `J(∇(f))`) evaluated at `x`, assuming `f` is called as `f(x)`.
 
 This method assumes that `isa(f(x), Real)`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function hessian(f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x))
-    ∇f = y -> gradient(f, y, cfg.gradient_config)
-    return jacobian(∇f, x, cfg.jacobian_config)
+function hessian(f, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {T,CHK}
+    CHK && checktag(T, f, x)
+    ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}())
+    return jacobian(∇f, x, cfg.jacobian_config, Val{false}())
 end
 
 """
-    ForwardDiff.hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x))
+    ForwardDiff.hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x), check=Val{true}())
 
 Compute `H(f)` (i.e. `J(∇(f))`) evaluated at `x` and store the result(s) in `result`,
 assuming `f` is called as `f(x)`.
 
 This method assumes that `isa(f(x), Real)`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x))
-    ∇f = y -> gradient(f, y, cfg.gradient_config)
-    jacobian!(result, ∇f, x, cfg.jacobian_config)
+function hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {T,CHK}
+    CHK && checktag(T, f, x)
+    ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}())
+    jacobian!(result, ∇f, x, cfg.jacobian_config, Val{false}())
     return result
 end
 
 """
-    ForwardDiff.hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, result, x))
+    ForwardDiff.hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, result, x), check=Val{true}())
 
 Exactly like `ForwardDiff.hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig)`, but
 because `isa(result, DiffResult)`, `cfg` is constructed as `HessianConfig(f, result, x)` instead of
 `HessianConfig(f, x)`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, result, x))
+function hessian!(result::DiffResult, f, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {T,CHK}
+    CHK && checktag(T, f, x)
     ∇f! = (y, z) -> begin
         inner_result = DiffResult(zero(eltype(y)), y)
-        gradient!(inner_result, f, z, cfg.gradient_config)
+        gradient!(inner_result, f, z, cfg.gradient_config, Val{false}())
         result = DiffResults.value!(result, value(DiffResults.value(inner_result)))
         return y
     end
-    jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config)
+    jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config, Val{false}())
     return result
 end
 
@@ -57,14 +66,13 @@ hessian!(result::MutableDiffResult, f, x::SArray) = hessian!(result, f, x, Hessi
 hessian!(result::ImmutableDiffResult, f, x::SArray, cfg::HessianConfig) = hessian!(result, f, x)
 
 function hessian!(result::ImmutableDiffResult, f::F, x::SArray{S,V}) where {F,S,V}
-    TJ = typeof(Tag((f,gradient),V))
-    d1 = dualize(TJ, x)
-    TG = typeof(Tag(f, eltype(d1)))
-    d2 = dualize(TG, d1)
+    T = typeof(Tag(f,V))
+    d1 = dualize(T, x)
+    d2 = dualize(T, d1)
     fd2 = f(d2)
-    val = value(TJ,value(TG,fd2))
-    grad = extract_gradient(TJ,value(TG,fd2), x)
-    hess = extract_jacobian(TJ,partials(TG,fd2), x)
+    val = value(T,value(T,fd2))
+    grad = extract_gradient(T,value(T,fd2), x)
+    hess = extract_jacobian(T,partials(T,fd2), x)
     result = DiffResults.hessian!(result, hess)
     result = DiffResults.gradient!(result, grad)
     result = DiffResults.value!(result, val)

src/jacobian.jl

@@ -3,14 +3,16 @@
 ###############
 
 """
-    ForwardDiff.jacobian(f, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f, x))
+    ForwardDiff.jacobian(f, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f, x), check=Val{true}())
 
 Return `J(f)` evaluated at `x`, assuming `f` is called as `f(x)`.
 
 This method assumes that `isa(f(x), AbstractArray)`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function jacobian(f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x)) where {T}
-    checktag(T, f, x)
+function jacobian(f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {T,CHK}
+    CHK && checktag(T, f, x)
     if chunksize(cfg) == length(x)
         return vector_mode_jacobian(f, x, cfg)
     else
@@ -19,13 +21,15 @@ function jacobian(f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f
 end
 
 """
-    ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x))
+    ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())
 
 Return `J(f!)` evaluated at `x`,  assuming `f!` is called as `f!(y, x)` where the result is
 stored in `y`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x)) where {T}
-    checktag(T, f!, x)    
+function jacobian(f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {T, CHK}
+    CHK && checktag(T, f!, x)    
     if chunksize(cfg) == length(x)
         return vector_mode_jacobian(f!, y, x, cfg)
     else
@@ -35,15 +39,17 @@ end
 
 
 """
-    ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f, x))
+    ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f, x), check=Val{true}())
 
 Compute `J(f)` evaluated at `x` and store the result(s) in `result`, assuming `f` is called
 as `f(x)`.
 
 This method assumes that `isa(f(x), AbstractArray)`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x)) where {T}
-    checktag(T, f, x)
+function jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK}
+    CHK && checktag(T, f, x)
     if chunksize(cfg) == length(x)
         vector_mode_jacobian!(result, f, x, cfg)
     else
@@ -53,15 +59,17 @@ function jacobian!(result::Union{AbstractArray,DiffResult}, f, x::AbstractArray,
 end
 
 """
-    ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x))
+    ForwardDiff.jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())
 
 Compute `J(f!)` evaluated at `x` and store the result(s) in `result`, assuming `f!` is
 called as `f!(y, x)` where the result is stored in `y`.
 
 This method assumes that `isa(f(x), AbstractArray)`.
+
+Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
-function jacobian!(result::Union{AbstractArray,DiffResult}, f!, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x)) where {T}
-    checktag(T, f!, x)
+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}
+    CHK && checktag(T, f!, x)
     if chunksize(cfg) == length(x)
         vector_mode_jacobian!(result, f!, y, x, cfg)
     else

test/GradientTest.jl

@@ -46,6 +46,11 @@ for c in (1, 2, 3), tag in (nothing, Tag(f, eltype(x)))
     @test isapprox(DiffResults.value(out), v)
 end
 
+cfgx = ForwardDiff.GradientConfig(sin, x)
+@test_throws ForwardDiff.InvalidTagException ForwardDiff.gradient(f, x, cfgx)
+@test ForwardDiff.gradient(f, x, cfgx, Val{false}()) == ForwardDiff.gradient(f,x)
+
+
 ########################
 # test vs. Calculus.jl #
 ########################
@@ -130,4 +135,5 @@ sresult3 = ForwardDiff.gradient!(sresult3, prod, sx, scfg)
 @test DiffResults.gradient(sresult2) == DiffResults.gradient(result)
 @test DiffResults.gradient(sresult3) == DiffResults.gradient(result)
 
+
 end # module