Order tags structurally and require unique tags in nested Duals

Project.toml

@@ -41,9 +41,10 @@ DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Calculus", "DiffTests", "IrrationalConstants", "JET", "SparseArrays", "StaticArrays", "Test", "InteractiveUtils"]
+test = ["Calculus", "DiffTests", "IrrationalConstants", "JET", "Serialization", "SparseArrays", "StaticArrays", "Test", "InteractiveUtils"]

docs/src/dev/how_it_works.md

@@ -43,7 +43,7 @@ of the function, propagating the derivative via multiplication. For example, `Ba
 can be overloaded on `Dual` like so:
 
 ```julia
-Base.sin(d::Dual{T}) where {T} = Dual{T}(sin(value(d)), cos(value(d)) * partials(d))
+Base.sin(d::Dual{T}) where {T} = Dual{T}(sin(value(T, d)), cos(value(T, d)) * partials(T, d))
 ```
 
 If we assume that a general function `f` is composed of entirely of these elementary

docs/src/user/advanced.md

@@ -128,8 +128,8 @@ aren't sensitive to the input and thus cause ForwardDiff to incorrectly return `
 ```julia-repl
 # the dual number's perturbation component is zero, so this
 # variable should not propagate derivative information
-julia> log(ForwardDiff.Dual{:tag}(0.0, 0.0))
-Dual{:tag}(-Inf,NaN) # oops, this NaN should be 0.0
+julia> log(ForwardDiff.Dual(0.0, 0.0))
+Dual{ForwardDiff.Tag{Nothing, Float64}}(-Inf,NaN) # oops, this NaN should be 0.0
 ```
 
 Here, ForwardDiff computes the derivative of `log(0.0)` as `NaN` and then propagates
