Exploit symmetry in the Hessian

ext/ForwardDiffStaticArraysExt.jl

@@ -5,9 +5,9 @@ using ForwardDiff.LinearAlgebra
 using ForwardDiff.DiffResults
 using ForwardDiff: Dual, partials, npartials, Partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk,
                    gradient, hessian, jacobian, gradient!, hessian!, jacobian!,
-                   extract_gradient!, extract_jacobian!, extract_value!,
-                   vector_mode_gradient, vector_mode_gradient!,
-                   vector_mode_jacobian, vector_mode_jacobian!, valtype, value
+                   extract_gradient!, extract_jacobian!, extract_value!, structural_linearindices,
+                   vector_mode_gradient, vector_mode_gradient!, outer_tag,
+                   vector_mode_jacobian, vector_mode_jacobian!, HESSIAN_ERROR, valtype, value
 using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult
 
 @generated function dualize(::Type{T}, x::StaticArray) where T
@@ -107,25 +107,58 @@ end
 end
 
 # Hessian
-ForwardDiff.hessian(f::F, x::StaticArray) where {F} = jacobian(Base.Fix1(gradient, f), x)
+@inline function extract_hessian(::Type{T}, ::Type{TO}, ydual::Dual{TO,<:Dual{T}}, x::StaticArray) where {T,TO}
+    H = extract_jacobian(T, partials(TO, ydual), x)
+    return typeof(H)(Symmetric(H, :U))
+end
+
+# A result that never picked up both perturbations has no second derivatives, and offers neither
+# the `length(x)` rows the method above reads nor, for an `f` ignoring its argument, any at all.
+@inline function extract_hessian(::Type{T}, ::Type{TO}, ydual, x::S) where {T,TO,S<:StaticArray}
+    R = StaticArrays.similar_type(S, valtype(T, valtype(TO, typeof(ydual))),
+                                  Size(length(x), length(x)))
+    return zero(R)
+end
+
+# The layers need distinct tags; see `ForwardDiff.outer_tag`.
+@inline function hessian_tags(f::F, x::StaticArray) where {F}
+    T = typeof(Tag(f, eltype(x)))
+    return T, outer_tag(T, Dual{T,eltype(x),length(x)})
+end
+
+@inline function ForwardDiff.hessian(f::F, x::StaticArray) where {F}
+    T, TO = hessian_tags(f, x)
+    ydual = f(dualize(TO, dualize(T, x)))
+    ydual isa Real || throw(HESSIAN_ERROR)
+    return extract_hessian(T, TO, ydual, x)
+end
+
 ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig) where {F} = hessian(f, x)
 ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig, ::Val) where {F} = hessian(f, x)
 
-ForwardDiff.hessian!(result::AbstractArray, f::F, x::StaticArray) where {F} = jacobian!(result, Base.Fix1(gradient, f), x)
+@inline function ForwardDiff.hessian!(result::AbstractArray, f::F, x::StaticArray) where {F}
+    T, TO = hessian_tags(f, x)
+    ydual = f(dualize(TO, dualize(T, x)))
+    ydual isa Real || throw(HESSIAN_ERROR)
+    H = ForwardDiff.reshape_hessian(result, x)
+    ForwardDiff.extract_hessian_chunk!(T, TO, H, ydual, structural_linearindices(x), 0, 0, length(x), length(x))
+    return result
+end
 
 ForwardDiff.hessian!(result::MutableDiffResult, f::F, x::StaticArray) where {F} = hessian!(result, f, x, HessianConfig(f, result, x))
 
 ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray, cfg::HessianConfig) where {F} = hessian!(result, f, x)
 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)))
+    T, TO = hessian_tags(f, x)
     d1 = dualize(T, x)
-    d2 = dualize(T, d1)
+    d2 = dualize(TO, 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)
+    fd2 isa Real || throw(HESSIAN_ERROR)
+    val = value(T,value(TO,fd2))
+    grad = extract_gradient(T,value(TO,fd2), x)
+    hess = extract_hessian(T,TO,fd2, x)
     result = DiffResults.hessian!(result, hess)
     result = DiffResults.gradient!(result, grad)
     result = DiffResults.value!(result, val)

src/apiutils.jl

@@ -70,6 +70,59 @@ function structural_eachindex(x::Diagonal, y::AbstractArray)
     return diagind(x)
 end
 
