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>

Author: ChrisRackauckas

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)