[Doc] Improve Capsule with Jinyang & Fix wrong tutorial level layout (#236)

* improve capsule tutorial with jinyang * fix wrong layout of second-level tutorial * delete transformer

[Doc] Improve Capsule with Jinyang & Fix wrong tutorial level layout (#236)
* improve capsule tutorial with jinyang * fix wrong layout of second-level tutorial * delete transformer
16cc5287 · VoVAllen · Minjie Wang · dafe4671 · 16cc5287 · 16cc5287
Commit 16cc5287 authored Dec 04, 2018 by VoVAllen Committed by Minjie Wang Dec 04, 2018
17 changed files
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
 # -*- coding: utf-8 -*-
 #
 # Configuration file for the Sphinx documentation builder.
 #
 # This file does only contain a selection of the most common options. For a
 # full list see the documentation:
 # http://www.sphinx-doc.org/en/master/config
 # -- Path setup --------------------------------------------------------------
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
 import os
 import sys
 sys.path.insert(0, os.path.abspath('../../python'))
 # -- Project information -----------------------------------------------------
 project = 'DGL'
 copyright = '2018, DGL Team'
 author = 'DGL Team'
 # The short X.Y version
 version = '0.0.1'
 # The full version, including alpha/beta/rc tags
 release = '0.0.1'
 # -- General configuration ---------------------------------------------------
 # If your documentation needs a minimal Sphinx version, state it here.
 #
 # needs_sphinx = '1.0'
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.autosummary',
    'sphinx.ext.coverage',
    'sphinx.ext.mathjax',
    'sphinx.ext.napoleon',
    'sphinx.ext.viewcode',
    'sphinx.ext.intersphinx',
    'sphinx.ext.graphviz',
    'sphinx_gallery.gen_gallery',
 ]
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ['_templates']
 # The suffix(es) of source filenames.
 # You can specify multiple suffix as a list of string:
 #
 source_suffix = ['.rst', '.md']
 # The master toctree document.
 master_doc = 'index'
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
 #
 # This is also used if you do content translation via gettext catalogs.
 # Usually you set "language" from the command line for these cases.
 language = None
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
 exclude_patterns = []
 # The name of the Pygments (syntax highlighting) style to use.
 pygments_style = None
 # -- Options for HTML output -------------------------------------------------
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
 html_theme = 'sphinx_rtd_theme'
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
 # documentation.
 #
 # html_theme_options = {}
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
 html_static_path = ['_static']
 # Custom sidebar templates, must be a dictionary that maps document names
 # to template names.
 #
 # The default sidebars (for documents that don't match any pattern) are
 # defined by theme itself.  Builtin themes are using these templates by
 # default: ``['localtoc.html', 'relations.html', 'sourcelink.html',
 # 'searchbox.html']``.
 #
 # html_sidebars = {}
 # -- Options for HTMLHelp output ---------------------------------------------
 # Output file base name for HTML help builder.
 htmlhelp_basename = 'dgldoc'
 # -- Options for LaTeX output ------------------------------------------------
 latex_elements = {
    # The paper size ('letterpaper' or 'a4paper').
    #
    # 'papersize': 'letterpaper',
    # The font size ('10pt', '11pt' or '12pt').
    #
    # 'pointsize': '10pt',
    # Additional stuff for the LaTeX preamble.
    #
    # 'preamble': '',
    # Latex figure (float) alignment
    #
    # 'figure_align': 'htbp',
 }
 # Grouping the document tree into LaTeX files. List of tuples
 # (source start file, target name, title,
 #  author, documentclass [howto, manual, or own class]).
 latex_documents = [
    (master_doc, 'dgl.tex', 'DGL Documentation',
     'DGL Team', 'manual'),
 ]
 # -- Options for manual page output ------------------------------------------
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
 man_pages = [
    (master_doc, 'dgl', 'DGL Documentation',
     [author], 1)
 ]
 # -- Options for Texinfo output ----------------------------------------------
 # Grouping the document tree into Texinfo files. List of tuples
 # (source start file, target name, title, author,
 #  dir menu entry, description, category)
 texinfo_documents = [
    (master_doc, 'dgl', 'DGL Documentation',
     author, 'dgl', 'Library for deep learning on graphs.',
     'Miscellaneous'),
 ]
 # -- Options for Epub output -------------------------------------------------
 # Bibliographic Dublin Core info.
 epub_title = project
 # The unique identifier of the text. This can be a ISBN number
 # or the project homepage.
 #
 # epub_identifier = ''
 # A unique identification for the text.
 #
 # epub_uid = ''
 # A list of files that should not be packed into the epub file.
 epub_exclude_files = ['search.html']
 # -- Extension configuration -------------------------------------------------
 autosummary_generate = True
 intersphinx_mapping = {
    'python': ('https://docs.python.org/{.major}'.format(sys.version_info), None),
    'numpy': ('http://docs.scipy.org/doc/numpy/', None),
    'scipy': ('http://docs.scipy.org/doc/scipy/reference', None),
    'matplotlib': ('http://matplotlib.org/', None),
    'networkx' : ('https://networkx.github.io/documentation/stable', None),
 }
 # sphinx gallery configurations
 from sphinx_gallery.sorting import FileNameSortKey
-examples_dirs = ['../../tutorials']  # path to find sources
+examples_dirs = ['../../tutorials/basics','../../tutorials/models']  # path to find sources
-gallery_dirs = ['tutorials']  # path to generate docs
+gallery_dirs = ['tutorials/basics','tutorials/models']  # path to generate docs
 reference_url = {
    'dgl' : None,
    'numpy': 'http://docs.scipy.org/doc/numpy/',
    'scipy': 'http://docs.scipy.org/doc/scipy/reference',
    'matplotlib': 'http://matplotlib.org/',
    'networkx' : 'https://networkx.github.io/documentation/stable',
 }
 sphinx_gallery_conf = {
    'backreferences_dir' : 'generated/backreferences',
    'doc_module' : ('dgl', 'numpy'),
    'examples_dirs' : examples_dirs,
    'gallery_dirs' : gallery_dirs,
    'within_subsection_order' : FileNameSortKey,
    'filename_pattern' : '.py',
 }
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -65,7 +65,8 @@ credit, see `here <https://www.dgl.ai/ack>`_.
   :caption: Tutorials
   :glob:
-   tutorials/index
+   tutorials/basics/index
+   tutorials/models/index
 .. toctree::
   :maxdepth: 2

--- a/tutorials/1_first.py
+++ b/tutorials/1_first.py
 """
 .. currentmodule:: dgl
 DGL at a Glance
 =========================
 **Author**: `Minjie Wang <https://jermainewang.github.io/>`_, Quan Gan, `Jake
 Zhao <https://cs.nyu.edu/~jakezhao/>`_, Zheng Zhang
 DGL is a Python package dedicated to deep learning on graphs, built atop
 existing tensor DL frameworks (e.g. Pytorch, MXNet) and simplifying the
 implementation of graph-based neural networks.
 The goal of this tutorial:
 - Understand how DGL enables computation on graph from a high level.
 - Train a simple graph neural network in DGL to classify nodes in a graph.
 At the end of this tutorial, we hope you get a brief feeling of how DGL works.
 *This tutorial assumes basic familiarity with pytorch.*
 """
 ###############################################################################
 # Step 0: Problem description
 # ---------------------------
 #
 # We start with the well-known "Zachary's karate club" problem. The karate club
 # is a social network which captures 34 members and document pairwise links
 # between members who interact outside the club.  The club later divides into
 # two communities led by the instructor (node 0) and the club president (node
 # 33). The network is visualized as follows with the color indicating the
 # community:
 #
 # .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/img/karate-club.png
 #    :align: center
 #
 # The task is to predict which side (0 or 33) each member tends to join given
 # the social network itself.
 ###############################################################################
 # Step 1: Creating a graph in DGL
 # -------------------------------
 # Creating the graph for Zachary's karate club goes as follows:
 import dgl
 def build_karate_club_graph():
    g = dgl.DGLGraph()
    # add 34 nodes into the graph; nodes are labeled from 0~33
    g.add_nodes(34)
    # all 78 edges as a list of tuples
    edge_list = [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2),
        (4, 0), (5, 0), (6, 0), (6, 4), (6, 5), (7, 0), (7, 1),
        (7, 2), (7, 3), (8, 0), (8, 2), (9, 2), (10, 0), (10, 4),
        (10, 5), (11, 0), (12, 0), (12, 3), (13, 0), (13, 1), (13, 2),
        (13, 3), (16, 5), (16, 6), (17, 0), (17, 1), (19, 0), (19, 1),
        (21, 0), (21, 1), (25, 23), (25, 24), (27, 2), (27, 23),
        (27, 24), (28, 2), (29, 23), (29, 26), (30, 1), (30, 8),
        (31, 0), (31, 24), (31, 25), (31, 28), (32, 2), (32, 8),
        (32, 14), (32, 15), (32, 18), (32, 20), (32, 22), (32, 23),
        (32, 29), (32, 30), (32, 31), (33, 8), (33, 9), (33, 13),
        (33, 14), (33, 15), (33, 18), (33, 19), (33, 20), (33, 22),
        (33, 23), (33, 26), (33, 27), (33, 28), (33, 29), (33, 30),
        (33, 31), (33, 32)]
    # add edges two lists of nodes: src and dst
    src, dst = tuple(zip(*edge_list))
    g.add_edges(src, dst)
    # edges are directional in DGL; make them bi-directional
    g.add_edges(dst, src)
    return g
 ###############################################################################
 # We can print out the number of nodes and edges in our newly constructed graph:
 G = build_karate_club_graph()
 print('We have %d nodes.' % G.number_of_nodes())
 print('We have %d edges.' % G.number_of_edges())
 ###############################################################################
 # We can also visualize the graph by converting it to a `networkx
 # <https://networkx.github.io/documentation/stable/>`_ graph:
 import networkx as nx
 # Since the actual graph is undirected, we convert it for visualization
 # purpose.
 nx_G = G.to_networkx().to_undirected()
 # Kamada-Kawaii layout usually looks pretty for arbitrary graphs
 pos = nx.kamada_kawai_layout(nx_G)
 nx.draw(nx_G, pos, with_labels=True, node_color=[[.7, .7, .7]])
 ###############################################################################
 # Step 2: assign features to nodes or edges
 # --------------------------------------------
 # Graph neural networks associate features with nodes and edges for training.
 # For our classification example, we assign each node's an input feature as a one-hot vector:
 # node :math:`v_i`'s feature vector is :math:`[0,\ldots,1,\dots,0]`,
 # where the :math:`i^{th}` position is one.
 #
 # In DGL, we can add features for all nodes at once, using a feature tensor that
 # batches node features along the first dimension. This code below adds the one-hot
 # feature for all nodes:
 import torch
 G.ndata['feat'] = torch.eye(34)
 ###############################################################################
 # We can print out the node features to verify:
 # print out node 2's input feature
 print(G.nodes[2].data['feat'])
 # print out node 10 and 11's input features
 print(G.nodes[[10, 11]].data['feat'])
 ###############################################################################
 # Step 3: define a Graph Convolutional Network (GCN)
 # --------------------------------------------------
 # To perform node classification, we use the Graph Convolutional Network
 # (GCN) developed by `Kipf and Welling <https://arxiv.org/abs/1609.02907>`_. Here
 # we provide the simpliest definition of a GCN framework, but we recommend the
 # reader to read the original paper for more details.
 #
 # - At layer :math:`l`, each node :math:`v_i^l` carries a feature vector :math:`h_i^l`.
 # - Each layer of the GCN tries to aggregate the features from :math:`u_i^{l}` where
 #   :math:`u_i`'s are neighborhood nodes to :math:`v` into the next layer representation at
 #   :math:`v_i^{l+1}`. This is followed by an affine transformation with some
 #   non-linearity.
 #
 # The above definition of GCN fits into a **message-passing** paradigm: each
 # node will update its own feature with information sent from neighboring
 # nodes. A graphical demonstration is displayed below.
 #
 # .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/1_first/mailbox.png
 #    :alt: mailbox
 #    :align: center
 #
 # Now, we show that the GCN layer can be easily implemented in DGL.
 import torch.nn as nn
 import torch.nn.functional as F
 # Define the message & reduce function
 # NOTE: we ignore the GCN's normalization constant c_ij for this tutorial.
 def gcn_message(edges):
    # The argument is a batch of edges.
    # This computes a (batch of) message called 'msg' using the source node's feature 'h'.
    return {'msg' : edges.src['h']}
 def gcn_reduce(nodes):
    # The argument is a batch of nodes.
    # This computes the new 'h' features by summing received 'msg' in each node's mailbox.
    return {'h' : torch.sum(nodes.mailbox['msg'], dim=1)}
 # Define the GCNLayer module
 class GCNLayer(nn.Module):
    def __init__(self, in_feats, out_feats):
        super(GCNLayer, self).__init__()
        self.linear = nn.Linear(in_feats, out_feats)
    def forward(self, g, inputs):
        # g is the graph and the inputs is the input node features
        # first set the node features
        g.ndata['h'] = inputs
        # trigger message passing on all edges 
        g.send(g.edges(), gcn_message)
        # trigger aggregation at all nodes
        g.recv(g.nodes(), gcn_reduce)
        # get the result node features
        h = g.ndata.pop('h')
        # perform linear transformation
        return self.linear(h)
 ###############################################################################
 # In general, the nodes send information computed via the *message functions*,
 # and aggregates incoming information with the *reduce functions*.
 #
 # We then define a deeper GCN model that contains two GCN layers:
 # Define a 2-layer GCN model
 class GCN(nn.Module):
    def __init__(self, in_feats, hidden_size, num_classes):
        super(GCN, self).__init__()
        self.gcn1 = GCNLayer(in_feats, hidden_size)
        self.gcn2 = GCNLayer(hidden_size, num_classes)
    def forward(self, g, inputs):
        h = self.gcn1(g, inputs)
        h = torch.relu(h)
        h = self.gcn2(g, h)
        return h
 # The first layer transforms input features of size of 34 to a hidden size of 5.
 # The second layer transforms the hidden layer and produces output features of
 # size 2, corresponding to the two groups of the karate club.
 net = GCN(34, 5, 2)
 ###############################################################################
 # Step 4: data preparation and initialization
 # -------------------------------------------
 #
 # We use one-hot vectors to initialize the node features. Since this is a
 # semi-supervised setting, only the instructor (node 0) and the club president
 # (node 33) are assigned labels. The implementation is available as follow.
 inputs = torch.eye(34)
 labeled_nodes = torch.tensor([0, 33])  # only the instructor and the president nodes are labeled
 labels = torch.tensor([0, 1])  # their labels are different
 ###############################################################################
 # Step 5: train then visualize
 # ----------------------------
 # The training loop is exactly the same as other PyTorch models.
 # We (1) create an optimizer, (2) feed the inputs to the model,
 # (3) calculate the loss and (4) use autograd to optimize the model.
 optimizer = torch.optim.Adam(net.parameters(), lr=0.01)
 all_logits = []
 for epoch in range(30):
    logits = net(G, inputs)
    # we save the logits for visualization later
    all_logits.append(logits.detach())
    logp = F.log_softmax(logits, 1)
    # we only compute loss for labeled nodes
    loss = F.nll_loss(logp[labeled_nodes], labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print('Epoch %d | Loss: %.4f' % (epoch, loss.item()))
 ###############################################################################
 # This is a rather toy example, so it does not even have a validation or test
 # set. Instead, Since the model produces an output feature of size 2 for each node, we can
 # visualize by plotting the output feature in a 2D space.
 # The following code animates the training process from initial guess
 # (where the nodes are not classified correctly at all) to the end
 # (where the nodes are linearly separable).
 import matplotlib.animation as animation
 import matplotlib.pyplot as plt
 def draw(i):
    cls1color = '#00FFFF'
    cls2color = '#FF00FF'
    pos = {}
    colors = []
    for v in range(34):
        pos[v] = all_logits[i][v].numpy()
        cls = pos[v].argmax()
        colors.append(cls1color if cls else cls2color)
    ax.cla()
    ax.axis('off')
    ax.set_title('Epoch: %d' % i)
    nx.draw_networkx(nx_G.to_undirected(), pos, node_color=colors,
            with_labels=True, node_size=300, ax=ax)
 fig = plt.figure(dpi=150)
 fig.clf()
 ax = fig.subplots()
 draw(0)  # draw the prediction of the first epoch
 plt.close()
 ###############################################################################
 # .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/1_first/karate0.png
 #    :height: 300px
 #    :width: 400px
 #    :align: center
 ###############################################################################
 # The following animation shows how the model correctly predicts the community
 # after a series of training epochs.
 ani = animation.FuncAnimation(fig, draw, frames=len(all_logits), interval=200)
 ###############################################################################
 # .. image:: https://s3.us-east-2.amazonaws.com/dgl.ai/tutorial/1_first/karate.gif
 #    :height: 300px
 #    :width: 400px
 #    :align: center
 ###############################################################################
 # Next steps
 # ----------
 # 
 # In the :doc:`next tutorial <2_basics>`, we will go through some more basics
 # of DGL, such as reading and writing node/edge features.
--- a/tutorials/2_basics.py
+++ b/tutorials/2_basics.py
 """
 .. currentmodule:: dgl
 DGL Basics
 ==========
 **Author**: `Minjie Wang <https://jermainewang.github.io/>`_, Quan Gan, Yu Gai,
 Zheng Zhang
 The Goal of this tutorial:
 * To create a graph.
 * To read and write node and edge representations.
 """
 ###############################################################################
 # Graph Creation
 # --------------
 # The design of :class:`DGLGraph` was influenced by other graph libraries. Indeed,
 # you can create a graph from networkx, and convert it into a :class:`DGLGraph` and
 # vice versa:
 import networkx as nx
 import dgl
 g_nx = nx.petersen_graph()
 g_dgl = dgl.DGLGraph(g_nx)
 import matplotlib.pyplot as plt
 plt.subplot(121)
 nx.draw(g_nx, with_labels=True)
 plt.subplot(122)
 nx.draw(g_dgl.to_networkx(), with_labels=True)
 plt.show()
 ###############################################################################
 # They are the same graph, except that :class:`DGLGraph` is *always* directional.
 #
 # One can also create a graph by calling DGL's own interface.
 # 
 # Now let's build a star graph. :class:`DGLGraph` nodes are consecutive range of
 # integers between 0 and :func:`number_of_nodes() <DGLGraph.number_of_nodes>`
 # and can grow by calling :func:`add_nodes <DGLGraph.add_nodes>`.
 # :class:`DGLGraph` edges are in order of their additions. Note that
 # edges are accessed in much the same way as nodes, with one extra feature
 # of *edge broadcasting*:
 import dgl
 import torch as th
 g = dgl.DGLGraph()
 g.add_nodes(10)
 # a couple edges one-by-one
 for i in range(1, 4):
    g.add_edge(i, 0)
 # a few more with a paired list
 src = list(range(5, 8)); dst = [0]*3
 g.add_edges(src, dst)
 # finish with a pair of tensors
 src = th.tensor([8, 9]); dst = th.tensor([0, 0])
 g.add_edges(src, dst)
 # edge broadcasting will do star graph in one go!
 g.clear(); g.add_nodes(10)
 src = th.tensor(list(range(1, 10)));
 g.add_edges(src, 0)
 import networkx as nx
 import matplotlib.pyplot as plt
 nx.draw(g.to_networkx(), with_labels=True)
 plt.show()
 ###############################################################################
 # Feature Assignment
 # ------------------
 # One can also assign features to nodes and edges of a :class:`DGLGraph`.  The
 # features are represented as dictionary of names (strings) and tensors,
 # called **fields**.
 #
 # The following code snippet assigns each node a vector (len=3).
 #
 # .. note::
 #
 #    DGL aims to be framework-agnostic, and currently it supports PyTorch and
 #    MXNet tensors. From now on, we use PyTorch as an example.
 import dgl
 import torch as th
 x = th.randn(10, 3)
 g.ndata['x'] = x
 ###############################################################################
 # :func:`ndata <DGLGraph.ndata>` is a syntax sugar to access states of all nodes,
 # states are stored
 # in a container ``data`` that hosts user defined dictionary.
 print(g.ndata['x'] == g.nodes[:].data['x'])
 # access node set with integer, list, or integer tensor
 g.nodes[0].data['x'] = th.zeros(1, 3)
 g.nodes[[0, 1, 2]].data['x'] = th.zeros(3, 3)
 g.nodes[th.tensor([0, 1, 2])].data['x'] = th.zeros(3, 3)
 ###############################################################################
 # Assigning edge features is in a similar fashion to that of node features,
 # except that one can also do it by specifying endpoints of the edges.
 g.edata['w'] = th.randn(9, 2)
 # access edge set with IDs in integer, list, or integer tensor
 g.edges[1].data['w'] = th.randn(1, 2)
 g.edges[[0, 1, 2]].data['w'] = th.zeros(3, 2)
 g.edges[th.tensor([0, 1, 2])].data['w'] = th.zeros(3, 2)
 # one can also access the edges by giving endpoints
 g.edges[1, 0].data['w'] = th.ones(1, 2)                 # edge 1 -> 0
 g.edges[[1, 2, 3], [0, 0, 0]].data['w'] = th.ones(3, 2) # edges [1, 2, 3] -> 0
 ###############################################################################
 # After assignments, each node/edge field will be associated with a scheme
 # containing the shape and data type (dtype) of its field value.
 print(g.node_attr_schemes())
 g.ndata['x'] = th.zeros((10, 4))
 print(g.node_attr_schemes())
 ###############################################################################
 # One can also remove node/edge states from the graph. This is particularly
 # useful to save memory during inference.
 g.ndata.pop('x')
 g.edata.pop('w')
 ###############################################################################
 # Multigraphs
 # ~~~~~~~~~~~
 # Many graph applications need multi-edges. To enable this, construct :class:`DGLGraph`
 # with ``multigraph=True``.
 g_multi = dgl.DGLGraph(multigraph=True)
 g_multi.add_nodes(10)
 g_multi.ndata['x'] = th.randn(10, 2)
 g_multi.add_edges(list(range(1, 10)), 0)
 g_multi.add_edge(1, 0) # two edges on 1->0
 g_multi.edata['w'] = th.randn(10, 2)
 g_multi.edges[1].data['w'] = th.zeros(1, 2)
 print(g_multi.edges())
 ###############################################################################
 # An edge in multi-graph cannot be uniquely identified using its incident nodes
 # :math:`u` and :math:`v`; query their edge ids use ``edge_id`` interface.
 eid_10 = g_multi.edge_id(1, 0)
 g_multi.edges[eid_10].data['w'] = th.ones(len(eid_10), 2)
 print(g_multi.edata['w'])
 ###############################################################################
 # .. note::
 #
 #    * Nodes and edges can be added but not removed; we will support removal in
 #      the future.
 #    * Updating a feature of different schemes raise error on indivdual node (or
 #      node subset).
 ###############################################################################
 # Next steps
 # ----------
 # In the :doc:`next tutorial <3_pagerank>`, we will go through the
 # DGL message passing interface by implementing PageRank.
--- a/tutorials/3_pagerank.py
+++ b/tutorials/3_pagerank.py
 """
 .. currentmodule:: dgl
 PageRank with DGL Message Passing
 =================================
 **Author**: `Minjie Wang <https://jermainewang.github.io/>`_, Quan Gan, Yu Gai,
 Zheng Zhang
 In this section we illustrate the usage of different levels of message
 passing API with PageRank on a small graph. In DGL, the message passing and
 feature transformations are all **User-Defined Functions** (UDFs).
 The goal of this tutorial: to implement PageRank using DGL message passing
 interface.
 """
 ###############################################################################
 # The PageRank Algorithm
 # ----------------------
 # In each iteration of PageRank, every node (web page) first scatters its
 # PageRank value uniformly to its downstream nodes. The new PageRank value of
 # each node is computed by aggregating the received PageRank values from its
 # neighbors, which is then adjusted by the damping factor:
 #
 # .. math::
 #
 #    PV(u) = \frac{1-d}{N} + d \times \sum_{v \in \mathcal{N}(u)}
 #    \frac{PV(v)}{D(v)}
 #
 # where :math:`N` is the number of nodes in the graph; :math:`D(v)` is the
 # out-degree of a node :math:`v`; and :math:`\mathcal{N}(u)` is the neighbor
 # nodes.
 ###############################################################################
 # A naive implementation
 # ----------------------
 # Let us first create a graph with 100 nodes with NetworkX and convert it to a
 # :class:`DGLGraph`:
 import networkx as nx
 import matplotlib.pyplot as plt
 import torch
 import dgl
 N = 100  # number of nodes
 DAMP = 0.85  # damping factor
 K = 10  # number of iterations
 g = nx.nx.erdos_renyi_graph(N, 0.1)
 g = dgl.DGLGraph(g)
 nx.draw(g.to_networkx(), node_size=50, node_color=[[.5, .5, .5,]])
 plt.show()
 ###############################################################################
 # According to the algorithm, PageRank consists of two phases in a typical
 # scatter-gather pattern. We first initialize the PageRank value of each node
 # to :math:`\frac{1}{N}` and store each node's out-degree as a node feature:
 g.ndata['pv'] = torch.ones(N) / N
 g.ndata['deg'] = g.out_degrees(g.nodes()).float()
 ###############################################################################
 # We then define the message function, which divides every node's PageRank
 # value by its out-degree and passes the result as message to its neighbors:
 def pagerank_message_func(edges):
    return {'pv' : edges.src['pv'] / edges.src['deg']}
 ###############################################################################
 # In DGL, the message functions are expressed as **Edge UDFs**.  Edge UDFs
 # take in a single argument ``edges``.  It has three members ``src``, ``dst``,
 # and ``data`` for accessing source node features, destination node features,
 # and edge features respectively.  Here, the function computes messages only
 # from source node features.
 #
 # Next, we define the reduce function, which removes and aggregates the
 # messages from its ``mailbox``, and computes its new PageRank value:
 def pagerank_reduce_func(nodes):
    msgs = torch.sum(nodes.mailbox['pv'], dim=1)
    pv = (1 - DAMP) / N + DAMP * msgs
    return {'pv' : pv}
 ###############################################################################
 # The reduce functions are **Node UDFs**.  Node UDFs have a single argument
 # ``nodes``, which has two members ``data`` and ``mailbox``.  ``data``
 # contains the node features while ``mailbox`` contains all incoming message
 # features, stacked along the second dimension (hence the ``dim=1`` argument).
 #
 # The message UDF works on a batch of edges, whereas the reduce UDF works on
 # a batch of edges but outputs a batch of nodes. Their relationships are as
 # follows:
 #
 # .. image:: https://i.imgur.com/kIMiuFb.png
 #
 # We register the message function and reduce function, which will be called
 # later by DGL.
 g.register_message_func(pagerank_message_func)
 g.register_reduce_func(pagerank_reduce_func)
 ###############################################################################
 # The algorithm is then very straight-forward. Here is the code for one
 # PageRank iteration:
 def pagerank_naive(g):
    # Phase #1: send out messages along all edges.
    for u, v in zip(*g.edges()):
        g.send((u, v))
    # Phase #2: receive messages to compute new PageRank values.
    for v in g.nodes():
        g.recv(v)
 ###############################################################################
 # Improvement with batching semantics
 # -----------------------------------
 # The above code does not scale to large graph because it iterates over all
 # the nodes. DGL solves this by letting user compute on a *batch* of nodes or
 # edges. For example, the following codes trigger message and reduce functions
 # on multiple nodes and edges at once.
 def pagerank_batch(g):
    g.send(g.edges())
    g.recv(g.nodes())
 ###############################################################################
 # Note that we are still using the same reduce function ``pagerank_reduce_func``,
 # where ``nodes.mailbox['pv']`` is a *single* tensor, stacking the incoming
 # messages along the second dimension.
 #
 # Naturally, one will wonder if this is even possible to perform reduce on all
 # nodes in parallel, since each node may have different number of incoming
 # messages and one cannot really "stack" tensors of different lengths together.
 # In general, DGL solves the problem by grouping the nodes by the number of
 # incoming messages, and calling the reduce function for each group.
 ###############################################################################
 # More improvement with higher level APIs
 # ---------------------------------------
 # DGL provides many routines that combines basic ``send`` and ``recv`` in
 # various ways. They are called **level-2 APIs**. For example, the PageRank
 # example can be further simplified as follows:
 def pagerank_level2(g):
    g.update_all()
 ###############################################################################
 # Besides ``update_all``, we also have ``pull``, ``push``, and ``send_and_recv``
-# in this level-2 category. Please refer to the :doc:`API reference <../api/python/graph>`
+# in this level-2 category. Please refer to the :doc:`API reference <../../api/python/graph>`
 # for more details.
 ###############################################################################
 # Even more improvement with DGL builtin functions
 # ------------------------------------------------
 # As some of the message and reduce functions are very commonly used, DGL also
 # provides **builtin functions**. For example, two builtin functions can be
 # used in the PageRank example.
 #
 # * :func:`dgl.function.copy_src(src, out) <function.copy_src>`
 #   is an edge UDF that computes the
 #   output using the source node feature data. User needs to specify the name of
 #   the source feature data (``src``) and the output name (``out``).
 # 
 # * :func:`dgl.function.sum(msg, out) <function.sum>` is a node UDF
 #   that sums the messages in
 #   the node's mailbox. User needs to specify the message name (``msg``) and the
 #   output name (``out``).
 #
 # For example, the PageRank example can be rewritten as following:
 import dgl.function as fn
 def pagerank_builtin(g):
    g.ndata['pv'] = g.ndata['pv'] / g.ndata['deg']
    g.update_all(message_func=fn.copy_src(src='pv', out='m'),
                 reduce_func=fn.sum(msg='m',out='m_sum'))
    g.ndata['pv'] = (1 - DAMP) / N + DAMP * g.ndata['m_sum']
 ###############################################################################
 # Here, we directly provide the UDFs to the :func:`update_all <DGLGraph.update_all>`
 # as its arguments.
 # This will override the previously registered UDFs.
 #
 # In addition to cleaner code, using builtin functions also gives DGL the
 # opportunity to fuse operations together, resulting in faster execution.  For
 # example, DGL will fuse the ``copy_src`` message function and ``sum`` reduce
 # function into one sparse matrix-vector (spMV) multiplication.
 #
 # `This section <spmv_>`_ describes why spMV can speed up the scatter-gather
 # phase in PageRank.  For more details about the builtin functions in DGL,
-# please read the :doc:`API reference <../api/python/function>`.
+# please read the :doc:`API reference <../../api/python/function>`.
 #
 # You can also download and run the codes to feel the difference.
 for k in range(K):
    # Uncomment the corresponding line to select different version.
    # pagerank_naive(g)
    # pagerank_batch(g)
    # pagerank_level2(g)
    pagerank_builtin(g)
 print(g.ndata['pv'])
 ###############################################################################
 # .. _spmv:
 #
 # Using spMV for PageRank
 # -----------------------
 # Using builtin functions allows DGL to understand the semantics of UDFs and
 # thus allows more efficient implementation for you. For example, in the case
 # of PageRank, one common trick to accelerate it is using its linear algebra
 # form.
 #
 # .. math::
 #
 #    \mathbf{R}^{k} = \frac{1-d}{N} \mathbf{1} + d \mathbf{A}*\mathbf{R}^{k-1}
 #
 # Here, :math:`\mathbf{R}^k` is the vector of the PageRank values of all nodes
 # at iteration :math:`k`; :math:`\mathbf{A}` is the sparse adjacency matrix
 # of the graph.
 # Computing this equation is quite efficient because there exists efficient
 # GPU kernel for the *sparse-matrix-vector-multiplication* (spMV). DGL
 # detects whether such optimization is available through the builtin
 # functions. If the certain combination of builtins can be mapped to a spMV
 # kernel (e.g. the pagerank example), DGL will use it automatically. As a
 # result, *we recommend using builtin functions whenever it is possible*.
 ###############################################################################
 # Next steps
 # ----------
-# Check out :doc:`GCN <models/1_gcn>` and :doc:`Capsule <models/2_capsule>`
+# Check out :doc:`GCN <../models/1_gnn/1_gcn>` and :doc:`Capsule <../models/4_old_wines/2_capsule>`
 # for more model implemenetations in DGL.
--- a/tutorials/README.txt
+++ b/tutorials/README.txt
 Basic Tutorials
 ===============
-These tutorials conver the basics of DGL.
+These tutorials cover the basics of DGL.
--- a/tutorials/models/1_gcn.py
+++ b/tutorials/models/1_gcn.py
 """
 .. _model-gcn:
 Graph Convolutional Network
 ====================================
 **Author:** `Qi Huang <https://github.com/HQ01>`_, `Minjie Wang  <https://jermainewang.github.io/>`_,
 Yu Gai, Quan Gan, Zheng Zhang
 This is a gentle introduction of using DGL to implement Graph Convolutional
 Networks (Kipf & Welling et al., `Semi-Supervised Classificaton with Graph
 Convolutional Networks <https://arxiv.org/pdf/1609.02907.pdf>`_). We build upon
-the :doc:`earlier tutorial <../3_pagerank>` on DGLGraph and demonstrate
+the :doc:`earlier tutorial <../../basics/3_pagerank>` on DGLGraph and demonstrate
 how DGL combines graph with deep neural network and learn structural representations.
 """
 ###############################################################################
 # Model Overview
 # ------------------------------------------
 # GCN from the perspective of message passing
 # ```````````````````````````````````````````````
 # We describe a layer of graph convolutional neural network from a message
 # passing perspective; the math can be found `here <math_>`_.
 # It boils down to the following step, for each node :math:`u`:
 # 
 # 1) Aggregate neighbors' representations :math:`h_{v}` to produce an
 # intermediate representation :math:`\hat{h}_u`.  2) Transform the aggregated
 # representation :math:`\hat{h}_{u}` with a linear projection followed by a
 # non-linearity: :math:`h_{u} = f(W_{u} \hat{h}_u)`.
 # 
 # We will implement step 1 with DGL message passing, and step 2 with the
 # ``apply_nodes`` method, whose node UDF will be a PyTorch ``nn.Module``.
 # 
 # GCN implementation with DGL
 # ``````````````````````````````````````````
 # We first define the message and reduce function as usual.  Since the
 # aggregation on a node :math:`u` only involves summing over the neighbors'
 # representations :math:`h_v`, we can simply use builtin functions:
 import dgl
 import dgl.function as fn
 import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
 from dgl import DGLGraph
 gcn_msg = fn.copy_src(src='h', out='m')
 gcn_reduce = fn.sum(msg='m', out='h')
 ###############################################################################
 # We then define the node UDF for ``apply_nodes``, which is a fully-connected layer:
 class NodeApplyModule(nn.Module):
    def __init__(self, in_feats, out_feats, activation):
        super(NodeApplyModule, self).__init__()
        self.linear = nn.Linear(in_feats, out_feats)
        self.activation = activation
    def forward(self, node):
        h = self.linear(node.data['h'])
        h = self.activation(h)
        return {'h' : h}
 ###############################################################################
 # We then proceed to define the GCN module. A GCN layer essentially performs
 # message passing on all the nodes then applies the `NodeApplyModule`. Note
 # that we omitted the dropout in the paper for simplicity.
 class GCN(nn.Module):
    def __init__(self, in_feats, out_feats, activation):
        super(GCN, self).__init__()
        self.apply_mod = NodeApplyModule(in_feats, out_feats, activation)
    def forward(self, g, feature):
        g.ndata['h'] = feature
        g.update_all(gcn_msg, gcn_reduce)
        g.apply_nodes(func=self.apply_mod)
        return g.ndata.pop('h')
 ###############################################################################
 # The forward function is essentially the same as any other commonly seen NNs
 # model in PyTorch.  We can initialize GCN like any ``nn.Module``. For example,
 # let's define a simple neural network consisting of two GCN layers. Suppose we
 # are training the classifier for the cora dataset (the input feature size is
 # 1433 and the number of classes is 7).
 class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.gcn1 = GCN(1433, 16, F.relu)
        self.gcn2 = GCN(16, 7, F.relu)
    def forward(self, g, features):
        x = self.gcn1(g, features)
        x = self.gcn2(g, x)
        return x
 net = Net()
 print(net)
 ###############################################################################
 # We load the cora dataset using DGL's built-in data module.
 from dgl.data import citation_graph as citegrh
 def load_cora_data():
    data = citegrh.load_cora()
    features = th.FloatTensor(data.features)
    labels = th.LongTensor(data.labels)
    mask = th.ByteTensor(data.train_mask)
    g = DGLGraph(data.graph)
    return g, features, labels, mask
 ###############################################################################
 # We then train the network as follows:
 import time
 import numpy as np
 g, features, labels, mask = load_cora_data()
 optimizer = th.optim.Adam(net.parameters(), lr=1e-3)
 dur = []
 for epoch in range(30):
    if epoch >=3:
        t0 = time.time()
    logits = net(g, features)
    logp = F.log_softmax(logits, 1)
    loss = F.nll_loss(logp[mask], labels[mask])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if epoch >=3:
        dur.append(time.time() - t0)
    print("Epoch {:05d} | Loss {:.4f} | Time(s) {:.4f}".format(
            epoch, loss.item(), np.mean(dur)))
 ###############################################################################
 # .. _math:
 #
 # GCN in one formula
 # ------------------
 # Mathematically, the GCN model follows this formula:
 # 
 # :math:`H^{(l+1)} = \sigma(\tilde{D}^{-\frac{1}{2}}\tilde{A}\tilde{D}^{-\frac{1}{2}}H^{(l)}W^{(l)})`
 # 
 # Here, :math:`H^{(l)}` denotes the :math:`l^{th}` layer in the network,
 # :math:`\sigma` is the non-linearity, and :math:`W` is the weight matrix for
 # this layer. :math:`D` and :math:`A`, as commonly seen, represent degree
 # matrix and adjacency matrix, respectively. The ~ is a renormalization trick
 # in which we add a self-connection to each node of the graph, and build the
 # corresponding degree and adjacency matrix.  The shape of the input
 # :math:`H^{(0)}` is :math:`N \times D`, where :math:`N` is the number of nodes
 # and :math:`D` is the number of input features. We can chain up multiple
 # layers as such to produce a node-level representation output with shape
 # :math`N \times F`, where :math:`F` is the dimension of the output node
 # feature vector.
 # 
 # The equation can be efficiently implemented using sparse matrix
 # multiplication kernels (such as Kipf's
 # `pygcn <https://github.com/tkipf/pygcn>`_ code). The above DGL implementation
 # in fact has already used this trick due to the use of builtin functions. To
-# understand what is under the hood, please read our tutorial on :doc:`PageRank <../3_pagerank>`.
+# understand what is under the hood, please read our tutorial on :doc:`PageRank <../../basics/3_pagerank>`.
--- a/tutorials/models/4_rgcn.py
+++ b/tutorials/models/4_rgcn.py
 """
 .. _model-rgcn:
 Relational Graph Convolutional Network Tutorial
 ================================================
 **Author:** Lingfan Yu, Mufei Li, Zheng Zhang
 The vanilla Graph Convolutional Network (GCN)
 (`paper <https://arxiv.org/pdf/1609.02907.pdf>`_,
 `DGL tutorial <http://doc.dgl.ai/tutorials/index.html>`_) exploits
 structural information of the dataset (i.e. the graph connectivity) to
 improve the extraction of node representations. Graph edges are left as
 untyped.
 A knowledge graph is made up by a collection of triples of the form
 (subject, relation, object). Edges thus encode important information and
 have their own embeddings to be learned. Furthermore, there may exist
 multiple edges among any given pair.
 A recent model Relational-GCN (R-GCN) from the paper
 `Modeling Relational Data with Graph Convolutional
 Networks <https://arxiv.org/pdf/1703.06103.pdf>`_ is one effort to
 generalize GCN to handle different relations between entities in knowledge
 base. This tutorial shows how to implement R-GCN with DGL.
 """
 ###############################################################################
 # R-GCN: a brief introduction
 # ---------------------------
 # In *statistical relational learning* (SRL), there are two fundamental
 # tasks:
 #
 # - **Entity classification**, i.e., assign types and categorical
 #   properties to entities.
 # - **Link prediction**, i.e., recover missing triples.
 #
 # In both cases, missing information are expected to be recovered from
 # neighborhood structure of the graph. Here is the example from the R-GCN
 # paper:
 #
 # "Knowing that Mikhail Baryshnikov was educated at the Vaganova Academy
 # implies both that Mikhail Baryshnikov should have the label person, and
 # that the triple (Mikhail Baryshnikov, lived in, Russia) must belong to the
 # knowledge graph."
 #
 # R-GCN solves these two problems using a common graph convolutional network
 # extended with multi-edge encoding to compute embedding of the entities, but
 # with different downstream processing:
 #
 # - Entity classification is done by attaching a softmax classifier at the
 #   final embedding of an entity (node). Training is through loss of standard
 #   cross-entropy.
 # - Link prediction is done by reconstructing an edge with an autoencoder
 #   architecture, using a parameterized score function. Training uses negative
 #   sampling.
 #
 # This tutorial will focus on the first task to show how to generate entity
 # representation. `Complete
 # code <https://github.com/jermainewang/dgl/tree/rgcn/examples/pytorch/rgcn>`_
 # for both tasks can be found in DGL's github repository.
 #
 # Key ideas of R-GCN
 # -------------------
 # Recall that in GCN, the hidden representation for each node :math:`i` at
 # :math:`(l+1)^{th}` layer is computed by:
 #
 # .. math:: h_i^{l+1} = \sigma\left(\sum_{j\in N_i}\frac{1}{c_i} W^{(l)} h_j^{(l)}\right)~~~~~~~~~~(1)\\
 #
 # where :math:`c_i` is a normalization constant.
 #
 # The key difference between R-GCN and GCN is that in R-GCN, edges can
 # represent different relations. In GCN, weight :math:`W^{(l)}` in equation
 # :math:`(1)` is shared by all edges in layer :math:`l`. In contrast, in
 # R-GCN, different edge types use different weights and only edges of the
 # same relation type :math:`r` are associated with the same projection weight
 # :math:`W_r^{(l)}`.
 #
 # So the hidden representation of entities in :math:`(l+1)^{th}` layer in
 # R-GCN can be formulated as the following equation:
 #
 # .. math:: h_i^{l+1} = \sigma\left(W_0^{(l)}h_i^{(l)}+\sum_{r\in R}\sum_{j\in N_i^r}\frac{1}{c_{i,r}}W_r^{(l)}h_j^{(l)}\right)~~~~~~~~~~(2)\\
 #
 # where :math:`N_i^r` denotes the set of neighbor indices of node :math:`i`
 # under relation :math:`r\in R` and :math:`c_{i,r}` is a normalization
 # constant. In entity classification, the R-GCN paper uses
 # :math:`c_{i,r}=|N_i^r|`.
 #
 # The problem of applying the above equation directly is rapid growth of
 # number of parameters, especially with highly multi-relational data. In
 # order to reduce model parameter size and prevent overfitting, the original
 # paper proposes to use basis decomposition:
 #
 # .. math:: W_r^{(l)}=\sum\limits_{b=1}^B a_{rb}^{(l)}V_b^{(l)}~~~~~~~~~~(3)\\
 #
 # Therefore, the weight :math:`W_r^{(l)}` is a linear combination of basis
 # transformation :math:`V_b^{(l)}` with coefficients :math:`a_{rb}^{(l)}`.
 # The number of bases :math:`B` is much smaller than the number of relations
 # in the knowledge base.
 #
 # .. note::
 #    Another weight regularization, block-decomposition, is implemented in
 #    the `link prediction <link-prediction_>`_.
 #
 # Implement R-GCN in DGL
 # ----------------------
 #
 # An R-GCN model is composed of several R-GCN layers. The first R-GCN layer
 # also serves as input layer and takes in features (e.g. description texts)
 # associated with node entity and project to hidden space. In this tutorial,
 # we only use entity id as entity feature.
 #
 # R-GCN Layers
 # ~~~~~~~~~~~~
 #
 # For each node, an R-GCN layer performs the following steps:
 #
 # - Compute outgoing message using node representation and weight matrix
 #   associated with the edge type (message function)
 # - Aggregate incoming messages and generate new node representations (reduce
 #   and apply function)
 #
 # The following is the definition of an R-GCN hidden layer.
 #
 # .. note::
 #    Each relation type is associated with a different weight. Therefore,
 #    the full weight matrix has three dimensions: relation, input_feature,
 #    output_feature.
 #
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from dgl import DGLGraph
 import dgl.function as fn
 from functools import partial
 class RGCNLayer(nn.Module):
    def __init__(self, in_feat, out_feat, num_rels, num_bases=-1, bias=None,
                 activation=None, is_input_layer=False):
        super(RGCNLayer, self).__init__()
        self.in_feat = in_feat
        self.out_feat = out_feat
        self.num_rels = num_rels
        self.num_bases = num_bases
        self.bias = bias
        self.activation = activation
        self.is_input_layer = is_input_layer
        # sanity check
        if self.num_bases <= 0 or self.num_bases > self.num_rels:
            self.num_bases = self.num_rels
        # weight bases in equation (3)
        self.weight = nn.Parameter(torch.Tensor(self.num_bases, self.in_feat,
                                                self.out_feat))
        if self.num_bases < self.num_rels:
            # linear combination coefficients in equation (3)
            self.w_comp = nn.Parameter(torch.Tensor(self.num_rels, self.num_bases))
        # add bias
        if self.bias:
            self.bias = nn.Parameter(torch.Tensor(out_feat))
        # init trainable parameters
        nn.init.xavier_uniform_(self.weight,
                                gain=nn.init.calculate_gain('relu'))
        if self.num_bases < self.num_rels:
            nn.init.xavier_uniform_(self.w_comp,
                                    gain=nn.init.calculate_gain('relu'))
        if self.bias:
            nn.init.xavier_uniform_(self.bias,
                                    gain=nn.init.calculate_gain('relu'))
    def forward(self, g):
        if self.num_bases < self.num_rels:
            # generate all weights from bases (equation (3))
            weight = self.weight.view(self.in_feat, self.num_bases, self.out_feat)
            weight = torch.matmul(self.w_comp, weight).view(self.num_rels,
                                                        self.in_feat, self.out_feat)
        else:
            weight = self.weight
        if self.is_input_layer:
            def message_func(edges):
                # for input layer, matrix multiply can be converted to be
                # an embedding lookup using source node id
                embed = weight.view(-1, self.out_feat)
                index = edges.data['rel_type'] * self.in_feat + edges.src['id']
                return {'msg': embed[index] * edges.data['norm']}
        else:
            def message_func(edges):
                w = weight[edges.data['rel_type']]
                msg = torch.bmm(edges.src['h'].unsqueeze(1), w).squeeze()
                msg = msg * edges.data['norm']
                return {'msg': msg}
        def apply_func(nodes):
            h = nodes.data['h']
            if self.bias:
                h = h + self.bias
            if self.activation:
                h = self.activation(h)
            return {'h': h}
        g.update_all(message_func, fn.sum(msg='msg', out='h'), apply_func)
 ###############################################################################
 # Define full R-GCN model
 # ~~~~~~~~~~~~~~~~~~~~~~~
 class Model(nn.Module):
    def __init__(self, num_nodes, h_dim, out_dim, num_rels,
                 num_bases=-1, num_hidden_layers=1):
        super(Model, self).__init__()
        self.num_nodes = num_nodes
        self.h_dim = h_dim
        self.out_dim = out_dim
        self.num_rels = num_rels
        self.num_bases = num_bases
        self.num_hidden_layers = num_hidden_layers
        # create rgcn layers
        self.build_model()
        # create initial features
        self.features = self.create_features()
    def build_model(self):
        self.layers = nn.ModuleList()
        # input to hidden
        i2h = self.build_input_layer()
        self.layers.append(i2h)
        # hidden to hidden
        for idx in range(self.num_hidden_layers):
            h2h = self.build_hidden_layer(idx)
            self.layers.append(h2h)
        # hidden to output
        h2o = self.build_output_layer()
        self.layers.append(h2o)
    # initialize feature for each node
    def create_features(self):
        features = torch.arange(self.num_nodes)
        return features
    def build_input_layer(self):
        return RGCNLayer(self.num_nodes, self.h_dim, self.num_rels, self.num_bases,
                         activation=F.relu, is_input_layer=True)
    def build_hidden_layer(self):
        return RGCNLayer(self.h_dim, self.h_dim, self.num_rels, self.num_bases,
                         activation=F.relu)
    def build_output_layer(self):
        return RGCNLayer(self.h_dim, self.out_dim, self.num_rels, self.num_bases,
                         activation=partial(F.softmax, dim=1))
    def forward(self, g):
        if self.features is not None:
            g.ndata['id'] = self.features
        for layer in self.layers:
            layer(g)
        return g.ndata.pop('h')
 ###############################################################################
 # Handle dataset
 # ~~~~~~~~~~~~~~~~
 # In this tutorial, we use AIFB dataset from R-GCN paper:
 # load graph data
 from dgl.contrib.data import load_data
 import numpy as np
 data = load_data(dataset='aifb')
 num_nodes = data.num_nodes
 num_rels = data.num_rels
 num_classes = data.num_classes
 labels = data.labels
 train_idx = data.train_idx
 # split training and validation set
 val_idx = train_idx[:len(train_idx) // 5]
 train_idx = train_idx[len(train_idx) // 5:]
 # edge type and normalization factor
 edge_type = torch.from_numpy(data.edge_type)
 edge_norm = torch.from_numpy(data.edge_norm).unsqueeze(1)
 labels = torch.from_numpy(labels).view(-1)
 ###############################################################################
 # Create graph and model
 # ~~~~~~~~~~~~~~~~~~~~~~~
 # configurations
 n_hidden = 16 # number of hidden units
 n_bases = -1 # use number of relations as number of bases
 n_hidden_layers = 0 # use 1 input layer, 1 output layer, no hidden layer
 n_epochs = 25 # epochs to train
 lr = 0.01 # learning rate
 l2norm = 0 # L2 norm coefficient
 # create graph
 g = DGLGraph()
 g.add_nodes(num_nodes)
 g.add_edges(data.edge_src, data.edge_dst)
 g.edata.update({'rel_type': edge_type, 'norm': edge_norm})
 # create model
 model = Model(len(g),
              n_hidden,
              num_classes,
              num_rels,
              num_bases=n_bases,
              num_hidden_layers=n_hidden_layers)
 ###############################################################################
 # Training loop
 # ~~~~~~~~~~~~~~~~
 # optimizer
 optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=l2norm)
 print("start training...")
 model.train()
 for epoch in range(n_epochs):
    optimizer.zero_grad()
    logits = model.forward(g)
    loss = F.cross_entropy(logits[train_idx], labels[train_idx])
    loss.backward()
    optimizer.step()
    train_acc = torch.sum(logits[train_idx].argmax(dim=1) == labels[train_idx])
    train_acc = train_acc.item() / len(train_idx)
    val_loss = F.cross_entropy(logits[val_idx], labels[val_idx])
    val_acc = torch.sum(logits[val_idx].argmax(dim=1) == labels[val_idx])
    val_acc = val_acc.item() / len(val_idx)
    print("Epoch {:05d} | ".format(epoch) +
          "Train Accuracy: {:.4f} | Train Loss: {:.4f} | ".format(
              train_acc, loss.item()) +
          "Validation Accuracy: {:.4f} | Validation loss: {:.4f}".format(
              val_acc, val_loss.item()))
 ###############################################################################
 # .. _link-prediction:
 #
 # The second task: Link prediction
 # --------------------------------
 # So far, we have seen how to use DGL to implement entity classification with
 # R-GCN model. In the knowledge base setting, representation generated by
 # R-GCN can be further used to uncover potential relations between nodes. In
 # R-GCN paper, authors feed the entity representations generated by R-GCN
 # into the `DistMult <https://arxiv.org/pdf/1412.6575.pdf>`_ prediction model
 # to predict possible relations.
 #
 # The implementation is similar to the above but with an extra DistMult layer
 # stacked on top of the R-GCN layers. You may find the complete
 # implementation of link prediction with R-GCN in our `example
 # code <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/rgcn/link_predict.py>`_.
--- a/tutorials/models/6_line_graph.py
+++ b/tutorials/models/6_line_graph.py
 """
 .. _model-line-graph:
 Line Graph Neural Network
 =========================
 **Author**: `Qi Huang <https://github.com/HQ01>`_, Yu Gai,
 `Minjie Wang <https://jermainewang.github.io/>`_, Zheng Zhang
 """
 ###########################################################################################
 # 
 # In :doc:`GCN <1_gcn>` , we demonstrate how to classify nodes on an input
 # graph in a semi-supervised setting, using graph convolutional neural network
 # as embedding mechanism for graph features.
 # In this tutorial, we shift our focus to community detection problem. The
 # task of community detection, i.e. graph clustering, consists of partitioning
 # the vertices in a graph into clusters in which nodes are more "similar" to
 # one another.
 #
 # To generalize GNN to supervised community detection, Chen et al. introduced
 # a line-graph based variation of graph neural network in 
 # `Supervised Community Detection with Line Graph Neural Networks <https://arxiv.org/abs/1705.08415>`__. 
 # One of the highlight of their model is
 # to augment the vanilla graph neural network(GNN) architecture to operate on
 # the line graph of edge adajcencies, defined with non-backtracking operator.
 #
 # In addition to its high performance, LGNN offers an opportunity to
 # illustrate how DGL can implement an advanced graph algorithm by flexibly
 # mixing vanilla tensor operations, sparse-matrix multiplication and message-
 # passing APIs.
 #
 # In the following sections, we will go through community detection, line
 # graph, LGNN, and its implementation.
 #
 # Supervised Community Detection Task on CORA
 # --------------------------------------------
 # Community Detection
 # ~~~~~~~~~~~~~~~~~~~~
 # In community detection task, we cluster "similar" nodes instead of
 # "labeling" them. The node similarity is typically described as higher inner
 # density in each cluster.
 #
 # What's the difference between community detection and node classification？
 # Comparing to node classification, community detection focuses on retrieving
 # cluster information in the graph, rather than assigning a specific label to
 # a node. For example, as long as a node is clusetered with its community
 # members, it doesn't matter whether the node is assigned as "community A",
 # or "community B", while assigning all "great movies" to label "bad movies"
 # will be a disaster in a movie network classification task.
 #
 # What's the difference then, between a community detection algorithm and
 # other clustering algorithm such as k-means? Community detection algorithm operates on
 # graph-structured data. Comparing to k-means, community detection leverages
 # graph structure, instead of simply clustering nodes based on their
 # features.
 #
 # CORA
 # ~~~~~
 # To be consistent with Graph Convolutional Network tutorial, 
 # we use `CORA <https://linqs.soe.ucsc.edu/data>`__ 
 # to illustrate a simple community detection task. To refresh our memory, 
 # CORA is a scientific publication dataset, with 2708 papers belonging to 7 
 # different mahcine learning sub-fields. Here, we formulate CORA as a 
 # directed graph, with each node being a paper, and each edge being a 
 # citation link (A->B means A cites B). Here is a visualization of the whole 
 # CORA dataset.
 #
 # .. figure:: https://i.imgur.com/X404Byc.png
 #    :alt: cora
 #    :height: 400px
 #    :width: 500px
 #    :align: center
 #
 # CORA naturally contains 7 "classes", and statistics below show that each
 # "class" does satisfy our assumption of community, i.e. nodes of same class
 # class have higher connection probability among them than with nodes of different class.
 # The following code snippet verifies that there are more intra-class edges
 # than inter-class:
 import torch
 import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
 import dgl
 from dgl.data import citation_graph as citegrh
 data = citegrh.load_cora()
 G = dgl.DGLGraph(data.graph)
 labels = th.tensor(data.labels)
 # find all the nodes labeled with class 0
 label0_nodes = th.nonzero(labels == 0).squeeze()
 # find all the edges pointing to class 0 nodes
 src, _ = G.in_edges(label0_nodes)
 src_labels = labels[src]
 # find all the edges whose both endpoints are in class 0
 intra_src = th.nonzero(src_labels == 0)
 print('Intra-class edges percent: %.4f' % (len(intra_src) / len(src_labels)))
 ###########################################################################################
 # Binary Community Subgraph from CORA -- a Toy Dataset
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # Without loss of generality, in this tutorial we limit the scope of our
 # task to binary community detection.
 # 
 # .. note::
 #
 #    To create a toy binary-community dataset from CORA, We first extract
 #    all two-class pairs from the original CORA 7 classes. For each pair, we
 #    treat each class as one community, and find the largest subgraph that
 #    at least contain one cross-community edge as the training example. As
 #    a result, there are a total of 21 training samples in this mini-dataset.
 #
 # Here we visualize one of the training samples and its community structure:
 import networkx as nx
 import matplotlib.pyplot as plt
 train_set = dgl.data.CoraBinary()
 G1, pmpd1, label1 = train_set[1]
 nx_G1 = G1.to_networkx()
 def visualize(labels, g):
    pos = nx.spring_layout(g, seed=1)
    plt.figure(figsize=(8, 8))
    plt.axis('off')
    nx.draw_networkx(g, pos=pos, node_size=50, cmap=plt.get_cmap('coolwarm'),
                     node_color=labels, edge_color='k',
                     arrows=False, width=0.5, style='dotted', with_labels=False)
 visualize(label1, nx_G1)
 ###########################################################################################
 # Interested readers can go to the original paper to see how to generalize
 # to multi communities case.
 #
 # Community Detection in a Supervised Setting
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # Community Detection problem could be tackled with both supervised and
 # unsupervised approaches. Same as the original paper, we formulate
 # Community Detection in a supervised setting as follows:
 #
 # - Each training example consists of :math:`(G, L)`, where :math:`G` is a
 #   directed graph :math:`(V, E)`. For each node :math:`v` in :math:`V`, we
 #   assign a ground truth community label :math:`z_v \in \{0,1\}`.
 # - The parameterized model :math:`f(G, \theta)` predicts a label set
 #   :math:`\tilde{Z} = f(G)` for nodes :math:`V`.
 # - For each example :math:`(G,L)`, the model learns to minimize a specially
 #   designed loss function (equivariant loss) :math:`L_{equivariant} =
 #   (\tilde{Z}，Z)`
 #
 # .. note::
 #
 #    In this supervised setting, the model naturally predicts a "label" for
 #    each community. However, community assignment should be equivariant to
 #    label permutations. To acheive this, in each forward process, we take
 #    the minimum among losses calcuated from all possible permutations of
 #    labels.
 #
 #    Mathematically, this means
 #    :math:`L_{equivariant} = \underset{\pi \in S_c} {min}-\log(\hat{\pi}, \pi)`,
 #    where :math:`S_c` is the set of all permutations of labels, and
 #    :math:`\hat{\pi}` is the set of predicted labels,
 #    :math:`- \log(\hat{\pi},\pi)` denotes negative log likelihood.
 #
 #    For instance, for a toy graph with node :math:`\{1,2,3,4\}` and
 #    community assignment :math:`\{A, A, A, B\}`, with each node's label
 #    :math:`l \in \{0,1\}`,The group of all possible permutations
 #    :math:`S_c = \{\{0,0,0,1\}, \{1,1,1,0\}\}`.
 # 
 # Line Graph Neural network: key ideas
 # ------------------------------------
 # An key innovation in this paper is the use of line-graph.
 # Unlike models in previous tutorials, message passing happens not only on the
 # original graph, e.g. the binary community subgraph from CORA, but also on the
 # line-graph associated with the original graph.
 #
 # What's a line-graph ?
 # ~~~~~~~~~~~~~~~~~~~~~
 # In graph theory, line graph is a graph representation that encodes the
 # edge adjacency sturcutre in the original graph.
 #
 # Specifically, a line-graph :math:`L(G)` turns an edge of the original graph `G`
 # into a node. This is illustrated with the graph below (taken from the
 # paper)
 # 
 # .. figure:: https://i.imgur.com/4WO5jEm.png
 #    :alt: lg
 #    :align: center
 #
 # Here, :math:`e_{A}:= （i\rightarrow j）` and :math:`e_{B}:= (j\rightarrow k)`
 # are two edges in the original graph :math:`G`. In line graph :math:`G_L`,
 # they correspond to nodes :math:`v^{l}_{A}, v^{l}_{B}`.
 #
 # The next natural question is, how to connect nodes in line-graph？ How to
 # connect two "edges"? Here, we use the following connection rule:
 #
 # Two nodes :math:`v^{l}_{A}`, :math:`v^{l}_{B}` in `lg` are connected if
 # the corresponding two edges :math:`e_{A}, e_{B}` in `g` share one and only 
 # one node:
 # :math:`e_{A}`'s destination node is :math:`e_{B}`'s source node
 # (:math:`j`).
 # 
 # .. note::
 #
 #    Mathematically, this definition corresponds to a notion called non-backtracking
 #    operator:
 #    :math:`B_{(i \rightarrow j), (\hat{i} \rightarrow \hat{j})}`
 #    :math:`= \begin{cases}
 #    1 \text{ if } j = \hat{i}, \hat{j} \neq i\\
 #    0 \text{ otherwise} \end{cases}`
 #    where an edge is formed if :math:`B_{node1, node2} = 1`.
 #
 #
 # One layer in LGNN -- algorithm sturcture
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 #
 # LGNN chains up a series of line-graph neural network layers. The graph
 # reprentation :math:`x` and its line-graph companion :math:`y` evolve with
 # the dataflow as follows,
 # 
 # .. figure:: https://i.imgur.com/bZGGIGp.png
 #    :alt: alg
 #    :align: center
 #
 # At the :math:`k`-th layer, the :math:`i`-th neuron of the :math:`l`-th
 # channel updates its embedding :math:`x^{(k+1)}_{i,l}` with:
 #
 # .. math::
 #    \begin{split}
 #    x^{(k+1)}_{i,l} ={}&\rho[x^{(k)}_{i}\theta^{(k)}_{1,l}
 #    +(Dx^{(k)})_{i}\theta^{(k)}_{2,l} \\
 #    &+\sum^{J-1}_{j=0}(A^{2^{j}}x^{k})_{i}\theta^{(k)}_{3+j,l}\\
 #    &+[\{\text{Pm},\text{Pd}\}y^{(k)}]_{i}\theta^{(k)}_{3+J,l}] \\
 #    &+\text{skip-connection}
 #    \qquad i \in V, l = 1,2,3, ... b_{k+1}/2
 #    \end{split}
 #
 # Then, the line-graph representation :math:`y^{(k+1)}_{i,l}` with,
 #
 # .. math::
 #
 #    \begin{split}
 #    y^{(k+1)}_{i',l^{'}} = {}&\rho[y^{(k)}_{i^{'}}\gamma^{(k)}_{1,l^{'}}+
 #    (D_{L(G)}y^{(k)})_{i^{'}}\gamma^{(k)}_{2,l^{'}}\\
 #    &+\sum^{J-1}_{j=0}(A_{L(G)}^{2^{j}}y^{k})_{i}\gamma^{(k)}_{3+j,l^{'}}\\
 #    &+[\{\text{Pm},\text{Pd}\}^{T}x^{(k+1)}]_{i^{'}}\gamma^{(k)}_{3+J,l^{'}}]\\
 #    &+\text{skip-connection}
 #    \qquad i^{'} \in V_{l}, l^{'} = 1,2,3, ... b^{'}_{k+1}/2
 #    \end{split}
 #
 # Where :math:`\text{skip-connection}` refers to performing the same operation without the non-linearity
 # :math:`\rho`, and with linear projection :math:`\theta_\{\frac{b_{k+1}}{2} + 1, ..., b_{k+1}-1, b_{k+1}\}`
 # and :math:`\gamma_\{\frac{b_{k+1}}{2} + 1, ..., b_{k+1}-1, b_{k+1}\}`.
 #
 # Implement LGNN in DGL
 # ---------------------
 # General idea
 # ~~~~~~~~~~~~
 # The above equations look intimidating. However, we observe the following:
 # 
 # - The two equations are symmetric and can be implemented as two instances
 #   of the same class with different parameters.
 #   Mainly, the first equation operates on graph representation :math:`x`,
 #   whereas the second operates on line-graph
 #   representation :math:`y`. Let us denote this abstraction as :math:`f`. Then
 #   the first is :math:`f(x,y; \theta_x)`, and the second
 #   is :math:`f(y,x, \theta_y)`. That is, they are parameterized to compute
 #   representations of the original graph and its
 #   companion line graph, respectively.
 #
 # - Each equation consists of 4 terms (take the first one as an example):
 #
 #   - :math:`x^{(k)}\theta^{(k)}_{1,l}`, a linear projection of previous
 #     layer's output :math:`x^{(k)}`, denote as :math:`\text{prev}(x)`.
 #   - :math:`(Dx^{(k)})\theta^{(k)}_{2,l}`, a linear projection of degree
 #     operator on :math:`x^{(k)}`, denote as :math:`\text{deg}(x)`.
 #   - :math:`\sum^{J-1}_{j=0}(A^{2^{j}}x^{(k)})\theta^{(k)}_{3+j,l}`,
 #     a summation of :math:`2^{j}` adjacency operator on :math:`x^{(k)}`,
 #     denote as :math:`\text{radius}(x)`
 #   - :math:`[\{Pm,Pd\}y^{(k)}]\theta^{(k)}_{3+J,l}`, fusing another
 #     graph's embedding information using incidence matrix
 #     :math:`\{Pm, Pd\}`, followed with a linear porjection,
 #     denote as :math:`\text{fuse}(y)`.
 #
 # - In addition, each of the terms are performed again with different
 #   parameters, and without the nonlinearity after the sum.
 #   Therefore, :math:`f` could be written as:
 # 
 #   .. math::
 #      \begin{split}
 #      f(x^{(k)},y^{(k)}) = {}\rho[&\text{prev}(x^{(k-1)}) + \text{deg}(x^{(k-1)}) +\text{radius}(x^{k-1})
 #      +\text{fuse}(y^{(k)})]\\
 #      +&\text{prev}(x^{(k-1)}) + \text{deg}(x^{(k-1)}) +\text{radius}(x^{k-1}) +\text{fuse}(y^{(k)})
 #      \end{split}
 #
 # - Two equations are chained up in the following order :
 # 
 #   .. math::
 #      \begin{split}
 #      x^{(k+1)} = {}& f(x^{(k)}, y^{(k)})\\
 #      y^{(k+1)} = {}& f(y^{(k)}, x^{(k+1)})
 #      \end{split}
 # 
 # With these observations, we proceed to implementation.
 # The important point is we are to use different strategies for these terms.
 # 
 # .. note::
 #    For a detailed explanation of :math:`\{Pm, Pd\}`, please go to `Advanced Topic`_.
 #
 # Implementing :math:`\text{prev}` and :math:`\text{deg}` as tensor operation
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # Since linear projection and degree operation are both simply matrix
 # multiplication, we can write them as PyTorch tensor operation.
 #
 # In ``__init__``, we define the projection variables:
 # 
 # ::
 # 
 #    self.linear_prev = nn.Linear(in_feats, out_feats)
 #    self.linear_deg = nn.Linear(in_feats, out_feats)
 # 
 #
 # In ``forward()``, :math:`\text{prev}` and :math:`\text{deg}` are the same
 # as any other PyTorch tensor operations.
 # 
 # ::
 # 
 #    prev_proj = self.linear_prev(feat_a)
 #    deg_proj = self.linear_deg(deg * feat_a)
 # 
 # Implementing :math:`\text{radius}` as message passing in DGL
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # As discussed in GCN tutorial, we can formulate one adjacency operator as
 # doing one step message passing. As a generalization, :math:`2^j` adjacency
 # operations can be formulated as performing :math:`2^j` step of message
 # passing. Therefore, the summation is equivalent to summing nodes'
 # representation of :math:`2^j, j=0, 1, 2..` step messsage passing, i.e.
 # gathering information in :math:`2^{j}` neighbourhood of each node.
 #
 # In ``__init__``, we define the projection variables used in each
 # :math:`2^j` steps of message passing:
 # 
 # ::
 # 
 #   self.linear_radius = nn.ModuleList(
 #           [nn.Linear(in_feats, out_feats) for i in range(radius)])
 #
 # In ``__forward__``, we use following function ``aggregate_radius()`` to
 # gather data from multiple hop. Note that the ``update_all`` is called
 # multiple times.
 # Return a list containing features gathered from multiple radius.
 import dgl.function as fn
 def aggregate_radius(radius, g, z):
    # initializing list to collect message passing result
    z_list = []
    g.ndata['z'] = z
    # pulling message from 1-hop neighbourhood
    g.update_all(fn.copy_src(src='z', out='m'), fn.sum(msg='m', out='z'))
    z_list.append(g.ndata['z'])
    for i in range(radius - 1):
        for j in range(2 ** i):
            #pulling message from 2^j neighborhood
            g.update_all(fn.copy_src(src='z', out='m'), fn.sum(msg='m', out='z'))
        z_list.append(g.ndata['z'])
    return z_list
 #########################################################################
 # Implementing :math:`\text{fuse}` as sparse matrix multiplication
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # :math:`\{Pm, Pd\}` is a sparse matrix with only two non-zero entries on
 # each column. Therefore, we construct it as a sparse matrix in the dataset,
 # and implement :math:`\text{fuse}` as a sparse matrix multiplication.
 #
 # in ``__forward__``:
 # 
 # ::
 # 
 #   fuse = self.linear_fuse(th.mm(pm_pd, feat_b))
 #
 # Completing :math:`f(x, y)`
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~
 # Finally, we sum up all the terms together, pass it to skip connection and
 # batch-norm.
 # 
 # ::
 #
 #   result = prev_proj + deg_proj + radius_proj + fuse
 # 
 # Then pass result to skip connection: 
 # 
 # ::
 # 
 #   result = th.cat([result[:, :n], F.relu(result[:, n:])], 1)
 # 
 # Then batch norm
 # 
 # ::
 # 
 #   result = self.bn(result) #Batch Normalization.
 # 
 #
 # Below is the complete code for one LGNN layer's abstraction :math:`f(x,y)`
 class LGNNCore(nn.Module):
    def __init__(self, in_feats, out_feats, radius):
        super(LGNNCore, self).__init__()
        self.out_feats = out_feats
        self.radius = radius
        self.linear_prev = nn.Linear(in_feats, out_feats)
        self.linear_deg = nn.Linear(in_feats, out_feats)
        self.linear_radius = nn.ModuleList(
                [nn.Linear(in_feats, out_feats) for i in range(radius)])
        self.linear_fuse = nn.Linear(in_feats, out_feats)
        self.bn = nn.BatchNorm1d(out_feats)
    def forward(self, g, feat_a, feat_b, deg, pm_pd):
        # term "prev"
        prev_proj = self.linear_prev(feat_a)
        # term "deg"
        deg_proj = self.linear_deg(deg * feat_a)
        # term "radius"
        # aggregate 2^j-hop features
        hop2j_list = aggregate_radius(self.radius, g, feat_a)
        # apply linear transformation
        hop2j_list = [linear(x) for linear, x in zip(self.linear_radius, hop2j_list)]
        radius_proj = sum(hop2j_list)
        # term "fuse"
        fuse = self.linear_fuse(th.mm(pm_pd, feat_b))
        # sum them together
        result = prev_proj + deg_proj + radius_proj + fuse
        # skip connection and batch norm
        n = self.out_feats // 2
        result = th.cat([result[:, :n], F.relu(result[:, n:])], 1)
        result = self.bn(result)
        return result
 ##############################################################################################################
 # Chain up LGNN abstractions as a LGNN layer
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # To implement:
 # 
 # .. math::
 #    \begin{split}
 #    x^{(k+1)} = {}& f(x^{(k)}, y^{(k)})\\
 #    y^{(k+1)} = {}& f(y^{(k)}, x^{(k+1)})
 #    \end{split}
 #
 # We chain up two ``LGNNCore`` instances with different parameter in the forward pass.
 class LGNNLayer(nn.Module):
    def __init__(self, in_feats, out_feats, radius):
        super(LGNNLayer, self).__init__()
        self.g_layer = LGNNCore(in_feats, out_feats, radius)
        self.lg_layer = LGNNCore(in_feats, out_feats, radius)
    def forward(self, g, lg, x, lg_x, deg_g, deg_lg, pm_pd):
        next_x = self.g_layer(g, x, lg_x, deg_g, pm_pd)
        pm_pd_y = th.transpose(pm_pd, 0, 1)
        next_lg_x = self.lg_layer(lg, lg_x, x, deg_lg, pm_pd_y)
        return next_x, next_lg_x
 ########################################################################################
 # Chain up LGNN layers
 # ~~~~~~~~~~~~~~~~~~~~
 # We then define an LGNN with three hidden layers.
 class LGNN(nn.Module):
    def __init__(self, radius):
        super(LGNN, self).__init__()
        self.layer1 = LGNNLayer(1, 16, radius)  # input is scalar feature
        self.layer2 = LGNNLayer(16, 16, radius)  # hidden size is 16
        self.layer3 = LGNNLayer(16, 16, radius)
        self.linear = nn.Linear(16, 2)  # predice two classes
    def forward(self, g, lg, pm_pd):
        # compute the degrees
        deg_g = g.in_degrees().float().unsqueeze(1)
        deg_lg = lg.in_degrees().float().unsqueeze(1)
        # use degree as the input feature
        x, lg_x = deg_g, deg_lg
        x, lg_x = self.layer1(g, lg, x, lg_x, deg_g, deg_lg, pm_pd)
        x, lg_x = self.layer2(g, lg, x, lg_x, deg_g, deg_lg, pm_pd)
        x, lg_x = self.layer3(g, lg, x, lg_x, deg_g, deg_lg, pm_pd)
        return self.linear(x)
 #########################################################################################
 # Training and Inference
 # -----------------------
 # We first load the data
 from torch.utils.data import DataLoader
 training_loader = DataLoader(train_set,
                             batch_size=1,
                             collate_fn=train_set.collate_fn,
                             drop_last=True)
 #######################################################################################
 # We then define the main training loop. Note that each training sample contains
 # three objects: a :class:`~dgl.DGLGraph`, a scipy sparse matrix ``pmpd`` and label
 # array in ``numpy.ndarray``. We first generate the line graph using:
 #
 # ::
 # 
 #   lg = g.line_graph(backtracking=False)
 #
 # Note that ``backtracking=False`` is required to correctly simulate non-backtracking
 # operation. We also define a utility function to convert the scipy sparse matrix to
 # torch sparse tensor.
 # create the model
 model = LGNN(radius=3)
 # define the optimizer
 optimizer = th.optim.Adam(model.parameters(), lr=1e-2)
 # a util function to convert a scipy.coo_matrix to torch.SparseFloat
 def sparse2th(mat):
    value = mat.data
    indices = th.LongTensor([mat.row, mat.col])
    tensor = th.sparse.FloatTensor(indices, th.from_numpy(value).float(), mat.shape)
    return tensor
 # train for 20 epochs
 for i in range(20):
    all_loss = []
    all_acc = []
    for [g, pmpd, label] in training_loader:
        # Generate the line graph.
        lg = g.line_graph(backtracking=False)
        # Create torch tensors
        pmpd = sparse2th(pmpd)
        label = th.from_numpy(label)
        # Forward
        z = model(g, lg, pmpd)
        # Calculate loss:
        # Since there are only two communities, there are only two permutations
        #  of the community labels.
        loss_perm1 = F.cross_entropy(z, label)
        loss_perm2 = F.cross_entropy(z, 1 - label)
        loss = th.min(loss_perm1, loss_perm2)
        # Calculate accuracy:
        _, pred = th.max(z, 1)
        acc_perm1 = (pred == label).float().mean()
        acc_perm2 = (pred == 1 - label).float().mean()
        acc = th.max(acc_perm1, acc_perm2)
        all_loss.append(loss.item())
        all_acc.append(acc.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    niters = len(all_loss)
    print("Epoch %d | loss %.4f | accuracy %.4f" % (i,
        sum(all_loss) / niters, sum(all_acc) / niters))
 #######################################################################################
 # Visualize training progress
 # -----------------------------
 # We visualize the network's community prediction on one training example,
 # together with the ground truth.
 pmpd1 = sparse2th(pmpd1)
 LG1 = G1.line_graph(backtracking=False)
 z = model(G1, LG1, pmpd1)
 _, pred = th.max(z, 1)
 visualize(pred, nx_G1)
 #######################################################################################
 # Compared with the ground truth. Note that the color might be reversed for the
 # two community as the model is to correctly predict the "partitioning".
 visualize(label1, nx_G1)
 #########################################
 # Here is an animation to better understand the process. (40 epochs)
 #
 # .. figure:: https://i.imgur.com/KDUyE1S.gif 
 #    :alt: lgnn-anim
 #
 # Advanced topic
 # --------------
 #
 # Batching
 # ~~~~~~~~
 # LGNN takes a collection of different graphs.
 # Thus, it's natural we use batching to explore parallelism.
 # Why is it not done?
 #
 # As it turned out, we moved batching into the dataloader itself.
 # In the ``collate_fn`` for PyTorch Dataloader, we batch graphs using DGL's
 # batched_graph API. To refresh our memory, DGL batches graphs by merging them
 # into a large graph, with each smaller graph's adjacency matrix being a block
 # along the diagonal of the large graph's adjacency matrix.  We concatentate
 # :math`\{Pm,Pd\}` as block diagonal matrix in corespondance to DGL batched
 # graph API.
 def collate_fn(batch):
    graphs, pmpds, labels = zip(*batch)
    batched_graphs = dgl.batch(graphs)
    batched_pmpds = sp.block_diag(pmpds)
    batched_labels = np.concatenate(labels, axis=0)
    return batched_graphs, batched_pmpds, batched_labels
 ######################################################################################
 # You can check out the complete code
 # `here <https://github.com/jermainewang/dgl/tree/master/examples/pytorch/line_graph>`_.
 # 
 # What's the business with :math:`\{Pm, Pd\}`?
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 # Rougly speaking, there is a relationship between how :math:`g` and
 # :math:`lg` (the line graph) working together with loopy brief propagation.
 # Here, we implement :math:`\{Pm, Pd\}` as scipy coo sparse matrix in the datset,
 # and stack them as tensors when batching. Another batching solution is to
 # treat :math:`\{Pm, Pd\}` as the adjacency matrix of a bipartie graph, which maps
 # line graph's feature to graph's, and vice versa.
--- a/tutorials/models/1_gnn/README.txt
+++ b/tutorials/models/1_gnn/README.txt
+.. _tutorials1-index:
+Graph Neural Network and its variant
+------------------------------------
+* **GCN** `[paper] <https://arxiv.org/abs/1609.02907>`__ `[tutorial] <models/1_gcn.html>`__
+  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/gcn/gcn.py>`__:
+  this is the vanilla GCN. The tutorial covers the basic uses of DGL APIs.
+* **GAT** `[paper] <https://arxiv.org/abs/1710.10903>`__
+  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/gat/gat.py>`__:
+  the key extension of GAT w.r.t vanilla GCN is deploying multi-head attention
+  among neighborhood of a node, thus greatly enhances the capacity and
+  expressiveness of the model.
+* **R-GCN** `[paper] <https://arxiv.org/abs/1703.06103>`__ `[tutorial] <models/4_rgcn.html>`__
+  [code (wip)]: the key
+  difference of RGNN is to allow multi-edges among two entities of a graph, and
+  edges with distinct relationships are encoded differently. This is an
+  interesting extension of GCN that can have a lot of applications of its own.
+* **LGNN** `[paper] <https://arxiv.org/abs/1705.08415>`__ `[tutorial (wip)]` `[code (wip)]`:
+  this model focuses on community detection by inspecting graph structures. It
+  uses representations of both the orignal graph and its line-graph companion. In
+  addition to demonstrate how an algorithm can harness multiple graphs, our
+  implementation shows how one can judiciously mix vanilla tensor operation,
+  sparse-matrix tensor operations, along with message-passing with DGL.
+* **SSE** `[paper] <http://proceedings.mlr.press/v80/dai18a/dai18a.pdf>`__ `[tutorial (wip)]`
+  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/mxnet/sse/sse_batch.py>`__:
+  the emphasize here is *giant* graph that cannot fit comfortably on one GPU
+  card. SSE is an example to illustrate the co-design of both algrithm and
+  system: sampling to guarantee asymptotic covergence while lowering the
+  complexity, and batching across samples for maximum parallelism.
\ No newline at end of file
--- a/tutorials/models/3_tree-lstm.py
+++ b/tutorials/models/3_tree-lstm.py
 """
 .. _model-tree-lstm:
 Tree LSTM DGL Tutorial
 =========================
 **Author**: Zihao Ye, Qipeng Guo, `Minjie Wang
 <https://jermainewang.github.io/>`_, `Jake Zhao
 <https://cs.nyu.edu/~jakezhao/>`_, Zheng Zhang
 """
 ##############################################################################
 #
 # Tree-LSTM structure was first introduced by Kai et. al in an ACL 2015 
 # paper: `Improved Semantic Representations From Tree-Structured Long
 # Short-Term Memory Networks <https://arxiv.org/pdf/1503.00075.pdf>`__.
 # The core idea is to introduce syntactic information for language tasks by 
 # extending the chain-structured LSTM to a tree-structured LSTM. The Dependency 
 # Tree/Constituency Tree techniques were leveraged to obtain a ''latent tree''.
 #
 # One, if not all, difficulty of training Tree-LSTMs is batching --- a standard 
 # technique in machine learning to accelerate optimization. However, since trees 
 # generally have different shapes by nature, parallization becomes non trivial. 
 # DGL offers an alternative: to pool all the trees into one single graph then 
 # induce the message passing over them guided by the structure of each tree.
 #
 # The task and the dataset
 # ------------------------
 # In this tutorial, we will use Tree-LSTMs for sentiment analysis.
 # We have wrapped the
 # `Stanford Sentiment Treebank <https://nlp.stanford.edu/sentiment/>`__ in
 # ``dgl.data``. The dataset provides a fine-grained tree level sentiment
 # annotation: 5 classes(very negative, negative, neutral, positive, and
 # very positive) that indicates the sentiment in current subtree. Non-leaf
 # nodes in constituency tree does not contain words, we use a special
 # ``PAD_WORD`` token to denote them, during the training/inferencing,
 # their embeddings would be masked to all-zero.
 #
 # .. figure:: https://i.loli.net/2018/11/08/5be3d4bfe031b.png
 #    :alt: 
 #
 # The figure displays one sample of the SST dataset, which is a
 # constituency parse tree with their nodes labeled with sentiment. To
 # speed up things, let's build a tiny set with 5 sentences and take a look
 # at the first one:
 #
 import dgl
 import dgl.data as data
 # Each sample in the dataset is a constituency tree. The leaf nodes
 # represent words. The word is a int value stored in the "x" field.
 # The non-leaf nodes has a special word PAD_WORD. The sentiment
 # label is stored in the "y" feature field.
 trainset = data.SST(mode='tiny')  # the "tiny" set has only 5 trees
 tiny_sst = trainset.trees
 num_vocabs = trainset.num_vocabs
 num_classes = trainset.num_classes
 vocab = trainset.vocab # vocabulary dict: key -> id
 inv_vocab = {v: k for k, v in vocab.items()} # inverted vocabulary dict: id -> word
 a_tree = tiny_sst[0]
 for token in a_tree.ndata['x'].tolist():
    if token != trainset.PAD_WORD:
        print(inv_vocab[token], end=" ")
 ##############################################################################
 # Step 1: batching
 # ----------------
 #
 # The first step is to throw all the trees into one graph, using
 # the :func:`~dgl.batched_graph.batch` API.
 #
 import networkx as nx
 import matplotlib.pyplot as plt
 graph = dgl.batch(tiny_sst)
 def plot_tree(g):
    # this plot requires pygraphviz package
    pos = nx.nx_agraph.graphviz_layout(g, prog='dot')
    nx.draw(g, pos, with_labels=False, node_size=10,
            node_color=[[.5, .5, .5]], arrowsize=4)
    plt.show()
 plot_tree(graph.to_networkx())
 ##############################################################################
 # You can read more about the definition of :func:`~dgl.batched_graph.batch`
 # (by clicking the API), or can skip ahead to the next step:
 # 
 # .. note::
 #
 #    **Definition**: a :class:`~dgl.batched_graph.BatchedDGLGraph` is a
 #    :class:`~dgl.DGLGraph` that unions a list of :class:`~dgl.DGLGraph`\ s. 
 #    
 #    - The union includes all the nodes,
 #      edges, and their features. The order of nodes, edges and features are
 #      preserved. 
 #     
 #        - Given that we have :math:`V_i` nodes for graph
 #          :math:`\mathcal{G}_i`, the node ID :math:`j` in graph
 #          :math:`\mathcal{G}_i` correspond to node ID
 #          :math:`j + \sum_{k=1}^{i-1} V_k` in the batched graph. 
 #    
 #        - Therefore, performing feature transformation and message passing on
 #          ``BatchedDGLGraph`` is equivalent to doing those
 #          on all ``DGLGraph`` constituents in parallel. 
 #
 #    - Duplicate references to the same graph are
 #      treated as deep copies; the nodes, edges, and features are duplicated,
 #      and mutation on one reference does not affect the other. 
 #    - Currently, ``BatchedDGLGraph`` is immutable in
 #      graph structure (i.e. one can't add
 #      nodes and edges to it). We need to support mutable batched graphs in
 #      (far) future. 
 #    - The ``BatchedDGLGraph`` keeps track of the meta
 #      information of the constituents so it can be
 #      :func:`~dgl.batched_graph.unbatch`\ ed to list of ``DGLGraph``\ s.
 #
 # For more details about the :class:`~dgl.batched_graph.BatchedDGLGraph`
 # module in DGL, you can click the class name.
 #
 # Step 2: Tree-LSTM Cell with message-passing APIs
 # ------------------------------------------------
 #
 # The authors proposed two types of Tree LSTM: Child-Sum
 # Tree-LSTMs, and :math:`N`-ary Tree-LSTMs. In this tutorial we focus 
 # on applying *Binary* Tree-LSTM to binarized constituency trees(this 
 # application is also known as *Constituency Tree-LSTM*). We use PyTorch 
 # as our backend framework to set up the network.
 #
 # In `N`-ary Tree LSTM, each unit at node :math:`j` maintains a hidden
 # representation :math:`h_j` and a memory cell :math:`c_j`. The unit
 # :math:`j` takes the input vector :math:`x_j` and the hidden
 # representations of the their child units: :math:`h_{jl}, 1\leq l\leq N` as
 # input, then update its new hidden representation :math:`h_j` and memory
 # cell :math:`c_j` by: 
 #
 # .. math::
 #
 #    i_j & = & \sigma\left(W^{(i)}x_j + \sum_{l=1}^{N}U^{(i)}_l h_{jl} + b^{(i)}\right),  & (1)\\
 #    f_{jk} & = & \sigma\left(W^{(f)}x_j + \sum_{l=1}^{N}U_{kl}^{(f)} h_{jl} + b^{(f)} \right), &  (2)\\
 #    o_j & = & \sigma\left(W^{(o)}x_j + \sum_{l=1}^{N}U_{l}^{(o)} h_{jl} + b^{(o)} \right), & (3)  \\
 #    u_j & = & \textrm{tanh}\left(W^{(u)}x_j + \sum_{l=1}^{N} U_l^{(u)}h_{jl} + b^{(u)} \right), & (4)\\
 #    c_j & = & i_j \odot u_j + \sum_{l=1}^{N} f_{jl} \odot c_{jl}, &(5) \\
 #    h_j & = & o_j \cdot \textrm{tanh}(c_j), &(6)  \\
 #
 # It can be decomposed into three phases: ``message_func``,
 # ``reduce_func`` and ``apply_node_func``.
 #
 # .. note::
 #    ``apply_node_func`` is a new node UDF we have not introduced before. In
 #    ``apply_node_func``, user specifies what to do with node features,
 #    without considering edge features and messages. In Tree-LSTM case,
 #    ``apply_node_func`` is a must, since there exists (leaf) nodes with
 #    :math:`0` incoming edges, which would not be updated via
 #    ``reduce_func``.
 #
 import torch as th
 import torch.nn as nn
 class TreeLSTMCell(nn.Module):
    def __init__(self, x_size, h_size):
        super(TreeLSTMCell, self).__init__()
        self.W_iou = nn.Linear(x_size, 3 * h_size, bias=False)
        self.U_iou = nn.Linear(2 * h_size, 3 * h_size, bias=False)
        self.b_iou = nn.Parameter(th.zeros(1, 3 * h_size))
        self.U_f = nn.Linear(2 * h_size, 2 * h_size)
    def message_func(self, edges):
        return {'h': edges.src['h'], 'c': edges.src['c']}
    def reduce_func(self, nodes):
        # concatenate h_jl for equation (1), (2), (3), (4)
        h_cat = nodes.mailbox['h'].view(nodes.mailbox['h'].size(0), -1)
        # equation (2)
        f = th.sigmoid(self.U_f(h_cat)).view(*nodes.mailbox['h'].size())
        # second term of equation (5)
        c = th.sum(f * nodes.mailbox['c'], 1)
        return {'iou': self.U_iou(h_cat), 'c': c}
    def apply_node_func(self, nodes):
        # equation (1), (3), (4)
        iou = nodes.data['iou'] + self.b_iou
        i, o, u = th.chunk(iou, 3, 1)
        i, o, u = th.sigmoid(i), th.sigmoid(o), th.tanh(u)
        # equation (5)
        c = i * u + nodes.data['c']
        # equation (6)
        h = o * th.tanh(c)
        return {'h' : h, 'c' : c}
 ##############################################################################
 # Step 3: define traversal
 # ------------------------
 #
 # After defining the message passing functions, we then need to induce the
 # right order to trigger them. This is a significant departure from models
 # such as GCN, where all nodes are pulling messages from upstream ones
 # *simultaneously*.
 #
 # In the case of Tree-LSTM, messages start from leaves of the tree, and
 # propogate/processed upwards until they reach the roots. A visulization
 # is as follows:
 #
 # .. figure:: https://i.loli.net/2018/11/09/5be4b5d2df54d.gif
 #    :alt:
 #
 # DGL defines a generator to perform the topological sort, each item is a
 # tensor recording the nodes from bottom level to the roots. One can
 # appreciate the degree of parallelism by inspecting the difference of the
 # followings:
 #
 print('Traversing one tree:')
 print(dgl.topological_nodes_generator(a_tree))
 print('Traversing many trees at the same time:')
 print(dgl.topological_nodes_generator(graph))
 ##############################################################################
 # We then call :meth:`~dgl.DGLGraph.prop_nodes` to trigger the message passing:
 import dgl.function as fn
 import torch as th
 graph.ndata['a'] = th.ones(graph.number_of_nodes(), 1)
 graph.register_message_func(fn.copy_src('a', 'a'))
 graph.register_reduce_func(fn.sum('a', 'a'))
 traversal_order = dgl.topological_nodes_generator(graph)
 graph.prop_nodes(traversal_order)
 # the following is a syntax sugar that does the same
 # dgl.prop_nodes_topo(graph)
 ##############################################################################
 # .. note::
 #
 #    Before we call :meth:`~dgl.DGLGraph.prop_nodes`, we must specify a
 #    `message_func` and `reduce_func` in advance, here we use built-in
 #    copy-from-source and sum function as our message function and reduce
 #    function for demonstration.
 #
 # Putting it together
 # -------------------
 #
 # Here is the complete code that specifies the ``Tree-LSTM`` class:
 #
 class TreeLSTM(nn.Module):
    def __init__(self,
                 num_vocabs,
                 x_size,
                 h_size,
                 num_classes,
                 dropout,
                 pretrained_emb=None):
        super(TreeLSTM, self).__init__()
        self.x_size = x_size
        self.embedding = nn.Embedding(num_vocabs, x_size)
        if pretrained_emb is not None:
            print('Using glove')
            self.embedding.weight.data.copy_(pretrained_emb)
            self.embedding.weight.requires_grad = True
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(h_size, num_classes)
        self.cell = TreeLSTMCell(x_size, h_size)
    def forward(self, batch, h, c):
        """Compute tree-lstm prediction given a batch.
        Parameters
        ----------
        batch : dgl.data.SSTBatch
            The data batch.
        h : Tensor
            Initial hidden state.
        c : Tensor
            Initial cell state.
        Returns
        -------
        logits : Tensor
            The prediction of each node.
        """
        g = batch.graph
        g.register_message_func(self.cell.message_func)
        g.register_reduce_func(self.cell.reduce_func)
        g.register_apply_node_func(self.cell.apply_node_func)
        # feed embedding
        embeds = self.embedding(batch.wordid * batch.mask)
        g.ndata['iou'] = self.cell.W_iou(self.dropout(embeds)) * batch.mask.float().unsqueeze(-1)
        g.ndata['h'] = h
        g.ndata['c'] = c
        # propagate
        dgl.prop_nodes_topo(g)
        # compute logits
        h = self.dropout(g.ndata.pop('h'))
        logits = self.linear(h)
        return logits
 ##############################################################################
 # Main Loop
 # ---------
 #
 # Finally, we could write a training paradigm in PyTorch:
 #
 from torch.utils.data import DataLoader
 import torch.nn.functional as F
 device = th.device('cpu')
 # hyper parameters
 x_size = 256
 h_size = 256
 dropout = 0.5
 lr = 0.05
 weight_decay = 1e-4
 epochs = 10
 # create the model
 model = TreeLSTM(trainset.num_vocabs,
                 x_size,
                 h_size,
                 trainset.num_classes,
                 dropout)
 print(model)
 # create the optimizer
 optimizer = th.optim.Adagrad(model.parameters(),
                          lr=lr,
                          weight_decay=weight_decay)
 train_loader = DataLoader(dataset=tiny_sst,
                          batch_size=5,
                          collate_fn=data.SST.batcher(device),
                          shuffle=False,
                          num_workers=0)
 # training loop
 for epoch in range(epochs):
    for step, batch in enumerate(train_loader):
        g = batch.graph
        n = g.number_of_nodes()
        h = th.zeros((n, h_size))
        c = th.zeros((n, h_size))
        logits = model(batch, h, c)
        logp = F.log_softmax(logits, 1)
        loss = F.nll_loss(logp, batch.label, reduction='sum') 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        pred = th.argmax(logits, 1)
        acc = float(th.sum(th.eq(batch.label, pred))) / len(batch.label)
        print("Epoch {:05d} | Step {:05d} | Loss {:.4f} | Acc {:.4f} |".format(
            epoch, step, loss.item(), acc))
 ##############################################################################
 # To train the model on full dataset with different settings(CPU/GPU,
 # etc.), please refer to our repo's
 # `example <https://github.com/jermainewang/dgl/tree/master/examples/pytorch/tree_lstm>`__.
 # Besides, we also provide an implementation of the Child-Sum Tree LSTM.
--- a/tutorials/models/2_small_graph/README.txt
+++ b/tutorials/models/2_small_graph/README.txt
+.. _tutorials2-index:
+Dealing with many small graphs
+------------------------------
+* **Tree-LSTM** `[paper] <https://arxiv.org/abs/1503.00075>`__ `[tutorial] <models/3_tree-lstm.html>`__
+  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/tree_lstm/tree_lstm.py>`__:
+  sentences of natural languages have inherent structures, which are thrown away
+  by treating them simply as sequences. Tree-LSTM is a powerful model that learns
+  the representation by leveraging prior syntactic structures (e.g. parse-tree).
+  The challenge to train it well is that simply by padding a sentence to the
+  maximum length no longer works, since trees of different sentences have
+  different sizes and topologies. DGL solves this problem by throwing the trees
+  into a bigger "container" graph, and use message-passing to explore maximum
+  parallelism. The key API we use is batching.
--- a/tutorials/models/5_dgmg.py
+++ b/tutorials/models/5_dgmg.py
 """
 .. _model-dgmg:
 Tutorial for Generative Models of Graphs
 ===========================================
 **Author**: `Mufei Li <https://github.com/mufeili>`_,
 `Lingfan Yu <https://github.com/ylfdq1118>`_, Zheng Zhang
 """
 ##############################################################################
 #
 # In earlier tutorials we have seen how learned embedding of a graph and/or
 # a node allow applications such as `semi-supervised classification for nodes
 # <http://docs.dgl.ai/tutorials/models/1_gcn.html#sphx-glr-tutorials-models-1-gcn-py>`__
 # or `sentiment analysis
 # <http://docs.dgl.ai/tutorials/models/3_tree-lstm.html#sphx-glr-tutorials-models-3-tree-lstm-py>`__.
 # Wouldn't it be interesting to predict the future evolution of the graph and
 # perform the analysis iteratively?
 #
 # We will need to generate a variety of graph samples, in other words, we need
 # **generative models** of graphs. Instead of and/or in addition to learning
 # node and edge features, we want to model the distribution of arbitrary graphs.
 # While general generative models can model the density function explicitly and
 # implicitly and generate samples at once or sequentially, we will only focus
 # on explicit generative models for sequential generation here. Typical applications
 # include drug/material discovery, chemical processes, proteomics, etc.
 #
 # Introduction
 # --------------------
 # The primitive actions of mutating a graph in DGL are nothing more than ``add_nodes``
 # and ``add_edges``. That is, if we were to draw a circle of 3 nodes,
 #
 # .. figure:: https://user-images.githubusercontent.com/19576924/48313438-78baf000-e5f7-11e8-931e-cd00ab34fa50.gif
 #    :alt:
 #
 # we can simply write the code as:
 #
 import dgl
 g = dgl.DGLGraph()
 g.add_nodes(1)              # Add node 0
 g.add_nodes(1)              # Add node 1
 # Edges in DGLGraph are directed by default.
 # For undirected edges, we add edges for both directions.
 g.add_edges([1, 0], [0, 1]) # Add edges (1, 0), (0, 1)
 g.add_nodes(1)              # Add node 2
 g.add_edges([2, 1], [1, 2]) # Add edges (2, 1), (1, 2)
 g.add_edges([2, 0], [0, 2]) # Add edges (2, 0), (0, 2)
 #######################################################################################
 # Real-world graphs are much more complex. There are many families of graphs,
 # with different sizes, topologies, node types, edge types, and the possibility
 # of multigraphs. Besides, a same graph can be generated in many different
 # orders. Regardless, the generative process entails a few steps:
 #
 # - Encode a changing graph,
 # - Perform actions stochastically,
 # - Collect error signals and optimize the model parameters (If we are training)
 #
 # When it comes to implementation, another important aspect is speed: how do we
 # parallelize the computation given that generating a graph is fundamentally a
 # sequential process?
 #
 # .. note::
 #
 #    To be sure, this is not necessarily a hard constraint, one can imagine
 #    that subgraphs can be built in parallel and then get assembled. But we
 #    will restrict ourselves to the sequential processes for this tutorial.
 #
 # In tutorial, we will first focus on how to train and generate one graph at
 # a time, exploring parallelism within the graph embedding operation, an
 # essential building block. We will end with a simple optimization that
 # delivers a 2x speedup by batching across graphs.
 #
 # DGMG: the main flow
 # --------------------
 # We pick DGMG (
 # `Learning Deep Generative Models of Graphs <https://arxiv.org/abs/1803.03324>`__
 # ) as an exercise to implement a graph generative model using DGL, primarily
 # because its algorithmic framework is general but also challenging to parallelize.
 #
 # .. note::
 #
 #    While it's possible for DGMG to handle complex graphs with typed nodes,
 #    typed edges and multigraphs, we only present a simplified version of it
 #    for generating graph topologies.
 #
 # DGMG generates a graph by following a state machine, which is basically a
 # two-level loop:  generate one node at a time, and connect it to a subset of
 # the existing nodes, one at a time. This is similar to language modeling: the
 # generative process is an iterative one that emits one word/character/sentence
 # at a time, conditioned on the sequence generated so far.
 #
 # At each time step, we either
 #      - add a new node to the graph, or
 #      - select two existing nodes and add an edge between them
 #
 # .. figure:: https://user-images.githubusercontent.com/19576924/48605003-7f11e900-e9b6-11e8-8880-87362348e154.png
 #    :alt:
 #
 # The Python code will look as follows; in fact, this is *exactly* how inference
 # with DGMG is implemented in DGL:
 #
 def forward_inference(self):
    stop = self.add_node_and_update()
    while (not stop) and (self.g.number_of_nodes() < self.v_max + 1):
        num_trials = 0
        to_add_edge = self.add_edge_or_not()
        while to_add_edge and (num_trials < self.g.number_of_nodes() - 1):
            self.choose_dest_and_update()
            num_trials += 1
            to_add_edge = self.add_edge_or_not()
        stop = self.add_node_and_update()
    return self.g
 #######################################################################################
 # Assume we have a pre-trained model for generating cycles of nodes 10 - 20, let's see
 # how it generates a cycle on the fly during inference. You can also use the code below
 # for creating animation with your own model.
 #
 # ::
 #
 #     import torch
 #     import matplotlib.animation as animation
 #     import matplotlib.pyplot as plt
 #     import networkx as nx
 #     from copy import deepcopy
 #
 #     if __name__ == '__main__':
 #         # pre-trained model saved with path ./model.pth
 #         model = torch.load('./model.pth')
 #         model.eval()
 #         g = model()
 #
 #         src_list = g.edges()[1]
 #         dest_list = g.edges()[0]
 #
 #         evolution = []
 #
 #         nx_g = nx.Graph()
 #         evolution.append(deepcopy(nx_g))
 #
 #         for i in range(0, len(src_list), 2):
 #             src = src_list[i].item()
 #             dest = dest_list[i].item()
 #             if src not in nx_g.nodes():
 #                 nx_g.add_node(src)
 #                 evolution.append(deepcopy(nx_g))
 #             if dest not in nx_g.nodes():
 #                 nx_g.add_node(dest)
 #                 evolution.append(deepcopy(nx_g))
 #             nx_g.add_edges_from([(src, dest), (dest, src)])
 #             evolution.append(deepcopy(nx_g))
 #
 #         def animate(i):
 #             ax.cla()
 #             g_t = evolution[i]
 #             nx.draw_circular(g_t, with_labels=True, ax=ax,
 #                              node_color=['#FEBD69'] * g_t.number_of_nodes())
 #
 #         fig, ax = plt.subplots()
 #         ani = animation.FuncAnimation(fig, animate,
 #                                       frames=len(evolution),
 #                                       interval=600)
 #
 # .. figure:: https://user-images.githubusercontent.com/19576924/48928548-2644d200-ef1b-11e8-8591-da93345382ad.gif
 #    :alt:
 #
 # DGMG: optimization objective
 # ------------------------------
 # Similar to language modeling, DGMG trains the model with *behavior cloning*,
 # or *teacher forcing*. Let's assume for each graph there exists a sequence of
 # *oracle actions* :math:`a_{1},\cdots,a_{T}` that generates it. What the model
 # does is to follow these actions, compute the joint probabilities of such
 # action sequences, and maximize them.
 #
 # By chain rule, the probability of taking :math:`a_{1},\cdots,a_{T}` is:
 #
 # .. math::
 #
 #    p(a_{1},\cdots, a_{T}) = p(a_{1})p(a_{2}|a_{1})\cdots p(a_{T}|a_{1},\cdots,a_{T-1}).\\
 #
 # The optimization objective is then simply the typical MLE loss:
 #
 # .. math::
 #
 #    -\log p(a_{1},\cdots,a_{T})=-\sum_{t=1}^{T}\log p(a_{t}|a_{1},\cdots, a_{t-1}).\\
 #
 def forward_train(self, actions):
    """
    - actions: list
        - Contains a_1, ..., a_T described above
    - self.prepare_for_train()
        - Initializes self.action_step to be 0, which will get
          incremented by 1 everytime it is called.
        - Initializes objects recording log p(a_t|a_1,...a_{t-1})
    Returns
    -------
    - self.get_log_prob(): log p(a_1, ..., a_T)
    """
    self.prepare_for_train()
    stop = self.add_node_and_update(a=actions[self.action_step])
    while not stop:
        to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
        while to_add_edge:
            self.choose_dest_and_update(a=actions[self.action_step])
            to_add_edge = self.add_edge_or_not(a=actions[self.action_step])
        stop = self.add_node_and_update(a=actions[self.action_step])
    return self.get_log_prob()
 #######################################################################################
 # The key difference between ``forward_train`` and ``forward_inference`` is
 # that the training process takes oracle actions as input, and returns log
 # probabilities for evaluating the loss.
 #
 # DGMG: the implementation
 # --------------------------
 # The ``DGMG`` class
 # ``````````````````````````
 # Below one can find the skeleton code for the model. We will gradually
 # fill in the details for each function.
 #
 import torch.nn as nn
 class DGMGSkeleton(nn.Module):
    def __init__(self, v_max):
        """
        Parameters
        ----------
        v_max: int
            Max number of nodes considered
        """
        super(DGMGSkeleton, self).__init__()
        # Graph configuration
        self.v_max = v_max
    def add_node_and_update(self, a=None):
        """Decide if to add a new node.
        If a new node should be added, update the graph."""
        return NotImplementedError
    def add_edge_or_not(self, a=None):
        """Decide if a new edge should be added."""
        return NotImplementedError
    def choose_dest_and_update(self, a=None):
        """Choose destination and connect it to the latest node.
        Add edges for both directions and update the graph."""
        return NotImplementedError
    def forward_train(self, actions):
        """Forward at training time. It records the probability
        of generating a ground truth graph following the actions."""
        return NotImplementedError
    def forward_inference(self):
        """Forward at inference time.
        It generates graphs on the fly."""
        return NotImplementedError
    def forward(self, actions=None):
        # The graph we will work on
        self.g = dgl.DGLGraph()
        # If there are some features for nodes and edges,
        # zero tensors will be set for those of new nodes and edges.
        self.g.set_n_initializer(dgl.frame.zero_initializer)
        self.g.set_e_initializer(dgl.frame.zero_initializer)
        if self.training:
            return self.forward_train(actions=actions)
        else:
            return self.forward_inference()
 #######################################################################################
 # Encoding a dynamic graph
 # ``````````````````````````
 # All the actions generating a graph are sampled from probability
 # distributions. In order to do that, we must project the structured data,
 # namely the graph, onto an Euclidean space. The challenge is that such
 # process, called *embedding*, needs to be repeated as the graphs mutate.
 #
 # Graph Embedding
 # ''''''''''''''''''''''''''
 # Let :math:`G=(V,E)` be an arbitrary graph. Each node :math:`v` has an
 # embedding vector :math:`\textbf{h}_{v} \in \mathbb{R}^{n}`. Similarly,
 # the graph has an embedding vector :math:`\textbf{h}_{G} \in \mathbb{R}^{k}`.
 # Typically, :math:`k > n` since a graph contains more information than
 # an individual node.
 #
 # The graph embedding is a weighted sum of node embeddings under a linear
 # transformation:
 #
 # .. math::
 #
 #    \textbf{h}_{G} =\sum_{v\in V}\text{Sigmoid}(g_m(\textbf{h}_{v}))f_{m}(\textbf{h}_{v}),\\
 #
 # The first term, :math:`\text{Sigmoid}(g_m(\textbf{h}_{v}))`, computes a
 # gating function and can be thought as how much the overall graph embedding
 # attends on each node. The second term :math:`f_{m}:\mathbb{R}^{n}\rightarrow\mathbb{R}^{k}`
 # maps the node embeddings to the space of graph embeddings.
 #
 # We implement graph embedding as a ``GraphEmbed`` class:
 #
 import torch
 class GraphEmbed(nn.Module):
    def __init__(self, node_hidden_size):
        super(GraphEmbed, self).__init__()
        # Setting from the paper
        self.graph_hidden_size = 2 * node_hidden_size
        # Embed graphs
        self.node_gating = nn.Sequential(
            nn.Linear(node_hidden_size, 1),
            nn.Sigmoid()
        )
        self.node_to_graph = nn.Linear(node_hidden_size,
                                       self.graph_hidden_size)
    def forward(self, g):
        if g.number_of_nodes() == 0:
            return torch.zeros(1, self.graph_hidden_size)
        else:
            # Node features are stored as hv in ndata.
            hvs = g.ndata['hv']
            return (self.node_gating(hvs) *
                    self.node_to_graph(hvs)).sum(0, keepdim=True)
 #######################################################################################
 # Update node embeddings via graph propagation
 # ''''''''''''''''''''''''''''''''''''''''''''
 #
 # The mechanism of updating node embeddings in DGMG is similar to that for
 # graph convolutional networks. For a node :math:`v` in the graph, its
 # neighbor :math:`u` sends a message to it with
 #
 # .. math::
 #
 #    \textbf{m}_{u\rightarrow v}=\textbf{W}_{m}\text{concat}([\textbf{h}_{v}, \textbf{h}_{u}, \textbf{x}_{u, v}]) + \textbf{b}_{m},\\
 #
 # where :math:`\textbf{x}_{u,v}` is the embedding of the edge between
 # :math:`u` and :math:`v`.
 #
 # After receiving messages from all its neighbors, :math:`v` summarizes them
 # with a node activation vector
 #
 # .. math::
 #
 #    \textbf{a}_{v} = \sum_{u: (u, v)\in E}\textbf{m}_{u\rightarrow v}\\
 #
 # and use this information to update its own feature:
 #
 # .. math::
 #
 #    \textbf{h}'_{v} = \textbf{GRU}(\textbf{h}_{v}, \textbf{a}_{v}).\\
 #
 # Performing all the operations above once for all nodes synchronously is
 # called one round of graph propagation. The more rounds of graph propagation
 # we perform, the longer distance messages travel throughout the graph.
 #
 # With dgl, we implement graph propagation with ``g.update_all``. Note that
 # the message notation here can be a bit confusing. While the authors refer
 # to :math:`\textbf{m}_{u\rightarrow v}` as messages, our message function
 # below only passes :math:`\text{concat}([\textbf{h}_{u}, \textbf{x}_{u, v}])`.
 # The operation :math:`\textbf{W}_{m}\text{concat}([\textbf{h}_{v}, \textbf{h}_{u}, \textbf{x}_{u, v}]) + \textbf{b}_{m}`
 # is then performed across all edges at once for efficiency consideration.
 #
 from functools import partial
 class GraphProp(nn.Module):
    def __init__(self, num_prop_rounds, node_hidden_size):
        super(GraphProp, self).__init__()
        self.num_prop_rounds = num_prop_rounds
        # Setting from the paper
        self.node_activation_hidden_size = 2 * node_hidden_size
        message_funcs = []
        node_update_funcs = []
        self.reduce_funcs = []
        for t in range(num_prop_rounds):
            # input being [hv, hu, xuv]
            message_funcs.append(nn.Linear(2 * node_hidden_size + 1,
                                           self.node_activation_hidden_size))
            self.reduce_funcs.append(partial(self.dgmg_reduce, round=t))
            node_update_funcs.append(
                nn.GRUCell(self.node_activation_hidden_size,
                           node_hidden_size))
        self.message_funcs = nn.ModuleList(message_funcs)
        self.node_update_funcs = nn.ModuleList(node_update_funcs)
    def dgmg_msg(self, edges):
        """For an edge u->v, return concat([h_u, x_uv])"""
        return {'m': torch.cat([edges.src['hv'],
                                edges.data['he']],
                               dim=1)}
    def dgmg_reduce(self, nodes, round):
        hv_old = nodes.data['hv']
        m = nodes.mailbox['m']
        message = torch.cat([
            hv_old.unsqueeze(1).expand(-1, m.size(1), -1), m], dim=2)
        node_activation = (self.message_funcs[round](message)).sum(1)
        return {'a': node_activation}
    def forward(self, g):
        if g.number_of_edges() > 0:
            for t in range(self.num_prop_rounds):
                g.update_all(message_func=self.dgmg_msg,
                             reduce_func=self.reduce_funcs[t])
                g.ndata['hv'] = self.node_update_funcs[t](
                     g.ndata['a'], g.ndata['hv'])
 #######################################################################################
 # Actions
 # ``````````````````````````
 # All actions are sampled from distributions parameterized using neural nets
 # and we introduce them in turn.
 #
 # Action 1: add nodes
 # ''''''''''''''''''''''''''
 #
 # Given the graph embedding vector :math:`\textbf{h}_{G}`, we evaluate
 #
 # .. math::
 #
 #    \text{Sigmoid}(\textbf{W}_{\text{add node}}\textbf{h}_{G}+b_{\text{add node}}),\\
 #
 # which is then used to parametrize a Bernoulli distribution for deciding whether
 # to add a new node.
 #
 # If a new node is to be added, we initialize its feature with
 #
 # .. math::
 #
 #    \textbf{W}_{\text{init}}\text{concat}([\textbf{h}_{\text{init}} , \textbf{h}_{G}])+\textbf{b}_{\text{init}},\\
 #
 # where :math:`\textbf{h}_{\text{init}}` is a learnable embedding module for
 # untyped nodes.
 #
 import torch.nn.functional as F
 from torch.distributions import Bernoulli
 def bernoulli_action_log_prob(logit, action):
    """Calculate the log p of an action with respect to a Bernoulli
    distribution. Use logit rather than prob for numerical stability."""
    if action == 0:
        return F.logsigmoid(-logit)
    else:
        return F.logsigmoid(logit)
 class AddNode(nn.Module):
    def __init__(self, graph_embed_func, node_hidden_size):
        super(AddNode, self).__init__()
        self.graph_op = {'embed': graph_embed_func}
        self.stop = 1
        self.add_node = nn.Linear(graph_embed_func.graph_hidden_size, 1)
        # If to add a node, initialize its hv
        self.node_type_embed = nn.Embedding(1, node_hidden_size)
        self.initialize_hv = nn.Linear(node_hidden_size + \
                                       graph_embed_func.graph_hidden_size,
                                       node_hidden_size)
        self.init_node_activation = torch.zeros(1, 2 * node_hidden_size)
    def _initialize_node_repr(self, g, node_type, graph_embed):
        """Whenver a node is added, initialize its representation."""
        num_nodes = g.number_of_nodes()
        hv_init = self.initialize_hv(
            torch.cat([
                self.node_type_embed(torch.LongTensor([node_type])),
                graph_embed], dim=1))
        g.nodes[num_nodes - 1].data['hv'] = hv_init
        g.nodes[num_nodes - 1].data['a'] = self.init_node_activation
    def prepare_training(self):
        self.log_prob = []
    def forward(self, g, action=None):
        graph_embed = self.graph_op['embed'](g)
        logit = self.add_node(graph_embed)
        prob = torch.sigmoid(logit)
        if not self.training:
            action = Bernoulli(prob).sample().item()
        stop = bool(action == self.stop)
        if not stop:
            g.add_nodes(1)
            self._initialize_node_repr(g, action, graph_embed)
        if self.training:
            sample_log_prob = bernoulli_action_log_prob(logit, action)
            self.log_prob.append(sample_log_prob)
        return stop
 #######################################################################################
 # Action 2: add edges
 # ''''''''''''''''''''''''''
 #
 # Given the graph embedding vector :math:`\textbf{h}_{G}` and the node
 # embedding vector :math:`\textbf{h}_{v}` for the latest node :math:`v`,
 # we evaluate
 #
 # .. math::
 #
 #    \text{Sigmoid}(\textbf{W}_{\text{add edge}}\text{concat}([\textbf{h}_{G}, \textbf{h}_{v}])+b_{\text{add edge}}),\\
 #
 # which is then used to parametrize a Bernoulli distribution for deciding
 # whether to add a new edge starting from :math:`v`.
 #
 class AddEdge(nn.Module):
    def __init__(self, graph_embed_func, node_hidden_size):
        super(AddEdge, self).__init__()
        self.graph_op = {'embed': graph_embed_func}
        self.add_edge = nn.Linear(graph_embed_func.graph_hidden_size + \
                                  node_hidden_size, 1)
    def prepare_training(self):
        self.log_prob = []
    def forward(self, g, action=None):
        graph_embed = self.graph_op['embed'](g)
        src_embed = g.nodes[g.number_of_nodes() - 1].data['hv']
        logit = self.add_edge(torch.cat(
            [graph_embed, src_embed], dim=1))
        prob = torch.sigmoid(logit)
        if self.training:
            sample_log_prob = bernoulli_action_log_prob(logit, action)
            self.log_prob.append(sample_log_prob)
        else:
            action = Bernoulli(prob).sample().item()
        to_add_edge = bool(action == 0)
        return to_add_edge
 #######################################################################################
 # Action 3: choosing destination
 # '''''''''''''''''''''''''''''''''
 #
 # When action 2 returns True, we need to choose a destination for the
 # latest node :math:`v`.
 #
 # For each possible destination :math:`u\in\{0, \cdots, v-1\}`, the
 # probability of choosing it is given by
 #
 # .. math::
 #
 #    \frac{\text{exp}(\textbf{W}_{\text{dest}}\text{concat}([\textbf{h}_{u}, \textbf{h}_{v}])+\textbf{b}_{\text{dest}})}{\sum_{i=0}^{v-1}\text{exp}(\textbf{W}_{\text{dest}}\text{concat}([\textbf{h}_{i}, \textbf{h}_{v}])+\textbf{b}_{\text{dest}})}\\
 #
 from torch.distributions import Categorical
 class ChooseDestAndUpdate(nn.Module):
    def __init__(self, graph_prop_func, node_hidden_size):
        super(ChooseDestAndUpdate, self).__init__()
        self.graph_op = {'prop': graph_prop_func}
        self.choose_dest = nn.Linear(2 * node_hidden_size, 1)
    def _initialize_edge_repr(self, g, src_list, dest_list):
        # For untyped edges, we only add 1 to indicate its existence.
        # For multiple edge types, we can use a one hot representation
        # or an embedding module.
        edge_repr = torch.ones(len(src_list), 1)
        g.edges[src_list, dest_list].data['he'] = edge_repr
    def prepare_training(self):
        self.log_prob = []
    def forward(self, g, dest):
        src = g.number_of_nodes() - 1
        possible_dests = range(src)
        src_embed_expand = g.nodes[src].data['hv'].expand(src, -1)
        possible_dests_embed = g.nodes[possible_dests].data['hv']
        dests_scores = self.choose_dest(
            torch.cat([possible_dests_embed,
                       src_embed_expand], dim=1)).view(1, -1)
        dests_probs = F.softmax(dests_scores, dim=1)
        if not self.training:
            dest = Categorical(dests_probs).sample().item()
        if not g.has_edge_between(src, dest):
            # For undirected graphs, we add edges for both directions
            # so that we can perform graph propagation.
            src_list = [src, dest]
            dest_list = [dest, src]
            g.add_edges(src_list, dest_list)
            self._initialize_edge_repr(g, src_list, dest_list)
            self.graph_op['prop'](g)
        if self.training:
            if dests_probs.nelement() > 1:
                self.log_prob.append(
                    F.log_softmax(dests_scores, dim=1)[:, dest: dest + 1])
 #######################################################################################
 # Putting it together
 # ``````````````````````````
 #
 # We are now ready to have a complete implementation of the model class.
 #
 class DGMG(DGMGSkeleton):
    def __init__(self, v_max, node_hidden_size,
                 num_prop_rounds):
        super(DGMG, self).__init__(v_max)
        # Graph embedding module
        self.graph_embed = GraphEmbed(node_hidden_size)
        # Graph propagation module
        self.graph_prop = GraphProp(num_prop_rounds,
                                    node_hidden_size)
        # Actions
        self.add_node_agent = AddNode(
            self.graph_embed, node_hidden_size)
        self.add_edge_agent = AddEdge(
            self.graph_embed, node_hidden_size)
        self.choose_dest_agent = ChooseDestAndUpdate(
            self.graph_prop, node_hidden_size)
        # Forward functions
        self.forward_train = partial(forward_train, self=self)
        self.forward_inference = partial(forward_inference, self=self)
    @property
    def action_step(self):
        old_step_count = self.step_count
        self.step_count += 1
        return old_step_count
    def prepare_for_train(self):
        self.step_count = 0
        self.add_node_agent.prepare_training()
        self.add_edge_agent.prepare_training()
        self.choose_dest_agent.prepare_training()
    def add_node_and_update(self, a=None):
        """Decide if to add a new node.
        If a new node should be added, update the graph."""
        return self.add_node_agent(self.g, a)
    def add_edge_or_not(self, a=None):
        """Decide if a new edge should be added."""
        return self.add_edge_agent(self.g, a)
    def choose_dest_and_update(self, a=None):
        """Choose destination and connect it to the latest node.
        Add edges for both directions and update the graph."""
        self.choose_dest_agent(self.g, a)
    def get_log_prob(self):
        add_node_log_p = torch.cat(self.add_node_agent.log_prob).sum()
        add_edge_log_p = torch.cat(self.add_edge_agent.log_prob).sum()
        choose_dest_log_p = torch.cat(self.choose_dest_agent.log_prob).sum()
        return add_node_log_p + add_edge_log_p + choose_dest_log_p
 #######################################################################################
 # Below is an animation where a graph is generated on the fly
 # after every 10 batches of training for the first 400 batches. One
 # can see how our model improves over time and begins generating cycles.
 #
 # .. figure:: https://user-images.githubusercontent.com/19576924/48929291-60fe3880-ef22-11e8-832a-fbe56656559a.gif
 #    :alt:
 #
 # For generative models, we can evaluate its performance by checking the percentage
 # of valid graphs among the graphs it generates on the fly.
 import torch.utils.model_zoo as model_zoo
 # Download a pre-trained model state dict for generating cycles with 10-20 nodes.
 state_dict = model_zoo.load_url('https://s3.us-east-2.amazonaws.com/dgl.ai/model/dgmg_cycles-5a0c40be.pth')
 model = DGMG(v_max=20, node_hidden_size=16, num_prop_rounds=2)
 model.load_state_dict(state_dict)
 model.eval()
 def is_valid(g):
    # Check if g is a cycle having 10-20 nodes.
    def _get_previous(i, v_max):
        if i == 0:
            return v_max
        else:
            return i - 1
    def _get_next(i, v_max):
        if i == v_max:
            return 0
        else:
            return i + 1
    size = g.number_of_nodes()
    if size < 10 or size > 20:
        return False
    for node in range(size):
        neighbors = g.successors(node)
        if len(neighbors) != 2:
            return False
        if _get_previous(node, size - 1) not in neighbors:
            return False
        if _get_next(node, size - 1) not in neighbors:
            return False
    return True
 num_valid = 0
 for i in range(100):
    g = model()
    num_valid += is_valid(g)
 del model
 print('Among 100 graphs generated, {}% are valid.'.format(num_valid))
 #######################################################################################
 # For the complete implementation, see `dgl DGMG example
 # <https://github.com/jermainewang/dgl/tree/master/examples/pytorch/dgmg>`__.
 #
 # Batched Graph Generation
 # ---------------------------
 #
 # Speeding up DGMG is hard since each graph can be generated with a
 # unique sequence of actions. One way to explore parallelism is to adopt
 # asynchronous gradient descent with multiple processes. Each of them
 # works on one graph at a time and the processes are loosely coordinated
 # by a parameter server. This is the approach that the authors adopted
 # and we can also use.
 #
 # DGL explores parallelism in the message-passing framework, on top of
 # the framework-provided tensor operation. The earlier tutorial already
 # does that in the message propagation and graph embedding phases, but
 # only within one graph. For a batch of graphs, a for loop is then needed:
 #
 # ::
 #
 #     for g in g_list:
 #     self.graph_prop(g)
 #
 # We can modify the code to work on a batch of graphs at once by replacing
 # these lines with the following. On CPU with a Mac machine, we instantly
 # enjoy a 6~7x reduction for the graph propagation part.
 # ::
 #
 #     bg = dgl.batch(g_list)
 #     self.graph_prop(bg)
 #     g_list = dgl.unbatch(bg)
 #
 # We have already used this trick of calling ``dgl.batch`` in the
 # `Tree-LSTM tutorial
 # <http://docs.dgl.ai/tutorials/models/3_tree-lstm.html#sphx-glr-tutorials-models-3-tree-lstm-py>`__
 # , and it is worth explaining one more time why this is so.
 #
 # By batching many small graphs, DGL internally maintains a large *container*
 # graph (``BatchedDGLGraph``) over which ``update_all`` propels message-passing
 # on all the edges and nodes.
 #
 # With ``dgl.batch``, we merge ``g_{1}, ..., g_{N}`` into one single giant
 # graph consisting of :math:`N` isolated small graphs. For example, if we
 # have two graphs with adjacency matrices
 #
 # ::
 #
 #     [0, 1]
 #     [1, 0]
 #
 #     [0, 1, 0]
 #     [1, 0, 0]
 #     [0, 1, 0]
 #
 # ``dgl.batch`` simply gives a graph whose adjacency matrix is
 #
 # ::
 #
 #     [0, 1, 0, 0, 0]
 #     [1, 0, 0, 0, 0]
 #     [0, 1, 0, 0, 0]
 #     [1, 0, 0, 0, 0]
 #     [0, 1, 0, 0, 0]
 #
 # In DGL, the message function is defined on the edges, thus batching scales
 # the processing of edge user-defined functions (UDFs) linearly.
 #
 # The reduce UDFs (i.e ``dgmg_reduce``) works on nodes, and each of them may
 # have different numbers of incoming edges. Using ``degree bucketing``, DGL
 # internally groups nodes with the same in-degrees and calls reduce UDF once
 # for each group. Thus, batching also reduces number of calls to these UDFs.
 #
 # The modification of the node/edge features of a ``BatchedDGLGraph`` object
 # does not take effect on the features of the original small graphs, so we
 # need to replace the old graph list with the new graph list
 # ``g_list = dgl.unbatch(bg)``.
 #
 # The complete code to the batched version can also be found in the example.
 # On our testbed, we get roughly 2x speed up comparing to the previous implementation
 #
--- a/tutorials/models/3_generative_model/README.txt
+++ b/tutorials/models/3_generative_model/README.txt
+.. _tutorials3-index:
+Generative models
+------------------------------
+* **DGMG** `[paper] <https://arxiv.org/abs/1803.03324>`__ `[tutorial] <models/5_dgmg.html>`__
+  `[code] <https://github.com/jermainewang/dgl/tree/master/examples/pytorch/dgmg>`__:
+  this model belongs to the important family that deals with structural
+  generation. DGMG is interesting because its state-machine approach is the most
+  general. It is also very challenging because, unlike Tree-LSTM, every sample
+  has a dynamic, probability-driven structure that is not available before
+  training. We are able to progressively leverage intra- and inter-graph
+  parallelism to steadily improve the performance.
+* **JTNN** `[paper] <https://arxiv.org/abs/1802.04364>`__ `[code (wip)]`: unlike DGMG, this
+  paper generates molecular graphs using the framework of variational
+  auto-encoder. Perhaps more interesting is its approach to build structure
+  hierarchically, in the case of molecular, with junction tree as the middle
+  scaffolding.
--- a/tutorials/models/2_capsule.py
+++ b/tutorials/models/2_capsule.py
@@ -4,62 +4,62 @@
 Capsule Network Tutorial
 ===========================
-**Author**: Jinjing Zhou, `Jake
+**Author**: Jinjing Zhou, `Jake Zhao <https://cs.nyu.edu/~jakezhao/>`_, Zheng Zhang, Jinyang Li
-Zhao <https://cs.nyu.edu/~jakezhao/>`_, Zheng Zhang
-It is perhaps a little surprising that some of the more classical models can
+It is perhaps a little surprising that some of the more classical models
-also be described in terms of graphs, offering a different perspective.
+can also be described in terms of graphs, offering a different
-This tutorial describes how this is done for the `capsule network <http://arxiv.org/abs/1710.09829>`__.
+perspective. This tutorial describes how this can be done for the
+`capsule network <http://arxiv.org/abs/1710.09829>`__.
 """
 #######################################################################################
 # Key ideas of Capsule
 # --------------------
 #
-# There are two key ideas that the Capsule model offers.
+# The Capsule model offers two key ideas.
 #
-# **Richer representations** In classic convolutional network, a scalar
+# **Richer representation** In classic convolutional networks, a scalar
-# value represents the activation of a given feature. Instead, a capsule
+# value represents the activation of a given feature. By contrast, a
-# outputs a vector, whose norm represents the probability of a feature,
+# capsule outputs a vector. The vector's length represents the probability
-# and the orientation its properties.
+# of a feature being present. The vector's orientation represents the
+# various properties of the feature (such as pose, deformation, texture
+# etc.).
 #
-# .. figure:: https://i.imgur.com/55Ovkdh.png
+# |image0|
-#    :alt:
 #
-# **Dynamic routing** To generalize max-pooling, there is another
+# **Dynamic routing** The output of a capsule is preferentially sent to
-# interesting proposed by the authors, as a representational more powerful
+# certain parents in the layer above based on how well the capsule's
-# way to construct higher level feature from its low levels. Consider a
+# prediction agrees with that of a parent. Such dynamic
-# capsule :math:`u_i`. The way :math:`u_i` is integrated to the next level
+# "routing-by-agreement" generalizes the static routing of max-pooling.
-# capsules take two steps:
 #
-# 1. :math:`u_i` projects differently to different higher level capsules
+# During training, routing is done iteratively; each iteration adjusts
-#    via a linear transformation: :math:`\hat{u}_{j|i} = W_{ij}u_i`.
+# "routing weights" between capsules based on their observed agreements,
-# 2. :math:`\hat{u}_{j|i}` routes to the higher level capsules by
+# in a manner similar to a k-means algorithm or `competitive
-#    spreading itself with a weighted sum, and the weight is dynamically
+# learning <https://en.wikipedia.org/wiki/Competitive_learning>`__.
-#    determined by iteratively modify the and checking against the
-#    "consistency" between :math:`\hat{u}_{j|i}` and :math:`v_j`, for any
-#    :math:`v_j`. Note that this is similar to a k-means algorithm or
-#    `competive
-#    learning <https://en.wikipedia.org/wiki/Competitive_learning>`__ in
-#    spirit. At the end of iterations, :math:`v_j` now integrates the
-#    lower level capsules.
 #
-# The full algorithm is the following: |image0|
+# In this tutorial, we show how capsule's dynamic routing algorithm can be
-#
+# naturally expressed as a graph algorithm. Our implementation is adapted
-# The dynamic routing step can be naturally expressed as a graph
+# from `Cedric
-# algorithm. This is the focus of this tutorial. Our implementation is
-# adapted from `Cedric
 # Chee <https://github.com/cedrickchee/capsule-net-pytorch>`__, replacing
-# only the routing layer, and achieving similar speed and accuracy.
+# only the routing layer. Our version achieves similar speed and accuracy.
 #
 # Model Implementation
-# -----------------------------------
+# ----------------------
-# Step 1: Setup and Graph Initialiation
+# Step 1: Setup and Graph Initialization
-# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#
+# The connectivity between two layers of capsules form a directed,
+# bipartite graph, as shown in the Figure below.
+#
+# |image1|
 #
-# The below figure shows the directed bipartitie graph built for capsules
+# Each node :math:`j` is associated with feature :math:`v_j`,
-# network. We denote :math:`b_{ij}`, :math:`\hat{u}_{j|i}` as edge
+# representing its capsule’s output. Each edge is associated with
-# features and :math:`v_j` as node features. |image1|
+# features :math:`b_{ij}` and :math:`\hat{u}_{j|i}`. :math:`b_{ij}`
+# determines routing weights, and :math:`\hat{u}_{j|i}` represents the
+# prediction of capsule :math:`i` for :math:`j`.
 #
+# Here's how we set up the graph and initialize node and edge features.
 import torch.nn as nn
 import torch as th
 import torch.nn.functional as F
@@ -88,32 +88,33 @@ def init_graph(in_nodes, out_nodes, f_size):
 #########################################################################################
 # Step 2: Define message passing functions
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-# Recall the following steps, and they are implemented in the class
-# ``DGLRoutingLayer`` as the followings:
 #
-# 1. Normalize over out edges
+# This is the pseudo code for Capsule's routing algorithm as given in the
+# paper:
+#
+# |image2|
+# We implement pseudo code lines 4-7 in the class `DGLRoutingLayer` as the following steps:
+#
+# 1. Calculate coupling coefficients:
 #
-#    -  Softmax over all out-edge of in-capsules
+#    -  Coefficients are the softmax over all out-edge of in-capsules:
-#       :math:`\textbf{c}_i = \text{softmax}(\textbf{b}_i)`.
+#       :math:`\textbf{c}_{i,j} = \text{softmax}(\textbf{b}_{i,j})`.
 #
-# 2. Weighted sum over all in-capsules
+# 2. Calculate weighted sum over all in-capsules:
 #
-#    -  Out-capsules equals weighted sum of in-capsules
+#    -  Output of a capsule is equal to the weighted sum of its in-capsules
 #       :math:`s_j=\sum_i c_{ij}\hat{u}_{j|i}`
 #
-# 3. Squash Operation
+# 3. Squash outputs:
 #
-#    -  Squashing function is to ensure that short capsule vectors get
+#    -  Squash the length of a capsule's output vector to range (0,1), so it can represent the probability (of some feature being present).
-#       shrunk to almost zero length while the long capsule vectors get
-#       shrunk to a length slightly below 1. Its norm is expected to
-#       represents probabilities at some levels.
 #    -  :math:`v_j=\text{squash}(s_j)=\frac{||s_j||^2}{1+||s_j||^2}\frac{s_j}{||s_j||}`
 #
-# 4. Update weights by agreement
+# 4. Update weights by the amount of agreement:
 #
-#    -  :math:`\hat{u}_{j|i}\cdot v_j` can be considered as agreement
+#    -  The scalar product :math:`\hat{u}_{j|i}\cdot v_j` can be considered as how well capsule :math:`i` agrees with :math:`j`. It is used to update
-#       between current capsule and updated capsule,
 #       :math:`b_{ij}=b_{ij}+\hat{u}_{j|i}\cdot v_j`
 class DGLRoutingLayer(nn.Module):
    def __init__(self, in_nodes, out_nodes, f_size):
        super(DGLRoutingLayer, self).__init__()
@@ -172,9 +173,9 @@ u_hat = th.randn(in_nodes * out_nodes, f_size)
 routing = DGLRoutingLayer(in_nodes, out_nodes, f_size)
 ############################################################################################################
-# We can visualize the behavior by monitoring the entropy of outgoing
+# We can visualize a capsule network's behavior by monitoring the entropy
-# weights, they should start high and then drop, as the assignment
+# of coupling coefficients. They should start high and then drop, as the
-# gradually concentrate:
+# weights gradually concentrate on fewer edges:
 entropy_list = []
 dist_list = []
@@ -193,9 +194,9 @@ plt.xlabel("Number of Routing")
 plt.xticks(np.arange(len(entropy_list)))
 plt.close()
 ############################################################################################################
+# |image3|
 #
-# .. figure:: https://i.imgur.com/dMvu7p3.png
+# Alternatively, we can also watch the evolution of histograms:
-#    :alt:
 import seaborn as sns
 import matplotlib.animation as animation
@@ -216,8 +217,10 @@ ani = animation.FuncAnimation(fig, dist_animate, frames=len(entropy_list), inter
 plt.close()
 ############################################################################################################
-# Alternatively, we can also watch the evolution of histograms: |image2|
+# |image4|
-# Or monitor the how lower level capcules gradually attach to one of the higher level ones:
+#
+# Or monitor the how lower level capsules gradually attach to one of the
+# higher level ones:
 import networkx as nx
 from networkx.algorithms import bipartite
@@ -251,14 +254,16 @@ ani2 = animation.FuncAnimation(fig2, weight_animate, frames=len(dist_list), inte
 plt.close()
 ############################################################################################################
-# |image3|
+# |image5|
 #
-# The full code of this visulization is provided at
+# The full code of this visualization is provided at
 # `link <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/capsule/simple_routing.py>`__; the complete
 # code that trains on MNIST is at `link <https://github.com/jermainewang/dgl/tree/tutorial/examples/pytorch/capsule>`__.
 #
-# .. |image0| image:: https://i.imgur.com/mv1W9Rv.png
+# .. |image0| image:: https://i.imgur.com/55Ovkdh.png
 # .. |image1| image:: https://i.imgur.com/9tc6GLl.png
-# .. |image2| image:: https://github.com/VoVAllen/DGL_Capsule/raw/master/routing_dist.gif
+# .. |image2| image:: https://i.imgur.com/mv1W9Rv.png
-# .. |image3| image:: https://github.com/VoVAllen/DGL_Capsule/raw/master/routing_vis.gif
+# .. |image3| image:: https://i.imgur.com/dMvu7p3.png
-#
+# .. |image4| image:: https://github.com/VoVAllen/DGL_Capsule/raw/master/routing_dist.gif
+# .. |image5| image:: https://github.com/VoVAllen/DGL_Capsule/raw/master/routing_vis.gif
--- a/tutorials/models/4_old_wines/README.txt
+++ b/tutorials/models/4_old_wines/README.txt
+.. _tutorials4-index:
+Old (new) wines in new bottle
+-----------------------------
+* **Capsule** `[paper] <https://arxiv.org/abs/1710.09829>`__ `[tutorial] <models/2_capsule.html>`__
+  `[code] <https://github.com/jermainewang/dgl/tree/master/examples/pytorch/capsule>`__: this new
+  computer vision model has two key ideas -- enhancing the feature representation
+  in a vector form (instead of a scalar) called *capsule*, and replacing
+  maxpooling with dynamic routing. The idea of dynamic routing is to integrate a
+  lower level capsule to one (or several) of a higher level one with
+  non-parametric message-passing. We show how the later can be nicely implemented
+  with DGL APIs.
+* **Transformer** `[paper] <https://arxiv.org/abs/1706.03762>`__ `[tutorial (wip)]` `[code (wip)]` and
+  **Universal Transformer** `[paper] <https://arxiv.org/abs/1807.03819>`__ `[tutorial (wip)]`
+  `[code (wip)]`: these
+  two models replace RNN with several layers of multi-head attention to encode
+  and discover structures among tokens of a sentence. These attention mechanisms
+  can similarly formulated as graph operations with message-passing.
--- a/tutorials/models/README.txt
+++ b/tutorials/models/README.txt
 Graph-based Neural Network Models
 =================================
 We developed DGL with a broad range of applications in mind. Building
 state-of-art models forces us to think hard on the most common and useful APIs,
 learn the hard lessons, and push the system design.
 We have prototyped altogether 10 different models, all of them are ready to run
 out-of-box and some of them are very new graph-based algorithms. In most of the
 cases, they demonstrate the performance, flexibility, and expressiveness of
 DGL. For where we still fall in short, these exercises point to future
 directions.
 We categorize the models below, providing links to the original code and
 tutorial when appropriate. As will become apparent, these models stress the use
 of different DGL APIs.
-Graph Neural Network and its variant
------------------------------------
-* **GCN** `[paper] <https://arxiv.org/abs/1609.02907>`__ `[tutorial] <models/1_gcn.html>`__
-  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/gcn/gcn.py>`__:
-  this is the vanilla GCN. The tutorial covers the basic uses of DGL APIs.
-* **GAT** `[paper] <https://arxiv.org/abs/1710.10903>`__
-  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/gat/gat.py>`__:
-  the key extension of GAT w.r.t vanilla GCN is deploying multi-head attention
-  among neighborhood of a node, thus greatly enhances the capacity and
-  expressiveness of the model.
-* **R-GCN** `[paper] <https://arxiv.org/abs/1703.06103>`__ `[tutorial] <models/4_rgcn.html>`__
-  [code (wip)]: the key
-  difference of RGNN is to allow multi-edges among two entities of a graph, and
-  edges with distinct relationships are encoded differently. This is an
-  interesting extension of GCN that can have a lot of applications of its own.
-* **LGNN** `[paper] <https://arxiv.org/abs/1705.08415>`__ `[tutorial (wip)]` `[code (wip)]`:
-  this model focuses on community detection by inspecting graph structures. It
-  uses representations of both the orignal graph and its line-graph companion. In
-  addition to demonstrate how an algorithm can harness multiple graphs, our
-  implementation shows how one can judiciously mix vanilla tensor operation,
-  sparse-matrix tensor operations, along with message-passing with DGL. 
-* **SSE** `[paper] <http://proceedings.mlr.press/v80/dai18a/dai18a.pdf>`__ `[tutorial (wip)]`
-  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/mxnet/sse/sse_batch.py>`__:
-  the emphasize here is *giant* graph that cannot fit comfortably on one GPU
-  card. SSE is an example to illustrate the co-design of both algrithm and
-  system: sampling to guarantee asymptotic covergence while lowering the
-  complexity, and batching across samples for maximum parallelism.
-Dealing with many small graphs
------------------------------
-* **Tree-LSTM** `[paper] <https://arxiv.org/abs/1503.00075>`__ `[tutorial] <models/3_tree-lstm.html>`__
-  `[code] <https://github.com/jermainewang/dgl/blob/master/examples/pytorch/tree_lstm/tree_lstm.py>`__:
-  sentences of natural languages have inherent structures, which are thrown away
-  by treating them simply as sequences. Tree-LSTM is a powerful model that learns
-  the representation by leveraging prior syntactic structures (e.g. parse-tree).
-  The challenge to train it well is that simply by padding a sentence to the
-  maximum length no longer works, since trees of different sentences have
-  different sizes and topologies. DGL solves this problem by throwing the trees
-  into a bigger "container" graph, and use message-passing to explore maximum
-  parallelism. The key API we use is batching.
-Generative models
------------------------------
-* **DGMG** `[paper] <https://arxiv.org/abs/1803.03324>`__ `[tutorial] <models/5_dgmg.html>`__
-  `[code] <https://github.com/jermainewang/dgl/tree/master/examples/pytorch/dgmg>`__:
-  this model belongs to the important family that deals with structural
-  generation. DGMG is interesting because its state-machine approach is the most
-  general. It is also very challenging because, unlike Tree-LSTM, every sample
-  has a dynamic, probability-driven structure that is not available before
-  training. We are able to progressively leverage intra- and inter-graph
-  parallelism to steadily improve the performance.
-* **JTNN** `[paper] <https://arxiv.org/abs/1802.04364>`__ `[code (wip)]`: unlike DGMG, this
-  paper generates molecular graphs using the framework of variational
-  auto-encoder. Perhaps more interesting is its approach to build structure
-  hierarchically, in the case of molecular, with junction tree as the middle
-  scaffolding.
-Old (new) wines in new bottle
-----------------------------
-* **Capsule** `[paper] <https://arxiv.org/abs/1710.09829>`__ `[tutorial] <models/2_capsule.html>`__
-  `[code] <https://github.com/jermainewang/dgl/tree/master/examples/pytorch/capsule>`__: this new
-  computer vision model has two key ideas -- enhancing the feature representation
-  in a vector form (instead of a scalar) called *capsule*, and replacing
-  maxpooling with dynamic routing. The idea of dynamic routing is to integrate a
-  lower level capsule to one (or several) of a higher level one with
-  non-parametric message-passing. We show how the later can be nicely implemented
-  with DGL APIs.
-* **Transformer** `[paper] <https://arxiv.org/abs/1706.03762>`__ `[tutorial (wip)]` `[code (wip)]` and
-  **Universal Transformer** `[paper] <https://arxiv.org/abs/1807.03819>`__ `[tutorial (wip)]`
-  `[code (wip)]`: these
-  two models replace RNN with several layers of multi-head attention to encode
-  and discover structures among tokens of a sentence. These attention mechanisms
-  can similarly formulated as graph operations with message-passing.