Repository navigation
Conversation
…ic 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>
… 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>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1080 +/- ##
==========================================
+ Coverage 98.21% 98.25% +0.03%
==========================================
Files 143 143
Lines 8246 8243 -3
==========================================
Hits 8099 8099
+ Misses 147 144 -3
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…-assembly Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
gdalle
left a comment
gdalle
left a comment
There was a problem hiding this comment.
Thanks, this looks like a worthwhile improvement! A few things to iron out though
|
|
||
| function hessian( | ||
| f::F, | ||
| prep::HVPGradientHessianPrep{SIG, <:BatchSizeSettings{1, false, true}}, |
There was a problem hiding this comment.
This method wasn't modified in the Jacobian file.
Essentially, there are currently three methods for Jacobian and Hessian:
- single batch
- multiple batches & batch size one
- general case
Which ones do we need to touch here?
There was a problem hiding this comment.
Single batch needs no change: it has one block, and stack_vec_col/row is already overloaded for static arrays. The general and B=1 methods, for both the Jacobian (pushforward and pullback) and the Hessian, now all go through stack_batches_col/row. So the B=1 Jacobians also return static matrices, which the static tests now check with chunksize = 1.
- 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>
Stacked on #1078.
Replaces
mapreduce(hcat/vcat, ...)in the multi-batch Jacobian and Hessian paths withmap+ a singlereduce(hcat/vcat, ...)through new internal hooks, which the StaticArrays extension overloads. Multi-batch Jacobians and Hessians of static arrays are now inferred and return static matrices (turning #1078's@test_brokeninto passing tests).Why
map+reduceinstead ofmapreducereduce(hcat, ::AbstractVector{<:AbstractVecOrMat})(andvcat) has a specialized Base method that allocates the result once.mapreduce(f, hcat, v)has no such method and falls back to pairwisehcat, re-copying the partial result at every level.Concatenating
Ablocks of sizem × B(Chairmarks):m× (A×B)mapreduce(f, hcat, v)reduce(hcat, map(f, v))End-to-end, prepared
jacobian(identity, prep, backend, x)with densexandchunksize = 12:length(x)AutoSimpleFiniteDiff, 100AutoSimpleFiniteDiff, 1000AutoReverseFromPrimitive(AutoSimpleFiniteDiff), 1000(
stackwas also considered: equivalent for columns, but rows need an extrapermutedimscopy.)🤖 Generated with Claude Code