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

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

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

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.

src/gradient.jl

@@ -49,44 +49,72 @@ 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 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)
     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. 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
 
-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)
+# 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}
+    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)
     end
     return result
 end
 
-function extract_gradient_chunk!(::Type{T}, result, dual, index, chunksize) where {T}
+function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) where {T}
     offset = index - 1
-    idxs = Iterators.drop(structural_eachindex(result), offset)
+    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 +126,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,22 +162,23 @@ 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)
+        zero_unseeded!(T, result, x, ydual)
+        extract_gradient_chunk!(T, result, x, ydual, 1, N)
         seed_zero_partials!(xdual, x, 1)
 
         # do middle chunks
         for c in middlechunks
             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)

src/jacobian.jl

@@ -92,71 +92,91 @@ 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 = 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)
-    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
+        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
 
-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 #
 ###############
 
 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
@@ -191,23 +211,25 @@ 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)
-        extract_jacobian_chunk!(T, out_reshaped, ydual, 1, N)
+        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)
 
         # do middle chunks
         for c in middlechunks
             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 +240,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
 

test/AllocationsTest.jl

@@ -1,6 +1,7 @@
 module AllocationsTest
 
 using ForwardDiff
+using LinearAlgebra
 using StaticArrays
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
@@ -50,6 +51,40 @@ 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_structured_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_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)
+    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_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))
+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