From 12c32f1454e9f391741b7bca21e79da52da77a58 Mon Sep 17 00:00:00 2001
From: runame <re393@cam.ac.uk>
Date: Sat, 11 Apr 2026 20:38:32 +0100
Subject: [PATCH] Remove stale preconditioner config names from `__init__.py`

`AmortizedPreconditionerConfig` was renamed to
`BaseShampooPreconditionerConfig` and `ShampooPreconditionerConfig` to
`ClassicShampooPreconditionerConfig` in `shampoo_types.py`, but
`__init__.py` still imported and re-exported the old names. Replace
them with the new names in the imports and `__all__`.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 distributed_shampoo/__init__.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/distributed_shampoo/__init__.py b/distributed_shampoo/__init__.py
index b767a52..767cf6a 100644
--- a/distributed_shampoo/__init__.py
+++ b/distributed_shampoo/__init__.py
@@ -34,7 +34,8 @@
 from distributed_shampoo.shampoo_types import (
     AdaGradPreconditionerConfig,
     AdamPreconditionerConfig,
-    AmortizedPreconditionerConfig,
+    BaseShampooPreconditionerConfig,
+    ClassicShampooPreconditionerConfig,
     DDPDistributedConfig,
     DefaultEigenvalueCorrectedShampooConfig,
     DefaultShampooConfig,
@@ -59,7 +60,6 @@
     RootInvShampooPreconditionerConfig,
     ScheduleFreeConfig,
     SGDPreconditionerConfig,
-    ShampooPreconditionerConfig,
     ShampooPT2CompileConfig,
     SignDescentPreconditionerConfig,
     SingleDeviceDistributedConfig,
@@ -85,14 +85,14 @@
     # `precision_config`.
     # `preconditioner_config` options.
     "PreconditionerConfig",  # Abstract base class.
-    "AmortizedPreconditionerConfig",  # Abstract base class (based on `PreconditionerConfig`).
-    "ShampooPreconditionerConfig",  # Abstract base class (based on `AmortizedPreconditionerConfig`).
-    "RootInvShampooPreconditionerConfig",  # Based on `ShampooPreconditionerConfig`.
+    "BaseShampooPreconditionerConfig",  # Abstract base class (based on `PreconditionerConfig`).
+    "ClassicShampooPreconditionerConfig",  # Abstract base class (based on `BaseShampooPreconditionerConfig`).
+    "RootInvShampooPreconditionerConfig",  # Based on `ClassicShampooPreconditionerConfig`.
     "DefaultShampooConfig",  # Default `RootInvShampooPreconditionerConfig` using `EigenConfig`.
     "RootInvKLShampooPreconditionerConfig",  # Based on `RootInvShampooPreconditionerConfig`.
-    "EigendecomposedShampooPreconditionerConfig",  # Based on `ShampooPreconditionerConfig`.
+    "EigendecomposedShampooPreconditionerConfig",  # Based on `ClassicShampooPreconditionerConfig`.
     "EigendecomposedKLShampooPreconditionerConfig",  # Based on `EigendecomposedShampooPreconditionerConfig`.
-    "EigenvalueCorrectedShampooPreconditionerConfig",  # Based on `AmortizedPreconditionerConfig`.
+    "EigenvalueCorrectedShampooPreconditionerConfig",  # Based on `BaseShampooPreconditionerConfig`.
     "DefaultEigenvalueCorrectedShampooConfig",  # Default `EigenvalueCorrectedShampooPreconditionerConfig` using `EighEigendecompositionConfig`.
     "DefaultSOAPConfig",  # Default `EigenvalueCorrectedShampooPreconditionerConfig` using `QREigendecompositionConfig`.
     "SpectralDescentPreconditionerConfig",  # Based on `PreconditionerConfig`.