From 71df5206720f924931708341a666c5bb44cef16e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Mon, 12 Oct 2026 11:37:50 +0200
Subject: [PATCH 1/5] test: check inference of batch size and Jacobian
 preparation for static arrays

Follow-up to #1076: lock in that `pick_batchsize` and `prepare_jacobian`
infer the batch size from the type for `SVector`, `SMatrix` and `MVector`,
with adaptive and fixed batch sizes. Mark multi-batch `jacobian` inference
as broken, since `mapreduce(hcat, ...)` changes the static size at every step.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../test/Core/Internals/batchsize.jl          | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/DifferentiationInterface/test/Core/Internals/batchsize.jl b/DifferentiationInterface/test/Core/Internals/batchsize.jl
index 5ee75cef4..9a08aa5c1 100644
--- a/DifferentiationInterface/test/Core/Internals/batchsize.jl
+++ b/DifferentiationInterface/test/Core/Internals/batchsize.jl
@@ -74,6 +74,27 @@ end
     ) isa BSS{4, false, true}
 end
 
+@testset "Static arrays" begin
+    # the batch size must be deduced from the type, for any static array shape
+    for x in (@SVector(zeros(6)), @SMatrix(zeros(2, 3)), @MVector(zeros(6)))
+        @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(), x)) isa BSS{6, true, true}
+        @test (@inferred pick_batchsize(AutoSimpleFiniteDiff(; chunksize = 4), x)) isa
+            BSS{4, false, false}
+        @test (@inferred pick_batchsize(AutoEnzyme(), x)) isa BSS{6, true, true}
+        @test (@inferred prepare_jacobian(identity, AutoSimpleFiniteDiff(), x)) isa
+            DI.PushforwardJacobianPrep{<:Any, BSS{6, true, true}}
+        @test (
+            @inferred prepare_jacobian(identity, AutoSimpleFiniteDiff(; chunksize = 4), x)
+        ) isa DI.PushforwardJacobianPrep{<:Any, BSS{4, false, false}}
+        @test (@inferred jacobian(identity, AutoSimpleFiniteDiff(), x)) isa StaticMatrix
+        # multiple batches are assembled with `mapreduce(hcat, ...)`, which changes the
+        # static size at every step
+        @test_broken (
+            @inferred jacobian(identity, AutoSimpleFiniteDiff(; chunksize = 4), x)
+        ) isa StaticMatrix
+    end
+end
+
 @testset "Thresholding" begin
     @test threshold_batchsize(AutoSimpleFiniteDiff(), 2) isa AutoSimpleFiniteDiff{nothing}
     @test threshold_batchsize(AutoSimpleFiniteDiff(; chunksize = 4), 2) isa

From 01e12741c4c90a5c92ffc59bbd264d80963891f7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Mon, 12 Oct 2026 16:36:37 +0200
Subject: [PATCH 2/5] fix: type-stable multi-batch Jacobian and Hessian
 assembly for static arrays

Compute all batch blocks with `map` and concatenate them once with
`reduce(hcat/vcat, ...)` via new internal hooks `stack_batches_col` and
`stack_batches_row`, instead of `mapreduce(hcat/vcat, ...)` with an
in-loop slice of the last batch. The StaticArrays extension overloads
the hooks to concatenate a statically sized tuple of blocks, so that
multi-batch Jacobians and Hessians of static arrays are inferred and
return static matrices.

For dense arrays this avoids the repeated reallocation of pairwise
`hcat`/`vcat` (e.g. 10x faster, 5x less memory for a 1000x1000
Jacobian with chunksize 12).

Also remove the unreachable unaligned branch in the single-batch
Jacobian methods: a single batch is always aligned.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ...DifferentiationInterfaceStaticArraysExt.jl | 33 +++++++++-
 .../src/first_order/jacobian.jl               | 60 ++++++-------------
 .../src/second_order/hessian.jl               | 22 +++----
 DifferentiationInterface/src/utils/linalg.jl  | 24 ++++++++
 .../test/Back/ForwardDiff/test.jl             | 16 ++++-
 .../test/Core/Internals/batchsize.jl          | 25 ++++++--
 .../test/Core/SimpleFiniteDiff/test.jl        | 10 ++++
 7 files changed, 127 insertions(+), 63 deletions(-)

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
index f450aaadc..309ae8199 100644
--- a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
+++ b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
@@ -2,7 +2,7 @@ module DifferentiationInterfaceStaticArraysExt
 
 using ADTypes: AutoForwardDiff, AutoEnzyme
 import DifferentiationInterface as DI
-using StaticArrays: SArray, StaticArray
+using StaticArrays: SArray, SOneTo, StaticArray
 
 function DI.stack_vec_col(t::NTuple{B, <:StaticArray}) where {B}
     return hcat(map(vec, t)...)
@@ -12,6 +12,37 @@ function DI.stack_vec_row(t::NTuple{B, <:StaticArray}) where {B}
     return vcat(transpose.(map(vec, t))...)
 end
 
+# The number of seeds `N` and of batches `A` are deduced from the seed type,
+# so that concatenating the blocks preserves static sizes.
+
+function static_batch_count(::Type{T}, B::Int) where {T <: StaticArray}
+    N = length(T)
+    A = cld(N, B)
+    # otherwise `hcat()` would silently return an empty `Vector{Any}`
+    A > 0 || throw(ArgumentError("Cannot concatenate zero batches"))
+    return N, A
+end
+
+function DI.stack_batches_col(
+        blocks::AbstractVector{<:StaticArray},
+        ::DI.BatchSizeSettings{B, singlebatch, aligned},
+        ::NTuple{B, T},
+    ) where {B, singlebatch, aligned, T <: StaticArray}
+    N, A = static_batch_count(T, B)
+    M = hcat(ntuple(a -> blocks[a], Val(A))...)
+    return aligned ? M : M[:, SOneTo(N)]
+end
+
+function DI.stack_batches_row(
+        blocks::AbstractVector{<:StaticArray},
+        ::DI.BatchSizeSettings{B, singlebatch, aligned},
+        ::NTuple{B, T},
+    ) where {B, singlebatch, aligned, T <: StaticArray}
+    N, A = static_batch_count(T, B)
+    M = vcat(ntuple(a -> blocks[a], Val(A))...)
+    return aligned ? M : M[SOneTo(N), :]
+end
+
 DI.ismutable_array(::Type{<:SArray}) = false
 
 function DI.pick_batchsize(::DI.AutoSimpleFiniteDiff{nothing}, x::StaticArray)
