From 503fbec9c40a3150b5be1fe5979124b7a18787b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 17 Aug 2026 17:24:29 +0200 Subject: [PATCH 01/10] Extract derivatives at the indices of `x`, not of the result 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 `x` and walk `structural_eachindex(x, result)`. Entries that receive no derivative are zeroed, which is their derivative; in chunk mode the first chunk does it. The `DiffResult` method splits on mutability, since an immutable result cannot be written to entry by entry (and only occurs for `StaticArray` inputs, all of whose entries are structural). 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. `structural_linearindices` maps a structural position to a linear index of `x` without materializing anything, and the single-broadcast path is kept when every index of `x` is structural, so no path allocates more than before. 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. Co-Authored-By: Claude Opus 5 (1M context) --- ext/ForwardDiffStaticArraysExt.jl | 6 +-- src/ForwardDiff.jl | 2 +- src/apiutils.jl | 9 +++++ src/gradient.jl | 52 +++++++++++++++++-------- src/jacobian.jl | 63 ++++++++++++++++++++----------- test/AllocationsTest.jl | 36 ++++++++++++++++++ test/GradientTest.jl | 46 ++++++++++++++++++++++ test/JacobianTest.jl | 47 +++++++++++++++++++++++ 8 files changed, 218 insertions(+), 43 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index bf0ef99a..2e9999d1 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -57,7 +57,7 @@ end @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_gradient!(T, result, f(dualize(T, x))) + return extract_gradient!(T, result, x, f(dualize(T, 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, x, ydual) 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, x, ydual) 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..c0560cc1 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -70,6 +70,15 @@ function structural_eachindex(x::Diagonal, y::AbstractArray) return diagind(x) end +# The linear indices of the seeded entries of `x`, in seeding order. Results that are not shaped like +# `x`, such as the columns of a Jacobian, are indexed by these. +structural_linearindices(x::AbstractArray) = eachindex(IndexLinear(), x) +structural_linearindices(x::Diagonal) = diagind(x) +function structural_linearindices(x::Union{LowerTriangular,UpperTriangular}) + lin = LinearIndices(x) + return (lin[idx] for idx in structural_eachindex(x)) +end + # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is # currently carrying and to initialize a freshly allocated work buffer, whose elements must all be # written before the target function reads them. diff --git a/src/gradient.jl b/src/gradient.jl index a5ef3dac..29695cf1 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -49,44 +49,64 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is # result extraction # ##################### -function extract_gradient!(::Type{T}, result::DiffResult, y::Real) where {T} +# Derivatives are only computed with respect to the structurally non-zero entries of `x`, since only +# those are seeded. The positions to write to therefore have to be taken from `x`, not from `result`: +# the two may have different structure, e.g. `DiffResults.HessianResult` allocates a dense gradient +# buffer even for a structured `x`. The remaining entries of `result` are zeroed, their derivative +# being zero, unless every index of `x` is structural. See #838. + +function extract_gradient!(::Type{T}, result::DiffResult, x, y::Real) 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, x, dual::Dual) where {T} + result = DiffResults.value!(result, value(T, dual)) + extract_gradient!(T, DiffResults.gradient(result), x, dual) + return result +end + +# Immutable results cannot be written to entry by entry. They only occur for `StaticArray` inputs, +# all of whose entries are structural, so copying the partials wholesale is correct. +function extract_gradient!(::Type{T}, result::ImmutableDiffResult, x, dual::Dual) 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) +extract_gradient!(::Type{T}, result::AbstractArray, x, y::Real) where {T} = fill!(result, zero(y)) +function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) where {T} + structural_length(x) == length(x) || fill!(result, zero(valtype(T, dual))) + idxs = structural_eachindex(x, result) for (i, idx) in zip(1:npartials(dual), idxs) result[idx] = partials(T, dual, i) end return result end -function extract_gradient_chunk!(::Type{T}, result, dual, index, chunksize) where {T} +# The first chunk zeroes `result`; the entries it does not fill are written by the later chunks. In +# chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk. +function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) where {T} offset = index - 1 - idxs = Iterators.drop(structural_eachindex(result), offset) + if iszero(offset) && structural_length(x) != length(x) + fill!(result, zero(valtype(T, dual))) + end + idxs = Iterators.drop(structural_eachindex(x, result), offset) for (i, idx) in zip(1:chunksize, idxs) result[idx] = partials(T, dual, i) end return result end -function extract_gradient_chunk!(::Type{T}, result::DiffResult, dual, index, chunksize) where {T} - extract_gradient_chunk!(T, DiffResults.gradient(result), dual, index, chunksize) +function extract_gradient_chunk!(::Type{T}, result::DiffResult, x, dual, index, chunksize) where {T} + extract_gradient_chunk!(T, DiffResults.gradient(result), x, dual, 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) +extract_gradient_chunk!(::Type, result, x, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR) +extract_gradient_chunk!(::Type, result::DiffResult, x, 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 +118,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, x, ydual) 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, x, ydual) return result end @@ -134,7 +154,7 @@ function chunk_mode_gradient_expr(result_definition::Expr) seed_zero_partials!(xdual, x, N + 1, xlen - N) ydual = f(xdual) $(result_definition) - extract_gradient_chunk!(T, result, ydual, 1, N) + extract_gradient_chunk!(T, result, x, ydual, 1, N) seed_zero_partials!(xdual, x, 1) # do middle chunks @@ -142,14 +162,14 @@ function chunk_mode_gradient_expr(result_definition::Expr) i = ((c - 1) * N + 1) seed!(xdual, x, i, seeds) ydual = f(xdual) - extract_gradient_chunk!(T, result, ydual, i, N) + extract_gradient_chunk!(T, result, x, ydual, i, N) seed_zero_partials!(xdual, x, i) end # do final chunk seed!(xdual, x, lastchunkindex, seeds, lastchunksize) ydual = f(xdual) - extract_gradient_chunk!(T, result, ydual, lastchunkindex, lastchunksize) + extract_gradient_chunk!(T, result, x, ydual, 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..df1a8630 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -92,28 +92,49 @@ 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) +# 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. + +function extract_jacobian!(::Type{T}, result::AbstractArray, x, ydual::AbstractArray) where {T} + out_reshaped = result isa AbstractMatrix ? result : reshape(result, length(ydual), length(x)) 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)) + n = structural_length(x) + if n == length(x) + out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n)) + else + fill!(out_reshaped, zero(valtype(T, eltype(ydual)))) + for (i, col) in zip(1:n, structural_linearindices(x)) + out_reshaped[:, col] .= partials_wrap.(ydual_reshaped, i) + end + end 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, x, ydual::AbstractArray) where {T} + extract_jacobian!(T, DiffResults.jacobian(result), x, ydual) return result end -function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) where {T} +function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize) where {T} ydual_reshaped = vec(ydual) offset = index - 1 irange = 1:chunksize - col = irange .+ offset # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - result[:, col] .= partials_wrap.(ydual_reshaped, transpose(irange)) + if structural_length(x) == length(x) + result[:, irange .+ offset] .= partials_wrap.(ydual_reshaped, transpose(irange)) + else + # The first chunk zeroes the columns of the structurally zero entries, which no chunk writes. + # In chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk. + iszero(offset) && fill!(result, zero(valtype(T, eltype(ydual)))) + idxs = Iterators.drop(structural_linearindices(x), offset) + for (i, col) in zip(irange, idxs) + result[:, col] .= partials_wrap.(ydual_reshaped, i) + end + end return result end @@ -125,38 +146,34 @@ reshape_jacobian(result::DiffResult, ydual, xdual) = reshape_jacobian(DiffResult ############### 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, x, ydual) 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, x, ydual) 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, x, ydual) 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, x, ydual) extract_value!(T, result, y, ydual) return result end @@ -192,7 +209,7 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex ydual isa AbstractArray || throw(JACOBIAN_ERROR) $(result_definition) out_reshaped = reshape_jacobian(result, ydual, xdual) - extract_jacobian_chunk!(T, out_reshaped, ydual, 1, N) + extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N) seed_zero_partials!(xdual, x, 1) # do middle chunks @@ -200,14 +217,14 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex i = ((c - 1) * N + 1) seed!(xdual, x, i, seeds) $(compute_ydual) - extract_jacobian_chunk!(T, out_reshaped, ydual, i, N) + extract_jacobian_chunk!(T, out_reshaped, x, ydual, i, N) seed_zero_partials!(xdual, x, i) end # do final chunk seed!(xdual, x, lastchunkindex, seeds, lastchunksize) $(compute_ydual) - extract_jacobian_chunk!(T, out_reshaped, ydual, lastchunkindex, lastchunksize) + extract_jacobian_chunk!(T, out_reshaped, x, ydual, lastchunkindex, lastchunksize) $(y_definition) @@ -218,14 +235,14 @@ end @eval function chunk_mode_jacobian(f::F, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(:(xdual = cfg.duals), :(ydual = f(xdual)), - :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), xlen)), + :(result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x))), :())) end @eval function chunk_mode_jacobian(f!::F, y, x, cfg::JacobianConfig{T,V,N}) where {F,T,V,N} $(jacobian_chunk_mode_expr(:((ydual, xdual) = cfg.duals), :(f!(seed_zero_partials!(ydual, y), xdual)), - :(result = similar(y, length(y), xlen)), + :(result = similar(y, length(y), length(x))), :(map!(d -> value(T,d), y, ydual)))) end diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 94e7cddd..f4bd2134 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -1,6 +1,7 @@ module AllocationsTest using ForwardDiff +using LinearAlgebra using StaticArrays include(joinpath(dirname(@__FILE__), "utils.jl")) @@ -50,6 +51,41 @@ end @test iszero(allocs_jacobian!()) end +# `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural +# position to an index of `x` must not allocate, whether or not `x` has structurally zero entries. +function allocs_gradient!(x, chunk) + f(z) = sum(abs2, z) + result = fill!(similar(x, Float64), 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_jacobian!(x, chunk) + f!(y, z) = (y[1] = sum(z); y[2] = 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 in (rand(6, 6), + LowerTriangular(rand(6, 6)), + UpperTriangular(rand(6, 6)), + Diagonal(rand(6, 6))) + # vector mode + chunk = ForwardDiff.Chunk{ForwardDiff.structural_length(x)}() + @test iszero(allocs_gradient!(x, chunk)) + @test iszero(allocs_jacobian!(x, chunk)) + + # chunk mode, where `jacobian!` allocates a fixed-size `reshape` wrapper for every input, so + # compare against a dense input of the same size instead of asserting zero + chunk = ForwardDiff.Chunk{2}() + @test iszero(allocs_gradient!(x, chunk)) + @test allocs_jacobian!(x, chunk) == allocs_jacobian!(rand(6, 6), chunk) +end + @testset "allocation-free nested StaticArray jacobian" begin # test that nested jacobians of StaticArrays do not allocate. # This is a regression test for issue #798, where the inner jacobian was allocating diff --git a/test/GradientTest.jl b/test/GradientTest.jl index bf121239..117e8873 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -275,6 +275,52 @@ end end end +# issue #838 +@testset "structured inputs: extraction positions" begin + # The seeds are laid out along `structural_eachindex(x)`, so the derivatives have to be written to + # the corresponding entries of the result and every other entry has to end up at zero. In + # particular this must not be read off `result`, which carries no structure to read it off of when + # it is dense -- as the gradient buffer of a `DiffResults.HessianResult` is even for a structured + # `x`. All results are prefilled with `NaN` so that entries left untouched are caught. + @testset "$T, n = $n" for T in (LowerTriangular, UpperTriangular, Diagonal), n in (3, 10) + M = rand(n, n) + x = T(randn(n, n)) + f = z -> dot(M, z) + expected = T(M) # zero derivative for the structurally zero entries + dense_expected = Matrix(expected) + val = f(x) + nstruct = ForwardDiff.structural_length(x) + + @testset "chunk size = $c" for c in unique((1, 2, nstruct)) + cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}()) + + 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) + end + 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/JacobianTest.jl b/test/JacobianTest.jl index b6d36180..9763d279 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -322,6 +322,53 @@ 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) + + @testset "chunk size = $c" for c in unique((1, 2, ForwardDiff.structural_length(x))) + 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 + + # `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 + 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))] From fea566c9030d56043685f81e13d2ff0d162e6dd0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 17 Aug 2026 21:17:53 +0200 Subject: [PATCH 02/10] Fix the `Diagonal` allocation test, and the `reshape` wrapper it uncovered The new allocation test failed on Julia <= 1.10 for `Diagonal` inputs, but none of the allocations came from extraction: the target function reduced with the no-function `sum(z)`, and `Base._sum(::Diagonal, ::Colon)` allocates there (32 bytes even for a `Diagonal{Float64}`). Reducing with `sum(f, z)` instead measures ForwardDiff rather than LinearAlgebra, and extraction turns out to be allocation-free for every input type on both 1.10 and 1.12. That left the chunk-mode comparison against a dense input, which was hiding a real cost: `reshape_jacobian` reshapes the result even when it already is a matrix, and since 1.11 `reshape` can no longer return its argument, so every chunk-mode `jacobian!` allocated an `Array` wrapper. `extract_jacobian!` had been given that short-circuit in #797; `reshape_jacobian` now shares it, with an explicit size check in place of the one `reshape` performed on the way past, and `extract_jacobian!` calls it instead of repeating the ternary. Both modes now reject a wrongly shaped matrix result with the same error, and the test can assert zero allocations outright. Co-Authored-By: Claude Opus 5 (1M context) --- src/jacobian.jl | 15 +++++++++++---- test/AllocationsTest.jl | 17 ++++++++--------- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index df1a8630..ddd258e9 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -97,7 +97,7 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is # structurally zero ones are zeroed instead. See #839. function extract_jacobian!(::Type{T}, result::AbstractArray, x, ydual::AbstractArray) where {T} - out_reshaped = result isa AbstractMatrix ? result : reshape(result, length(ydual), length(x)) + out_reshaped = reshape_jacobian(result, ydual, x) ydual_reshaped = vec(ydual) # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) @@ -138,8 +138,15 @@ function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize) 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) +# A matrix is used as is: reshaping it would allocate a wrapper on Julia >= 1.11, where `reshape` +# can no longer return its argument. The size is checked instead, as `reshape` did on the way past. +function reshape_jacobian(result::AbstractMatrix, ydual, x) + size(result) == (length(ydual), length(x)) || throw(DimensionMismatch( + lazy"cannot store the $(length(ydual))x$(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 # @@ -208,7 +215,7 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex $(compute_ydual) ydual isa AbstractArray || throw(JACOBIAN_ERROR) $(result_definition) - out_reshaped = reshape_jacobian(result, ydual, xdual) + out_reshaped = reshape_jacobian(result, ydual, x) extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N) seed_zero_partials!(xdual, x, 1) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index f4bd2134..53eb5854 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -53,7 +53,7 @@ end # `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural # position to an index of `x` must not allocate, whether or not `x` has structurally zero entries. -function allocs_gradient!(x, chunk) +function allocs_structured_gradient!(x, chunk) f(z) = sum(abs2, z) result = fill!(similar(x, Float64), false) cfg = ForwardDiff.GradientConfig(f, x, chunk) @@ -61,8 +61,8 @@ function allocs_gradient!(x, chunk) return @allocated ForwardDiff.gradient!(result, f, x, cfg) end -function allocs_jacobian!(x, chunk) - f!(y, z) = (y[1] = sum(z); y[2] = sum(abs2, z); y) +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) @@ -76,14 +76,13 @@ end Diagonal(rand(6, 6))) # vector mode chunk = ForwardDiff.Chunk{ForwardDiff.structural_length(x)}() - @test iszero(allocs_gradient!(x, chunk)) - @test iszero(allocs_jacobian!(x, chunk)) + @test iszero(allocs_structured_gradient!(x, chunk)) + @test iszero(allocs_structured_jacobian!(x, chunk)) - # chunk mode, where `jacobian!` allocates a fixed-size `reshape` wrapper for every input, so - # compare against a dense input of the same size instead of asserting zero + # chunk mode chunk = ForwardDiff.Chunk{2}() - @test iszero(allocs_gradient!(x, chunk)) - @test allocs_jacobian!(x, chunk) == allocs_jacobian!(rand(6, 6), chunk) + @test iszero(allocs_structured_gradient!(x, chunk)) + @test iszero(allocs_structured_jacobian!(x, chunk)) end @testset "allocation-free nested StaticArray jacobian" begin From 07e42721f07bdcbdd611128c0d502461ad0bfef8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 17 Aug 2026 21:49:10 +0200 Subject: [PATCH 03/10] Zero the unwritten entries in the sweep, not in the chunk `extract_gradient_chunk!`/`extract_jacobian_chunk!` recognised the first chunk by `index == 1` and used it to zero the whole result, which is neither part of extracting one chunk nor something a chunk can decide: the entries at stake -- those of the structural zeros of `x` -- belong to no chunk in particular. Both sweeps now do it once up front, and the chunk functions only write their chunk. The gradient's half is shared with `extract_gradient!` as `zero_unseeded!`, which also fixes the condition. `structural_length(x) != length(x)` zeroed a result that is itself structured, whose every stored entry the sweep goes on to write; comparing against `structural_length(result)` skips that, and leaves the mismatched-structure cases erroring at the same write as before. The Jacobian keeps its own test, since its result is not shaped like `x` and what has to be covered there is columns. Also drops the `map!` that `vector_mode_jacobian(f!, ...)` ran before `extract_jacobian!`, which reads only `ydual`, and that the `map!` after it repeats. Co-Authored-By: Claude Opus 5 (1M context) --- src/gradient.jl | 27 ++++++++++++++++++--------- src/jacobian.jl | 6 ++---- 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/src/gradient.jl b/src/gradient.jl index 29695cf1..d3c68dbc 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -53,7 +53,8 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is # those are seeded. The positions to write to therefore have to be taken from `x`, not from `result`: # the two may have different structure, e.g. `DiffResults.HessianResult` allocates a dense gradient # buffer even for a structured `x`. The remaining entries of `result` are zeroed, their derivative -# being zero, unless every index of `x` is structural. See #838. +# being zero, unless the seeded entries of `x` already account for every entry `result` stores. See +# #838. function extract_gradient!(::Type{T}, result::DiffResult, x, y::Real) where {T} result = DiffResults.value!(result, y) @@ -68,17 +69,29 @@ function extract_gradient!(::Type{T}, result::MutableDiffResult, x, dual::Dual) return result end -# Immutable results cannot be written to entry by entry. They only occur for `StaticArray` inputs, -# all of whose entries are structural, so copying the partials wholesale is correct. +# 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, x, dual::Dual) where {T} result = DiffResults.value!(result, value(T, dual)) result = DiffResults.gradient!(result, partials(T, dual)) return result end +# Zeroes the entries that receive no derivative. In chunk mode the sweep calls this once up front, +# since the entries that no chunk writes belong to none of them in particular. +function zero_unseeded!(::Type{T}, result::AbstractArray, x, dual) where {T} + structural_length(x) == structural_length(result) || fill!(result, zero(valtype(T, dual))) + return result +end +function zero_unseeded!(::Type{T}, result::MutableDiffResult, x, dual) where {T} + zero_unseeded!(T, DiffResults.gradient(result), x, dual) + return result +end + extract_gradient!(::Type{T}, result::AbstractArray, x, y::Real) where {T} = fill!(result, zero(y)) function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) where {T} - structural_length(x) == length(x) || fill!(result, zero(valtype(T, dual))) + zero_unseeded!(T, result, x, dual) idxs = structural_eachindex(x, result) for (i, idx) in zip(1:npartials(dual), idxs) result[idx] = partials(T, dual, i) @@ -86,13 +99,8 @@ function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) wher return result end -# The first chunk zeroes `result`; the entries it does not fill are written by the later chunks. In -# chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk. function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) where {T} offset = index - 1 - if iszero(offset) && structural_length(x) != length(x) - fill!(result, zero(valtype(T, dual))) - end idxs = Iterators.drop(structural_eachindex(x, result), offset) for (i, idx) in zip(1:chunksize, idxs) result[idx] = partials(T, dual, i) @@ -154,6 +162,7 @@ function chunk_mode_gradient_expr(result_definition::Expr) seed_zero_partials!(xdual, x, N + 1, xlen - N) ydual = f(xdual) $(result_definition) + zero_unseeded!(T, result, x, ydual) extract_gradient_chunk!(T, result, x, ydual, 1, N) seed_zero_partials!(xdual, x, 1) diff --git a/src/jacobian.jl b/src/jacobian.jl index ddd258e9..752c817b 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -127,9 +127,6 @@ function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize) if structural_length(x) == length(x) result[:, irange .+ offset] .= partials_wrap.(ydual_reshaped, transpose(irange)) else - # The first chunk zeroes the columns of the structurally zero entries, which no chunk writes. - # In chunk mode `structural_length(x) > chunksize`, so `index == 1` only for the first chunk. - iszero(offset) && fill!(result, zero(valtype(T, eltype(ydual)))) idxs = Iterators.drop(structural_linearindices(x), offset) for (i, col) in zip(irange, idxs) result[:, col] .= partials_wrap.(ydual_reshaped, i) @@ -163,7 +160,6 @@ end function vector_mode_jacobian(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) result = similar(y, length(y), length(x)) extract_jacobian!(T, result, x, ydual) map!(d -> value(T,d), y, ydual) @@ -216,6 +212,8 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex ydual isa AbstractArray || throw(JACOBIAN_ERROR) $(result_definition) out_reshaped = reshape_jacobian(result, ydual, x) + # zero the columns of the structurally zero entries of `x`, which no chunk of the sweep writes + structural_length(x) == length(x) || fill!(out_reshaped, zero(valtype(T, eltype(ydual)))) extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N) seed_zero_partials!(xdual, x, 1) From ec3caa40d342a6509ecf7b1ed9a49d0f7595b063 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 17 Aug 2026 22:55:02 +0200 Subject: [PATCH 04/10] Cover the Hessian, and the cases the extraction tests left out `hessian` inherits both conventions through `jacobian(gradient(f), x)`, so it changes shape for a structured `x` and `hessian!(DiffResults.HessianResult(x), f, x)` starts working, none of which was asserted. The new testset pins the value, the gradient and the Hessian of a function whose second derivative is `1 + (a == b)` on the structural entries, so the reference is exact and the hard-zero rows and columns are checked rather than approximated. Against master it fails everywhere: the smaller chunks throw, and the full-length one gets the wrong shape and cannot take a `HessianResult` at all. Also: the out-of-place `gradient` is shaped like `x`, so its zeros are the structural ones and its type is worth asserting; a structured result cannot hold the gradient of a dense `x`, which now throws where it used to write to the wrong entries; a Jacobian result that is not a matrix is reshaped; and the `f!` form takes a `JacobianResult` too. Both structured testsets take the chunk sizes from `length(sidx)` rather than `ForwardDiff.structural_length`, keeping the reference data in the test. Co-Authored-By: Claude Opus 5 (1M context) --- test/GradientTest.jl | 10 ++++++++++ test/HessianTest.jl | 42 ++++++++++++++++++++++++++++++++++++++++++ test/JacobianTest.jl | 12 +++++++++++- 3 files changed, 63 insertions(+), 1 deletion(-) diff --git a/test/GradientTest.jl b/test/GradientTest.jl index 117e8873..8f618fe5 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -294,6 +294,11 @@ end @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 @@ -317,6 +322,11 @@ end # 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 diff --git a/test/HessianTest.jl b/test/HessianTest.jl index 8be72ee5..01d29143 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -156,6 +156,48 @@ 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) + + @testset "chunk size = $c" for c in unique((1, 2, length(sidx))) + cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}()) + + H = ForwardDiff.hessian(f, x, cfg) + @test size(H) == (L, L) + @test H == expected + + out = fill(NaN, L, L) + @test ForwardDiff.hessian!(out, f, x, cfg) === out + @test out == expected + + # `DiffResults.HessianResult` allocates a dense gradient buffer even for a structured `x` + result = DiffResults.HessianResult(x) + result = ForwardDiff.hessian!(result, f, x, + ForwardDiff.HessianConfig(f, result, x, ForwardDiff.Chunk{c}())) + @test DiffResults.value(result) ≈ val + @test DiffResults.gradient(result) == grad + @test DiffResults.hessian(result) == expected + end +end + @testset "branches in dot" begin # https://github.com/JuliaDiff/ForwardDiff.jl/issues/551 H = [1 2 3; 4 5 6; 7 8 9]; diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index 9763d279..0117d844 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -340,7 +340,7 @@ end expected[2, sidx] .= 2 .* x[sidx] val = g(x) - @testset "chunk size = $c" for c in unique((1, 2, ForwardDiff.structural_length(x))) + @testset "chunk size = $c" for c in unique((1, 2, length(sidx))) cfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{c}()) J = ForwardDiff.jacobian(g, x, cfg) @test size(J) == (2, length(x)) @@ -350,6 +350,11 @@ end @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) @@ -366,6 +371,11 @@ end 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 From 27033bbeedf810d72d822a93001d866cd9ce6d71 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 18 Aug 2026 17:23:08 +0200 Subject: [PATCH 05/10] Take `x` last, and extract a Jacobian column in one place `extract_gradient!`, `extract_gradient_chunk!`, `extract_jacobian!` and `extract_jacobian_chunk!` take `x` after the derivatives, like every other internal function that pairs an output with an input: `reshape_jacobian(result, ydual, x)`, which `extract_jacobian!` calls on its first line, `extract_value!(T, out, y, ydual)`, and the `extract_jacobian(T, ydual, x::StaticArray)` of the StaticArrays extension. `x` is annotated `AbstractArray` there, so a call written against an older signature fails on the method rather than inside `structural_length`. `extract_jacobian!` duplicated the structured branch of `extract_jacobian_chunk!` at `offset == 0` and now delegates to it, and both Jacobian sweeps zero through one helper. `structural_linearindices(x)` becomes `structural_columns(out, x)`: what a caller needs are the columns of `out` that receive a derivative, and every method returns those, `axes(out, 2)` or a subset of it, without materializing anything. It checks its arguments the way `structural_eachindex` does, whose three bare `DimensionMismatch()` now name the sizes they expected. The chunk mode gradient sweep rejects an array-valued `f` up front, as the Jacobian sweep does, in place of the `dual::AbstractArray` methods of `extract_gradient_chunk!`: `zero_unseeded!` runs before them and would fail on `zero(::Type{<:AbstractArray})` before dispatch got there. `zero_unseeded!` also dispatches on `DiffResult` rather than `MutableDiffResult`, since a `StaticArray` gradient buffer makes the result immutable even when the buffer itself can be written to entry by entry, as an `MVector` can. Co-Authored-By: Claude Opus 5 (1M context) --- ext/ForwardDiffStaticArraysExt.jl | 6 ++-- src/apiutils.jl | 55 +++++++++++++++++++++--------- src/gradient.jl | 55 ++++++++++++++++-------------- src/jacobian.jl | 56 ++++++++++++++----------------- 4 files changed, 98 insertions(+), 74 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 2e9999d1..6d29f78c 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -57,7 +57,7 @@ end @inline function ForwardDiff.vector_mode_gradient!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) - return extract_gradient!(T, result, x, f(dualize(T, x))) + return extract_gradient!(T, result, f(dualize(T, x)), x) end # Jacobian @@ -87,13 +87,13 @@ end function extract_jacobian(::Type{T}, ydual::AbstractArray, x::StaticArray) where T result = similar(ydual, valtype(T, eltype(ydual)), length(ydual), length(x)) - return extract_jacobian!(T, result, x, ydual) + return extract_jacobian!(T, result, ydual, x) end @inline function ForwardDiff.vector_mode_jacobian!(result, f::F, x::StaticArray) where {F} T = typeof(Tag(f, eltype(x))) ydual = f(dualize(T, x)) - result = extract_jacobian!(T, result, x, ydual) + result = extract_jacobian!(T, result, ydual, x) result = extract_value!(T, result, ydual) return result end diff --git a/src/apiutils.jl b/src/apiutils.jl index c0560cc1..081e40fd 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -48,35 +48,60 @@ function structural_eachindex(x::AbstractArray, y::AbstractArray) end function structural_eachindex(x::UpperTriangular, y::AbstractArray) require_one_based_indexing(x, y) - if size(x) != size(y) - throw(DimensionMismatch()) - end + check_matching_size(x, y) n = size(x, 1) return (CartesianIndex(i, j) for j in 1:n for i in 1:j) end function structural_eachindex(x::LowerTriangular, y::AbstractArray) require_one_based_indexing(x, y) - if size(x) != size(y) - throw(DimensionMismatch()) - end + check_matching_size(x, y) n = size(x, 1) return (CartesianIndex(i, j) for j in 1:n for i in j:n) end function structural_eachindex(x::Diagonal, y::AbstractArray) require_one_based_indexing(x, y) - if size(x) != size(y) - throw(DimensionMismatch()) - end + check_matching_size(x, y) return diagind(x) end -# The linear indices of the seeded entries of `x`, in seeding order. Results that are not shaped like -# `x`, such as the columns of a Jacobian, are indexed by these. -structural_linearindices(x::AbstractArray) = eachindex(IndexLinear(), x) -structural_linearindices(x::Diagonal) = diagind(x) -function structural_linearindices(x::Union{LowerTriangular,UpperTriangular}) +# The two arrays are indexed by the same indices, so they have to have the same size. This is the +# error a `gradient!` into a container that is not shaped like `x` aborts with, so it names the sizes. +function check_matching_size(x::AbstractArray, y::AbstractArray) + size(x) == size(y) || throw(DimensionMismatch( + lazy"expected an array of size $(size(x)), got an array of size $(size(y))")) + return nothing +end + +# The columns of the Jacobian `out` that receive derivatives, in seeding order. Column `j` holds the +# derivatives with respect to `x[j]`, so a seeded entry writes to the column at its position in the +# linear order of `x`. Every entry of an array is seeded unless one of the methods below applies. +function structural_columns(out::AbstractMatrix, x::AbstractArray) + require_one_based_indexing(out, x) + check_matching_columns(out, x) + return axes(out, 2) +end +# The seeded columns are runs of increasing length, so they are no range. Deriving their order from +# `structural_eachindex` rather than recomputing it keeps a single source of truth: the two have to +# agree entry by entry, or the derivatives land in the wrong columns. +function structural_columns(out::AbstractMatrix, x::Union{LowerTriangular,UpperTriangular}) + require_one_based_indexing(out, x) + check_matching_columns(out, x) + cols = axes(out, 2) lin = LinearIndices(x) - return (lin[idx] for idx in structural_eachindex(x)) + return (cols[lin[idx]] for idx in structural_eachindex(x)) +end +# `diagind` is already a range of linear positions, so it can select the columns directly. +function structural_columns(out::AbstractMatrix, x::Diagonal) + require_one_based_indexing(out, x) + check_matching_columns(out, x) + return axes(out, 2)[diagind(x)] +end + +# A column of the Jacobian belongs to an entry of `x`, so there have to be as many as `x` has entries. +function check_matching_columns(out::AbstractMatrix, x::AbstractArray) + size(out, 2) == length(x) || throw(DimensionMismatch( + lazy"expected a matrix with $(length(x)) columns, got a matrix with $(size(out, 2)) columns")) + return nothing end # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is diff --git a/src/gradient.jl b/src/gradient.jl index d3c68dbc..b05c5fa8 100644 --- a/src/gradient.jl +++ b/src/gradient.jl @@ -53,45 +53,50 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is # those are seeded. The positions to write to therefore have to be taken from `x`, not from `result`: # the two may have different structure, e.g. `DiffResults.HessianResult` allocates a dense gradient # buffer even for a structured `x`. The remaining entries of `result` are zeroed, their derivative -# being zero, unless the seeded entries of `x` already account for every entry `result` stores. See -# #838. +# being zero, unless `x` has as many seeded entries as `result` has entries. See #838. -function extract_gradient!(::Type{T}, result::DiffResult, x, y::Real) where {T} +function extract_gradient!(::Type{T}, result::DiffResult, y::Real, x) where {T} result = DiffResults.value!(result, y) grad = DiffResults.gradient(result) fill!(grad, zero(y)) return result end -function extract_gradient!(::Type{T}, result::MutableDiffResult, x, dual::Dual) where {T} +function extract_gradient!(::Type{T}, result::MutableDiffResult, dual::Dual, x) where {T} result = DiffResults.value!(result, value(T, dual)) - extract_gradient!(T, DiffResults.gradient(result), x, dual) + extract_gradient!(T, DiffResults.gradient(result), dual, x) return result end # Immutable results cannot be written to entry by entry. Copying the partials wholesale is correct # as long as every entry of `x` is seeded, which holds for the `StaticArray` gradient buffers that # are the only source of such results; anything else throws on the length mismatch. -function extract_gradient!(::Type{T}, result::ImmutableDiffResult, x, dual::Dual) where {T} +function extract_gradient!(::Type{T}, result::ImmutableDiffResult, dual::Dual, x) where {T} result = DiffResults.value!(result, value(T, dual)) result = DiffResults.gradient!(result, partials(T, dual)) return result end -# Zeroes the entries that receive no derivative. In chunk mode the sweep calls this once up front, -# since the entries that no chunk writes belong to none of them in particular. -function zero_unseeded!(::Type{T}, result::AbstractArray, x, dual) where {T} +# Zeroes `result` unless extraction is going to write every entry of it; the written ones are +# overwritten immediately after. In chunk mode the sweep calls this once up front, since the entries +# that no chunk writes belong to none of them in particular. Comparing counts rather than matching +# entries up is exact for every pair that can hold the gradient: a differently structured `result` +# with the same count, an `UpperTriangular` one for a `LowerTriangular` `x`, cannot. `dual` is passed +# for its value type, which unlike `eltype(result)` is a number type even for an `Any` result. +function zero_unseeded!(::Type{T}, result::AbstractArray, dual, x) where {T} structural_length(x) == structural_length(result) || fill!(result, zero(valtype(T, dual))) return result end -function zero_unseeded!(::Type{T}, result::MutableDiffResult, x, dual) where {T} - zero_unseeded!(T, DiffResults.gradient(result), x, dual) +# Dispatched on `DiffResult`, not on `MutableDiffResult`: a `StaticArray` gradient buffer makes the +# result immutable even when the buffer itself can be written to, as an `MVector` can. +function zero_unseeded!(::Type{T}, result::DiffResult, dual, x) where {T} + zero_unseeded!(T, DiffResults.gradient(result), dual, x) return result end -extract_gradient!(::Type{T}, result::AbstractArray, x, y::Real) where {T} = fill!(result, zero(y)) -function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) where {T} - zero_unseeded!(T, result, x, dual) +extract_gradient!(::Type{T}, result::AbstractArray, y::Real, x) where {T} = fill!(result, zero(y)) +function extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual, x) where {T} + zero_unseeded!(T, result, dual, x) idxs = structural_eachindex(x, result) for (i, idx) in zip(1:npartials(dual), idxs) result[idx] = partials(T, dual, i) @@ -99,7 +104,7 @@ function extract_gradient!(::Type{T}, result::AbstractArray, x, dual::Dual) wher return result end -function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) where {T} +function extract_gradient_chunk!(::Type{T}, result, dual, x, index, chunksize) where {T} offset = index - 1 idxs = Iterators.drop(structural_eachindex(x, result), offset) for (i, idx) in zip(1:chunksize, idxs) @@ -108,14 +113,11 @@ function extract_gradient_chunk!(::Type{T}, result, x, dual, index, chunksize) w return result end -function extract_gradient_chunk!(::Type{T}, result::DiffResult, x, dual, index, chunksize) where {T} - extract_gradient_chunk!(T, DiffResults.gradient(result), x, dual, index, chunksize) +function extract_gradient_chunk!(::Type{T}, result::DiffResult, dual, x, index, chunksize) where {T} + extract_gradient_chunk!(T, DiffResults.gradient(result), dual, x, index, chunksize) return result end -extract_gradient_chunk!(::Type, result, x, dual::AbstractArray, index, chunksize) = throw(GRAD_ERROR) -extract_gradient_chunk!(::Type, result::DiffResult, x, 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)?") ############### @@ -126,12 +128,12 @@ function vector_mode_gradient(f::F, x, cfg::GradientConfig{T}) where {T, F} ydual = vector_mode_dual_eval!(f, cfg, x) ydual isa Real || throw(GRAD_ERROR) result = similar(x, valtype(T, ydual)) - return extract_gradient!(T, result, x, ydual) + return extract_gradient!(T, result, ydual, x) end function vector_mode_gradient!(result, f::F, x, cfg::GradientConfig{T}) where {T, F} ydual = vector_mode_dual_eval!(f, cfg, x) - result = extract_gradient!(T, result, x, ydual) + result = extract_gradient!(T, result, ydual, x) return result end @@ -161,9 +163,10 @@ function chunk_mode_gradient_expr(result_definition::Expr) seed!(xdual, x, 1, seeds) seed_zero_partials!(xdual, x, N + 1, xlen - N) ydual = f(xdual) + ydual isa Real || throw(GRAD_ERROR) $(result_definition) - zero_unseeded!(T, result, x, ydual) - extract_gradient_chunk!(T, result, x, ydual, 1, N) + zero_unseeded!(T, result, ydual, x) + extract_gradient_chunk!(T, result, ydual, x, 1, N) seed_zero_partials!(xdual, x, 1) # do middle chunks @@ -171,14 +174,14 @@ function chunk_mode_gradient_expr(result_definition::Expr) i = ((c - 1) * N + 1) seed!(xdual, x, i, seeds) ydual = f(xdual) - extract_gradient_chunk!(T, result, x, ydual, i, N) + extract_gradient_chunk!(T, result, ydual, x, i, N) seed_zero_partials!(xdual, x, i) end # do final chunk seed!(xdual, x, lastchunkindex, seeds, lastchunksize) ydual = f(xdual) - extract_gradient_chunk!(T, result, x, ydual, lastchunkindex, lastchunksize) + extract_gradient_chunk!(T, result, ydual, x, lastchunkindex, lastchunksize) # get the value, this is a no-op unless result is a DiffResult extract_value!(T, result, ydual) diff --git a/src/jacobian.jl b/src/jacobian.jl index 752c817b..74f88635 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -96,29 +96,28 @@ 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. -function extract_jacobian!(::Type{T}, result::AbstractArray, x, ydual::AbstractArray) where {T} +# Zeroes the whole Jacobian unless every column is going to be written; the written ones are +# overwritten immediately after. In chunk mode the sweep calls this once up front, since the columns +# that no chunk writes belong to none of them in particular. +function zero_unseeded_columns!(::Type{T}, out::AbstractArray, ydual, x) where {T} + structural_length(x) == length(x) || fill!(out, zero(valtype(T, eltype(ydual)))) + return out +end + +# Vector mode is a single chunk that covers every seeded entry of `x`, so it extracts like the sweep. +function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, x::AbstractArray) where {T} out_reshaped = reshape_jacobian(result, ydual, x) - ydual_reshaped = vec(ydual) - # Use closure to avoid GPU broadcasting with Type - partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - n = structural_length(x) - if n == length(x) - out_reshaped .= partials_wrap.(ydual_reshaped, transpose(1:n)) - else - fill!(out_reshaped, zero(valtype(T, eltype(ydual)))) - for (i, col) in zip(1:n, structural_linearindices(x)) - out_reshaped[:, col] .= partials_wrap.(ydual_reshaped, i) - end - end + zero_unseeded_columns!(T, out_reshaped, ydual, x) + extract_jacobian_chunk!(T, out_reshaped, ydual, x, 1, structural_length(x)) return result end -function extract_jacobian!(::Type{T}, result::MutableDiffResult, x, ydual::AbstractArray) where {T} - extract_jacobian!(T, DiffResults.jacobian(result), x, ydual) +function extract_jacobian!(::Type{T}, result::MutableDiffResult, ydual::AbstractArray, x::AbstractArray) where {T} + extract_jacobian!(T, DiffResults.jacobian(result), ydual, x) return result end -function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize) where {T} +function extract_jacobian_chunk!(::Type{T}, result, ydual, x::AbstractArray, index, chunksize) where {T} ydual_reshaped = vec(ydual) offset = index - 1 irange = 1:chunksize @@ -127,19 +126,17 @@ function extract_jacobian_chunk!(::Type{T}, result, x, ydual, index, chunksize) if structural_length(x) == length(x) result[:, irange .+ offset] .= partials_wrap.(ydual_reshaped, transpose(irange)) else - idxs = Iterators.drop(structural_linearindices(x), offset) - for (i, col) in zip(irange, idxs) + cols = Iterators.drop(structural_columns(result, x), offset) + for (i, col) in zip(irange, cols) result[:, col] .= partials_wrap.(ydual_reshaped, i) end end return result end -# A matrix is used as is: reshaping it would allocate a wrapper on Julia >= 1.11, where `reshape` -# can no longer return its argument. The size is checked instead, as `reshape` did on the way past. function reshape_jacobian(result::AbstractMatrix, ydual, x) size(result) == (length(ydual), length(x)) || throw(DimensionMismatch( - lazy"cannot store the $(length(ydual))x$(length(x)) Jacobian in a result of size $(size(result))")) + 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)) @@ -153,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, x, ydual) + extract_jacobian!(T, result, ydual, x) extract_value!(T, result, ydual) return result end @@ -161,14 +158,14 @@ 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, x, ydual) + extract_jacobian!(T, result, ydual, x) map!(d -> value(T,d), y, ydual) return result end function vector_mode_jacobian!(result, f::F, x, cfg::JacobianConfig{T}) where {F,T} ydual = vector_mode_dual_eval!(f, cfg, x) - extract_jacobian!(T, result, x, ydual) + extract_jacobian!(T, result, ydual, x) extract_value!(T, result, ydual) return result end @@ -176,7 +173,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, x, ydual) + extract_jacobian!(T, result, ydual, x) extract_value!(T, result, y, ydual) return result end @@ -212,9 +209,8 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex ydual isa AbstractArray || throw(JACOBIAN_ERROR) $(result_definition) out_reshaped = reshape_jacobian(result, ydual, x) - # zero the columns of the structurally zero entries of `x`, which no chunk of the sweep writes - structural_length(x) == length(x) || fill!(out_reshaped, zero(valtype(T, eltype(ydual)))) - extract_jacobian_chunk!(T, out_reshaped, x, ydual, 1, N) + zero_unseeded_columns!(T, out_reshaped, ydual, x) + extract_jacobian_chunk!(T, out_reshaped, ydual, x, 1, N) seed_zero_partials!(xdual, x, 1) # do middle chunks @@ -222,14 +218,14 @@ function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Ex i = ((c - 1) * N + 1) seed!(xdual, x, i, seeds) $(compute_ydual) - extract_jacobian_chunk!(T, out_reshaped, x, ydual, i, N) + extract_jacobian_chunk!(T, out_reshaped, ydual, x, i, N) seed_zero_partials!(xdual, x, i) end # do final chunk seed!(xdual, x, lastchunkindex, seeds, lastchunksize) $(compute_ydual) - extract_jacobian_chunk!(T, out_reshaped, x, ydual, lastchunkindex, lastchunksize) + extract_jacobian_chunk!(T, out_reshaped, ydual, x, lastchunkindex, lastchunksize) $(y_definition) From 3063bff9bb20ac140b2e289e7c23f16fa07a81f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 18 Aug 2026 17:23:19 +0200 Subject: [PATCH 06/10] Cover the wrongly shaped results, the partial final chunks and the static buffer A wrongly shaped `jacobian!` result and a `gradient!` result that is not shaped like `x` both throw now, the latter for a dense `x` too, and neither was asserted. Nor was the `MVector` gradient buffer of an `ImmutableDiffResult`, which chunk mode writes entry by entry. The structured Jacobian and Hessian testsets only ran chunk sizes that divide the number of seeded entries, so the final chunk was always a full one; they gain a size that leaves a remainder. `SeedTest` pins the columns `structural_columns` returns against the index sets it already writes out by hand, since the Jacobian misplaces its derivatives if they disagree with `structural_eachindex`. The allocation test measures a dense result as well, which is the case that zeroes before extracting, and the JET tests cover a structured input, which reaches the generator branches. Co-Authored-By: Claude Opus 5 (1M context) --- test/AllocationsTest.jl | 30 +++++++++++++++--------------- test/GradientTest.jl | 24 ++++++++++++++++++++++++ test/HessianTest.jl | 3 ++- test/JacobianTest.jl | 21 ++++++++++++++++++++- test/QATest.jl | 9 +++++++++ test/SeedTest.jl | 4 ++++ 6 files changed, 74 insertions(+), 17 deletions(-) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 53eb5854..4d3f4d92 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -53,9 +53,9 @@ end # `extract_gradient!`/`extract_jacobian!` take their positions from `x`, so mapping a structural # position to an index of `x` must not allocate, whether or not `x` has structurally zero entries. -function allocs_structured_gradient!(x, chunk) +function allocs_structured_gradient!(result, x, chunk) f(z) = sum(abs2, z) - result = fill!(similar(x, Float64), false) + fill!(result, false) cfg = ForwardDiff.GradientConfig(f, x, chunk) ForwardDiff.gradient!(result, f, x, cfg) # warmup return @allocated ForwardDiff.gradient!(result, f, x, cfg) @@ -70,19 +70,19 @@ function allocs_structured_jacobian!(x, chunk) return @allocated ForwardDiff.jacobian!(result, f!, y, x, cfg) end -@testset "Test gradient!/jacobian! allocations for $(nameof(typeof(x)))" for x in (rand(6, 6), - LowerTriangular(rand(6, 6)), - UpperTriangular(rand(6, 6)), - Diagonal(rand(6, 6))) - # vector mode - chunk = ForwardDiff.Chunk{ForwardDiff.structural_length(x)}() - @test iszero(allocs_structured_gradient!(x, chunk)) - @test iszero(allocs_structured_jacobian!(x, chunk)) - - # chunk mode - chunk = ForwardDiff.Chunk{2}() - @test iszero(allocs_structured_gradient!(x, chunk)) - @test iszero(allocs_structured_jacobian!(x, chunk)) +@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 diff --git a/test/GradientTest.jl b/test/GradientTest.jl index 8f618fe5..f80468da 100644 --- a/test/GradientTest.jl +++ b/test/GradientTest.jl @@ -331,6 +331,30 @@ end end end +@testset "result not shaped like x" begin + # The extraction positions are indices of `x`, which a result of a different shape cannot be + # indexed by, dense `x` included. + x = randn(4) + f = z -> dot(z, z) + @testset "chunk size = $c" for c in (2, 4) + cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.gradient!(fill(NaN, 5), f, x, cfg) + result = DiffResults.DiffResult(NaN, fill(NaN, 5)) + @test_throws DimensionMismatch ForwardDiff.gradient!(result, f, x, cfg) + end +end + +@testset "mutable gradient buffer in an immutable result" begin + # A `StaticArray` buffer makes the result immutable, but an `MVector` can still be written to + # entry by entry, which is how the chunk mode sweep fills it. + x = randn(6) + f = z -> dot(z, z) + result = DiffResults.DiffResult(NaN, @MVector fill(NaN, 6)) + @test result isa DiffResults.ImmutableDiffResult + ForwardDiff.gradient!(result, f, x, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{2}())) + @test DiffResults.gradient(result) ≈ 2 .* x +end + # issue #769 @testset "functions with `Dual` output" begin x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))] diff --git a/test/HessianTest.jl b/test/HessianTest.jl index 01d29143..28fd03ec 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -177,7 +177,8 @@ end grad = zeros(3, 3) grad[sidx] .= x[sidx] .+ sum(x) - @testset "chunk size = $c" for c in unique((1, 2, length(sidx))) + # 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) diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index 0117d844..5d2b901a 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -340,7 +340,8 @@ end expected[2, sidx] .= 2 .* x[sidx] val = g(x) - @testset "chunk size = $c" for c in unique((1, 2, length(sidx))) + # `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)) @@ -379,6 +380,24 @@ end end end +@testset "wrongly shaped result" begin + # A matrix result is used as is, so it has to have the shape of the Jacobian and not merely as + # many entries. Results of other shapes are reshaped and only have to match in length. + x = randn(4) + g = z -> [sum(z), sum(abs2, z)] + g! = (y, z) -> (y[1] = sum(z); y[2] = sum(abs2, z); y) + @testset "chunk size = $c" for c in (2, 4) + cfg = ForwardDiff.JacobianConfig(g, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.jacobian!(fill(NaN, 4, 2), g, x, cfg) + result = DiffResults.DiffResult(fill(NaN, 2), fill(NaN, 4, 2)) + @test_throws DimensionMismatch ForwardDiff.jacobian!(result, g, x, cfg) + + y = fill(NaN, 2) + cfg! = ForwardDiff.JacobianConfig(g!, y, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.jacobian!(fill(NaN, 4, 2), g!, y, x, cfg!) + end +end + # issue #769 @testset "functions with `Dual` output" begin x = [Dual{OuterTestTag}(Dual{TestTag}(1.3, 2.1), Dual{TestTag}(0.3, -2.4))] diff --git a/test/QATest.jl b/test/QATest.jl index 860ccdb0..925bbfb8 100644 --- a/test/QATest.jl +++ b/test/QATest.jl @@ -1,6 +1,7 @@ module QATest using ForwardDiff +using LinearAlgebra using Test using JET: @test_opt @@ -11,6 +12,14 @@ using JET: @test_opt @test_opt ForwardDiff.gradient(only, [1.0], ForwardDiff.GradientConfig(only, [1.0], ForwardDiff.Chunk{1}())) @test_opt ForwardDiff.jacobian(identity, [1.0], ForwardDiff.JacobianConfig(identity, [1.0], ForwardDiff.Chunk{1}())) @test_opt ForwardDiff.hessian(only, [1.0], ForwardDiff.HessianConfig(only, [1.0], ForwardDiff.Chunk{1}())) + + # extraction iterates the structural positions of `x` for these + @testset "$(nameof(typeof(x)))" for x in (LowerTriangular(rand(3, 3)), + UpperTriangular(rand(3, 3)), + Diagonal(rand(3, 3))) + @test_opt ForwardDiff.gradient(first, x, ForwardDiff.GradientConfig(first, x, ForwardDiff.Chunk{2}())) + @test_opt ForwardDiff.jacobian(vec, x, ForwardDiff.JacobianConfig(vec, x, ForwardDiff.Chunk{2}())) + end end end # module diff --git a/test/SeedTest.jl b/test/SeedTest.jl index 02b821c3..a6b3944a 100644 --- a/test/SeedTest.jl +++ b/test/SeedTest.jl @@ -53,6 +53,10 @@ end @test collect(ForwardDiff.structural_eachindex(duals, x)) == sidx @test ForwardDiff.structural_length(x) == nstruct + # the columns a Jacobian receives derivatives in, in the order the seeds are laid out in + @test collect(ForwardDiff.structural_columns(zeros(2, length(x)), x)) == LinearIndices(x)[sidx] + @test_throws DimensionMismatch ForwardDiff.structural_columns(zeros(2, length(x) + 1), x) + # `count` defaults to N fill_marker!(duals, x, sidx, marker) ForwardDiff.seed_zero_partials!(duals, x, 4) From 3c74cfeb41be6db4cd2f59b193d9aed5f026af99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 18 Aug 2026 22:40:11 +0200 Subject: [PATCH 07/10] Index the Hessian by the linear indices of `x` Both axes of a Hessian belong to the entries of `x`, the way the columns of a Jacobian do, so `structural_columns` supplies the positions for both and the result is `length(x)` by `length(x)` with hard zeros in the rows and columns of the structurally zero entries. Without this the sweep numbered its rows and columns by seeding position, which is a shape and an order that only ForwardDiff knows, and which #838/#839 had just removed from `gradient` and `jacobian`. The sweep indexes into the positions of two blocks at once and re-reads each row block once per column block, so unlike `jacobian!` it cannot walk them lazily and materializes them once per call instead. For everything but the triangular wrappers they are a range already and `_indexable` is a no-op. Co-Authored-By: Claude Opus 5 --- ext/ForwardDiffStaticArraysExt.jl | 4 ++- src/apiutils.jl | 8 +++++ src/hessian.jl | 51 +++++++++++++++++++++++-------- 3 files changed, 49 insertions(+), 14 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index f571d3ea..80ab009d 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -132,7 +132,9 @@ ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig, ::Val) where {F} = ydual = f(dualize(T, dualize(T, x))) ydual isa Real || throw(HESSIAN_ERROR) H = result isa AbstractMatrix ? result : reshape(result, length(x), length(x)) - ForwardDiff.extract_hessian_chunk!(T, H, ydual, 0, 0, length(x), length(x)) + # a `StaticArray` has no structurally zero entries, so the positions are the columns of `H` + positions = ForwardDiff.structural_columns(H, x) + ForwardDiff.extract_hessian_chunk!(T, H, positions, ydual, 0, 0, length(x), length(x)) return result end diff --git a/src/apiutils.jl b/src/apiutils.jl index 7ac1cfc3..16253dfe 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -104,6 +104,14 @@ function check_matching_columns(out::AbstractMatrix, x::AbstractArray) return nothing end +# `structural_columns` is lazy, since walking it once from the front is all `jacobian!` ever needs. +# The Hessian sweep reads the positions of two blocks at once and re-reads each row block once per +# column block, so it indexes into them instead, and materializes them first. The ranges of +# unstructured inputs and of `Diagonal` are already indexable and pass through unchanged; only the +# triangular wrappers pay, one `structural_length(x)`-element vector against an n²-entry result. +_indexable(idxs::AbstractArray) = idxs +_indexable(idxs) = collect(idxs) + # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is # currently carrying and to initialize a freshly allocated work buffer, whose elements must all be # written before the target function reads them. diff --git a/src/hessian.jl b/src/hessian.jl index 61238e6d..1cd143c7 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -9,6 +9,14 @@ Return `H(f)` evaluated at `x`, assuming `f` is called as `f(x)`. The returned Hessian is exactly symmetric: its two triangles are filled from the same derivative values. +Both axes of the result are indexed by the linear indices of `x`, so it is +`length(x)` by `length(x)`. For an `x` with structurally zero entries, such as a +`LowerTriangular`, `UpperTriangular` or `Diagonal` matrix, only the entries that are not +structurally zero are differentiated; the rows and columns of the others are zero. Note +that this makes the result of `hessian(f, ::Diagonal)` quadratic in `length(x)` and hence +quartic in the size of the diagonal — differentiate with respect to the diagonal vector +instead if that matters. + This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. @@ -27,6 +35,10 @@ Compute `H(f)` evaluated at `x` and store the result(s) in `result`, assuming `f called as `f(x)`. The stored Hessian is exactly symmetric: its two triangles are filled from the same derivative values. +`result` has to hold `length(x)^2` entries, indexed as described for +`ForwardDiff.hessian`; a matrix `result` is written to as is and hence has to be +`length(x)` by `length(x)`. + This method assumes that `isa(f(x), Real)`. Set `check` to `Val{false}()` to disable tag checking. This can lead to perturbation confusion, so should be used with care. @@ -34,8 +46,8 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(result, x) CHK && checktag(T, f, x) - xlen = structural_length(x) - H = result isa AbstractMatrix ? result : reshape(result, xlen, xlen) + hlen = length(x) + H = result isa AbstractMatrix ? result : reshape(result, hlen, hlen) symmetric_hessian!(H, f, x, cfg, nothing) return result end @@ -53,9 +65,9 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) - xlen = structural_length(x) + hlen = length(x) hess = DiffResults.hessian(result) - H = hess isa AbstractMatrix ? hess : reshape(hess, xlen, xlen) + H = hess isa AbstractMatrix ? hess : reshape(hess, hlen, hlen) _, ydual = symmetric_hessian!(H, f, x, cfg, DiffResults.gradient(result)) result = DiffResults.value!(result, value(T, value(T, ydual))) return result @@ -68,15 +80,18 @@ end const HESSIAN_ERROR = DimensionMismatch("hessian(f, x) expects that f(x) is a real number. Perhaps you meant jacobian(f, x)?") # Copy a block from the nested partials and fill its transpose. On diagonal blocks, read -# only the upper triangle so the result is exactly symmetric. -function extract_hessian_chunk!(::Type{T}, H, ydual, roffset, coffset, rsize, csize) where {T} +# only the upper triangle so the result is exactly symmetric. `positions` maps a seeding position to +# the row and column of `H` it belongs to, i.e. to the linear index of `x` it was seeded from. +function extract_hessian_chunk!(::Type{T}, H, positions, ydual, roffset, coffset, rsize, csize) where {T} for r in 1:rsize + i = positions[roffset + r] drow = partials(T, ydual, r) cstart = roffset == coffset ? r : 1 for c in cstart:csize + j = positions[coffset + c] h = partials(T, drow, c) - H[roffset + r, coffset + c] = h - H[coffset + c, roffset + r] = h + H[i, j] = h + H[j, i] = h end end return H @@ -92,6 +107,9 @@ extract_hessian_gradient_chunk!(::Type{T}, grad, ydual, x, index, chunksize) whe function symmetric_hessian_expr(result_definition::Expr) return quote xlen = structural_length(x) + # Only the structurally non-zero entries of `x` are seeded, but both axes of the result are + # indexed by the linear indices of `x`, as the columns of a Jacobian are. See #839. + hlen = length(x) if xlen < N throw(ArgumentError(lazy"chunk size cannot be greater than ForwardDiff.structural_length(x) ($(N) > $(structural_length(x)))")) end @@ -111,9 +129,16 @@ function symmetric_hessian_expr(result_definition::Expr) ydual1 = f(xdual) ydual1 isa Real || throw(HESSIAN_ERROR) $(result_definition) - # the entries that no chunk of the sweep writes belong to none of them in particular + # `H` is square and both its axes are indexed by the entries of `x`, so the columns a + # Jacobian would receive derivatives in are also the rows and columns this sweep writes. + positions = _indexable(structural_columns(H, x)) + # The entries that no block of the sweep writes belong to none of them in particular: for + # `H` the rows and columns of the structurally zero entries of `x`, for `grad` their + # entries. `value(T, ydual1)` is passed for its value type, which unlike `eltype` is a + # number type even for a result that stores `Any`. + zero_unseeded_columns!(T, H, value(T, ydual1), x) grad === nothing || zero_unseeded!(T, grad, value(T, ydual1), x) - extract_hessian_chunk!(T, H, ydual1, 0, 0, N, N) + extract_hessian_chunk!(T, H, positions, ydual1, 0, 0, N, N) extract_hessian_gradient_chunk!(T, grad, ydual1, x, 1, N) nblocks > 1 && seed_hessian_chunk!(xdual, x, 1, nothing, nothing) @@ -127,13 +152,13 @@ function symmetric_hessian_expr(result_definition::Expr) poffset = (p - 1) * N seed_hessian_chunk!(xdual, x, poffset + 1, iseeds, nothing) ydual = f(xdual) - extract_hessian_chunk!(T, H, ydual, qoffset, poffset, qsize, N) + extract_hessian_chunk!(T, H, positions, ydual, qoffset, poffset, qsize, N) seed_hessian_chunk!(xdual, x, poffset + 1, nothing, nothing) end # The diagonal block adds q's inner seeds while retaining its outer seeds. seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, oseeds, qsize) ydual = f(xdual) - extract_hessian_chunk!(T, H, ydual, qoffset, qoffset, qsize, qsize) + extract_hessian_chunk!(T, H, positions, ydual, qoffset, qoffset, qsize, qsize) extract_hessian_gradient_chunk!(T, grad, ydual, x, qoffset + 1, qsize) seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, nothing, qsize) end @@ -143,7 +168,7 @@ function symmetric_hessian_expr(result_definition::Expr) end @eval function symmetric_hessian(f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N} - $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), xlen, xlen)))) + $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), hlen, hlen)))) end @eval function symmetric_hessian!(H, f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N} From 01c6cb7cad21a2f4449d53d6973ddb5c12c48116 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 18 Aug 2026 22:53:59 +0200 Subject: [PATCH 08/10] Restore the result shape check dropped with `extract_jacobian!` Writing the result entry by entry lost the validation the broadcast used to do on the way past, so a matrix of the wrong size was partially filled and returned instead of rejected: `hessian!(fill(NaN, 4, 4), f, rand(3))` came back with a row and a column of `NaN`. `reshape_hessian` mirrors #840's `reshape_jacobian`, down to short-circuiting on a matrix rather than letting `reshape` allocate a wrapper for it on 1.11 and later. Co-Authored-By: Claude Opus 5 --- ext/ForwardDiffStaticArraysExt.jl | 2 +- src/hessian.jl | 15 ++++++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 80ab009d..46f8c519 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -131,7 +131,7 @@ ForwardDiff.hessian(f::F, x::StaticArray, cfg::HessianConfig, ::Val) where {F} = T = typeof(Tag(f, eltype(x))) ydual = f(dualize(T, dualize(T, x))) ydual isa Real || throw(HESSIAN_ERROR) - H = result isa AbstractMatrix ? result : reshape(result, length(x), length(x)) + H = ForwardDiff.reshape_hessian(result, x) # a `StaticArray` has no structurally zero entries, so the positions are the columns of `H` positions = ForwardDiff.structural_columns(H, x) ForwardDiff.extract_hessian_chunk!(T, H, positions, ydual, 0, 0, length(x), length(x)) diff --git a/src/hessian.jl b/src/hessian.jl index 1cd143c7..a0e634dc 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -46,8 +46,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function hessian!(result::AbstractArray, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(result, x) CHK && checktag(T, f, x) - hlen = length(x) - H = result isa AbstractMatrix ? result : reshape(result, hlen, hlen) + H = reshape_hessian(result, x) symmetric_hessian!(H, f, x, cfg, nothing) return result end @@ -65,9 +64,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba function hessian!(result::DiffResult, f::F, x::AbstractArray, cfg::HessianConfig{T} = HessianConfig(f, result, x), ::Val{CHK}=Val{true}()) where {F,T,CHK} require_one_based_indexing(x) CHK && checktag(T, f, x) - hlen = length(x) - hess = DiffResults.hessian(result) - H = hess isa AbstractMatrix ? hess : reshape(hess, hlen, hlen) + H = reshape_hessian(result, x) _, ydual = symmetric_hessian!(H, f, x, cfg, DiffResults.gradient(result)) result = DiffResults.value!(result, value(T, value(T, ydual))) return result @@ -102,6 +99,14 @@ extract_hessian_gradient_chunk!(::Type{T}, ::Nothing, ydual, x, index, chunksize extract_hessian_gradient_chunk!(::Type{T}, grad, ydual, x, index, chunksize) where {T} = extract_gradient_chunk!(T, grad, value(T, ydual), x, index, chunksize) +function reshape_hessian(result::AbstractMatrix, x) + size(result) == (length(x), length(x)) || throw(DimensionMismatch( + lazy"cannot store the $(length(x))×$(length(x)) Hessian in a result of size $(size(result))")) + return result +end +reshape_hessian(result::AbstractArray, x) = reshape(result, length(x), length(x)) +reshape_hessian(result::DiffResult, x) = reshape_hessian(DiffResults.hessian(result), x) + # Evaluate one pair of chunks at a time using nested duals. Only one triangle of block # pairs is evaluated; the other is filled by symmetry (see #836). function symmetric_hessian_expr(result_definition::Expr) From f567163f07ac95f9b9b96d54c813475ab8141497 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 19 Aug 2026 00:01:13 +0200 Subject: [PATCH 09/10] Read the same triangle at every chunk size MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The diagonal blocks read the nested partials at `outer <= inner` and the off-diagonal ones at `outer > inner`, which are equal in exact arithmetic and not in floating point, so the Hessian depended on the chunk size in the last bit: at n = 16, every chunk size below 16 disagreed with vector mode and with the `StaticArray` path, which reads `outer <= inner` as well. Swapping which layer the column block seeds costs nothing — the same evaluations, the same seed writes — and makes all of them bitwise equal. `seed_hessian_chunk!` built the zero partials of both layers whether or not it needed them, which an isbits `V` optimizes away and `BigFloat` does not: 288 bytes per call at N = 3, now zero when both seeds are supplied. Co-Authored-By: Claude Opus 5 --- ext/ForwardDiffStaticArraysExt.jl | 2 + src/apiutils.jl | 14 ++++--- src/hessian.jl | 13 ++++--- test/AllocationsTest.jl | 32 +++++++++++++-- test/HessianTest.jl | 65 +++++++++++++++++++++++++++++++ test/QATest.jl | 2 + test/SeedTest.jl | 19 +++++++++ 7 files changed, 132 insertions(+), 15 deletions(-) diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 46f8c519..efb9aea6 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -112,6 +112,8 @@ end return typeof(H)(Symmetric(H, :U)) end +# `partials` is empty both for an empty `x` and for an `f` that ignores its argument, and neither +# reaches `extract_jacobian`, whose generated method would build a matrix with no columns to fill. @inline function extract_hessian(::Type{T}, ydual::Partials{0}, x::S) where {T,S<:StaticArray} R = StaticArrays.similar_type(S, valtype(T, eltype(ydual)), Size(length(x), length(x))) return zero(R) diff --git a/src/apiutils.jl b/src/apiutils.jl index 16253dfe..661c3483 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -171,16 +171,20 @@ function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index, end end -# Seed a chunk in either layer of nested duals. A `nothing` seed clears that layer. +# Seed a chunk in either layer of nested duals. A `nothing` seed clears that layer. This is not +# `seed_zero_partials!` with both layers cleared: that builds the inner value at the exact type of +# the array's `V`, which the nested buffer's `Dual{T,V,N}` is not, so it would be a `MethodError`. function seed_hessian_chunk!(duals::AbstractArray{Dual{T,Dual{T,V,N},N}}, x, index, iseeds::Union{Nothing,NTuple{N,Partials{N,V}}}, oseeds::Union{Nothing,NTuple{N,Partials{N,Dual{T,V,N}}}}, chunksize = N) where {T,V,N} - izero = zero(Partials{N,V}) - ozero = zero(Partials{N,Dual{T,V,N}}) + # `iseeds === nothing` is a compile-time constant here, so building each layer's accessor in its + # own branch means the zero partials it does not need are never constructed. That is free for an + # isbits `V` but not, say, for `BigFloat`. + ipartials = iseeds === nothing ? Returns(zero(Partials{N,V})) : Base.Fix1(getindex, iseeds) + opartials = oseeds === nothing ? Returns(zero(Partials{N,Dual{T,V,N}})) : Base.Fix1(getindex, oseeds) idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), chunksize) return _seed!(duals, x, idxs) do value, i - inner = Dual{T,V,N}(value, iseeds === nothing ? izero : iseeds[i]) - Dual{T,Dual{T,V,N},N}(inner, oseeds === nothing ? ozero : oseeds[i]) + Dual{T,Dual{T,V,N},N}(Dual{T,V,N}(value, ipartials(i)), opartials(i)) end end diff --git a/src/hessian.jl b/src/hessian.jl index a0e634dc..d3dcfe54 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -150,17 +150,18 @@ function symmetric_hessian_expr(result_definition::Expr) for q in 2:nblocks qoffset = (q - 1) * N qsize = min(N, xlen - qoffset) - # Off-diagonal blocks: p seeds columns and q seeds rows. The outer seeds for q - # remain unchanged throughout this loop. - seed_hessian_chunk!(xdual, x, qoffset + 1, nothing, oseeds, qsize) + # Off-diagonal blocks: q seeds columns and p seeds rows. Every block of the sweep is + # read from the same triangle of the nested partials, so that the result does not + # depend on the chunk size. The inner seeds for q remain unchanged throughout this loop. + seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, nothing, qsize) for p in 1:(q - 1) poffset = (p - 1) * N - seed_hessian_chunk!(xdual, x, poffset + 1, iseeds, nothing) + seed_hessian_chunk!(xdual, x, poffset + 1, nothing, oseeds) ydual = f(xdual) - extract_hessian_chunk!(T, H, positions, ydual, qoffset, poffset, qsize, N) + extract_hessian_chunk!(T, H, positions, ydual, poffset, qoffset, N, qsize) seed_hessian_chunk!(xdual, x, poffset + 1, nothing, nothing) end - # The diagonal block adds q's inner seeds while retaining its outer seeds. + # The diagonal block adds q's outer seeds while retaining its inner seeds. seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, oseeds, qsize) ydual = f(xdual) extract_hessian_chunk!(T, H, positions, ydual, qoffset, qoffset, qsize, qsize) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index c317c34d..cf15f3e9 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -8,7 +8,7 @@ include(joinpath(dirname(@__FILE__), "utils.jl")) convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,Float64,8},4},2}, 1.3) -@testset "Test seed!/seed_zero_partials! allocations" begin +@testset "Test seed!/seed_zero_partials!/seed_hessian_chunk! allocations" begin x = rand(1000) cfg = ForwardDiff.GradientConfig(nothing, x) duals = cfg.duals @@ -35,8 +35,11 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F iseeds = hcfg.jacobian_config.seeds oseeds = hcfg.gradient_config.seeds allocs_hseed!(args...) = @allocated ForwardDiff.seed_hessian_chunk!(args...) - allocs_hseed!(hduals, x, 1, iseeds, oseeds) - @test iszero(allocs_hseed!(hduals, x, 1, iseeds, oseeds)) + # all four seed combinations, since each builds the zero partials of a different set of layers + for (is, os) in ((iseeds, oseeds), (iseeds, nothing), (nothing, oseeds), (nothing, nothing)) + allocs_hseed!(hduals, x, 1, is, os) + @test iszero(allocs_hseed!(hduals, x, 1, is, os)) + end allocs_hseed!(hduals, x, 1, nothing, nothing, 4) @test iszero(allocs_hseed!(hduals, x, 1, nothing, nothing, 4)) @@ -80,7 +83,20 @@ function allocs_structured_jacobian!(x, chunk) return @allocated ForwardDiff.jacobian!(result, f!, y, x, cfg) end -@testset "Test gradient!/jacobian! allocations for $(nameof(typeof(x)))" for (x, nstruct) in ( +# Unlike `gradient!`/`jacobian!`, which walk the structural positions once from the front, the +# Hessian sweep indexes into them, so it materializes them first and cannot be allocation-free for +# the inputs whose positions are lazy. Pin what that costs rather than dropping the check. +function allocs_structured_hessian!(H, x, chunk) + f(z) = sum(abs2, z) + fill!(H, false) + cfg = ForwardDiff.HessianConfig(f, x, chunk) + ForwardDiff.hessian!(H, f, x, cfg) # warmup + return @allocated ForwardDiff.hessian!(H, f, x, cfg) +end + +allocs_positions(H, x) = @allocated ForwardDiff._indexable(ForwardDiff.structural_columns(H, x)) + +@testset "Test gradient!/jacobian!/hessian! allocations for $(nameof(typeof(x)))" for (x, nstruct) in ( (rand(6, 6), 36), (LowerTriangular(rand(6, 6)), 21), (UpperTriangular(rand(6, 6)), 21), @@ -93,6 +109,14 @@ end @test iszero(allocs_structured_gradient!(result, x, chunk)) @test iszero(allocs_structured_jacobian!(x, chunk)) end + + # zero for the inputs whose positions are a range already, and otherwise exactly the one vector + # they are collected into: nothing else in the sweep may allocate + H = zeros(length(x), length(x)) + allocs_positions(H, x) # warmup + for chunk_size in (2, nstruct) + @test allocs_structured_hessian!(H, x, ForwardDiff.Chunk{chunk_size}()) == allocs_positions(H, x) + end end @testset "allocation-free nested StaticArray jacobian" begin diff --git a/test/HessianTest.jl b/test/HessianTest.jl index ab42ad53..3df3a8ba 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -211,6 +211,16 @@ end @test ForwardDiff.hessian!(out, f, x, cfg) === out @test out == expected + # a result that is not a matrix is reshaped, so it only has to match in length + flat = fill(NaN, L^2) + @test ForwardDiff.hessian!(flat, f, x, cfg) === flat + @test reshape(flat, L, L) == expected + + # a wrongly shaped result is rejected rather than partially filled, the structural shape + # this used to have included + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, L + 1, L + 1), f, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, length(sidx), length(sidx)), f, x, cfg) + # `DiffResults.HessianResult` allocates a dense gradient buffer even for a structured `x` result = DiffResults.HessianResult(x) result = ForwardDiff.hessian!(result, f, x, @@ -221,6 +231,61 @@ end end end +@testset "chunk size independence" begin + # Every block of the sweep reads the same triangle of the nested partials, so the result is not + # merely close across chunk sizes but bitwise equal, and equal to the vector mode of the + # `StaticArray` path, which reads that triangle too. + n = 16 + f = z -> sum(sin(z[i]) * exp(z[mod1(i + 1, n)]) / (1 + z[mod1(i + 2, n)]^2) for i in eachindex(z)) + x = collect(range(0.1, 0.9; length=n)) + reference = ForwardDiff.hessian(f, SVector{n}(x)) + @test reference == transpose(reference) + @testset "chunk size = $c" for c in (1, 2, 3, 5, 7, 11, n) + H = ForwardDiff.hessian(f, x, ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}())) + @test H == reference + end +end + +@testset "wrongly shaped result" begin + # A matrix result is written to as is, so it has to have the shape of the Hessian and not merely + # as many entries. Results of other shapes are reshaped and only have to match in length. + x = randn(3) + f = z -> sum(abs2, z) + sx = SVector{3}(x) + @testset "chunk size = $c" for c in (2, 3) + cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 4, 4), f, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 10), f, x, cfg) + result = DiffResults.DiffResult(NaN, fill(NaN, 3), fill(NaN, 4, 4)) + @test_throws DimensionMismatch ForwardDiff.hessian!(result, f, x, cfg) + end + # the `StaticArray` methods bypass the sweep and reshape the result themselves + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 4, 4), f, sx) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 10), f, sx) + # a result that is not a matrix goes through `reshape` + flat = fill(NaN, 9) + @test ForwardDiff.hessian!(flat, f, sx) === flat + @test reshape(flat, 3, 3) == 2I(3) +end + +@testset "empty input" begin + f = z -> 1.0 + sum(z) + @test ForwardDiff.hessian(f, Float64[]) == zeros(0, 0) + @test ForwardDiff.hessian(f, SVector{0,Float64}()) == zeros(0, 0) + # `DiffResults.HessianResult` cannot be built from an empty `x`, so assemble one by hand + result = DiffResults.DiffResult(NaN, Float64[], Matrix{Float64}(undef, 0, 0)) + result = ForwardDiff.hessian!(result, f, Float64[]) + @test DiffResults.value(result) == 1.0 + @test isempty(DiffResults.hessian(result)) +end + +@testset "f(x) is not a real number" begin + x = randn(3) + sx = SVector{3}(x) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 3, 3), identity, sx) + @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.HessianResult(sx), identity, sx) +end + @testset "BigFloat with an unassigned input entry" begin x = Vector{BigFloat}(undef, 10) hole = 5 diff --git a/test/QATest.jl b/test/QATest.jl index 925bbfb8..48fc3c0d 100644 --- a/test/QATest.jl +++ b/test/QATest.jl @@ -19,6 +19,8 @@ using JET: @test_opt Diagonal(rand(3, 3))) @test_opt ForwardDiff.gradient(first, x, ForwardDiff.GradientConfig(first, x, ForwardDiff.Chunk{2}())) @test_opt ForwardDiff.jacobian(vec, x, ForwardDiff.JacobianConfig(vec, x, ForwardDiff.Chunk{2}())) + # the Hessian sweep materializes those positions, so `_indexable` has to stay concrete + @test_opt ForwardDiff.hessian(first, x, ForwardDiff.HessianConfig(first, x, ForwardDiff.Chunk{2}())) end end diff --git a/test/SeedTest.jl b/test/SeedTest.jl index 7e0c6a4a..5a7becbc 100644 --- a/test/SeedTest.jl +++ b/test/SeedTest.jl @@ -110,6 +110,25 @@ end ForwardDiff.seed_hessian_chunk!(duals, x, 4, nothing, nothing) @test all(idx -> iszero(ForwardDiff.partials(ForwardDiff.value(duals[idx]))), sidx) @test all(idx -> iszero(ForwardDiff.partials(duals[idx])), sidx) + + # The off-diagonal blocks of the sweep seed one layer at a time, so check that each seed lands + # in the layer it was asked for and that the other one is left cleared, rather than only that + # something was written. + izero = zero(eltype(iseeds)) + ozero = zero(eltype(oseeds)) + @testset "$(iseeds === nothing ? "outer" : "inner") layer only" for (is, os) in + ((iseeds, nothing), (nothing, oseeds)) + ForwardDiff.seed_hessian_chunk!(duals, x, 4, is, os) + for (i, idx) in enumerate(sidx) + chunkpos = i - 3 + inner = ForwardDiff.partials(ForwardDiff.value(duals[idx])) + outer = ForwardDiff.partials(duals[idx]) + wanted = 1 <= chunkpos <= 3 + @test inner == (wanted && is !== nothing ? is[chunkpos] : izero) + @test outer == (wanted && os !== nothing ? os[chunkpos] : ozero) + end + ForwardDiff.seed_hessian_chunk!(duals, x, 4, nothing, nothing) + end end end # module From 6279027bb4772ee29f99e0c2381dc7326adf95ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 19 Aug 2026 01:52:12 +0200 Subject: [PATCH 10/10] Walk the structural positions of the Hessian sweep The sweep needs no random access into the positions its blocks are seeded at. For a fixed column block the row blocks are walked in order, and within a block each row resumes the column walk from the state saved at the start of the block, so the walk costs what `seed_hessian_chunk!` already pays to seed the same block and no more. A diagonal block reads its columns from wherever its rows have got to, which is exactly the upper triangle it wants. So the lazy positions of the triangular wrappers never have to be materialized, and `hessian!` is allocation-free for every input kind, as `gradient!` and `jacobian!` are. Three comments, while nearby: - Reading every block at the same nesting order is what fixes the entries of `H` across chunk sizes, but it does not fix them for an `f` that branches on the seeds: `<` on `Dual` breaks ties on the partials, so the two layers of `max(x[1], x[2])` at an exact tie resolve to different arguments. The off-diagonal comment names the mechanism rather than the consequence. - `zero_unseeded_columns!` reads `ydual` for its value type, so the sweep, which zeroes rows as well as columns, can hand it a scalar dual where a Jacobian hands it an array of them. Say so at both ends. - `length(x)` belongs to the one generated body that allocates a result. Tests: - The zero partials of a `Float64` `x` are isbits and free to build, so pin the seed combination that needs neither layer's on a `BigFloat` one, where building them costs 288 bytes at N = 3. - The structured testset's `f` has an exactly representable Hessian, so its `==` says nothing about floating point. Cover the three wrappers in the chunk size independence testset, against their largest chunk. - Cover the sweep's own `HESSIAN_ERROR`, an `f` that ignores a non-empty `x`, and a `DiffResult` holding a Hessian that is not a matrix. Co-Authored-By: Claude Opus 5 (1M context) --- src/apiutils.jl | 8 -------- src/hessian.jl | 44 +++++++++++++++++++++++++++-------------- src/jacobian.jl | 4 +++- test/AllocationsTest.jl | 24 +++++++++++++--------- test/HessianTest.jl | 44 ++++++++++++++++++++++++++++++++++++++--- test/QATest.jl | 2 +- 6 files changed, 89 insertions(+), 37 deletions(-) diff --git a/src/apiutils.jl b/src/apiutils.jl index 661c3483..7ab61f21 100644 --- a/src/apiutils.jl +++ b/src/apiutils.jl @@ -104,14 +104,6 @@ function check_matching_columns(out::AbstractMatrix, x::AbstractArray) return nothing end -# `structural_columns` is lazy, since walking it once from the front is all `jacobian!` ever needs. -# The Hessian sweep reads the positions of two blocks at once and re-reads each row block once per -# column block, so it indexes into them instead, and materializes them first. The ranges of -# unstructured inputs and of `Diagonal` are already indexable and pass through unchanged; only the -# triangular wrappers pay, one `structural_length(x)`-element vector against an n²-entry result. -_indexable(idxs::AbstractArray) = idxs -_indexable(idxs) = collect(idxs) - # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is # currently carrying and to initialize a freshly allocated work buffer, whose elements must all be # written before the target function reads them. diff --git a/src/hessian.jl b/src/hessian.jl index d3dcfe54..08894bfc 100644 --- a/src/hessian.jl +++ b/src/hessian.jl @@ -79,17 +79,31 @@ const HESSIAN_ERROR = DimensionMismatch("hessian(f, x) expects that f(x) is a re # Copy a block from the nested partials and fill its transpose. On diagonal blocks, read # only the upper triangle so the result is exactly symmetric. `positions` maps a seeding position to # the row and column of `H` it belongs to, i.e. to the linear index of `x` it was seeded from. +# +# It is walked, not indexed, so that the lazy positions of the triangular wrappers never have to be +# materialized. Walking to the start of a block costs what `seed_hessian_chunk!` already pays to +# seed the same block, and no more, since each row resumes the column walk from a saved state +# instead of restarting it. function extract_hessian_chunk!(::Type{T}, H, positions, ydual, roffset, coffset, rsize, csize) where {T} + diagonal = roffset == coffset + rows = Iterators.drop(positions, roffset) + cols = Iterators.drop(positions, coffset) + rnext = iterate(rows) + cfirst = iterate(cols) for r in 1:rsize - i = positions[roffset + r] + i, rstate = something(rnext) drow = partials(T, ydual, r) - cstart = roffset == coffset ? r : 1 + # A diagonal block reads only the upper triangle, whose row `r` begins at column `r` — the + # position the row walk has just reached. Any other block begins at its first column. + cstart, cnext = diagonal ? (r, (i, rstate)) : (1, something(cfirst)) for c in cstart:csize - j = positions[coffset + c] + j, cstate = cnext h = partials(T, drow, c) H[i, j] = h H[j, i] = h + c == csize || (cnext = something(iterate(cols, cstate))) end + r == rsize || (rnext = something(iterate(rows, rstate))) end return H end @@ -112,9 +126,6 @@ reshape_hessian(result::DiffResult, x) = reshape_hessian(DiffResults.hessian(res function symmetric_hessian_expr(result_definition::Expr) return quote xlen = structural_length(x) - # Only the structurally non-zero entries of `x` are seeded, but both axes of the result are - # indexed by the linear indices of `x`, as the columns of a Jacobian are. See #839. - hlen = length(x) if xlen < N throw(ArgumentError(lazy"chunk size cannot be greater than ForwardDiff.structural_length(x) ($(N) > $(structural_length(x)))")) end @@ -134,13 +145,15 @@ function symmetric_hessian_expr(result_definition::Expr) ydual1 = f(xdual) ydual1 isa Real || throw(HESSIAN_ERROR) $(result_definition) - # `H` is square and both its axes are indexed by the entries of `x`, so the columns a - # Jacobian would receive derivatives in are also the rows and columns this sweep writes. - positions = _indexable(structural_columns(H, x)) + # Only the structurally non-zero entries of `x` are seeded, but both axes of `H` are indexed + # by the linear indices of `x`, as the columns of a Jacobian are, so the columns a Jacobian + # would receive derivatives in are the rows and columns this sweep writes. See #839. + positions = structural_columns(H, x) # The entries that no block of the sweep writes belong to none of them in particular: for # `H` the rows and columns of the structurally zero entries of `x`, for `grad` their - # entries. `value(T, ydual1)` is passed for its value type, which unlike `eltype` is a - # number type even for a result that stores `Any`. + # entries. Both take `value(T, ydual1)` for its value type rather than reading `eltype` off + # the result, which is no number type at all for a result that stores `Any` — a scalar dual + # is as good an argument as the array of them `zero_unseeded_columns!` usually gets. zero_unseeded_columns!(T, H, value(T, ydual1), x) grad === nothing || zero_unseeded!(T, grad, value(T, ydual1), x) extract_hessian_chunk!(T, H, positions, ydual1, 0, 0, N, N) @@ -150,9 +163,10 @@ function symmetric_hessian_expr(result_definition::Expr) for q in 2:nblocks qoffset = (q - 1) * N qsize = min(N, xlen - qoffset) - # Off-diagonal blocks: q seeds columns and p seeds rows. Every block of the sweep is - # read from the same triangle of the nested partials, so that the result does not - # depend on the chunk size. The inner seeds for q remain unchanged throughout this loop. + # Off-diagonal blocks: q seeds columns and p seeds rows. Every block reads the nested + # partials with the earlier of its two seeding positions in the outer layer, so an entry + # of `H` is read at the same nesting order at every chunk size. The inner seeds for q + # remain unchanged throughout this loop. seed_hessian_chunk!(xdual, x, qoffset + 1, iseeds, nothing, qsize) for p in 1:(q - 1) poffset = (p - 1) * N @@ -174,7 +188,7 @@ function symmetric_hessian_expr(result_definition::Expr) end @eval function symmetric_hessian(f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N} - $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), hlen, hlen)))) + $(symmetric_hessian_expr(:(H = similar(x, valtype(T, valtype(T, typeof(ydual1))), length(x), length(x))))) end @eval function symmetric_hessian!(H, f::F, x, cfg::HessianConfig{T,V,N}, grad) where {F,T,V,N} diff --git a/src/jacobian.jl b/src/jacobian.jl index 74f88635..1204c556 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -98,7 +98,9 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is # 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. +# that no chunk writes belong to none of them in particular. `ydual` is only read for the type its +# values have, so the Hessian sweep, which zeroes rows as well as columns, passes a single dual +# rather than the array of them a Jacobian has. function zero_unseeded_columns!(::Type{T}, out::AbstractArray, ydual, x) where {T} structural_length(x) == length(x) || fill!(out, zero(valtype(T, eltype(ydual)))) return out diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index cf15f3e9..307f17a5 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -43,6 +43,17 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F allocs_hseed!(hduals, x, 1, nothing, nothing, 4) @test iszero(allocs_hseed!(hduals, x, 1, nothing, nothing, 4)) + # An isbits `V` has free zero partials, so the loop above passes whether or not they are built + # only when they are needed. A `V` whose zeros are heap-allocated shows the difference: seeding + # both layers needs neither layer's zeros, so it has to allocate nothing at all. + xbig = big.(rand(9)) + hcfgbig = ForwardDiff.HessianConfig(nothing, xbig, ForwardDiff.Chunk{3}()) + hdualsbig = hcfgbig.gradient_config.duals + ibig = hcfgbig.jacobian_config.seeds + obig = hcfgbig.gradient_config.seeds + allocs_hseed!(hdualsbig, xbig, 1, ibig, obig) + @test iszero(allocs_hseed!(hdualsbig, xbig, 1, ibig, obig)) + allocs_convert_test_574() = @allocated convert_test_574() allocs_convert_test_574() @test iszero(allocs_convert_test_574()) @@ -83,9 +94,9 @@ function allocs_structured_jacobian!(x, chunk) return @allocated ForwardDiff.jacobian!(result, f!, y, x, cfg) end -# Unlike `gradient!`/`jacobian!`, which walk the structural positions once from the front, the -# Hessian sweep indexes into them, so it materializes them first and cannot be allocation-free for -# the inputs whose positions are lazy. Pin what that costs rather than dropping the check. +# The Hessian sweep reads the positions of two blocks at once and re-reads each row block once per +# column block, so unlike `gradient!`/`jacobian!` it cannot simply walk them once from the front. It +# still must not materialize them: it resumes each walk from a saved state instead. function allocs_structured_hessian!(H, x, chunk) f(z) = sum(abs2, z) fill!(H, false) @@ -94,8 +105,6 @@ function allocs_structured_hessian!(H, x, chunk) return @allocated ForwardDiff.hessian!(H, f, x, cfg) end -allocs_positions(H, x) = @allocated ForwardDiff._indexable(ForwardDiff.structural_columns(H, x)) - @testset "Test gradient!/jacobian!/hessian! allocations for $(nameof(typeof(x)))" for (x, nstruct) in ( (rand(6, 6), 36), (LowerTriangular(rand(6, 6)), 21), @@ -110,12 +119,9 @@ allocs_positions(H, x) = @allocated ForwardDiff._indexable(ForwardDiff.structura @test iszero(allocs_structured_jacobian!(x, chunk)) end - # zero for the inputs whose positions are a range already, and otherwise exactly the one vector - # they are collected into: nothing else in the sweep may allocate H = zeros(length(x), length(x)) - allocs_positions(H, x) # warmup for chunk_size in (2, nstruct) - @test allocs_structured_hessian!(H, x, ForwardDiff.Chunk{chunk_size}()) == allocs_positions(H, x) + @test iszero(allocs_structured_hessian!(H, x, ForwardDiff.Chunk{chunk_size}())) end end diff --git a/test/HessianTest.jl b/test/HessianTest.jl index 3df3a8ba..3d59effe 100644 --- a/test/HessianTest.jl +++ b/test/HessianTest.jl @@ -232,9 +232,9 @@ end end @testset "chunk size independence" begin - # Every block of the sweep reads the same triangle of the nested partials, so the result is not - # merely close across chunk sizes but bitwise equal, and equal to the vector mode of the - # `StaticArray` path, which reads that triangle too. + # Every block of the sweep reads the nested partials at the same nesting order, so the result is + # not merely close across chunk sizes but bitwise equal, and equal to the vector mode of the + # `StaticArray` path, which reads them that way too. n = 16 f = z -> sum(sin(z[i]) * exp(z[mod1(i + 1, n)]) / (1 + z[mod1(i + 2, n)]^2) for i in eachindex(z)) x = collect(range(0.1, 0.9; length=n)) @@ -244,6 +244,23 @@ end H = ForwardDiff.hessian(f, x, ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}())) @test H == reference end + + # A structured `x` reads its positions block by block rather than in one pass, and the coupled + # `f` of the structured testset above has an exactly representable Hessian, so its `==` says + # nothing about floating point. There is no `StaticArray` wrapper to compare against, so the + # largest chunk size is the reference. + g = z -> sum(sin(z[i]) * exp(z[j]) / (1 + z[k]^2) + for i in eachindex(z), j in eachindex(z), k in eachindex(z) if i <= j <= k) + @testset "$(nameof(typeof(w)))" for w in (LowerTriangular(rand(4, 4) .+ 1), + UpperTriangular(rand(4, 4) .+ 1), + Diagonal(rand(4) .+ 1)) + nstruct = ForwardDiff.structural_length(w) + reference = ForwardDiff.hessian(g, w, ForwardDiff.HessianConfig(g, w, ForwardDiff.Chunk{nstruct}())) + @test reference == transpose(reference) + @testset "chunk size = $c" for c in 1:(nstruct - 1) + @test ForwardDiff.hessian(g, w, ForwardDiff.HessianConfig(g, w, ForwardDiff.Chunk{c}())) == reference + end + end end @testset "wrongly shaped result" begin @@ -258,6 +275,10 @@ end @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 10), f, x, cfg) result = DiffResults.DiffResult(NaN, fill(NaN, 3), fill(NaN, 4, 4)) @test_throws DimensionMismatch ForwardDiff.hessian!(result, f, x, cfg) + # a `DiffResult` holding a Hessian that is not a matrix is reshaped like a bare one + result = DiffResults.DiffResult(NaN, fill(NaN, 3), fill(NaN, 9)) + result = ForwardDiff.hessian!(result, f, x, cfg) + @test reshape(DiffResults.hessian(result), 3, 3) == 2I(3) end # the `StaticArray` methods bypass the sweep and reshape the result themselves @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 4, 4), f, sx) @@ -284,6 +305,23 @@ end sx = SVector{3}(x) @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 3, 3), identity, sx) @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.HessianResult(sx), identity, sx) + # the sweep checks the first evaluation itself, before it has a result to write into + @testset "chunk size = $c" for c in (2, 3) + cfg = ForwardDiff.HessianConfig(identity, x, ForwardDiff.Chunk{c}()) + @test_throws DimensionMismatch ForwardDiff.hessian(identity, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(fill(NaN, 3, 3), identity, x, cfg) + @test_throws DimensionMismatch ForwardDiff.hessian!(DiffResults.HessianResult(x), identity, x, cfg) + end +end + +@testset "f(x) ignores x" begin + # `partials` is empty, the second case the `Partials{0}` `StaticArray` method covers + f = z -> 1.0 + @test ForwardDiff.hessian(f, SVector{3}(randn(3))) == zeros(3, 3) + @testset "chunk size = $c" for c in (2, 3) + x = randn(3) + @test ForwardDiff.hessian(f, x, ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}())) == zeros(3, 3) + end end @testset "BigFloat with an unassigned input entry" begin diff --git a/test/QATest.jl b/test/QATest.jl index 48fc3c0d..17acd3bc 100644 --- a/test/QATest.jl +++ b/test/QATest.jl @@ -19,7 +19,7 @@ using JET: @test_opt Diagonal(rand(3, 3))) @test_opt ForwardDiff.gradient(first, x, ForwardDiff.GradientConfig(first, x, ForwardDiff.Chunk{2}())) @test_opt ForwardDiff.jacobian(vec, x, ForwardDiff.JacobianConfig(vec, x, ForwardDiff.Chunk{2}())) - # the Hessian sweep materializes those positions, so `_indexable` has to stay concrete + # the Hessian sweep walks those positions by hand, so their iteration state has to stay concrete @test_opt ForwardDiff.hessian(first, x, ForwardDiff.HessianConfig(first, x, ForwardDiff.Chunk{2}())) end end