+function check_structural_size(duals, x)
+    if size(duals) != size(x)
+        throw(DimensionMismatch(lazy"the config was built for an array of size $(size(duals)) and cannot be used with an array of size $(size(x))"))
+    end
+    return nothing
+end
+
+# The positions of `structural_eachindex`, in the same order, as linear indices of `x`. The two
+# argument form is only ever given a config's work buffer and the input it is used with.
+structural_linearindices(x::AbstractArray) = structural_linearindices(x, x)
+function structural_linearindices(duals::AbstractArray, x::AbstractArray)
+    require_one_based_indexing(duals, x)
+    check_structural_size(duals, x)
+    return Base.OneTo(length(duals))
+end
+function structural_linearindices(duals::UpperTriangular, x::AbstractArray)
+    require_one_based_indexing(duals, x)
+    check_structural_size(duals, x)
+    n = size(duals, 1)
+    indices = Vector{Int}(undef, structural_length(duals))
+    k = idx = 0
+    for j in 1:n
+        for _ in 1:j
+            indices[k += 1] = (idx += 1)
+        end
+        idx += n - j
+    end
+    return indices
+end
+function structural_linearindices(duals::LowerTriangular, x::AbstractArray)
+    require_one_based_indexing(duals, x)
+    check_structural_size(duals, x)
+    n = size(duals, 1)
+    indices = Vector{Int}(undef, structural_length(duals))
+    k = idx = 0
+    for j in 1:n
+        for _ in j:n
+            indices[k += 1] = (idx += 1)
+        end
+        idx += j
+    end
+    return indices
+end
+function structural_linearindices(duals::Diagonal, x::AbstractArray)
+    require_one_based_indexing(duals, x)
+    check_structural_size(duals, x)
+    n = size(duals, 1)
+    return range(1; step = n + 1, length = n)
+end
+
+# The `count` positions starting at structural position `index`.
+structural_chunk(indices, index, count) = view(indices, index:(index + count - 1))
+
 # 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.
@@ -88,16 +141,32 @@ end
 
 function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N}
     seed = zero(Partials{N,V})
+    return _seed!(duals, x, idxs) do value, _
+        Dual{T,V,N}(value, seed)
+    end
+end
+
+# `Base._unsetindex!` is implemented for `Array` alone: for a linear index its `AbstractArray`
+# fallback recurses forever, and it has no `CartesianIndex` method at all.
+_unsetindex!(duals::Array, idx) = Base._unsetindex!(duals, idx)
+_unsetindex!(duals::AbstractArray, idx) = throw(ArgumentError(LazyString(
+    "cannot differentiate at an input with an unassigned entry at index ", idx,
+    ": that would leave an entry of the ", nameof(typeof(duals)),
+    " work buffer unassigned, which is only possible for an Array")))
+
+# Write a sequence of duals while preserving unassigned entries in arrays whose element type is not
+# stored inline. `make_dual` receives the primal value and its one-based position in `idxs`.
+@inline function _seed!(make_dual::F, duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {F,T,V,N}
     if isbitstype(V)
-        for idx in idxs
-            duals[idx] = Dual{T,V,N}(x[idx], seed)
+        for (i, idx) in enumerate(idxs)
+            duals[idx] = make_dual(x[idx], i)
         end
     else
-        for idx in idxs
+        for (i, idx) in enumerate(idxs)
             if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seed)
+                duals[idx] = make_dual(x[idx], i)
             else
-                Base._unsetindex!(duals, idx)
+                _unsetindex!(duals, idx)
             end
         end
     end
@@ -106,38 +175,32 @@ end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, 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])
-        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])
-            else
-                Base._unsetindex!(duals, idx)
-            end
-        end
+    idxs = Iterators.take(structural_eachindex(duals, x), N)
+    return _seed!(duals, x, idxs) do value, i
+        Dual{T,V,N}(value, seeds[i])
     end
-    return duals
 end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, 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])
-        end
-    else
-        for (i, idx) in zip(1:chunksize, idxs)
-            if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
-            else
-                Base._unsetindex!(duals, idx)
-            end
-        end
+    idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), offset), chunksize)
+    return _seed!(duals, x, idxs) do value, i
+        Dual{T,V,N}(value, seeds[i])
+    end
+end
+
+# Seed a chunk in either layer of nested duals. A `nothing` seed clears that layer;
+# `seed_zero_partials!` cannot, as it would pass the primal where a nested `Dual` is wanted.
+function seed_hessian_chunk!(duals::AbstractArray{Dual{TO,Dual{T,V,N},N}}, x, indices, index,
+                             iseeds::Union{Nothing,NTuple{N,Partials{N,V}}},
+                             oseeds::Union{Nothing,NTuple{N,Partials{N,Dual{T,V,N}}}},
+                             chunksize = N) where {TO,T,V,N}
+    izero = iseeds === nothing ? zero(Partials{N,V}) : nothing
+    ozero = oseeds === nothing ? zero(Partials{N,Dual{T,V,N}}) : nothing
+    idxs = structural_chunk(indices, index, chunksize)
+    return _seed!(duals, x, idxs) do value, i
+        inner = Dual{T,V,N}(value, iseeds === nothing ? izero : iseeds[i])
+        Dual{TO,Dual{T,V,N},N}(inner, oseeds === nothing ? ozero : oseeds[i])
     end
