Avoid reshape allocation in extract_jacobian! for Matrix results (#797)

* 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>

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

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

* Update Project.toml

---------

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: 3 people

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"

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 = 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)

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