exploit symmetry in the hessian

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

@@ -5,7 +5,9 @@
 """
     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)`.
+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.
 
 This method assumes that `isa(f(x), Real)`.
 
@@ -14,15 +16,16 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function hessian(f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T,CHK}
     require_one_based_indexing(x)
     CHK && checktag(T, f, x)
-    ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}())
-    return jacobian(∇f, x, cfg.jacobian_config, Val{false}())
+    H, _ = symmetric_hessian(f, x, cfg, nothing)
+    return H
 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.
 
 This method assumes that `isa(f(x), Real)`.
 
@@ -31,41 +34,116 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK}
     require_one_based_indexing(result, x)
     CHK && checktag(T, f, x)
-    ∇f = y -> gradient(f, y, cfg.gradient_config, Val{false}())
-    jacobian!(result, ∇f, x, cfg.jacobian_config, Val{false}())
+    xlen = structural_length(x)
+    H = result isa AbstractMatrix ? result : reshape(result, xlen, xlen)
+    symmetric_hessian!(H, f, x, cfg, nothing)
     return result
 end
 
-
-# We use this struct below instead of an
-# equivalent closure in order to avoid
-# JuliaLang/julia#15276-related performance
-# issues. See #316.
-mutable struct InnerGradientForHess{R,C,F}
-    result::R
-    cfg::C
-    f::F
-end
-
-function (g::InnerGradientForHess)(y, z)
-    inner_result = DiffResult(zero(eltype(y)), y)
-    gradient!(inner_result, g.f, z, g.cfg.gradient_config, Val{false}())
-    g.result = DiffResults.value!(g.result, value(DiffResults.value(inner_result)))
-    return y
-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.
 """
 function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,CHK}
+    require_one_based_indexing(x)
     CHK && checktag(T, f, x)
-    ∇f! = InnerGradientForHess(result, cfg, f)
-    jacobian!(DiffResults.hessian(result), ∇f!, DiffResults.gradient(result), x, cfg.jacobian_config, Val{false}())
-    return ∇f!.result
+    xlen = structural_length(x)
+    hess = DiffResults.hessian(result)
+    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
+end
+
+############################
+# symmetric Hessian kernel #
+############################
+
+const HESSIAN_ERROR = DimensionMismatch("hessian(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?")
+
+# 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}
+    for r in 1:rsize
+        drow = partials(T, ydual, r)
+        cstart = roffset == coffset ? r : 1
+        for c in cstart:csize
+            h = partials(T, drow, c)
+            H[roffset + r, coffset + c] = h
+            H[coffset + c, roffset + r] = h
+        end
+    end
+    return H
+end
+
+# The inner partials of a diagonal block contain the corresponding gradient chunk.
+extract_hessian_gradient_chunk!(::Type{T}, ::Nothing, ydual, index, chunksize) where {T} = nothing
+extract_hessian_gradient_chunk!(::Type{T}, grad, ydual, index, chunksize) where {T} =
+    extract_gradient_chunk!(T, grad, value(T, ydual), index, chunksize)
+
+# Evaluate one pair of chunks at a time using nested duals. Only one triangle of block
+# pairs is evaluated; the other is filled by symmetry (see #836).
+function symmetric_hessian_expr(result_definition::Expr)
+    return quote
+        xlen = structural_length(x)
+        if xlen < N
+            throw(ArgumentError(lazy"chunk size cannot be greater than ForwardDiff.structural_length(x) ($(N) > $(structural_length(x)))"))
+        end
+
+        # `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
+
+        # 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)
+        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. 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)
+                ydual = f(xdual)
+                extract_hessian_chunk!(T, H, ydual, qoffset, poffset, qsize, N)
+                seed_hessian_chunk!(xdual, x, poffset + 1, nothing, nothing)
+            end
+            # 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)
+            extract_hessian_gradient_chunk!(T, grad, ydual, qoffset + 1, qsize)
+            seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, nothing, qsize)
+        end
+
+        return H, ydual1
+    end
+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, 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}
+    $(symmetric_hessian_expr(:()))
 end

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)