Extract derivatives at the indices of x, not of the result

[devmotion]

Fixes #838 and #839. Both are the same omission from #739: seeding became structure-aware, extraction did not.

gradient! wrote to the wrong entries (#838)

extract_gradient!/extract_gradient_chunk! took their positions from structural_eachindex(result) while the seeds are laid out along structural_eachindex(x). They now take x and walk structural_eachindex(x, result). Entries that receive no derivative are zeroed -- that is their derivative -- by zero_unseeded!, skipped when the seeded entries of x already account for every entry result stores, which covers both an unstructured x and the similar(x) result that gradient allocates for a structured one.

The (DiffResult, Dual) method splits on mutability: MutableDiffResult is written entry by entry, ImmutableDiffResult keeps the wholesale DiffResults.gradient! copy, since it cannot be written entry by entry and only arises from StaticArray gradient buffers, for which every entry of x is seeded.

julia> x = UpperTriangular([1.0 2.0 3.0; 0.0 4.0 5.0; 0.0 0.0 6.0]);

julia> f(z) = sum(abs2, z) / 2;

julia> out = fill(NaN, 3, 3); ForwardDiff.gradient!(out, f, x); out
3×3 Matrix{Float64}:      # master:  1.0  3.0  NaN
 1.0  2.0  3.0            #          2.0  5.0  NaN
 0.0  4.0  5.0            #          4.0  6.0  NaN
 0.0  0.0  6.0

julia> ForwardDiff.gradient!(DiffResults.GradientResult(x), f, x);   # master: ArgumentError

This also fixes the mis-scattered gradient of hessian!(::DiffResult, f, x), whose buffer DiffResults.HessianResult allocates densely even for a structured x (the case from the #837 review).

Chunked jacobian threw (#839)

The result was allocated with structural_length(x) columns while reshape_jacobian asked for length(xdual). The allocation is the side that was wrong: the Jacobian is indexed by the linear indices of x -- column j holds ∂f(x)[i]/∂x[j], exactly as the docstring says -- with hard zeros in the columns of the structural zeros. So the four allocations now use length(x), and the redundant n argument of extract_jacobian! is dropped (it was always structural_length(x)).

structural_columns(out, x) returns the columns of out that receive a derivative, in seeding order, without materializing anything: axes(out, 2) where every entry of x is seeded, axes(out, 2)[diagind(x)] for a Diagonal, and a lazy generator over structural_eachindex for the triangular wrappers, whose seeded columns are runs of increasing length and hence no range. It checks its arguments the way structural_eachindex does.

julia> g(z) = [sum(z), sum(abs2, z)];

julia> size(ForwardDiff.jacobian(g, x, ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{2}())))
(2, 9)                    # master: DimensionMismatch

Zeroing belongs to the sweep, not to a chunk

The chunk extractors first recognised the first chunk by index == 1 and used it to zero the whole result. That is neither part of extracting one chunk nor something a chunk can decide -- the entries at stake, those of the structural zeros of x, belong to no chunk in particular -- and it made the generated expressions depend on an invariant that only a comment carried. Both sweeps now do it once up front and the chunk functions only write their chunk.

reshape_jacobian allocated a wrapper on every chunked call

Falling out of the allocation test below. reshape_jacobian reshaped the result even when it already was a matrix, and since 1.12 reshape can no longer return its argument, so every chunk-mode jacobian! allocated an Array wrapper -- for dense inputs too, not just structured ones. extract_jacobian! was given that short-circuit in #797; reshape_jacobian now shares it, with an explicit size check that is stricter than the one reshape was performing on the way past -- that only ruled out a wrong total length, so a matrix of the right length but the wrong shape was silently reinterpreted -- and extract_jacobian! calls it rather than repeating the ternary. Also drops the map! that vector_mode_jacobian(f!, ...) ran before extract_jacobian!, which reads only ydual, and that the map! after it repeats.

x last, and each job in one place

extract_gradient!, extract_gradient_chunk!, extract_jacobian! and extract_jacobian_chunk! take x after the derivatives, like every other internal function that pairs an output with an input: reshape_jacobian(result, ydual, x), which extract_jacobian! calls on its first line, extract_value!(T, out, y, ydual), and the extension's extract_jacobian(T, ydual, x::StaticArray). x is annotated AbstractArray in the Jacobian extractors, so a call written against an older signature fails on the method rather than inside structural_length.

extract_jacobian! duplicated the structured branch of extract_jacobian_chunk! at offset == 0 and now delegates to it, and both Jacobian sweeps zero through one helper, as the gradient sweep already did.

The chunk mode gradient sweep rejects an array-valued f up front, as the Jacobian sweep does, in place of the dual::AbstractArray methods of extract_gradient_chunk!: zero_unseeded! runs before them and dies on zero(::Type{<:AbstractArray}) before dispatch gets there, so gradient on an array-valued f reported a MethodError instead of pointing at jacobian. zero_unseeded! also dispatches on DiffResult rather than MutableDiffResult: a StaticArray gradient buffer makes the result immutable even when the buffer itself can be written to entry by entry, as the MVector of DiffResult(0.0, @MVector zeros(n)) can, which is what chunk mode does with it.

Breaking for structured inputs

Both fixes follow the convention argued for in #839/#837 -- index the result by the indices of x, hard zeros off the structure -- which is what gradient(f, x) has always returned. Consequences:

• jacobian gains the zero columns: (2, 6) → (2, 9) for UpperTriangular(3×3).
• hessian inherits both conventions through jacobian(∇f, x) and becomes length(x) × length(x) with hard-zero rows and columns, instead of mixing linear rows with structural columns: (9, 6) → (9, 9). hessian!(DiffResults.HessianResult(x), f, x) starts working (it allocates length(x)^2), as does jacobian!(DiffResults.JacobianResult(y, x), ...).
• For a Diagonal the Jacobian/Hessian now scale with length(x) = n², so differentiating with respect to the diagonal vector is the better choice there.
• gradient! into a container whose size differs from a structured x now throws DimensionMismatch instead of packing the derivatives in structural order (gradient!(zeros(6), f, UpperTriangular(3×3))). The packed order is a ForwardDiff-internal detail. A dense x is affected only by a length mismatch, which eachindex(x, result) rejects where the entries the two had in common used to be filled (gradient!(zeros(5), f, rand(4))); a result of the same length but a different shape, such as a matrix for a vector input, keeps working.
• jacobian! into a matrix whose shape is not length(y) × length(x) now throws DimensionMismatch in chunk mode, where it used to be reshaped whenever the total length happened to match. Vector mode already used such a result as is; the two modes now agree.

Tests

New testsets in GradientTest.jl (#838), JacobianTest.jl (#839) and HessianTest.jl (both, inherited) over the three wrappers × sizes × every relevant chunk size, covering a dense result, a result shaped like x, DiffResults.GradientResult/JacobianResult/HessianResult, a dense DiffResult gradient buffer, a Jacobian result that is not a matrix, and both the f and f! Jacobian forms. Results are prefilled with NaN so that entries left untouched fail rather than pass, and the expected nonzero positions are written out by hand so a bug in the position mapping cannot hide inside the reference. The Hessian test differentiates a function whose second derivative is 1 + (a == b) on the structural entries, so the reference is exact and the hard-zero rows and columns are asserted rather than approximated; against master that testset gets 0 passed, 3 failed, 4 errored. Also asserted: the out-of-place gradient returns the structure of x, and a structured result cannot hold the gradient of a dense x.

AllocationsTest.jl guards the new branches. Its first version failed on Julia ≤ 1.10 for Diagonal, but none of the allocations came from extraction: the target function reduced with the no-function sum(z), and Base._sum(::Diagonal, ::Colon) allocates there -- 32 bytes even for a Diagonal{Float64}. Reducing with sum(f, z) measures ForwardDiff rather than LinearAlgebra, and with the reshape_jacobian fix above every path asserts zero outright, on 1.10 and on 1.12, for all four input types in both modes.

Asserted on top of that: a wrongly shaped jacobian! result and a gradient! result that is not shaped like x, a dense x included, both throw; chunk mode fills the MVector gradient buffer of an ImmutableDiffResult; the structured Jacobian and Hessian testsets run a chunk size that leaves a partial final chunk, which every size they had divided; SeedTest pins the columns of structural_columns against the index sets it already writes out by hand, since the derivatives land in the wrong columns if they disagree with structural_eachindex; the allocation test measures a dense result as well, which is the case that zeroes before extracting; and the JET tests cover a structured input, which reaches the generator branches. (On arm64 macOS with Julia 1.10.12, the older Test jacobian! allocations testset reports 16 bytes both here and on 569af35 -- they come from the target function's own broadcast over Vector{Dual}, not from extraction, and 1.12 reports zero.)

Beyond the suite: a sweep over the three wrappers × n ∈ {1,2,3,5,8} × every chunk size × every result container agrees with the closed-form expectation, with the Hessian symmetric and its structural rows/columns exactly zero; empty inputs, flat Vector results and BigFloat inputs with unassigned entries behave. On the GPU side the dense path -- the only one a GPU array reaches, since the new branch requires structural_length(x) != length(x) -- is unchanged, partials_wrap and the single fused broadcast intact; a JLArray run fails identically on this branch and on master, at the scalar indexing in seed! that #816 addresses.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 98.82353% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.24%. Comparing base (569af35) to head (3063bff).

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

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #840      +/-   ##
==========================================
+ Coverage   90.68%   91.24%   +0.55%     
==========================================
  Files          11       11              
  Lines        1052     1085      +33     
==========================================
+ Hits          954      990      +36     
+ Misses         98       95       -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.