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, 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, ydual, length(x))
+    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, ydual, length(x))
+    result = extract_jacobian!(T, result, ydual, x)
     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

@@ -48,28 +48,62 @@ 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 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 (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
 # 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,45 +49,75 @@ 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 `x` has as many seeded entries as `result` has entries. See #838.
+
+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::DiffResult, 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), 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, dual::Dual, x) 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 `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
+# 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, 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)
     end
     return result
 end
 
-function extract_gradient_chunk!(::Type{T}, result, 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(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, 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, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR)
-extract_gradient_chunk!(::Type, result::DiffResult, 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 +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, 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, ydual)
+    result = extract_gradient!(T, result, ydual, x)
     return result
 end
 
@@ -133,23 +163,25 @@ 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)
-        extract_gradient_chunk!(T, result, 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
         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, 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, 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)

src/jacobian.jl

@@ -92,71 +92,88 @@ 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)
-    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))
+# 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.
+
+# 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)
+    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, ydual::AbstractArray, n) where {T}
-    extract_jacobian!(T, DiffResults.jacobian(result), ydual, n)
+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, 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
-    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
+        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
 
-reshape_jacobian(result, ydual, xdual) = reshape(result, length(ydual), length(xdual))
-reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResults.jacobian(result), ydual, xdual)
+function reshape_jacobian(result::AbstractMatrix, ydual, x)
+    size(result) == (length(ydual), length(x)) || throw(DimensionMismatch(
+        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))
+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, ydual, x)
     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, 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}
-    N = chunksize(cfg)
     ydual = vector_mode_dual_eval!(f, cfg, x)
-    extract_jacobian!(T, result, ydual, N)
+    extract_jacobian!(T, result, ydual, x)
     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, ydual, x)
     extract_value!(T, result, y, ydual)
     return result
 end
@@ -191,23 +208,24 @@ 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_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
         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, 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, ydual, lastchunkindex, lastchunksize)
+        extract_jacobian_chunk!(T, out_reshaped, ydual, x, lastchunkindex, lastchunksize)
 
         $(y_definition)
 
@@ -218,14 +236,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