f-dangel / curvlinops / 23207457263

91%

push

github

[REF] Pass `_model_func` to `make_batch_*`, align `(params, X)` convention (#270)

* [REF] Pass _model_func to make_batch_* functions, align (params, X) convention

- Update all make_batch_* functions (GGN, Hessian, Jacobian, EF, GGN
  diagonal) to accept a functional callable (params_dict, X) -> prediction
  instead of a Module. Call sites pass self._model_func directly,
  eliminating redundant make_functional_call wrapping.
- Same for make_functional_flattened_model_and_loss.
- Align IO collector (with_param_io, with_kfac_io) to (params, x)
  convention, matching the rest of the codebase. Remove the arg-swapping
  wrapper in KFAC/EKFAC FX backends.
- Migrate _check_supports_batched_and_unbatched_inputs to use
  _model_func via partial.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [DOC] Add changelog entry for make_batch_* callable refactor

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [FIX] Address review: fix type aliases, return type, placeholder selection

- Update ParamIOFunction/KFACIOFunction type aliases to (params, x) order
- Fix return type of make_batch_transposed_jacobian_matrix_product to
  include MutableMapping
- Use placeholders[:len(named_params)] instead of placeholders[:-1] to
  correctly handle pytree x inputs

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: f-dangel <felix.dangel@concordia.ca>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

17 of 17 new or added lines in 8 files covered. (100.0%)

2329 of 2552 relevant lines covered (91.26%)

0.91 hits per line