Apply the second-round symmetric-Hessian review: output shape, shape check, chunk-size independence

[devmotion]

Stacked on #837, applying my second round of review as code rather than as prose, since you suggested that would save us both a round-trip. Merge it into kc/symmetric_hessian and it becomes part of #837. Draft until #840 lands, since it depends on it.

Read commits 1–4; commit 0 is a merge of #840. Because the base is kc/symmetric_hessian, the Files view shows #840's diff as well as mine — the per-commit view does not. Your two commits are untouched and never rewritten.

Why #840 has to come first

#840 (fixing #838/#839, the two issues I filed off the last review) settles that gradient, jacobian and hessian results for a structured x are indexed by the linear indices of x, with hard zeros at the structurally zero entries — so hessian(f, ::UpperTriangular{3×3}) is 9×9, not 6×6 — and pins it with a HessianTest.jl testset. On master hessian inherits that through jacobian(∇f, x); the new kernel numbers rows and columns by seeding position instead, an order only ForwardDiff knows. That testset fails with 16 failures on the merge commit and passes from commit 1 on.

The merge is textually clean apart from one testset collision, but not semantically: #840 gave extract_gradient_chunk! an x argument (last, as of its 27033bb) and stopped it zeroing what it does not write. Both are fixed inside the merge commit so that every commit builds.

One thing the merge deletes from my review: I said extract_hessian_gradient_chunk! must not delegate to extract_gradient_chunk!, because that read positions off result. #840 fixed that, so delegating is now the right thing to do and the gradient of hessian!(HessianResult(x), f, x) lands correctly for free.

The commits

1. Index the Hessian by the linear indices of x. Both axes of a Hessian belong to entries of x the way a Jacobian's columns do, so Extract derivatives at the indices of x, and store those indices in the config #840's structural_columns supplies the positions for both and no new mapping is needed.
2. Restore the result-shape check. Writing entry by entry lost the validation the old broadcast did on the way past: on 0c04e90, hessian!(fill(NaN, 4, 4), f, rand(3)) returns with a row and column of NaN where master threw. reshape_hessian mirrors Extract derivatives at the indices of x, and store those indices in the config #840's reshape_jacobian rather than introducing a different helper.
3. Read the same triangle at every chunk size. Diagonal blocks read the nested partials at outer <= inner, off-diagonal ones at outer > inner — equal in exact arithmetic, not in floating point. Swapping which layer the column block seeds costs nothing and makes them all bitwise equal, which also makes the existing symmetric_static == hessian(symmetry_f, x) assertion meaningful rather than passing only because n = 9 stays under DEFAULT_CHUNK_THRESHOLD. The swap also moves the cheaper of the two seedings into the hot loop, since the inner zero is a Partials{N,V} and the outer one a Partials{N,Dual{T,V,N}}: for BigFloat at N = 3 the call that runs O(nblocks²) times drops from 192 to 96 bytes, the one that runs O(nblocks) times rises from 96 to 192. Same commit: seed_hessian_chunk! built both layers' zero partials whether or not it needed them (free for isbits, 288 → 0 bytes per call for BigFloat at N = 3), the Partials{0} and seed_zero_partials!-MethodError comments, and the tests.
4. Walk the structural positions instead of indexing them. The sweep needs no random access: for a fixed column block the row blocks are walked in order, and within a block each row resumes the column walk from the state saved at the block's start rather than restarting it, so the walk costs what seed_hessian_chunk! already pays to seed the same block and no more. A diagonal block reads its columns from wherever its rows have got to, which is exactly the upper triangle it wants. The lazy positions of the triangular wrappers therefore never have to be materialized, and hessian! is allocation-free for every input kind, as gradient! and jacobian! are. Results are bitwise unchanged and JET stays clean. Same commit: the comment fixes below and the remaining test gaps.

Everything else from the review is a test: #837's structured testset is replaced by #840's coupled one (the old dot(weights, abs2.(x)) has a strictly diagonal Hessian, so it could not detect a wrong offset), plus the extension's reshape branch and both HESSIAN_ERROR throws, empty inputs, the mixed seed forms in SeedTest asserting which layer got which seed, and the AllocationsTest rename.

