AdaMuon impl w/ a few other ideas based on recent reading

Author: rwightman

tests/test_optim.py

@@ -402,6 +402,14 @@ def test_muon(optimizer):
     _test_model(optimizer, dict(lr=1e-3))
 
 
+@pytest.mark.parametrize('optimizer',  ['adamuon', 'nadamuon'])
+def test_adamuon(optimizer):
+    _test_rosenbrock(
+        lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
+    )
+    _test_model(optimizer, dict(lr=1e-3))
+
+
 @pytest.mark.parametrize('optimizer',  ['adopt', 'adoptw'])
 def test_adopt(optimizer):
     _test_rosenbrock(

timm/optim/_optim_factory.py

@@ -897,6 +897,24 @@ def _register_other_optimizers(registry: OptimizerRegistry) -> None:
             has_betas=True,
             defaults={'nesterov': True}
         ),
+        OptimInfo(
+            name='adamuon',
+            opt_class=Muon,
+            description='AdaMuon: Muon with adaptive second moment estimation on orthogonalized directions',
+            has_momentum=True,
+            has_eps=True,
+            has_betas=True,
+            defaults={'algo': 'adamuon'}
+        ),
+        OptimInfo(
+            name='nadamuon',
+            opt_class=Muon,
+            description='AdaMuon with Nesterov momentum and NAdamW fallback for 1D params',
+            has_momentum=True,
+            has_eps=True,
+            has_betas=True,
+            defaults={'algo': 'adamuon', 'nesterov': True}
+        ),
         OptimInfo(
             name='novograd',
             opt_class=NvNovoGrad,