Add FisherType and KFACType to init to make them easy to import

f-dangel · Jun 12, 2024 · 08cc6d8 · 08cc6d8
1 parent 188f84e
commit 08cc6d8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/curvlinops/__init__.py b/curvlinops/__init__.py
@@ -12,7 +12,7 @@
     NeumannInverseLinearOperator,
 )
 from curvlinops.jacobian import JacobianLinearOperator, TransposedJacobianLinearOperator
-from curvlinops.kfac import KFACLinearOperator
+from curvlinops.kfac import FisherType, KFACLinearOperator, KFACType
 from curvlinops.norm.hutchinson import HutchinsonSquaredFrobeniusNormEstimator
 from curvlinops.papyan2020traces.spectrum import (
     LanczosApproximateLogSpectrumCached,
@@ -33,6 +33,9 @@
     "KFACLinearOperator",
     "JacobianLinearOperator",
     "TransposedJacobianLinearOperator",
+    # Enums
+    "FisherType",
+    "KFACType",
     # inversion
     "CGInverseLinearOperator",
     "LSMRInverseLinearOperator",