From 0a48378ac756472c3b39730a8ecddef2d4449751 Mon Sep 17 00:00:00 2001
From: runame <re393@cam.ac.uk>
Date: Tue, 7 Jan 2025 21:24:38 +0800
Subject: [PATCH] Fix imports in init

---
 distributed_shampoo/__init__.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/distributed_shampoo/__init__.py b/distributed_shampoo/__init__.py
index 1e255a0..c850983 100644
--- 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.