[Doc] udf (#195)

* Docstring for udf

* Track udf docs

* Improve

* Improved

* Delete udf.py

* Improve

docs/source/api/python/index.rst

@@ -9,3 +9,4 @@ API Reference
    function
    traversal
    propagate
+   udf

docs/source/api/python/udf.rst

+.. _apiudf:
+
+User-defined function related data structures
+==================================================
+
+.. currentmodule:: dgl.udf
+
+EdgeBatch
+---------
+
+The class that can represent a batch of edges.
+
+.. autosummary::
+    :toctree: ../../generated/
+
+    EdgeBatch.src
+    EdgeBatch.dst
+    EdgeBatch.data
+    EdgeBatch.edges
+    EdgeBatch.batch_size
+
+NodeBatch
+---------
+
+The class that can represent a batch of nodes.
+
+.. autosummary::
+    :toctree: ../../generated/
+
+    NodeBatch.data
+    NodeBatch.mailbox
+    NodeBatch.nodes
+    NodeBatch.batch_size

python/dgl/udf.py

 """User-defined function related data structures."""
 from __future__ import absolute_import
 
-from collections import Mapping
-
 from .base import ALL, is_all
 from . import backend as F
 from . import utils
 
 class EdgeBatch(object):
-    """The object that represents a batch of edges.
+    """The class that can represent a batch of edges.
 
     Parameters
     ----------
@@ -16,12 +14,15 @@ class EdgeBatch(object):
         The graph object.
     edges : tuple of utils.Index
         The edge tuple (u, v, eid). eid can be ALL
-    src_data : dict of tensors
-        The src node features
-    edge_data : dict of tensors
-        The edge features.
+    src_data : dict
+        The src node features, in the form of ``dict``
+        with ``str`` keys and ``tensor`` values
+    edge_data : dict
+        The edge features, in the form of ``dict`` with
+        ``str`` keys and ``tensor`` values
     dst_data : dict of tensors
-        The dst node features
+        The dst node features, in the form of ``dict``
+        with ``str`` keys and ``tensor`` values
     """
     def __init__(self, g, edges, src_data, edge_data, dst_data):
         self._g = g
@@ -36,8 +37,8 @@ class EdgeBatch(object):
 
         Returns
         -------
-        dict of str to tensors
-            The feature data.
+        dict with str keys and tensor values
+            Features of the source nodes.
         """
         return self._src_data
 
@@ -47,8 +48,8 @@ class EdgeBatch(object):
 
         Returns
         -------
-        dict of str to tensors
-            The feature data.
+        dict with str keys and tensor values
+            Features of the destination nodes.
         """
         return self._dst_data
 
@@ -58,18 +59,21 @@ class EdgeBatch(object):
 
         Returns
         -------
-        dict of str to tensors
-            The feature data.
+        dict with str keys and tensor values
+            Features of the edges.
         """
         return self._edge_data
 
     def edges(self):
         """Return the edges contained in this batch.
-        
+
         Returns
         -------
-        tuple of tensors
-            The edge tuple (u, v, eid).
+        tuple of three tensors
+            The edge tuple :math:`(src, dst, eid)`. :math:`src[i],
+            dst[i], eid[i]` separately specifies the source node,
+            destination node and the edge id for the ith edge
+            in the batch.
         """
         if is_all(self._edges[2]):
             self._edges[2] = utils.toindex(F.arange(
@@ -78,7 +82,12 @@ class EdgeBatch(object):
         return (u.tousertensor(), v.tousertensor(), eid.tousertensor())
 
     def batch_size(self):
-        """Return the number of edges in this edge batch."""
+        """Return the number of edges in this edge batch.
+
+        Returns
+        -------
+        int
+        """
         return len(self._edges[0])
 
     def __len__(self):
@@ -86,7 +95,7 @@ class EdgeBatch(object):
         return self.batch_size()
 
 class NodeBatch(object):
-    """The object that represents a batch of nodes.
+    """The class that can represent a batch of nodes.
 
     Parameters
     ----------
@@ -94,10 +103,12 @@ class NodeBatch(object):
         The graph object.
     nodes : utils.Index or ALL
         The node ids.
-    data : dict of tensors
-        The node features
-    msgs : dict of tensors, optional
-        The messages.
+    data : dict
+        The node features, in the form of ``dict``
+        with ``str`` keys and ``tensor`` values
+    msgs : dict, optional
+        The messages, , in the form of ``dict``
+        with ``str`` keys and ``tensor`` values
     """
     def __init__(self, g, nodes, data, msgs=None):
         self._g = g
@@ -111,8 +122,8 @@ class NodeBatch(object):
 
         Returns
         -------
-        dict of str to tensors
-            The feature data.
+        dict with str keys and tensor values
+            Features of the nodes.
         """
         return self._data
 
@@ -120,18 +131,19 @@ class NodeBatch(object):
     def mailbox(self):
         """Return the received messages.
 
-        If no messages received, a None will be returned.
+        If no messages received, a ``None`` will be returned.
 
         Returns
         -------
-        dict of str to tensors
-            The message data.
+        dict or None
+            The messages nodes received. If dict, the keys are
+            ``str`` and the values are ``tensor``.
         """
         return self._msgs
 
     def nodes(self):
         """Return the nodes contained in this batch.
-        
+
         Returns
         -------
         tensor
@@ -143,7 +155,12 @@ class NodeBatch(object):
         return self._nodes.tousertensor()
 
     def batch_size(self):
-        """Return the number of nodes in this node batch."""
+        """Return the number of nodes in this batch.
+
+        Returns
+        -------
+        int
+        """
         if is_all(self._nodes):
             return self._g.number_of_nodes()
         else: