Update docstring for MuonHyperball (#99)

* update docstring for MuonHyperball

Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Tim Lau <timlautk@gmail.com>

Author: skyw

docs/apidocs/orthogonalized-optimizers.md

@@ -27,12 +27,18 @@ emerging_optimizers.orthogonalized_optimizers
 .. autoclass:: Scion
     :members:
 
-:hidden:`Mop`
+:hidden:`MOP`
 ~~~~~~~~~~~~~~~
 
 .. autoclass:: MOP
     :members:
 
+:hidden:`MuonHyperball`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: MuonHyperball
+    :members:
+
 
 :hidden:`Newton-Schulz`
 ~~~~~~~~~~~~~~~~~~~~~~~~

emerging_optimizers/orthogonalized_optimizers/muon_hyperball.py

@@ -41,6 +41,17 @@ class MuonHyperball(muon.Muon):
 
     See :class:`~emerging_optimizers.orthogonalized_optimizers.muon.Muon` for full documentation
     of the base Muon optimizer.
+
+
+    Args:
+        *args: Arguments passed to Muon.
+        hyperball_eps: Epsilon for numerical stability in normalization.
+            Default: ``1e-8``.
+        hyperball_radius: Fixed radius for the hyperball. If ``None`` (default),
+            uses each parameter's initial Frobenius norm as its radius. If specified, all
+            parameters will be rescaled to have this radius at initialization.
+        **kwargs: Keyword arguments passed to Muon.
+
     """
 
     def __init__(
@@ -50,19 +61,6 @@ def __init__(
         hyperball_radius: float | None = None,
         **kwargs: Any,
     ) -> None:
-        """Initialize MuonHyperball optimizer.
-
-        Args:
-            *args: Arguments passed to Muon.
-            **kwargs: Keyword arguments passed to Muon.
-
-        Keyword args:
-            hyperball_eps (float, optional): Epsilon for numerical stability in normalization.
-                Default: ``1e-8``.
-            hyperball_radius (float, optional): Fixed radius for the hyperball. If ``None`` (default),
-                uses each parameter's initial Frobenius norm as its radius. If specified, all
-                parameters will be rescaled to have this radius at initialization.
-        """
         self.hyperball_eps = hyperball_eps
         self.hyperball_radius = hyperball_radius
         super().__init__(*args, **kwargs)