From 503fbec9c40a3150b5be1fe5979124b7a18787b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Mon, 17 Aug 2026 17:24:29 +0200
Subject: [PATCH 1/6] Extract derivatives at the indices of `x`, not of the
 result
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Since #739 only the structurally non-zero entries of an input are seeded,
but extraction was not updated to match, so the derivatives were written
to positions taken from the result container instead of from `x`.

`extract_gradient!`/`extract_gradient_chunk!` now take `x` and walk
`structural_eachindex(x, result)`. Entries that receive no derivative are
zeroed, which is their derivative; in chunk mode the first chunk does it.
The `DiffResult` method splits on mutability, since an immutable result
cannot be written to entry by entry (and only occurs for `StaticArray`
inputs, all of whose entries are structural). Fixes #838, where a dense
result got the derivatives at linear positions `1:structural_length(x)`
and a `DiffResults.GradientResult` threw, and with it the mis-scattered
gradient of `hessian!(::DiffResult, ...)`.

The Jacobian is indexed by the linear indices of `x`: column `j` holds
`∂f(x)[i]/∂x[j]`, as documented, with hard zeros in the columns of the
structural zeros. Its allocations therefore use `length(x)` rather than
`structural_length(x)`, which is what `reshape_jacobian` expected all
along, so chunk mode stops throwing. Fixes #839. `structural_linearindices`
maps a structural position to a linear index of `x` without materializing
anything, and the single-broadcast path is kept when every index of `x` is
structural, so no path allocates more than before.

This changes the shape of the result for structured inputs: `jacobian`
gains the zero columns and `hessian` inherits both conventions through
`jacobian(∇f, x)`, becoming `length(x) x length(x)` with hard-zero rows
and columns instead of mixing linear and structural indices. That also
makes `hessian!(DiffResults.HessianResult(x), ...)` work. For a `Diagonal`
the result now scales with `length(x)`, so differentiating with respect to
the diagonal vector is the better choice there.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 ext/ForwardDiffStaticArraysExt.jl |  6 +--
 src/ForwardDiff.jl                |  2 +-
 src/apiutils.jl                   |  9 +++++
 src/gradient.jl                   | 52 +++++++++++++++++--------
 src/jacobian.jl                   | 63 ++++++++++++++++++++-----------
 test/AllocationsTest.jl           | 36 ++++++++++++++++++
 test/GradientTest.jl              | 46 ++++++++++++++++++++++
 test/JacobianTest.jl              | 47 +++++++++++++++++++++++
 8 files changed, 218 insertions(+), 43 deletions(-)

diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl
index bf0ef99a..2e9999d1 100644
--- a/ext/ForwardDiffStaticArraysExt.jl
+++ b/ext/ForwardDiffStaticArraysExt.jl
@@ -57,7 +57,7 @@ end
 
 @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F}
     T = typeof(Tag(f, eltype(x)))
-    return extract_gradient!(T, result, f(dualize(T, x)))
+    return extract_gradient!(T, result, x, f(dualize(T, x)))
 end
 
 # Jacobian
@@ -87,13 +87,13 @@ end
 
 function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
     result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))
-    return extract_jacobian!(T, result, ydual, length(x))
+    return extract_jacobian!(T, result, x, ydual)
 end
 
 @inline function ForwardDiff.vector_mode_jacobian!(result, f::F, x::StaticArray) where {F}
     T = typeof(Tag(f, eltype(x)))
     ydual = f(dualize(T, x))
-    result = extract_jacobian!(T, result, ydual, length(x))
+    result = extract_jacobian!(T, result, x, ydual)
     result = extract_value!(T, result, ydual)
     return result
 end
diff --git a/src/ForwardDiff.jl b/src/ForwardDiff.jl
index b16b986b..3ccf9403 100644
--- a/src/ForwardDiff.jl
+++ b/src/ForwardDiff.jl
@@ -1,7 +1,7 @@
 module ForwardDiff
 
 using DiffRules, DiffResults
-using DiffResults: DiffResult, MutableDiffResult
+using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult
 using Preferences
 using Random
 using LinearAlgebra
diff --git a/src/apiutils.jl b/src/apiutils.jl
index 0615fdb3..c0560cc1 100644
--- a/src/apiutils.jl
+++ b/src/apiutils.jl
@@ -70,6 +70,15 @@ function structural_eachindex(x::Diagonal, y::AbstractArray)
     return diagind(x)
 end
 
+# The linear indices of the seeded entries of `x`, in seeding order. Results that are not shaped like
+# `x`, such as the columns of a Jacobian, are indexed by these.
+structural_linearindices(x::AbstractArray) = eachindex(IndexLinear(), x)
+structural_linearindices(x::Diagonal) = diagind(x)
+function structural_linearindices(x::Union{LowerTriangular,UpperTriangular})
+    lin = LinearIndices(x)
+    return (lin[idx] for idx in structural_eachindex(x))
+end
+
 # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is
 # currently carrying and to initialize a freshly allocated work buffer, whose elements must all be
 # written before the target function reads them.
diff --git a/src/gradient.jl b/src/gradient.jl
index a5ef3dac..29695cf1 100644
--- a/src/gradient.jl
+++ b/src/gradient.jl
@@ -49,44 +49,64 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is
 # result extraction #
 #####################
 
-function extract_gradient!(::Type{T}, result::DiffResult, y::Real) where {T}
+# Derivatives are only computed with respect to the structurally non-zero entries of `x`, since only
+# those are seeded. The positions to write to therefore have to be taken from `x`, not from `result`:
+# the two may have different structure, e.g. `DiffResults.HessianResult` allocates a dense gradient
+# buffer even for a structured `x`. The remaining entries of `result` are zeroed, their derivative
+# being zero, unless every index of `x` is structural. See #838.
+
+function extract_gradient!(::Type{T}, result::DiffResult, x, y::Real) where {T}
     result = DiffResults.value!(result, y)
     grad = DiffResults.gradient(result)
     fill!(grad, zero(y))
     return result
 end
 
-function extract_gradient!(::Type{T}, result::DiffResult, dual::Dual) where {T}
+function extract_gradient!(::Type{T}, result::MutableDiffResult, x, dual::Dual) where {T}
+    result = DiffResults.value!(result, value(T, dual))
+    extract_gradient!(T, DiffResults.gradient(result), x, dual)
+    return result
+end
+
+# Immutable results cannot be written to entry by entry. They only occur for `StaticArray` inputs,
+# all of whose entries are structural, so copying the partials wholesale is correct.
+function extract_gradient!(::Type{T}, result::ImmutableDiffResult, x, dual::Dual) where {T}
     result = DiffResults.value!(result, value(T, dual))
     result = DiffResults.gradient!(result, partials(T, dual))
     return result
 end
 
-extract_gradient!(::Type{T}, result::AbstractArray, y::Real) where {T} = fill!(result, zero(y))
-function extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual) where {T}
-    idxs = structural_eachindex(result)
+extract_gradient!(::Type{T}, result::AbstractArray, x, y::Real) where {T} = fill!(result, zero(y))
+function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) where {T}
+    structural_length(x) == length(x) || fill!(result, zero(valtype(T, dual)))
+    idxs = structural_eachindex(x, result)
     for (i, idx) in zip(1:npartials(dual), idxs)
         result[idx] = partials(T, dual, i)
     end
     return result
 end
 
-function extract_gradient_chunk!(::Type{T}, result, dual, index, chunksize) where {T}
+# The first chunk zeroes `result`; the entries it does not fill are written by the later chunks. In
+# chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk.
+function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) where {T}
     offset = index - 1
-    idxs = Iterators.drop(structural_eachindex(result), offset)
+    if iszero(offset) && structural_length(x) != length(x)
+        fill!(result, zero(valtype(T, dual)))
+    end
+    idxs = Iterators.drop(structural_eachindex(x, result), offset)
     for (i, idx) in zip(1:chunksize, idxs)
         result[idx] = partials(T, dual, i)
     end
     return result
 end
 
-function extract_gradient_chunk!(::Type{T}, result::DiffResult, dual, index, chunksize) where {T}
-    extract_gradient_chunk!(T, DiffResults.gradient(result), dual, index, chunksize)
+function extract_gradient_chunk!(::Type{T}, result::DiffResult, x, dual, index, chunksize) where {T}
+    extract_gradient_chunk!(T, DiffResults.gradient(result), x, dual, index, chunksize)
     return result
 end
 
-extract_gradient_chunk!(::Type, result, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR)
-extract_gradient_chunk!(::Type, result::DiffResult, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR)
+extract_gradient_chunk!(::Type, result, x, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR)
+extract_gradient_chunk!(::Type, result::DiffResult, x, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR)
 
 const GRAD_ERROR = DimensionMismatch("gradient(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?")
 
