[devmotion]

Stacked on #1078.

Replaces mapreduce(hcat/vcat, ...) in the multi-batch Jacobian and Hessian paths with new internal hooks stack_batches_col/row(block, batch_size_settings, seed_example), which take the function computing the block of each batch:

• The generic methods collect the blocks with a comprehension and concatenate them with a single reduce(hcat/vcat, ...). A comprehension rather than map, since map fails to infer when the block function calls map itself.
• The StaticArrays extension overloads them to concatenate an ntuple of blocks, with the number of batches deduced from the seed type. It also overloads the new vec_block with similar_type, so that flattened blocks of e.g. an MMatrix are MVectors and the result is an MMatrix instead of a SizedMatrix backed by a Matrix.

Multi-batch Jacobians and Hessians of static arrays are now inferred and return static matrices (turning #1078's @test_broken into passing tests). Tests cover all batch methods (B = 1, aligned, unaligned, single batch) for pushforward and pullback Jacobians (out-of-place and in-place, real and complex) and Hessians, plus JET checks over the static DIT scenarios. The in-place Hessian with AutoReverseFromPrimitive already fails JET on main and is marked broken.

Why reduce instead of mapreduce

reduce(hcat, ::AbstractVector{<:AbstractVecOrMat}) (and vcat) has a specialized Base method that allocates the result once. mapreduce(f, hcat, v) has no such method and falls back to pairwise hcat, re-copying the partial result at every level.

Concatenating A blocks of size m × B (Chairmarks):

m × (A × B)	mapreduce(f, hcat, v)	reduce(hcat, map(f, v))
10 × (3 × 4)	89 ns, 3.0 KB	93 ns, 2.4 KB
100 × (9 × 12)	12.4 µs, 573 KB	3.5 µs, 191 KB
1000 × (84 × 12)	78 ms, 358 MB	0.36 ms, 16 MB

End-to-end, prepared jacobian(identity, prep, backend, x) with dense x and chunksize = 12 (Julia 1.11):

backend, length(x)	before	after
AutoSimpleFiniteDiff, 100	31.9 µs, 1.16 MB	29.3 µs, 0.87 MB
AutoSimpleFiniteDiff, 1000	12.2 ms, 402 MB	2.7 ms, 74 MB
AutoReverseFromPrimitive(AutoSimpleFiniteDiff), 1000	721 ms, 5.23 GB	419 ms, 4.90 GB

(stack was also considered: equivalent for columns, but rows need an extra permutedims copy.)

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 98.87640% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 98.27%. Comparing base (e64b3a4) to head (f460474).

Files with missing lines	Patch %	Lines
...iationInterface/test/Core/SimpleFiniteDiff/test.jl	93.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1080      +/-   ##
==========================================
+ Coverage   98.21%   98.27%   +0.05%     
==========================================
  Files         143      143              
  Lines        8246     8274      +28     
==========================================
+ Hits         8099     8131      +32     
+ Misses        147      143       -4     

Flag	Coverage Δ
DI	98.12% <98.87%> (+0.07%)	⬆️
DIT	98.67% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 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.

[gdalle]

Thanks, this looks like a worthwhile improvement! A few things to iron out though

[gdalle]

Thanks! I'm good with the changes to the source code, a few more tests and this will be good to go!

[gdalle]

Can we add non-regression type stability tests here, possibly filtering the in-place scenarios?

[devmotion]

I added JET checks (type_stability = :prepared) over all static Jacobian and Hessian scenarios, in-place and out-of-place, for both multi-batch backends. They use chunk size 2, so they cover the aligned path. All four batch methods (B = 1, aligned, unaligned, single batch) are covered by a new @inferred matrix in test/Core/Internals/batchsize.jl: pushforward and pullback Jacobians (out-of-place and in-place, real and complex) and Hessians.

Every JET check passes except the in-place Hessian with AutoReverseFromPrimitive, which fails on main too and is marked broken = true. On main, 44 of these 96 checks fail; with this PR, only those 4 do.

The JET checks are guarded with if JET.JET_AVAILABLE rather than safetypestab, since aviatesk/JET.jl#699 and aviatesk/JET.jl#742 are both closed. Locally, 1.12 with JET 0.12.3 gives the same results as 1.11. The guard has to wrap the whole block, because JET's stubs ignore skip.

These checks exposed a regression in this PR: map over the blocks failed to infer when the block function itself calls map. Fixing it changed the hook you approved: stack_batches_col/row(block, batch_size_settings, seed_example) now takes the block function. The generic method collects the blocks with a comprehension; the StaticArrays method uses ntuple, with no intermediate vector.

A new vec_block uses similar_type, so matrix-shaped static inputs no longer produce SizedMatrix{Matrix}; SMatrix pullback Jacobians and Hessians are now MMatrix. On 1.11, all 224 prepared combinations I benchmarked infer and return the same results as main. Static inputs allocate the same as on main or less; Vector inputs allocate 2 more, for the block vector.

[gdalle]

Can we test the Jacobian here too?

[devmotion]

I added "Static Jacobian with multiple batches". It covers pushforward and pullback Jacobians, out-of-place and in-place, at chunk sizes that hit every batch method. It goes through AutoForwardFromPrimitive / AutoReverseFromPrimitive, because a plain AutoForwardDiff calls ForwardDiff.jacobian directly. The Hessian testset now also covers the single-batch case (chunk size 6).