Avoid reshape allocation in extract_jacobian! for Matrix results by ChrisRackauckas-Claude · Pull Request #797 · JuliaDiff/ForwardDiff.jl

ChrisRackauckas-Claude · 2026-03-17T13:55:13Z

Summary

extract_jacobian! calls reshape(result, length(ydual), n) which allocates a 48-byte ReshapedArray wrapper on every call. Under --check-bounds=yes (which Pkg.test() always uses), this allocation cannot be elided by the compiler.

For implicit ODE/SDE solvers that call jacobian! multiple times per step via the NL solver, this adds up. For example, SKenCarp in StochasticDiffEq.jl calls nlsolve! 3 times per step, each triggering a Jacobian computation, resulting in 3 × 48 = 144 bytes/step of unnecessary allocations.

Fix: Add a specialized extract_jacobian! method for Matrix result with AbstractVector ydual that uses direct loop indexing instead of reshape+broadcast. This is zero-alloc under --check-bounds=yes and produces identical results.

MWE

julia>using ForwardDiff
julia> T = ForwardDiff.Tag{Nothing, Float64};
julia> ydual = ForwardDiff.Dual{T}.([1.0, 2.0, 3.0],
[ForwardDiff.Partials((1.0, 2.0, 3.0)),
ForwardDiff.Partials((4.0, 5.0, 6.0)),
ForwardDiff.Partials((7.0, 8.0, 9.0))]);
julia> result =zeros(3, 3);
# Under --check-bounds=yes:# Before: 48 bytes per call# After: 0 bytes per call
julia>@allocated ForwardDiff.extract_jacobian!(T, result, ydual, 3)
0

Test plan

New allocation test in test/AllocationsTest.jl (passes under --check-bounds=yes)
New correctness test verifying extracted values match expected Jacobian
Full test suite passes (9036/9036 tests, including under --check-bounds=yes)

🤖 Generated with Claude Code

devmotion · 2026-03-17T13:58:41Z

+function extract_jacobian!(::Type{T}, result::Matrix, ydual::AbstractVector, n) where {T}
+ for j in 1:n
+ for i in eachindex(ydual)
+ result[i, j] = partials(T, ydual[i], j)


There's no guarantee that result has the correct dimensions? Additionally, i might not be a valid index for result.

codecov · 2026-03-17T14:03:07Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 90.75%. Comparing base (ff0d903) to head (ac91112).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@ Coverage Diff @@## master #797 +/- ##
=======================================
Coverage 90.75% 90.75% =======================================
Files 11 11 Lines 1071 1071 =======================================
Hits 972 972 Misses 99 99

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

devmotion · 2026-03-17T14:06:53Z

+ T = ForwardDiff.Tag{Nothing, Float64}
+ N = 3
+ ydual = ForwardDiff.Dual{T}.(
+ [1.0, 2.0, 3.0],
+ [ForwardDiff.Partials((1.0, 2.0, 3.0)),
+ ForwardDiff.Partials((4.0, 5.0, 6.0)),
+ ForwardDiff.Partials((7.0, 8.0, 9.0))]
+ )
+ result = zeros(3, 3)
+
+ allocs_extract!() = @allocated ForwardDiff.extract_jacobian!(T, result, ydual, N)


Move these definitions into the function body to avoid closing over these variables?

`extract_jacobian!` called `reshape(result, length(ydual), n)` which allocates a 48-byte ReshapedArray wrapper. Under `--check-bounds=yes` (used by Pkg.test), this allocation cannot be elided by the compiler, causing 48 bytes per jacobian! call. For implicit ODE/SDE solvers that call jacobian! multiple times per step, this adds up (e.g. 144 bytes/step for SKenCarp with 3 NL solver iterations). Add `_maybe_reshape` that returns the array as-is when it already has the target shape, avoiding the wrapper allocation. Falls back to `reshape` when dimensions don't match. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ChrisRackauckas · 2026-03-17T14:30:01Z

@devmotion this should be simpler?

devmotion · 2026-03-17T15:07:02Z

+ if size(result) == (m, n)
+ return result
+ else
+ return reshape(result, m, n)
+ end


