From 8e906dd6137fe1ae12b9bc7305f0bd409e232bf4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 20 Aug 2026 16:29:36 +0200 Subject: [PATCH 1/4] Extract derivatives at the indices of `x`, and store those indices in the config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since #739 only the structurally non-zero entries of an input are seeded, but extraction was not updated to match, so the derivatives were written to positions taken from the result container instead of from `x`. `extract_gradient!`/`extract_gradient_chunk!` now take the positions of `x`. Entries that receive no derivative are zeroed, which is their derivative, and the sweep does that once up front: those entries -- the ones belonging to the structural zeros of `x` -- belong to no chunk in particular. The `DiffResult` method splits on mutability, since an immutable result cannot be written to entry by entry, and it dispatches on `DiffResult` rather than `MutableDiffResult`, because a `StaticArray` gradient buffer makes the result immutable even when the buffer itself can be written to entry by entry, as an `MVector` can. Fixes #838, where a dense result got the derivatives at linear positions `1:structural_length(x)` and a `DiffResults.GradientResult` threw, and with it the mis-scattered gradient of `hessian!(::DiffResult, ...)`. The Jacobian is indexed by the linear indices of `x`: column `j` holds `∂f(x)[i]/∂x[j]`, as documented, with hard zeros in the columns of the structural zeros. Its allocations therefore use `length(x)` rather than `structural_length(x)`, which is what `reshape_jacobian` expected all along, so chunk mode stops throwing. Fixes #839. This changes the shape of the result for structured inputs: `jacobian` gains the zero columns and `hessian` inherits both conventions through `jacobian(∇f, x)`, becoming `length(x) x length(x)` with hard-zero rows and columns instead of mixing linear and structural indices. That also makes `hessian!(DiffResults.HessianResult(x), ...)` work. For a `Diagonal` the result now scales with `length(x)`, so differentiating with respect to the diagonal vector is the better choice there. `reshape_jacobian` keeps the short-circuit `extract_jacobian!` was given in #797, with an explicit size check in place of the one `reshape` performed on the way past: since 1.12 `reshape` can no longer return its argument, so every chunk-mode `jacobian!` allocated an `Array` wrapper. Both modes now reject a wrongly shaped matrix result with the same error. Also drops the `map!` that `vector_mode_jacobian(f!, ...)` ran before `extract_jacobian!`, which reads only `ydual`, and that the `map!` after it repeats. Seeding, meanwhile, took the positions to seed from the config's work buffer while extraction took them from `x`, so reusing a config with an input of the same length and a different structure seeded one set of positions and extracted another -- and left the buffer entries outside the input's structure uninitialized for the target function to read. Fixes #842. `structural_eachindex` and `structural_columns` are replaced by `structural_indices`, which returns the linear indices of the seeded entries in seeding order, and each config stores one per work buffer it owns, built from the buffer rather than from `x`. Every API entry then calls `checkstructure` next to `checktag`: it compares `structural_kind`, which is O(1) and a compile-time constant, so it can run unconditionally, and no size or count comparison could replace it -- `LowerTriangular(n, n)` and `UpperTriangular(n, n)` agree on `size` and on `structural_length` alike. Storing the positions also removes the `Iterators.drop` walk that re-traversed them from the front to reach each chunk, three times per gradient chunk. That walk was 35-49% of `gradient!` for an `UpperTriangular` input, and dense inputs paid it too, `Iterators.drop` having no range specialization: `gradient!` is now 1.4-1.9x faster and `jacobian!` 1.2-1.5x, dense and structured alike. Because the positions are linear indices, the Jacobian's structured and dense extraction paths collapse into the one fused broadcast the dense path already used, and neither chunk extractor takes `x` any more. The unassigned-entry branch of seeding is fixed with them, the second half of #842: it called `Base._unsetindex!(duals, idx)`, for which Base has no `CartesianIndex` method at all and whose `AbstractArray` fallback for a linear index recurses forever. `adjoint`, `transpose`, `PermutedDimsArray` and `view` inputs work now, their buffers being plain `Array`s, and the three wrappers `similar` preserves raise an `ArgumentError` naming the entry instead of a `MethodError` or a `StackOverflowError`, since only an `Array` can hold an unassigned entry at all. The new config type parameter is not shimmed: code that pins the parameter count of these types has to be updated. A config for an offset-indexed input now throws at construction rather than at the first `gradient`/`jacobian` call. Co-Authored-By: Claude Opus 5 (1M context) --- ext/ForwardDiffStaticArraysExt.jl | 8 +- src/ForwardDiff.jl | 2 +- src/apiutils.jl | 206 +++++++++++++++++++----------- src/config.jl | 44 +++++-- src/derivative.jl | 6 +- src/gradient.jl | 85 ++++++++---- src/jacobian.jl | 105 ++++++++------- test/AllocationsTest.jl | 60 +++++++-- test/DerivativeTest.jl | 30 +++++ test/GradientTest.jl | 141 ++++++++++++++++++++ test/HessianTest.jl | 66 ++++++++++ test/JacobianTest.jl | 112 ++++++++++++++++ test/QATest.jl | 9 ++ test/SeedTest.jl | 84 +++++++++--- 14 files changed, 759 insertions(+), 199 deletions(-) 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..1b75d29c 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,161 @@ 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 # +######################## + +# Names which `structural_indices` method an array picks, so that two arrays can be checked for +# having the same structural positions without comparing the positions themselves. Add a method here +# whenever one is added below. +structural_kind(::AbstractArray) = nothing +structural_kind(::LowerTriangular) = LowerTriangular +structural_kind(::UpperTriangular) = UpperTriangular +structural_kind(::Diagonal) = Diagonal + +# The linear indices of the entries that are seeded, in seeding order. Configs hold one of these per +# work buffer, so a sweep can index straight to a chunk instead of walking a lazy iterator from the +# front. +# +# The positions are linear indices of the array itself, so one position indexes the input, the work +# buffer and the result alike, and is a Jacobian column number as it stands -- the convention the +# outputs already follow. +# +# Only the triangles are materialized; the other two are ranges already. +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 + +# The positions of the `count` entries starting at structural position `index`. Allocation-free: a +# view of a range is a range, a view of a `Vector` a `SubArray`. A window overrunning the end is a +# `BoundsError` rather than a silently truncated chunk, i.e. silently missing derivatives. +structural_chunk(indices, index, count) = view(indices, index:(index + count - 1)) + +# Checks that the structural positions of `x` are positions of `y` as well. A dense pair may differ +# in shape -- extracting a gradient into a result of the same length has always worked -- but the +# positions of a structured `x` come from its size, so there `y` has to share it. +function check_structural_indices(x::AbstractArray, y::AbstractArray) + require_one_based_indexing(y) + length(x) == length(y) || throw(DimensionMismatch( + lazy"expected an array with $(length(x)) elements, got an array with $(length(y)) elements")) + return nothing +end +function check_structural_indices(x::Union{LowerTriangular,UpperTriangular,Diagonal}, + y::AbstractArray) + require_one_based_indexing(y) + return check_matching_size(x, y) +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 + +# A config's positions were built for its work buffer, which `similar` gave the structure and the size +# of the array the config was constructed for, so they are the positions of `x` too exactly when the +# two are structurally interchangeable. Comparing the kinds is O(1) and a compile-time constant, so +# this can run on every 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 structural_eachindex(x::Diagonal, y::AbstractArray) - require_one_based_indexing(x, y) - if size(x) != size(y) - throw(DimensionMismatch()) + +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. Only an `Array` owns the storage its +# entries live in, and `Base._unsetindex!` is implemented for nothing else -- a structured wrapper is +# a view onto a parent, so it has no slot of its own to unset. Base's `AbstractArray` fallback +# recurses until the stack runs out, so say so instead. +_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 (i, idx) in enumerate(idxs) + duals[idx] = make_dual(x[idx], i) + end + else + for (i, idx) in enumerate(idxs) + if isassigned(x, idx) + duals[idx] = make_dual(x[idx], i) + else + _unsetindex!(duals, idx) + end + end end - return diagind(x) + return duals 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)) +seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, indices) where {T,V,N} = + _seed_zero_partials!(duals, x, indices) # 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, +# `seed!(duals, x, indices, index, seeds, chunksize)`. +function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, indices, 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) + return _seed_zero_partials!(duals, x, structural_chunk(indices, index, count)) end function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N} seed = zero(Partials{N,V}) - if isbitstype(V) - for idx in idxs - duals[idx] = Dual{T,V,N}(x[idx], seed) - end - else - for idx in idxs - if isassigned(x, idx) - duals[idx] = Dual{T,V,N}(x[idx], seed) - else - Base._unsetindex!(duals, idx) - end - end + 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, +function seed!(duals::AbstractArray{Dual{T,V,N}}, x, indices, 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 + return _seed!(duals, x, structural_chunk(indices, 1, N)) do value, i + Dual{T,V,N}(value, seeds[i]) end - return duals end -function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index, +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..1a583d84 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,25 @@ 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 +# The positions belonging to the *input*, whichever constructor built the config: the `f!(y, x)` one +# holds a `(y, x)` pair. +input_indices(cfg::JacobianConfig) = cfg.indices +input_indices(cfg::JacobianConfig{<:Any,<:Any,<:Any,<:Tuple}) = cfg.indices[2] + 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 +269,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..d14a98b1 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,69 @@ 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 the config, i.e. 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)) + extract_gradient!(T, DiffResults.gradient(result), dual, x, indices) + 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, indices) 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) - for (i, idx) in zip(1:npartials(dual), idxs) - result[idx] = partials(T, dual, i) - end +# Readies `result` for extraction, once per sweep: checks that it can be indexed by the positions of +# `x`, then zeroes it unless every entry is going to be written anyway. Comparing counts is exact +# here -- a `result` with the same count but a different structure, an `UpperTriangular` one for a +# `LowerTriangular` `x`, cannot hold the gradient at all. `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 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 -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) +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 +124,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, 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) + result = extract_gradient!(T, result, ydual, x, cfg.indices) return result end @@ -127,29 +153,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..939c734d 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. + +# 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, + 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 # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - result[:, col] .= partials_wrap.(ydual_reshaped, transpose(irange)) + cols = structural_chunk(indices, index, chunksize) + 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, input_indices(cfg)) 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, input_indices(cfg)) 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, input_indices(cfg)) 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, input_indices(cfg)) 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..e4574b17 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,23 @@ 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 + # the windowed 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 +52,40 @@ 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 + @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..6e847569 100644 --- a/test/DerivativeTest.jl +++ b/test/DerivativeTest.jl @@ -106,6 +106,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..5604fc59 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -275,6 +275,147 @@ 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 ForwardDiff.gradient!(T(fill(NaN, n, n)), f, dense_x, dense_cfg) + 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..85cdc5a7 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -322,6 +322,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 From 47a07532a88eb5e8ee6b2c3479bf7f0f8d87d2ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 20 Aug 2026 23:32:01 +0200 Subject: [PATCH 2/4] Name the structural kinds, and make the result check a subset relation `structural_kind` returned the wrapper type itself, with `nothing` for the dense case. It names an index set rather than a wrapper, so it now returns one of four singletons under a `StructuralKind` supertype: the fallback is typed, the dense case reads as "every position" rather than as an absence, and there is no `Type{UnionAll}` dispatch. With the kinds named, `check_structural_indices` can ask the question it actually needs: does the result store every position `x` does? That is a subset relation, not equality -- an `UpperTriangular` result can hold the gradient of a `Diagonal` input -- so `structural_issubset` spells the four pairs out. This also lets the check catch a structured result for a dense `x` on the ForwardDiff side, where it used to fall through to `setindex!` on a triangular wrapper. The size check goes with it. It gained nothing: the case it targeted, a result with the wrong number of entries, is already a length mismatch, and it forbade for a structured `x` what the dense path has always allowed -- a result of the same length and a different shape. Also here: - `vector_mode_gradient!` and `vector_mode_jacobian!` reject a wrongly shaped `f(x)` up front, as the other six sweeps do, instead of failing later with a `MethodError`. - `extract_gradient!(::ImmutableDiffResult, ::Dual, ...)` checks the buffer it copies the partials into wholesale, which is only correct when every entry of `x` is seeded. - `zero_unseeded!` and `zero_unseeded_columns!` return `nothing`, since a caller reassigning their result would silently drop it for an immutable buffer. - `_seed_zero_partials!`, the un-windowed `seed!` body, `check_matching_size` and `input_indices` fold into their callers. - Comments trimmed to what the code does not already say. Co-Authored-By: Claude Opus 5 (1M context) --- src/apiutils.jl | 107 +++++++++++++++------------------------- src/config.jl | 5 -- src/gradient.jl | 27 +++++----- src/jacobian.jl | 20 ++++---- test/AllocationsTest.jl | 2 - test/GradientTest.jl | 12 ++++- test/JacobianTest.jl | 1 + 7 files changed, 75 insertions(+), 99 deletions(-) diff --git a/src/apiutils.jl b/src/apiutils.jl index 1b75d29c..daf9a1f6 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -45,23 +45,23 @@ end # structural positions # ######################## -# Names which `structural_indices` method an array picks, so that two arrays can be checked for -# having the same structural positions without comparing the positions themselves. Add a method here -# whenever one is added below. -structural_kind(::AbstractArray) = nothing -structural_kind(::LowerTriangular) = LowerTriangular -structural_kind(::UpperTriangular) = UpperTriangular -structural_kind(::Diagonal) = Diagonal - -# The linear indices of the entries that are seeded, in seeding order. Configs hold one of these per -# work buffer, so a sweep can index straight to a chunk instead of walking a lazy iterator from the -# front. -# -# The positions are linear indices of the array itself, so one position indexes the input, the work -# buffer and the result alike, and is a Jacobian column number as it stands -- the convention the -# outputs already follow. -# -# Only the triangles are materialized; the other two are ranges already. +# 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)) @@ -81,38 +81,31 @@ function structural_indices(x::LowerTriangular) return [i + (j - 1) * n for j in 1:n for i in j:n] end -# The positions of the `count` entries starting at structural position `index`. Allocation-free: a -# view of a range is a range, a view of a `Vector` a `SubArray`. A window overrunning the end is a -# `BoundsError` rather than a silently truncated chunk, i.e. silently missing derivatives. +# 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)) -# Checks that the structural positions of `x` are positions of `y` as well. A dense pair may differ -# in shape -- extracting a gradient into a result of the same length has always worked -- but the -# positions of a structured `x` come from its size, so there `y` has to share it. +# 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 -function check_structural_indices(x::Union{LowerTriangular,UpperTriangular,Diagonal}, - y::AbstractArray) - require_one_based_indexing(y) - return check_matching_size(x, y) -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 -# A config's positions were built for its work buffer, which `similar` gave the structure and the size -# of the array the config was constructed for, so they are the positions of `x` too exactly when the -# two are structurally interchangeable. Comparing the kinds is O(1) and a compile-time constant, so -# this can run on every call. +# 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)), @@ -134,10 +127,8 @@ end # seeding # ########### -# Mirrors an unassigned entry of `x` into the work buffer. Only an `Array` owns the storage its -# entries live in, and `Base._unsetindex!` is implemented for nothing else -- a structured wrapper is -# a view onto a parent, so it has no slot of its own to unset. Base's `AbstractArray` fallback -# recurses until the stack runs out, so say so instead. +# 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, @@ -164,34 +155,18 @@ _unsetindex!(duals::AbstractArray, idx) = throw(ArgumentError(LazyString( return duals 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, indices) where {T,V,N} = - _seed_zero_partials!(duals, x, indices) - -# 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, indices, index, seeds, chunksize)`. -function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, indices, index, - count = N) where {T,V,N} - return _seed_zero_partials!(duals, x, structural_chunk(indices, index, count)) -end - -function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N} +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 end -function seed!(duals::AbstractArray{Dual{T,V,N}}, x, indices, - seeds::NTuple{N,Partials{N,V}}) where {T,V,N} - return _seed!(duals, x, structural_chunk(indices, 1, N)) do value, i - Dual{T,V,N}(value, seeds[i]) - end -end +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} diff --git a/src/config.jl b/src/config.jl index 1a583d84..3db333f6 100644 --- a/src/config.jl +++ b/src/config.jl @@ -199,11 +199,6 @@ function JacobianConfig(f::F, return JacobianConfig{T,X,N,typeof(duals),typeof(indices)}(seeds, duals, indices) end -# The positions belonging to the *input*, whichever constructor built the config: the `f!(y, x)` one -# holds a `(y, x)` pair. -input_indices(cfg::JacobianConfig) = cfg.indices -input_indices(cfg::JacobianConfig{<:Any,<:Any,<:Any,<:Tuple}) = cfg.indices[2] - checktag(::JacobianConfig{T},f,x) where {T} = checktag(T,f,x) Base.eltype(::Type{JacobianConfig{T,V,N,D,I}}) where {T,V,N,D,I} = Dual{T,V,N} diff --git a/src/gradient.jl b/src/gradient.jl index d14a98b1..34a5ef84 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -52,9 +52,9 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is ##################### # Only the structurally non-zero entries of `x` are seeded, so only those have a derivative to -# extract. The positions come from the config, i.e. 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. +# 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) @@ -69,30 +69,26 @@ function extract_gradient!(::Type{T}, result::MutableDiffResult, dual::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. +# 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 -# Readies `result` for extraction, once per sweep: checks that it can be indexed by the positions of -# `x`, then zeroes it unless every entry is going to be written anyway. Comparing counts is exact -# here -- a `result` with the same count but a different structure, an `UpperTriangular` one for a -# `LowerTriangular` `x`, cannot hold the gradient at all. `dual` is passed for its value type, which -# unlike `eltype(result)` is a number type even for an `Any` result. +# 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 result + return nothing 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. +# `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 result + return nothing end extract_gradient!(::Type{T}, result::AbstractArray, y::Real, x, indices) where {T} = @@ -129,6 +125,7 @@ end function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T, F} ydual = vector_mode_dual_eval!(f, cfg, x) + ydual isa Real || throw(GRAD_ERROR) result = extract_gradient!(T, result, ydual, x, cfg.indices) return result end diff --git a/src/jacobian.jl b/src/jacobian.jl index 939c734d..51852f69 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -100,15 +100,14 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is # 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. +# 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 out + return nothing end -# Vector mode is a single chunk that covers every seeded entry of `x`, so it extracts like the sweep. +# 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) @@ -128,9 +127,9 @@ end function extract_jacobian_chunk!(::Type{T}, result, ydual, indices, index, chunksize) where {T} ydual_reshaped = vec(ydual) irange = 1:chunksize + cols = structural_chunk(indices, index, chunksize) # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - cols = structural_chunk(indices, index, chunksize) result[:, cols] .= partials_wrap.(ydual_reshaped, transpose(irange)) return result end @@ -151,7 +150,7 @@ function vector_mode_jacobian(f::F, x, cfg::JacobianConfig{T}) where {F,T} ydual = vector_mode_dual_eval!(f, cfg, x) ydual isa AbstractArray || throw(JACOBIAN_ERROR) result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x)) - extract_jacobian!(T, result, ydual, x, input_indices(cfg)) + extract_jacobian!(T, result, ydual, x, cfg.indices) extract_value!(T, result, ydual) return result end @@ -159,14 +158,15 @@ end function vector_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T}) where {F,T} ydual = vector_mode_dual_eval!(f!, cfg, y, x) result = similar(y, length(y), length(x)) - extract_jacobian!(T, result, ydual, x, input_indices(cfg)) + 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} ydual = vector_mode_dual_eval!(f, cfg, x) - extract_jacobian!(T, result, ydual, x, input_indices(cfg)) + ydual isa AbstractArray || throw(JACOBIAN_ERROR) + extract_jacobian!(T, result, ydual, x, cfg.indices) extract_value!(T, result, ydual) return result end @@ -174,7 +174,7 @@ end function vector_mode_jacobian!(result, f!::F, y, x, cfg::JacobianConfig{T}) where {F,T} ydual = vector_mode_dual_eval!(f!, cfg, y, x) map!(d -> value(T,d), y, ydual) - extract_jacobian!(T, result, ydual, x, input_indices(cfg)) + extract_jacobian!(T, result, ydual, x, cfg.indices[2]) extract_value!(T, result, y, ydual) return result end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index e4574b17..2caa2337 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -21,8 +21,6 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F allocs_seed!(duals, x, indices, 1, seeds) @test iszero(allocs_seed!(duals, x, indices, 1, seeds)) - # the windowed 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, indices) @test iszero(allocs_szp!(duals, x, indices)) diff --git a/test/GradientTest.jl b/test/GradientTest.jl index 5604fc59..2f30cd3f 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 @@ -326,7 +327,16 @@ end # 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) + @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 + out = UpperTriangular(fill(NaN, n, n)) + ForwardDiff.gradient!(out, f, x, cfg) + @test out == UpperTriangular(dense_expected) + end end end end diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index 85cdc5a7..efbeed1b 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -240,6 +240,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 From 6682b3651de04c752459c659c7ccddd5bfb5ba75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 21 Aug 2026 00:06:26 +0200 Subject: [PATCH 3/4] Cover the three patch lines left untested Each of the three is a pre-existing gap that only entered the diff because threading `indices` through changed its signature: - `structural_issubset(::MainDiagonal, ::LowerTriangle)`: the case of a result whose structure contains that of `x` ran `UpperTriangular` alone, so the structured gradient testset now runs both triangles. - `eltype(::Type{DerivativeConfig})`: `DerivativeTest.jl` never asked a config for its element type, unlike `GradientTest.jl` and `JacobianTest.jl`. - `extract_jacobian(::Type, ::AbstractArray, ::StaticArray)`: `_diff` returns a `StaticArray`, so the `@generated` method always won and the whole fallback body was dead. An `f` returning an `Array` takes it, and gets an `Array` Jacobian. Co-Authored-By: Claude Opus 5 (1M context) --- test/DerivativeTest.jl | 2 ++ test/GradientTest.jl | 8 +++++--- test/JacobianTest.jl | 4 ++++ 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/test/DerivativeTest.jl b/test/DerivativeTest.jl index 6e847569..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) diff --git a/test/GradientTest.jl b/test/GradientTest.jl index 2f30cd3f..9d7d75c2 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -333,9 +333,11 @@ end # a result whose structure *contains* that of `x` is fine, and gets hard zeros off it if T === Diagonal - out = UpperTriangular(fill(NaN, n, n)) - ForwardDiff.gradient!(out, f, x, cfg) - @test out == UpperTriangular(dense_expected) + 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 diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index efbeed1b..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 From ae5bdb86d60c17381487fc2f226d6621903c5493 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 21 Aug 2026 11:38:14 +0200 Subject: [PATCH 4/4] Cover the extraction paths of the `StaticArray` extension The extension has no config to cache the structural positions in, so it derives them per extraction call. Nothing pinned that down: the one static allocation test nests `jacobian`, which the `@generated` methods answer without ever reaching the extraction in `src`. `gradient!`/`jacobian!` do reach it, so they are measured here, into an `Array` as well as a mutable static buffer -- extraction writes the result through a view, and a `view` of an `MArray` at a `UnitRange` is no static array. Co-Authored-By: Claude Opus 5 (1M context) --- test/AllocationsTest.jl | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 2caa2337..5e21ef62 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -84,6 +84,32 @@ end 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