Address symmetric Hessian review feedback

Author: KristofferC

ext/ForwardDiffStaticArraysExt.jl

@@ -7,7 +7,7 @@ using ForwardDiff: Dual, partials, npartials, Partials, GradientConfig, Jacobian
                    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
+                   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,11 +107,34 @@ end
 end
 
 # Hessian
-ForwardDiff.hessian(f::F, x::StaticArray) where {F} = jacobian(Base.Fix1(gradient, f), x)
+@inline function extract_hessian(::Type{T}, ydual::Partials, x::StaticArray) where {T}
+    H = extract_jacobian(T, ydual, x)
+    return typeof(H)(Symmetric(H, :U))
+end
+
+@inline function extract_hessian(::Type{T}, ydual::Partials{0}, x::S) where {T,S<:StaticArray}
+    R = StaticArrays.similar_type(S, valtype(T, eltype(ydual)), Size(length(x), length(x)))
+    return zero(R)
+end
+
+@inline function ForwardDiff.hessian(f::F, x::StaticArray) where {F}
+    T = typeof(Tag(f, eltype(x)))
+    ydual = f(dualize(T, dualize(T, x)))
+    ydual isa Real || throw(HESSIAN_ERROR)
+    return extract_hessian(T, partials(T, 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 = typeof(Tag(f, eltype(x)))
+    ydual = f(dualize(T, dualize(T, x)))
+    ydual isa Real || throw(HESSIAN_ERROR)
+    H = result isa AbstractMatrix ? result : reshape(result, length(x), length(x))
+    ForwardDiff.extract_hessian_chunk!(T, H, ydual, 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))
 
@@ -123,9 +146,10 @@ function ForwardDiff.hessian!(result::ImmutableDiffResult, f::F, x::StaticArray)
     d1 = dualize(T, x)
     d2 = dualize(T, d1)
     fd2 = f(d2)
+    fd2 isa Real || throw(HESSIAN_ERROR)
     val = value(T,value(T,fd2))
     grad = extract_gradient(T,value(T,fd2), x)
-    hess = extract_jacobian(T,partials(T,fd2), x)
+    hess = extract_hessian(T,partials(T,fd2), x)
     result = DiffResults.hessian!(result, hess)
     result = DiffResults.gradient!(result, grad)
     result = DiffResults.value!(result, val)

src/apiutils.jl

@@ -88,14 +88,22 @@ 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
+
+# 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)
             end
@@ -106,38 +114,31 @@ 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.
+function seed_hessian_chunk!(duals::AbstractArray{Dual{T,Dual{T,V,N},N}}, x, index,
+                             iseeds::Union{Nothing,NTuple{N,Partials{N,V}}},
+                             oseeds::Union{Nothing,NTuple{N,Partials{N,Dual{T,V,N}}}},
+                             chunksize = N) where {T,V,N}
+    izero = zero(Partials{N,V})
+    ozero = zero(Partials{N,Dual{T,V,N}})
+    idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), chunksize)
+    return _seed!(duals, x, idxs) do value, i
+        inner = Dual{T,V,N}(value, iseeds === nothing ? izero : iseeds[i])
+        Dual{T,Dual{T,V,N},N}(inner, oseeds === nothing ? ozero : oseeds[i])
     end
-    return duals
 end

src/config.jl

@@ -207,10 +207,9 @@ Return a `HessianConfig` instance based on the type of `f` and type/shape of the
 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
@@ -234,7 +233,9 @@ Return a `HessianConfig` instance based on the type of `f`, types/storage in `re
 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`.
+`ForwardDiff.hessian` and `ForwardDiff.hessian!`. It is interchangeable with a config
+constructed via `ForwardDiff.HessianConfig(f, x, chunk)`; this constructor retains the
+result-aware form 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

src/hessian.jl

@@ -6,6 +6,9 @@
     ForwardDiff.hessian(f, x::AbstractArray, cfg::HessianConfig = HessianConfig(f, x), check=Val{true}())
 
 Return `H(f)` evaluated at `x`, assuming `f` is called as `f(x)`.
+The returned Hessian is exactly symmetric: its two triangles are filled from the same
+derivative values. Its size is `(structural_length(x), structural_length(x))`; for
+structured matrix inputs, only structurally stored elements are treated as variables.
 
 This method assumes that `isa(f(x), Real)`.
 
