diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index bf0ef99a..010af29b 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -5,7 +5,7 @@ using ForwardDiff.LinearAlgebra using ForwardDiff.DiffResults using ForwardDiff: Dual, partials, npartials, Partials, GradientConfig, JacobianConfig, HessianConfig, Tag, Chunk, gradient, hessian, jacobian, gradient!, hessian!, jacobian!, - extract_gradient!, extract_jacobian!, extract_value!, + extract_gradient!, extract_jacobian!, extract_value!, structural_indices, vector_mode_gradient, vector_mode_gradient!, vector_mode_jacobian, vector_mode_jacobian!, valtype, value using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult @@ -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, structural_indices(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, structural_indices(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, structural_indices(x)) 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..daf9a1f6 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -20,14 +20,15 @@ end function vector_mode_dual_eval!(f::F, cfg::Union{JacobianConfig,GradientConfig}, x) where {F} xdual = cfg.duals - seed!(xdual, x, cfg.seeds) + seed!(xdual, x, cfg.indices, cfg.seeds) return f(xdual) end function vector_mode_dual_eval!(f!::F, cfg::JacobianConfig, y, x) where {F} ydual, xdual = cfg.duals - seed!(xdual, x, cfg.seeds) - seed_zero_partials!(ydual, y) + yindices, xindices = cfg.indices + seed!(xdual, x, xindices, cfg.seeds) + seed_zero_partials!(ydual, y, yindices) f!(ydual, xdual) return ydual end @@ -40,104 +41,136 @@ end return Expr(:tuple, [:(single_seed(Partials{N,V}, Val{$i}())) for i in 1:N]...) end -# Only seed indices that are structurally non-zero -structural_eachindex(x::AbstractArray) = structural_eachindex(x, x) -function structural_eachindex(x::AbstractArray, y::AbstractArray) - require_one_based_indexing(x, y) - eachindex(x, y) +######################## +# structural positions # +######################## + +# The set of linear positions an array stores, so that two arrays can be compared without comparing +# their positions themselves. `structural_indices` below enumerates the positions of each kind. +abstract type StructuralKind end +struct AllEntries <: StructuralKind end +struct LowerTriangle <: StructuralKind end +struct UpperTriangle <: StructuralKind end +struct MainDiagonal <: StructuralKind end + +structural_kind(::AbstractArray) = AllEntries() +structural_kind(::LowerTriangular) = LowerTriangle() +structural_kind(::UpperTriangular) = UpperTriangle() +structural_kind(::Diagonal) = MainDiagonal() + +# The linear indices of the entries that are seeded, in seeding order. Being linear indices of the +# array itself, one position indexes the input, the work buffer and the result alike, and is a +# Jacobian column number as it stands. Configs hold one per work buffer, so that a sweep indexes +# straight to a chunk. +function structural_indices(x::AbstractArray) + require_one_based_indexing(x) + return Base.OneTo(length(x)) end -function structural_eachindex(x::UpperTriangular, y::AbstractArray) - require_one_based_indexing(x, y) - if size(x) != size(y) - throw(DimensionMismatch()) - end +function structural_indices(x::Diagonal) + require_one_based_indexing(x) + return diagind(x) +end +function structural_indices(x::UpperTriangular) + require_one_based_indexing(x) n = size(x, 1) - return (CartesianIndex(i, j) for j in 1:n for i in 1:j) + return [i + (j - 1) * n 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 +function structural_indices(x::LowerTriangular) + require_one_based_indexing(x) n = size(x, 1) - return (CartesianIndex(i, j) for j in 1:n for i in j:n) + return [i + (j - 1) * n 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 - return diagind(x) + +# The positions of the `count` entries starting at structural position `index`. Allocation-free, and +# a window overrunning the end is a `BoundsError` rather than a silently truncated chunk. +structural_chunk(indices, index, count) = view(indices, index:(index + count - 1)) + +# Does an array of kind `outer` store every position of an array of kind `inner`? Add a method here +# whenever a kind is added above. +structural_issubset(inner::StructuralKind, outer::StructuralKind) = inner === outer +structural_issubset(::StructuralKind, ::AllEntries) = true +structural_issubset(::MainDiagonal, ::LowerTriangle) = true +structural_issubset(::MainDiagonal, ::UpperTriangle) = true + +# Checks that the structural positions of `x` are positions of `y` as well. Being linear indices, +# they only constrain how many entries `y` has, not its shape. +function check_structural_indices(x::AbstractArray, y::AbstractArray) + require_one_based_indexing(y) + structural_issubset(structural_kind(x), structural_kind(y)) || throw(ArgumentError(LazyString( + "an array of type ", nameof(typeof(y)), " does not store every entry of an array of type ", + nameof(typeof(x)), ": the two are structured differently"))) + length(x) == length(y) || throw(DimensionMismatch( + lazy"expected an array with $(length(x)) elements, got an array with $(length(y)) elements")) + 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. -seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} = - _seed_zero_partials!(duals, x, structural_eachindex(duals, x)) - -# Zeroes the partials of `count` elements starting at structural position `index`. Chunk mode only -# needs to clear the chunk it just seeded, so writing through to the end of the array would be O(n) -# redundant work per chunk, i.e. O(n^2/N) per sweep. `count` mirrors the `chunksize` argument of -# `seed!(duals, x, index, seeds, chunksize)`. -function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, index, - count = N) where {T,V,N} - idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), count) - return _seed_zero_partials!(duals, x, idxs) +# The config's positions were built for its work buffer, so they fit `x` exactly when the two are +# structurally interchangeable. The kind comparison is a compile-time constant, hence free per call. +function checkstructure(duals::AbstractArray, x::AbstractArray) + structural_kind(duals) === structural_kind(x) || throw(ArgumentError(LazyString( + "the config was built for an array of type ", nameof(typeof(duals)), + " and cannot be used with an array of type ", nameof(typeof(x)), + ": the two are structured differently"))) + return check_structural_indices(duals, x) end -function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N} - seed = zero(Partials{N,V}) +checkstructure(cfg::AbstractConfig, x) = checkstructure(cfg.duals, x) + +# The `f!(y, x)` configs hold a buffer for the output as well, and it is seeded too. +function checkstructure(cfg::AbstractConfig, y, x) + ydual, xdual = cfg.duals + checkstructure(ydual, y) + return checkstructure(xdual, x) +end + +########### +# seeding # +########### + +# Mirrors an unassigned entry of `x` into the work buffer. `Base._unsetindex!` is implemented for +# `Array` alone, a structured wrapper being a view onto a parent with no slot of its own to unset. +_unsetindex!(duals::Array, idx) = Base._unsetindex!(duals, idx) +_unsetindex!(duals::AbstractArray, idx) = throw(ArgumentError(LazyString( + "cannot differentiate at an input with an unassigned entry at index ", idx, + ": that would leave an entry of the ", nameof(typeof(duals)), + " work buffer unassigned, which is only possible for an Array"))) + +# The two seeding operations below differ only in the `Dual` they build for the `i`th position of +# their window, so they share the walk. The `isbitstype` branch keeps the common case a plain store: +# `isassigned` is a `try`/`catch` for most array types, and a bits element type is never unassigned. +@inline function _seed!(make_dual::F, duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {F,T,V,N} if isbitstype(V) - for idx in idxs - duals[idx] = Dual{T,V,N}(x[idx], seed) + for (i, idx) in enumerate(idxs) + duals[idx] = make_dual(x[idx], i) end else - for idx in idxs + for (i, idx) in enumerate(idxs) if isassigned(x, idx) - duals[idx] = Dual{T,V,N}(x[idx], seed) + duals[idx] = make_dual(x[idx], i) else - Base._unsetindex!(duals, idx) + _unsetindex!(duals, idx) end end end return duals end -function seed!(duals::AbstractArray{Dual{T,V,N}}, x, - seeds::NTuple{N,Partials{N,V}}) where {T,V,N} - if isbitstype(V) - for (i, idx) in zip(1:N, structural_eachindex(duals, x)) - duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) - end - else - for (i, idx) in zip(1:N, structural_eachindex(duals, x)) - if isassigned(x, idx) - duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) - else - Base._unsetindex!(duals, idx) - end - end +function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N} + seed = zero(Partials{N,V}) + return _seed!(duals, x, idxs) do value, _ + Dual{T,V,N}(value, seed) end - return duals end -function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index, +seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, indices, index, count = N) where {T,V,N} = + seed_zero_partials!(duals, x, structural_chunk(indices, index, count)) + +seed!(duals::AbstractArray{Dual{T,V,N}}, x, indices, + seeds::NTuple{N,Partials{N,V}}) where {T,V,N} = seed!(duals, x, indices, 1, seeds) + +function seed!(duals::AbstractArray{Dual{T,V,N}}, x, indices, index, seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N} - offset = index - 1 - idxs = Iterators.drop(structural_eachindex(duals, x), offset) - if isbitstype(V) - for (i, idx) in zip(1:chunksize, idxs) - duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) - end - else - for (i, idx) in zip(1:chunksize, idxs) - if isassigned(x, idx) - duals[idx] = Dual{T,V,N}(x[idx], seeds[i]) - else - Base._unsetindex!(duals, idx) - end - end + return _seed!(duals, x, structural_chunk(indices, index, chunksize)) do value, i + Dual{T,V,N}(value, seeds[i]) end - return duals end diff --git a/src/config.jl b/src/config.jl index 3c6c97e3..3db333f6 100644 --- a/src/config.jl +++ b/src/config.jl @@ -59,8 +59,11 @@ Base.eltype(cfg::AbstractConfig) = eltype(typeof(cfg)) # DerivativeConfig # #################### -struct DerivativeConfig{T,D} <: AbstractConfig{1} +# `indices` holds the structural positions of `duals`, so that the sweeps do not derive them per +# chunk. They are linear indices throughout. +struct DerivativeConfig{T,D,I} <: AbstractConfig{1} duals::D + indices::I end """ @@ -84,19 +87,21 @@ function DerivativeConfig(f::F, x::X, tag::T = Tag(f, X)) where {F,X<:Real,Y<:Real,T} duals = similar(y, Dual{T,Y,1}) - return DerivativeConfig{T,typeof(duals)}(duals) + indices = structural_indices(duals) + return DerivativeConfig{T,typeof(duals),typeof(indices)}(duals, indices) end checktag(::DerivativeConfig{T},f,x) where {T} = checktag(T,f,x) -Base.eltype(::Type{DerivativeConfig{T,D}}) where {T,D} = eltype(D) +Base.eltype(::Type{DerivativeConfig{T,D,I}}) where {T,D,I} = eltype(D) ################## # GradientConfig # ################## -struct GradientConfig{T,V,N,D} <: AbstractConfig{N} +struct GradientConfig{T,V,N,D,I} <: AbstractConfig{N} seeds::NTuple{N,Partials{N,V}} duals::D + indices::I end """ @@ -120,19 +125,23 @@ function GradientConfig(f::F, ::T = Tag(f, V)) where {F,V,N,T} seeds = construct_seeds(Partials{N,V}) duals = similar(x, Dual{T,V,N}) - return GradientConfig{T,V,N,typeof(duals)}(seeds, duals) + indices = structural_indices(duals) + return GradientConfig{T,V,N,typeof(duals),typeof(indices)}(seeds, duals, indices) end checktag(::GradientConfig{T},f,x) where {T} = checktag(T,f,x) -Base.eltype(::Type{GradientConfig{T,V,N,D}}) where {T,V,N,D} = Dual{T,V,N} +Base.eltype(::Type{GradientConfig{T,V,N,D,I}}) where {T,V,N,D,I} = Dual{T,V,N} ################## # JacobianConfig # ################## -struct JacobianConfig{T,V,N,D} <: AbstractConfig{N} +# `indices` mirrors `duals`: the structural positions of the one work buffer of an `f(x)`, or a +# `(y, x)` pair of position vectors for the two buffers of an `f!(y, x)`. +struct JacobianConfig{T,V,N,D,I} <: AbstractConfig{N} seeds::NTuple{N,Partials{N,V}} duals::D + indices::I end """ @@ -157,7 +166,8 @@ function JacobianConfig(f::F, ::T = Tag(f, V)) where {F,V,N,T} seeds = construct_seeds(Partials{N,V}) duals = similar(x, Dual{T,V,N}) - return JacobianConfig{T,V,N,typeof(duals)}(seeds, duals) + indices = structural_indices(duals) + return JacobianConfig{T,V,N,typeof(duals),typeof(indices)}(seeds, duals, indices) end """ @@ -185,19 +195,20 @@ function JacobianConfig(f::F, yduals = similar(y, Dual{T,Y,N}) xduals = similar(x, Dual{T,X,N}) duals = (yduals, xduals) - return JacobianConfig{T,X,N,typeof(duals)}(seeds, duals) + indices = (structural_indices(yduals), structural_indices(xduals)) + return JacobianConfig{T,X,N,typeof(duals),typeof(indices)}(seeds, duals, indices) end checktag(::JacobianConfig{T},f,x) where {T} = checktag(T,f,x) -Base.eltype(::Type{JacobianConfig{T,V,N,D}}) where {T,V,N,D} = Dual{T,V,N} +Base.eltype(::Type{JacobianConfig{T,V,N,D,I}}) where {T,V,N,D,I} = Dual{T,V,N} ################# # HessianConfig # ################# -struct HessianConfig{T,V,N,DG,DJ} <: AbstractConfig{N} - jacobian_config::JacobianConfig{T,V,N,DJ} - gradient_config::GradientConfig{T,Dual{T,V,N},N,DG} +struct HessianConfig{T,V,N,DG,DJ,IG,IJ} <: AbstractConfig{N} + jacobian_config::JacobianConfig{T,V,N,DJ,IJ} + gradient_config::GradientConfig{T,Dual{T,V,N},N,DG,IG} end """ @@ -253,5 +264,5 @@ function HessianConfig(f::F, end checktag(::HessianConfig{T},f,x) where {T} = checktag(T,f,x) -Base.eltype(::Type{HessianConfig{T,V,N,DG,DJ}}) where {T,V,N,DG,DJ} = +Base.eltype(::Type{HessianConfig{T,V,N,DG,DJ,IG,IJ}}) where {T,V,N,DG,DJ,IG,IJ} = Dual{T,Dual{T,V,N},N} diff --git a/src/derivative.jl b/src/derivative.jl index 0c8a6c05..8ea83690 100644 --- a/src/derivative.jl +++ b/src/derivative.jl @@ -26,8 +26,9 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F, T, CHK} require_one_based_indexing(y) CHK && checktag(T, f!, x) + checkstructure(cfg, y) ydual = cfg.duals - seed_zero_partials!(ydual, y) + seed_zero_partials!(ydual, y, cfg.indices) f!(ydual, Dual{T}(x, one(x))) map!(value, y, ydual) return extract_derivative(T, ydual) @@ -64,8 +65,9 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba cfg::DerivativeConfig{T} = DerivativeConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F, T, CHK} result isa DiffResult ? require_one_based_indexing(y) : require_one_based_indexing(result, y) CHK && checktag(T, f!, x) + checkstructure(cfg, y) ydual = cfg.duals - seed_zero_partials!(ydual, y) + seed_zero_partials!(ydual, y, cfg.indices) f!(ydual, Dual{T}(x, one(x))) result = extract_value!(T, result, y, ydual) result = extract_derivative!(T, result, ydual) diff --git a/src/gradient.jl b/src/gradient.jl index a5ef3dac..34a5ef84 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -16,6 +16,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function gradient(f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T, CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) + checkstructure(cfg, x) if chunksize(cfg) == structural_length(x) return vector_mode_gradient(f, x, cfg) else @@ -35,6 +36,7 @@ This method assumes that `isa(f(x), Real)`. function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK, F} result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x) CHK && checktag(T, f, x) + checkstructure(cfg, x) if chunksize(cfg) == structural_length(x) vector_mode_gradient!(result, f, x, cfg) else @@ -49,45 +51,65 @@ 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} +# Only the structurally non-zero entries of `x` are seeded, so only those have a derivative to +# extract. The positions come from `x`, not from `result`: the two may have different structure, as +# they do when `DiffResults.HessianResult` allocates a dense gradient buffer for a structured `x`. +# Entries of `result` that no position covers are zeroed. See #838. + +function extract_gradient!(::Type{T}, result::DiffResult, y::Real, x, indices) 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, indices) where {T} result = DiffResults.value!(result, value(T, dual)) - result = DiffResults.gradient!(result, partials(T, dual)) + extract_gradient!(T, DiffResults.gradient(result), dual, x, indices) 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) - for (i, idx) in zip(1:npartials(dual), idxs) - result[idx] = partials(T, dual, i) - end +# Immutable results cannot be written to entry by entry, so the partials are copied wholesale, which +# is only correct when every entry of `x` is seeded. +function extract_gradient!(::Type{T}, result::ImmutableDiffResult, dual::Dual, x, indices) where {T} + check_structural_indices(x, DiffResults.gradient(result)) + result = DiffResults.value!(result, value(T, dual)) + result = DiffResults.gradient!(result, partials(T, dual)) return result end -function extract_gradient_chunk!(::Type{T}, result, dual, index, chunksize) where {T} - offset = index - 1 - idxs = Iterators.drop(structural_eachindex(result), offset) - for (i, idx) in zip(1:chunksize, idxs) +# Zeroes `result` unless every entry is going to be written anyway. `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} + check_structural_indices(x, result) + structural_length(x) == structural_length(result) || fill!(result, zero(valtype(T, dual))) + return nothing +end +# `DiffResult`, not `MutableDiffResult`: `DiffResult(v, ::MVector)` is an `ImmutableDiffResult`. +function zero_unseeded!(::Type{T}, result::DiffResult, dual, x) where {T} + zero_unseeded!(T, DiffResults.gradient(result), dual, x) + return nothing +end + +extract_gradient!(::Type{T}, result::AbstractArray, y::Real, x, indices) where {T} = + fill!(result, zero(y)) +function extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual, x, indices) where {T} + zero_unseeded!(T, result, dual, x) + return extract_gradient_chunk!(T, result, dual, indices, 1, npartials(dual)) +end + +function extract_gradient_chunk!(::Type{T}, result, dual, indices, index, chunksize) where {T} + for (i, idx) in enumerate(structural_chunk(indices, index, chunksize)) 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, indices, index, chunksize) where {T} + extract_gradient_chunk!(T, DiffResults.gradient(result), dual, indices, 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 +120,13 @@ 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, cfg.indices) 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) + ydual isa Real || throw(GRAD_ERROR) + result = extract_gradient!(T, result, ydual, x, cfg.indices) return result end @@ -127,29 +150,32 @@ function chunk_mode_gradient_expr(result_definition::Expr) # seed work vectors xdual = cfg.duals seeds = cfg.seeds + indices = cfg.indices # do first chunk manually to calculate output type. Seeding the first chunk and zeroing the # remaining elements partitions `xdual`, so every element is initialized exactly once. - seed!(xdual, x, 1, seeds) - seed_zero_partials!(xdual, x, N + 1, xlen - N) + seed!(xdual, x, indices, 1, seeds) + seed_zero_partials!(xdual, x, indices, N + 1, xlen - N) ydual = f(xdual) + ydual isa Real || throw(GRAD_ERROR) $(result_definition) - extract_gradient_chunk!(T, result, ydual, 1, N) - seed_zero_partials!(xdual, x, 1) + zero_unseeded!(T, result, ydual, x) + extract_gradient_chunk!(T, result, ydual, indices, 1, N) + seed_zero_partials!(xdual, x, indices, 1) # do middle chunks for c in middlechunks i = ((c - 1) * N + 1) - seed!(xdual, x, i, seeds) + seed!(xdual, x, indices, i, seeds) ydual = f(xdual) - extract_gradient_chunk!(T, result, ydual, i, N) - seed_zero_partials!(xdual, x, i) + extract_gradient_chunk!(T, result, ydual, indices, i, N) + seed_zero_partials!(xdual, x, indices, i) end # do final chunk - seed!(xdual, x, lastchunkindex, seeds, lastchunksize) + seed!(xdual, x, indices, lastchunkindex, seeds, lastchunksize) ydual = f(xdual) - extract_gradient_chunk!(T, result, ydual, lastchunkindex, lastchunksize) + extract_gradient_chunk!(T, result, ydual, indices, 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..51852f69 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -18,6 +18,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function jacobian(f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) + checkstructure(cfg, x) if chunksize(cfg) == structural_length(x) return vector_mode_jacobian(f, x, cfg) else @@ -36,6 +37,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function jacobian(f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T, CHK} require_one_based_indexing(y, x) CHK && checktag(T, f!, x) + checkstructure(cfg, y, x) if chunksize(cfg) == structural_length(x) return vector_mode_jacobian(f!, y, x, cfg) else @@ -57,6 +59,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T, CHK} result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x) CHK && checktag(T, f, x) + checkstructure(cfg, x) if chunksize(cfg) == structural_length(x) vector_mode_jacobian!(result, f, x, cfg) else @@ -78,6 +81,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} result isa DiffResult ? require_one_based_indexing(y, x) : require_one_based_indexing(result, y, x) CHK && checktag(T, f!, x) + checkstructure(cfg, y, x) if chunksize(cfg) == structural_length(x) vector_mode_jacobian!(result, f!, y, x, cfg) else @@ -92,71 +96,85 @@ 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. + +# The columns that no chunk writes belong to none of them in particular, so the sweep zeroes once up +# front rather than each chunk zeroing its own. +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 nothing +end + +# Vector mode is a single chunk that covers every seeded entry of `x`. +function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, x::AbstractArray, + indices) where {T} + out_reshaped = reshape_jacobian(result, ydual, x) + zero_unseeded_columns!(T, out_reshaped, ydual, x) + extract_jacobian_chunk!(T, out_reshaped, ydual, indices, 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, indices) where {T} + extract_jacobian!(T, DiffResults.jacobian(result), ydual, x, indices) return result end -function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) where {T} +# The positions are linear indices of `x`, which is what a Jacobian column is numbered by, so they +# select the columns directly -- for a dense `x` `structural_chunk` even hands back a `UnitRange`. +function extract_jacobian_chunk!(::Type{T}, result, ydual, indices, index, chunksize) where {T} ydual_reshaped = vec(ydual) - offset = index - 1 irange = 1:chunksize - col = irange .+ offset + cols = structural_chunk(indices, index, chunksize) # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - result[:, col] .= partials_wrap.(ydual_reshaped, transpose(irange)) + result[:, cols] .= partials_wrap.(ydual_reshaped, transpose(irange)) 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, cfg.indices) 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, cfg.indices[2]) 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) + ydual isa AbstractArray || throw(JACOBIAN_ERROR) + extract_jacobian!(T, result, ydual, x, cfg.indices) 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, cfg.indices[2]) extract_value!(T, result, y, ydual) return result end @@ -186,28 +204,29 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex # do first chunk manually to calculate output type. Seeding the first chunk and zeroing the # remaining elements partitions `xdual`, so every element is initialized exactly once. - seed!(xdual, x, 1, seeds) - seed_zero_partials!(xdual, x, N + 1, xlen - N) + seed!(xdual, x, indices, 1, seeds) + seed_zero_partials!(xdual, x, indices, N + 1, xlen - N) $(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) - seed_zero_partials!(xdual, x, 1) + out_reshaped = reshape_jacobian(result, ydual, x) + zero_unseeded_columns!(T, out_reshaped, ydual, x) + extract_jacobian_chunk!(T, out_reshaped, ydual, indices, 1, N) + seed_zero_partials!(xdual, x, indices, 1) # do middle chunks for c in middlechunks i = ((c - 1) * N + 1) - seed!(xdual, x, i, seeds) + seed!(xdual, x, indices, i, seeds) $(compute_ydual) - extract_jacobian_chunk!(T, out_reshaped, ydual, i, N) - seed_zero_partials!(xdual, x, i) + extract_jacobian_chunk!(T, out_reshaped, ydual, indices, i, N) + seed_zero_partials!(xdual, x, indices, i) end # do final chunk - seed!(xdual, x, lastchunkindex, seeds, lastchunksize) + seed!(xdual, x, indices, lastchunkindex, seeds, lastchunksize) $(compute_ydual) - extract_jacobian_chunk!(T, out_reshaped, ydual, lastchunkindex, lastchunksize) + extract_jacobian_chunk!(T, out_reshaped, ydual, indices, lastchunkindex, lastchunksize) $(y_definition) @@ -216,29 +235,29 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex 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), + $(jacobian_chunk_mode_expr(:((xdual = cfg.duals; indices = cfg.indices)), :(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)), + $(jacobian_chunk_mode_expr(:(((ydual, xdual) = cfg.duals; (yindices, indices) = cfg.indices)), + :(f!(seed_zero_partials!(ydual, y, yindices), xdual)), + :(result = similar(y, length(y), length(x))), :(map!(d -> value(T,d), y, ydual)))) end @eval function chunk_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} - $(jacobian_chunk_mode_expr(:(xdual = cfg.duals), + $(jacobian_chunk_mode_expr(:((xdual = cfg.duals; indices = cfg.indices)), :(ydual = f(xdual)), :(), :(extract_value!(T, result, ydual)))) end @eval function chunk_mode_jacobian!(result, 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)), + $(jacobian_chunk_mode_expr(:(((ydual, xdual) = cfg.duals; (yindices, indices) = cfg.indices)), + :(f!(seed_zero_partials!(ydual, y, yindices), xdual)), :(), :(extract_value!(T, result, y, ydual)))) end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 94e7cddd..5e21ef62 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")) @@ -12,22 +13,21 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F cfg = ForwardDiff.GradientConfig(nothing, x) duals = cfg.duals seeds = cfg.seeds + indices = cfg.indices allocs_seed!(args...) = @allocated ForwardDiff.seed!(args...) - allocs_seed!(duals, x, seeds) - @test iszero(allocs_seed!(duals, x, seeds)) - allocs_seed!(duals, x, 1, seeds) - @test iszero(allocs_seed!(duals, x, 1, seeds)) + allocs_seed!(duals, x, indices, seeds) + @test iszero(allocs_seed!(duals, x, indices, seeds)) + allocs_seed!(duals, x, indices, 1, seeds) + @test iszero(allocs_seed!(duals, x, indices, 1, seeds)) - # the 4-arg form passes `count` as a runtime value, so it catches an inference regression at the - # `_seed_zero_partials!` boundary that the forms defaulting `count` to `N` could hide allocs_szp!(args...) = @allocated ForwardDiff.seed_zero_partials!(args...) - allocs_szp!(duals, x) - @test iszero(allocs_szp!(duals, x)) - allocs_szp!(duals, x, 1) - @test iszero(allocs_szp!(duals, x, 1)) - allocs_szp!(duals, x, 1, 4) - @test iszero(allocs_szp!(duals, x, 1, 4)) + allocs_szp!(duals, x, indices) + @test iszero(allocs_szp!(duals, x, indices)) + allocs_szp!(duals, x, indices, 1) + @test iszero(allocs_szp!(duals, x, indices, 1)) + allocs_szp!(duals, x, indices, 1, 4) + @test iszero(allocs_szp!(duals, x, indices, 1, 4)) allocs_convert_test_574() = @allocated convert_test_574() allocs_convert_test_574() @@ -50,6 +50,66 @@ end @test iszero(allocs_jacobian!()) end +# Seeding and extraction take their positions from the config, so the `structural_chunk` view of a +# chunk must not allocate, whether or not `x` has structurally zero entries. +function allocs_structured_gradient!(result, x, chunk) + f(z) = sum(abs2, z) + fill!(result, 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, 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 + +# The `StaticArray` extension has no config to cache the structural positions in, so it derives them +# per extraction call. +function allocs_static_gradient!(result, x) + f(z) = sum(abs2, z) + fill!(result, false) + ForwardDiff.gradient!(result, f, x) # warmup + return @allocated ForwardDiff.gradient!(result, f, x) +end + +function allocs_static_jacobian!(result, x) + f(z) = z .* z + fill!(result, false) + ForwardDiff.jacobian!(result, f, x) # warmup + return @allocated ForwardDiff.jacobian!(result, f, x) +end + +@testset "Test StaticArray gradient!/jacobian! allocations for size $(size(x))" for x in ( + @SVector(rand(4)), @SMatrix(rand(2, 2)), + ) + # extraction writes the result through a view, which is no static array for a mutable buffer + @test iszero(allocs_static_gradient!(zeros(size(x)), x)) + @test iszero(allocs_static_gradient!(similar(x), x)) + @test iszero(allocs_static_jacobian!(zeros(4, 4), x)) + @test iszero(allocs_static_jacobian!(@MMatrix(zeros(4, 4)), x)) +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/DerivativeTest.jl b/test/DerivativeTest.jl index d66a7cc7..b06ed73b 100644 --- a/test/DerivativeTest.jl +++ b/test/DerivativeTest.jl @@ -55,6 +55,8 @@ end cfg = ForwardDiff.DerivativeConfig(f!, y, x) d = ForwardDiff.derivative(f, x) + @test eltype(cfg) === ForwardDiff.Dual{ForwardDiff.Tag{typeof(f!),typeof(x)},eltype(y),1} + fill!(y, 0.0) @test isapprox(ForwardDiff.derivative(f!, y, x), d) @test isapprox(v, y) @@ -106,6 +108,36 @@ end @test_throws DomainError ForwardDiff.derivative(x -> x^0.5, -1.0) end +# issue #842 +@testset "structured output buffer" begin + # A `DerivativeConfig` seeds the positions of its own output buffer, so it fits an output of the + # same structure and no other. + A = Float64[1 4 7; 2 5 8; 3 6 9] + outputs = (A, LinearAlgebra.LowerTriangular(A), LinearAlgebra.UpperTriangular(A), + LinearAlgebra.Diagonal(LinearAlgebra.diag(A))) + x = 2.0 + + @testset "$(nameof(typeof(Y)))" for Y in outputs + f! = (out, t) -> (out .= t .* Y) # writes exactly the entries `out` stores + cfg = ForwardDiff.DerivativeConfig(f!, Y, x) + + y = zero(Y) + @test ForwardDiff.derivative(f!, y, x, cfg) == Y + @test y == x .* Y + @test ForwardDiff.derivative!(zero(Y), f!, zero(Y), x, cfg) == Y + + # each output has a different structure, so `Z === Y` is the one the config fits + for Z in outputs + Z === Y && continue + msg = "ArgumentError: the config was built for an array of type " * + "$(nameof(typeof(Y))) and cannot be used with an array of type " * + "$(nameof(typeof(Z)))" + @test_throws msg ForwardDiff.derivative(f!, zero(Z), x, cfg) + @test_throws msg ForwardDiff.derivative!(zero(Z), f!, zero(Z), x, cfg) + end + end +end + @testset "dimension error for derivative" begin @test_throws DimensionMismatch ForwardDiff.derivative(sum, fill(2pi, 3)) end diff --git a/test/GradientTest.jl b/test/GradientTest.jl index bf121239..9d7d75c2 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -177,6 +177,7 @@ end @test_throws DimensionMismatch ForwardDiff.gradient(identity, 2pi) # input @test_throws DimensionMismatch ForwardDiff.gradient(identity, fill(2pi, 2)) # vector_mode_gradient @test_throws DimensionMismatch ForwardDiff.gradient(identity, fill(2pi, 10^6)) # chunk_mode_gradient + @test_throws DimensionMismatch ForwardDiff.gradient!(fill(NaN, 2), identity, fill(2pi, 2)) # vector_mode_gradient! end # Issue 548 @@ -275,6 +276,158 @@ end end end +# issue #838 +@testset "structured inputs: extraction positions" begin + # The seeds are laid out along `structural_indices(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}()) + + # 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 + + 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) + + # 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: an array of type $(nameof(T)) does not store every entry of an " * + "array of type Array" ForwardDiff.gradient!( + T(fill(NaN, n, n)), f, dense_x, dense_cfg) + + # a result whose structure *contains* that of `x` is fine, and gets hard zeros off it + if T === Diagonal + for S in (LowerTriangular, UpperTriangular) + out = S(fill(NaN, n, n)) + ForwardDiff.gradient!(out, f, x, cfg) + @test out == S(dense_expected) + end + end + end + end +end + +# issue #842 +@testset "config reused for a differently structured input" begin + # A config seeds the positions of its own work buffer, so an input of a different structure would + # be seeded at one set of positions and extracted at another. + n = 4 + A = randn(n, n) + inputs = (A, LowerTriangular(A), UpperTriangular(A), Diagonal(diag(A))) + f = z -> sum(abs2, z) / 2 # ∇f(z) == z + + # `n` is vector mode for the `Diagonal` input and chunk mode for the others + @testset "chunk size = $c" for c in (2, n) + @testset "$(nameof(typeof(xcfg))) config" for xcfg in inputs + cfg = ForwardDiff.GradientConfig(f, xcfg, ForwardDiff.Chunk{c}()) + # each input has a different structure, so `x === xcfg` is the one the config fits + for x in inputs + if x === xcfg + @test ForwardDiff.gradient(f, x, cfg) == x + else + msg = "ArgumentError: the config was built for an array of type " * + "$(nameof(typeof(xcfg))) and cannot be used with an array of type " * + "$(nameof(typeof(x)))" + @test_throws msg ForwardDiff.gradient(f, x, cfg) + @test_throws msg ForwardDiff.gradient!(similar(x), f, x, cfg) + end + end + end + + # an input of the wrong length keeps failing as a `DimensionMismatch` + cfg = ForwardDiff.GradientConfig(f, A, ForwardDiff.Chunk{c}()) + @test_throws "DimensionMismatch: expected an array with $(n^2) elements, got an array " * + "with $(n * (n + 1)) elements" ForwardDiff.gradient(f, randn(n + 1, n), cfg) + end +end + +# issue #842 +@testset "unassigned input entry" begin + # `Base._unsetindex!` takes a linear index into an `Array`, so mirroring the hole into the work + # buffer works precisely when that buffer is one. `similar` makes it one for all four inputs here. + M = Matrix{BigFloat}(undef, 3, 3) + for i in eachindex(M) + i == 5 || (M[i] = BigFloat(i)) + end + used = [i for i in eachindex(M) if i != 5] + f = z -> sum(i -> z[i]^2, used) / 2 # never reads the hole + + @testset "$(nameof(typeof(x))), chunk size = $c" for + x in (M, adjoint(M), PermutedDimsArray(M, (2, 1)), view(M, :, :)), c in (2, 9) + grad = ForwardDiff.gradient(f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}())) + @test grad[used] == x[used] + @test iszero(grad[5]) + end + + # a buffer `similar` gives the structure of `x` cannot hold the hole + @testset "$(nameof(typeof(x))), chunk size = $c" for + x in (LowerTriangular(M), UpperTriangular(M)), c in (2, 6) + @test_throws "ArgumentError: cannot differentiate at an input with an unassigned entry " * + "at index 5" ForwardDiff.gradient( + f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}())) + 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 8be72ee5..9c72faf3 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -156,6 +156,72 @@ 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) + + # 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) + @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 + +# issue #842, which `hessian` inherits through both of its sub-configs +@testset "config reused for a differently structured input" begin + n = 3 + A = randn(n, n) + inputs = (A, LowerTriangular(A), UpperTriangular(A), Diagonal(diag(A))) + f = z -> (sum(abs2, z) + sum(z)^2) / 2 + + @testset "chunk size = $c" for c in (2, n) + @testset "$(nameof(typeof(xcfg))) config" for xcfg in inputs + cfg = ForwardDiff.HessianConfig(f, xcfg, ForwardDiff.Chunk{c}()) + # each input has a different structure, so `x === xcfg` is the one the config fits + for x in inputs + x === xcfg && continue + msg = "ArgumentError: the config was built for an array of type " * + "$(nameof(typeof(xcfg))) and cannot be used with an array of type " * + "$(nameof(typeof(x)))" + @test_throws msg ForwardDiff.hessian(f, x, cfg) + @test_throws msg ForwardDiff.hessian!(fill(NaN, n^2, n^2), f, x, cfg) + end + end + 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 b6d36180..7cee2b21 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -186,6 +186,10 @@ for T in (StaticArrays.SArray, StaticArrays.MArray) @test ForwardDiff.jacobian(_diff, sx, scfg, Val{false}()) == actual @test ForwardDiff.jacobian(_diff, sx, scfg, Val{false}()) isa StaticArray + _densediff(A) = collect(diff(A; dims=1)) + @test ForwardDiff.jacobian(_densediff, sx) == actual + @test ForwardDiff.jacobian(_densediff, sx) isa Matrix + out = similar(x, 6, 9) ForwardDiff.jacobian!(out, _diff, sx) @test out == actual @@ -240,6 +244,7 @@ end @test_throws DimensionMismatch ForwardDiff.jacobian(identity, 2pi) # input @test_throws DimensionMismatch ForwardDiff.jacobian(sum, fill(2pi, 2)) # vector_mode_jacobian @test_throws DimensionMismatch ForwardDiff.jacobian(sum, fill(2pi, 10^6)) # chunk_mode_jacobian + @test_throws DimensionMismatch ForwardDiff.jacobian!(fill(NaN, 1, 2), sum, fill(2pi, 2)) # vector_mode_jacobian! end @testset "eigen" begin @@ -322,6 +327,118 @@ 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) + + # `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)) + @test J == expected + + out = fill(NaN, 2, length(x)) + @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) + @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 + 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 + +# issue #842 +@testset "config reused for a differently structured input" begin + # A config seeds the positions of its own work buffers, so an input of a different structure would + # be seeded at one set of positions and extracted at another. + n = 4 + A = randn(n, n) + inputs = (A, LowerTriangular(A), UpperTriangular(A), Diagonal(diag(A))) + 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, n) + @testset "$(nameof(typeof(xcfg))) config" for xcfg in inputs + cfg = ForwardDiff.JacobianConfig(g, xcfg, ForwardDiff.Chunk{c}()) + cfg! = ForwardDiff.JacobianConfig(g!, zeros(2), xcfg, ForwardDiff.Chunk{c}()) + # each input has a different structure, so `x === xcfg` is the one the config fits + for x in inputs + x === xcfg && continue + out = fill(NaN, 2, n^2) + msg = "ArgumentError: the config was built for an array of type " * + "$(nameof(typeof(xcfg))) and cannot be used with an array of type " * + "$(nameof(typeof(x)))" + @test_throws msg ForwardDiff.jacobian(g, x, cfg) + @test_throws msg ForwardDiff.jacobian!(out, g, x, cfg) + @test_throws msg ForwardDiff.jacobian(g!, fill(NaN, 2), x, cfg!) + @test_throws msg ForwardDiff.jacobian!(out, g!, fill(NaN, 2), x, cfg!) + end + end + + # the buffer an `f!(y, x)` config holds for the output is checked too + cfg! = ForwardDiff.JacobianConfig(nothing, Diagonal(zeros(2, 2)), A, ForwardDiff.Chunk{c}()) + @test_throws "ArgumentError: the config was built for an array of type Diagonal and " * + "cannot be used with an array of type Array" ForwardDiff.jacobian( + g!, zeros(2, 2), A, cfg!) + 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 c883f0bb..bfaec716 100644 --- a/test/QATest.jl +++ b/test/QATest.jl @@ -1,6 +1,7 @@ module QATest using ForwardDiff +using LinearAlgebra using Test import JET @@ -12,6 +13,14 @@ if !isdefined(JET, :JET_AVAILABLE) || JET.JET_AVAILABLE JET.@test_opt ForwardDiff.gradient(only, [1.0], ForwardDiff.GradientConfig(only, [1.0], ForwardDiff.Chunk{1}())) JET.@test_opt ForwardDiff.jacobian(identity, [1.0], ForwardDiff.JacobianConfig(identity, [1.0], ForwardDiff.Chunk{1}())) JET.@test_opt ForwardDiff.hessian(only, [1.0], ForwardDiff.HessianConfig(only, [1.0], ForwardDiff.Chunk{1}())) + + # seeding and extraction index the structural positions held by the config for these + @testset "$(nameof(typeof(x)))" for x in (LowerTriangular(rand(3, 3)), + UpperTriangular(rand(3, 3)), + Diagonal(rand(3, 3))) + JET.@test_opt ForwardDiff.gradient(first, x, ForwardDiff.GradientConfig(first, x, ForwardDiff.Chunk{2}())) + JET.@test_opt ForwardDiff.jacobian(vec, x, ForwardDiff.JacobianConfig(vec, x, ForwardDiff.Chunk{2}())) + end end end diff --git a/test/SeedTest.jl b/test/SeedTest.jl index 02b821c3..7c4f1792 100644 --- a/test/SeedTest.jl +++ b/test/SeedTest.jl @@ -14,14 +14,12 @@ include("utils.jl") # and check exactly which positions lost their marker. # # The expected structural index sets are written out by hand rather than obtained from -# `structural_eachindex`, so a bug in that iterator cannot hide inside the assertions depending on -# it; one test ties the two together. Order is significant: `index` and `count` are positions along -# the sequence, not array indices. The sets are heterogeneous by design — `Vector` and `Diagonal` -# enumerate linear indices (the latter via `diagind`), `UpperTriangular` enumerates `CartesianIndex` -# in column-major order. +# `structural_indices`, so a bug there cannot hide inside the assertions depending on it; one test +# ties the two together. They are linear indices of `x` in seeding order, so `index` and `count` are +# positions along the sequence, not array indices. const SEED_CASES = ( (rand(10), collect(1:10)), - (UpperTriangular(rand(5, 5)), [CartesianIndex(i, j) for j in 1:5 for i in 1:j]), + (UpperTriangular(rand(5, 5)), [i + (j - 1) * 5 for j in 1:5 for i in 1:j]), (Diagonal(rand(6, 6)), collect(1:7:36)), ) @@ -44,50 +42,94 @@ end @testset "seed_zero_partials!: $(nameof(typeof(x)))" for (x, sidx) in SEED_CASES cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{3}()) - duals, seeds = cfg.duals, cfg.seeds + duals, seeds, indices = cfg.duals, cfg.seeds, cfg.indices N = ForwardDiff.npartials(eltype(duals)) marker = Partials(ntuple(i -> Float64(i), N)) nstruct = length(sidx) - # everything below counts positions along `sidx`, so pin it to the implementation once - @test collect(ForwardDiff.structural_eachindex(duals, x)) == sidx + # everything below counts positions along `sidx`, so pin it to the implementation once. The + # config builds its positions from its own work buffer, which has the structure of `x`. + @test ForwardDiff.structural_indices(x) == sidx + @test indices == sidx @test ForwardDiff.structural_length(x) == nstruct # `count` defaults to N fill_marker!(duals, x, sidx, marker) - ForwardDiff.seed_zero_partials!(duals, x, 4) + ForwardDiff.seed_zero_partials!(duals, x, indices, 4) @test zeroed_positions(duals, sidx) == collect(4:(4 + N - 1)) @test values_match(duals, x) - # an explicit `count` narrows the window; a `count` overrunning the end is clamped by - # `Iterators.take` rather than throwing; a zero-width window is a no-op, which is what makes + # an explicit `count` narrows the window; a zero-width window is a no-op, which is what makes # `xlen - N` safe as the `count` of chunk mode's tail clear - @testset "index=$index count=$count" for (index, count, expected) in - ((4, 2, 4:5), - (nstruct - 1, N, (nstruct - 1):nstruct), - (1, 0, 1:0)) + @testset "index=$index count=$count" for (index, count, expected) in ((4, 2, 4:5), (1, 0, 1:0)) fill_marker!(duals, x, sidx, marker) - ForwardDiff.seed_zero_partials!(duals, x, index, count) + ForwardDiff.seed_zero_partials!(duals, x, indices, index, count) @test zeroed_positions(duals, sidx) == collect(expected) @test values_match(duals, x) end - # the 2-arg form clears every structural position + # a window overrunning the end is an error rather than a silently truncated chunk + @test_throws BoundsError ForwardDiff.seed_zero_partials!(duals, x, indices, nstruct - 1, N) + + # the form without a window clears every structural position fill_marker!(duals, x, sidx, marker) - ForwardDiff.seed_zero_partials!(duals, x) + ForwardDiff.seed_zero_partials!(duals, x, indices) @test zeroed_positions(duals, sidx) == collect(1:nstruct) @test values_match(duals, x) # `seed!` and `seed_zero_partials!` must agree on what "the chunk at `index`" is, or chunk mode # would leave stale seeds behind. `duals` enters each iteration fully cleared. @testset "round-trips seed! at index=$index" for index in unique((1, 4, nstruct - N + 1)) - ForwardDiff.seed!(duals, x, index, seeds) + ForwardDiff.seed!(duals, x, indices, index, seeds) @test zeroed_positions(duals, sidx) == [i for i in 1:nstruct if !(index <= i <= index + N - 1)] - ForwardDiff.seed_zero_partials!(duals, x, index) + ForwardDiff.seed_zero_partials!(duals, x, indices, index) @test zeroed_positions(duals, sidx) == collect(1:nstruct) @test values_match(duals, x) end + + # the form without an `index` seeds the first chunk, as vector mode needs + ForwardDiff.seed_zero_partials!(duals, x, indices) + ForwardDiff.seed!(duals, x, indices, seeds) + @test zeroed_positions(duals, sidx) == collect((N + 1):nstruct) + @test values_match(duals, x) +end + +# An unassigned entry of `x` is mirrored into the work buffer rather than read, which only an `Array` +# buffer can do. See #842. +@testset "unassigned entries" begin + # `x` is only read below, so every case shares these. Their first entry stays unassigned. + M = Matrix{BigFloat}(undef, 2, 2) + for i in 2:4 + M[i] = big(i) + end + v = Vector{BigFloat}(undef, 2) + v[2] = big(2) + + @testset "$(nameof(typeof(x)))" for x in + (M, adjoint(M), PermutedDimsArray(M, (2, 1)), view(M, :, :)) + cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{2}()) + duals, indices = cfg.duals, cfg.indices + + ForwardDiff.seed_zero_partials!(duals, x, indices) + @test !isassigned(duals, 1) + @test all(i -> ForwardDiff.value(duals[i]) == x[i], 2:4) + + ForwardDiff.seed!(duals, x, indices, cfg.seeds) + @test !isassigned(duals, 1) + end + + @testset "$(nameof(typeof(x)))" for x in + (UpperTriangular(M), LowerTriangular(M), Diagonal(v)) + cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{2}()) + duals, indices = cfg.duals, cfg.indices + + @test_throws ArgumentError ForwardDiff.seed_zero_partials!(duals, x, indices) + @test_throws ArgumentError ForwardDiff.seed!(duals, x, indices, cfg.seeds) + @test_throws "ArgumentError: cannot differentiate at an input with an unassigned entry " * + "at index 1: that would leave an entry of the $(nameof(typeof(x))) work " * + "buffer unassigned" ForwardDiff.seed!(duals, x, indices, cfg.seeds) + end end end # module