@@ -98,12 +118,12 @@ function vector_mode_gradient(f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa Real || throw(GRAD_ERROR)
     result = similar(x, valtype(T, ydual))
-    return extract_gradient!(T, result, ydual)
+    return extract_gradient!(T, result, x, ydual)
 end
 
 function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval!(f, cfg, x)
-    result = extract_gradient!(T, result, ydual)
+    result = extract_gradient!(T, result, x, ydual)
     return result
 end
 
@@ -134,7 +154,7 @@ function chunk_mode_gradient_expr(result_definition::Expr)
         seed_zero_partials!(xdual, x, N + 1, xlen - N)
         ydual = f(xdual)
         $(result_definition)
-        extract_gradient_chunk!(T, result, ydual, 1, N)
+        extract_gradient_chunk!(T, result, x, ydual, 1, N)
         seed_zero_partials!(xdual, x, 1)
 
         # do middle chunks
@@ -142,14 +162,14 @@ function chunk_mode_gradient_expr(result_definition::Expr)
             i = ((c - 1) * N + 1)
             seed!(xdual, x, i, seeds)
             ydual = f(xdual)
-            extract_gradient_chunk!(T, result, ydual, i, N)
+            extract_gradient_chunk!(T, result, x, ydual, i, N)
             seed_zero_partials!(xdual, x, i)
         end
 
         # do final chunk
         seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
         ydual = f(xdual)
-        extract_gradient_chunk!(T, result, ydual, lastchunkindex, lastchunksize)
+        extract_gradient_chunk!(T, result, x, ydual, lastchunkindex, lastchunksize)
 
         # get the value, this is a no-op unless result is a DiffResult
         extract_value!(T, result, ydual)
diff --git a/src/jacobian.jl b/src/jacobian.jl
index f14a6a7b..df1a8630 100644
--- a/src/jacobian.jl
+++ b/src/jacobian.jl
@@ -92,28 +92,49 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is
 # result extraction #
 #####################
 
-function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T}
-    out_reshaped = result isa AbstractMatrix ? result : reshape(result, length(ydual), n)
+# The Jacobian is indexed by the linear indices of `x`: column `j` holds the derivatives with respect
+# to `x[j]`. Only the seeded entries of `x` have a derivative to extract, so the columns of the
+# structurally zero ones are zeroed instead. See #839.
+
+function extract_jacobian!(::Type{T}, result::AbstractArray, x, ydual::AbstractArray) where {T}
+    out_reshaped = result isa AbstractMatrix ? result : reshape(result, length(ydual), length(x))
     ydual_reshaped = vec(ydual)
     # Use closure to avoid GPU broadcasting with Type
     partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
-    out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n))
+    n = structural_length(x)
+    if n == length(x)
+        out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n))
+    else
+        fill!(out_reshaped, zero(valtype(T, eltype(ydual))))
+        for (i, col) in zip(1:n, structural_linearindices(x))
+            out_reshaped[:, col] .= partials_wrap.(ydual_reshaped, i)
+        end
+    end
     return result
 end
 
-function extract_jacobian!(::Type{T}, result::MutableDiffResult, ydual::AbstractArray, n) where {T}
-    extract_jacobian!(T, DiffResults.jacobian(result), ydual, n)
+function extract_jacobian!(::Type{T}, result::MutableDiffResult, x, ydual::AbstractArray) where {T}
+    extract_jacobian!(T, DiffResults.jacobian(result), x, ydual)
     return result
 end
 
-function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) where {T}
+function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize) where {T}
     ydual_reshaped = vec(ydual)
     offset = index - 1
     irange = 1:chunksize
-    col = irange .+ offset
     # Use closure to avoid GPU broadcasting with Type
     partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
-    result[:, col] .= partials_wrap.(ydual_reshaped, transpose(irange))
+    if structural_length(x) == length(x)
+        result[:, irange .+ offset] .= partials_wrap.(ydual_reshaped, transpose(irange))
+    else
+        # The first chunk zeroes the columns of the structurally zero entries, which no chunk writes.
+        # In chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk.
+        iszero(offset) && fill!(result, zero(valtype(T, eltype(ydual))))
+        idxs = Iterators.drop(structural_linearindices(x), offset)
+        for (i, col) in zip(irange, idxs)
+            result[:, col] .= partials_wrap.(ydual_reshaped, i)
+        end
+    end
     return result
 end
 
@@ -125,38 +146,34 @@ reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResult
 ###############
 
 function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T}) where {F,T}
-    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa AbstractArray || throw(JACOBIAN_ERROR)
-    result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), N)
-    extract_jacobian!(T, result, ydual, N)
+    result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))
+    extract_jacobian!(T, result, x, ydual)
     extract_value!(T, result, ydual)
     return result
 end
 
 function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
-    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     map!(d -> value(T,d), y, ydual)
-    result = similar(y, length(y), N)
-    extract_jacobian!(T, result, ydual, N)
+    result = similar(y, length(y), length(x))
+    extract_jacobian!(T, result, x, ydual)
     map!(d -> value(T,d), y, ydual)
     return result
 end
 
 function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T}) where {F,T}
-    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f, cfg, x)
-    extract_jacobian!(T, result, ydual, N)
+    extract_jacobian!(T, result, x, ydual)
     extract_value!(T, result, ydual)
     return result
 end
 
 function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
-    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     map!(d -> value(T,d), y, ydual)
-    extract_jacobian!(T, result, ydual, N)
+    extract_jacobian!(T, result, x, ydual)
     extract_value!(T, result, y, ydual)
     return result
 end
@@ -192,7 +209,7 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex
         ydual isa AbstractArray || throw(JACOBIAN_ERROR)
         $(result_definition)
         out_reshaped = reshape_jacobian(result, ydual, xdual)
-        extract_jacobian_chunk!(T, out_reshaped, ydual, 1, N)
+        extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N)
         seed_zero_partials!(xdual, x, 1)
 
         # do middle chunks
@@ -200,14 +217,14 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex
             i = ((c - 1) * N + 1)
             seed!(xdual, x, i, seeds)
             $(compute_ydual)
-            extract_jacobian_chunk!(T, out_reshaped, ydual, i, N)
+            extract_jacobian_chunk!(T, out_reshaped, x, ydual, i, N)
             seed_zero_partials!(xdual, x, i)
         end
 
         # do final chunk
         seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
         $(compute_ydual)
-        extract_jacobian_chunk!(T, out_reshaped, ydual, lastchunkindex, lastchunksize)
+        extract_jacobian_chunk!(T, out_reshaped, x, ydual, lastchunkindex, lastchunksize)
 
         $(y_definition)
 
@@ -218,14 +235,14 @@ end
 @eval function chunk_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
     $(jacobian_chunk_mode_expr(:(xdual = cfg.duals),
                                :(ydual = f(xdual)),
-                               :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)),
+                               :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))),
                                :()))
 end
 
 @eval function chunk_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N}
     $(jacobian_chunk_mode_expr(:((ydual, xdual) = cfg.duals),
                                :(f!(seed_zero_partials!(ydual, y), xdual)),
-                               :(result = similar(y, length(y), xlen)),
+                               :(result = similar(y, length(y), length(x))),
                                :(map!(d -> value(T,d), y, ydual))))
 end
 
diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl
index 94e7cddd..f4bd2134 100644
--- a/test/AllocationsTest.jl
+++ b/test/AllocationsTest.jl
@@ -1,6 +1,7 @@
 module AllocationsTest
 
 using ForwardDiff
+using LinearAlgebra
 using StaticArrays
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
@@ -50,6 +51,41 @@ end
     @test iszero(allocs_jacobian!())
 end
 
+# `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural
+# position to an index of `x` must not allocate, whether or not `x` has structurally zero entries.
+function allocs_gradient!(x, chunk)
+    f(z) = sum(abs2, z)
+    result = fill!(similar(x, Float64), false)
+    cfg = ForwardDiff.GradientConfig(f, x, chunk)
+    ForwardDiff.gradient!(result, f, x, cfg)  # warmup
+    return @allocated ForwardDiff.gradient!(result, f, x, cfg)
+end
+
+function allocs_jacobian!(x, chunk)
+    f!(y, z) = (y[1] = sum(z); y[2] = sum(abs2, z); y)
+    y = zeros(2)
+    result = zeros(2, length(x))
+    cfg = ForwardDiff.JacobianConfig(f!, y, x, chunk)
+    ForwardDiff.jacobian!(result, f!, y, x, cfg)  # warmup
+    return @allocated ForwardDiff.jacobian!(result, f!, y, x, cfg)
+end
+
+@testset "Test gradient!/jacobian! allocations for $(nameof(typeof(x)))" for x in (rand(6, 6),
+                                                                                  LowerTriangular(rand(6, 6)),
+                                                                                  UpperTriangular(rand(6, 6)),
+                                                                                  Diagonal(rand(6, 6)))
+    # vector mode
+    chunk = ForwardDiff.Chunk{ForwardDiff.structural_length(x)}()
+    @test iszero(allocs_gradient!(x, chunk))
+    @test iszero(allocs_jacobian!(x, chunk))
+
+    # chunk mode, where `jacobian!` allocates a fixed-size `reshape` wrapper for every input, so
+    # compare against a dense input of the same size instead of asserting zero
+    chunk = ForwardDiff.Chunk{2}()
+    @test iszero(allocs_gradient!(x, chunk))
+    @test allocs_jacobian!(x, chunk) == allocs_jacobian!(rand(6, 6), chunk)
+end
+
 @testset "allocation-free nested StaticArray jacobian" begin
     # test that nested jacobians of StaticArrays do not allocate. 
     # This is a regression test for issue #798, where the inner jacobian was allocating
