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 (ceac9e1).

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 ceac9e1.

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, structured-input dimensions, 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 ceac9e1 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.