[ADD] Minimal prototype for KFAC (#43)

* [ADD] Prototype for KFAC linear operator

* [DOC] Progress on documentation

* [DOC] Describe KFAC and its limitations

* [FIX] Name of fixture

* [FIX] Darglint

* [FIX] Darglint

See terrencepreilly/darglint#53

* [DOC] Show supported layers in error message

* [DOC] Improve correctness

curvlinops/__init__.py

@@ -7,6 +7,7 @@
 from curvlinops.hessian import HessianLinearOperator
 from curvlinops.inverse import CGInverseLinearOperator, NeumannInverseLinearOperator
 from curvlinops.jacobian import JacobianLinearOperator, TransposedJacobianLinearOperator
+from curvlinops.kfac import KFACLinearOperator
 from curvlinops.papyan2020traces.spectrum import (
     LanczosApproximateLogSpectrumCached,
     LanczosApproximateSpectrumCached,
@@ -22,6 +23,7 @@
     "GGNLinearOperator",
     "EFLinearOperator",
     "FisherMCLinearOperator",
+    "KFACLinearOperator",
     "JacobianLinearOperator",
     "TransposedJacobianLinearOperator",
     "CGInverseLinearOperator",

curvlinops/_base.py

@@ -1,6 +1,7 @@
 """Contains functionality to analyze Hessian & GGN via matrix-free multiplication."""
 
 from typing import Callable, Iterable, List, Optional, Tuple, Union
+from warnings import warn
 
 from backpack.utils.convert_parameters import vector_to_parameter_list
 from numpy import (
@@ -254,6 +255,13 @@ def _preprocess(self, x: ndarray) -> List[Tensor]:
         Returns:
             Vector in list format.
         """
+        if x.dtype != self.dtype:
+            warn(
+                f"Input vector is {x.dtype}, while linear operator is {self.dtype}. "
+                + f"Converting to {self.dtype}."
+            )
+            x = x.astype(self.dtype)
+
         x_torch = from_numpy(x).to(self._device)
         return vector_to_parameter_list(x_torch, self._params)
 

curvlinops/examples/utils.py

@@ -31,5 +31,7 @@ def report_nonclose(
     else:
         for a1, a2 in zip(array1.flatten(), array2.flatten()):
             if not isclose(a1, a2, atol=atol, rtol=rtol, equal_nan=equal_nan):
-                print(f"{a1} ≠ {a2}")
+                print(f"{a1} ≠ {a2} (ratio {a1 / a2:.5f})")
+        print(f"Max: {array1.max():.5f}, {array2.max():.5f}")
+        print(f"Min: {array1.min():.5f}, {array2.min():.5f}")
         raise ValueError("Compared arrays don't match.")