refactor: extract composite optimizer helpers and unify naming

- add astrai/optim/composite.py with shared step/zero_grad/state_dict/param_groups helpers and OptimizerFactory
- rename MuonMix to MuonAdamW (matches registered name muon_adamw) and file to muon_adamw.py
- use @OptimizerFactory.register decorator in each optimizer module instead of post-import registration in __init__
- fix closure being invoked once per sub-optimizer in MuonAdamW.step (now exactly once via composite_step)
- NoraNAdamW.step now forwards closure correctly

astrai/optim/__init__.py

@@ -1,15 +1,13 @@
 """Optimizer implementations and factory registration."""
 
-from torch.optim import Optimizer
-
-from astrai.factory import BaseFactory
-
-
-class OptimizerFactory(BaseFactory[Optimizer]):
-    """Factory for built-in training optimizers."""
-
-
-from astrai.optim.muon_mix import MuonMix
+from astrai.optim.composite import (
+    OptimizerFactory,
+    composite_state_dict,
+    composite_step,
+    composite_zero_grad,
+    refresh_param_groups,
+)
+from astrai.optim.muon_adamw import MuonAdamW
 from astrai.optim.nora_nadamw import (
     NAdamW,
     Nora,
@@ -20,17 +18,18 @@ from astrai.optim.nora_nadamw import (
     partition_optimizer_parameters,
 )
 
-OptimizerFactory.register("nora_nadamw")(NoraNAdamW)
-OptimizerFactory.register("muon_adamw")(MuonMix)
-
 __all__ = [
-    "MuonMix",
+    "MuonAdamW",
     "NAdamW",
     "Nora",
     "NoraNAdamW",
     "OptimizerFactory",
     "OptimizerParameterGroups",
+    "composite_state_dict",
+    "composite_step",
+    "composite_zero_grad",
     "nora_direction",
     "nora_lr_scale",
     "partition_optimizer_parameters",
+    "refresh_param_groups",
 ]