@@ -166,8 +166,8 @@ julia> set_preferences!(UUID("f6369f11-7733-5829-9624-2563aa707210"), "nansafe_m
 
 julia> using ForwardDiff
 
-julia> log(ForwardDiff.Dual{:tag}(0.0, 0.0))
-Dual{:tag}(-Inf,0.0)
+julia> log(ForwardDiff.Dual(0.0, 0.0))
+Dual{ForwardDiff.Tag{Nothing, Float64}}(-Inf,0.0)
 ```
 
 In the future, we plan on allowing users and downstream library authors to dynamically

ext/ForwardDiffStaticArraysExt.jl

@@ -132,13 +132,16 @@ ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray, cfg::Hes
 ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray, cfg::HessianConfig, ::Val) where {F} = hessian!(result, f, x)
 
 function ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray) where {F}
-    T = typeof(Tag(f, eltype(x)))
-    d1 = dualize(T, x)
-    d2 = dualize(T, d1)
+    T1 = typeof(Tag(f, eltype(x)))
+    d1 = dualize(T1, x)
+    T2 = typeof(Tag(f, eltype(d1)))
+    d2 = dualize(T2, d1)
     fd2 = f(d2)
-    val = value(T,value(T,fd2))
-    grad = extract_gradient(T,value(T,fd2), x)
-    hess = extract_jacobian(T,partials(T,fd2), x)
+    # Hessian = Jacobian (w.r.t. `T1`) of the gradient (w.r.t. `T2`), as in `hessian!` for other arrays
+    ∇fd2 = extract_gradient(T2, fd2, d1)
+    val = value(T1, value(T2, fd2))
+    grad = map(Base.Fix1(value, T1), ∇fd2)
+    hess = extract_jacobian(T1, ∇fd2, x)
     result = DiffResults.hessian!(result, hess)
     result = DiffResults.gradient!(result, grad)
     result = DiffResults.value!(result, val)

src/ForwardDiff.jl

@@ -21,6 +21,7 @@ include("derivative.jl")
 include("gradient.jl")
 include("jacobian.jl")
 include("hessian.jl")
+include("deprecated.jl")
 
 export DiffResults
 

src/apiutils.jl

@@ -20,14 +20,14 @@ end
 
 function vector_mode_dual_eval!(f::F, cfg::Union{JacobianConfig,GradientConfig}, x) where {F}
     xdual = cfg.duals
-    seed!(xdual, x, cfg.seeds)
+    seed!(eltype(cfg), xdual, x, cfg.seeds)
     return f(xdual)
 end
 
-function vector_mode_dual_eval!(f!::F, cfg::JacobianConfig, y, x) where {F}
+function vector_mode_dual_eval!(f!::F, cfg::JacobianConfig{T,V,N}, y, x) where {F,T,V,N}
     ydual, xdual = cfg.duals
-    seed!(xdual, x, cfg.seeds)
-    seed_zero_partials!(ydual, y)
+    seed!(eltype(cfg), xdual, x, cfg.seeds)
+    seed_zero_partials!(Dual{T,eltype(y),N}, ydual, y)
     f!(ydual, xdual)
     return ydual
 end
@@ -40,6 +40,24 @@ end
     return Expr(:tuple, [:(single_seed(Partials{N,V}, Val{$i}())) for i in 1:N]...)
 end
 
+# Seeds `x` with tag `T` and partials `p`, converted to the type of `x`. Layers of `x` with greater
+# tags are kept outside, so nested `Dual`s stay sorted even if `T` is not greater than all tags in `x`.
+@inline seed_dual(::Type{T}, x, p::Partials{N}) where {T,N} = Dual{T}(x, convert(Partials{N,typeof(x)}, p))
+@inline function seed_dual(::Type{T}, x::Dual{S}, p::Partials{N}) where {T,S,N}
+    p = convert(Partials{N,typeof(x)}, p)
+    T ≺ S || return Dual{T}(x, p)
+    # the seeds are constants, so their partials w.r.t. `S` are zero
+    q = map_partials(y -> value(S, y), valtype(S, eltype(p)), p)
+    return Dual{S}(seed_dual(T, value(S, x), q), map(y -> seed_dual(T, y, zero(q)), partials(S, x).values))
+end
+
+# Type of `seed_dual(T, x, p)` for `x::V` and `p::Partials{N,V}`. If `V` is abstract, each element is
+# seeded with its own type, so only `Real` is a bound.
+seed_type(::Type{Dual{T,V,N}}) where {T,V,N} = isconcretetype(V) ? Dual{T,V,N} : Real
+function seed_type(::Type{Dual{T,Dual{S,W,M},N}}) where {T,S,W,M,N}
+    return T ≺ S ? Dual{S,seed_type(Dual{T,W,N}),M} : Dual{T,Dual{S,W,M},N}
+end
+
 # Only seed indices that are structurally non-zero
 structural_eachindex(x::AbstractArray) = structural_eachindex(x, x)
 function structural_eachindex(x::AbstractArray, y::AbstractArray)
@@ -73,29 +91,29 @@ end
 # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is
 # currently carrying and to initialize a freshly allocated work buffer, whose elements must all be
 # written before the target function reads them.
-seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} =
-    _seed_zero_partials!(duals, x, structural_eachindex(duals, x))
+seed_zero_partials!(::Type{D}, duals::AbstractArray, x) where {D<:Dual} =
+    _seed_zero_partials!(D, duals, x, structural_eachindex(duals, x))
 
 # Zeroes the partials of `count` elements starting at structural position `index`. Chunk mode only
 # needs to clear the chunk it just seeded, so writing through to the end of the array would be O(n)
 # redundant work per chunk, i.e. O(n^2/N) per sweep. `count` mirrors the `chunksize` argument of
-# `seed!(duals, x, index, seeds, chunksize)`.
-function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, index,
+# `seed!(D, duals, x, index, seeds, chunksize)`.
+function seed_zero_partials!(::Type{Dual{T,V,N}}, duals::AbstractArray, x, index,
                              count = N) where {T,V,N}
     idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), count)
-    return _seed_zero_partials!(duals, x, idxs)
+    return _seed_zero_partials!(Dual{T,V,N}, duals, x, idxs)
 end
 
-function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N}
+function _seed_zero_partials!(::Type{Dual{T,V,N}}, duals::AbstractArray, x, idxs) where {T,V,N}
     seed = zero(Partials{N,V})
     if isbitstype(V)
         for idx in idxs
-            duals[idx] = Dual{T,V,N}(x[idx], seed)
+            duals[idx] = seed_dual(T, x[idx], seed)
         end
     else
         for idx in idxs
             if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seed)
+                duals[idx] = seed_dual(T, x[idx], seed)
             else
                 Base._unsetindex!(duals, idx)
             end
@@ -104,16 +122,16 @@ function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where
     return duals
 end
 
-function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
+function seed!(::Type{Dual{T,V,N}}, duals::AbstractArray, x,
                seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
     if isbitstype(V)
         for (i, idx) in zip(1:N, structural_eachindex(duals, x))
-            duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+            duals[idx] = seed_dual(T, x[idx], seeds[i])
         end
     else
         for (i, idx) in zip(1:N, structural_eachindex(duals, x))
             if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+                duals[idx] = seed_dual(T, x[idx], seeds[i])
             else
                 Base._unsetindex!(duals, idx)
             end
@@ -122,18 +140,18 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
     return duals
 end
 
-function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
+function seed!(::Type{Dual{T,V,N}}, duals::AbstractArray, x, index,
                seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
     offset = index - 1
     idxs = Iterators.drop(structural_eachindex(duals, x), offset)
     if isbitstype(V)
         for (i, idx) in zip(1:chunksize, idxs)
-            duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+            duals[idx] = seed_dual(T, x[idx], seeds[i])
         end
     else
         for (i, idx) in zip(1:chunksize, idxs)
             if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+                duals[idx] = seed_dual(T, x[idx], seeds[i])
             else
                 Base._unsetindex!(duals, idx)
             end

src/config.jl

@@ -5,25 +5,42 @@
 struct Tag{F,V}
 end
 
-const TAGCOUNT = Threads.Atomic{UInt}(0)
-
-# each tag is assigned a unique number
-# tags which depend on other tags will be larger
-@generated function tagcount(::Type{Tag{F,V}}) where {F,V}
-    :($(Threads.atomic_add!(TAGCOUNT, UInt(1))))
+Tag(f::F, ::Type{V}) where {F,V} = Tag{F,V}()
+
+# Encodes a type (or type parameter) as a sequence of strings that depends only on its
+# structure. Distinct objects have distinct keys, except a type and its redefinition in the
+# same session. The key of a parameter is a strict subsequence of the key of the type.
+function typekey!(key::Vector{String}, x::DataType)
+    mod = parentmodule(x)
+    push!(key, "T", string(Base.PkgId(Base.moduleroot(mod)).uuid), string(fullname(mod)), String(nameof(x)), string(length(x.parameters)))
+    foreach(p -> typekey!(key, p), x.parameters)
+    return key
 end
-
-function Tag(f::F, ::Type{V}) where {F,V}
-    tagcount(Tag{F,V}) # trigger generated function
-    Tag{F,V}()
+typekey!(key::Vector{String}, x::Symbol) = push!(key, "S", String(x))
+function typekey!(key::Vector{String}, x)
+    typekey!(push!(key, "V"), typeof(x))
+    if isprimitivetype(typeof(x))
+        push!(key, bytes2hex(reinterpret(UInt8, [x])))
+    else
+        for i in 1:nfields(x)
+            if isdefined(x, i)
+                typekey!(key, getfield(x, i))
+            else
+                push!(key, "#undef")
+            end
+        end
+    end
+    return key
 end
 
-Tag(::Nothing, ::Type{V}) where {V} = nothing
+# `A ≺ B` compares `(rank, key)` lexicographically. The rank strictly increases from a type to
+# any type containing it, so every tag is greater than the tags occurring in its parameters.
+@generated tagid(::Type{T}) where {T} = (key = typekey!(String[], T); (length(key), key...))
 
-
-@inline function ≺(::Type{Tag{F1,V1}}, ::Type{Tag{F2,V2}}) where {F1,V1,F2,V2}
-    tagcount(Tag{F1,V1}) < tagcount(Tag{F2,V2})
-end
+# Nested `Dual`s store the greatest tag outermost. A tag `Tag{F,V}` is greater than all tags
+# in the type `V` of its input and in the type `F` of its function, so seeding it outermost
+# keeps nested `Dual`s sorted. The comparison of the constant IDs is evaluated at compile time.
+≺(::Type{A}, ::Type{B}) where {A,B} = isless(tagid(A), tagid(B))
 
 struct InvalidTagException{E,O} <: Exception
 end
@@ -39,6 +56,9 @@ checktag(::Type{Tag{F,V}}, f::F, x::AbstractArray{V}) where {F,V} = true
 # no easy way to check Jacobian tag used with Hessians as multiple functions may be used
 checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT<:Tuple,VT,F,V} = true
 
+# tag of `nothing` configs, for any function
+checktag(::Type{Tag{FT,VT}}, f::F, x::AbstractArray{V}) where {FT<:Nothing,VT,F,V} = true
+
 # custom tag: you're on your own.
 checktag(z, f, x) = true
 
@@ -83,7 +103,7 @@ function DerivativeConfig(f::F,
                           y::AbstractArray{Y},
                           x::X,
                           tag::T = Tag(f, X)) where {F,X<:Real,Y<:Real,T}
-    duals = similar(y, Dual{T,Y,1})
+    duals = similar(y, seed_type(Dual{T,Y,1}))
     return DerivativeConfig{T,typeof(duals)}(duals)
 end
 
@@ -119,7 +139,7 @@ function GradientConfig(f::F,
                         ::Chunk{N} = Chunk(x),
                         ::T = Tag(f, V)) where {F,V,N,T}
     seeds = construct_seeds(Partials{N,V})
-    duals = similar(x, Dual{T,V,N})
+    duals = similar(x, seed_type(Dual{T,V,N}))
     return GradientConfig{T,V,N,typeof(duals)}(seeds, duals)
 end
 
@@ -156,7 +176,7 @@ function JacobianConfig(f::F,
                         ::Chunk{N} = Chunk(x),
                         ::T = Tag(f, V)) where {F,V,N,T}
     seeds = construct_seeds(Partials{N,V})
-    duals = similar(x, Dual{T,V,N})
+    duals = similar(x, seed_type(Dual{T,V,N}))
     return JacobianConfig{T,V,N,typeof(duals)}(seeds, duals)
 end
 
@@ -182,8 +202,8 @@ function JacobianConfig(f::F,
                         ::Chunk{N} = Chunk(x),
                         ::T = Tag(f, X)) where {F,Y,X,N,T}
     seeds = construct_seeds(Partials{N,X})
-    yduals = similar(y, Dual{T,Y,N})
-    xduals = similar(x, Dual{T,X,N})
+    yduals = similar(y, seed_type(Dual{T,Y,N}))
+    xduals = similar(x, seed_type(Dual{T,X,N}))
     duals = (yduals, xduals)
     return JacobianConfig{T,X,N,typeof(duals)}(seeds, duals)
 end
@@ -195,9 +215,9 @@ Base.eltype(::Type{JacobianConfig{T,V,N,D}}) where {T,V,N,D} = Dual{T,V,N}
 # HessianConfig #
 #################
 
-struct HessianConfig{T,V,N,DG,DJ} <: AbstractConfig{N}
+struct HessianConfig{T,V,N,DJ,G<:GradientConfig} <: AbstractConfig{N}
     jacobian_config::JacobianConfig{T,V,N,DJ}
-    gradient_config::GradientConfig{T,Dual{T,V,N},N,DG}
+    gradient_config::G
 end
 
 """
@@ -223,7 +243,7 @@ function HessianConfig(f::F,
                        chunk::Chunk = Chunk(x),
                        tag = Tag(f, V)) where {F,V}
     jacobian_config = JacobianConfig(f, x, chunk, tag)
-    gradient_config = GradientConfig(f, jacobian_config.duals, chunk, tag)
+    gradient_config = GradientConfig(f, jacobian_config.duals, chunk, Tag{F,eltype(jacobian_config)}())
     return HessianConfig(jacobian_config, gradient_config)
 end
 
@@ -248,10 +268,9 @@ function HessianConfig(f::F,
                        chunk::Chunk = Chunk(x),
                        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)
+    gradient_config = GradientConfig(f, jacobian_config.duals[2], chunk, Tag{F,eltype(jacobian_config)}())
     return HessianConfig(jacobian_config, gradient_config)
 end
 
 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}
+Base.eltype(::Type{HessianConfig{T,V,N,DJ,G}}) where {T,V,N,DJ,G} = eltype(G)

src/deprecated.jl

@@ -0,0 +1,45 @@
+# Accessors without a tag extract the outermost layer of a nested `Dual`, which depends on
+# the order of the tags rather than on the derivative the caller is interested in.
+
+function depwarn_untagged(f::Symbol, replacement::String)
+    Base.depwarn("`ForwardDiff.$f` without a tag is deprecated, use `$replacement` with the tag `T` of the derivative instead.", f)
+end
+
+@inline outer_partials(x, i) = zero(x)
+@inline outer_partials(d::Dual{T}, i) where {T} = partials(T, d, i)
+@inline outer_partials(x, i, j, k...) = outer_partials(outer_partials(x, i), j, k...)
+
+function value(x)
+    depwarn_untagged(:value, "ForwardDiff.value(T, x)")
+    return x
+end
+function value(d::Dual{T}) where {T}
+    depwarn_untagged(:value, "ForwardDiff.value(T, d)")
+    return value(T, d)
+end
+
+function partials(x)
+    depwarn_untagged(:partials, "ForwardDiff.partials(T, x)")
+    return Partials{0,typeof(x)}(tuple())
+end
+function partials(d::Dual{T}) where {T}
+    depwarn_untagged(:partials, "ForwardDiff.partials(T, d)")
+    return partials(T, d)
+end
+function partials(x, i, j...)
+    depwarn_untagged(:partials, "ForwardDiff.partials(T, x, i)")
+    return outer_partials(x, i, j...)
+end
+function partials(::Type{T}, x, i, j, k...) where {T}
+    Base.depwarn("`ForwardDiff.partials(T, x, i, j...)` is deprecated, use `ForwardDiff.partials(S, ForwardDiff.partials(T, x, i), j)` with the tag `S` of the inner derivative instead.", :partials)
+    return outer_partials(partials(T, x, i), j, k...)
+end
+
+function npartials(::Dual{T,V,N}) where {T,V,N}
+    depwarn_untagged(:npartials, "ForwardDiff.npartials(T, typeof(d))")
+    return N
+end
+function npartials(::Type{Dual{T,V,N}}) where {T,V,N}
+    depwarn_untagged(:npartials, "ForwardDiff.npartials(T, D)")
+    return N
+end