@@ -21,8 +24,10 @@ end
 """
     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)`.
+Compute `H(f)` evaluated at `x` and store the result(s) in `result`, assuming `f` is
+called as `f(x)`. The stored Hessian is exactly symmetric: its two triangles are filled
+from the same derivative values. `result` must store a
+`structural_length(x)`-by-`structural_length(x)` matrix.
 
 This method assumes that `isa(f(x), Real)`.
 
@@ -32,17 +37,18 @@ function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianCon
     require_one_based_indexing(result, x)
     CHK && checktag(T, f, x)
     xlen = structural_length(x)
-    H = result isa AbstractMatrix && size(result) == (xlen, xlen) ? result : reshape(result, xlen, xlen)
+    H = result isa AbstractMatrix ? result : reshape(result, xlen, xlen)
     symmetric_hessian!(H, f, x, cfg, nothing)
     return result
 end
 
 """
     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)`.
+Exactly like `ForwardDiff.hessian!(result::AbstractArray, f, x::AbstractArray, cfg::HessianConfig)`,
+but also stores the value and gradient in `result`. The default `cfg` is constructed as
+`HessianConfig(f, result, x)`, though a config constructed as `HessianConfig(f, x)` may also
+be used.
 
 Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care.
 """
@@ -51,7 +57,7 @@ function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig
     CHK && checktag(T, f, x)
     xlen = structural_length(x)
     hess = DiffResults.hessian(result)
-    H = hess isa AbstractMatrix && size(hess) == (xlen, xlen) ? hess : reshape(hess, xlen, xlen)
+    H = hess isa AbstractMatrix ? hess : reshape(hess, xlen, xlen)
     _, ydual = symmetric_hessian!(H, f, x, cfg, DiffResults.gradient(result))
     result = DiffResults.value!(result, value(T, value(T, ydual)))
     return result
@@ -63,32 +69,6 @@ end
 
 const HESSIAN_ERROR = DimensionMismatch("hessian(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?")
 
-# Seed a chunk in either layer of the nested duals. A `nothing` seed clears that layer.
-function seed_hessian_chunk!(duals::AbstractArray{Dual{T,Dual{T,V,N},N}}, x, index,
-                             iseeds::Union{Nothing,NTuple{N,Partials{N,V}}},
-                             oseeds::Union{Nothing,NTuple{N,Partials{N,Dual{T,V,N}}}},
-                             chunksize = N) where {T,V,N}
-    izero = zero(Partials{N,V})
-    ozero = zero(Partials{N,Dual{T,V,N}})
-    idxs = Iterators.drop(structural_eachindex(duals, x), index - 1)
-    if isbitstype(V)
-        for (i, idx) in zip(1:chunksize, idxs)
-            inner = Dual{T,V,N}(x[idx], iseeds === nothing ? izero : iseeds[i])
-            duals[idx] = Dual{T,Dual{T,V,N},N}(inner, oseeds === nothing ? ozero : oseeds[i])
-        end
-    else
-        for (i, idx) in zip(1:chunksize, idxs)
-            if isassigned(x, idx)
-                inner = Dual{T,V,N}(x[idx], iseeds === nothing ? izero : iseeds[i])
-                duals[idx] = Dual{T,Dual{T,V,N},N}(inner, oseeds === nothing ? ozero : oseeds[i])
-            else
-                Base._unsetindex!(duals, idx)
-            end
-        end
-    end
-    return duals
-end
-
 # Copy a block from the nested partials and fill its transpose. On diagonal blocks, read
 # only the upper triangle so the result is exactly symmetric.
 function extract_hessian_chunk!(::Type{T}, H, ydual, roffset, coffset, rsize, csize) where {T}
@@ -118,38 +98,39 @@ function symmetric_hessian_expr(result_definition::Expr)
             throw(ArgumentError(lazy"chunk size cannot be greater than ForwardDiff.structural_length(x) ($(N) > $(structural_length(x)))"))
         end
 
-        nblocks = xlen == 0 ? 1 : div(xlen + N - 1, N)
+        # `N == 0` only for empty inputs, which still need one evaluation to determine the
+        # output type and value.
+        nblocks = xlen == 0 ? 1 : cld(xlen, N)
 
         xdual = cfg.gradient_config.duals
         iseeds = cfg.jacobian_config.seeds
         oseeds = cfg.gradient_config.seeds
 
-        # Keep all unseeded blocks at zero between evaluations.
-        seed_hessian_chunk!(xdual, x, 1, nothing, nothing, xlen)
-
-        # The first evaluation determines the output type.
+        # The first evaluation determines the output type. Seeding the first block and clearing
+        # the untouched tail partitions the fresh buffer, so every element is initialized once.
         seed_hessian_chunk!(xdual, x, 1, iseeds, oseeds)
+        seed_hessian_chunk!(xdual, x, N + 1, nothing, nothing, xlen - N)
         ydual1 = f(xdual)
         ydual1 isa Real || throw(HESSIAN_ERROR)
         $(result_definition)
         extract_hessian_chunk!(T, H, ydual1, 0, 0, N, N)
         extract_hessian_gradient_chunk!(T, grad, ydual1, 1, N)
-        seed_hessian_chunk!(xdual, x, 1, nothing, nothing)
+        nblocks > 1 && seed_hessian_chunk!(xdual, x, 1, nothing, nothing)
 
         for q in 2:nblocks
             qoffset = (q - 1) * N
             qsize = min(N, xlen - qoffset)
-            # Off-diagonal blocks: p seeds columns and q seeds rows.
+            # Off-diagonal blocks: p seeds columns and q seeds rows. The outer seeds for q
+            # remain unchanged throughout this loop.
+            seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, oseeds, qsize)
             for p in 1:(q - 1)
                 poffset = (p - 1) * N
                 seed_hessian_chunk!(xdual, x, poffset + 1, iseeds, nothing)
-                seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, oseeds, qsize)
                 ydual = f(xdual)
                 extract_hessian_chunk!(T, H, ydual, qoffset, poffset, qsize, N)
                 seed_hessian_chunk!(xdual, x, poffset + 1, nothing, nothing)
-                seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, nothing, qsize)
             end
-            # Diagonal blocks seed both layers.
+            # The diagonal block adds q's inner seeds while retaining its outer seeds.
             seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, oseeds, qsize)
             ydual = f(xdual)
             extract_hessian_chunk!(T, H, ydual, qoffset, qoffset, qsize, qsize)
@@ -162,7 +143,7 @@ function symmetric_hessian_expr(result_definition::Expr)
 end
 
 @eval function symmetric_hessian(f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N}
-    $(symmetric_hessian_expr(:(H = similar(x, typeof(value(T, value(T, ydual1))), xlen, xlen))))
+    $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), xlen, xlen))))
 end
 
 @eval function symmetric_hessian!(H, f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N}

test/AllocationsTest.jl

@@ -29,6 +29,16 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F
     allocs_szp!(duals, x, 1, 4)
     @test iszero(allocs_szp!(duals, x, 1, 4))
 
+    hcfg = ForwardDiff.HessianConfig(nothing, x)
+    hduals = hcfg.gradient_config.duals
+    iseeds = hcfg.jacobian_config.seeds
+    oseeds = hcfg.gradient_config.seeds
+    allocs_hseed!(args...) = @allocated ForwardDiff.seed_hessian_chunk!(args...)
+    allocs_hseed!(hduals, x, 1, iseeds, oseeds)
+    @test iszero(allocs_hseed!(hduals, x, 1, iseeds, oseeds))
+    allocs_hseed!(hduals, x, 1, nothing, nothing, 4)
+    @test iszero(allocs_hseed!(hduals, x, 1, nothing, nothing, 4))
+
     allocs_convert_test_574() = @allocated convert_test_574()
     allocs_convert_test_574()
     @test iszero(allocs_convert_test_574())

test/GradientTest.jl

@@ -56,6 +56,7 @@ 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_throws ArgumentError ForwardDiff.gradient(f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{length(x) + 1}()))
 
 
 ########################
@@ -115,6 +116,10 @@ end
     ForwardDiff.gradient!(out, prod, sx, scfg)
     @test out == actual
 
+    out = similar(x)
+    ForwardDiff.gradient!(out, prod, sx, scfg, Val{false}())
+    @test out == actual
+
     result = DiffResults.GradientResult(x)
     result = ForwardDiff.gradient!(result, prod, x)