Apply the second-round symmetric-Hessian review: output shape, shape check, chunk-size independence

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
@@ -112,6 +112,8 @@ end
     return typeof(H)(Symmetric(H, :U))
 end
 
+# `partials` is empty both for an empty `x` and for an `f` that ignores its argument, and neither
+# reaches `extract_jacobian`, whose generated method would build a matrix with no columns to fill.
 @inline function extract_hessian(::Type{T}, ydual::Partials{0}, x::S) where {T,S<:StaticArray}
     R = StaticArrays.similar_type(S, valtype(T, eltype(ydual)), Size(length(x), length(x)))
     return zero(R)
@@ -131,8 +133,10 @@ ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig, ::Val) where {F} =
     T = typeof(Tag(f, eltype(x)))
     ydual = f(dualize(T, dualize(T, x)))
     ydual isa Real || throw(HESSIAN_ERROR)
-    H = result isa AbstractMatrix ? result : reshape(result, length(x), length(x))
-    ForwardDiff.extract_hessian_chunk!(T, H, ydual, 0, 0, length(x), length(x))
+    H = ForwardDiff.reshape_hessian(result, x)
+    # a `StaticArray` has no structurally zero entries, so the positions are the columns of `H`
+    positions = ForwardDiff.structural_columns(H, x)
+    ForwardDiff.extract_hessian_chunk!(T, H, positions, ydual, 0, 0, length(x), length(x))
     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.
@@ -129,16 +163,20 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
     end
 end
 
-# Seed a chunk in either layer of nested duals. A `nothing` seed clears that layer.
+# Seed a chunk in either layer of nested duals. A `nothing` seed clears that layer. This is not
+# `seed_zero_partials!` with both layers cleared: that builds the inner value at the exact type of
+# the array's `V`, which the nested buffer's `Dual{T,V,N}` is not, so it would be a `MethodError`.
 function seed_hessian_chunk!(duals::AbstractArray{Dual{T,Dual{T,V,N},N}}, x, index,
                              iseeds::Union{Nothing,NTuple{N,Partials{N,V}}},
                              oseeds::Union{Nothing,NTuple{N,Partials{N,Dual{T,V,N}}}},
                              chunksize = N) where {T,V,N}
-    izero = zero(Partials{N,V})
-    ozero = zero(Partials{N,Dual{T,V,N}})
+    # `iseeds === nothing` is a compile-time constant here, so building each layer's accessor in its
+    # own branch means the zero partials it does not need are never constructed. That is free for an
+    # isbits `V` but not, say, for `BigFloat`.
+    ipartials = iseeds === nothing ? Returns(zero(Partials{N,V})) : Base.Fix1(getindex, iseeds)
+    opartials = oseeds === nothing ? Returns(zero(Partials{N,Dual{T,V,N}})) : Base.Fix1(getindex, oseeds)
     idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), chunksize)
     return _seed!(duals, x, idxs) do value, i
-        inner = Dual{T,V,N}(value, iseeds === nothing ? izero : iseeds[i])
-        Dual{T,Dual{T,V,N},N}(inner, oseeds === nothing ? ozero : oseeds[i])
+        Dual{T,Dual{T,V,N},N}(Dual{T,V,N}(value, ipartials(i)), opartials(i))
     end
 end

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)