exploit symmetry in the hessian

[KristofferC]

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

┌─────────────────────┬───────────────────┬───────────────────────┐
│      function       │         n         │        speedup        │
├─────────────────────┼───────────────────┼───────────────────────┤
│ rosenbrock / ackley │ 10 (single chunk) │ ~1.0x (no regression) │
├─────────────────────┼───────────────────┼───────────────────────┤
│ rosenbrock / ackley │ 30                │ 1.45x / 1.57x         │
├─────────────────────┼───────────────────┼───────────────────────┤
│ rosenbrock / ackley │ 100               │ 1.85x / 1.94x         │
└─────────────────────┴───────────────────┴───────────────────────┘

Fixes #836 cc @gdalle

Basically one-shotted by Claude and then decringified with gpt.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.30%. Comparing base (569af35) to head (0c04e90).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #837      +/-   ##
==========================================
+ Coverage   90.68%   91.30%   +0.61%     
==========================================
  Files          11       11              
  Lines        1052     1115      +63     
==========================================
+ Hits          954     1018      +64     
+ Misses         98       97       -1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[devmotion]

Thanks, this is a nice improvement. I went through it with the help of Claude.

The core of it holds up. I reproduced the speedup locally with rosenbrock and the default chunk size: 49.3 -> 31.8 µs at n = 30 and 1985 -> 1084 µs at n = 100, with allocations going from 127472 to 16256 bytes. I also diffed the results against master for 3 functions x n ∈ {1,2,3,4,5,7,8,13} x chunk sizes {1,2,3,n} x all four hessian/hessian! methods (values, gradients and Hessians) and found no differences, and HessianTest.jl, MiscTest.jl, AllocationsTest.jl and ConfusionTest.jl all pass with the new implementation. JET.report_opt is clean for hessian! with both a Matrix and a flat Vector result. The zero-invariant is maintained (every seeded chunk is cleared before the next evaluation), every diagonal block is visited so the gradient is written in full, and I'm glad to see InnerGradientForHess and the JuliaLang/julia#15276 workaround from #316 go.

The failing Julia pre jobs are unrelated: the only failure is the @test_opt ForwardDiff.jacobian(identity, ...) in test/QATest.jl, which fails on #835 as well. Xref #828.

A few comments.

The single-chunk case became slower

Since there is no vector-mode path anymore, HessianConfig(f, x, Chunk{n}()) now runs the block kernel with nblocks == 1, and it makes three passes over xdual where master made one: the initial full zeroing, the re-write of block 1 on top of it, and the trailing clear of block 1 for a loop that never runs.

seed_hessian_chunk!(xdual, x, 1, nothing, nothing, xlen)  # zeroes everything
seed_hessian_chunk!(xdual, x, 1, iseeds, oseeds)          # ... then overwrites block 1

chunk_mode_gradient and jacobian_chunk_mode_expr partition the buffer instead, so that every element is written exactly once (that was the point of a337ee6). Doing the same here:

seed_hessian_chunk!(xdual, x, 1, iseeds, oseeds)
seed_hessian_chunk!(xdual, x, N + 1, nothing, nothing, xlen - N)
...
nblocks > 1 && seed_hessian_chunk!(xdual, x, 1, nothing, nothing)

turns a ~6% regression into a ~12-16% improvement over master, and gives identical results for every chunk size I tried:

forced single chunk	master	PR	with the partition
sum, n=40	79.3 µs	84.1 µs	67.9 µs
sum(abs2, ·), n=40	102 µs	108 µs	91.4 µs

It matters more than it looks, because the buffer is n x (N+1)^2 words and users do pass large explicit chunk sizes.

While you are in there: q's outer seeds are re-written and cleared on every p iteration although they never change within a q. Seeding q once before the p loop and letting the diagonal block re-seed it with both layers is worth another ~5% when f is cheap, and reads better.

The StaticArrays path is left behind