diff --git a/test/GradientTest.jl b/test/GradientTest.jl
index bf121239..117e8873 100644
--- a/test/GradientTest.jl
+++ b/test/GradientTest.jl
@@ -275,6 +275,52 @@ end
     end
 end
 
+# issue #838
+@testset "structured inputs: extraction positions" begin
+    # The seeds are laid out along `structural_eachindex(x)`, so the derivatives have to be written to
+    # the corresponding entries of the result and every other entry has to end up at zero. In
+    # particular this must not be read off `result`, which carries no structure to read it off of when
+    # it is dense -- as the gradient buffer of a `DiffResults.HessianResult` is even for a structured
+    # `x`. All results are prefilled with `NaN` so that entries left untouched are caught.
+    @testset "$T, n = $n" for T in (LowerTriangular, UpperTriangular, Diagonal), n in (3, 10)
+        M = rand(n, n)
+        x = T(randn(n, n))
+        f = z -> dot(M, z)
+        expected = T(M)                     # zero derivative for the structurally zero entries
+        dense_expected = Matrix(expected)
+        val = f(x)
+        nstruct = ForwardDiff.structural_length(x)
+
+        @testset "chunk size = $c" for c in unique((1, 2, nstruct))
+            cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}())
+
+            out = fill(NaN, n, n)
+            @test ForwardDiff.gradient!(out, f, x, cfg) === out
+            @test out == dense_expected
+
+            out = T(fill(NaN, n, n))
+            ForwardDiff.gradient!(out, f, x, cfg)
+            @test out == expected
+
+            # gradient buffer shaped like `x`, cf. `DiffResults.GradientResult`
+            result = DiffResults.GradientResult(x)
+            result = ForwardDiff.gradient!(result, f, x, cfg)
+            @test DiffResults.gradient(result) == expected
+            @test DiffResults.value(result) ≈ val
+
+            # dense gradient buffer, cf. `DiffResults.HessianResult`
+            result = DiffResults.DiffResult(NaN, fill(NaN, n, n))
+            result = ForwardDiff.gradient!(result, f, x, cfg)
+            @test DiffResults.gradient(result) == dense_expected
+            @test DiffResults.value(result) ≈ val
+
+            # the result has to be shaped like `x`, packing into the structural positions is not
+            # supported since their order is an implementation detail
+            @test_throws DimensionMismatch ForwardDiff.gradient!(fill(NaN, nstruct), f, x, cfg)
+        end
+    end
+end
+
 # issue #769
 @testset "functions with `Dual` output" begin
     x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]
diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl
index b6d36180..9763d279 100644
--- a/test/JacobianTest.jl
+++ b/test/JacobianTest.jl
@@ -322,6 +322,53 @@ end
     end
 end
 
+# issue #839
+@testset "structured inputs: $(nameof(typeof(x)))" for (x, sidx) in (
+        # The Jacobian is indexed by the linear indices of `x`: column `j` holds the derivatives with
+        # respect to `x[j]`, and the columns of the structurally zero entries are zero. The nonzero
+        # columns are written out by hand so that a bug in the position mapping cannot hide inside the
+        # reference. Only the full-length chunk worked before, the others threw.
+        (LowerTriangular(randn(4, 4)), [i + 4 * (j - 1) for j in 1:4 for i in j:4]),
+        (UpperTriangular(randn(4, 4)), [i + 4 * (j - 1) for j in 1:4 for i in 1:j]),
+        (Diagonal(randn(4, 4)),        collect(1:5:16)),
+    )
+    g = z -> [sum(z), sum(abs2, z)]
+    g! = (y, z) -> (y[1] = sum(z); y[2] = sum(abs2, z); y)
+
+    expected = zeros(2, length(x))
+    expected[1, sidx] .= 1
+    expected[2, sidx] .= 2 .* x[sidx]
+    val = g(x)
+
+    @testset "chunk size = $c" for c in unique((1, 2, ForwardDiff.structural_length(x)))
+        cfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{c}())
+        J = ForwardDiff.jacobian(g, x, cfg)
+        @test size(J) == (2, length(x))
+        @test J == expected
+
+        out = fill(NaN, 2, length(x))
+        @test ForwardDiff.jacobian!(out, g, x, cfg) === out
+        @test out == expected
+
+        # `DiffResults.JacobianResult` allocates `length(x)` columns, which is what is needed
+        result = DiffResults.JacobianResult(similar(val), x)
+        result = ForwardDiff.jacobian!(result, g, x, cfg)
+        @test DiffResults.jacobian(result) == expected
+        @test DiffResults.value(result) ≈ val
+
+        # in-place target function
+        cfg! = ForwardDiff.JacobianConfig(g!, similar(val), x, ForwardDiff.Chunk{c}())
+        y = fill(NaN, 2)
+        @test ForwardDiff.jacobian(g!, y, x, cfg!) == expected
+        @test y ≈ val
+        out = fill(NaN, 2, length(x))
+        y = fill(NaN, 2)
+        ForwardDiff.jacobian!(out, g!, y, x, cfg!)
+        @test out == expected
+        @test y ≈ val
+    end
+end
+
 # issue #769
 @testset "functions with `Dual` output" begin
     x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]

From fea566c9030d56043685f81e13d2ff0d162e6dd0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Mon, 17 Aug 2026 21:17:53 +0200
Subject: [PATCH 2/6] Fix the `Diagonal` allocation test, and the `reshape`
 wrapper it uncovered

The new allocation test failed on Julia <= 1.10 for `Diagonal` inputs, 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)`
instead measures ForwardDiff rather than LinearAlgebra, and extraction turns
out to be allocation-free for every input type on both 1.10 and 1.12.

That left the chunk-mode comparison against a dense input, which was hiding a
real cost: `reshape_jacobian` reshapes the result even when it already is a
matrix, and since 1.11 `reshape` can no longer return its argument, so every
chunk-mode `jacobian!` allocated an `Array` wrapper. `extract_jacobian!` had
been given that short-circuit in #797; `reshape_jacobian` now shares it, with
an explicit size check in place of the one `reshape` performed on the way
past, and `extract_jacobian!` calls it instead of repeating the ternary. Both
modes now reject a wrongly shaped matrix result with the same error, and the
test can assert zero allocations outright.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 src/jacobian.jl         | 15 +++++++++++----
 test/AllocationsTest.jl | 17 ++++++++---------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/jacobian.jl b/src/jacobian.jl
index df1a8630..ddd258e9 100644
--- a/src/jacobian.jl
+++ b/src/jacobian.jl
@@ -97,7 +97,7 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is
 # structurally zero ones are zeroed instead. See #839.
 
 function extract_jacobian!(::Type{T}, result::AbstractArray, x, ydual::AbstractArray) where {T}
-    out_reshaped = result isa AbstractMatrix ? result : reshape(result, length(ydual), length(x))
+    out_reshaped = reshape_jacobian(result, ydual, x)
     ydual_reshaped = vec(ydual)
     # Use closure to avoid GPU broadcasting with Type
     partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
@@ -138,8 +138,15 @@ function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize)
     return result
 end
 
-reshape_jacobian(result, ydual, xdual) = reshape(result, length(ydual), length(xdual))
-reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResults.jacobian(result), ydual, xdual)
+# A matrix is used as is: reshaping it would allocate a wrapper on Julia >= 1.11, where `reshape`
+# can no longer return its argument. The size is checked instead, as `reshape` did on the way past.
+function reshape_jacobian(result::AbstractMatrix, ydual, x)
+    size(result) == (length(ydual), length(x)) || throw(DimensionMismatch(
+        lazy"cannot store the $(length(ydual))x$(length(x)) Jacobian in a result of size $(size(result))"))
+    return result
+end
+reshape_jacobian(result::AbstractArray, ydual, x) = reshape(result, length(ydual), length(x))
+reshape_jacobian(result::DiffResult, ydual, x) = reshape_jacobian(DiffResults.jacobian(result), ydual, x)
 
 ###############
 # vector mode #
@@ -208,7 +215,7 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex
         $(compute_ydual)
         ydual isa AbstractArray || throw(JACOBIAN_ERROR)
         $(result_definition)
-        out_reshaped = reshape_jacobian(result, ydual, xdual)
+        out_reshaped = reshape_jacobian(result, ydual, x)
         extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N)
         seed_zero_partials!(xdual, x, 1)
 
diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl
index f4bd2134..53eb5854 100644
--- a/test/AllocationsTest.jl
+++ b/test/AllocationsTest.jl
@@ -53,7 +53,7 @@ end
 
 # `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural
 # position to an index of `x` must not allocate, whether or not `x` has structurally zero entries.
