diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index bf0ef99a..6d29f78c 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 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..081e40fd 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. 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/jacobian.jl b/src/jacobian.jl index f14a6a7b..74f88635 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -92,71 +92,88 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is # result extraction # ##################### -function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T} - out_reshaped = result isa AbstractMatrix ? result : reshape(result, length(ydual), n) - ydual_reshaped = vec(ydual) - # Use closure to avoid GPU broadcasting with Type - partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n)) +# The Jacobian is indexed by the linear indices of `x`: column `j` holds the derivatives with respect +# to `x[j]`. Only the seeded entries of `x` have a derivative to extract, so the columns of the +# structurally zero ones are zeroed instead. See #839. + +# Zeroes the whole Jacobian unless every column is going to be written; the written ones are +# overwritten immediately after. In chunk mode the sweep calls this once up front, since the columns +# that no chunk writes belong to none of them in particular. +function zero_unseeded_columns!(::Type{T}, out::AbstractArray, ydual, x) where {T} + structural_length(x) == length(x) || fill!(out, zero(valtype(T, eltype(ydual)))) + return out +end + +# Vector mode is a single chunk that covers every seeded entry of `x`, so it extracts like the sweep. +function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, x::AbstractArray) where {T} + out_reshaped = reshape_jacobian(result, ydual, x) + zero_unseeded_columns!(T, out_reshaped, ydual, x) + extract_jacobian_chunk!(T, out_reshaped, ydual, x, 1, structural_length(x)) return result end -function extract_jacobian!(::Type{T}, result::MutableDiffResult, ydual::AbstractArray, n) where {T} - extract_jacobian!(T, DiffResults.jacobian(result), ydual, n) +function extract_jacobian!(::Type{T}, result::MutableDiffResult, ydual::AbstractArray, x::AbstractArray) where {T} + extract_jacobian!(T, DiffResults.jacobian(result), ydual, x) return result end -function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) where {T} +function extract_jacobian_chunk!(::Type{T}, result, ydual, x::AbstractArray, index, chunksize) where {T} ydual_reshaped = vec(ydual) offset = index - 1 irange = 1:chunksize - col = irange .+ offset # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - result[:, col] .= partials_wrap.(ydual_reshaped, transpose(irange)) + if structural_length(x) == length(x) + result[:, irange .+ offset] .= partials_wrap.(ydual_reshaped, transpose(irange)) + else + cols = Iterators.drop(structural_columns(result, x), offset) + for (i, col) in zip(irange, cols) + result[:, col] .= partials_wrap.(ydual_reshaped, i) + end + end return result end -reshape_jacobian(result, ydual, xdual) = reshape(result, length(ydual), length(xdual)) -reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResults.jacobian(result), ydual, xdual) +function reshape_jacobian(result::AbstractMatrix, ydual, x) + size(result) == (length(ydual), length(x)) || throw(DimensionMismatch( + lazy"cannot store the $(length(ydual))×$(length(x)) Jacobian in a result of size $(size(result))")) + return result +end +reshape_jacobian(result::AbstractArray, ydual, x) = reshape(result, length(ydual), length(x)) +reshape_jacobian(result::DiffResult, ydual, x) = reshape_jacobian(DiffResults.jacobian(result), ydual, x) ############### # vector mode # ############### function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T}) where {F,T} - N = chunksize(cfg) ydual = vector_mode_dual_eval!(f, cfg, x) ydual isa AbstractArray || throw(JACOBIAN_ERROR) - result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), N) - extract_jacobian!(T, result, ydual, N) + result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x)) + extract_jacobian!(T, result, ydual, x) extract_value!(T, result, ydual) return result end function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T}) where {F,T} - N = chunksize(cfg) ydual = vector_mode_dual_eval!(f!, cfg, y, x) - map!(d -> value(T,d), y, ydual) - result = similar(y, length(y), N) - extract_jacobian!(T, result, ydual, N) + result = similar(y, length(y), length(x)) + extract_jacobian!(T, result, ydual, x) map!(d -> value(T,d), y, ydual) return result end function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T}) where {F,T} - N = chunksize(cfg) ydual = vector_mode_dual_eval!(f, cfg, x) - extract_jacobian!(T, result, ydual, N) + extract_jacobian!(T, result, ydual, x) extract_value!(T, result, ydual) return result end function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) where {F,T} - N = chunksize(cfg) ydual = vector_mode_dual_eval!(f!, cfg, y, x) map!(d -> value(T,d), y, ydual) - extract_jacobian!(T, result, ydual, N) + extract_jacobian!(T, result, ydual, x) extract_value!(T, result, y, ydual) return result end @@ -191,8 +208,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 +218,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 +236,14 @@ end @eval function chunk_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(:(xdual = cfg.duals), :(ydual = f(xdual)), - :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)), + :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))), :())) end @eval function chunk_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(:((ydual, xdual) = cfg.duals), :(f!(seed_zero_partials!(ydual, y), xdual)), - :(result = similar(y, length(y), xlen)), + :(result = similar(y, length(y), length(x))), :(map!(d -> value(T,d), y, ydual)))) end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 94e7cddd..4d3f4d92 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")) @@ -50,6 +51,40 @@ end @test iszero(allocs_jacobian!()) end +# `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural +# position to an index of `x` must not allocate, whether or not `x` has structurally zero entries. +function allocs_structured_gradient!(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 + @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 bf121239..f80468da 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -275,6 +275,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 8be72ee5..28fd03ec 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -156,6 +156,49 @@ 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 + @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..5d2b901a 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -322,6 +322,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..925bbfb8 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,14 @@ 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}())) + end end end # module diff --git a/test/SeedTest.jl b/test/SeedTest.jl index 02b821c3..a6b3944a 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)