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_linearindices(x) maps a structural position to a linear index of x without materializing anything: eachindex(IndexLinear(), x) and diagind(x) where a range already exists, a lazy generator over structural_eachindex for the triangular wrappers.

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.11 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 in place of the one reshape was performing on the way past, 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.

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. Dense inputs are unaffected: eachindex(x, result) only compares linear ranges, so e.g. a matrix result 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.

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 95.71429% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.00%. Comparing base (569af35) to head (ec3caa4).

Files with missing lines	Patch %	Lines
ext/ForwardDiffStaticArraysExt.jl	66.66%	1 Missing ⚠️
src/apiutils.jl	80.00%	1 Missing ⚠️
src/gradient.jl	96.42%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #840      +/-   ##
==========================================
+ Coverage   90.68%   91.00%   +0.31%     
==========================================
  Files          11       11              
  Lines        1052     1078      +26     
==========================================
+ Hits          954      981      +27     
+ 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.