-function allocs_gradient!(x, chunk)
+function allocs_structured_gradient!(x, chunk)
     f(z) = sum(abs2, z)
     result = fill!(similar(x, Float64), false)
     cfg = ForwardDiff.GradientConfig(f, x, chunk)
@@ -61,8 +61,8 @@ function allocs_gradient!(x, chunk)
     return @allocated ForwardDiff.gradient!(result, f, x, cfg)
 end
 
-function allocs_jacobian!(x, chunk)
-    f!(y, z) = (y[1] = sum(z); y[2] = sum(abs2, z); y)
+function allocs_structured_jacobian!(x, chunk)
+    f!(y, z) = (y[1] = sum(abs2, z); y[2] = sqrt(sum(abs2, z)); y)
     y = zeros(2)
     result = zeros(2, length(x))
     cfg = ForwardDiff.JacobianConfig(f!, y, x, chunk)
@@ -76,14 +76,13 @@ end
                                                                                   Diagonal(rand(6, 6)))
     # vector mode
     chunk = ForwardDiff.Chunk{ForwardDiff.structural_length(x)}()
-    @test iszero(allocs_gradient!(x, chunk))
-    @test iszero(allocs_jacobian!(x, chunk))
+    @test iszero(allocs_structured_gradient!(x, chunk))
+    @test iszero(allocs_structured_jacobian!(x, chunk))
 
-    # chunk mode, where `jacobian!` allocates a fixed-size `reshape` wrapper for every input, so
-    # compare against a dense input of the same size instead of asserting zero
+    # chunk mode
     chunk = ForwardDiff.Chunk{2}()
-    @test iszero(allocs_gradient!(x, chunk))
-    @test allocs_jacobian!(x, chunk) == allocs_jacobian!(rand(6, 6), chunk)
+    @test iszero(allocs_structured_gradient!(x, chunk))
+    @test iszero(allocs_structured_jacobian!(x, chunk))
 end
 
 @testset "allocation-free nested StaticArray jacobian" begin

From 07e42721f07bdcbdd611128c0d502461ad0bfef8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Mon, 17 Aug 2026 21:49:10 +0200
Subject: [PATCH 3/6] Zero the unwritten entries in the sweep, not in the chunk

`extract_gradient_chunk!`/`extract_jacobian_chunk!` recognised the first chunk
by `index == 1` and used it to zero the whole result, which 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. Both
sweeps now do it once up front, and the chunk functions only write their chunk.

The gradient's half is shared with `extract_gradient!` as `zero_unseeded!`,
which also fixes the condition. `structural_length(x) != length(x)` zeroed a
result that is itself structured, whose every stored entry the sweep goes on to
write; comparing against `structural_length(result)` skips that, and leaves the
mismatched-structure cases erroring at the same write as before. The Jacobian
keeps its own test, since its result is not shaped like `x` and what has to be
covered there is columns.

Also drops the `map!` that `vector_mode_jacobian(f!, ...)` ran before
`extract_jacobian!`, which reads only `ydual`, and that the `map!` after it
repeats.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 src/gradient.jl | 27 ++++++++++++++++++---------
 src/jacobian.jl |  6 ++----
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/src/gradient.jl b/src/gradient.jl
index 29695cf1..d3c68dbc 100644
--- a/src/gradient.jl
+++ b/src/gradient.jl
@@ -53,7 +53,8 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is
 # those are seeded. The positions to write to therefore have to be taken from `x`, not from `result`:
 # the two may have different structure, e.g. `DiffResults.HessianResult` allocates a dense gradient
 # buffer even for a structured `x`. The remaining entries of `result` are zeroed, their derivative
-# being zero, unless every index of `x` is structural. See #838.
+# being zero, unless the seeded entries of `x` already account for every entry `result` stores. See
+# #838.
 
 function extract_gradient!(::Type{T}, result::DiffResult, x, y::Real) where {T}
     result = DiffResults.value!(result, y)
@@ -68,17 +69,29 @@ function extract_gradient!(::Type{T}, result::MutableDiffResult, x, dual::Dual)
     return result
 end
 
-# Immutable results cannot be written to entry by entry. They only occur for `StaticArray` inputs,
-# all of whose entries are structural, so copying the partials wholesale is correct.
+# Immutable results cannot be written to entry by entry. Copying the partials wholesale is correct
+# as long as every entry of `x` is seeded, which holds for the `StaticArray` gradient buffers that
+# are the only source of such results; anything else throws on the length mismatch.
 function extract_gradient!(::Type{T}, result::ImmutableDiffResult, x, dual::Dual) where {T}
     result = DiffResults.value!(result, value(T, dual))
     result = DiffResults.gradient!(result, partials(T, dual))
     return result
 end
 
+# Zeroes the entries that receive no derivative. In chunk mode the sweep calls this once up front,
+# since the entries that no chunk writes belong to none of them in particular.
+function zero_unseeded!(::Type{T}, result::AbstractArray, x, dual) where {T}
+    structural_length(x) == structural_length(result) || fill!(result, zero(valtype(T, dual)))
+    return result
+end
+function zero_unseeded!(::Type{T}, result::MutableDiffResult, x, dual) where {T}
+    zero_unseeded!(T, DiffResults.gradient(result), x, dual)
+    return result
+end
+
 extract_gradient!(::Type{T}, result::AbstractArray, x, y::Real) where {T} = fill!(result, zero(y))
 function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) where {T}
-    structural_length(x) == length(x) || fill!(result, zero(valtype(T, dual)))
+    zero_unseeded!(T, result, x, dual)
     idxs = structural_eachindex(x, result)
     for (i, idx) in zip(1:npartials(dual), idxs)
         result[idx] = partials(T, dual, i)
@@ -86,13 +99,8 @@ function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) wher
     return result
 end
 
-# The first chunk zeroes `result`; the entries it does not fill are written by the later chunks. In
-# chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk.
 function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) where {T}
     offset = index - 1
-    if iszero(offset) && structural_length(x) != length(x)
-        fill!(result, zero(valtype(T, dual)))
-    end
     idxs = Iterators.drop(structural_eachindex(x, result), offset)
     for (i, idx) in zip(1:chunksize, idxs)
         result[idx] = partials(T, dual, i)
@@ -154,6 +162,7 @@ function chunk_mode_gradient_expr(result_definition::Expr)
         seed_zero_partials!(xdual, x, N + 1, xlen - N)
         ydual = f(xdual)
         $(result_definition)
+        zero_unseeded!(T, result, x, ydual)
         extract_gradient_chunk!(T, result, x, ydual, 1, N)
         seed_zero_partials!(xdual, x, 1)
 
diff --git a/src/jacobian.jl b/src/jacobian.jl
index ddd258e9..752c817b 100644
--- a/src/jacobian.jl
+++ b/src/jacobian.jl
@@ -127,9 +127,6 @@ function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize)
     if structural_length(x) == length(x)
         result[:, irange .+ offset] .= partials_wrap.(ydual_reshaped, transpose(irange))
     else
-        # The first chunk zeroes the columns of the structurally zero entries, which no chunk writes.
-        # In chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk.
-        iszero(offset) && fill!(result, zero(valtype(T, eltype(ydual))))
         idxs = Iterators.drop(structural_linearindices(x), offset)
         for (i, col) in zip(irange, idxs)
             result[:, col] .= partials_wrap.(ydual_reshaped, i)
@@ -163,7 +160,6 @@ end
 
 function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
     ydual = vector_mode_dual_eval!(f!, cfg, y, x)
-    map!(d -> value(T,d), y, ydual)
     result = similar(y, length(y), length(x))
     extract_jacobian!(T, result, x, ydual)
     map!(d -> value(T,d), y, ydual)
@@ -216,6 +212,8 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex
         ydual isa AbstractArray || throw(JACOBIAN_ERROR)
         $(result_definition)
         out_reshaped = reshape_jacobian(result, ydual, x)
+        # zero the columns of the structurally zero entries of `x`, which no chunk of the sweep writes
+        structural_length(x) == length(x) || fill!(out_reshaped, zero(valtype(T, eltype(ydual))))
         extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N)
         seed_zero_partials!(xdual, x, 1)
 

From ec3caa40d342a6509ecf7b1ed9a49d0f7595b063 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Mon, 17 Aug 2026 22:55:02 +0200
Subject: [PATCH 4/6] Cover the Hessian, and the cases the extraction tests
 left out

`hessian` inherits both conventions through `jacobian(gradient(f), x)`, so it
changes shape for a structured `x` and `hessian!(DiffResults.HessianResult(x),
f, x)` starts working, none of which was asserted. The new testset pins the
value, the gradient and the Hessian of 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 checked rather than approximated. Against
master it fails everywhere: the smaller chunks throw, and the full-length one
gets the wrong shape and cannot take a `HessianResult` at all.