ext/ForwardDiffStaticArraysExt.jl still routes hessian(f, x::StaticArray) and hessian!(result::AbstractArray, f, x::StaticArray) through jacobian(Base.Fix1(gradient, f), x), and that path is not exactly symmetric. Maximum abs(H - H') on master:

f	n=3	n=5	n=8
x -> sum(sin(x[i])/(1+x[i+1]^2) ...)	5.6e-17	1.1e-16	2.2e-16
x -> reduce(*, cumsum(x)) + atan(x[1], x[end])	8.9e-16	2.8e-14	7.3e-12

So after this PR hessian(f, ::Vector) would be exactly symmetric while hessian(f, ::SVector) would not, which is awkward to document. IMO we should fix the extension in the same PR. The shape is already there in hessian!(result::ImmutableDiffResult, f, x::StaticArray), which does the nested dualize with a single tag, matching what HessianConfig does. Adding a generated extract_hessian that emits only the N(N+1)/2 upper-triangle entries and reuses each one for the mirrored position, then

function ForwardDiff.hessian(f::F, x::StaticArray) where {F}
    T = typeof(Tag(f, eltype(x)))
    return extract_hessian(T, partials(T, f(dualize(T, dualize(T, x)))), x)
end

plus extract_hessian_chunk! for the mutating method and a one-word swap in the ImmutableDiffResult method, is enough. I tried it and all 590 assertions in HessianTest.jl still pass, including isa StaticArray, the comparisons against the Array path, and iszero(hessian_allocs()) from #720. For SVector inputs the timings are unchanged, but hessian(prod, ::SMatrix{3,3}) goes from 249 ns to 123 ns because the intermediate gradient array disappears. As a bonus the extension then uses the same tag for both layers as HessianConfig does, instead of building a separate Tag(Fix1(gradient, f), V).

hessian!(result::MutableDiffResult, f, x::StaticArray) needs no change, it forwards to the generic method.

Structured inputs silently change shape

For UpperTriangular, LowerTriangular and Diagonal inputs the result shape changes, because it is now structural_length(x)^2 instead of length(x) x structural_length(x):

julia> size(ForwardDiff.hessian(z -> sum(abs2, z), UpperTriangular(rand(3, 3))))
(9, 6)   # master
(6, 6)   # this PR

The new shape is the right one (master mixes linear and structural indices, which I would call a bug), and I checked that the values are correct. But we have no Hessian test for structured inputs at all, so this is completely untested, and it is breaking for hessian!(result::AbstractArray, ...). Could you add a testset along the lines of the LowerTriangular, UpperTriangular and Diagonal one in test/GradientTest.jl and mention the change in the PR description?

Docs

• The exact symmetry of the result is the whole point of the PR but is not documented. It should be stated in both the hessian and hessian! docstrings.
• The hessian!(result::AbstractArray, ...) docstring still says H(f) is J(∇(f)). That was removed from hessian but not here.
• Since the kernel only uses cfg.jacobian_config.seeds now, the two HessianConfig constructors have become interchangeable and hessian!(::DiffResult, f, x, HessianConfig(f, x)) works. Worth documenting. As a follow-up, jacobian_config.duals is now dead weight (a pair of buffers for the DiffResult constructor), only needed transiently so that GradientConfig has something to similar.

Smaller things

• result isa AbstractMatrix && size(result) == (xlen, xlen) ? result : reshape(result, xlen, xlen): Base._reshape already returns the parent unchanged when the dims match, so the size check only turns a compile-time branch into a runtime one. extract_jacobian! just uses result isa AbstractMatrix ? result : reshape(...).
• seed_hessian_chunk! is the third copy of the isbitstype / isassigned / Base._unsetindex! loop, after _seed_zero_partials! and seed!. Could we factor the shared shape into one helper in apiutils.jl that takes a function producing the dual? That would also move seed_hessian_chunk! next to its siblings, where SeedTest.jl can reach it.
• nblocks = xlen == 0 ? 1 : div(xlen + N - 1, N) is cld(xlen, N) plus a guard for N == 0, which only happens for empty x. A short comment would help.
• typeof(value(T, value(T, ydual1))) is spelled valtype(T, valtype(T, typeof(ydual1))) in chunk_mode_gradient.
• The comment "Keep all unseeded blocks at zero between evaluations" is on the line that initialises the whole buffer. The initialisation is load-bearing (a fresh similar may contain #undef), so I would separate the two statements.
• The error for non-Real output changed from the gradient message to the new HESSIAN_ERROR, which is better, but it is only checked on the first evaluation and there is no test for it.

Coverage

The -0.24% is 9 lines. From the Codecov line report:

• src/hessian.jl 80-87: the whole non-isbits branch of seed_hessian_chunk!, including the Base._unsetindex! path. A BigFloat testset modelled on the one in test/JacobianTest.jl covers it: a Vector{BigFloat}(undef, 10) with one unassigned entry that f never reads, at chunk sizes (1, 2, 10). Unlike the Jacobian case the position of the hole does not matter here, since the new kernel clears every chunk it seeds including the last. If we factor out the shared loop as suggested above, extending SeedTest.jl instead would be nicer.
• src/hessian.jl:118: the "chunk size cannot be greater than ..." ArgumentError. One @test_throws ArgumentError with Chunk{length(x) + 1}(). The same line in src/gradient.jl is uncovered on master too, so that could be added at the same time.
• ext/ForwardDiffStaticArraysExt.jl 42 and 70: the 5-argument gradient!(..., cfg, ::Val) and jacobian!(..., cfg, ::Val) methods for StaticArray. They were only ever reached because the old Hessian called gradient!/jacobian! with Val{false}(), and the new implementation calls neither, so they need direct tests in the StaticArray sections of GradientTest.jl and JacobianTest.jl.

Unrelated to the coverage numbers, but while adding tests: hessian(f, x::StaticArray, cfg) and friends discard cfg and the Val argument entirely, so they do not throw InvalidTagException the way the Array methods do. That is pre-existing and consistent across gradient/jacobian/hessian in the extension, so not for this PR, but it is a real perturbation confusion hole and probably deserves its own issue.

[KristofferC]

Addressed the review feedback in 0c04e90.

Highlights:

• Partitioned the initial Hessian buffer seeding, avoiding the redundant passes in the single-chunk case.
• Seed each row block outer perturbation once across its off-diagonal evaluations.
• Factored the assignment-safe dual-writing loop into a shared helper and moved Hessian seeding next to the other seeding utilities.
• Made the StaticArrays path exactly symmetric while preserving SMatrix/MMatrix return types. The implementation reuses the existing static Jacobian extractor and the core Hessian block extractor rather than adding another generated-code subsystem.
• Documented exact symmetry and interchangeable HessianConfig constructors.
• Added coverage for structured inputs, unassigned BigFloat entries, oversized chunks, non-Real outputs, config interchangeability, exact StaticArray symmetry, and the previously uncovered five-argument StaticArray methods.

Benchmarks compare 0091edd with 0c04e90 on Julia 1.12.7 using warm BenchmarkTools medians (samples=10_000, evals=1):

Benchmark	Before	After	Improvement
Single chunk: sum, n=40	70.291 µs	56.958 µs	19.0%
Single chunk: sum(abs2), n=40	91.333 µs	79.750 µs	12.7%
Default chunks: sum(abs2), n=100	381.750 µs	362.333 µs	5.1%
Default chunks: Rosenbrock, n=30	24.916 µs	23.625 µs	5.2%
Default chunks: Rosenbrock, n=100	811.333 µs	801.459 µs	1.2%
hessian(prod, ::SMatrix{3,3})	167 ns	42 ns	74.9%

Memory and allocation counts were unchanged. The complete test suite passes: 9,396/9,396 tests.

[devmotion]

Thanks, this addresses everything from the last round. Suite is green here (9396/9396), and I re-checked the kernel against references for structured inputs at chunk sizes 1/2/xlen, empty Vector/SVector, constant f, BigFloat with an unassigned entry, SubArray, Dual-valued input, and both nesting directions. All correct and exactly symmetric. Again with Claude's help; the suggestions below are prototyped and run against the suite, not just read.

I want to revisit the (6, 6) shape

I said the new shape was the right one. I no longer think so. The structural ordering is a ForwardDiff-internal detail — for UpperTriangular it is column-major over the upper triangle — so a caller cannot say what H[4, 5] refers to without reading structural_eachindex. gradient does not have this problem: its result is shaped like x, so g[i, j] belongs to x[i, j]. (jacobian uses structural columns, but that is weak precedent — chunked jacobian on a structured input throws, see the bottom.)

I would rather have length(x) x length(x), indexed by the linear indices of x, with hard zeros in the rows and columns of the structural zeros:

size(hessian(f, UpperTriangular(rand(3, 3))))  # (9, 9), rows/cols 2, 3, 6 all zero
size(hessian(f, Diagonal(rand(3))))            # (9, 9)

It is also the smaller change relative to master, which already used linear indices for the rows and only got the columns wrong. And it makes DiffResults.HessianResult(x) the correct buffer — that allocates length(x)^2 — so hessian!(::DiffResult, f, x) starts working for structured inputs instead of erroring.

The cost is Diagonal: hessian(f, ::Diagonal{n}) becomes n^2 x n^2, so Diagonal(100) is 800 MB. I think that is fine if the docstring says so and points at differentiating the diagonal vector instead.

Implementation is a structural_linearindices(x) map next to structural_eachindex (Base.OneTo(length(x)) for plain arrays, diagind for Diagonal, a small Vector{Int} for triangular), threaded into extract_hessian_chunk!, plus hlen == xlen || fill!(H, zero(eltype(H))). One non-obvious extra: extract_hessian_gradient_chunk! has to take its positions from that map rather than delegating to extract_gradient_chunk!, because the latter reads structural_eachindex(grad) and HessianResult(x) hands back a dense size(x) buffer — otherwise the newly reachable DiffResult path scatters the 6 derivatives into linear positions 1:6. With that, 48/48 checks pass across the three wrappers x n in (3, 5) x chunk in (1, 2, xlen), and hessian! on plain vectors is unchanged in both time and allocations.

hessian! no longer checks the result shape

The element-wise write loop dropped the validation that reshape_jacobian and the broadcast in extract_jacobian! used to provide. Vectors are still checked by reshape, matrices are not:

call	master	0c04e90
hessian!(fill(NaN, 4, 4), f, rand(3))	DimensionMismatch	no error, row/col 4 left NaN
hessian!(HessianResult(UpperTriangular(randn(3, 3))), f, x)	DimensionMismatch	no error, rows 7:9 are uninitialized memory
hessian!(fill(NaN, 4, 4), f, SVector(1., 2., 3.))	DimensionMismatch	no error, row/col 4 left NaN

A shared helper used by both hessian! methods and the extension:

function hessian_result_matrix(result, hlen)
    H = result isa AbstractMatrix ? result : reshape(result, hlen, hlen)
    axes(H) == (Base.OneTo(hlen), Base.OneTo(hlen)) ||
        throw(DimensionMismatch(lazy"hessian result has size $(size(H)), expected ($hlen, $hlen)"))
    return H
end

This also turns the remaining breakage from the shape change into a clear error rather than a partial fill.

One triangle everywhere, for free

Diagonal blocks read outer <= inner, off-diagonal blocks read outer > inner, and the StaticArrays path reads outer <= inner everywhere. So the result is not reproducible across chunk sizes: at n = 16, chunk 1, 2, 3, 5, 7, 11 each differ from chunk 16 and from the SVector path (maxdiff ~1e-16). Swapping which layer block q carries fixes it:

seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, nothing, qsize)   # was: nothing, oseeds
for p in 1:(q - 1)
    poffset = (p - 1) * N
    seed_hessian_chunk!(xdual, x, poffset + 1, nothing, oseeds)      # was: iseeds, nothing
    ydual = f(xdual)
    extract_hessian_chunk!(T, H, ydual, poffset, qoffset, N, qsize)  # was: qoffset, poffset, qsize, N
    seed_hessian_chunk!(xdual, x, poffset + 1, nothing, nothing)
end

Same number of evaluations and seed writes, and q is still seeded once outside the loop. Afterwards every chunk size gives bitwise-identical results, identical to the StaticArrays path. It also makes symmetric_static == hessian(symmetry_f, x) robust — right now that only passes because n = 9 is below DEFAULT_CHUNK_THRESHOLD, so the array path uses a single block.

The structured testset cannot fail

dot(weights, abs2.(x)) has a strictly diagonal Hessian and expected is diagm(...), so the testset cannot detect wrong row/column offsets, and it cannot detect a structure-preserving H container either — writing 0.0 off-diagonal into a Diagonal or UpperTriangular is legal. That matters because n = 5 with the triangular types is the only structured case that reaches the chunked path (xlen = 15, Chunk(x) gives N = 8), i.e. exactly the new off-diagonal code, and it is only ever compared against zeros. Could you use something coupled, e.g. z -> exp(dot(w, z)) with expected = f * wv * wv', and loop over explicit chunk sizes instead of relying on Chunk(x)?

Smaller things

• Please don't try to fold the four seed_hessian_chunk!(..., nothing, nothing, ...) calls into seed_zero_partials!. It looks equivalent but Dual{T,V,N}(::V, ::Partials{N,V}) takes the value at its exact type, so it MethodErrors on the nested buffer.
• seed_hessian_chunk! materialises both izero and ozero even when the seeds are supplied. Free for isbits V; 288 bytes per call for BigFloat at N = 3, identical for all three seed combinations.
• extract_hessian(::Type{T}, ::Partials{0}, x) deserves a comment. It is load-bearing for a constant f as well as for empty x, and neither is obvious from the signature.
• Coverage: the extension's hessian! reshape branch and its two new HESSIAN_ERROR throws are untested, as are empty inputs. SeedTest covers (iseeds, oseeds) and (nothing, nothing) but not the two mixed forms the off-diagonal loop actually uses, and it only asserts nonzero-ness rather than which layer received which seed.
• The AllocationsTest testset is still called "Test seed!/seed_zero_partials! allocations".
• The PR description still doesn't mention the structured-input shape change.

Unrelated, filed separately

Both are on master and predate this PR, but the second one is why I don't trust jacobian as precedent above:

1. gradient! writes to the wrong entries for structured inputs #838 — gradient! takes its target positions from structural_eachindex(result) rather than from x. gradient!(fill(NaN, 3, 3), f, UpperTriangular(...)) silently writes the 6 derivatives into linear positions 1:6 at every chunk size, and gradient!(GradientResult(x), f, x) throws ArgumentError: cannot set index (2, 1) in the lower triangular part from DiffResults' linear copyto!.
2. Chunked jacobian throws for Diagonal/LowerTriangular/UpperTriangular inputs #839 — chunked jacobian on structured inputs throws DimensionMismatch — reshape_jacobian uses length(xdual) while the result has structural_length(x) columns. Vector mode uses chunksize(cfg) and works.

[KristofferC]

Maybe it is easier if you just point your bot on this branch instead of going through me, would save us both some tokens :).