laplax-org / laplax / 21840287371

67%

push

github

Fisher implementation (#59)

* Moved loss hessian mvp from ggn to own .py file

* Wrote fisher mv function that dispatches between fisher types

* Moved loss hessian input logic to curv/loss.py

* Fixed typos in previous commit

* Added gradients for fisher to loss.py

* First idea of empirical fisher calculation

* Reformatted according to Ruff checks

* Ruff formatting

* Wrote tests for BCE gradient

* Wrote test for cross entropy loss gradient

* Fixed bug in cross entropy loss gradient calculation

* Added loss grad test cases for MSE and custom callable

* Ruff formatting

* Started implmntation of emp fisher test case

* Two small fixes in empirical fisher implementation

* Added test case for empirical fisher

* Ruff formatting

* Added function for fisher computation _with_ data

* Fixed inconsistent naming in docstring of create_ggn_mv

* Split empirical and MC fisher entry point into two different functions

* Fixed math in docstrings for empirical fisher

* Defined function for MC fisher mvp

* Implemented scaffolding of mc fisher mv computation

* Reworked emp fisher implementation to vmap data over entire computation

* Cleaned up empirical fisher implementation

* Neat wrapper for jax.linear_transpose when called with linop with only 1 arg

* Test for emp fisher with fn ouput dim > 1

* Fixed small shape bug in emp fisher calculation

* Added emp fisher test for output shape >1

* Extracted some common functionality for empirical and MC fisher

* Found out why test_emp_fisher_on_quadratic_fn failed for parameters away from optimum

* Structure for MC Fisher calculation

* Docstring for the fisher calculation inner function

* Implemented label sampling from MSE loss

* Test for label sampling from MSE loss

* Sampling from CE Loss for MC fisher

* Ruff formatting

* Small change in docstring

* Created function for MC fisher computation _with_ data

* Simplified shapes in fisher computation

* Simplified forward pass calculation a bit

... (continued)

205 of 217 new or added lines in 6 files covered. (94.47%)

1194 of 1777 relevant lines covered (67.19%)

0.67 hits per line