Also: the out-of-place `gradient` is shaped like `x`, so its zeros are the
structural ones and its type is worth asserting; a structured result cannot
hold the gradient of a dense `x`, which now throws where it used to write to
the wrong entries; a Jacobian result that is not a matrix is reshaped; and the
`f!` form takes a `JacobianResult` too.

Both structured testsets take the chunk sizes from `length(sidx)` rather than
`ForwardDiff.structural_length`, keeping the reference data in the test.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 test/GradientTest.jl | 10 ++++++++++
 test/HessianTest.jl  | 42 ++++++++++++++++++++++++++++++++++++++++++
 test/JacobianTest.jl | 12 +++++++++++-
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/test/GradientTest.jl b/test/GradientTest.jl
index 117e8873..8f618fe5 100644
--- a/test/GradientTest.jl
+++ b/test/GradientTest.jl
@@ -294,6 +294,11 @@ end
         @testset "chunk size = $c" for c in unique((1, 2, nstruct))
             cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}())
 
+            # the allocated result is shaped like `x`, so its zeros are the structural ones
+            grad = ForwardDiff.gradient(f, x, cfg)
+            @test grad isa T
+            @test grad == expected
+
             out = fill(NaN, n, n)
             @test ForwardDiff.gradient!(out, f, x, cfg) === out
             @test out == dense_expected
@@ -317,6 +322,11 @@ end
             # the result has to be shaped like `x`, packing into the structural positions is not
             # supported since their order is an implementation detail
             @test_throws DimensionMismatch ForwardDiff.gradient!(fill(NaN, nstruct), f, x, cfg)
+
+            # every entry of a dense `x` is seeded, so a structured result cannot hold its gradient
+            dense_x = Matrix(x)
+            dense_cfg = ForwardDiff.GradientConfig(f, dense_x, ForwardDiff.Chunk{c}())
+            @test_throws ArgumentError ForwardDiff.gradient!(T(fill(NaN, n, n)), f, dense_x, dense_cfg)
         end
     end
 end
diff --git a/test/HessianTest.jl b/test/HessianTest.jl
index 8be72ee5..01d29143 100644
--- a/test/HessianTest.jl
+++ b/test/HessianTest.jl
@@ -156,6 +156,48 @@ for T in (StaticArrays.SArray, StaticArrays.MArray)
     @test DiffResults.hessian(sresult3) == DiffResults.hessian(result)
 end
 
+# issues #838 and #839, which `hessian` inherits through `jacobian(gradient(f), x)`
+@testset "structured inputs: $(nameof(W))" for (W, sidx) in (
+        # both axes are indexed by the linear indices of `x`, hard zeros off the structure
+        (LowerTriangular, [i + 3 * (j - 1) for j in 1:3 for i in j:3]),
+        (UpperTriangular, [i + 3 * (j - 1) for j in 1:3 for i in 1:j]),
+        (Diagonal,        1:4:9),
+    )
+    x = W(randn(3, 3))
+    # d²f/dx[a]dx[b] is `1 + (a == b)` for structural `a`, `b`, and zero everywhere else
+    f = z -> (sum(abs2, z) + sum(z)^2) / 2
+    L = length(x)
+
+    expected = zeros(L, L)
+    expected[sidx, sidx] .= 1
+    for k in sidx
+        expected[k, k] += 1
+    end
+    val = f(x)
+    grad = zeros(3, 3)
+    grad[sidx] .= x[sidx] .+ sum(x)
+
+    @testset "chunk size = $c" for c in unique((1, 2, length(sidx)))
+        cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}())
+
+        H = ForwardDiff.hessian(f, x, cfg)
+        @test size(H) == (L, L)
+        @test H == expected
+
+        out = fill(NaN, L, L)
+        @test ForwardDiff.hessian!(out, f, x, cfg) === out
+        @test out == expected
+
+        # `DiffResults.HessianResult` allocates a dense gradient buffer even for a structured `x`
+        result = DiffResults.HessianResult(x)
+        result = ForwardDiff.hessian!(result, f, x,
+                                      ForwardDiff.HessianConfig(f, result, x, ForwardDiff.Chunk{c}()))
+        @test DiffResults.value(result) ≈ val
+        @test DiffResults.gradient(result) == grad
+        @test DiffResults.hessian(result) == expected
+    end
+end
+
 @testset "branches in dot" begin
     # https://github.com/JuliaDiff/ForwardDiff.jl/issues/551
     H = [1 2 3; 4 5 6; 7 8 9];
diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl
index 9763d279..0117d844 100644
--- a/test/JacobianTest.jl
+++ b/test/JacobianTest.jl
@@ -340,7 +340,7 @@ end
     expected[2, sidx] .= 2 .* x[sidx]
     val = g(x)
 
-    @testset "chunk size = $c" for c in unique((1, 2, ForwardDiff.structural_length(x)))
+    @testset "chunk size = $c" for c in unique((1, 2, length(sidx)))
         cfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{c}())
         J = ForwardDiff.jacobian(g, x, cfg)
         @test size(J) == (2, length(x))
@@ -350,6 +350,11 @@ end
         @test ForwardDiff.jacobian!(out, g, x, cfg) === out
         @test out == expected
 
+        # a result that is not a matrix is reshaped to one
+        out = fill(NaN, 2 * length(x))
+        @test ForwardDiff.jacobian!(out, g, x, cfg) === out
+        @test reshape(out, 2, length(x)) == expected
+
         # `DiffResults.JacobianResult` allocates `length(x)` columns, which is what is needed
         result = DiffResults.JacobianResult(similar(val), x)
         result = ForwardDiff.jacobian!(result, g, x, cfg)
@@ -366,6 +371,11 @@ end
         ForwardDiff.jacobian!(out, g!, y, x, cfg!)
         @test out == expected
         @test y ≈ val
+        result = DiffResults.JacobianResult(similar(val), x)
+        y = fill(NaN, 2)
+        result = ForwardDiff.jacobian!(result, g!, y, x, cfg!)
+        @test DiffResults.jacobian(result) == expected
+        @test DiffResults.value(result) ≈ val
     end
 end
 

From 27033bbeedf810d72d822a93001d866cd9ce6d71 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= <d.widmann@pumas.ai>
Date: Tue, 18 Aug 2026 17:23:08 +0200
Subject: [PATCH 5/6] Take `x` last, and extract a Jacobian column 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
`extract_jacobian(T, ydual, x::StaticArray)` of the StaticArrays extension. `x` is
annotated `AbstractArray` there, 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.
`structural_linearindices(x)` becomes `structural_columns(out, x)`: what a caller needs
are the columns of `out` that receive a derivative, and every method returns those,
`axes(out, 2)` or a subset of it, without materializing anything. It checks its
arguments the way `structural_eachindex` does, whose three bare `DimensionMismatch()`
now name the sizes they expected.

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 would fail on `zero(::Type{<:AbstractArray})`
before dispatch got there. `zero_unseeded!` also dispatches on `DiffResult` rather than
`MutableDiffResult`, since a `StaticArray` gradient buffer makes the result immutable
even when the buffer itself can be written to entry by entry, as an `MVector` can.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 ext/ForwardDiffStaticArraysExt.jl |  6 ++--
 src/apiutils.jl                   | 55 +++++++++++++++++++++---------
 src/gradient.jl                   | 55 ++++++++++++++++--------------
 src/jacobian.jl                   | 56 ++++++++++++++-----------------
 4 files changed, 98 insertions(+), 74 deletions(-)

diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl
index 2e9999d1..6d29f78c 100644
--- a/ext/ForwardDiffStaticArraysExt.jl
+++ b/ext/ForwardDiffStaticArraysExt.jl
@@ -57,7 +57,7 @@ end
 
 @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F}
     T = typeof(Tag(f, eltype(x)))
-    return extract_gradient!(T, result, x, f(dualize(T, x)))
+    return extract_gradient!(T, result, f(dualize(T, x)), x)
 end
 
 # Jacobian
@@ -87,13 +87,13 @@ end
 
 function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T
     result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))
-    return extract_jacobian!(T, result, x, ydual)
+    return extract_jacobian!(T, result, ydual, x)
 end
 
 @inline function ForwardDiff.vector_mode_jacobian!(result, f::F, x::StaticArray) where {F}
     T = typeof(Tag(f, eltype(x)))
     ydual = f(dualize(T, x))
-    result = extract_jacobian!(T, result, x, ydual)
+    result = extract_jacobian!(T, result, ydual, x)
     result = extract_value!(T, result, ydual)
     return result
 end
diff --git a/src/apiutils.jl b/src/apiutils.jl
index c0560cc1..081e40fd 100644
--- a/src/apiutils.jl
+++ b/src/apiutils.jl
@@ -48,35 +48,60 @@ function structural_eachindex(x::AbstractArray, y::AbstractArray)
 end
 function structural_eachindex(x::UpperTriangular, y::AbstractArray)
     require_one_based_indexing(x, y)
-    if size(x) != size(y)
-        throw(DimensionMismatch())
-    end
+    check_matching_size(x, y)
     n = size(x, 1)
     return (CartesianIndex(i, j) for j in 1:n for i in 1:j)
 end
 function structural_eachindex(x::LowerTriangular, y::AbstractArray)
     require_one_based_indexing(x, y)
