diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 26abf43b..efb9aea6 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -57,7 +57,7 @@ end @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_gradient!(T, result, f(dualize(T, x))) + return extract_gradient!(T, result, 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 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 1d54d7bb..7ab61f21 100644 --- a/src/apiutils.jl +++ b/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 diff --git a/src/gradient.jl b/src/gradient.jl index a5ef3dac..b05c5fa8 100644 --- a/src/gradient.jl +++ b/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,8 +163,10 @@ function chunk_mode_gradient_expr(result_definition::Expr) seed!(xdual, x, 1, seeds) seed_zero_partials!(xdual, x, N + 1, xlen - N) ydual = f(xdual) + ydual isa Real || throw(GRAD_ERROR) $(result_definition) - 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 @@ -142,14 +174,14 @@ function chunk_mode_gradient_expr(result_definition::Expr) i = ((c - 1) * N + 1) seed!(xdual, x, i, seeds) ydual = f(xdual) - extract_gradient_chunk!(T, result, 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) diff --git a/src/hessian.jl b/src/hessian.jl index 489b9014..08894bfc 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -9,6 +9,14 @@ Return `H(f)` evaluated at `x`, assuming `f` is called as `f(x)`. The returned Hessian is exactly symmetric: its two triangles are filled from the same derivative values. +Both axes of the result are indexed by the linear indices of `x`, so it is +`length(x)` by `length(x)`. For an `x` with structurally zero entries, such as a +`LowerTriangular`, `UpperTriangular` or `Diagonal` matrix, only the entries that are not +structurally zero are differentiated; the rows and columns of the others are zero. Note +that this makes the result of `hessian(f, ::Diagonal)` quadratic in `length(x)` and hence +quartic in the size of the diagonal — differentiate with respect to the diagonal vector +instead if that matters. + This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. @@ -27,6 +35,10 @@ Compute `H(f)` evaluated at `x` and store the result(s) in `result`, assuming `f called as `f(x)`. The stored Hessian is exactly symmetric: its two triangles are filled from the same derivative values. +`result` has to hold `length(x)^2` entries, indexed as described for +`ForwardDiff.hessian`; a matrix `result` is written to as is and hence has to be +`length(x)` by `length(x)`. + This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. @@ -34,8 +46,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(result, x) CHK && checktag(T, f, x) - xlen = structural_length(x) - H = result isa AbstractMatrix ? result : reshape(result, xlen, xlen) + H = reshape_hessian(result, x) symmetric_hessian!(H, f, x, cfg, nothing) return result end @@ -53,9 +64,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) - xlen = structural_length(x) - hess = DiffResults.hessian(result) - H = hess isa AbstractMatrix ? hess : reshape(hess, xlen, xlen) + H = reshape_hessian(result, x) _, ydual = symmetric_hessian!(H, f, x, cfg, DiffResults.gradient(result)) result = DiffResults.value!(result, value(T, value(T, ydual))) return result @@ -68,24 +77,49 @@ end const HESSIAN_ERROR = DimensionMismatch("hessian(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?") # Copy a block from the nested partials and fill its transpose. On diagonal blocks, read -# only the upper triangle so the result is exactly symmetric. -function extract_hessian_chunk!(::Type{T}, H, ydual, roffset, coffset, rsize, csize) where {T} +# only the upper triangle so the result is exactly symmetric. `positions` maps a seeding position to +# the row and column of `H` it belongs to, i.e. to the linear index of `x` it was seeded from. +# +# It is walked, not indexed, so that the lazy positions of the triangular wrappers never have to be +# materialized. Walking to the start of a block costs what `seed_hessian_chunk!` already pays to +# seed the same block, and no more, since each row resumes the column walk from a saved state +# instead of restarting it. +function extract_hessian_chunk!(::Type{T}, H, positions, ydual, roffset, coffset, rsize, csize) where {T} + diagonal = roffset == coffset + rows = Iterators.drop(positions, roffset) + cols = Iterators.drop(positions, coffset) + rnext = iterate(rows) + cfirst = iterate(cols) for r in 1:rsize + i, rstate = something(rnext) drow = partials(T, ydual, r) - cstart = roffset == coffset ? r : 1 + # A diagonal block reads only the upper triangle, whose row `r` begins at column `r` — the + # position the row walk has just reached. Any other block begins at its first column. + cstart, cnext = diagonal ? (r, (i, rstate)) : (1, something(cfirst)) for c in cstart:csize + j, cstate = cnext h = partials(T, drow, c) - H[roffset + r, coffset + c] = h - H[coffset + c, roffset + r] = h + H[i, j] = h + H[j, i] = h + c == csize || (cnext = something(iterate(cols, cstate))) end + r == rsize || (rnext = something(iterate(rows, rstate))) end return H end # The inner partials of a diagonal block contain the corresponding gradient chunk. -extract_hessian_gradient_chunk!(::Type{T}, ::Nothing, ydual, index, chunksize) where {T} = nothing -extract_hessian_gradient_chunk!(::Type{T}, grad, ydual, index, chunksize) where {T} = - extract_gradient_chunk!(T, grad, value(T, ydual), index, chunksize) +extract_hessian_gradient_chunk!(::Type{T}, ::Nothing, ydual, x, index, chunksize) where {T} = nothing +extract_hessian_gradient_chunk!(::Type{T}, grad, ydual, x, index, chunksize) where {T} = + extract_gradient_chunk!(T, grad, value(T, ydual), x, index, chunksize) + +function reshape_hessian(result::AbstractMatrix, x) + size(result) == (length(x), length(x)) || throw(DimensionMismatch( + lazy"cannot store the $(length(x))×$(length(x)) Hessian in a result of size $(size(result))")) + return result +end +reshape_hessian(result::AbstractArray, x) = reshape(result, length(x), length(x)) +reshape_hessian(result::DiffResult, x) = reshape_hessian(DiffResults.hessian(result), x) # Evaluate one pair of chunks at a time using nested duals. Only one triangle of block # pairs is evaluated; the other is filled by symmetry (see #836). @@ -111,28 +145,41 @@ function symmetric_hessian_expr(result_definition::Expr) ydual1 = f(xdual) ydual1 isa Real || throw(HESSIAN_ERROR) $(result_definition) - extract_hessian_chunk!(T, H, ydual1, 0, 0, N, N) - extract_hessian_gradient_chunk!(T, grad, ydual1, 1, N) + # Only the structurally non-zero entries of `x` are seeded, but both axes of `H` are indexed + # by the linear indices of `x`, as the columns of a Jacobian are, so the columns a Jacobian + # would receive derivatives in are the rows and columns this sweep writes. See #839. + positions = structural_columns(H, x) + # The entries that no block of the sweep writes belong to none of them in particular: for + # `H` the rows and columns of the structurally zero entries of `x`, for `grad` their + # entries. Both take `value(T, ydual1)` for its value type rather than reading `eltype` off + # the result, which is no number type at all for a result that stores `Any` — a scalar dual + # is as good an argument as the array of them `zero_unseeded_columns!` usually gets. + zero_unseeded_columns!(T, H, value(T, ydual1), x) + grad === nothing || zero_unseeded!(T, grad, value(T, ydual1), x) + extract_hessian_chunk!(T, H, positions, ydual1, 0, 0, N, N) + extract_hessian_gradient_chunk!(T, grad, ydual1, x, 1, N) nblocks > 1 && seed_hessian_chunk!(xdual, x, 1, nothing, nothing) for q in 2:nblocks qoffset = (q - 1) * N qsize = min(N, xlen - qoffset) - # Off-diagonal blocks: p seeds columns and q seeds rows. The outer seeds for q + # Off-diagonal blocks: q seeds columns and p seeds rows. Every block reads the nested + # partials with the earlier of its two seeding positions in the outer layer, so an entry + # of `H` is read at the same nesting order at every chunk size. The inner seeds for q # remain unchanged throughout this loop. - seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, oseeds, qsize) + seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, nothing, qsize) for p in 1:(q - 1) poffset = (p - 1) * N - seed_hessian_chunk!(xdual, x, poffset + 1, iseeds, nothing) + seed_hessian_chunk!(xdual, x, poffset + 1, nothing, oseeds) ydual = f(xdual) - extract_hessian_chunk!(T, H, ydual, qoffset, poffset, qsize, N) + extract_hessian_chunk!(T, H, positions, ydual, poffset, qoffset, N, qsize) seed_hessian_chunk!(xdual, x, poffset + 1, nothing, nothing) end - # The diagonal block adds q's inner seeds while retaining its outer seeds. + # The diagonal block adds q's outer seeds while retaining its inner seeds. seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, oseeds, qsize) ydual = f(xdual) - extract_hessian_chunk!(T, H, ydual, qoffset, qoffset, qsize, qsize) - extract_hessian_gradient_chunk!(T, grad, ydual, qoffset + 1, qsize) + extract_hessian_chunk!(T, H, positions, ydual, qoffset, qoffset, qsize, qsize) + extract_hessian_gradient_chunk!(T, grad, ydual, x, qoffset + 1, qsize) seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, nothing, qsize) end @@ -141,7 +188,7 @@ function symmetric_hessian_expr(result_definition::Expr) end @eval function symmetric_hessian(f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N} - $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), xlen, xlen)))) + $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), length(x), length(x))))) end @eval function symmetric_hessian!(H, f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N} diff --git a/src/jacobian.jl b/src/jacobian.jl index f14a6a7b..1204c556 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -92,71 +92,90 @@ 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. `ydual` is only read for the type its +# values have, so the Hessian sweep, which zeroes rows as well as columns, passes a single dual +# rather than the array of them a Jacobian has. +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,8 +210,9 @@ 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 @@ -200,14 +220,14 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex i = ((c - 1) * N + 1) seed!(xdual, x, i, seeds) $(compute_ydual) - extract_jacobian_chunk!(T, out_reshaped, ydual, i, N) + extract_jacobian_chunk!(T, out_reshaped, 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 +238,14 @@ end @eval function chunk_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(:(xdual = cfg.duals), :(ydual = f(xdual)), - :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)), + :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))), :())) end @eval function chunk_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(:((ydual, xdual) = cfg.duals), :(f!(seed_zero_partials!(ydual, y), xdual)), - :(result = similar(y, length(y), xlen)), + :(result = similar(y, length(y), length(x))), :(map!(d -> value(T,d), y, ydual)))) end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 3a59a5ad..307f17a5 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -1,13 +1,14 @@ module AllocationsTest using ForwardDiff +using LinearAlgebra using StaticArrays include(joinpath(dirname(@__FILE__), "utils.jl")) convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,Float64,8},4},2}, 1.3) -@testset "Test seed!/seed_zero_partials! allocations" begin +@testset "Test seed!/seed_zero_partials!/seed_hessian_chunk! allocations" begin x = rand(1000) cfg = ForwardDiff.GradientConfig(nothing, x) duals = cfg.duals @@ -34,11 +35,25 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F iseeds = hcfg.jacobian_config.seeds oseeds = hcfg.gradient_config.seeds allocs_hseed!(args...) = @allocated ForwardDiff.seed_hessian_chunk!(args...) - allocs_hseed!(hduals, x, 1, iseeds, oseeds) - @test iszero(allocs_hseed!(hduals, x, 1, iseeds, oseeds)) + # all four seed combinations, since each builds the zero partials of a different set of layers + for (is, os) in ((iseeds, oseeds), (iseeds, nothing), (nothing, oseeds), (nothing, nothing)) + allocs_hseed!(hduals, x, 1, is, os) + @test iszero(allocs_hseed!(hduals, x, 1, is, os)) + end allocs_hseed!(hduals, x, 1, nothing, nothing, 4) @test iszero(allocs_hseed!(hduals, x, 1, nothing, nothing, 4)) + # An isbits `V` has free zero partials, so the loop above passes whether or not they are built + # only when they are needed. A `V` whose zeros are heap-allocated shows the difference: seeding + # both layers needs neither layer's zeros, so it has to allocate nothing at all. + xbig = big.(rand(9)) + hcfgbig = ForwardDiff.HessianConfig(nothing, xbig, ForwardDiff.Chunk{3}()) + hdualsbig = hcfgbig.gradient_config.duals + ibig = hcfgbig.jacobian_config.seeds + obig = hcfgbig.gradient_config.seeds + allocs_hseed!(hdualsbig, xbig, 1, ibig, obig) + @test iszero(allocs_hseed!(hdualsbig, xbig, 1, ibig, obig)) + allocs_convert_test_574() = @allocated convert_test_574() allocs_convert_test_574() @test iszero(allocs_convert_test_574()) @@ -60,6 +75,56 @@ end @test iszero(allocs_jacobian!()) end +# `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural +# position to an index of `x` must not allocate, whether or not `x` has structurally zero entries. +function allocs_structured_gradient!(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 + +# The Hessian sweep reads the positions of two blocks at once and re-reads each row block once per +# column block, so unlike `gradient!`/`jacobian!` it cannot simply walk them once from the front. It +# still must not materialize them: it resumes each walk from a saved state instead. +function allocs_structured_hessian!(H, x, chunk) + f(z) = sum(abs2, z) + fill!(H, false) + cfg = ForwardDiff.HessianConfig(f, x, chunk) + ForwardDiff.hessian!(H, f, x, cfg) # warmup + return @allocated ForwardDiff.hessian!(H, f, x, cfg) +end + +@testset "Test gradient!/jacobian!/hessian! 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 + + H = zeros(length(x), length(x)) + for chunk_size in (2, nstruct) + @test iszero(allocs_structured_hessian!(H, x, ForwardDiff.Chunk{chunk_size}())) + end +end + @testset "allocation-free nested StaticArray jacobian" begin # test that nested jacobians of StaticArrays do not allocate. # This is a regression test for issue #798, where the inner jacobian was allocating diff --git a/test/GradientTest.jl b/test/GradientTest.jl index c9967812..85e10dae 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -280,6 +280,86 @@ end end end +# issue #838 +@testset "structured inputs: extraction positions" begin + # The seeds are laid out along `structural_eachindex(x)`, so the derivatives have to be written to + # the corresponding entries of the result and every other entry has to end up at zero. In + # particular this must not be read off `result`, which carries no structure to read it off of when + # it is dense -- as the gradient buffer of a `DiffResults.HessianResult` is even for a structured + # `x`. All results are prefilled with `NaN` so that entries left untouched are caught. + @testset "$T, n = $n" for T in (LowerTriangular, UpperTriangular, Diagonal), n in (3, 10) + M = rand(n, n) + x = T(randn(n, n)) + f = z -> dot(M, z) + expected = T(M) # zero derivative for the structurally zero entries + dense_expected = Matrix(expected) + val = f(x) + nstruct = ForwardDiff.structural_length(x) + + @testset "chunk size = $c" for c in unique((1, 2, nstruct)) + cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}()) + + # 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 ForwardDiff.gradient!(T(fill(NaN, n, n)), f, dense_x, dense_cfg) + end + end +end + +@testset "result not shaped like x" begin + # The extraction positions are indices of `x`, which a result of a different shape cannot be + # indexed by, dense `x` included. + x = randn(4) + f = z -> dot(z, z) + @testset "chunk size = $c" for c in (2, 4) + cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.gradient!(fill(NaN, 5), f, x, cfg) + result = DiffResults.DiffResult(NaN, fill(NaN, 5)) + @test_throws DimensionMismatch ForwardDiff.gradient!(result, f, x, cfg) + end +end + +@testset "mutable gradient buffer in an immutable result" begin + # A `StaticArray` buffer makes the result immutable, but an `MVector` can still be written to + # entry by entry, which is how the chunk mode sweep fills it. + x = randn(6) + f = z -> dot(z, z) + result = DiffResults.DiffResult(NaN, @MVector fill(NaN, 6)) + @test result isa DiffResults.ImmutableDiffResult + ForwardDiff.gradient!(result, f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{2}())) + @test DiffResults.gradient(result) ≈ 2 .* x +end + # issue #769 @testset "functions with `Dual` output" begin x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))] diff --git a/test/HessianTest.jl b/test/HessianTest.jl index 119fd14d..3d59effe 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -178,25 +178,149 @@ for T in (StaticArrays.SArray, StaticArrays.MArray) @test DiffResults.hessian(sresult3) == DiffResults.hessian(result) end -@testset "LowerTriangular, UpperTriangular and Diagonal" begin - for n in (3, 5), T in (LowerTriangular, UpperTriangular, Diagonal) - x = T(randn(n, n)) - xlen = ForwardDiff.structural_length(x) - weights = reshape(collect(1.0:n^2), n, n) - objective = x -> dot(weights, abs2.(x)) - expected = diagm(2 .* [weights[idx] for idx in ForwardDiff.structural_eachindex(x)]) - - H = ForwardDiff.hessian(objective, x) - @test size(H) == (xlen, xlen) +# 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, xlen, xlen) - ForwardDiff.hessian!(out, objective, x) + out = fill(NaN, L, L) + @test ForwardDiff.hessian!(out, f, x, cfg) === out @test out == expected - flat = fill(NaN, xlen^2) - ForwardDiff.hessian!(flat, objective, x) - @test reshape(flat, xlen, xlen) == expected + # a result that is not a matrix is reshaped, so it only has to match in length + flat = fill(NaN, L^2) + @test ForwardDiff.hessian!(flat, f, x, cfg) === flat + @test reshape(flat, L, L) == expected + + # a wrongly shaped result is rejected rather than partially filled, the structural shape + # this used to have included + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, L + 1, L + 1), f, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, length(sidx), length(sidx)), f, x, cfg) + + # `DiffResults.HessianResult` allocates a dense gradient buffer even for a structured `x` + result = DiffResults.HessianResult(x) + result = ForwardDiff.hessian!(result, f, x, + ForwardDiff.HessianConfig(f, result, x, ForwardDiff.Chunk{c}())) + @test DiffResults.value(result) ≈ val + @test DiffResults.gradient(result) == grad + @test DiffResults.hessian(result) == expected + end +end + +@testset "chunk size independence" begin + # Every block of the sweep reads the nested partials at the same nesting order, so the result is + # not merely close across chunk sizes but bitwise equal, and equal to the vector mode of the + # `StaticArray` path, which reads them that way too. + n = 16 + f = z -> sum(sin(z[i]) * exp(z[mod1(i + 1, n)]) / (1 + z[mod1(i + 2, n)]^2) for i in eachindex(z)) + x = collect(range(0.1, 0.9; length=n)) + reference = ForwardDiff.hessian(f, SVector{n}(x)) + @test reference == transpose(reference) + @testset "chunk size = $c" for c in (1, 2, 3, 5, 7, 11, n) + H = ForwardDiff.hessian(f, x, ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}())) + @test H == reference + end + + # A structured `x` reads its positions block by block rather than in one pass, and the coupled + # `f` of the structured testset above has an exactly representable Hessian, so its `==` says + # nothing about floating point. There is no `StaticArray` wrapper to compare against, so the + # largest chunk size is the reference. + g = z -> sum(sin(z[i]) * exp(z[j]) / (1 + z[k]^2) + for i in eachindex(z), j in eachindex(z), k in eachindex(z) if i <= j <= k) + @testset "$(nameof(typeof(w)))" for w in (LowerTriangular(rand(4, 4) .+ 1), + UpperTriangular(rand(4, 4) .+ 1), + Diagonal(rand(4) .+ 1)) + nstruct = ForwardDiff.structural_length(w) + reference = ForwardDiff.hessian(g, w, ForwardDiff.HessianConfig(g, w, ForwardDiff.Chunk{nstruct}())) + @test reference == transpose(reference) + @testset "chunk size = $c" for c in 1:(nstruct - 1) + @test ForwardDiff.hessian(g, w, ForwardDiff.HessianConfig(g, w, ForwardDiff.Chunk{c}())) == reference + end + end +end + +@testset "wrongly shaped result" begin + # A matrix result is written to as is, so it has to have the shape of the Hessian and not merely + # as many entries. Results of other shapes are reshaped and only have to match in length. + x = randn(3) + f = z -> sum(abs2, z) + sx = SVector{3}(x) + @testset "chunk size = $c" for c in (2, 3) + cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 4, 4), f, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 10), f, x, cfg) + result = DiffResults.DiffResult(NaN, fill(NaN, 3), fill(NaN, 4, 4)) + @test_throws DimensionMismatch ForwardDiff.hessian!(result, f, x, cfg) + # a `DiffResult` holding a Hessian that is not a matrix is reshaped like a bare one + result = DiffResults.DiffResult(NaN, fill(NaN, 3), fill(NaN, 9)) + result = ForwardDiff.hessian!(result, f, x, cfg) + @test reshape(DiffResults.hessian(result), 3, 3) == 2I(3) + end + # the `StaticArray` methods bypass the sweep and reshape the result themselves + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 4, 4), f, sx) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 10), f, sx) + # a result that is not a matrix goes through `reshape` + flat = fill(NaN, 9) + @test ForwardDiff.hessian!(flat, f, sx) === flat + @test reshape(flat, 3, 3) == 2I(3) +end + +@testset "empty input" begin + f = z -> 1.0 + sum(z) + @test ForwardDiff.hessian(f, Float64[]) == zeros(0, 0) + @test ForwardDiff.hessian(f, SVector{0,Float64}()) == zeros(0, 0) + # `DiffResults.HessianResult` cannot be built from an empty `x`, so assemble one by hand + result = DiffResults.DiffResult(NaN, Float64[], Matrix{Float64}(undef, 0, 0)) + result = ForwardDiff.hessian!(result, f, Float64[]) + @test DiffResults.value(result) == 1.0 + @test isempty(DiffResults.hessian(result)) +end + +@testset "f(x) is not a real number" begin + x = randn(3) + sx = SVector{3}(x) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 3, 3), identity, sx) + @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.HessianResult(sx), identity, sx) + # the sweep checks the first evaluation itself, before it has a result to write into + @testset "chunk size = $c" for c in (2, 3) + cfg = ForwardDiff.HessianConfig(identity, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.hessian(identity, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 3, 3), identity, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.HessianResult(x), identity, x, cfg) + end +end + +@testset "f(x) ignores x" begin + # `partials` is empty, the second case the `Partials{0}` `StaticArray` method covers + f = z -> 1.0 + @test ForwardDiff.hessian(f, SVector{3}(randn(3))) == zeros(3, 3) + @testset "chunk size = $c" for c in (2, 3) + x = randn(3) + @test ForwardDiff.hessian(f, x, ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}())) == zeros(3, 3) end end diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index adc63cd7..499c6dbb 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -326,6 +326,82 @@ 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 + +@testset "wrongly shaped result" begin + # A matrix result is used as is, so it has to have the shape of the Jacobian and not merely as + # many entries. Results of other shapes are reshaped and only have to match in length. + x = randn(4) + g = z -> [sum(z), sum(abs2, z)] + g! = (y, z) -> (y[1] = sum(z); y[2] = sum(abs2, z); y) + @testset "chunk size = $c" for c in (2, 4) + cfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.jacobian!(fill(NaN, 4, 2), g, x, cfg) + result = DiffResults.DiffResult(fill(NaN, 2), fill(NaN, 4, 2)) + @test_throws DimensionMismatch ForwardDiff.jacobian!(result, g, x, cfg) + + y = fill(NaN, 2) + cfg! = ForwardDiff.JacobianConfig(g!, y, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.jacobian!(fill(NaN, 4, 2), g!, y, x, cfg!) + end +end + # issue #769 @testset "functions with `Dual` output" begin x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))] diff --git a/test/QATest.jl b/test/QATest.jl index 860ccdb0..17acd3bc 100644 --- a/test/QATest.jl +++ b/test/QATest.jl @@ -1,6 +1,7 @@ module QATest using ForwardDiff +using LinearAlgebra using Test using JET: @test_opt @@ -11,6 +12,16 @@ using JET: @test_opt @test_opt ForwardDiff.gradient(only, [1.0], ForwardDiff.GradientConfig(only, [1.0], ForwardDiff.Chunk{1}())) @test_opt ForwardDiff.jacobian(identity, [1.0], ForwardDiff.JacobianConfig(identity, [1.0], ForwardDiff.Chunk{1}())) @test_opt ForwardDiff.hessian(only, [1.0], ForwardDiff.HessianConfig(only, [1.0], ForwardDiff.Chunk{1}())) + + # extraction iterates the structural positions of `x` for these + @testset "$(nameof(typeof(x)))" for x in (LowerTriangular(rand(3, 3)), + UpperTriangular(rand(3, 3)), + Diagonal(rand(3, 3))) + @test_opt ForwardDiff.gradient(first, x, ForwardDiff.GradientConfig(first, x, ForwardDiff.Chunk{2}())) + @test_opt ForwardDiff.jacobian(vec, x, ForwardDiff.JacobianConfig(vec, x, ForwardDiff.Chunk{2}())) + # the Hessian sweep walks those positions by hand, so their iteration state has to stay concrete + @test_opt ForwardDiff.hessian(first, x, ForwardDiff.HessianConfig(first, x, ForwardDiff.Chunk{2}())) + end end end # module diff --git a/test/SeedTest.jl b/test/SeedTest.jl index 90ef1858..5a7becbc 100644 --- a/test/SeedTest.jl +++ b/test/SeedTest.jl @@ -53,6 +53,10 @@ end @test collect(ForwardDiff.structural_eachindex(duals, x)) == sidx @test ForwardDiff.structural_length(x) == nstruct + # the columns a Jacobian receives derivatives in, in the order the seeds are laid out in + @test collect(ForwardDiff.structural_columns(zeros(2, length(x)), x)) == LinearIndices(x)[sidx] + @test_throws DimensionMismatch ForwardDiff.structural_columns(zeros(2, length(x) + 1), x) + # `count` defaults to N fill_marker!(duals, x, sidx, marker) ForwardDiff.seed_zero_partials!(duals, x, 4) @@ -106,6 +110,25 @@ end ForwardDiff.seed_hessian_chunk!(duals, x, 4, nothing, nothing) @test all(idx -> iszero(ForwardDiff.partials(ForwardDiff.value(duals[idx]))), sidx) @test all(idx -> iszero(ForwardDiff.partials(duals[idx])), sidx) + + # The off-diagonal blocks of the sweep seed one layer at a time, so check that each seed lands + # in the layer it was asked for and that the other one is left cleared, rather than only that + # something was written. + izero = zero(eltype(iseeds)) + ozero = zero(eltype(oseeds)) + @testset "$(iseeds === nothing ? "outer" : "inner") layer only" for (is, os) in + ((iseeds, nothing), (nothing, oseeds)) + ForwardDiff.seed_hessian_chunk!(duals, x, 4, is, os) + for (i, idx) in enumerate(sidx) + chunkpos = i - 3 + inner = ForwardDiff.partials(ForwardDiff.value(duals[idx])) + outer = ForwardDiff.partials(duals[idx]) + wanted = 1 <= chunkpos <= 3 + @test inner == (wanted && is !== nothing ? is[chunkpos] : izero) + @test outer == (wanted && os !== nothing ? os[chunkpos] : ozero) + end + ForwardDiff.seed_hessian_chunk!(duals, x, 4, nothing, nothing) + end end end # module