Do not seed structural zeros by devmotion · Pull Request #739 · JuliaDiff/ForwardDiff.jl

devmotion · 2025-04-02T08:47:19Z

Fixes#738.

Since I had made #695 I felt somewhat responsible for #738. Defining seed! and extract_gradient! for triangular matrices seems to fix the example.

Probably this should b extended to additional matrix types.

mcabbott · 2025-04-02T13:39:05Z

+@testset "LowerTriangular and UpperTriangular" begin
+ M = rand(3, 3)
+ for T in (LowerTriangular, UpperTriangular)
+ @test ForwardDiff.gradient(sum, T(randn(3, 3))) == T(ones(3, 3))


Unfortunately this fails if I increase the size of the matrix above the chunk size:
julia>@testset"LowerTriangular and UpperTriangular"begin M =rand(3, 3) for T in (LowerTriangular, UpperTriangular) @test ForwardDiff.gradient(sum, T(randn(3, 3))) ==T(ones(3, 3)) @test ForwardDiff.gradient(x ->dot(M, x), T(randn(3, 3))) ==T(M) endend Test Summary:| Pass Total Time LowerTriangular and UpperTriangular |441.0s Test.DefaultTestSet("LowerTriangular and UpperTriangular", Any[], 4, false, false, true, 1.74360103628159e9, 1.743601037319704e9, false, "REPL[8]") julia>@testset"LowerTriangular and UpperTriangular"begin M =rand(10, 10) for T in (LowerTriangular, UpperTriangular) @test ForwardDiff.gradient(sum, T(randn(10, 10))) ==T(ones(10, 10)) @test ForwardDiff.gradient(x ->dot(M, x), T(randn(10, 10))) ==T(M) endend LowerTriangular and UpperTriangular: Error During Test at REPL[9]:4 Test threw exception Expression: ForwardDiff.gradient(sum, T(randn(10, 10))) ==T(ones(10, 10)) ArgumentError: cannot set index in the upper triangular part (1, 2) of an LowerTriangular matrix to a nonzero value (Dual{ForwardDiff.Tag{typeof(sum), Float64}}(0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0,0.0)) Stacktrace: [1] throw_nonzeroerror(T::Type, x::Any, i::Int64, j::Int64) @ LinearAlgebra ~/.julia/juliaup/julia-1.11.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:295 [2] setindex! @ ~/.julia/juliaup/julia-1.11.4+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/triangular.jl:326 [inlined] [3] _unsafe_setindex! @ ./reshapedarray.jl:297 [inlined] [4] setindex! @ ./reshapedarray.jl:286 [inlined] [5] setindex! @ ./subarray.jl:372 [inlined] [6] macro expansion @ ./broadcast.jl:973 [inlined] [7] macro expansion @ ./simdloop.jl:77 [inlined] [8] copyto! @ ./broadcast.jl:972 [inlined] [9] copyto! @ ./broadcast.jl:925 [inlined] [10] materialize! @ ./broadcast.jl:883 [inlined] [11] materialize! @ ./broadcast.jl:880 [inlined] [12] seed!(duals::LowerTriangular{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sum), Float64}, Float64, 12}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sum), Float64}, Float64, 12}}}, x::LowerTriangular{Float64, Matrix{Float64}}, index::Int64, seeds::NTuple{12, ForwardDiff.Partials{12, Float64}}, chunksize::Int64) @ ForwardDiff ~/.julia/packages/ForwardDiff/L0kjR/src/apiutils.jl:69 [13] seed! @ ~/.julia/packages/ForwardDiff/L0kjR/src/apiutils.jl:66 [inlined] [14] chunk_mode_gradient(f::typeof(sum), x::LowerTriangular{Float64, Matrix{Float64}}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{typeof(sum), Float64}, Float64, 12, LowerTriangular{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sum), Float64}, Float64, 12}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sum), Float64}, Float64, 12}}}})

Ah good catch 👍

Another catch is: This PR doesn't (yet) fix the fact that still for every element of the matrices seeds are generated - even though (IMO) we should limit it to the structurally non-zero entries.

I have forgotten how all this code works... but 5-arg seed! implies the iteration is happening elsewhere?
The other thing worth checking would be GradientConfig and DiffResult things from https://juliadiff.org/ForwardDiff.jl/stable/user/advanced/ as I'm not sure what paths they take.

