f-dangel / curvlinops / 24036335706

91%

push

github

[OPT] Make (E)KFAC's `matvec` compile-friendly and benchmark (#290)

* [WIP] Make KFAC matvec compile-friendly (0 graph breaks)

- Replace einops.einsum with torch.einsum in KroneckerProductLinearOperator
- Call _matmat instead of @ in BlockDiagonalLinearOperator (dynamo can't
  trace __matmul__ on user-defined objects)
- Replace numpy.cumsum with pure Python in split_list
- Add KFAC compile tests and benchmark entries

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

* [OPT] Make KFAC and EKFAC matvec compile-friendly (0 graph breaks)

Replace einops.einsum with torch.einsum in KroneckerProductLinearOperator
(inductor can't lower einops). Use _matmat/_adjoint_matmat instead of @
in BlockDiagonalLinearOperator and EighDecomposedLinearOperator (dynamo
can't trace __matmul__ on user-defined types). Replace numpy.cumsum in
split_list with pure Python (avoids data-dependent graph breaks).

Add EKFAC to _IS_COMPILABLE set and re-run benchmarks for all problems.
Add docstring note explaining why explain() is used instead of
fullgraph=True in compile tests.

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

* [ADD] EKFAC compile tests and re-run KFAC/EKFAC benchmarks

Add tests verifying 0 graph breaks for EKFAC matvec (hooks and fx).
Re-run KFAC/EKFAC memory benchmarks with corrected protocol (no peak
stats reset). Regenerate compiled time and memory plots.

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

* [DOC] Add changelog entry for KFAC/EKFAC compile support

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

* [FIX] Address Copilot review: Q.T -> Q.mH, fix num_per_example_loss_terms

Use Q.mH instead of Q.T in EighDecomposedLinearOperator for complex
correctness. Fix num_per_example_loss_terms from 1 to 2 in compile
tests (model output dim is 2).

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

* [FIX] Use 'Z' as column subscript in einsum to avoid ... (continued)

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

2424 of 2657 relevant lines covered (91.23%)

0.91 hits per line