f-dangel / curvlinops / 22835830184

91%

main: 92%

push

github

[REF] Refactor vmap convention: X/loss_args split, -1 dims, typed c

Three related improvements to the torch.func convention:

1. Split data arguments into X (model input) and loss_args (loss
   function arguments as a tuple). This makes the data flow explicit:
   f consumes X, c consumes loss_args. Adding extra loss arguments
   (e.g. generator for MC) no longer changes vmap in_dims.

2. 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.

3. Type the loss function c as Callable[[Tensor, tuple], Tensor]
   instead of Callable[..., Tensor]. The c returned by
   make_functional_model_and_loss now natively accepts
   (prediction, loss_args) instead of (prediction, y).

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

13 of 13 new or added lines in 4 files covered. (100.0%)

2208 of 2419 relevant lines covered (91.28%)

0.91 hits per line