From 0091edd55c4ea8fd0acd93cdfc883e357e963058 Mon Sep 17 00:00:00 2001 From: Kristoffer Carlsson Date: Sun, 16 Aug 2026 10:05:43 +0200 Subject: [PATCH 1/2] exploit symmetry in the hessian instead of relying on the jacobian of gradient for the hessian explicitly seed dual numbers and only calculate the upper triangular part when chunking, gives ~2x speedup as input length becomes big --- src/hessian.jl | 151 ++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 125 insertions(+), 26 deletions(-) diff --git a/src/hessian.jl b/src/hessian.jl index 9c755c9a..59c99e94 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -5,7 +5,7 @@ """ 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)`. This method assumes that `isa(f(x), Real)`. @@ -14,8 +14,8 @@ 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 """ @@ -31,29 +31,12 @@ 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 && size(result) == (xlen, xlen) ? 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}()) @@ -64,8 +47,124 @@ because `isa(result, DiffResult)`, `cfg` is constructed as `HessianConfig(f, res 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 && size(hess) == (xlen, xlen) ? 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)?") + +# 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} + 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 + + nblocks = xlen == 0 ? 1 : div(xlen + N - 1, 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. + seed_hessian_chunk!(xdual, x, 1, iseeds, oseeds) + 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) + + for q in 2:nblocks + qoffset = (q - 1) * N + qsize = min(N, xlen - qoffset) + # Off-diagonal blocks: p seeds columns and q seeds rows. + 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. + 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, typeof(value(T, value(T, 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 From 0c04e90a91aa52e02c3ee90dc7ee4e2ec52db5d3 Mon Sep 17 00:00:00 2001 From: Kristoffer Carlsson Date: Mon, 17 Aug 2026 11:57:17 +0200 Subject: [PATCH 2/2] Address symmetric Hessian review feedback --- ext/ForwardDiffStaticArraysExt.jl | 32 +++++++++++++-- src/apiutils.jl | 63 +++++++++++++++-------------- src/config.jl | 11 ++--- src/hessian.jl | 67 +++++++++++-------------------- test/AllocationsTest.jl | 10 +++++ test/GradientTest.jl | 5 +++ test/HessianTest.jl | 66 ++++++++++++++++++++++++++++++ test/JacobianTest.jl | 4 ++ test/SeedTest.jl | 18 +++++++++ 9 files changed, 192 insertions(+), 84 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index bf0ef99a..26abf43b 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/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) diff --git a/src/apiutils.jl b/src/apiutils.jl index 0615fdb3..1d54d7bb 100644 --- a/src/apiutils.jl +++ b/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 diff --git a/src/config.jl b/src/config.jl index 3c6c97e3..58c145f3 100644 --- a/src/config.jl +++ b/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 diff --git a/src/hessian.jl b/src/hessian.jl index 59c99e94..489b9014 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -6,6 +6,8 @@ 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. This method assumes that `isa(f(x), Real)`. @@ -21,8 +23,9 @@ 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)`. @@ -32,7 +35,7 @@ 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 @@ -40,9 +43,10 @@ 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 +55,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 +67,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 +96,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 +141,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} diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 94e7cddd..3a59a5ad 100644 --- a/test/AllocationsTest.jl +++ b/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()) diff --git a/test/GradientTest.jl b/test/GradientTest.jl index bf121239..c9967812 100644 --- a/test/GradientTest.jl +++ b/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) diff --git a/test/HessianTest.jl b/test/HessianTest.jl index 8be72ee5..119fd14d 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -51,11 +51,21 @@ h = [-66.0 -40.0 0.0; @test isapprox(DiffResults.value(out), v) @test isapprox(DiffResults.gradient(out), g) @test isapprox(DiffResults.hessian(out), h) + + # The result-aware and result-independent config constructors are interchangeable. + out = DiffResults.HessianResult(x) + ForwardDiff.hessian!(out, f, x, cfg) + @test isapprox(DiffResults.value(out), v) + @test isapprox(DiffResults.gradient(out), g) + @test isapprox(DiffResults.hessian(out), h) end cfgx = ForwardDiff.HessianConfig(sin, x) @test_throws ForwardDiff.InvalidTagException ForwardDiff.hessian(f, x, cfgx) @test ForwardDiff.hessian(f, x, cfgx, Val{false}()) == ForwardDiff.hessian(f,x) +@test_throws ArgumentError ForwardDiff.hessian(f, x, ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{length(x) + 1}())) +@test_throws DimensionMismatch ForwardDiff.hessian(identity, x) +@test_throws DimensionMismatch ForwardDiff.hessian!(similar(x, 3, 3), identity, x) ######################## @@ -108,10 +118,22 @@ for T in (StaticArrays.SArray, StaticArrays.MArray) @test ForwardDiff.hessian(prod, sx, scfg, Val{false}()) == actual @test ForwardDiff.hessian(prod, sx, scfg, Val{false}()) isa StaticArray + symmetry_f(z) = sum(sin(z[i]) / (1 + z[mod1(i + 1, length(z))]^2) for i in eachindex(z)) + symmetric_static = ForwardDiff.hessian(symmetry_f, sx) + @test symmetric_static == transpose(symmetric_static) + @test symmetric_static == ForwardDiff.hessian(symmetry_f, x) + @test all(iszero, ForwardDiff.hessian(Returns(2.0), sx)) + @test_throws DimensionMismatch ForwardDiff.hessian(identity, sx) + out = similar(x, 9, 9) ForwardDiff.hessian!(out, prod, sx) @test out == actual + out = similar(x, 9, 9) + ForwardDiff.hessian!(out, symmetry_f, sx) + @test out == symmetric_static + @test out == transpose(out) + out = similar(x, 9, 9) ForwardDiff.hessian!(out, prod, sx, cfg) @test out == actual @@ -156,6 +178,50 @@ for T in (StaticArrays.SArray, StaticArrays.MArray) @test DiffResults.hessian(sresult3) == DiffResults.hessian(result) end +@testset "LowerTriangular, UpperTriangular and Diagonal" begin + for n in (3, 5), T in (LowerTriangular, UpperTriangular, Diagonal) + x = T(randn(n, n)) + xlen = ForwardDiff.structural_length(x) + weights = reshape(collect(1.0:n^2), n, n) + objective = x -> dot(weights, abs2.(x)) + expected = diagm(2 .* [weights[idx] for idx in ForwardDiff.structural_eachindex(x)]) + + H = ForwardDiff.hessian(objective, x) + @test size(H) == (xlen, xlen) + @test H == expected + + out = fill(NaN, xlen, xlen) + ForwardDiff.hessian!(out, objective, x) + @test out == expected + + flat = fill(NaN, xlen^2) + ForwardDiff.hessian!(flat, objective, x) + @test reshape(flat, xlen, xlen) == expected + end +end + +@testset "BigFloat with an unassigned input entry" begin + x = Vector{BigFloat}(undef, 10) + hole = 5 + for i in eachindex(x) + i == hole || (x[i] = BigFloat(i)) + end + used = [i for i in eachindex(x) if i != hole] + f(x) = sum(abs2(x[i]) for i in used) + expected = zeros(BigFloat, 10, 10) + for i in used + expected[i, i] = 2 + end + + @test !isassigned(x, hole) + for chunksize in (1, 2, 10) + cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{chunksize}()) + H = ForwardDiff.hessian(f, x, cfg) + @test H isa Matrix{BigFloat} + @test H == expected + end +end + @testset "branches in dot" begin # https://github.com/JuliaDiff/ForwardDiff.jl/issues/551 H = [1 2 3; 4 5 6; 7 8 9]; diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index b6d36180..adc63cd7 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -198,6 +198,10 @@ for T in (StaticArrays.SArray, StaticArrays.MArray) ForwardDiff.jacobian!(out, _diff, sx, scfg) @test out == actual + out = similar(x, 6, 9) + ForwardDiff.jacobian!(out, _diff, sx, scfg, Val{false}()) + @test out == actual + result = DiffResults.JacobianResult(similar(x, 6), x) result = ForwardDiff.jacobian!(result, _diff, x) diff --git a/test/SeedTest.jl b/test/SeedTest.jl index 02b821c3..90ef1858 100644 --- a/test/SeedTest.jl +++ b/test/SeedTest.jl @@ -90,4 +90,22 @@ end end end +@testset "seed_hessian_chunk!: $(nameof(typeof(x)))" for (x, sidx) in SEED_CASES + cfg = ForwardDiff.HessianConfig(nothing, x, ForwardDiff.Chunk{3}()) + duals = cfg.gradient_config.duals + iseeds = cfg.jacobian_config.seeds + oseeds = cfg.gradient_config.seeds + nstruct = length(sidx) + + ForwardDiff.seed_hessian_chunk!(duals, x, 1, nothing, nothing, nstruct) + ForwardDiff.seed_hessian_chunk!(duals, x, 4, iseeds, oseeds) + @test [i for (i, idx) in enumerate(sidx) if !iszero(ForwardDiff.partials(ForwardDiff.value(duals[idx])))] == collect(4:6) + @test [i for (i, idx) in enumerate(sidx) if !iszero(ForwardDiff.partials(duals[idx]))] == collect(4:6) + @test all(idx -> ForwardDiff.value(ForwardDiff.value(duals[idx])) == x[idx], eachindex(x)) + + ForwardDiff.seed_hessian_chunk!(duals, x, 4, nothing, nothing) + @test all(idx -> iszero(ForwardDiff.partials(ForwardDiff.value(duals[idx]))), sidx) + @test all(idx -> iszero(ForwardDiff.partials(duals[idx])), sidx) +end + end # module