-    if size(x) != size(y)
-        throw(DimensionMismatch())
-    end
+    check_matching_size(x, y)
     n = size(x, 1)
     return (CartesianIndex(i, j) for j in 1:n for i in j:n)
 end
 function structural_eachindex(x::Diagonal, y::AbstractArray)
     require_one_based_indexing(x, y)
-    if size(x) != size(y)
-        throw(DimensionMismatch())
-    end
+    check_matching_size(x, y)
     return diagind(x)
 end
 
-# The linear indices of the seeded entries of `x`, in seeding order. Results that are not shaped like
-# `x`, such as the columns of a Jacobian, are indexed by these.
-structural_linearindices(x::AbstractArray) = eachindex(IndexLinear(), x)
-structural_linearindices(x::Diagonal) = diagind(x)
-function structural_linearindices(x::Union{LowerTriangular,UpperTriangular})
+# The two arrays are indexed by the same indices, so they have to have the same size. This is the
+# error a `gradient!` into a container that is not shaped like `x` aborts with, so it names the sizes.
+function check_matching_size(x::AbstractArray, y::AbstractArray)
+    size(x) == size(y) || throw(DimensionMismatch(
+        lazy"expected an array of size $(size(x)), got an array of size $(size(y))"))
+    return nothing
+end
+
+# The columns of the Jacobian `out` that receive derivatives, in seeding order. Column `j` holds the
+# derivatives with respect to `x[j]`, so a seeded entry writes to the column at its position in the
+# linear order of `x`. Every entry of an array is seeded unless one of the methods below applies.
+function structural_columns(out::AbstractMatrix, x::AbstractArray)
+    require_one_based_indexing(out, x)
+    check_matching_columns(out, x)
+    return axes(out, 2)
+end
+# The seeded columns are runs of increasing length, so they are no range. Deriving their order from
+# `structural_eachindex` rather than recomputing it keeps a single source of truth: the two have to
+# agree entry by entry, or the derivatives land in the wrong columns.
+function structural_columns(out::AbstractMatrix, x::Union{LowerTriangular,UpperTriangular})
+    require_one_based_indexing(out, x)
+    check_matching_columns(out, x)
+    cols = axes(out, 2)
     lin = LinearIndices(x)
-    return (lin[idx] for idx in structural_eachindex(x))
+    return (cols[lin[idx]] for idx in structural_eachindex(x))
+end
+# `diagind` is already a range of linear positions, so it can select the columns directly.
+function structural_columns(out::AbstractMatrix, x::Diagonal)
+    require_one_based_indexing(out, x)
+    check_matching_columns(out, x)
+    return axes(out, 2)[diagind(x)]
+end
+
+# A column of the Jacobian belongs to an entry of `x`, so there have to be as many as `x` has entries.
+function check_matching_columns(out::AbstractMatrix, x::AbstractArray)
+    size(out, 2) == length(x) || throw(DimensionMismatch(
+        lazy"expected a matrix with $(length(x)) columns, got a matrix with $(size(out, 2)) columns"))
+    return nothing
 end
 
 # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is
diff --git a/src/gradient.jl b/src/gradient.jl
index d3c68dbc..b05c5fa8 100644
--- a/src/gradient.jl
+++ b/src/gradient.jl
@@ -53,45 +53,50 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is
 # those are seeded. The positions to write to therefore have to be taken from `x`, not from `result`:
 # the two may have different structure, e.g. `DiffResults.HessianResult` allocates a dense gradient
 # buffer even for a structured `x`. The remaining entries of `result` are zeroed, their derivative
-# being zero, unless the seeded entries of `x` already account for every entry `result` stores. See
-# #838.
+# being zero, unless `x` has as many seeded entries as `result` has entries. See #838.
 
-function extract_gradient!(::Type{T}, result::DiffResult, x, y::Real) where {T}
+function extract_gradient!(::Type{T}, result::DiffResult, y::Real, x) where {T}
     result = DiffResults.value!(result, y)
     grad = DiffResults.gradient(result)
     fill!(grad, zero(y))
     return result
 end
 
-function extract_gradient!(::Type{T}, result::MutableDiffResult, x, dual::Dual) where {T}
+function extract_gradient!(::Type{T}, result::MutableDiffResult, dual::Dual, x) where {T}
     result = DiffResults.value!(result, value(T, dual))
-    extract_gradient!(T, DiffResults.gradient(result), x, dual)
+    extract_gradient!(T, DiffResults.gradient(result), dual, x)
     return result
 end
 
 # Immutable results cannot be written to entry by entry. Copying the partials wholesale is correct
 # as long as every entry of `x` is seeded, which holds for the `StaticArray` gradient buffers that
 # are the only source of such results; anything else throws on the length mismatch.
-function extract_gradient!(::Type{T}, result::ImmutableDiffResult, x, dual::Dual) where {T}
+function extract_gradient!(::Type{T}, result::ImmutableDiffResult, dual::Dual, x) where {T}
     result = DiffResults.value!(result, value(T, dual))
     result = DiffResults.gradient!(result, partials(T, dual))
     return result
 end
 
-# Zeroes the entries that receive no derivative. In chunk mode the sweep calls this once up front,
-# since the entries that no chunk writes belong to none of them in particular.
-function zero_unseeded!(::Type{T}, result::AbstractArray, x, dual) where {T}
+# Zeroes `result` unless extraction is going to write every entry of it; the written ones are
+# overwritten immediately after. In chunk mode the sweep calls this once up front, since the entries
+# that no chunk writes belong to none of them in particular. Comparing counts rather than matching
+# entries up is exact for every pair that can hold the gradient: a differently structured `result`
+# with the same count, an `UpperTriangular` one for a `LowerTriangular` `x`, cannot. `dual` is passed
+# for its value type, which unlike `eltype(result)` is a number type even for an `Any` result.
+function zero_unseeded!(::Type{T}, result::AbstractArray, dual, x) where {T}
     structural_length(x) == structural_length(result) || fill!(result, zero(valtype(T, dual)))
     return result
 end
-function zero_unseeded!(::Type{T}, result::MutableDiffResult, x, dual) where {T}
-    zero_unseeded!(T, DiffResults.gradient(result), x, dual)
+# Dispatched on `DiffResult`, not on `MutableDiffResult`: a `StaticArray` gradient buffer makes the
+# result immutable even when the buffer itself can be written to, as an `MVector` can.
+function zero_unseeded!(::Type{T}, result::DiffResult, dual, x) where {T}
+    zero_unseeded!(T, DiffResults.gradient(result), dual, x)
     return result
 end
 
-extract_gradient!(::Type{T}, result::AbstractArray, x, y::Real) where {T} = fill!(result, zero(y))
-function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) where {T}
-    zero_unseeded!(T, result, x, dual)
+extract_gradient!(::Type{T}, result::AbstractArray, y::Real, x) where {T} = fill!(result, zero(y))
+function extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual, x) where {T}
+    zero_unseeded!(T, result, dual, x)
     idxs = structural_eachindex(x, result)
     for (i, idx) in zip(1:npartials(dual), idxs)
         result[idx] = partials(T, dual, i)
@@ -99,7 +104,7 @@ function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) wher
     return result
 end
 
-function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) where {T}
+function extract_gradient_chunk!(::Type{T}, result, dual, x, index, chunksize) where {T}
     offset = index - 1
     idxs = Iterators.drop(structural_eachindex(x, result), offset)
     for (i, idx) in zip(1:chunksize, idxs)
@@ -108,14 +113,11 @@ function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) w
     return result
 end
 
-function extract_gradient_chunk!(::Type{T}, result::DiffResult, x, dual, index, chunksize) where {T}
-    extract_gradient_chunk!(T, DiffResults.gradient(result), x, dual, index, chunksize)
+function extract_gradient_chunk!(::Type{T}, result::DiffResult, dual, x, index, chunksize) where {T}
+    extract_gradient_chunk!(T, DiffResults.gradient(result), dual, x, index, chunksize)
     return result
 end
 
-extract_gradient_chunk!(::Type, result, x, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR)
-extract_gradient_chunk!(::Type, result::DiffResult, x, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR)
-
 const GRAD_ERROR = DimensionMismatch("gradient(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?")
 
 ###############
@@ -126,12 +128,12 @@ function vector_mode_gradient(f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa Real || throw(GRAD_ERROR)
     result = similar(x, valtype(T, ydual))
-    return extract_gradient!(T, result, x, ydual)
+    return extract_gradient!(T, result, ydual, x)
 end
 
 function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T, F}
     ydual = vector_mode_dual_eval!(f, cfg, x)
-    result = extract_gradient!(T, result, x, ydual)
+    result = extract_gradient!(T, result, ydual, x)
     return result
 end
 
@@ -161,9 +163,10 @@ function chunk_mode_gradient_expr(result_definition::Expr)
         seed!(xdual, x, 1, seeds)
         seed_zero_partials!(xdual, x, N + 1, xlen - N)
         ydual = f(xdual)
