Avoid reshape allocation in extract_jacobian! for Matrix results

[ChrisRackauckas-Claude]

Summary

extract_jacobian! calls reshape(result, length(ydual), n) which allocates a 48-byte ReshapedArray wrapper on every call. Under --check-bounds=yes (which Pkg.test() always uses), this allocation cannot be elided by the compiler.

For implicit ODE/SDE solvers that call jacobian! multiple times per step via the NL solver, this adds up. For example, SKenCarp in StochasticDiffEq.jl calls nlsolve! 3 times per step, each triggering a Jacobian computation, resulting in 3 × 48 = 144 bytes/step of unnecessary allocations.

Fix: Add a specialized extract_jacobian! 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.

MWE

julia> using ForwardDiff

julia> T = ForwardDiff.Tag{Nothing, Float64};

julia> 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))]);

julia> result = zeros(3, 3);

# Under --check-bounds=yes:
# Before: 48 bytes per call
# After:  0 bytes per call
julia> @allocated ForwardDiff.extract_jacobian!(T, result, ydual, 3)
0

Test plan

• New allocation test in test/AllocationsTest.jl (passes under --check-bounds=yes)
• New correctness test verifying extracted values match expected Jacobian
• Full test suite passes (9036/9036 tests, including under --check-bounds=yes)

🤖 Generated with Claude Code

[devmotion]

There's no guarantee that result has the correct dimensions? Additionally, i might not be a valid index for result.

[codecov]

Codecov Report

❌ Patch coverage is 80.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 90.69%. Comparing base (ff0d903) to head (84b9290).

Files with missing lines	Patch %	Lines
src/jacobian.jl	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #797      +/-   ##
==========================================
- Coverage   90.75%   90.69%   -0.06%     
==========================================
  Files          11       11              
  Lines        1071     1075       +4     
==========================================
+ Hits          972      975       +3     
- Misses         99      100       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[devmotion]

Move these definitions into the function body to avoid closing over these variables?

[ChrisRackauckas]

@devmotion this should be simpler?

[devmotion]

This is not type-stable in general?

From searching through the codebase, my impression is that in quite a few cases we call extract_jacobian! internally with a matrix of the correct dimensions constructed in the previous line. In these cases, we don't even need this check and can always operate with result.

AFAICT the only case where we might have to reshape is when users provide an output array.

Maybe we could reshape user-provided arrays to matrices higher up in the call stack, maybe even in the user-facing function directly, and then never reshape in this internal function and only accept matrices.

The reshaping of the user-provided arrays could eg be done unconditionally for non-Matrix input (and other types for which we don't know whether the function in this draft would be type-stable) and only for Matrix (and other known to be type-stable types) we would use the conditional reshaping.