v0.3.0 pre-release cleanup

CONTRIBUTING.md

@@ -49,6 +49,38 @@ Although common, **mixed case is not allowed** in any code. "Mixed case" here me
 
 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.
 
+## Optimizer API
+
+`closure` is **not supported** on any optimizer's `step` in this repo. The parameter is kept on the signature for compatibility with `torch.optim.Optimizer`, but every `step` implementation raises `ValueError("closure is not supported")` when one is passed and returns `None` otherwise. New optimizers must do the same — do not add code paths that call `closure()` or return its result.
+
+### Argument-order convention
+
+When designing or modifying any public function in `emerging_optimizers/` — optimizer `__init__`, `calculate_*_update`, helpers — follow this precedence for argument order:
+
+1. **PyTorch native first.** Match `torch.optim` / `torch.optim._functional` signatures. For Adam-family signatures that means `betas` and `eps` adjacent (`Adam(lr, betas, eps, weight_decay, ...)`).
+2. **HuggingFace next.** For knobs PyTorch doesn't expose, follow HF's order. The clearest case is `correct_bias`: HF's `transformers.AdamW` puts it after the standard Adam scalars (`lr, betas, eps, weight_decay`), so in our update functions `correct_bias` goes **after** `eps`.
+3. **Repo-specific extras last.** Anything not in PyTorch or HF — `nesterov` on Adam, the scalar-int `step`, warmup schedulers, `alpha`, `scale_log2`, `use_shape_scaling`, etc. — goes at the end.
+
+Concrete shape for the scalar `calculate_*_update` family:
+
+```python
+def calculate_*_update(
+    grad,                   # positional: tensors only
+    <state buffers>,
+    *,                      # everything below is keyword-only
+    betas (or momentum),
+    eps,
+    correct_bias,
+    [nesterov],
+    step,
+    <other repo-specific extras>,
+):
+```
+
+`step` is a repo-specific argument and goes in the trailing bucket. (PyTorch's functional adam takes `state_steps` as a `List[Tensor]` grouped with the buffers — that is *not* the same as our scalar `step: int`, so the PyTorch precedent does not transfer.)
+
+The `*` keyword-only separator matches the convention used by every optimizer `__init__` in `emerging_optimizers/` (Muon, Soap, PSGDPro, AdaptiveMuon, etc.) and by `torch.optim._functional.adam` — only tensors are positional, scalars and flags must be passed by name. mypy will reject positional scalar passes, so the rule is enforced at type-check time.
+
 ## Test
 
 All tests should be placed under [tests](tests). We aim for 100% test coverage for this tiny project.

emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py

@@ -274,12 +274,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Single optimization step.
 
         Args:
-            closure: A closure that reevaluates the model and returns the loss.
+            closure: Unsupported; must be ``None``.
         """
-        if closure is None:
-            loss = None
-        else:
-            loss = closure()
+        if closure is not None:
+            raise ValueError("closure is not supported")
 
         for group in self.param_groups:
             self._init_group(group)
@@ -330,4 +328,4 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 p.add_(update, alpha=-group["lr"])
                 self.post_weight_update_fn_inplace(p)
 
-        return loss
+        return None

emerging_optimizers/orthogonalized_optimizers/orthogonalized_optimizer.py

@@ -158,12 +158,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
 
         Args:
-            closure: A closure that reevaluates the model and returns the loss.
+            closure: Unsupported; must be ``None``.
         """
-        if closure is None:
-            loss = None
-        else:
-            loss = closure()
+        if closure is not None:
+            raise ValueError("closure is not supported")
 
         for group in self.param_groups:
             self._init_group(group)
@@ -200,7 +198,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 p.add_(orth_grad, alpha=-group["lr"])
                 self.post_weight_update_fn_inplace(p)
 
-        return loss
+        return None
 
     def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor:
         """Orthogonalize the momentum.

emerging_optimizers/psgd/psgd.py

@@ -105,12 +105,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
 
         Args:
-            closure: A closure that reevaluates the model and returns the loss.
+            closure: Unsupported; must be ``None``.
         """
-        if closure is None:
-            loss = None
-        else:
-            loss = closure()
+        if closure is not None:
+            raise ValueError("closure is not supported")
 
         for group in self.param_groups:
             for p in group["params"]:
@@ -162,7 +160,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 # Apply weight update
                 p.add_(update, alpha=-group["lr"])
 
-        return loss
+        return None
 
 
 def _init_psgd_kron_states(

emerging_optimizers/riemannian_optimizers/normalized_optimizer.py

@@ -58,13 +58,14 @@ def __init__(
         lr: float = 1e-3,
         momentum: float = 0.9,
         weight_decay: float = 0.0,
+        *,
         weight_decay_method: opt_mixin.WeightDecayT = "decoupled",
         dim: int = 0,
         eps: float = 1e-8,
     ) -> None:
         if lr < 0.0:
             raise ValueError(f"Invalid learning rate: {lr}")
-        if momentum < 0.0 or momentum >= 1.0:
+        if not 0.0 <= momentum < 1.0:
             raise ValueError(f"Invalid momentum value: {momentum}")
         if weight_decay < 0.0:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
@@ -93,9 +94,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
 
         Args:
-            closure: A closure that reevaluates the model and returns the loss.
+            closure: Unsupported; must be ``None``.
         """
-        loss = closure() if closure is not None else None
+        if closure is not None:
+            raise ValueError("closure is not supported")
 
         for group in self.param_groups:
             lr = group["lr"]
@@ -129,7 +131,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 # Retraction back to the manifold, the hyper-sphere
                 torch.nn.functional.normalize(param, p=2.0, dim=dim, eps=eps, out=param)
 
-        return loss
+        return None
 
 
 @registry.register_optimizer("oblique_adam")
@@ -147,6 +149,7 @@ def __init__(
         lr: float = 1e-3,
         betas: tuple[float, float] = (0.9, 0.99),
         weight_decay: float = 0.0,
+        *,
         weight_decay_method: opt_mixin.WeightDecayT = "decoupled",
         dim: int = 0,
         eps: float = 1e-8,
@@ -165,9 +168,9 @@ def __init__(
         """
         if lr < 0.0:
             raise ValueError(f"Invalid learning rate: {lr}")
-        if betas[0] < 0.0 or betas[0] >= 1.0:
+        if not 0.0 <= betas[0] < 1.0:
             raise ValueError(f"Invalid beta1 value: {betas[0]}")
-        if betas[1] < 0.0 or betas[1] >= 1.0:
+        if not 0.0 <= betas[1] < 1.0:
             raise ValueError(f"Invalid beta2 value: {betas[1]}")
         if weight_decay < 0.0:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
@@ -197,9 +200,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
 
         Args:
-            closure: A closure that reevaluates the model and returns the loss.
+            closure: Unsupported; must be ``None``.
         """
-        loss = closure() if closure is not None else None
+        if closure is not None:
+            raise ValueError("closure is not supported")
 
         for group in self.param_groups:
             lr = group["lr"]
@@ -256,7 +260,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 # Retraction back to the manifold, i.e. the hyper-sphere
                 torch.nn.functional.normalize(param, p=2.0, dim=dim, eps=eps, out=param)
 
-        return loss
+        return None
 
 
 def _compute_riemannian_grad(param: torch.Tensor, grad_like: torch.Tensor, dim: int) -> torch.Tensor:

emerging_optimizers/scalar_optimizers/laprop.py

@@ -77,19 +77,20 @@ def __init__(
         betas: tuple[float, float] = (0.9, 0.999),
         eps: float = 1e-8,
         weight_decay: float = 0.0,
+        *,
         correct_bias: bool = True,
         frob_normalize: bool = False,
         weight_decay_method: WeightDecayT = "decoupled",
     ) -> None:
-        if not 0.0 <= lr:
+        if lr < 0.0:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= betas[0] < 1.0:
             raise ValueError(f"Invalid beta at index 0: {betas[0]}")
         if not 0.0 <= betas[1] < 1.0:
             raise ValueError(f"Invalid beta at index 1: {betas[1]}")
-        if not 0.0 <= eps:
+        if eps < 0.0:
             raise ValueError(f"Invalid epsilon value: {eps}")
-        if not 0.0 <= weight_decay:
+        if weight_decay < 0.0:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
         if frob_normalize and weight_decay != 0.0:
             logging.warning("LaProp with frob_normalize=True is intended to be used with weight_decay=0.0.")
@@ -139,15 +140,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
             If you need the original gradient after this call, clone it beforehand.
 
         Args:
-            closure: A closure that reevaluates the model and returns the loss (optional).
-
-        Returns:
-            The loss from the closure, if provided.
+            closure: Unsupported; must be ``None``.
         """
-        if closure is None:
-            loss = None
-        else:
-            loss = closure()
+        if closure is not None:
+            raise ValueError("closure is not supported")
 
         for group in self.param_groups:
             self._init_group(group)
@@ -173,14 +169,14 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                     grad,
                     state["exp_avg"],
                     state["exp_avg_sq"],
-                    correct_bias,
-                    betas,
-                    state["step"],
-                    eps,
+                    betas=betas,
+                    eps=eps,
+                    correct_bias=correct_bias,
+                    step=state["step"],
                 )
                 p.data.add_(update, alpha=-lr)
                 if self.frob_normalize:
                     assert pre_norm is not None
                     p.data.mul_(pre_norm / p.data.norm().clamp_min(eps))
 
-        return loss
+        return None

emerging_optimizers/scalar_optimizers/lion.py

@@ -71,15 +71,16 @@ def __init__(
         lr: float = 1e-4,
         betas: tuple[float, float] = (0.9, 0.99),
         weight_decay: float = 0.01,
+        *,
         weight_decay_method: WeightDecayT = "decoupled",
     ) -> None:
-        if not 0.0 <= lr:
+        if lr < 0.0:
             raise ValueError(f"Invalid learning rate: {lr}")
         if not 0.0 <= betas[0] < 1.0:
             raise ValueError(f"Invalid beta at index 0: {betas[0]}")
         if not 0.0 <= betas[1] < 1.0:
             raise ValueError(f"Invalid beta at index 1: {betas[1]}")
-        if not 0.0 <= weight_decay:
+        if weight_decay < 0.0:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
         defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
         self.weight_decay_method = weight_decay_method
@@ -124,15 +125,10 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
             If you need the original gradient after this call, clone it beforehand.
 
         Args:
-            closure: A closure that reevaluates the model and returns the loss (optional).
-
-        Returns:
-            The loss from the closure, if provided.
+            closure: Unsupported; must be ``None``.
         """
-        loss = None
         if closure is not None:
-            with torch.enable_grad():
-                loss = closure()
+            raise ValueError("closure is not supported")
 
         for group in self.param_groups:
             self._init_group(group)
@@ -154,7 +150,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 # Lion update: sign(beta1 * m + (1-beta1) * g)
                 # Note: different betas per param-group will each trigger a one-time
                 # torch.compile recompilation of calculate_lion_update.
-                update = calculate_lion_update(grad, exp_avg, betas)
+                update = calculate_lion_update(grad, exp_avg, betas=betas)
                 p.data.add_(update, alpha=-lr)
 
-        return loss
+        return None

emerging_optimizers/scalar_optimizers/update_functions/adam.py

@@ -26,11 +26,12 @@ def calculate_adam_update(
     grad: torch.Tensor,
     exp_avg: torch.Tensor,
     exp_avg_sq: torch.Tensor,
+    *,
     betas: tuple[float, float],
+    eps: float,
     correct_bias: bool,
     nesterov: bool,
     step: int,
-    eps: float,
 ) -> torch.Tensor:
     """Performs the Adam update.
 
@@ -47,16 +48,16 @@ def calculate_adam_update(
 
     Args:
         grad: The gradient tensor.
-        exp_avg: The accumulated first moment of the gradient.
-        exp_avg_sq: The accumulated second moment of the gradient.
-        betas: The EMA beta coefficients for the Adam update.
-        correct_bias: Whether to correct the bias of the Adam update.
-        nesterov: Whether to use nesterov momentum.
-        step: The current step of the optimizer, used to compute the bias correction terms.
-        eps: The epsilon for the Adam second moment update.
+        exp_avg: The accumulated first moment of the gradient (modified in place).
+        exp_avg_sq: The accumulated second moment of the gradient (modified in place).
+        betas: The EMA beta coefficients ``(beta1, beta2)``.
+        eps: Epsilon for the second-moment denominator.
+        correct_bias: Whether to apply Adam-style bias correction.
+        nesterov: Whether to use Nesterov momentum.
+        step: Current optimizer step (1-based), used for bias correction.
 
     Returns:
-        The Adam-update.
+        The Adam update.
     """
 
     beta1, beta2 = betas