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 <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

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)
 

test/AllocationsTest.jl

@@ -28,4 +28,29 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F
     @test iszero(allocs_convert_test_574())
 end
 
+@testset "Test extract_jacobian! allocations" begin
+    # extract_jacobian! should not allocate when result is a Matrix.
+    # Previously, reshape() inside extract_jacobian! allocated a wrapper
+    # object that could not be elided under --check-bounds=yes.
+    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)
+    allocs_extract!()  # warmup
+    @test iszero(allocs_extract!())
+
+    # Verify correctness: result[i,j] should be ∂yᵢ/∂xⱼ
+    ForwardDiff.extract_jacobian!(T, result, ydual, N)
+    @test result == [1.0 2.0 3.0;
+                     4.0 5.0 6.0;
+                     7.0 8.0 9.0]
+end
+
 end