Extract derivatives at the indices of x, and store those indices in the config

[devmotion]

Fixes #838, fixes #839 and fixes #842 — the same omission from #739 three times over: seeding became structure-aware, extraction and the config did not.

gradient! wrote to the wrong entries (#838)

extract_gradient!/extract_gradient_chunk! took their positions from the result, while the seeds are laid out along the structural positions of x. They now take the positions of x, which the config holds (see #842 below), and entries that receive no derivative are zeroed — that is their derivative.

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

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

Reusing a config across structures was silently wrong (#842)

seed! derived the positions to seed from the config's work buffer, while extraction derived them from x. A config reused with an input of the same length and a different structure therefore seeded one set of positions and extracted another — silently — and left the buffer entries outside the input's structure uninitialized for the target function to read (UndefRefError for a non-bits element type, garbage for Float64).

Taking the positions from x instead would leave those buffer entries never written at all. So structural_indices is derived from the buffer and stored on the config — a new indices field, one position vector per buffer a config owns — and every API entry now checks the config against the arrays it is used with, calling checkstructure next to checktag. A size or count comparison could not replace that check: LowerTriangular(n, n) and UpperTriangular(n, n) agree on size and on structural_length, yet their positions are disjoint off the diagonal.

julia> A = collect(reshape(1.0:9.0, 3, 3));

julia> ForwardDiff.gradient(f, A, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{2}()))
ERROR: ArgumentError: the config was built for an array of type UpperTriangular and cannot be
used with an array of type Array: the two are structured differently
# master: [1.0 4.0 5.0; 7.0 8.0 9.0; 0.0 0.0 0.0], silently

Separately, seeding's unassigned-entry branch only ever worked for a dense Array: Base._unsetindex! has no CartesianIndex method at all, and its AbstractArray fallback for a linear index recurses forever. With the positions now linear, adjoint, transpose, PermutedDimsArray and non-strided view inputs behave like a Matrix, their buffers being plain Arrays; the wrappers similar preserves raise an ArgumentError naming the entry, in place of a MethodError (the triangles) or a StackOverflowError (Diagonal).

It is also faster

The lazy position iterators were walked from the front to reach each chunk (Iterators.drop, which has no range specialization), three times per gradient chunk. Storing them turns that into view(indices, index:(index + count - 1)). Dense inputs paid the same walk over eachindex(duals, x), so they speed up too — best-of-200 gradient!/jacobian!, µs:

x	n	chunk	gradient!	jacobian!
Matrix	20	2	122.0 → 68.5	175.4 → 134.8
Matrix	40	2	1885 → 1108	2728 → 2198
Matrix	40	12	411 → 282	636 → 550
UpperTriangular	20	2	70.4 → 37.7	110.7 → 76.2
UpperTriangular	40	2	1092 → 581	1698 → 1168
UpperTriangular	40	12	253 → 168	389 → 294

Triangular buffers are now walked by linear index rather than by CartesianIndex, which is also what lets a position double as a Jacobian column number.

Breaking changes

The result conventions follow #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! requires the result to have the same length as x and to store every entry x does: gradient!(zeros(6), f, UpperTriangular(3×3)) throws DimensionMismatch rather than packing the derivatives in structural order, that order being an internal detail, and a structured result for a dense x throws ArgumentError. Shape is not constrained beyond the length, so zeros(9) for a 3×3 input and an UpperTriangular result for a Diagonal one both work.
• 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.
• Reusing a config across structures throws rather than computing the wrong derivatives (Reusing a GradientConfig/JacobianConfig with a differently structured input silently computes wrong derivatives #842 above).
• The configs gained a type parameter for the indices field, deliberately not shimmed, since pinning the parameter count means depending on internals. Two known cases: SimpleNonlinearSolve constructs ForwardDiff.JacobianConfig{tag,V,N,typeof(duals)}(seeds, duals) directly, a hard error, and Pumas types two struct fields as JacobianConfig{TJ,V,NJ,DJ} / GradientConfig{TG,Dual{…},NG,DG}, which stop being concrete.
• A config for an offset-indexed input throws at construction rather than at the first gradient/jacobian call, since structural_indices requires one-based indexing.

Smaller changes

Zeroing belongs to the sweep, not to a chunk. The chunk extractors recognised the first chunk by index == 1 and used it to zero the whole result — neither part of extracting one chunk nor something a chunk can decide, the entries at stake belonging to no chunk in particular. Both sweeps now zero once up front.

reshape_jacobian allocated a wrapper on every chunked call. It reshaped the result even when that was already a matrix, and since 1.12 reshape can no longer return its argument. It now shares the short-circuit extract_jacobian! was given in #797, with an explicit size check stricter than what reshape performed on the way past — that only ruled out a wrong total length, so a matrix of the right length and the wrong shape was silently reinterpreted.

x last, one job per function. extract_gradient! and extract_jacobian! take x after the derivatives, like every other internal function pairing an output with an input; the chunk extractors drop it, since they only write positions, which they now receive. extract_jacobian! duplicated the structured branch of extract_jacobian_chunk! at offset == 0 and now delegates to it, and its redundant n argument is gone. Both gradient sweeps reject an array-valued f up front, as the Jacobian sweeps do, so gradient on one points at jacobian instead of raising a MethodError from zero(::Type{<:AbstractArray}).

Tests

New testsets in GradientTest.jl (#838), JacobianTest.jl (#839) and HessianTest.jl (both, inherited) cover the three wrappers × sizes × every relevant chunk size against a dense result, a result shaped like x, DiffResults.GradientResult/JacobianResult/HessianResult, a dense DiffResult gradient buffer, and both Jacobian forms. Results are prefilled with NaN so untouched entries 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; against master the Hessian testset gets 0 passed, 3 failed, 4 errored. For #842, a config-reuse testset in each of DerivativeTest.jl, GradientTest.jl, JacobianTest.jl and HessianTest.jl runs the full four-kinds × four-kinds grid in both modes — including the LowerTriangular/UpperTriangular pair no size or count check can catch — and asserts the error messages, not just the types. AllocationsTest.jl asserts zero allocations outright for all four input types in both modes, which is also what pins the stored positions to being indexed through stack-allocated views.

Further test detail

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. The config-reuse testsets also assert that a wrong length still fails as a DimensionMismatch and that the fitting config still returns the right answer. DerivativeTest.jl had no structure coverage at all before, so it also pins the value and the derivative for a structured y.

The unassigned-entry cases are covered end-to-end in GradientTest.jl and per-operation in SeedTest.jl, which now pins structural_indices against the index sets it writes out by hand, asserts that a window overrunning the end is a BoundsError rather than a silently truncated chunk, and covers all four Array-buffer input kinds plus the three wrappers.

The first version of the allocation test 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. (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}, and 1.12 reports zero.)

Also asserted: the out-of-place gradient returns the structure of x; a structured result cannot hold the gradient of a dense x, while a result whose structure contains that of x can; a wrongly shaped jacobian! result and a gradient! result of the wrong length both throw; chunk mode fills the MVector gradient buffer of an ImmutableDiffResult; the structured Jacobian and Hessian testsets run a chunk size leaving a partial final chunk, which every size they had divided; and the JET tests cover a structured input.

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 is unchanged: partials_wrap is kept, and result[:, cols] .= … targets a contiguous view of the whole array there, cols being a UnitRange over every column. 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

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.54%. Comparing base (b742809) to head (ae5bdb8).

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #840      +/-   ##
==========================================
+ Coverage   90.68%   91.54%   +0.85%     
==========================================
  Files          11       11              
  Lines        1052     1076      +24     
==========================================
+ Hits          954      985      +31     
+ Misses         98       91       -7     

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

[andreasnoack]

Looks like your PRs need some JET related adjustments

[KristofferC]

Having these huge comments on almost every function getting changed is not great. It's a good practice to read through the comments and really think about if they have a reasonable noise-to-signal ratio and are written so a human actually understands them.

[devmotion]

I know 🙂

I noticed a few more issues (eg #842) so I'm still iterating quite a bit on the PR locally (actually it has diverged completely). It's not useful at all to spend time on reviewing it right now.