update docs for newly introduced functions (#48)

Signed-off-by: Hao Wu <[email protected]>

Author: skyw

docs/apidocs/index.md

@@ -9,5 +9,7 @@ NeMo Emerging Optimizers API reference provides comprehensive technical document
 utils.md
 orthogonalized-optimizers.md
 soap.md
+riemannian-optimizers.md
+psgd.md
 scalar-optimizers.md
 ```

docs/apidocs/psgd.md

@@ -0,0 +1,16 @@
+```{eval-rst}
+.. role:: hidden
+    :class: hidden-section
+
+emerging_optimizers.psgd
+========================================
+
+.. automodule:: emerging_optimizers.psgd.procrustes_step
+    :members:
+
+.. automodule:: emerging_optimizers.psgd.psgd_kron_contractions
+    :members:
+
+.. automodule:: emerging_optimizers.psgd.psgd_utils
+    :members:
+```

docs/apidocs/riemannian-optimizers.md

@@ -0,0 +1,10 @@
+```{eval-rst}
+.. role:: hidden
+    :class: hidden-section
+
+emerging_optimizers.riemannian_optimizers
+=========================================
+
+.. automodule:: emerging_optimizers.riemannian_optimizers.normalized_optimizer
+    :members:
+```

emerging_optimizers/riemannian_optimizers/normalized_optimizer.py

@@ -68,9 +68,9 @@ def __init__(
     @torch.no_grad()  # type: ignore[misc]
     def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
+
         Args:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
+            closure: A closure that reevaluates the model and returns the loss.
         """
         loss = closure() if closure is not None else None
 
@@ -157,9 +157,9 @@ def __init__(
     @torch.no_grad()  # type: ignore[misc]
     def step(self, closure: Callable[[], float] | None = None) -> float | None:
         """Performs a single optimization step.
+
         Args:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
+            closure: A closure that reevaluates the model and returns the loss.
         """
         loss = closure() if closure is not None else None