fix: type-stable multi-batch Jacobian and Hessian assembly for static arrays

DifferentiationInterface/ext/DifferentiationInterfaceStaticArraysExt/DifferentiationInterfaceStaticArraysExt.jl

@@ -2,14 +2,52 @@ module DifferentiationInterfaceStaticArraysExt
 
 using ADTypes: AutoForwardDiff, AutoEnzyme
 import DifferentiationInterface as DI
-using StaticArrays: SArray, StaticArray
+using StaticArrays: SArray, Size, SOneTo, StaticArray, similar_type
+
+# 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}
-    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,
+# 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(
+        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(block, Val(A))...)
+    if aligned
+        return M
+    else
+        return M[:, SOneTo(N)]
+    end
+end
+
+function DI.stack_batches_row(
+        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(block, Val(A))...)
+    if aligned
+        return M
+    else
+        return M[SOneTo(N), :]
+    end
 end
 
 DI.ismutable_array(::Type{<:SArray}) = false

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(
@@ -379,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
+    return stack_batches_col(batch_size_settings, seed_example) do a
         dy = only(
             pushforward(
                 f_or_f!y...,
@@ -396,26 +390,24 @@ function _jacobian_aux(
                 contexts...,
             )
         )
-        return vec(dy)
+        return vec_block(dy)
     end
-    return jac
 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
+    return stack_batches_col(batch_size_settings, seed_example) do a
         dy_batch = pushforward(
             f_or_f!y...,
             pushforward_prep_same,
@@ -424,37 +416,25 @@ 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
 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(
@@ -464,56 +444,49 @@ 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
+    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...
             )
         )
+        # each block is a row of the Jacobian
         if eltype(x) <: Complex
-            return map(conj, vec(dx))
+            return transpose(map(conj, vec_block(dx)))
         else
-            return vec(dx)
+            return transpose(vec_block(dx))
         end
     end
-    return jac
 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
+    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...
         )
         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
 end
 
 function _jacobian_aux!(

DifferentiationInterface/src/second_order/hessian.jl

@@ -159,44 +159,36 @@ 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
+    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 hess
 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
+    return stack_batches_col(batch_size_settings, seed_example) 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
 end
 
 function hessian!(

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)

DifferentiationInterface/src/utils/linalg.jl

@@ -1,6 +1,46 @@
+"""
+    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(block, batch_size_settings, seed_example)
+    stack_batches_row(block, batch_size_settings, seed_example)
+
+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(
+        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
+        return M[:, 1:(batch_size_settings.N)]
+    end
+end
+
+function stack_batches_row(
+        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
+        return M[1:(batch_size_settings.N), :]
+    end
+end
+
 """
     ismutable_array(x)