Fixed docstrings to always use numpy format, and to always document init functions after function definition.
For TRAK, the X^t X tensor from the paper might not be invertible. this is why they added the option of adding a constant diagonal regularization matrix before matrix inversion. i added this option to our wrapper, a very straightforward change.
Fixed docstrings to always use numpy format, and to always document init functions after function definition.
For TRAK, the X^t X tensor from the paper might not be invertible. this is why they added the option of adding a constant diagonal regularization matrix before matrix inversion. i added this option to our wrapper, a very straightforward change.