fix arguments

optax/contrib/_muon.py

@@ -111,6 +111,7 @@ def muon(
     ] = (3.4445, -4.7750, 2.0315),
     newton_schulz_steps: Optional[int] = 5,
     mumentum: float = 0.95,
+    eps: float = 1e-8,
     mu_dtype: Optional[Any] = None,
     *,
     nesterov: bool = True,
@@ -134,6 +135,7 @@ def muon(
     newton_schulz_steps: Number of Newton-schulz iterations.
     mumentum: Exponential decay rate to track the first moment of past
       gradients.
+    eps: Term added to the denominator to improve numerical stability.
     mu_dtype: Data type of the momentum accumulator.
     nesterov: Whether to use Nesterov momentum.
 
@@ -145,6 +147,7 @@ def muon(
           newton_schulz_coeffs,
           newton_schulz_steps,
           mumentum,
+          eps,
           mu_dtype,
           nesterov=nesterov,
       ),