add more docs + black

docs/conf.py

@@ -40,7 +40,7 @@
     'sphinx.ext.coverage',
     'sphinx.ext.mathjax',
     'sphinx.ext.viewcode',
-    'sphinx.ext.githubpages',
+    'sphinx.ext.intersphinx',
     'sphinx.ext.napoleon',
 ]
 
@@ -349,3 +349,10 @@
 # If true, do not generate a @detailmenu in the "Top" node's menu.
 #
 # texinfo_no_detailmenu = False
+
+# Example configuration for intersphinx: refer to the Python standard library.
+intersphinx_mapping = {
+    'numpy': ('http://docs.scipy.org/doc/numpy/', None),
+    'python': ('https://docs.python.org/', None),
+    'torch': ('https://pytorch.org/docs/master/', None),
+}

docs/manifolds.rst

@@ -4,5 +4,4 @@ Manifolds
 .. currentmodule:: geoopt.manifolds
 
 .. automodule:: geoopt.manifolds
-   :members:
-   :imported-members: True
+   :members: Stiefel, Euclidean

geoopt/manifolds/euclidean.py

@@ -5,9 +5,7 @@
 
 class Euclidean(Manifold):
     """
-    Euclidean manifold
-
-    An unconstrained manifold
+    Simple Euclidean manifold
     """
 
     name = "Euclidean"

geoopt/manifolds/stiefel.py

@@ -7,25 +7,28 @@
 __all__ = ["Stiefel", "EuclideanStiefel", "CanonicalStiefel"]
 
 
-class Stiefel(Manifold):
-    r"""
+_stiefel_doc = r"""
     Manifold induced by the following matrix constraint:
 
     .. math::
 
-        X^\top X = I
-        X \in \mathrm{R}^{n\times m}
+        X^\top X = I\\
+        X \in \mathrm{R}^{n\times m}\\
         n \ge m
+"""
+
+
+class Stiefel(Manifold):
+    __doc__ = r"""
+    {}
 
     Parameters
     ----------
     canonical : bool
         Use canonical inner product instead of euclidean one (defaults to canonical)
-
-    Notes
-    -----
-    works with batch sized tensors
-    """
+    """.format(
+        _stiefel_doc
+    )
     ndim = 2
 
     def __new__(cls, canonical=True):
@@ -76,6 +79,13 @@ def _projx(self, x):
 
 
 class CanonicalStiefel(Stiefel):
+    __doc__ = r"""Stiefel Manifold with Canonical inner product
+
+    {}
+    """.format(
+        _stiefel_doc
+    )
+
     name = "Stiefel(canonical)"
     reversible = True
 
@@ -131,6 +141,13 @@ def _retr(self, x, u, t):
 
 
 class EuclideanStiefel(Stiefel):
+    __doc__ = r"""Stiefel Manifold with Euclidean inner product
+
+    {}
+    """.format(
+        _stiefel_doc
+    )
+
     name = "Stiefel(euclidean)"
     reversible = False
 

geoopt/optim/radam.py

@@ -7,7 +7,7 @@
 
 
 class RiemannianAdam(OptimMixin, torch.optim.Adam):
-    r"""Riemannian Adam
+    r"""Riemannian Adam with the same API as :class:`torch.optim.Adam`
 
     Parameters
     ----------
@@ -29,9 +29,16 @@ class RiemannianAdam(OptimMixin, torch.optim.Adam):
         algorithm from the paper `On the Convergence of Adam and Beyond`_
         (default: False)
 
+    Other Parameters
+    ----------------
+    stabilize : int
+        Stabilize parameters if they are off-manifold due to numerical
+        reasons every ``stabilize`` steps (default: ``None`` -- no stabilize)
+
 
     .. _On the Convergence of Adam and Beyond:
         https://openreview.net/forum?id=ryQu7f-RZ
+
     """
 
     def step(self, closure=None):

geoopt/optim/rsgd.py

@@ -9,7 +9,7 @@
 
 
 class RiemannianSGD(OptimMixin, torch.optim.Optimizer):
-    r"""Riemannian Stochastic Gradient Descent
+    r"""Riemannian Stochastic Gradient Descent with the same API as :class:`torch.optim.SGD`
 
     Parameters
     ----------
@@ -26,6 +26,12 @@ class RiemannianSGD(OptimMixin, torch.optim.Optimizer):
         dampening for momentum (default: 0)
     nesterov : bool (optional)
         enables Nesterov momentum (default: False)
+
+    Other Parameters
+    ----------------
+    stabilize : int
+        Stabilize parameters if they are off-manifold due to numerical
+        reasons every ``stabilize`` steps (default: ``None`` -- no stabilize)
     """
 
     def __init__(

geoopt/tensor.py

@@ -6,8 +6,12 @@
 
 
 class ManifoldTensor(torch.Tensor):
-    """A regular tensor that has information about its manifold.
-    It is a very tiny wrapper over regular tensor so that all API is the same
+    """Same as :class:`torch.Tensor` that has information about its manifold.
+
+    Other Parameters
+    ----------------
+    manifold : :class:`geoopt.Manifold`
+        A manifold for the tensor, (default: :class:`geoopt.Euclidean`)
     """
 
     def __new__(cls, *args, manifold=Euclidean(), requires_grad=False, **kwargs):
@@ -50,6 +54,16 @@ def __repr__(self):
 
 
 class ManifoldParameter(ManifoldTensor, torch.nn.Parameter):
+    """Same as :class:`torch.nn.Parameter` that has information about its manifold.
+    It should be used within :class:`torch.nn.Module` to be recognized
+    in parameter collection.
+
+    Other Parameters
+    ----------------
+    manifold : :class:`geoopt.Manifold` (optional)
+        A manifold for the tensor if ``data`` is not a :class:`geoopt.ManifoldTensor`
+    """
+
     def __new__(cls, data=None, manifold=None, requires_grad=True):
         if data is None:
             data = ManifoldTensor(manifold=manifold)