This is not type-stable in general?
From searching through the codebase, my impression is that in quite a few cases we call extract_jacobian! internally with a matrix of the correct dimensions constructed in the previous line. In these cases, we don't even need this check and can always operate with result.
AFAICT the only case where we might have to reshape is when users provide an output array.
Maybe we could reshape user-provided arrays to matrices higher up in the call stack, maybe even in the user-facing function directly, and then never reshape in this internal function and only accept matrices.
The reshaping of the user-provided arrays could eg be done unconditionally for non-Matrix input (and other types for which we don't know whether the function in this draft would be type-stable) and only for Matrix (and other known to be type-stable types) we would use the conditional reshaping.

Address review feedback: replace the type-unstable _maybe_reshape helper with two dispatch-based extract_jacobian! methods: - AbstractMatrix: skip reshape entirely (zero-alloc, hot path for DiffEq) - AbstractArray: reshape unconditionally (type-stable for non-matrix inputs) Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ChrisRackauckas-Claude · 2026-03-24T10:44:40Z

@devmotion Good point about the type instability of _maybe_reshape. Replaced it with dispatch-based methods in 0964e46:

extract_jacobian!(::Type{T}, result::AbstractMatrix, ...) — skips reshape entirely (zero-alloc, covers the hot path where callers construct a matrix internally or users pass a Matrix)
extract_jacobian!(::Type{T}, result::AbstractArray, ...) — reshapes unconditionally (type-stable for non-matrix user-provided arrays)

Both paths are now type-stable. Full test suite passes (9035/9035) and the allocation test confirms 0 bytes for the Matrix path under --check-bounds=yes.

devmotion · 2026-03-24T10:55:02Z


+# Specialized method for AbstractMatrix: no reshape needed, avoids ReshapedArray allocation
+# that cannot be elided under --check-bounds=yes.
+function extract_jacobian!(::Type{T}, result::AbstractMatrix, ydual::AbstractArray, n) where {T}


Move to the function below to minimize the diff to a one-line change?

Address review: instead of a separate dispatch method for AbstractMatrix, use a conditional in the existing method to skip reshape when result is already a matrix. This minimizes the diff while preserving the allocation fix under --check-bounds=yes. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

devmotion

If you bump the version number, we can tag a release when the PR is merged.

ChrisRackauckas · 2026-03-24T12:23:36Z

done

…vered 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) <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>

devmotion reviewed Mar 17, 2026
View reviewed changes

ChrisRackauckas-Claude force-pushed the fix-extract-jacobian-reshape-alloc branch from cb850f0 to 526aecdCompare March 17, 2026 14:15

ChrisRackauckas-Claude force-pushed the fix-extract-jacobian-reshape-alloc branch from 526aecd to 84b9290Compare March 17, 2026 14:29

ChrisRackauckas requested a review from devmotion March 17, 2026 14:29

ChrisRackauckas-Claude mentioned this pull request Mar 17, 2026
Tighten SKenCarp allocation test to == 0 after ForwardDiff.jl#797 SciML/StochasticDiffEq.jl#699
Closed

devmotion reviewed Mar 17, 2026
View reviewed changes

devmotion reviewed Mar 24, 2026
View reviewed changes

Update Project.toml
ac91112

devmotion approved these changes Mar 24, 2026
View reviewed changes

devmotion merged commit 1295777 into JuliaDiff:masterMar 24, 2026
37 of 52 checks passed

ChrisRackauckas-Claude mentioned this pull request Jul 16, 2026
Eliminate dense continuation step allocations on Julia 1.12 SciML/NonlinearSolve.jl#1061
Closed

devmotion mentioned this pull request Aug 17, 2026
Extract derivatives at the indices of x, and store those indices in the config #840
Open

Navigation Menu

Uh oh!

Avoid reshape allocation in extract_jacobian! for Matrix results - #797

Avoid reshape allocation in extract_jacobian! for Matrix results#797
devmotion merged 4 commits into
JuliaDiff:masterfrom
ChrisRackauckas-Claude:fix-extract-jacobian-reshape-alloc

ChrisRackauckas-Claude commented Mar 17, 2026

Uh oh!

devmotionMar 17, 2026

Uh oh!

codecovBot commented Mar 17, 2026•
edited
Loading

Uh oh!

devmotionMar 17, 2026

Uh oh!

ChrisRackauckas commented Mar 17, 2026

Uh oh!

devmotionMar 17, 2026

Uh oh!

ChrisRackauckas-Claude commented Mar 24, 2026

Uh oh!

devmotionMar 24, 2026

Uh oh!

devmotion left a comment

Uh oh!

ChrisRackauckas commented Mar 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

ChrisRackauckas-Claude commented Mar 17, 2026

Summary

MWE

Test plan

Uh oh!

devmotionMar 17, 2026

Choose a reason for hiding this comment

Uh oh!

codecovBot commented Mar 17, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

devmotionMar 17, 2026

Choose a reason for hiding this comment

Uh oh!

ChrisRackauckas commented Mar 17, 2026

Uh oh!

devmotionMar 17, 2026

Choose a reason for hiding this comment

Uh oh!

ChrisRackauckas-Claude commented Mar 24, 2026

Uh oh!

devmotionMar 24, 2026

Choose a reason for hiding this comment

Uh oh!

devmotion left a comment

Choose a reason for hiding this comment

Uh oh!

ChrisRackauckas commented Mar 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecovBot commented Mar 17, 2026•
edited
Loading