diff --git a/DifferentiationInterface/src/first_order/jacobian.jl b/DifferentiationInterface/src/first_order/jacobian.jl
index 4c64b45bc..8b261878c 100644
--- a/DifferentiationInterface/src/first_order/jacobian.jl
+++ b/DifferentiationInterface/src/first_order/jacobian.jl
@@ -354,22 +354,16 @@ end
 
 function _jacobian_aux(
         f_or_f!y::FY,
-        prep::PushforwardJacobianPrep{SIG, <:BatchSizeSettings{B, true, aligned}},
+        prep::PushforwardJacobianPrep{SIG, <:BatchSizeSettings{B, true}},
         backend::AbstractADType,
         x,
         contexts::Vararg{Context, C},
-    ) where {FY, SIG, B, aligned, C}
-    (; batch_size_settings, batched_seeds, pushforward_prep) = prep
-    (; B_last) = batch_size_settings
+    ) where {FY, SIG, B, C}
+    (; batched_seeds, pushforward_prep) = prep
     dy_batch = pushforward(
         f_or_f!y..., pushforward_prep, backend, x, only(batched_seeds), contexts...
     )
-    block = stack_vec_col(dy_batch)
-    if aligned
-        return block
-    else
-        return block[:, 1:B_last]
-    end
+    return stack_vec_col(dy_batch)
 end
 
 function _jacobian_aux(
@@ -403,19 +397,18 @@ end
 
 function _jacobian_aux(
         f_or_f!y::FY,
-        prep::PushforwardJacobianPrep{SIG, <:BatchSizeSettings{B, false, aligned}},
+        prep::PushforwardJacobianPrep{SIG, <:BatchSizeSettings{B, false}},
         backend::AbstractADType,
         x,
         contexts::Vararg{Context, C},
-    ) where {FY, SIG, B, aligned, C}
+    ) where {FY, SIG, B, C}
     (; batch_size_settings, batched_seeds, seed_example, pushforward_prep) = prep
-    (; A, B_last) = batch_size_settings
 
     pushforward_prep_same = prepare_pushforward_same_point(
         f_or_f!y..., pushforward_prep, backend, x, seed_example, contexts...
     )
 
-    jac = mapreduce(hcat, eachindex(batched_seeds)) do a
+    blocks = map(eachindex(batched_seeds)) do a
         dy_batch = pushforward(
             f_or_f!y...,
             pushforward_prep_same,
@@ -424,37 +417,26 @@ function _jacobian_aux(
             batched_seeds[a],
             contexts...,
         )
-        block = stack_vec_col(dy_batch)
-        if !aligned && a == A
-            return block[:, 1:B_last]
-        else
-            return block
-        end
+        return stack_vec_col(dy_batch)
     end
-    return jac
+    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux(
         f_or_f!y::FY,
-        prep::PullbackJacobianPrep{SIG, <:BatchSizeSettings{B, true, aligned}},
+        prep::PullbackJacobianPrep{SIG, <:BatchSizeSettings{B, true}},
         backend::AbstractADType,
         x,
         contexts::Vararg{Context, C},
-    ) where {FY, SIG, B, aligned, C}
-    (; batch_size_settings, batched_seeds, pullback_prep) = prep
-    (; B_last) = batch_size_settings
+    ) where {FY, SIG, B, C}
+    (; batched_seeds, pullback_prep) = prep
     dx_batch = pullback(
         f_or_f!y..., pullback_prep, backend, x, only(batched_seeds), contexts...
     )
     if eltype(x) <: Complex
         dx_batch = map(conj, dx_batch)
     end
-    block = stack_vec_row(dx_batch)
-    if aligned
-        return block
-    else
-        return block[1:B_last, :]
-    end
+    return stack_vec_row(dx_batch)
 end
 
 function _jacobian_aux(
@@ -487,33 +469,27 @@ end
 
 function _jacobian_aux(
         f_or_f!y::FY,
-        prep::PullbackJacobianPrep{SIG, <:BatchSizeSettings{B, false, aligned}},
+        prep::PullbackJacobianPrep{SIG, <:BatchSizeSettings{B, false}},
         backend::AbstractADType,
         x,
         contexts::Vararg{Context, C},
-    ) where {FY, SIG, B, aligned, C}
+    ) where {FY, SIG, B, C}
     (; batch_size_settings, batched_seeds, seed_example, pullback_prep) = prep
-    (; A, B_last) = batch_size_settings
 
     pullback_prep_same = prepare_pullback_same_point(
         f_or_f!y..., pullback_prep, backend, x, seed_example, contexts...
     )
 
-    jac = mapreduce(vcat, eachindex(batched_seeds)) do a
+    blocks = map(eachindex(batched_seeds)) do a
         dx_batch = pullback(
             f_or_f!y..., pullback_prep_same, backend, x, batched_seeds[a], contexts...
         )
         if eltype(x) <: Complex
             dx_batch = map(conj, dx_batch)
         end
-        block = stack_vec_row(dx_batch)
-        if !aligned && a == A
-            return block[1:B_last, :]
-        else
-            return block
-        end
+        return stack_vec_row(dx_batch)
     end
