f-dangel / curvlinops / 22835615977

91%

main: 92%

push

github

[REF] Split vmap data args into X and loss_args, simplify dims to -1

Separate model input X from loss function arguments loss_args in
make_ggn_vector_product and all downstream callers. This makes the
data flow explicit: f consumes X, c consumes loss_args. Adding extra
loss args (e.g. generator for MC) no longer changes vmap in_dims.

Also replace tuple(p.ndim for p in params) with scalar -1 in all vmap
in_dims/out_dims. A scalar -1 broadcasts to "last axis" per tensor
leaf in the pytree, which is the trailing matrix-column dimension.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

7 of 7 new or added lines in 3 files covered. (100.0%)

2203 of 2414 relevant lines covered (91.26%)

0.91 hits per line