[Doc] Fix docs for sparse optimizer (dmlc#2680)

Co-authored-by: Ubuntu <[email protected]> Co-authored-by: Minjie Wang <[email protected]>
YHGui · Feb 19, 2021 · 453d358 · 453d358
1 parent 9e04a52
commit 453d358
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 3 deletions.
diff --git a/docs/source/api/python/nn.pytorch.rst b/docs/source/api/python/nn.pytorch.rst
@@ -268,6 +268,9 @@ SegmentedKNNGraph
     :members:
     :show-inheritance:
 
+NodeEmbedding Module
+----------------------------------------
+
 NodeEmbedding
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -122,6 +122,7 @@ Getting Started
    api/python/dgl.function
    api/python/nn
    api/python/dgl.ops
+   api/python/dgl.optim
    api/python/dgl.sampling
    api/python/udf
 

diff --git a/python/dgl/nn/pytorch/sparse_emb.py b/python/dgl/nn/pytorch/sparse_emb.py
@@ -22,6 +22,8 @@ class NodeEmbedding: # NodeEmbedding
     ``torch.distributed.TCPStore`` to share meta-data information across multiple gpu processes.
     It use the local address of '127.0.0.1:12346' to initialize the TCPStore.
 
+    NOTE: The support of NodeEmbedding is experimental.
+
     Parameters
     ----------
     num_embeddings : int

diff --git a/python/dgl/optim/pytorch/__init__.py b/python/dgl/optim/pytorch/__init__.py
@@ -1,2 +1,2 @@
-"""dgl optims for pytorch."""
+"""dgl sparse optimizer for pytorch."""
 from .sparse_optim import SparseAdagrad, SparseAdam
diff --git a/python/dgl/optim/pytorch/sparse_optim.py b/python/dgl/optim/pytorch/sparse_optim.py
@@ -237,6 +237,8 @@ class SparseAdagrad(SparseGradOptimizer):
     :math:`G_{t,i,j}=G_{t-1,i,j} + g_{t,i,j}^2` and :math:`g_{t,i,j}` is the gradient of
     the dimension :math:`j` of embedding :math:`i` at step :math:`t`.
 
+    NOTE: The support of sparse Adagrad optimizer is experimental.
+
     Parameters
     ----------
     params : list[dgl.nn.NodeEmbedding]
@@ -335,6 +337,8 @@ class SparseAdam(SparseGradOptimizer):
     :math:`g_{t,i,j}` is the gradient of the dimension :math:`j` of embedding :math:`i`
     at step :math:`t`.
 
+    NOTE: The support of sparse Adam optimizer is experimental.
+
     Parameters
     ----------
     params : list[dgl.nn.NodeEmbedding]
@@ -348,8 +352,8 @@ class SparseAdam(SparseGradOptimizer):
         The term added to the denominator to improve numerical stability
         Default: 1e-8
 
-    Examples:
-
+    Examples
+    --------
     >>> def initializer(emb):
             th.nn.init.xavier_uniform_(emb)
             return emb