-    return duals
 end

src/config.jl

@@ -195,22 +195,27 @@ 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}
-    jacobian_config::JacobianConfig{T,V,N,DJ}
-    gradient_config::GradientConfig{T,Dual{T,V,N},N,DG}
+struct HessianConfig{T,TO,V,N,D} <: AbstractConfig{N}
+    iseeds::NTuple{N,Partials{N,V}}
+    oseeds::NTuple{N,Partials{N,Dual{T,V,N}}}
+    duals::D
 end
 
+# The layers need distinct tags, or `value`/`partials` inside `f` cannot tell them apart (#845).
+# `tagcount` fixes the ordering here rather than at the first comparison, as `Tag` does.
+outer_tag(::Type{T}, ::Type{D}) where {T,D} = (tagcount(Tag{T,D}); Tag{T,D})
+outer_tag(::Type{Nothing}, ::Type) = Nothing
+
 """
     ForwardDiff.HessianConfig(f, x::AbstractArray, chunk::Chunk = Chunk(x))
 
 Return a `HessianConfig` instance based on the type of `f` and type/shape of the input
 vector `x`.
 
 The returned `HessianConfig` instance contains all the work buffers required by
-`ForwardDiff.hessian` and `ForwardDiff.hessian!`. For the latter, the buffers are
-configured for the case where the `result` argument is an `AbstractArray`. If
-it is a `DiffResult`, the `HessianConfig` should instead be constructed via
-`ForwardDiff.HessianConfig(f, result, x, chunk)`.
+`ForwardDiff.hessian` and `ForwardDiff.hessian!`, including when the latter stores into a
+`DiffResult`. The `ForwardDiff.HessianConfig(f, result, x, chunk)` constructor may also be
+used with any of these methods.
 
 If `f` is `nothing` instead of the actual target function, then the returned instance can
 be used with any target function. However, this will reduce ForwardDiff's ability to catch
@@ -220,11 +225,13 @@ This constructor does not store/modify `x`.
 """
 function HessianConfig(f::F,
                        x::AbstractArray{V},
-                       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)
-    return HessianConfig(jacobian_config, gradient_config)
+                       ::Chunk{N} = Chunk(x),
+                       ::T = Tag(f, V)) where {F,V,N,T}
+    iseeds = construct_seeds(Partials{N,V})
+    oseeds = construct_seeds(Partials{N,Dual{T,V,N}})
+    TO = outer_tag(T, Dual{T,V,N})
+    duals = similar(x, Dual{TO,Dual{T,V,N},N})
+    return HessianConfig{T,TO,V,N,typeof(duals)}(iseeds, oseeds, duals)
 end
 
 """
@@ -233,25 +240,20 @@ end
 Return a `HessianConfig` instance based on the type of `f`, types/storage in `result`, and
 type/shape of the input vector `x`.
 
-The returned `HessianConfig` instance contains all the work buffers required by
-`ForwardDiff.hessian!` for the case where the `result` argument is an `DiffResult`.
+Equivalent to `ForwardDiff.HessianConfig(f, x, chunk)`: the work buffers do not depend on
+`result`. The result-aware form is retained for compatibility.
 
 If `f` is `nothing` instead of the actual target function, then the returned instance can
 be used with any target function. However, this will reduce ForwardDiff's ability to catch
 and prevent perturbation confusion (see https://github.com/JuliaDiff/ForwardDiff.jl/issues/83).
 
-This constructor does not store/modify `x`.
+This constructor does not store/modify `result` or `x`.
 """
-function HessianConfig(f::F,
-                       result::DiffResult,
-                       x::AbstractArray{V},
-                       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)
-    return HessianConfig(jacobian_config, gradient_config)
-end
+HessianConfig(f::F,
+              ::DiffResult,
+              x::AbstractArray{V},
+              chunk::Chunk = Chunk(x),
+              tag = Tag(f, V)) where {F,V} = HessianConfig(f, x, chunk, tag)
 
 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,TO,V,N,D}}) where {T,TO,V,N,D} = Dual{TO,Dual{T,V,N},N}