Three comments that were not quite right

• The off-diagonal comment said reading one triangle means "the result does not depend on the chunk size", unconditionally. What commit 3 fixes is the nesting order every entry of H is read at; the result follows only for an f that does not branch on the seeds. < on Dual is lexicographic — values first, partials on a tie — so at the outer layer a tie is broken by comparing inner Duals and at the inner layer by comparing plain numbers. f(z) = max(z[1], z[2]) * z[3] at z = [1, 1, 2] is still chunk-size dependent for that reason. The comment now names the mechanism.
• zero_unseeded_columns! is called from the sweep with a scalar dual where its only other callers pass an array of them; it works because it reads ydual for its value type and eltype is the identity on numbers. Stated at the helper and at the call site.
• length(x) was emitted into both generated bodies and read by the one that allocates a result.

Verified

• Full suite 10111/10111, zero failures.
• 130 structured configurations (Lower/UpperTriangular/Diagonal × n ∈ {3,5,8} × every chunk size) against a reference that differentiates the vector of structural entries and scatters it: exact symmetry, exactly-zero non-structural rows and columns.
• n = 16 bitwise identical across chunk sizes 1, 2, 3, 5, 7, 11, 16 and to the SVector path; the three wrappers bitwise identical across every chunk size, on an f whose Hessian is not exactly representable. Both are now testsets rather than side checks.
• Against Extract derivatives at the indices of x, and store those indices in the config #840's branch for the three wrappers: same shapes, agreeing to better than 1e-12 relative — the two compute the Hessian by different algorithms, so last-bit differences are expected.
• hessian! allocates nothing, for dense inputs and all three wrappers, at every chunk size.
• The four-combination seed_hessian_chunk! allocation test is on a Float64x, whose zero partials are isbits and free to build whether or not they are needed, so it cannot see commit 3's laziness. It is now also pinned on a BigFloatx; reverting the laziness makes that one assertion fail at 288 bytes and leaves the Float64 loop passing, which is how I checked it bites.
• Writing both triangles inside the extraction loop rather than mirroring one at the end: 1.3–2.2× faster in 7 of 8 configurations over n ∈ {64, 256, 1024, 2048} × N ∈ {8, 64}, since the value is still in a register for the second store while a second pass re-reads n²/2 entries. Extraction is 0.11–1.6% of hessian! either way, for f = sum(abs2, ·).

Left for you

• Version bump. Exact symmetry, the new shape for structured inputs and the restored DimensionMismatch are all user-visible, so 1.5.0 rather than a patch — but that is your release, and Extract derivatives at the indices of x, and store those indices in the config #840 is breaking for structured inputs on its own.
• exploit symmetry in the hessian #837's description still doesn't mention the structured-input shape change.
• Retargeting exploit symmetry in the hessian #837's base onto devmotion/structured-extraction-838-839, or merging Extract derivatives at the indices of x, and store those indices in the config #840 into kc/symmetric_hessian directly, would make exploit symmetry in the hessian #837's own diff readable again. Both are your call; I didn't want to touch your branch.
• jacobian's docstring deserves the same note about Diagonal costing length(x)² that I added to hessian's — Extract derivatives at the indices of x, and store those indices in the config #840's to fix, not this PR's.
• require_one_based_indexing is not applied to DiffResult Hessian buffers. Pre-existing and shared with jacobian!; not widened into here.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 99.19355% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.91%. Comparing base (0c04e90) to head (6279027).

Files with missing lines	Patch %	Lines
ext/ForwardDiffStaticArraysExt.jl	83.33%	1 Missing ⚠️

Additional details and impacted files

@@ Coverage Diff @@## kc/symmetric_hessian #841 +/- ##
========================================================
+ Coverage 91.30% 91.91% +0.61% 
========================================================
Files 11 11 Lines 1115 1162 +47 ========================================================
+ Hits 1018 1068 +50 + Misses 97 94 -3 

☔ 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]

Superseded by #843, which carries #837's commits on top of #840 instead of merging #840 into a branch based on #837, so the diff shows the Hessian work alone.