diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index bf0ef99a..bafd266d 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/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,11 +107,43 @@ 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)) @@ -119,13 +151,14 @@ 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))) + 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) diff --git a/src/apiutils.jl b/src/apiutils.jl index 0615fdb3..c7869c95 100644 --- a/src/apiutils.jl +++ b/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 diff --git a/src/config.jl b/src/config.jl index 3c6c97e3..96dd73ee 100644 --- a/src/config.jl +++ b/src/config.jl @@ -195,11 +195,17 @@ 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)) @@ -207,10 +213,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 @@ -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} diff --git a/src/hessian.jl b/src/hessian.jl index 9c755c9a..8bc403da 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -5,7 +5,11 @@ """ 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)`. +Multidimensional arrays are flattened in iteration order: the array `H(f)` has shape +`length(x) × length(x)`, and its elements are `H(f)[j,k] = ∂²f(x)/∂x[j]∂x[k]`. +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 +18,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 +36,179 @@ 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}()) + symmetric_hessian!(reshape_hessian(result, x), 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} +function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T,TO} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,TO,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 + _, ydual = symmetric_hessian!(reshape_hessian(result, x), f, x, cfg, + DiffResults.gradient(result)) + result = DiffResults.value!(result, value(T, value(TO, 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)?") + +# Mirrors `reshape_jacobian`. The sweep writes the result entry by entry, so nothing else checks it. +function reshape_hessian(result::AbstractMatrix, x) + require_one_based_indexing(result) + if size(result) != (length(x), length(x)) + throw(DimensionMismatch(lazy"cannot store the $(length(x))×$(length(x)) Hessian in a result of size $(size(result))")) + end + return result +end +function reshape_hessian(result::AbstractArray, x) + require_one_based_indexing(result) + if length(result) != length(x)^2 + throw(DimensionMismatch(lazy"cannot store the $(length(x))×$(length(x)) Hessian in a result of length $(length(result))")) + end + return reshape(result, length(x), length(x)) +end +function reshape_hessian(result::DiffResult, x) + structural_eachindex(DiffResults.gradient(result), x) + return reshape_hessian(DiffResults.hessian(result), x) +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. `indices` maps a block +# position to its row and column, both being linear indices of `x`. +function extract_hessian_chunk!(::Type{T}, ::Type{TO}, H, ydual::Dual{TO,<:Dual{T}}, indices, roffset, coffset, rsize, csize) where {T,TO} + rows = structural_chunk(indices, roffset + 1, rsize) + cols = structural_chunk(indices, coffset + 1, csize) + for r in 1:rsize + drow = partials(TO, ydual, r) + i = rows[r] + cstart = roffset == coffset ? r : 1 + for c in cstart:csize + h = partials(T, drow, c) + j = cols[c] + H[i, j] = h + H[j, i] = h + end + end + return H +end + +# Without both perturbations the block is zero. +function extract_hessian_chunk!(::Type{T}, ::Type{TO}, H, ydual, indices, roffset, coffset, rsize, csize) where {T,TO} + rows = structural_chunk(indices, roffset + 1, rsize) + cols = structural_chunk(indices, coffset + 1, csize) + h = zero(valtype(T, valtype(TO, typeof(ydual)))) + for j in cols, i in rows + H[i, j] = h + H[j, i] = h + end + return H +end + +# The inner partials of a diagonal block contain the corresponding gradient chunk. +# TODO: delegate to `extract_gradient_chunk!` once it takes its positions from `x` (#838). +extract_hessian_gradient_chunk!(::Type{T}, ::Type{TO}, ::Nothing, ydual, indices, index, chunksize) where {T,TO} = nothing +function extract_hessian_gradient_chunk!(::Type{T}, ::Type{TO}, grad, ydual, indices, index, chunksize) where {T,TO} + dual = value(TO, ydual) + for (i, idx) in enumerate(structural_chunk(indices, index, chunksize)) + grad[idx] = partials(T, dual, i) + end + return grad +end + +# 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 the number of differentiated entries of x ($(N) > $(xlen))")) + 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.duals + iseeds = cfg.iseeds + oseeds = cfg.oseeds + indices = structural_linearindices(xdual, x) + + # 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, indices, 1, iseeds, oseeds) + seed_hessian_chunk!(xdual, x, indices, N + 1, nothing, nothing, xlen - N) + ydual1 = f(xdual) + ydual1 isa Real || throw(HESSIAN_ERROR) + Vout = valtype(T, valtype(TO, typeof(ydual1))) + $(result_definition) + # A second derivative needs both perturbations, a first derivative only the inner one: + # what the result does not carry vanishes identically. + zero_hessian = !(ydual1 isa Dual{TO,<:Dual{T}}) + zero_gradient = zero_hessian && !(ydual1 isa Dual{T}) + + # Zero what no block writes: a derivative that vanishes, and the rows and columns of + # the structural zeros of `x`, which are not variables. + if zero_hessian || xlen != length(x) + fill!(H, zero(Vout)) + end + if grad !== nothing && (zero_gradient || xlen != length(x)) + fill!(grad, zero(Vout)) + end + # off-diagonal blocks find second derivatives, diagonal ones also the gradient + if zero_hessian && (zero_gradient || grad === nothing) + return H, ydual1 + end + extract_hessian_chunk!(T, TO, H, ydual1, indices, 0, 0, N, N) + extract_hessian_gradient_chunk!(T, TO, grad, ydual1, indices, 1, N) + if nblocks > 1 + seed_hessian_chunk!(xdual, x, indices, 1, nothing, nothing) + end + + for q in 2:nblocks + qoffset = (q - 1) * N + qsize = min(N, xlen - qoffset) + if !zero_hessian + # Outer-i inner-j and outer-j inner-i round differently, so the outer layer always + # takes the earlier position -- else the result would depend on the chunk size. + # q's inner seeds remain unchanged throughout this loop. + seed_hessian_chunk!(xdual, x, indices, qoffset + 1, iseeds, nothing, qsize) + for p in 1:(q - 1) + poffset = (p - 1) * N + seed_hessian_chunk!(xdual, x, indices, poffset + 1, nothing, oseeds) + ydual = f(xdual) + extract_hessian_chunk!(T, TO, H, ydual, indices, poffset, qoffset, N, qsize) + seed_hessian_chunk!(xdual, x, indices, poffset + 1, nothing, nothing) + end + end + # The diagonal block adds q's outer seeds while retaining its inner seeds. + seed_hessian_chunk!(xdual, x, indices, qoffset + 1, iseeds, oseeds, qsize) + ydual = f(xdual) + extract_hessian_chunk!(T, TO, H, ydual, indices, qoffset, qoffset, qsize, qsize) + extract_hessian_gradient_chunk!(T, TO, grad, ydual, indices, qoffset + 1, qsize) + seed_hessian_chunk!(xdual, x, indices, qoffset + 1, nothing, nothing, qsize) + end + + return H, ydual1 + end +end + +@eval function symmetric_hessian(f::F, x, cfg::HessianConfig{T,TO,V,N}, grad) where {F,T,TO,V,N} + $(symmetric_hessian_expr(:(H = similar(x, Vout, length(x), length(x))))) +end + +@eval function symmetric_hessian!(H, f::F, x, cfg::HessianConfig{T,TO,V,N}, grad) where {F,T,TO,V,N} + $(symmetric_hessian_expr(:())) end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 94e7cddd..3a51885c 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -7,7 +7,7 @@ include(joinpath(dirname(@__FILE__), "utils.jl")) convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,Float64,8},4},2}, 1.3) -@testset "Test seed!/seed_zero_partials! allocations" begin +@testset "Test seeding allocations" begin x = rand(1000) cfg = ForwardDiff.GradientConfig(nothing, x) duals = cfg.duals @@ -29,6 +29,27 @@ 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.duals + iseeds = hcfg.iseeds + oseeds = hcfg.oseeds + hindices = ForwardDiff.structural_linearindices(hduals, x) + allocs_hseed!(args...) = @allocated ForwardDiff.seed_hessian_chunk!(args...) + for i in (iseeds, nothing), o in (oseeds, nothing) + allocs_hseed!(hduals, x, hindices, 1, i, o) + @test iszero(allocs_hseed!(hduals, x, hindices, 1, i, o)) + end + allocs_hseed!(hduals, x, hindices, 1, nothing, nothing, 4) + @test iszero(allocs_hseed!(hduals, x, hindices, 1, nothing, nothing, 4)) + + # a zero is free for an isbits value type, so only a `BigFloat` catches one being built for a + # layer that was given seeds + bx = BigFloat.(x) + bcfg = ForwardDiff.HessianConfig(nothing, bx, ForwardDiff.Chunk{3}()) + bindices = ForwardDiff.structural_linearindices(bcfg.duals, bx) + allocs_hseed!(bcfg.duals, bx, bindices, 1, bcfg.iseeds, bcfg.oseeds) + @test iszero(allocs_hseed!(bcfg.duals, bx, bindices, 1, bcfg.iseeds, bcfg.oseeds)) + allocs_convert_test_574() = @allocated convert_test_574() allocs_convert_test_574() @test iszero(allocs_convert_test_574()) @@ -50,6 +71,21 @@ end @test iszero(allocs_jacobian!()) end +@testset "Test hessian! allocations" begin + # the sweep is allocation-free only as long as the positions of a dense input stay a range: + # a `structural_linearindices` or `structural_chunk` returning an array would show up here + function allocs_hessian!(n, c) + f(z) = sum(abs2, z) + sum(z)^3 + x = randn(n) + result = zeros(n, n) + cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}()) + ForwardDiff.hessian!(result, f, x, cfg) # warmup + return @allocated ForwardDiff.hessian!(result, f, x, cfg) + end + @test iszero(allocs_hessian!(40, 6)) # seven blocks, the last one partial + @test iszero(allocs_hessian!(10, 10)) # a single block +end + @testset "allocation-free nested StaticArray jacobian" begin # test that nested jacobians of StaticArrays do not allocate. # This is a regression test for issue #798, where the inner jacobian was allocating diff --git a/test/GradientTest.jl b/test/GradientTest.jl index bf121239..50f96656 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) @@ -275,6 +280,14 @@ end end end +# a structured work buffer cannot preserve an unassigned entry: it can only be left unassigned in +# an `Array` +@testset "$(nameof(T)) with an unassigned entry" for T in (LowerTriangular, UpperTriangular) + x = T(Matrix{BigFloat}(undef, 3, 3)) + @test_throws "ArgumentError: cannot differentiate at an input with an unassigned entry at index CartesianIndex(1, 1): that would leave an entry of the $(nameof(T)) work buffer unassigned" ForwardDiff.gradient(sum, x) +end +@test_throws "ArgumentError: cannot differentiate at an input with an unassigned entry at index 1: that would leave an entry of the Diagonal work buffer unassigned" ForwardDiff.gradient(sum, Diagonal(Vector{BigFloat}(undef, 3))) + # issue #769 @testset "functions with `Dual` output" begin x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))] diff --git a/test/HessianTest.jl b/test/HessianTest.jl index 8be72ee5..91db825e 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -23,6 +23,8 @@ h = [-66.0 -40.0 0.0; -40.0 130.0 -80.0; 0.0 -80.0 200.0] +hessian_error = "DimensionMismatch: hessian(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?" + @testset "running hardcoded test with chunk size = $c and tag = $(repr(tag))" for c in HESSIAN_CHUNK_SIZES, tag in (nothing, Tag((f,ForwardDiff.gradient), eltype(x))) cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}(), tag) resultcfg = ForwardDiff.HessianConfig(f, DiffResults.HessianResult(x), x, ForwardDiff.Chunk{c}(), tag) @@ -51,11 +53,22 @@ 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: chunk size cannot be greater than the number of differentiated entries of x (4 > 3)" ForwardDiff.hessian(f, x, ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{length(x) + 1}())) +@test_throws hessian_error ForwardDiff.hessian(identity, x) +@test_throws hessian_error ForwardDiff.hessian!(similar(x, 3, 3), identity, x) +@test_throws "DimensionMismatch: the config was built for an array of size (3,) and cannot be used with an array of size (4,)" ForwardDiff.hessian(f, rand(4), ForwardDiff.HessianConfig(f, x)) ######################## @@ -108,10 +121,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 hessian_error 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 +181,247 @@ for T in (StaticArrays.SArray, StaticArrays.MArray) @test DiffResults.hessian(sresult3) == DiffResults.hessian(result) end +@testset "empty input" begin + f = z -> 1.0 + for x in (Float64[], SVector{0,Float64}()) + H = ForwardDiff.hessian(f, x) + @test size(H) == (0, 0) + @test ForwardDiff.hessian!(fill(NaN, 0, 0), f, x) == H + end +end + +@testset "a StaticArray result that is not a matrix" begin + sx = SVector(1.0, 2.0, 3.0) + flat = fill(NaN, 9) + @test ForwardDiff.hessian!(flat, prod, sx) === flat + @test reshape(flat, 3, 3) == ForwardDiff.hessian(prod, sx) +end + +@testset "a result of the wrong shape" begin + @testset "$(nameof(typeof(x)))" for x in (randn(3), SVector(1.0, 2.0, 3.0)) + @test_throws "DimensionMismatch: cannot store the 3×3 Hessian in a result of size (4, 4)" ForwardDiff.hessian!(fill(NaN, 4, 4), sum, x) + @test_throws "DimensionMismatch: cannot store the 3×3 Hessian in a result of size (2, 2)" ForwardDiff.hessian!(fill(NaN, 2, 2), sum, x) + @test_throws "DimensionMismatch: cannot store the 3×3 Hessian in a result of length 8" ForwardDiff.hessian!(fill(NaN, 8), sum, x) + end + result = DiffResults.DiffResult(0.0, randn(3), fill(NaN, 4, 4)) + @test_throws "DimensionMismatch: cannot store the 3×3 Hessian in a result of size (4, 4)" ForwardDiff.hessian!(result, sum, randn(3)) + + # the gradient buffer is written by linear index of `x`, so it is checked against `x` as well + @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.DiffResult(0.0, fill(NaN, 4), fill(NaN, 3, 3)), sum, randn(3)) + @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.DiffResult(0.0, fill(NaN, 2), fill(NaN, 3, 3)), sum, randn(3)) + @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.DiffResult(0.0, fill(NaN, 4, 3), fill(NaN, 9, 9)), sum, UpperTriangular(randn(3, 3))) +end + +@testset "a gradient buffer that is flat rather than shaped like x" begin + X = randn(3, 3) + f = z -> sum(abs2, z) + result = ForwardDiff.hessian!(DiffResults.DiffResult(0.0, fill(NaN, 9), fill(NaN, 9, 9)), f, X) + @test DiffResults.gradient(result) == vec(2 .* X) + @test DiffResults.hessian(result) == ForwardDiff.hessian(f, X) +end + +@testset "an array-valued f is not a Hessian" begin + sx = SVector(1.0, 2.0, 3.0) + @test_throws hessian_error ForwardDiff.hessian(identity, sx) + @test_throws hessian_error ForwardDiff.hessian!(fill(NaN, 3, 3), identity, sx) + @test_throws hessian_error ForwardDiff.hessian!(DiffResults.HessianResult(sx), identity, sx) + @test_throws hessian_error ForwardDiff.hessian(identity, [1.0, 2.0, 3.0]) +end + +@testset "the result does not depend on the chunk size" begin + n = 16 + v = randn(n) + # the mixed partials of `log(sum(exp, ·))` round differently in the two orders; `sum(z)^3`, + # `exp(sum(z))`, `prod(z)` and `sum(sin, z) * sum(cos, z)` do not, and would pass regardless + f = z -> log(sum(exp, z)) + expected = Matrix(ForwardDiff.hessian(f, SVector{n}(v))) + + @testset "chunk size = $c" for c in (1, 2, 3, 5, 7, 11, n) + H = ForwardDiff.hessian(f, v, ForwardDiff.HessianConfig(f, v, ForwardDiff.Chunk{c}())) + @test H == expected + @test H == transpose(H) + end +end + +# https://github.com/JuliaDiff/ForwardDiff.jl/issues/845 +@testset "f that inspects the layers of its argument" begin + # `ForwardDiff.value` drops the outer perturbation of an intermediate, not of the result + f = z -> sum(abs2, z) + ForwardDiff.value(z[1]) * z[2] + # the mixed derivative of `value(z[1]) * z[2]` is 1 in one order and 0 in the other, and + # only one triangle of block pairs is evaluated + expected_hessian = [2.0 0.0 0.0; 0.0 2.0 0.0; 0.0 0.0 2.0] + + @testset "$(nameof(typeof(x))), chunk size = $c" for x in ( + [1.0, 2.0, 3.0], SVector(1.0, 2.0, 3.0), MVector(1.0, 2.0, 3.0), + ), c in HESSIAN_CHUNK_SIZES + cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}()) + + @test ForwardDiff.hessian(f, x) == expected_hessian + @test ForwardDiff.hessian(f, x, cfg) == expected_hessian + + out = fill(NaN, 3, 3) + @test ForwardDiff.hessian!(out, f, x, cfg) === out + @test out == expected_hessian + + result = ForwardDiff.hessian!(DiffResults.HessianResult(x), f, x, cfg) + @test DiffResults.value(result) == 16.0 + @test DiffResults.gradient(result) == [4.0, 5.0, 6.0] + @test DiffResults.hessian(result) == expected_hessian + end +end + +# https://github.com/JuliaDiff/ForwardDiff.jl/issues/845 +# https://github.com/JuliaDiff/ForwardDiff.jl/issues/846 +@testset "a result that does not carry both perturbations" begin + xs = ([1.0, 2.0, 3.0], SVector(1.0, 2.0, 3.0), MVector(1.0, 2.0, 3.0)) + + @testset "the inner perturbation only: $(nameof(typeof(x)))" for x in xs + # the result itself lost its outer layer, so no second derivative survives -- but + # the gradient does, in the inner one + f = z -> ForwardDiff.value(sum(abs2, z)) + @test all(iszero, ForwardDiff.hessian(f, x)) + @test all(iszero, ForwardDiff.hessian!(fill(NaN, 3, 3), f, x)) + result = ForwardDiff.hessian!(DiffResults.HessianResult(x), f, x) + @test DiffResults.value(result) == 14.0 + @test DiffResults.gradient(result) == [2.0, 4.0, 6.0] + @test all(iszero, DiffResults.hessian(result)) + end + + @testset "neither perturbation: $(nameof(typeof(x)))" for x in xs + f = Returns(2.0) + @test all(iszero, ForwardDiff.hessian(f, x)) + @test all(iszero, ForwardDiff.hessian!(fill(NaN, 3, 3), f, x)) + result = ForwardDiff.hessian!(DiffResults.HessianResult(x), f, x) + @test DiffResults.value(result) == 2.0 + @test all(iszero, DiffResults.gradient(result)) + @test all(iszero, DiffResults.hessian(result)) + end + + @testset "an enclosing tag only: $(nameof(typeof(x)))" for x in xs + # `f` does not depend on `z`, so its result carries the `derivative` tag alone + @test ForwardDiff.derivative(a -> ForwardDiff.hessian(z -> a * 2.0, x)[1, 1], 1.0) == 0.0 + + ForwardDiff.derivative(1.0) do a + # a buffer that can hold the enclosing tag is written in full, one that cannot errors + H = fill(a * 111.0, 3, 3) + ForwardDiff.hessian!(H, z -> a * 2.0, x) + @test all(iszero, H) + @test_throws MethodError ForwardDiff.hessian!(fill(111.0, 3, 3), z -> a * 2.0, x) + return zero(a) + end + end +end + +@testset "no block is evaluated for a derivative the result cannot carry" begin + x = [1.0, 2.0, 3.0] + chunk = ForwardDiff.Chunk{1}() + evaluations = Ref(0) + + function second_order(z) + evaluations[] += 1 + return sum(abs2, z) + end + function first_order(z) + evaluations[] += 1 + return ForwardDiff.value(sum(abs2, z)) + end + function constant(z) + evaluations[] += 1 + return 2.0 + end + + # one evaluation per diagonal block and one per pair of distinct blocks + evaluations[] = 0 + ForwardDiff.hessian(second_order, x, ForwardDiff.HessianConfig(second_order, x, chunk)) + @test evaluations[] == 6 + + # without the outer perturbation only the diagonal blocks contribute, and only a gradient + evaluations[] = 0 + ForwardDiff.hessian(first_order, x, ForwardDiff.HessianConfig(first_order, x, chunk)) + @test evaluations[] == 1 + + evaluations[] = 0 + ForwardDiff.hessian!(DiffResults.HessianResult(x), first_order, x, + ForwardDiff.HessianConfig(first_order, x, chunk)) + @test evaluations[] == 3 + + # a constant still needs the evaluation that determines the output type + evaluations[] = 0 + ForwardDiff.hessian!(DiffResults.HessianResult(x), constant, x, + ForwardDiff.HessianConfig(constant, x, chunk)) + @test evaluations[] == 1 +end + +@testset "nested differentiation" begin + f = z -> sum(w -> w^3, z) + @testset "$(nameof(typeof(x)))" for x in ([1.0, 2.0, 3.0], SVector(1.0, 2.0, 3.0)) + @test ForwardDiff.derivative(a -> ForwardDiff.hessian(f, a .* x)[1, 1], 1.0) == 6.0 + end +end + +@testset "$(nameof(W)), n = $n" for n in (3, 5), (W, sidx) in ( + (LowerTriangular, [i + n * (j - 1) for j in 1:n for i in j:n]), + (UpperTriangular, [i + n * (j - 1) for j in 1:n for i in 1:j]), + (Diagonal, 1:(n + 1):n^2), +) + x = W(randn(n, n)) + L = length(x) + # d²f/dx[a]dx[b] is `w[a] * w[b]`, which no permutation of the structural positions reproduces + w = zeros(L) + w[sidx] = 1:length(sidx) + f = z -> dot(w, z)^2 / 2 + + expected = w * transpose(w) + grad = reshape(w .* dot(w, x), n, n) + + # `length(sidx) - 1` makes the final chunk a partial one + @testset "chunk size = $c" for c in unique((1, 2, length(sidx) - 1, length(sidx))) + cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}()) + + H = ForwardDiff.hessian(f, x, cfg) + @test H isa Matrix + @test size(H) == (L, L) + @test H == expected + + out = fill(NaN, L, L) + @test ForwardDiff.hessian!(out, f, x, cfg) === out + @test out == expected + + flat = fill(NaN, L^2) + @test ForwardDiff.hessian!(flat, f, x, cfg) === flat + @test reshape(flat, L, L) == expected + + # `DiffResults.HessianResult` allocates a dense gradient buffer even for a structured `x` + result = DiffResults.HessianResult(x) + result = ForwardDiff.hessian!(result, f, x, cfg) + @test DiffResults.value(result) ≈ f(x) + @test DiffResults.gradient(result) == grad + @test DiffResults.hessian(result) == 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..6a6454ca 100644 --- a/test/SeedTest.jl +++ b/test/SeedTest.jl @@ -16,13 +16,19 @@ include("utils.jl") # The expected structural index sets are written out by hand rather than obtained from # `structural_eachindex`, so a bug in that iterator cannot hide inside the assertions depending on # it; one test ties the two together. Order is significant: `index` and `count` are positions along -# the sequence, not array indices. The sets are heterogeneous by design — `Vector` and `Diagonal` -# enumerate linear indices (the latter via `diagind`), `UpperTriangular` enumerates `CartesianIndex` -# in column-major order. +# the sequence, not array indices. The `sidx` sets are heterogeneous by design — `Vector` and +# `Diagonal` enumerate linear indices (the latter via `diagind`), `UpperTriangular` enumerates +# `CartesianIndex` in column-major order — while `lidx` is the same positions as linear indices. const SEED_CASES = ( - (rand(10), collect(1:10)), - (UpperTriangular(rand(5, 5)), [CartesianIndex(i, j) for j in 1:5 for i in 1:j]), - (Diagonal(rand(6, 6)), collect(1:7:36)), + (rand(10), + 1:10, + 1:10), + (UpperTriangular(rand(5, 5)), + [CartesianIndex(i, j) for j in 1:5 for i in 1:j], + [i + 5 * (j - 1) for j in 1:5 for i in 1:j]), + (Diagonal(rand(6, 6)), + 1:7:36, + 1:7:36), ) # Positions within `sidx` whose partials are zero. @@ -42,7 +48,7 @@ function fill_marker!(duals, x, sidx, marker) return duals end -@testset "seed_zero_partials!: $(nameof(typeof(x)))" for (x, sidx) in SEED_CASES +@testset "seed_zero_partials!: $(nameof(typeof(x)))" for (x, sidx, _) in SEED_CASES cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{3}()) duals, seeds = cfg.duals, cfg.seeds N = ForwardDiff.npartials(eltype(duals)) @@ -90,4 +96,23 @@ end end end +@testset "seed_hessian_chunk!: $(nameof(typeof(x)))" for (x, sidx, lidx) in SEED_CASES + (; duals, iseeds, oseeds) = ForwardDiff.HessianConfig(nothing, x, ForwardDiff.Chunk{3}()) + nstruct = length(sidx) + indices = ForwardDiff.structural_linearindices(duals, x) + + # the windows below are positions along `indices`, so pin it to the implementation once + @test indices == lidx + + ForwardDiff.seed_hessian_chunk!(duals, x, indices, 1, nothing, nothing, nstruct) + ForwardDiff.seed_hessian_chunk!(duals, x, indices, 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, indices, 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