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
 
 @generatedfunctiondualize(::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)
+@inlinefunctionextract_hessian(::Type{T}, ydual::Partials, x::StaticArray) where {T}
+ H =extract_jacobian(T, ydual, x)
+returntypeof(H)(Symmetric(H, :U))
+end
+
+@inlinefunctionextract_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)))
+returnzero(R)
+end
+
+@inlinefunction 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)
+returnextract_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)
+@inlinefunction 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`.
+@inlinefunction_seed!(make_dual::F, duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {F,T,V,N}
 ifisbitstype(V)
-for idx in idxs
- duals[idx] =Dual{T,V,N}(x[idx], seed)
+for(i, idx)inenumerate(idxs)
+ duals[idx] =make_dual(x[idx], i)
 end
 else
-for idx in idxs
+for(i, idx)inenumerate(idxs)
 ifisassigned(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
 
 functionseed!(duals::AbstractArray{Dual{T,V,N}}, x,
  seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
-ifisbitstype(V)
-for (i, idx) inzip(1:N, structural_eachindex(duals, x))
- duals[idx] =Dual{T,V,N}(x[idx], seeds[i])
-end
-else
-for (i, idx) inzip(1:N, structural_eachindex(duals, x))
-ifisassigned(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
 
 functionseed!(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)
-ifisbitstype(V)
-for (i, idx) inzip(1:chunksize, idxs)
- duals[idx] =Dual{T,V,N}(x[idx], seeds[i])
-end
-else
-for (i, idx) inzip(1:chunksize, idxs)
-ifisassigned(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.
+functionseed_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,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,17 +35,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 +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.
-functionseed_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)
-ifisbitstype(V)
-for (i, idx) inzip(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) inzip(1:chunksize, idxs)
-ifisassigned(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.
 functionextract_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 in2: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 in1:(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
 
 @evalfunctionsymmetric_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
 
 @evalfunctionsymmetric_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)
 @testiszero(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)
+@testiszero(allocs_hseed!(hduals, x, 1, iseeds, oseeds))
+allocs_hseed!(hduals, x, 1, nothing, nothing, 4)
+@testiszero(allocs_hseed!(hduals, x, 1, nothing, nothing, 4))
+
 allocs_convert_test_574() =@allocatedconvert_test_574()
 allocs_convert_test_574()
 @testiszero(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)