From 71df5206720f924931708341a666c5bb44cef16e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 12 Oct 2026 11:37:50 +0200 Subject: [PATCH] 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 --- .../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