f-dangel / curvlinops / 13650823525

90%

push

github

[ADD] `EKFACLinearOperator` (#173)

* Add EKFAC and its inverse

* Add EKFACLinearOperator to API

* Fix flake8 and darglint

* Fix in-place bug in inverse KFAC and add inverse EKFAC tests

* Fix in-place damping bug, make backwards compatible, and refactor

* Refactor KFAC and EKFAC computation API

* Fix docstrings formatting

* Fix type hint

* Let dict access fail loudly

* Only unfold conv activations in backward pass

* Fix case when no activations are cached

* Simplify conditional

* Make block_diagonal function more general

* Add more KFAC test cases

* Temp fix for tests

* Add _rearrange_for_larger_than_2d_output method, restrict EKFAC to 2d outputs, and support type-2

* Remove testing helpers

* Fix isort

* Set weights_only=False in torch.load to support torch 2.6

* Fix black

* Restrict test cases for EKFAC inverse

* Refactor EKFACLinearOperator into separate file

* Add EKFAC tests

* Fix flake8

* Increase tolerance

* Add KFAC vs. EKFAC approximation quality test

* Minor fixes

* Fix darglint

* Increase MC samples

* Increase MC samples further

* I give up

* Remove superfluous shape argument

* Remove more references to shape

* Modify docs

* Remove outdated warning in docstrings

* Address review comments

* Minor simplifications

* Fix inverse and tests

* Refactor excluding and shuffling params in tests

* Add missing arg

* Address EKFAC review comments

* Remove a few to_scipy calls

* Improve block_diagonal and flatten option in test models

* Remove all numpy/scipy from EKFAC tests

* Fix test function names

* Refactor (E)KFAC property tests

* Fix type hint

* Again fix type hint

* Refactor test_save_and_load_state_dict

* Add support for MC Fisher EKFAC vs KFAC test case

* Add rtol and atol option to _test_property

* Add docstrings

* Replace eye with eye_like

* Avoid MC samples

* Minor fixes

* Refactor _test_ekfac_closer_to_exact_than_kfac

* Attempt to add enums to docs

* Fix imports

* Refactor _test_from_... (continued)

293 of 310 new or added lines in 4 files covered. (94.52%)

1630 of 1730 relevant lines covered (94.22%)

0.94 hits per line