Cleanup before v0.2 release

CONTRIBUTING.md

@@ -5,7 +5,7 @@
 Use [abseil-py](https://github.com/abseil/abseil-py/tree/main)'s **logging**, **testing** and **flags** instead of Python's own **logging**, **unittest** and **argparse**.
 
 We use [uv](https://docs.astral.sh/uv/) for managing dependencies. For reproducible builds, our project tracks the generated `uv.lock` file in the repository.
-On a weekly basis, the CI attemps an update of the lock file to test against upstream dependencies.
+On a weekly basis, the CI attempts an update of the lock file to test against upstream dependencies.
 
 New required dependencies can be added by `uv add $DEPENDENCY`.
 
@@ -47,7 +47,7 @@ We generally follow [Google's style guides](https://google.github.io/styleguide/
 
 Although common, **mixed case is not allowed** in any code.
 
-Run pre-commit at local before submitting merge request. You can also read [.pre-commit-config.yaml]( .pre-commit-config.yaml) to understand what are being forced. The **flake8** and **mypy** settings are inherited from PyTorch.
+Run pre-commit at local before submitting merge request. You can also read [.pre-commit-config.yaml]( .pre-commit-config.yaml) to understand what are being forced. The **mypy** settings are inherited from PyTorch.
 
 ## Test
 

README.md

@@ -60,7 +60,7 @@ Refer to tests for usage of different optimizers, e.g.  [`tests/test_orthogonali
 
 ### Integration with Megatron Core
 
-Integration with Megatron Core is available in **dev** branch, e.g. [muon.py](https://github.com/NVIDIA/Megatron-LM/blob/dev/megatron/core/optimizer/muon.py)
+Integration with Megatron Core is available in **dev** branch, e.g. [emerging_optimizers.py](https://github.com/NVIDIA/Megatron-LM/blob/dev/megatron/core/optimizer/emerging_optimizers.py)
 
 ## Benchmarks
 

docs/apidocs/soap.md

@@ -22,7 +22,7 @@ emerging_optimizers.soap
 
 .. autofunction:: update_kronecker_factors_kl_shampoo
 
-.. autofunction:: update_eigenbasis_and_momentum
+.. autofunction:: update_eigenbasis_and_exp_avgs
 
 
 :hidden:`REKLS`

docs/conf.py

@@ -30,7 +30,7 @@
 project = "Emerging-Optimizers"
 copyright = "2025, NVIDIA Corporation"
 author = "NVIDIA Corporation"
-release = "0.1.0"
+release = "0.2.0"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

emerging_optimizers/orthogonalized_optimizers/muon_hyperball.py

@@ -17,12 +17,14 @@
 
 import torch
 
+from emerging_optimizers import registry
 from emerging_optimizers.orthogonalized_optimizers import muon
 
 
 __all__ = ["MuonHyperball"]
 
 
+@registry.register_optimizer("muon_hyperball")
 class MuonHyperball(muon.Muon):
     """Muon optimizer with hyperball-style norm-preserving weight updates.
 

emerging_optimizers/riemannian_optimizers/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from emerging_optimizers.riemannian_optimizers.normalized_optimizer import *

emerging_optimizers/riemannian_optimizers/normalized_optimizer.py

@@ -21,7 +21,13 @@
 import torch
 from torch.optim.optimizer import Optimizer
 
+from emerging_optimizers import registry
 
+
+__all__ = ["ObliqueSGD", "ObliqueAdam"]
+
+
+@registry.register_optimizer("oblique_sgd")
 class ObliqueSGD(Optimizer):
     """SGD optimizer for row- or column-normalized 2D parameters on oblique manifolds.
 
@@ -120,6 +126,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         return loss
 
 
+@registry.register_optimizer("oblique_adam")
 class ObliqueAdam(Optimizer):
     """Adam optimizer for row- or column-normalized 2D parameters on oblique manifolds.
 

emerging_optimizers/soap/soap.py

@@ -36,7 +36,7 @@
     "precondition",
     "init_kronecker_factors",
     "update_kronecker_factors",
-    "update_eigenbasis_and_momentum",
+    "update_eigenbasis_and_exp_avgs",
 ]
 
 
@@ -194,7 +194,6 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 if p.grad is None:
                     continue
 
-                # TODO(skyw): Fix the double cast. It is casted once in _init_group and once here.
                 grad = p.grad.to(torch.float32)
                 state = self.state[p]
 
@@ -252,11 +251,11 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                     )
                     if not skip_update:
                         with utils.fp32_matmul_precision(self.qr_fp32_matmul_prec):
-                            updated_eigenbasis_list, exp_avg, exp_avg_sq = update_eigenbasis_and_momentum(
+                            updated_eigenbasis_list, exp_avg, exp_avg_sq = update_eigenbasis_and_exp_avgs(
                                 kronecker_factor_list=kronecker_factor_list,
                                 eigenbasis_list=eigenbasis_list,
                                 exp_avg_sq=state["exp_avg_sq"],
-                                momentum=state["exp_avg"],
+                                exp_avg=state["exp_avg"],
                                 use_eigh=use_eigh,
                                 power_iter_steps=self.power_iter_steps,
                             )
@@ -436,22 +435,22 @@ def update_kronecker_factors_kl_shampoo(
 
 
 @torch.no_grad()  # type: ignore[misc]
-def update_eigenbasis_and_momentum(
+def update_eigenbasis_and_exp_avgs(
     kronecker_factor_list: list[torch.Tensor],
     eigenbasis_list: list[torch.Tensor],
     exp_avg_sq: torch.Tensor,
-    momentum: torch.Tensor,
+    exp_avg: torch.Tensor,
     use_eigh: bool = False,
     power_iter_steps: int = 1,
 ) -> tuple[list[torch.Tensor], torch.Tensor, torch.Tensor]:
-    """Updates the eigenbases using QR decomposition and power iteration or eigh.
+    """Updates the eigenbases and moving averages.
 
     This function performs an update of the eigenbases (QL and QR)
     used for preconditioning. It follows these steps:
 
-    1. Projects momentum back to the original basis
+    1. Projects exp_avg back to the original basis
     2. Updates the eigenbases using QR decomposition and power iteration (orthogonal iteration)
-    3. Projects momentum back to the new eigenbasis
+    3. Projects exp_avg back to the new eigenbasis
 
     Args:
         kronecker_factor_list: List of preconditioner matrices (L and R) that define
@@ -460,7 +459,7 @@ def update_eigenbasis_and_momentum(
             used for preconditioning. These will be updated by this function.
         exp_avg_sq: Inner Adam's second moment tensor, used for scaling the preconditioner updates.
             This tensor is modified in-place.
-        momentum: Inner Adam's first moment tensor, used for tracking gradient momentum.
+        exp_avg: Inner Adam's first moment tensor, used for tracking gradient momentum.
             This tensor is modified in-place.
         use_eigh: Whether to use full symmetric eigendecomposition (eigh) to compute the eigenbasis.
             If False, use orthogonal iteration to compute the eigenbasis.
@@ -470,23 +469,24 @@ def update_eigenbasis_and_momentum(
     Returns:
         A tuple containing:
             - Updated list of eigenbases (QL and QR)
-            - Updated momentum tensor projected to the new eigenbasis
+            - Updated exp_avg tensor projected to the new eigenbasis
+            - Updated exp_avg_sq tensor
 
     Example:
         >>> L = torch.randn(10, 10)
         >>> R = torch.randn(20, 20)
         >>> QL = torch.randn(10, 10)
         >>> QR = torch.randn(20, 20)
         >>> exp_avg_sq = torch.randn(10, 20)
-        >>> momentum = torch.randn(10, 20)
-        >>> updated_eigenbases = update_eigenbasis(
-        ...     [L, R], [QL, QR], exp_avg_sq, momentum)
+        >>> exp_avg = torch.randn(10, 20)
+        >>> updated_eigenbasis_list, updated_exp_avg, updated_exp_avg_sq = update_eigenbasis_and_exp_avgs(
+        ...     [L, R], [QL, QR], exp_avg_sq, exp_avg)
 
     """
-    # Step 1: Project momentum back to the original basis
+    # Step 1: Project exp_avg back to the original basis
     torch.cuda.nvtx.range_push("eigenbasis update step 1: precondition")
-    momentum = precondition(
-        momentum,
+    exp_avg = precondition(
+        exp_avg,
         eigenbasis_list,
         dims=[[0], [1]],
     )
@@ -508,16 +508,16 @@ def update_eigenbasis_and_momentum(
         )
     torch.cuda.nvtx.range_pop()
 
-    # Step 3: Project momentum to the new eigenbasis using the updated eigenbases
-    torch.cuda.nvtx.range_push("eigenbasis update step 3: project momentum")
-    momentum = precondition(
-        momentum,
+    # Step 3: Project exp_avg to the new eigenbasis using the updated eigenbases
+    torch.cuda.nvtx.range_push("eigenbasis update step 3: project exp_avg")
+    exp_avg = precondition(
+        exp_avg,
         updated_eigenbasis_list,
         dims=[[0], [0]],
     )
     torch.cuda.nvtx.range_pop()
 
-    return updated_eigenbasis_list, momentum, exp_avg_sq
+    return updated_eigenbasis_list, exp_avg, exp_avg_sq
 
 
 @torch.no_grad()  # type: ignore[misc]

pyproject.toml

@@ -76,8 +76,6 @@ docs = [
 ]
 test = [
     "coverage>=7.8.1",
-    "flake8>=7.2.0",
-    "pylint>=3.3.7",
     "triton>=3.4.0",
 ]
 dev = [

tests/test_soap.py

@@ -273,12 +273,12 @@ def test_clip_update_rms(self, max_rms: float) -> None:
         N=[4, 8, 33],
         use_eigh=[True, False],
     )
-    def test_update_eigenbasis_and_momentum(self, M: int, N: int, use_eigh: bool) -> None:
-        """Tests that update_eigenbasis_and_momentum returns valid outputs.
+    def test_update_eigenbasis_and_exp_avgs(self, M: int, N: int, use_eigh: bool) -> None:
+        """Tests that update_eigenbasis_and_exp_avgs returns valid outputs.
 
         Verifies output shapes, eigenbasis orthogonality, and that the round-trip
         projection (original → eigenbasis → original → new eigenbasis) preserves the
-        norm of momentum.
+        norm of exp_avg.
         """
         # Create symmetric positive definite kronecker factors
         g = torch.randn(M, N, device=self.device)
@@ -292,22 +292,22 @@ def test_update_eigenbasis_and_momentum(self, M: int, N: int, use_eigh: bool) ->
         eigenbasis_list = [Q_L, Q_R]
 
         exp_avg_sq = torch.abs(torch.randn(M, N, device=self.device))
-        momentum = torch.randn(M, N, device=self.device)
-        momentum_norm_before = torch.linalg.norm(momentum)
+        exp_avg = torch.randn(M, N, device=self.device)
+        exp_avg_norm_before = torch.linalg.norm(exp_avg)
 
-        updated_eigenbasis_list, updated_momentum, updated_exp_avg_sq = soap.update_eigenbasis_and_momentum(
+        updated_eigenbasis_list, updated_exp_avg, updated_exp_avg_sq = soap.update_eigenbasis_and_exp_avgs(
             kronecker_factor_list=kronecker_factor_list,
             eigenbasis_list=eigenbasis_list,
             exp_avg_sq=exp_avg_sq,
-            momentum=momentum,
+            exp_avg=exp_avg,
             use_eigh=use_eigh,
         )
 
         # Check output shapes
         self.assertEqual(len(updated_eigenbasis_list), 2)
         self.assertEqual(updated_eigenbasis_list[0].shape, (M, M))
         self.assertEqual(updated_eigenbasis_list[1].shape, (N, N))
-        self.assertEqual(updated_momentum.shape, (M, N))
+        self.assertEqual(updated_exp_avg.shape, (M, N))
         self.assertEqual(updated_exp_avg_sq.shape, (M, N))
 
         # Check eigenbasis orthogonality
@@ -321,13 +321,13 @@ def test_update_eigenbasis_and_momentum(self, M: int, N: int, use_eigh: bool) ->
                 msg="Updated eigenbasis is not orthogonal.",
             )
 
-        # Momentum is projected via orthogonal transforms, so norm should be preserved
+        # exp_avg is projected via orthogonal transforms, so norm should be preserved
         torch.testing.assert_close(
-            torch.linalg.norm(updated_momentum),
-            momentum_norm_before,
+            torch.linalg.norm(updated_exp_avg),
+            exp_avg_norm_before,
             atol=1e-5,
             rtol=1e-5,
-            msg="Momentum norm not preserved after eigenbasis update.",
+            msg="exp_avg norm not preserved after eigenbasis update.",
         )
 
     @parameterized.parameters(