Fix imports in init

facebookresearch · Jan 7, 2025 · 0a48378 · 0a48378
1 parent c51e4e6
commit 0a48378
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/distributed_shampoo/__init__.py b/distributed_shampoo/__init__.py
@@ -33,9 +33,13 @@
     CoupledHigherOrderConfig,
     CoupledNewtonConfig,
     DefaultEigenConfig,
+    DefaultEighEigenvectorConfig,
     EigenConfig,
+    EigenvalueDecompositionConfig,
     EigenvectorConfig,
+    EighEigenvectorConfig,
     MatrixFunctionConfig,
+    QRConfig,
     RootInvConfig,
 )
 
@@ -66,13 +70,14 @@
     "DefaultSOAPConfig",  # Default `EigenvalueCorrectedShampooPreconditionerConfig` using `QRConfig`.
     # matrix functions configs.
     "MatrixFunctionConfig",  # Abstract base class.
+    "EigenvalueDecompositionConfig",  # Abstract base class (based on `MatrixFunctionConfig`).
     "RootInvConfig",  # Abstract base class (based on `MatrixFunctionConfig`).
-    "EigenConfig",  # Based on `RootInvConfig`.
+    "EigenConfig",  # Based on `RootInvConfig` and `EigenvalueDecompositionConfig`.
     "DefaultEigenConfig",  # Default `RootInvConfig` using `EigenConfig`.
     "CoupledNewtonConfig",  # Based on `RootInvConfig`.
     "CoupledHigherOrderConfig",  # Based on `RootInvConfig`.
     "EigenvectorConfig",  # Abstract base class (based on `MatrixFunctionConfig`).
-    "EighEigenvectorConfig",  # Based on `EigenvectorConfig`.
+    "EighEigenvectorConfig",  # Based on `EigenvectorConfig` and `EigenvalueDecompositionConfig`.
     "DefaultEighEigenvectorConfig",  # Default `EigenvectorConfig` using `EighEigenvectorConfig`.
     "QRConfig",  # Based on `EigenvectorConfig`.
     # Other utilities.