From 84b92909ba75db5e93baf26e485bc4c03637650b Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 17 Mar 2026 09:54:55 -0400 Subject: [PATCH 1/4] Avoid reshape allocation in extract_jacobian! for Matrix results `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 Co-Authored-By: Claude Opus 4.6 (1M context) --- src/jacobian.jl | 12 +++++++++++- test/AllocationsTest.jl | 16 ++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index b12142a3..b82f260f 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -93,7 +93,7 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is ##################### function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T} - out_reshaped = reshape(result, length(ydual), n) + out_reshaped = _maybe_reshape(result, length(ydual), n) ydual_reshaped = vec(ydual) # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) @@ -117,6 +117,16 @@ function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) whe return result end +# Avoid allocating a ReshapedArray wrapper when `result` already has the target shape. +# reshape() always allocates a wrapper that cannot be elided under --check-bounds=yes. +@inline function _maybe_reshape(result::AbstractArray, m, n) + if size(result) == (m, n) + return result + else + return reshape(result, m, n) + end +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) diff --git a/test/AllocationsTest.jl b/test/AllocationsTest.jl index 9e7c38be..36df8eba 100644 --- a/test/AllocationsTest.jl +++ b/test/AllocationsTest.jl @@ -28,4 +28,20 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F @test iszero(allocs_convert_test_574()) end +@testset "Test jacobian! allocations" begin + # jacobian! should not allocate when called with a pre-allocated result Matrix. + # Previously, reshape() inside extract_jacobian! allocated a wrapper + # object that could not be elided under --check-bounds=yes. + function allocs_jacobian!() + f!(y, x) = (y .= x .^ 2) + x = [1.0, 2.0, 3.0] + y = similar(x) + result = zeros(3, 3) + cfg = ForwardDiff.JacobianConfig(f!, y, x) + ForwardDiff.jacobian!(result, f!, y, x, cfg) # warmup + return @allocated ForwardDiff.jacobian!(result, f!, y, x, cfg) + end + @test iszero(allocs_jacobian!()) +end + end From 0964e468137cd0f541704d9e818d2d2735feea71 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 24 Mar 2026 06:44:18 -0400 Subject: [PATCH 2/4] Use dispatch instead of conditional _maybe_reshape for type stability 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 Co-Authored-By: Claude Opus 4.6 (1M context) --- src/jacobian.jl | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index b82f260f..4372dc02 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -92,8 +92,19 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is # result extraction # ##################### +# 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} + ydual_reshaped = vec(ydual) + # Use closure to avoid GPU broadcasting with Type + partials_wrap(ydual, nrange) = partials(T, ydual, nrange) + result .= partials_wrap.(ydual_reshaped, transpose(1:n)) + return result +end + +# General method for non-matrix arrays: reshape unconditionally (type-stable). function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T} - out_reshaped = _maybe_reshape(result, length(ydual), n) + out_reshaped = reshape(result, length(ydual), n) ydual_reshaped = vec(ydual) # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) @@ -117,16 +128,6 @@ function extract_jacobian_chunk!(::Type{T}, result, ydual, index, chunksize) whe return result end -# Avoid allocating a ReshapedArray wrapper when `result` already has the target shape. -# reshape() always allocates a wrapper that cannot be elided under --check-bounds=yes. -@inline function _maybe_reshape(result::AbstractArray, m, n) - if size(result) == (m, n) - return result - else - return reshape(result, m, n) - end -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) From 5a3dda9cd48fee779b08d66e079ab2c25534faac Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 24 Mar 2026 07:27:57 -0400 Subject: [PATCH 3/4] Fold AbstractMatrix optimization into general method as 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 Co-Authored-By: Claude Opus 4.6 (1M context) --- src/jacobian.jl | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/src/jacobian.jl b/src/jacobian.jl index 4372dc02..b8ce58fb 100644 --- a/src/jacobian.jl +++ b/src/jacobian.jl @@ -92,19 +92,8 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is # result extraction # ##################### -# 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} - ydual_reshaped = vec(ydual) - # Use closure to avoid GPU broadcasting with Type - partials_wrap(ydual, nrange) = partials(T, ydual, nrange) - result .= partials_wrap.(ydual_reshaped, transpose(1:n)) - return result -end - -# General method for non-matrix arrays: reshape unconditionally (type-stable). function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArray, n) where {T} - out_reshaped = reshape(result, length(ydual), n) + out_reshaped = result isa AbstractMatrix ? result : reshape(result, length(ydual), n) ydual_reshaped = vec(ydual) # Use closure to avoid GPU broadcasting with Type partials_wrap(ydual, nrange) = partials(T, ydual, nrange) From ac91112face5a472f9a0a88b8074a95076edfb56 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Tue, 24 Mar 2026 10:49:41 -0100 Subject: [PATCH 4/4] Update Project.toml --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index d0138ce0..fa3e86f8 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ForwardDiff" uuid = "f6369f11-7733-5829-9624-2563aa707210" -version = "1.3.2" +version = "1.3.3" [deps] CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"