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 a specialized 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.

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

@@ -101,6 +101,17 @@ function extract_jacobian!(::Type{T}, result::AbstractArray, ydual::AbstractArra
     return result
 end
 
+# Avoid reshape allocation for the common case where result is already a correctly-shaped Matrix.
+# reshape() allocates a wrapper object that cannot be elided under --check-bounds=yes.
+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)
+        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)
     return result

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