-    return jac
+    return stack_batches_row(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux!(
diff --git a/DifferentiationInterface/src/second_order/hessian.jl b/DifferentiationInterface/src/second_order/hessian.jl
index 84bbcad20..8c5c99d37 100644
--- a/DifferentiationInterface/src/second_order/hessian.jl
+++ b/DifferentiationInterface/src/second_order/hessian.jl
@@ -159,44 +159,38 @@ function hessian(
         contexts::Vararg{Context, C},
     ) where {F, SIG, C}
     check_prep(f, prep, backend, x, contexts...)
-    (; batched_seeds, seed_example, hvp_prep) = prep
+    (; batch_size_settings, batched_seeds, seed_example, hvp_prep) = prep
 
     hvp_prep_same = prepare_hvp_same_point(
         f, hvp_prep, backend, x, seed_example, contexts...
     )
 
-    hess = mapreduce(hcat, eachindex(batched_seeds)) do a
+    blocks = map(eachindex(batched_seeds)) do a
         dg = only(hvp(f, hvp_prep_same, backend, x, batched_seeds[a], contexts...))
         return vec(dg)
     end
-    return hess
+    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function hessian(
         f::F,
-        prep::HVPGradientHessianPrep{SIG, <:BatchSizeSettings{B, false, aligned}},
+        prep::HVPGradientHessianPrep{SIG, <:BatchSizeSettings{B, false}},
         backend::AbstractADType,
         x,
         contexts::Vararg{Context, C},
-    ) where {F, SIG, B, aligned, C}
+    ) where {F, SIG, B, C}
     check_prep(f, prep, backend, x, contexts...)
     (; batch_size_settings, batched_seeds, seed_example, hvp_prep) = prep
-    (; A, B_last) = batch_size_settings
 
     hvp_prep_same = prepare_hvp_same_point(
         f, hvp_prep, backend, x, seed_example, contexts...
     )
 
-    hess = mapreduce(hcat, eachindex(batched_seeds)) do a
+    blocks = map(eachindex(batched_seeds)) do a
         dg_batch = hvp(f, hvp_prep_same, backend, x, batched_seeds[a], contexts...)
-        block = stack_vec_col(dg_batch)
-        if !aligned && a == A
-            return block[:, 1:B_last]
-        else
-            return block
-        end
+        return stack_vec_col(dg_batch)
     end
-    return hess
+    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function hessian!(
diff --git a/DifferentiationInterface/src/utils/linalg.jl b/DifferentiationInterface/src/utils/linalg.jl
index b27706a5e..5d591e21e 100644
--- a/DifferentiationInterface/src/utils/linalg.jl
+++ b/DifferentiationInterface/src/utils/linalg.jl
@@ -1,6 +1,30 @@
 stack_vec_col(t::NTuple) = stack(vec, t; dims = 2)
 stack_vec_row(t::NTuple) = stack(vec, t; dims = 1)
 
+"""
+    stack_batches_col(blocks, batch_size_settings, seed_example)
+    stack_batches_row(blocks, batch_size_settings, seed_example)
+
+Concatenate the blocks computed from all batches along columns (resp. rows), and drop the trailing entries of the last batch if it was padded by wrapping around the seeds.
+
+The `seed_example` is only used for dispatch, e.g. to deduce the number of seeds statically.
+
+Unlike `mapreduce(hcat, ...)`, `reduce(hcat, ...)` allocates the result only once.
+"""
+function stack_batches_col(
+        blocks::AbstractVector, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
+    ) where {B, singlebatch, aligned}
+    M = reduce(hcat, blocks)
+    return aligned ? M : M[:, 1:(batch_size_settings.N)]
+end
+
+function stack_batches_row(
+        blocks::AbstractVector, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
+    ) where {B, singlebatch, aligned}
+    M = reduce(vcat, blocks)
+    return aligned ? M : M[1:(batch_size_settings.N), :]
+end
+
 """
     ismutable_array(x)
 
diff --git a/DifferentiationInterface/test/Back/ForwardDiff/test.jl b/DifferentiationInterface/test/Back/ForwardDiff/test.jl
index c829e0555..c272b971d 100644
--- a/DifferentiationInterface/test/Back/ForwardDiff/test.jl
+++ b/DifferentiationInterface/test/Back/ForwardDiff/test.jl
@@ -6,7 +6,7 @@ using DifferentiationInterface, DifferentiationInterfaceTest
 import DifferentiationInterface as DI
 import DifferentiationInterfaceTest as DIT
 using ForwardDiff: ForwardDiff
-using StaticArrays: StaticArrays, @SVector
+using StaticArrays: StaticArrays, @MVector, @SMatrix, @SVector, StaticMatrix
 using JLArrays: JLArrays
 using Test
 import JET, Chairmarks
@@ -94,6 +94,20 @@ end
             )
         ) isa DI.BatchSizeSettings{5}
     end
+
+    @testset "Static Hessian with multiple batches" begin
+        f(x) = sum(abs2, x) * sum(x)  # non-diagonal Hessian
+        for x in (@SVector(rand(6)), @SMatrix(rand(2, 3)), @MVector(rand(6)))
+            H_ref = hessian(f, AutoForwardDiff(), Array(x))
+            # B = 1 and unaligned use different methods
+            @testset "Chunksize $chunksize" for chunksize in (1, 3, 4)
+                backend = SecondOrder(DI.AutoSimpleFiniteDiff(; chunksize), AutoForwardDiff())
+                H = @inferred hessian(f, backend, x)
+                @test H isa StaticMatrix
+                @test H ≈ H_ref rtol = 1.0e-5
+            end
+        end
+    end
 end
 
 @testset verbose = true "Overloaded inputs" begin
diff --git a/DifferentiationInterface/test/Core/Internals/batchsize.jl b/DifferentiationInterface/test/Core/Internals/batchsize.jl
index 9a08aa5c1..3f4c63739 100644
--- a/DifferentiationInterface/test/Core/Internals/batchsize.jl
+++ b/DifferentiationInterface/test/Core/Internals/batchsize.jl
@@ -87,12 +87,27 @@ end
             @inferred prepare_jacobian(identity, AutoSimpleFiniteDiff(; chunksize = 4), x)
         ) isa DI.PushforwardJacobianPrep{<:Any, BSS{4, false, false}}
         @test (@inferred jacobian(identity, AutoSimpleFiniteDiff(), x)) isa StaticMatrix
-        # multiple batches are assembled with `mapreduce(hcat, ...)`, which changes the
-        # static size at every step
-        @test_broken (
-            @inferred jacobian(identity, AutoSimpleFiniteDiff(; chunksize = 4), x)
-        ) isa StaticMatrix
     end
+
+    # multiple batches must be concatenated without losing static sizes
+    f(x) = vcat(vec(x), vec(x) .^ 2)  # non-square Jacobian
+    for x in (@SVector(rand(6)), @SMatrix(rand(2, 3)), @MVector(rand(6)))
+        J_ref = jacobian(f, AutoSimpleFiniteDiff(), Array(x))
+        @testset "Pushforward, chunksize $chunksize" for chunksize in (3, 4)  # aligned, unaligned
+            J = @inferred jacobian(f, AutoSimpleFiniteDiff(; chunksize), x)
+            @test J isa StaticMatrix
+            @test J ≈ J_ref
+        end
+        @testset "Pullback, chunksize $chunksize" for chunksize in (4, 5)  # aligned, unaligned
+            backend = DI.AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize))
+            J = jacobian(f, backend, x)
+            @test J isa StaticMatrix
+            @test J ≈ J_ref
+        end
+    end
+    @test_throws ArgumentError("Cannot concatenate zero batches") jacobian(
+        identity, AutoSimpleFiniteDiff(; chunksize = 2), @SVector(zeros(0))
+    )
 end
 
 @testset "Thresholding" begin
diff --git a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
index 1ff483e17..321f7195f 100644
--- a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
+++ b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
@@ -1,4 +1,5 @@
 using DifferentiationInterface, DifferentiationInterfaceTest
+import DifferentiationInterfaceTest as DIT
 using DifferentiationInterface:
     AutoSimpleFiniteDiff,
     AutoForwardFromPrimitive,
@@ -168,6 +169,15 @@ end
         vcat(static_scenarios(), gpu_scenarios());
         logging = LOGGING,
     )
+    # multiple batches
+    test_differentiation(
+        [
+            AutoSimpleFiniteDiff(; chunksize = 2),
+            AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize = 2)),
+        ],
+        filter(s -> DIT.operator(s) in (:jacobian, :hessian), static_scenarios());
+        logging = LOGGING,
+    )
 end;
 
 @testset "Array format preservation in wrong mode" begin

From 3209372b644d70ca76c4ff7203064639472f2bc1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Mon, 19 Oct 2026 16:45:07 +0200
Subject: [PATCH 3/5] refactor: address review on multi-batch assembly

- split `aligned ? M : M[...]` into if/else blocks for coverage
- drop the unused `B_last` field from `BatchSizeSettings`
- route the B=1 Jacobian methods through `stack_batches_col/row` too,
  so they also return static matrices for static arrays

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../DifferentiationInterfaceStaticArraysExt.jl | 18 ++++++++++++++----
 .../src/first_order/jacobian.jl                | 17 +++++++++--------
 .../src/utils/batchsize.jl                     | 11 ++---------
 DifferentiationInterface/src/utils/linalg.jl   | 12 ++++++++++--
 .../test/Core/Internals/batchsize.jl           | 17 +++++++++--------
 5 files changed, 44 insertions(+), 31 deletions(-)

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
index 309ae8199..ff27f8bd2 100644
--- a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
+++ b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
@@ -2,7 +2,8 @@ module DifferentiationInterfaceStaticArraysExt
 
 using ADTypes: AutoForwardDiff, AutoEnzyme
 import DifferentiationInterface as DI
-using StaticArrays: SArray, SOneTo, StaticArray
+using LinearAlgebra: Transpose
+using StaticArrays: SArray, SOneTo, StaticArray, StaticVector
 
 function DI.stack_vec_col(t::NTuple{B, <:StaticArray}) where {B}
     return hcat(map(vec, t)...)
@@ -30,17 +31,26 @@ function DI.stack_batches_col(
     ) where {B, singlebatch, aligned, T <: StaticArray}
     N, A = static_batch_count(T, B)
     M = hcat(ntuple(a -> blocks[a], Val(A))...)
-    return aligned ? M : M[:, SOneTo(N)]
+    if aligned
+        return M
+    else
+        return M[:, SOneTo(N)]
+    end
 end
 
+# rows are either blocks of a batch or transposed vectors for batch size one
 function DI.stack_batches_row(
-        blocks::AbstractVector{<:StaticArray},
+        blocks::AbstractVector{<:Union{StaticArray, Transpose{<:Any, <:StaticVector}}},
         ::DI.BatchSizeSettings{B, singlebatch, aligned},
         ::NTuple{B, T},
     ) where {B, singlebatch, aligned, T <: StaticArray}
     N, A = static_batch_count(T, B)
     M = vcat(ntuple(a -> blocks[a], Val(A))...)
-    return aligned ? M : M[SOneTo(N), :]
+    if aligned
+        return M
+    else
+        return M[SOneTo(N), :]
+    end
 end
 
 DI.ismutable_array(::Type{<:SArray}) = false
diff --git a/DifferentiationInterface/src/first_order/jacobian.jl b/DifferentiationInterface/src/first_order/jacobian.jl
index 8b261878c..ede03aa81 100644
--- a/DifferentiationInterface/src/first_order/jacobian.jl
+++ b/DifferentiationInterface/src/first_order/jacobian.jl
@@ -373,13 +373,13 @@ function _jacobian_aux(
         x,
         contexts::Vararg{Context, C},
     ) where {FY, SIG, C}
-    (; batched_seeds, seed_example, pushforward_prep) = prep
+    (; batch_size_settings, batched_seeds, seed_example, pushforward_prep) = prep
 
     pushforward_prep_same = prepare_pushforward_same_point(
         f_or_f!y..., pushforward_prep, backend, x, seed_example, contexts...
     )
 
-    jac = stack(eachindex(batched_seeds); dims = 2) do a
+    blocks = map(eachindex(batched_seeds)) do a
         dy = only(
             pushforward(
                 f_or_f!y...,
@@ -392,7 +392,7 @@ function _jacobian_aux(
         )
         return vec(dy)
     end
-    return jac
+    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux(
@@ -446,25 +446,26 @@ function _jacobian_aux(
         x,
         contexts::Vararg{Context, C},
     ) where {FY, SIG, C}
-    (; batched_seeds, seed_example, pullback_prep) = prep
+    (; batch_size_settings, batched_seeds, seed_example, pullback_prep) = prep
 
     pullback_prep_same = prepare_pullback_same_point(
         f_or_f!y..., pullback_prep, backend, x, seed_example, contexts...
     )
 
-    jac = stack(eachindex(batched_seeds); dims = 1) do a
+    blocks = map(eachindex(batched_seeds)) do a
         dx = only(
             pullback(
                 f_or_f!y..., pullback_prep_same, backend, x, batched_seeds[a], contexts...
             )
         )
+        # each block is a row of the Jacobian
         if eltype(x) <: Complex
-            return map(conj, vec(dx))
+            return transpose(map(conj, vec(dx)))
         else
-            return vec(dx)
+            return transpose(vec(dx))
         end
     end
-    return jac
+    return stack_batches_row(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux(
diff --git a/DifferentiationInterface/src/utils/batchsize.jl b/DifferentiationInterface/src/utils/batchsize.jl
index 64c4288b4..20f8a03c6 100644
--- a/DifferentiationInterface/src/utils/batchsize.jl
+++ b/DifferentiationInterface/src/utils/batchsize.jl
@@ -13,23 +13,16 @@ Configuration for the batch size deduced from a backend and a sample array of le
 
   - `N::Int`: array length
   - `A::Int`: number of batches `A = div(N, B, RoundUp)`
-  - `B_last::Int`: size of the last batch (if `aligned` is `false`)
 """
 struct BatchSizeSettings{B, singlebatch, aligned}
     N::Int
     A::Int
-    B_last::Int
 end
 
 function BatchSizeSettings{B, singlebatch, aligned}(N::Integer) where {B, singlebatch, aligned}
     B > N > 0 && throw(ArgumentError("Batch size $B larger than input size $N"))
-    if B == N == 0
-        A = B_last = 0
-    else
-        A = div(N, B, RoundUp)
-        B_last = N % B
-    end
-    return BatchSizeSettings{B, singlebatch, aligned}(N, A, B_last)
+    A = B == N == 0 ? 0 : div(N, B, RoundUp)
+    return BatchSizeSettings{B, singlebatch, aligned}(N, A)
 end
 
 function BatchSizeSettings(B::Int, N::Int)
diff --git a/DifferentiationInterface/src/utils/linalg.jl b/DifferentiationInterface/src/utils/linalg.jl
index 5d591e21e..4d963a3d4 100644
--- a/DifferentiationInterface/src/utils/linalg.jl
+++ b/DifferentiationInterface/src/utils/linalg.jl
@@ -15,14 +15,22 @@ function stack_batches_col(
         blocks::AbstractVector, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
     ) where {B, singlebatch, aligned}
     M = reduce(hcat, blocks)
-    return aligned ? M : M[:, 1:(batch_size_settings.N)]
+    if aligned
+        return M
+    else
+        return M[:, 1:(batch_size_settings.N)]
+    end
 end
 
 function stack_batches_row(
         blocks::AbstractVector, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
     ) where {B, singlebatch, aligned}
     M = reduce(vcat, blocks)
-    return aligned ? M : M[1:(batch_size_settings.N), :]
+    if aligned
+        return M
+    else
+        return M[1:(batch_size_settings.N), :]
+    end
 end
 
 """
diff --git a/DifferentiationInterface/test/Core/Internals/batchsize.jl b/DifferentiationInterface/test/Core/Internals/batchsize.jl
index 3f4c63739..269814b44 100644
--- a/DifferentiationInterface/test/Core/Internals/batchsize.jl
+++ b/DifferentiationInterface/test/Core/Internals/batchsize.jl
@@ -26,12 +26,12 @@ BSS = BatchSizeSettings
 end
 
 @testset "Constructor" begin
-    @test BSS(1, 0) === BSS{1, false, true}(0, 0, 0)
-    @test BSS(0, 0) === BSS{0, true, true}(0, 0, 0)
-    @test BSS(1, 1) === BSS{1, true, true}(1, 1, 0)
-    @test BSS(2, 2) === BSS{2, true, true}(2, 1, 0)
-    @test BSS(4, 8) === BSS{4, false, true}(8, 2, 0)
-    @test BSS(8, 13) === BSS{8, false, false}(13, 2, 5)
+    @test BSS(1, 0) === BSS{1, false, true}(0, 0)
+    @test BSS(0, 0) === BSS{0, true, true}(0, 0)
+    @test BSS(1, 1) === BSS{1, true, true}(1, 1)
+    @test BSS(2, 2) === BSS{2, true, true}(2, 1)
+    @test BSS(4, 8) === BSS{4, false, true}(8, 2)
+    @test BSS(8, 13) === BSS{8, false, false}(13, 2)
     @test_throws ArgumentError BSS(3, 2)
 end
 
@@ -93,12 +93,13 @@ end
     f(x) = vcat(vec(x), vec(x) .^ 2)  # non-square Jacobian
     for x in (@SVector(rand(6)), @SMatrix(rand(2, 3)), @MVector(rand(6)))
         J_ref = jacobian(f, AutoSimpleFiniteDiff(), Array(x))
-        @testset "Pushforward, chunksize $chunksize" for chunksize in (3, 4)  # aligned, unaligned
+        # B = 1, aligned, unaligned use different methods
+        @testset "Pushforward, chunksize $chunksize" for chunksize in (1, 3, 4)
             J = @inferred jacobian(f, AutoSimpleFiniteDiff(; chunksize), x)
             @test J isa StaticMatrix
             @test J ≈ J_ref
         end
-        @testset "Pullback, chunksize $chunksize" for chunksize in (4, 5)  # aligned, unaligned
+        @testset "Pullback, chunksize $chunksize" for chunksize in (1, 4, 5)
             backend = DI.AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize))
             J = jacobian(f, backend, x)
             @test J isa StaticMatrix

From dcf6496be7b303e146a0bba83e9d5ec971576495 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Fri, 30 Oct 2026 17:04:28 +0100
Subject: [PATCH 4/5] fix: infer multi-batch assembly and avoid SizedArray
 blocks

Pass a block function to `stack_batches_col/row` instead of a vector of
blocks. The generic method collects the blocks with a comprehension,
since `map` fails to infer when the block function calls `map` itself
(e.g. pullbacks via pushforwards). The StaticArrays methods build the
blocks with `ntuple` and allocate no intermediate vector.

Flatten blocks with `vec_block`, which uses `similar_type` for static
arrays: `vec` of an `MMatrix` is a `SizedVector` view, and concatenating
those allocates a `SizedMatrix` backed by a `Matrix`.

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, with inference checks, a ForwardDiff Jacobian
through `AutoForward/ReverseFromPrimitive`, and JET checks over the
static DIT scenarios. The in-place Hessian with a pullback-based inner
gradient is marked broken, as it already failed to infer before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ...DifferentiationInterfaceStaticArraysExt.jl | 27 ++++----
 .../src/first_order/jacobian.jl               | 18 +++---
 .../src/second_order/hessian.jl               |  8 +--
 DifferentiationInterface/src/utils/linalg.jl  | 26 +++++---
 .../test/Back/ForwardDiff/test.jl             | 37 ++++++++++-
 .../test/Core/Internals/batchsize.jl          | 62 ++++++++++++++++---
 .../test/Core/SimpleFiniteDiff/test.jl        | 43 +++++++++++--
 7 files changed, 164 insertions(+), 57 deletions(-)

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
index ff27f8bd2..1076774b8 100644
--- a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
+++ b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
@@ -2,15 +2,17 @@ module DifferentiationInterfaceStaticArraysExt
 
 using ADTypes: AutoForwardDiff, AutoEnzyme
 import DifferentiationInterface as DI
-using LinearAlgebra: Transpose
-using StaticArrays: SArray, SOneTo, StaticArray, StaticVector
+using StaticArrays: SArray, Size, SOneTo, StaticArray, similar_type
+
+# a flat vector of the same kind (e.g. `MVector` for an `MMatrix`) can be concatenated without heap allocations
+DI.vec_block(a::StaticArray) = similar_type(a, Size(length(a)))(a)
 
 function DI.stack_vec_col(t::NTuple{B, <:StaticArray}) where {B}
-    return hcat(map(vec, t)...)
+    return hcat(map(DI.vec_block, t)...)
 end
 
 function DI.stack_vec_row(t::NTuple{B, <:StaticArray}) where {B}
-    return vcat(transpose.(map(vec, t))...)
+    return vcat(transpose.(map(DI.vec_block, t))...)
 end
 
 # The number of seeds `N` and of batches `A` are deduced from the seed type,
@@ -25,12 +27,10 @@ function static_batch_count(::Type{T}, B::Int) where {T <: StaticArray}
 end
 
 function DI.stack_batches_col(
-        blocks::AbstractVector{<:StaticArray},
-        ::DI.BatchSizeSettings{B, singlebatch, aligned},
-        ::NTuple{B, T},
-    ) where {B, singlebatch, aligned, T <: StaticArray}
+        block::F, ::DI.BatchSizeSettings{B, singlebatch, aligned}, ::NTuple{B, T}
+    ) where {F, B, singlebatch, aligned, T <: StaticArray}
     N, A = static_batch_count(T, B)
-    M = hcat(ntuple(a -> blocks[a], Val(A))...)
+    M = hcat(ntuple(block, Val(A))...)
     if aligned
         return M
     else
@@ -38,14 +38,11 @@ function DI.stack_batches_col(
     end
 end
 
-# rows are either blocks of a batch or transposed vectors for batch size one
 function DI.stack_batches_row(
-        blocks::AbstractVector{<:Union{StaticArray, Transpose{<:Any, <:StaticVector}}},
-        ::DI.BatchSizeSettings{B, singlebatch, aligned},
-        ::NTuple{B, T},
-    ) where {B, singlebatch, aligned, T <: StaticArray}
+        block::F, ::DI.BatchSizeSettings{B, singlebatch, aligned}, ::NTuple{B, T}
+    ) where {F, B, singlebatch, aligned, T <: StaticArray}
     N, A = static_batch_count(T, B)
-    M = vcat(ntuple(a -> blocks[a], Val(A))...)
+    M = vcat(ntuple(block, Val(A))...)
     if aligned
         return M
     else
diff --git a/DifferentiationInterface/src/first_order/jacobian.jl b/DifferentiationInterface/src/first_order/jacobian.jl
index ede03aa81..9e1cfc88e 100644
--- a/DifferentiationInterface/src/first_order/jacobian.jl
+++ b/DifferentiationInterface/src/first_order/jacobian.jl
@@ -379,7 +379,7 @@ function _jacobian_aux(
         f_or_f!y..., pushforward_prep, backend, x, seed_example, contexts...
     )
 
-    blocks = map(eachindex(batched_seeds)) do a
+    return stack_batches_col(batch_size_settings, seed_example) do a
         dy = only(
             pushforward(
                 f_or_f!y...,
@@ -390,9 +390,8 @@ function _jacobian_aux(
                 contexts...,
             )
         )
-        return vec(dy)
+        return vec_block(dy)
     end
-    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux(
@@ -408,7 +407,7 @@ function _jacobian_aux(
         f_or_f!y..., pushforward_prep, backend, x, seed_example, contexts...
     )
 
-    blocks = map(eachindex(batched_seeds)) do a
+    return stack_batches_col(batch_size_settings, seed_example) do a
         dy_batch = pushforward(
             f_or_f!y...,
             pushforward_prep_same,
@@ -419,7 +418,6 @@ function _jacobian_aux(
         )
         return stack_vec_col(dy_batch)
     end
-    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux(
@@ -452,7 +450,7 @@ function _jacobian_aux(
         f_or_f!y..., pullback_prep, backend, x, seed_example, contexts...
     )
 
-    blocks = map(eachindex(batched_seeds)) do a
+    return stack_batches_row(batch_size_settings, seed_example) do a
         dx = only(
             pullback(
                 f_or_f!y..., pullback_prep_same, backend, x, batched_seeds[a], contexts...
@@ -460,12 +458,11 @@ function _jacobian_aux(
         )
         # each block is a row of the Jacobian
         if eltype(x) <: Complex
-            return transpose(map(conj, vec(dx)))
+            return transpose(map(conj, vec_block(dx)))
         else
-            return transpose(vec(dx))
+            return transpose(vec_block(dx))
         end
     end
-    return stack_batches_row(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux(
@@ -481,7 +478,7 @@ function _jacobian_aux(
         f_or_f!y..., pullback_prep, backend, x, seed_example, contexts...
     )
 
-    blocks = map(eachindex(batched_seeds)) do a
+    return stack_batches_row(batch_size_settings, seed_example) do a
         dx_batch = pullback(
             f_or_f!y..., pullback_prep_same, backend, x, batched_seeds[a], contexts...
         )
@@ -490,7 +487,6 @@ function _jacobian_aux(
         end
         return stack_vec_row(dx_batch)
     end
-    return stack_batches_row(blocks, batch_size_settings, seed_example)
 end
 
 function _jacobian_aux!(
diff --git a/DifferentiationInterface/src/second_order/hessian.jl b/DifferentiationInterface/src/second_order/hessian.jl
index 8c5c99d37..d2577acf5 100644
--- a/DifferentiationInterface/src/second_order/hessian.jl
+++ b/DifferentiationInterface/src/second_order/hessian.jl
@@ -165,11 +165,10 @@ function hessian(
         f, hvp_prep, backend, x, seed_example, contexts...
     )
 
-    blocks = map(eachindex(batched_seeds)) do a
+    return stack_batches_col(batch_size_settings, seed_example) do a
         dg = only(hvp(f, hvp_prep_same, backend, x, batched_seeds[a], contexts...))
-        return vec(dg)
+        return vec_block(dg)
     end
-    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function hessian(
@@ -186,11 +185,10 @@ function hessian(
         f, hvp_prep, backend, x, seed_example, contexts...
     )
 
-    blocks = map(eachindex(batched_seeds)) do a
+    return stack_batches_col(batch_size_settings, seed_example) do a
         dg_batch = hvp(f, hvp_prep_same, backend, x, batched_seeds[a], contexts...)
         return stack_vec_col(dg_batch)
     end
-    return stack_batches_col(blocks, batch_size_settings, seed_example)
 end
 
 function hessian!(
diff --git a/DifferentiationInterface/src/utils/linalg.jl b/DifferentiationInterface/src/utils/linalg.jl
index 4d963a3d4..85520fd71 100644
--- a/DifferentiationInterface/src/utils/linalg.jl
+++ b/DifferentiationInterface/src/utils/linalg.jl
@@ -1,20 +1,28 @@
+"""
+    vec_block(a)
+
+Flatten `a` into a vector suitable for concatenation into a Jacobian or Hessian block.
+"""
+vec_block(a) = vec(a)
+
 stack_vec_col(t::NTuple) = stack(vec, t; dims = 2)
 stack_vec_row(t::NTuple) = stack(vec, t; dims = 1)
 
 """
-    stack_batches_col(blocks, batch_size_settings, seed_example)
-    stack_batches_row(blocks, batch_size_settings, seed_example)
+    stack_batches_col(block, batch_size_settings, seed_example)
+    stack_batches_row(block, batch_size_settings, seed_example)
 
-Concatenate the blocks computed from all batches along columns (resp. rows), and drop the trailing entries of the last batch if it was padded by wrapping around the seeds.
+Concatenate the blocks `block(a)` computed from all batches `a` along columns (resp. rows), and drop the trailing entries of the last batch if it was padded by wrapping around the seeds.
 
 The `seed_example` is only used for dispatch, e.g. to deduce the number of seeds statically.
 
 Unlike `mapreduce(hcat, ...)`, `reduce(hcat, ...)` allocates the result only once.
 """
 function stack_batches_col(
-        blocks::AbstractVector, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
-    ) where {B, singlebatch, aligned}
-    M = reduce(hcat, blocks)
+        block::F, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
+    ) where {F, B, singlebatch, aligned}
+    # a comprehension, since `map` fails to infer when `block` calls `map` itself
+    M = reduce(hcat, [block(a) for a in 1:(batch_size_settings.A)])
     if aligned
         return M
     else
@@ -23,9 +31,9 @@ function stack_batches_col(
 end
 
 function stack_batches_row(
-        blocks::AbstractVector, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
-    ) where {B, singlebatch, aligned}
-    M = reduce(vcat, blocks)
+        block::F, batch_size_settings::BatchSizeSettings{B, singlebatch, aligned}, _
+    ) where {F, B, singlebatch, aligned}
+    M = reduce(vcat, [block(a) for a in 1:(batch_size_settings.A)])
     if aligned
         return M
     else
diff --git a/DifferentiationInterface/test/Back/ForwardDiff/test.jl b/DifferentiationInterface/test/Back/ForwardDiff/test.jl
index c272b971d..e8e4b0bce 100644
--- a/DifferentiationInterface/test/Back/ForwardDiff/test.jl
+++ b/DifferentiationInterface/test/Back/ForwardDiff/test.jl
@@ -95,12 +95,45 @@ end
         ) isa DI.BatchSizeSettings{5}
     end
 
+    @testset "Static Jacobian with multiple batches" begin
+        f(x) = vcat(vec(x), vec(x) .^ 2)  # non-square Jacobian
+        function f!(y, x)
+            y[1:length(x)] .= vec(x)
+            y[(length(x) + 1):end] .= vec(x) .^ 2
+            return nothing
+        end
+        # a plain `AutoForwardDiff` calls `ForwardDiff.jacobian` and skips the batch assembly
+        @testset "$(typeof(x))" for x in (@SVector(rand(6)), @SMatrix(rand(2, 3)), @MVector(rand(6)))
+            y = StaticArrays.MVector{12, Float64}(undef)
+            J_ref = jacobian(f, AutoForwardDiff(), Array(x))
+            # B = 1, aligned, unaligned and single batch use different methods
+            @testset "Pushforward, chunksize $chunksize" for chunksize in (1, 3, 4, 6)
+                backend = DI.AutoForwardFromPrimitive(AutoForwardDiff(; chunksize))
+                J = @inferred jacobian(f, backend, x)
+                @test J isa StaticMatrix
+                @test J ≈ J_ref
+                J = @inferred jacobian(f!, y, backend, x)
+                @test J isa StaticMatrix
+                @test J ≈ J_ref
+            end
+            @testset "Pullback, chunksize $chunksize" for chunksize in (1, 4, 5, 12)
+                backend = DI.AutoReverseFromPrimitive(AutoForwardDiff(; chunksize))
+                J = @inferred jacobian(f, backend, x)
+                @test J isa StaticMatrix
+                @test J ≈ J_ref
+                J = @inferred jacobian(f!, y, backend, x)
+                @test J isa StaticMatrix
+                @test J ≈ J_ref
+            end
+        end
+    end
+
     @testset "Static Hessian with multiple batches" begin
         f(x) = sum(abs2, x) * sum(x)  # non-diagonal Hessian
         for x in (@SVector(rand(6)), @SMatrix(rand(2, 3)), @MVector(rand(6)))
             H_ref = hessian(f, AutoForwardDiff(), Array(x))
-            # B = 1 and unaligned use different methods
-            @testset "Chunksize $chunksize" for chunksize in (1, 3, 4)
+            # B = 1, unaligned and single batch use different methods
+            @testset "Chunksize $chunksize" for chunksize in (1, 3, 4, 6)
                 backend = SecondOrder(DI.AutoSimpleFiniteDiff(; chunksize), AutoForwardDiff())
                 H = @inferred hessian(f, backend, x)
                 @test H isa StaticMatrix
diff --git a/DifferentiationInterface/test/Core/Internals/batchsize.jl b/DifferentiationInterface/test/Core/Internals/batchsize.jl
index 269814b44..49990e25c 100644
--- a/DifferentiationInterface/test/Core/Internals/batchsize.jl
+++ b/DifferentiationInterface/test/Core/Internals/batchsize.jl
@@ -7,6 +7,7 @@ using DifferentiationInterface:
     reasonable_batchsize,
     threshold_batchsize
 import DifferentiationInterface as DI
+using LinearAlgebra: I
 using StaticArrays
 using Test
 
@@ -89,22 +90,65 @@ end
         @test (@inferred jacobian(identity, AutoSimpleFiniteDiff(), x)) isa StaticMatrix
     end
 
-    # multiple batches must be concatenated without losing static sizes
+    # multiple batches must be concatenated without losing static sizes or inference
     f(x) = vcat(vec(x), vec(x) .^ 2)  # non-square Jacobian
-    for x in (@SVector(rand(6)), @SMatrix(rand(2, 3)), @MVector(rand(6)))
+    function f!(y, x)
+        y[1:length(x)] .= vec(x)
+        y[(length(x) + 1):end] .= vec(x) .^ 2
+        return nothing
+    end
+    g(x) = sum(abs2, x) * sum(x)  # non-diagonal Hessian
+    has_type(J, x) = x isa StaticArray ? J isa Union{SArray, MArray} : J isa Matrix
+    # without preparation, the batch size is only inferred for static arrays
+    function inferred_jacobian(f_or_f!y, backend, x)
+        x isa StaticArray && @inferred jacobian(f_or_f!y..., backend, x)
+        prep = prepare_jacobian(f_or_f!y..., backend, x)
+        return @inferred jacobian(f_or_f!y..., prep, backend, x)
+    end
+    function inferred_hessian(f, backend, x)
+        x isa StaticArray && @inferred hessian(f, backend, x)
+        prep = prepare_hessian(f, backend, x)
+        return @inferred hessian(f, prep, backend, x)
+    end
+    @testset "$(typeof(x))" for x in (
+            @SVector(rand(6)),
+            @SMatrix(rand(2, 3)),
+            @MVector(rand(6)),
+            rand(6),
+            @SVector(rand(ComplexF64, 6)),
+        )
+        y = x isa StaticArray ? MVector{12, eltype(x)}(undef) : similar(x, 12)
         J_ref = jacobian(f, AutoSimpleFiniteDiff(), Array(x))
-        # B = 1, aligned, unaligned use different methods
-        @testset "Pushforward, chunksize $chunksize" for chunksize in (1, 3, 4)
-            J = @inferred jacobian(f, AutoSimpleFiniteDiff(; chunksize), x)
-            @test J isa StaticMatrix
+        # B = 1, aligned, unaligned and single batch use different methods
+        @testset "Pushforward, chunksize $chunksize" for chunksize in (1, 3, 4, 6)
+            backend = AutoSimpleFiniteDiff(; chunksize)
+            J = inferred_jacobian((f,), backend, x)
+            @test has_type(J, x)
+            @test J ≈ J_ref
+            J = inferred_jacobian((f!, y), backend, x)
+            @test has_type(J, x)
             @test J ≈ J_ref
         end
-        @testset "Pullback, chunksize $chunksize" for chunksize in (1, 4, 5)
+        @testset "Pullback, chunksize $chunksize" for chunksize in (1, 4, 5, 12)
             backend = DI.AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize))
-            J = jacobian(f, backend, x)
-            @test J isa StaticMatrix
+            J = inferred_jacobian((f,), backend, x)
+            @test has_type(J, x)
+            @test J ≈ J_ref
+            J = inferred_jacobian((f!, y), backend, x)
+            @test has_type(J, x)
             @test J ≈ J_ref
         end
+        if eltype(x) <: Real
+            H_ref = 2 * (sum(x) * I + (vec(x) .+ vec(x)'))
+            @testset "Hessian, chunksize $chunksize" for chunksize in (1, 3, 4, 6)
+                backend = SecondOrder(
+                    AutoSimpleFiniteDiff(; chunksize), AutoSimpleFiniteDiff()
+                )
+                H = inferred_hessian(g, backend, x)
+                @test has_type(H, x)
+                @test H ≈ H_ref rtol = 1.0e-4
+            end
+        end
     end
     @test_throws ArgumentError("Cannot concatenate zero batches") jacobian(
         identity, AutoSimpleFiniteDiff(; chunksize = 2), @SVector(zeros(0))
diff --git a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
index 321f7195f..95185c024 100644
--- a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
+++ b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
@@ -7,6 +7,7 @@ using DifferentiationInterface:
     DenseSparsityDetector
 using SparseMatrixColorings
 using JLArrays, StaticArrays
+using JET: JET
 using Test
 
 backends = [ #
@@ -170,14 +171,44 @@ end
         logging = LOGGING,
     )
     # multiple batches
+    forward_backend = AutoSimpleFiniteDiff(; chunksize = 2)
+    reverse_backend = AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize = 2))
+    multibatch_scenarios = filter(
+        s -> DIT.operator(s) in (:jacobian, :hessian), static_scenarios()
+    )
     test_differentiation(
-        [
-            AutoSimpleFiniteDiff(; chunksize = 2),
-            AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize = 2)),
-        ],
-        filter(s -> DIT.operator(s) in (:jacobian, :hessian), static_scenarios());
-        logging = LOGGING,
+        [forward_backend, reverse_backend], multibatch_scenarios; logging = LOGGING
     )
+
+    # JET < 0.12 does not define `JET_AVAILABLE`
+    if !isdefined(JET, :JET_AVAILABLE) || JET.JET_AVAILABLE
+        inplace_hessian(s) = DIT.operator(s) == :hessian && DIT.operator_place(s) == :in
+        test_differentiation(
+            forward_backend,
+            multibatch_scenarios;
+            correctness = false,
+            type_stability = :prepared,
+            logging = LOGGING,
+        )
+        test_differentiation(
+            reverse_backend,
+            filter(!inplace_hessian, multibatch_scenarios);
+            correctness = false,
+            type_stability = :prepared,
+            logging = LOGGING,
+        )
+        @testset "Type stability (broken)" for scen in filter(inplace_hessian, multibatch_scenarios)
+            # the inner pullback-based gradient fails to infer in `hessian!`
+            (; f, x, res1, res2, prep_args) = deepcopy(scen)
+            prep = prepare_hessian(f, reverse_backend, prep_args.x)
+            JET.@test_opt broken = true hessian!(
+                f, DIT.mysimilar(res2), prep, reverse_backend, x
+            )
+            JET.@test_opt broken = true value_gradient_and_hessian!(
+                f, DIT.mysimilar(res1), DIT.mysimilar(res2), prep, reverse_backend, x
+            )
+        end
+    end
 end;
 
 @testset "Array format preservation in wrong mode" begin

From f46047475be7cbdb6ec84dbbc13383a2c8ec7a3c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20M=C3=BCller-Widmann?=
 <devmotion@users.noreply.github.com>
Date: Fri, 30 Oct 2026 17:13:40 +0100
Subject: [PATCH 5/5] chore: clarify comments and pass contexts in broken JET
 checks

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../DifferentiationInterfaceStaticArraysExt.jl    |  2 +-
 .../test/Core/SimpleFiniteDiff/test.jl            | 15 ++++++++++-----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
index 1076774b8..eb8157d99 100644
--- a/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
+++ b/DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl
@@ -4,7 +4,7 @@ using ADTypes: AutoForwardDiff, AutoEnzyme
 import DifferentiationInterface as DI
 using StaticArrays: SArray, Size, SOneTo, StaticArray, similar_type
 
-# a flat vector of the same kind (e.g. `MVector` for an `MMatrix`) can be concatenated without heap allocations
+# a flat vector of the same kind (e.g. `MVector` for an `MMatrix`) keeps concatenated blocks `SArray`s or `MArray`s
 DI.vec_block(a::StaticArray) = similar_type(a, Size(length(a)))(a)
 
 function DI.stack_vec_col(t::NTuple{B, <:StaticArray}) where {B}
diff --git a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
index 95185c024..05e9da732 100644
--- a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
+++ b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
@@ -198,14 +198,19 @@ end
             logging = LOGGING,
         )
         @testset "Type stability (broken)" for scen in filter(inplace_hessian, multibatch_scenarios)
-            # the inner pullback-based gradient fails to infer in `hessian!`
-            (; f, x, res1, res2, prep_args) = deepcopy(scen)
-            prep = prepare_hessian(f, reverse_backend, prep_args.x)
+            (; f, x, res1, res2, contexts, prep_args) = deepcopy(scen)
+            prep = prepare_hessian(f, reverse_backend, prep_args.x, prep_args.contexts...)
             JET.@test_opt broken = true hessian!(
-                f, DIT.mysimilar(res2), prep, reverse_backend, x
+                f, DIT.mysimilar(res2), prep, reverse_backend, x, contexts...
             )
             JET.@test_opt broken = true value_gradient_and_hessian!(
-                f, DIT.mysimilar(res1), DIT.mysimilar(res2), prep, reverse_backend, x
+                f,
+                DIT.mysimilar(res1),
+                DIT.mysimilar(res2),
+                prep,
+                reverse_backend,
+                x,
+                contexts...,
             )
         end
     end