Yeah, seed! (and extract_gradient_chunk!) with > 3 arguments are called if gradients are computed in multiple chunks:
ForwardDiff.jl/src/gradient.jl
Lines 118 to 142 in fd2d4a9
# seed work vectors
xdual = cfg.duals
seeds = cfg.seeds
seed!(xdual, x)
# do first chunk manually to calculate output type
seed!(xdual, x, 1, seeds)
ydual =f(xdual)
$(result_definition)
extract_gradient_chunk!(T, result, ydual, 1, N)
seed!(xdual, x, 1)
# do middle chunks
for c in middlechunks
i = ((c -1) * N +1)
seed!(xdual, x, i, seeds)
ydual =f(xdual)
extract_gradient_chunk!(T, result, ydual, i, N)
seed!(xdual, x, i)
end
# do final chunk
seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
ydual =f(xdual)
extract_gradient_chunk!(T, result, ydual, lastchunkindex, lastchunksize)

Those look good. Some chance this is flawed copy-pasting at my end, but for chunk mode I see strange effects:
julia> ForwardDiff.gradient(x ->begin@show(eltype(x)); sum(x) end, UpperTriangular(rand(5, 5))) eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 8} eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 8} eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 8} eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#27#28", Float64}, Float64, 8} 5×5 UpperTriangular{Float64, Matrix{Float64}}:1.01.01.01.01.0⋅1.01.01.01.0⋅⋅1.01.01.0⋅⋅⋅1.01.0⋅⋅⋅⋅1.0 julia>sum(ans) 15.0 julia> ForwardDiff.gradient(x ->begin@show(eltype(x)); sum(x) end, Diagonal(rand(11))); # okeltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#21#22", Float64}, Float64, 11} julia> ForwardDiff.gradient(x ->begin@show(eltype(x)); sum(x) end, Diagonal(rand(13))); # weirdeltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#23#24", Float64}, Float64, 7} eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#23#24", Float64}, Float64, 7} eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#23#24", Float64}, Float64, 7} [....>15 more]

Hmm, strange. I get
julia> ForwardDiff.gradient(x ->begin@show(eltype(x)); sum(x) end, UpperTriangular(rand(5, 5))) eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#40#41", Float64}, Float64, 8} eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#40#41", Float64}, Float64, 8} 5×5 UpperTriangular{Float64, Matrix{Float64}}:1.01.01.01.01.0⋅1.01.01.01.0⋅⋅1.01.01.0⋅⋅⋅1.01.0⋅⋅⋅⋅1.0 julia>sum(ans) 15.0 julia> ForwardDiff.gradient(x ->begin@show(eltype(x)); sum(x) end, Diagonal(rand(11))); eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#42#43", Float64}, Float64, 11} julia> ForwardDiff.gradient(x ->begin@show(eltype(x)); sum(x) end, Diagonal(rand(13))); eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#44#45", Float64}, Float64, 7} eltype(x) = ForwardDiff.Dual{ForwardDiff.Tag{var"#44#45", Float64}, Float64, 7}
which seems expected given
ForwardDiff.jl/src/prelude.jl
Lines 10 to 13 in fd2d4a9
functionChunk(input_length::Integer, threshold::Integer= DEFAULT_CHUNK_THRESHOLD)
N =pickchunksize(input_length, threshold)
Base.@nif12 d->(N == d) d->(Chunk{d}()) d->(Chunk{N}())
end

The only test failures are the ones already present on the master branch (possibly related to some recent changes of acosh or NaNMath.acosh?).

My weird results above must have been my mistake, lazy to check out the branch...
We could add tests of how many executions are used in chunk mode (not sure whether there are any such now) but not essential.

I added a test in d2a8af7.