+        ydual isa Real || throw(GRAD_ERROR)
         $(result_definition)
-        zero_unseeded!(T, result, x, ydual)
-        extract_gradient_chunk!(T, result, x, ydual, 1, N)
+        zero_unseeded!(T, result, ydual, x)
+        extract_gradient_chunk!(T, result, ydual, x, 1, N)
         seed_zero_partials!(xdual, x, 1)
 
         # do middle chunks
@@ -171,14 +174,14 @@ function chunk_mode_gradient_expr(result_definition::Expr)
             i = ((c - 1) * N + 1)
             seed!(xdual, x, i, seeds)
             ydual = f(xdual)
-            extract_gradient_chunk!(T, result, x, ydual, i, N)
+            extract_gradient_chunk!(T, result, ydual, x, i, N)
             seed_zero_partials!(xdual, x, i)
         end
 
         # do final chunk
         seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
         ydual = f(xdual)
-        extract_gradient_chunk!(T, result, x, ydual, lastchunkindex, lastchunksize)
+        extract_gradient_chunk!(T, result, ydual, x, lastchunkindex, lastchunksize)
 
         # get the value, this is a no-op unless result is a DiffResult
         extract_value!(T, result, ydual)
diff --git a/src/jacobian.jl b/src/jacobian.jl
index 752c817b..74f88635 100644
--- a/src/jacobian.jl
+++ b/src/jacobian.jl
@@ -96,29 +96,28 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is
 # to `x[j]`. Only the seeded entries of `x` have a derivative to extract, so the columns of the
 # structurally zero ones are zeroed instead. See #839.
 
-function extract_jacobian!(::Type{T}, result::AbstractArray, x, ydual::AbstractArray) where {T}
+# Zeroes the whole Jacobian unless every column is going to be written; the written ones are
+# overwritten immediately after. In chunk mode the sweep calls this once up front, since the columns
+# that no chunk writes belong to none of them in particular.
+function zero_unseeded_columns!(::Type{T}, out::AbstractArray, ydual, x) where {T}
+    structural_length(x) == length(x) || fill!(out, zero(valtype(T, eltype(ydual))))
+    return out
+end
+
+# Vector mode is a single chunk that covers every seeded entry of `x`, so it extracts like the sweep.
+function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, x::AbstractArray) where {T}
     out_reshaped = reshape_jacobian(result, ydual, x)
-    ydual_reshaped = vec(ydual)
-    # Use closure to avoid GPU broadcasting with Type
-    partials_wrap(ydual, nrange) = partials(T, ydual, nrange)
-    n = structural_length(x)
-    if n == length(x)
-        out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n))
-    else
-        fill!(out_reshaped, zero(valtype(T, eltype(ydual))))
-        for (i, col) in zip(1:n, structural_linearindices(x))
-            out_reshaped[:, col] .= partials_wrap.(ydual_reshaped, i)
-        end
-    end
+    zero_unseeded_columns!(T, out_reshaped, ydual, x)
+    extract_jacobian_chunk!(T, out_reshaped, ydual, x, 1, structural_length(x))
     return result
 end
 
-function extract_jacobian!(::Type{T}, result::MutableDiffResult, x, ydual::AbstractArray) where {T}
-    extract_jacobian!(T, DiffResults.jacobian(result), x, ydual)
+function extract_jacobian!(::Type{T}, result::MutableDiffResult, ydual::AbstractArray, x::AbstractArray) where {T}
+    extract_jacobian!(T, DiffResults.jacobian(result), ydual, x)
     return result
 end
 
-function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize) where {T}
+function extract_jacobian_chunk!(::Type{T}, result, ydual, x::AbstractArray, index, chunksize) where {T}
     ydual_reshaped = vec(ydual)
     offset = index - 1
     irange = 1:chunksize
@@ -127,19 +126,17 @@ function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize)
     if structural_length(x) == length(x)
         result[:, irange .+ offset] .= partials_wrap.(ydual_reshaped, transpose(irange))
     else
-        idxs = Iterators.drop(structural_linearindices(x), offset)
-        for (i, col) in zip(irange, idxs)
+        cols = Iterators.drop(structural_columns(result, x), offset)
+        for (i, col) in zip(irange, cols)
             result[:, col] .= partials_wrap.(ydual_reshaped, i)
         end
     end
     return result
 end
 
-# A matrix is used as is: reshaping it would allocate a wrapper on Julia >= 1.11, where `reshape`
-# can no longer return its argument. The size is checked instead, as `reshape` did on the way past.
 function reshape_jacobian(result::AbstractMatrix, ydual, x)
     size(result) == (length(ydual), length(x)) || throw(DimensionMismatch(
-        lazy"cannot store the $(length(ydual))x$(length(x)) Jacobian in a result of size $(size(result))"))
+        lazy"cannot store the $(length(ydual))×$(length(x)) Jacobian in a result of size $(size(result))"))
     return result
 end
 reshape_jacobian(result::AbstractArray, ydual, x) = reshape(result, length(ydual), length(x))
@@ -153,7 +150,7 @@ function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T}) where {F,T}
     ydual = vector_mode_dual_eval!(f, cfg, x)
     ydual isa AbstractArray || throw(JACOBIAN_ERROR)
     result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))
-    extract_jacobian!(T, result, x, ydual)
+    extract_jacobian!(T, result, ydual, x)
     extract_value!(T, result, ydual)
     return result
 end
@@ -161,14 +158,14 @@ end
 function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
     ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     result = similar(y, length(y), length(x))
-    extract_jacobian!(T, result, x, ydual)
+    extract_jacobian!(T, result, ydual, x)
     map!(d -> value(T,d), y, ydual)
     return result
 end
 
 function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T}) where {F,T}
     ydual = vector_mode_dual_eval!(f, cfg, x)
-    extract_jacobian!(T, result, x, ydual)
+    extract_jacobian!(T, result, ydual, x)
     extract_value!(T, result, ydual)
     return result
 end
@@ -176,7 +173,7 @@ end
 function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) where {F,T}
     ydual = vector_mode_dual_eval!(f!, cfg, y, x)
     map!(d -> value(T,d), y, ydual)
-    extract_jacobian!(T, result, x, ydual)
+    extract_jacobian!(T, result, ydual, x)
     extract_value!(T, result, y, ydual)
     return result
 end
@@ -212,9 +209,8 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex
         ydual isa AbstractArray || throw(JACOBIAN_ERROR)
         $(result_definition)
         out_reshaped = reshape_jacobian(result, ydual, x)
-        # zero the columns of the structurally zero entries of `x`, which no chunk of the sweep writes
-        structural_length(x) == length(x) || fill!(out_reshaped, zero(valtype(T, eltype(ydual))))
-        extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N)
+        zero_unseeded_columns!(T, out_reshaped, ydual, x)
+        extract_jacobian_chunk!(T, out_reshaped, ydual, x, 1, N)
         seed_zero_partials!(xdual, x, 1)
 
         # do middle chunks
@@ -222,14 +218,14 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex
             i = ((c - 1) * N + 1)
             seed!(xdual, x, i, seeds)
             $(compute_ydual)
-            extract_jacobian_chunk!(T, out_reshaped, x, ydual, i, N)
+            extract_jacobian_chunk!(T, out_reshaped, ydual, x, i, N)
             seed_zero_partials!(xdual, x, i)
         end
 
         # do final chunk
         seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
         $(compute_ydual)
-        extract_jacobian_chunk!(T, out_reshaped, x, ydual, lastchunkindex, lastchunksize)
+        extract_jacobian_chunk!(T, out_reshaped, ydual, x, lastchunkindex, lastchunksize)
 
         $(y_definition)
 

From 3063bff9bb20ac140b2e289e7c23f16fa07a81f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= <d.widmann@pumas.ai>
Date: Tue, 18 Aug 2026 17:23:19 +0200
Subject: [PATCH 6/6] Cover the wrongly shaped results, the partial final
 chunks and the static buffer

A wrongly shaped `jacobian!` result and a `gradient!` result that is not shaped like `x`
both throw now, the latter for a dense `x` too, and neither was asserted. Nor was the
`MVector` gradient buffer of an `ImmutableDiffResult`, which chunk mode writes entry by
entry.

The structured Jacobian and Hessian testsets only ran chunk sizes that divide the number
of seeded entries, so the final chunk was always a full one; they gain a size that leaves
a remainder. `SeedTest` pins the columns `structural_columns` returns against the index
sets it already writes out by hand, since the Jacobian misplaces its derivatives 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.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 test/AllocationsTest.jl | 30 +++++++++++++++---------------
 test/GradientTest.jl    | 24 ++++++++++++++++++++++++
 test/HessianTest.jl     |  3 ++-
 test/JacobianTest.jl    | 21 ++++++++++++++++++++-
 test/QATest.jl          |  9 +++++++++
 test/SeedTest.jl        |  4 ++++
 6 files changed, 74 insertions(+), 17 deletions(-)

diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl
index 53eb5854..4d3f4d92 100644
--- a/test/AllocationsTest.jl
+++ b/test/AllocationsTest.jl
@@ -53,9 +53,9 @@ end
 
 # `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural
 # position to an index of `x` must not allocate, whether or not `x` has structurally zero entries.
-function allocs_structured_gradient!(x, chunk)
+function allocs_structured_gradient!(result, x, chunk)
     f(z) = sum(abs2, z)
-    result = fill!(similar(x, Float64), false)
+    fill!(result, false)
     cfg = ForwardDiff.GradientConfig(f, x, chunk)
     ForwardDiff.gradient!(result, f, x, cfg)  # warmup
     return @allocated ForwardDiff.gradient!(result, f, x, cfg)
@@ -70,19 +70,19 @@ function allocs_structured_jacobian!(x, chunk)
     return @allocated ForwardDiff.jacobian!(result, f!, y, x, cfg)
 end
 
-@testset "Test gradient!/jacobian! allocations for $(nameof(typeof(x)))" for x in (rand(6, 6),
-                                                                                  LowerTriangular(rand(6, 6)),
-                                                                                  UpperTriangular(rand(6, 6)),
-                                                                                  Diagonal(rand(6, 6)))
-    # vector mode
-    chunk = ForwardDiff.Chunk{ForwardDiff.structural_length(x)}()
-    @test iszero(allocs_structured_gradient!(x, chunk))
-    @test iszero(allocs_structured_jacobian!(x, chunk))
-
-    # chunk mode
-    chunk = ForwardDiff.Chunk{2}()
-    @test iszero(allocs_structured_gradient!(x, chunk))
-    @test iszero(allocs_structured_jacobian!(x, chunk))
+@testset "Test gradient!/jacobian! allocations for $(nameof(typeof(x)))" for (x, nstruct) in (
+        (rand(6, 6),                  36),
+        (LowerTriangular(rand(6, 6)), 21),
+        (UpperTriangular(rand(6, 6)), 21),
+        (Diagonal(rand(6, 6)),         6),
+    )
+    # A result shaped like `x` receives a derivative in every entry it stores, a dense one has the
+    # entries off the structure of `x` zeroed as well. The chunk sizes cover chunk and vector mode.
+    for result in (similar(x), zeros(size(x))), chunk_size in (2, nstruct)
+        chunk = ForwardDiff.Chunk{chunk_size}()
+        @test iszero(allocs_structured_gradient!(result, x, chunk))
+        @test iszero(allocs_structured_jacobian!(x, chunk))
+    end
 end
 
 @testset "allocation-free nested StaticArray jacobian" begin
diff --git a/test/GradientTest.jl b/test/GradientTest.jl
index 8f618fe5..f80468da 100644
--- a/test/GradientTest.jl
+++ b/test/GradientTest.jl
@@ -331,6 +331,30 @@ end
     end
 end
 
+@testset "result not shaped like x" begin
+    # The extraction positions are indices of `x`, which a result of a different shape cannot be
+    # indexed by, dense `x` included.
+    x = randn(4)
+    f = z -> dot(z, z)
+    @testset "chunk size = $c" for c in (2, 4)
+        cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}())
+        @test_throws DimensionMismatch ForwardDiff.gradient!(fill(NaN, 5), f, x, cfg)
+        result = DiffResults.DiffResult(NaN, fill(NaN, 5))
+        @test_throws DimensionMismatch ForwardDiff.gradient!(result, f, x, cfg)
+    end
+end
+
+@testset "mutable gradient buffer in an immutable result" begin
+    # A `StaticArray` buffer makes the result immutable, but an `MVector` can still be written to
+    # entry by entry, which is how the chunk mode sweep fills it.
+    x = randn(6)
+    f = z -> dot(z, z)
+    result = DiffResults.DiffResult(NaN, @MVector fill(NaN, 6))
+    @test result isa DiffResults.ImmutableDiffResult
+    ForwardDiff.gradient!(result, f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{2}()))
+    @test DiffResults.gradient(result) ≈ 2 .* x
+end
+
 # issue #769
 @testset "functions with `Dual` output" begin
     x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]
diff --git a/test/HessianTest.jl b/test/HessianTest.jl
index 01d29143..28fd03ec 100644
--- a/test/HessianTest.jl
+++ b/test/HessianTest.jl
@@ -177,7 +177,8 @@ end
     grad = zeros(3, 3)
     grad[sidx] .= x[sidx] .+ sum(x)
 
-    @testset "chunk size = $c" for c in unique((1, 2, length(sidx)))
+    # one chunk size below the full length, so that the final chunk is a partial one
+    @testset "chunk size = $c" for c in unique((1, 2, length(sidx) - 1, length(sidx)))
         cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}())
 
         H = ForwardDiff.hessian(f, x, cfg)
diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl
index 0117d844..5d2b901a 100644
--- a/test/JacobianTest.jl
+++ b/test/JacobianTest.jl
@@ -340,7 +340,8 @@ end
     expected[2, sidx] .= 2 .* x[sidx]
     val = g(x)
 
-    @testset "chunk size = $c" for c in unique((1, 2, length(sidx)))
+    # `length(sidx)` is 10 or 4, so a chunk size of 3 leaves a partial final chunk
+    @testset "chunk size = $c" for c in unique((1, 2, 3, length(sidx)))
         cfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{c}())
         J = ForwardDiff.jacobian(g, x, cfg)
         @test size(J) == (2, length(x))
@@ -379,6 +380,24 @@ end
     end
 end
 
+@testset "wrongly shaped result" begin
+    # A matrix result is used as is, so it has to have the shape of the Jacobian and not merely as
+    # many entries. Results of other shapes are reshaped and only have to match in length.
+    x = randn(4)
+    g = z -> [sum(z), sum(abs2, z)]
+    g! = (y, z) -> (y[1] = sum(z); y[2] = sum(abs2, z); y)
+    @testset "chunk size = $c" for c in (2, 4)
+        cfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{c}())
+        @test_throws DimensionMismatch ForwardDiff.jacobian!(fill(NaN, 4, 2), g, x, cfg)
+        result = DiffResults.DiffResult(fill(NaN, 2), fill(NaN, 4, 2))
+        @test_throws DimensionMismatch ForwardDiff.jacobian!(result, g, x, cfg)
+
+        y = fill(NaN, 2)
+        cfg! = ForwardDiff.JacobianConfig(g!, y, x, ForwardDiff.Chunk{c}())
+        @test_throws DimensionMismatch ForwardDiff.jacobian!(fill(NaN, 4, 2), g!, y, x, cfg!)
+    end
+end
+
 # issue #769
 @testset "functions with `Dual` output" begin
     x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))]
diff --git a/test/QATest.jl b/test/QATest.jl
index 860ccdb0..925bbfb8 100644
--- a/test/QATest.jl
+++ b/test/QATest.jl
@@ -1,6 +1,7 @@
 module QATest
 
 using ForwardDiff
+using LinearAlgebra
 using Test
 
 using JET: @test_opt
@@ -11,6 +12,14 @@ using JET: @test_opt
     @test_opt ForwardDiff.gradient(only, [1.0], ForwardDiff.GradientConfig(only, [1.0], ForwardDiff.Chunk{1}()))
     @test_opt ForwardDiff.jacobian(identity, [1.0], ForwardDiff.JacobianConfig(identity, [1.0], ForwardDiff.Chunk{1}()))
     @test_opt ForwardDiff.hessian(only, [1.0], ForwardDiff.HessianConfig(only, [1.0], ForwardDiff.Chunk{1}()))
+
+    # extraction iterates the structural positions of `x` for these
+    @testset "$(nameof(typeof(x)))" for x in (LowerTriangular(rand(3, 3)),
+                                              UpperTriangular(rand(3, 3)),
+                                              Diagonal(rand(3, 3)))
+        @test_opt ForwardDiff.gradient(first, x, ForwardDiff.GradientConfig(first, x, ForwardDiff.Chunk{2}()))
+        @test_opt ForwardDiff.jacobian(vec, x, ForwardDiff.JacobianConfig(vec, x, ForwardDiff.Chunk{2}()))
+    end
 end
 
 end # module
diff --git a/test/SeedTest.jl b/test/SeedTest.jl
index 02b821c3..a6b3944a 100644
--- a/test/SeedTest.jl
+++ b/test/SeedTest.jl
@@ -53,6 +53,10 @@ end
     @test collect(ForwardDiff.structural_eachindex(duals, x)) == sidx
     @test ForwardDiff.structural_length(x) == nstruct
 
+    # the columns a Jacobian receives derivatives in, in the order the seeds are laid out in
+    @test collect(ForwardDiff.structural_columns(zeros(2, length(x)), x)) == LinearIndices(x)[sidx]
+    @test_throws DimensionMismatch ForwardDiff.structural_columns(zeros(2, length(x) + 1), x)
+
     # `count` defaults to N
     fill_marker!(duals, x, sidx, marker)
     ForwardDiff.seed_zero_partials!(duals, x, 4)