f-dangel / curvlinops / 24062381256

91%

main: 91%

Build Type

Pull #292

github

Committed by Felix Dangel

Commit Message

[REF] Add trace_kfac_batch and simplify KFAC FX computer

Introduce trace_kfac_batch() top-level function that combines IO tracing
and batch computation tracing in one call, analogous to
make_batch_hessian_vector_product for Hessian. This eliminates the
two-pass data iteration (separate _trace_io_functions and
_trace_batch_functions) from MakeFxKFACComputer.

Move _trace_io_functions, _trace_batch_functions, and
_compute_layer_output_grads to MakeFxEKFACComputer, which still needs
the two-step pattern for eigenvalue correction.

Also replace FakeTensorMode context managers with make_fx's
_allow_non_fake_inputs parameter, and add einops string comments
to the rearrange replacements in _base.py.

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

Pull Request Pull Request #292: [OPT] Make KFAC FX precompute compile-friendly and reduce tracing memory

Coverage Stats

137 of 137 new or added lines in 6 files covered. (100.0%)

12 existing lines in 3 files now uncovered.

2493 of 2726 relevant lines covered (91.45%)

0.91 hits per line