… init Keeps the O(n^2/N) fix from the previous commits, but revises the API around it. Rename `unseed!` to `seed_zero_partials!`. Of the 14 call sites the rename touched, 10 are initializing a work buffer that was never seeded -- `cfg.duals` is raw `similar` memory, and `ydual` in the `f!` paths never carried a perturbation -- and only 4 are clearing a chunk. `unseed!` names the minority case and asserts a prior state that usually does not exist; `seed_zero_partials!` names what the function does (write `x`'s values, zero the perturbations), which holds at every site. Give the windowed method an explicit `count = N` rather than an implicit `take(..., N)`, mirroring the `chunksize = N` argument that `seed!(duals, x, index, seeds, chunksize)` already has, so the seeding and zeroing paths take the same window arguments. Factor the shared loop into `_seed_zero_partials!`. The two public methods differed only in the index iterator they pass. Initialize the chunk-mode work buffer as a disjoint partition: seed chunk 1, then zero only the untouched tail (`N + 1`, `xlen - N`), instead of clearing the whole array and overwriting the first N elements. `xlen > N` always holds in chunk mode, since `chunksize(cfg) == structural_length(x)` routes to vector mode, so the two windows are disjoint and together cover `1:xlen` -- the buffer is still fully written before the first `f` call, which is what makes `cfg` reuse safe. This also lets the four `work_array_definition` quotes in jacobian.jl drop to just unpacking `cfg.duals`, removing a line duplicated 4x, and gives `count` an in-package caller. Tests. Add test/SeedTest.jl: the windowed clear is only ever called on a chunk that was just seeded, so clearing too much is harmless and no test written against the public API can tell a bounded implementation from an unbounded one -- which is why this regression survived since JuliaDiff#739. The new tests pin the window directly, and were validated by mutation: restoring the unbounded sweep fails 4 of them. Expected structural index sets are written out by hand and `structural_eachindex` is pinned to them once, so a bug in that iterator cannot hide inside the assertions that depend on it; `values_match` compares over `eachindex(x)` for the same reason. Extend the allocation test with the 4-arg form, where `count` is a runtime value and so catches an inference regression at the `_seed_zero_partials!` boundary that the `count`-defaulting forms could hide. Consolidate the second BigFloat block in test/JacobianTest.jl into a loop over the position of the unassigned entry. Existing coverage only ever placed it in the last chunk, which is never cleared, leaving the `Base._unsetindex!` branch of the windowed path unreached; `hole = 5` puts it in a middle chunk. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

… the config Since #739 only the structurally non-zero entries of an input are seeded, but extraction was not updated to match, so the derivatives were written to positions taken from the result container instead of from `x`. `extract_gradient!`/`extract_gradient_chunk!` now take the positions of `x`. Entries that receive no derivative are zeroed, which is their derivative, and the sweep does that once up front: those entries -- the ones belonging to the structural zeros of `x` -- belong to no chunk in particular. The `DiffResult` method splits on mutability, since an immutable result cannot be written to entry by entry, and it dispatches on `DiffResult` rather than `MutableDiffResult`, because a `StaticArray` gradient buffer makes the result immutable even when the buffer itself can be written to entry by entry, as an `MVector` can. Fixes#838, where a dense result got the derivatives at linear positions `1:structural_length(x)` and a `DiffResults.GradientResult` threw, and with it the mis-scattered gradient of `hessian!(::DiffResult, ...)`. The Jacobian is indexed by the linear indices of `x`: column `j` holds `∂f(x)[i]/∂x[j]`, as documented, with hard zeros in the columns of the structural zeros. Its allocations therefore use `length(x)` rather than `structural_length(x)`, which is what `reshape_jacobian` expected all along, so chunk mode stops throwing. Fixes#839. This changes the shape of the result for structured inputs: `jacobian` gains the zero columns and `hessian` inherits both conventions through `jacobian(∇f, x)`, becoming `length(x) x length(x)` with hard-zero rows and columns instead of mixing linear and structural indices. That also makes `hessian!(DiffResults.HessianResult(x), ...)` work. For a `Diagonal` the result now scales with `length(x)`, so differentiating with respect to the diagonal vector is the better choice there. `reshape_jacobian` keeps the short-circuit `extract_jacobian!` was given in #797, with an explicit size check in place of the one `reshape` performed on the way past: since 1.11 `reshape` can no longer return its argument, so every chunk-mode `jacobian!` allocated an `Array` wrapper. Both modes now reject a wrongly shaped matrix result with the same error. Also drops the `map!` that `vector_mode_jacobian(f!, ...)` ran before `extract_jacobian!`, which reads only `ydual`, and that the `map!` after it repeats. Seeding, meanwhile, took the positions to seed from the config's work buffer while extraction took them from `x`, so reusing a config with an input of the same length and a different structure seeded one set of positions and extracted another -- and left the buffer entries outside the input's structure uninitialized for the target function to read. Fixes#842. `structural_eachindex` and `structural_columns` are replaced by `structural_indices`, which returns the linear indices of the seeded entries in seeding order, and each config stores one per work buffer it owns, built from the buffer rather than from `x`. Every API entry then calls `checkstructure` next to `checktag`: it compares `structural_kind`, which is O(1) and a compile-time constant, so it can run unconditionally, and no size or count comparison could replace it -- `LowerTriangular(n, n)` and `UpperTriangular(n, n)` agree on `size` and on `structural_length` alike. Storing the positions also removes the `Iterators.drop` walk that re-traversed them from the front to reach each chunk, three times per gradient chunk. That walk was 35-49% of `gradient!` for an `UpperTriangular` input, and dense inputs paid it too, `Iterators.drop` having no range specialization: `gradient!` is now 1.4-1.9x faster and `jacobian!` 1.2-1.5x, dense and structured alike. Because the positions are linear indices, the Jacobian's structured and dense extraction paths collapse into the one fused broadcast the dense path already used, and neither chunk extractor takes `x` any more. The unassigned-entry branch of seeding is fixed with them, the second half of `CartesianIndex` method at all and whose `AbstractArray` fallback for a linear index recurses forever. `adjoint`, `transpose`, `PermutedDimsArray` and `view` inputs work now, their buffers being plain `Array`s, and the three wrappers `similar` preserves raise an `ArgumentError` naming the entry instead of a `MethodError` or a `StackOverflowError`, since only an `Array` can hold an unassigned entry at all. The new config type parameter is not shimmed: code that pins the parameter count of these types has to be updated. A config for an offset-indexed input now throws at construction rather than at the first `gradient`/`jacobian` call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

… the config Since #739 only the structurally non-zero entries of an input are seeded, but extraction was not updated to match, so the derivatives were written to positions taken from the result container instead of from `x`. `extract_gradient!`/`extract_gradient_chunk!` now take the positions of `x`. Entries that receive no derivative are zeroed, which is their derivative, and the sweep does that once up front: those entries -- the ones belonging to the structural zeros of `x` -- belong to no chunk in particular. The `DiffResult` method splits on mutability, since an immutable result cannot be written to entry by entry, and it dispatches on `DiffResult` rather than `MutableDiffResult`, because a `StaticArray` gradient buffer makes the result immutable even when the buffer itself can be written to entry by entry, as an `MVector` can. Fixes#838, where a dense result got the derivatives at linear positions `1:structural_length(x)` and a `DiffResults.GradientResult` threw, and with it the mis-scattered gradient of `hessian!(::DiffResult, ...)`. The Jacobian is indexed by the linear indices of `x`: column `j` holds `∂f(x)[i]/∂x[j]`, as documented, with hard zeros in the columns of the structural zeros. Its allocations therefore use `length(x)` rather than `structural_length(x)`, which is what `reshape_jacobian` expected all along, so chunk mode stops throwing. Fixes#839. This changes the shape of the result for structured inputs: `jacobian` gains the zero columns and `hessian` inherits both conventions through `jacobian(∇f, x)`, becoming `length(x) x length(x)` with hard-zero rows and columns instead of mixing linear and structural indices. That also makes `hessian!(DiffResults.HessianResult(x), ...)` work. For a `Diagonal` the result now scales with `length(x)`, so differentiating with respect to the diagonal vector is the better choice there. `reshape_jacobian` keeps the short-circuit `extract_jacobian!` was given in #797, with an explicit size check in place of the one `reshape` performed on the way past: since 1.12 `reshape` can no longer return its argument, so every chunk-mode `jacobian!` allocated an `Array` wrapper. Both modes now reject a wrongly shaped matrix result with the same error. Also drops the `map!` that `vector_mode_jacobian(f!, ...)` ran before `extract_jacobian!`, which reads only `ydual`, and that the `map!` after it repeats. Seeding, meanwhile, took the positions to seed from the config's work buffer while extraction took them from `x`, so reusing a config with an input of the same length and a different structure seeded one set of positions and extracted another -- and left the buffer entries outside the input's structure uninitialized for the target function to read. Fixes#842. `structural_eachindex` and `structural_columns` are replaced by `structural_indices`, which returns the linear indices of the seeded entries in seeding order, and each config stores one per work buffer it owns, built from the buffer rather than from `x`. Every API entry then calls `checkstructure` next to `checktag`: it compares `structural_kind`, which is O(1) and a compile-time constant, so it can run unconditionally, and no size or count comparison could replace it -- `LowerTriangular(n, n)` and `UpperTriangular(n, n)` agree on `size` and on `structural_length` alike. Storing the positions also removes the `Iterators.drop` walk that re-traversed them from the front to reach each chunk, three times per gradient chunk. That walk was 35-49% of `gradient!` for an `UpperTriangular` input, and dense inputs paid it too, `Iterators.drop` having no range specialization: `gradient!` is now 1.4-1.9x faster and `jacobian!` 1.2-1.5x, dense and structured alike. Because the positions are linear indices, the Jacobian's structured and dense extraction paths collapse into the one fused broadcast the dense path already used, and neither chunk extractor takes `x` any more. The unassigned-entry branch of seeding is fixed with them, the second half of #842: it called `Base._unsetindex!(duals, idx)`, for which Base has no `CartesianIndex` method at all and whose `AbstractArray` fallback for a linear index recurses forever. `adjoint`, `transpose`, `PermutedDimsArray` and `view` inputs work now, their buffers being plain `Array`s, and the three wrappers `similar` preserves raise an `ArgumentError` naming the entry instead of a `MethodError` or a `StackOverflowError`, since only an `Array` can hold an unassigned entry at all. The new config type parameter is not shimmed: code that pins the parameter count of these types has to be updated. A config for an offset-indexed input now throws at construction rather than at the first `gradient`/`jacobian` call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Fix gradient with LowerTriangular and UpperTriangular input
f3340a3

mcabbott reviewed Apr 2, 2025
View reviewed changes

devmotion added 2 commits April 2, 2025 16:17

Only seed structurally non-zero entries
715f4bc

Base vector vs chunk mode on structural length
161d227

devmotion changed the title ~~Fix gradient with LowerTriangular and UpperTriangular input~~Do not seed structural zerosApr 2, 2025

Fix tests
cf9f4bb

mcabbott approved these changes Apr 2, 2025
View reviewed changes

Check number of function evaluations and chunk sizes
d2a8af7

devmotion merged commit fce7f76 into masterApr 2, 2025

odow mentioned this pull request Apr 3, 2025
Prep for v1.38.1 jump-dev/MathOptInterface.jl#2707
Merged
5 tasks

devmotion deleted the dw/lower_upper_triangular branch April 3, 2025 07:55

penelopeysm mentioned this pull request Apr 3, 2025
add ForwardDiff@1 TuringLang/Bijectors.jl#378
Merged

thomvet mentioned this pull request Aug 8, 2025
Scalar indexing error when computing Jacobian with a CUDA.CuArray#759
Open

devmotion mentioned this pull request Aug 19, 2025
add back the iteration based seeding for Array#634
Closed

KristofferC mentioned this pull request Aug 19, 2025
fix: regression in non-fast scalar indexing support #760
Open
2 tasks

devmotion mentioned this pull request Aug 3, 2026
Unseed only the chunk, and split unseed! out of seed! #821
Merged

Navigation Menu

Uh oh!

Do not seed structural zeros - #739

Do not seed structural zeros#739
devmotion merged 5 commits into
masterfrom
dw/lower_upper_triangular

devmotion commented Apr 2, 2025

Uh oh!

mcabbottApr 2, 2025

Uh oh!

devmotionApr 2, 2025

Uh oh!

devmotionApr 2, 2025

Uh oh!

mcabbottApr 2, 2025

Uh oh!

devmotionApr 2, 2025•
edited
Loading

Uh oh!

mcabbottApr 2, 2025•
edited
Loading

Uh oh!

devmotionApr 2, 2025

Uh oh!

devmotionApr 2, 2025

Uh oh!

mcabbottApr 2, 2025

Uh oh!

devmotionApr 2, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

	# seed work vectors
	xdual = cfg.duals
	seeds = cfg.seeds
	seed!(xdual, x)

	# do first chunk manually to calculate output type
	seed!(xdual, x, 1, seeds)
	ydual =f(xdual)
	$(result_definition)
	extract_gradient_chunk!(T, result, ydual, 1, N)
	seed!(xdual, x, 1)

	# do middle chunks
	for c in middlechunks
	i = ((c -1) * N +1)
	seed!(xdual, x, i, seeds)
	ydual =f(xdual)
	extract_gradient_chunk!(T, result, ydual, i, N)
	seed!(xdual, x, i)
	end

	# do final chunk
	seed!(xdual, x, lastchunkindex, seeds, lastchunksize)
	ydual =f(xdual)
	extract_gradient_chunk!(T, result, ydual, lastchunkindex, lastchunksize)

	functionChunk(input_length::Integer, threshold::Integer= DEFAULT_CHUNK_THRESHOLD)
	N =pickchunksize(input_length, threshold)
	Base.@nif12 d->(N == d) d->(Chunk{d}()) d->(Chunk{N}())
	end

Uh oh!

Conversation

devmotion commented Apr 2, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

devmotionApr 2, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mcabbottApr 2, 2025• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

devmotionApr 2, 2025•
edited
Loading

mcabbottApr 2, 2025•
edited
Loading