chore: import upstream snapshot with attribution

python/dgl/nn/pytorch/conv/__init__.py

@@ -0,0 +1,78 @@
+"""Torch modules for graph convolutions."""
+# pylint: disable= no-member, arguments-differ, invalid-name
+
+from .agnnconv import AGNNConv
+from .appnpconv import APPNPConv
+from .atomicconv import AtomicConv
+from .cfconv import CFConv
+from .chebconv import ChebConv
+from .cugraph_gatconv import CuGraphGATConv
+from .cugraph_relgraphconv import CuGraphRelGraphConv
+from .cugraph_sageconv import CuGraphSAGEConv
+from .densechebconv import DenseChebConv
+from .densegraphconv import DenseGraphConv
+from .densesageconv import DenseSAGEConv
+from .dgnconv import DGNConv
+from .dotgatconv import DotGatConv
+from .edgeconv import EdgeConv
+from .edgegatconv import EdgeGATConv
+from .egatconv import EGATConv
+from .egnnconv import EGNNConv
+from .gatconv import GATConv
+from .gatedgcnconv import GatedGCNConv
+from .gatedgraphconv import GatedGraphConv
+from .gatv2conv import GATv2Conv
+from .gcn2conv import GCN2Conv
+from .ginconv import GINConv
+from .gineconv import GINEConv
+from .gmmconv import GMMConv
+from .graphconv import EdgeWeightNorm, GraphConv
+from .grouprevres import GroupRevRes
+from .hgtconv import HGTConv
+from .nnconv import NNConv
+from .pnaconv import PNAConv
+from .relgraphconv import RelGraphConv
+from .sageconv import SAGEConv
+from .sgconv import SGConv
+from .tagconv import TAGConv
+from .twirlsconv import TWIRLSConv, TWIRLSUnfoldingAndAttention
+
+__all__ = [
+    "GraphConv",
+    "EdgeWeightNorm",
+    "GATConv",
+    "GATv2Conv",
+    "EGATConv",
+    "EdgeGATConv",
+    "TAGConv",
+    "RelGraphConv",
+    "SAGEConv",
+    "SGConv",
+    "APPNPConv",
+    "GINConv",
+    "GINEConv",
+    "GatedGraphConv",
+    "GatedGCNConv",
+    "GMMConv",
+    "ChebConv",
+    "AGNNConv",
+    "NNConv",
+    "DenseGraphConv",
+    "DenseSAGEConv",
+    "DenseChebConv",
+    "EdgeConv",
+    "AtomicConv",
+    "CFConv",
+    "DotGatConv",
+    "TWIRLSConv",
+    "TWIRLSUnfoldingAndAttention",
+    "GCN2Conv",
+    "HGTConv",
+    "GroupRevRes",
+    "EGNNConv",
+    "PNAConv",
+    "DGNConv",
+    "CuGraphGATConv",
+    "CuGraphRelGraphConv",
+    "CuGraphSAGEConv",
+]

python/dgl/nn/pytorch/conv/agnnconv.py

@@ -0,0 +1,160 @@
+"""Torch Module for Attention-based Graph Neural Network layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+from torch.nn import functional as F
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import expand_as_pair
+from ...functional import edge_softmax
+
+
+class AGNNConv(nn.Module):
+    r"""Attention-based Graph Neural Network layer from `Attention-based Graph Neural Network for
+    Semi-Supervised Learning <https://arxiv.org/abs/1803.03735>`__
+
+    .. math::
+        H^{l+1} = P H^{l}
+
+    where :math:`P` is computed as:
+
+    .. math::
+        P_{ij} = \mathrm{softmax}_i ( \beta \cdot \cos(h_i^l, h_j^l))
+
+    where :math:`\beta` is a single scalar parameter.
+
+    Parameters
+    ----------
+    init_beta : float, optional
+        The :math:`\beta` in the formula, a single scalar parameter.
+    learn_beta : bool, optional
+        If True, :math:`\beta` will be learnable parameter.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Default: ``False``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import AGNNConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> conv = AGNNConv()
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+            [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+            [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+            [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+            [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
+            [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]],
+        grad_fn=<BinaryReduceBackward>)
+    """
+
+    def __init__(
+        self, init_beta=1.0, learn_beta=True, allow_zero_in_degree=False
+    ):
+        super(AGNNConv, self).__init__()
+        self._allow_zero_in_degree = allow_zero_in_degree
+        if learn_beta:
+            self.beta = nn.Parameter(th.Tensor([init_beta]))
+        else:
+            self.register_buffer("beta", th.Tensor([init_beta]))
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat):
+        r"""
+
+        Description
+        -----------
+        Compute AGNN layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, *)` :math:`N` is the
+            number of nodes, and :math:`*` could be of any shape.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, *)` and :math:`(N_{out}, *)`, the :math:`*` in the later
+            tensor must equal the previous one.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, *)` where :math:`*`
+            should be the same as input shape.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            feat_src, feat_dst = expand_as_pair(feat, graph)
+
+            graph.srcdata["h"] = feat_src
+            graph.srcdata["norm_h"] = F.normalize(feat_src, p=2, dim=-1)
+            if isinstance(feat, tuple) or graph.is_block:
+                graph.dstdata["norm_h"] = F.normalize(feat_dst, p=2, dim=-1)
+            # compute cosine distance
+            graph.apply_edges(fn.u_dot_v("norm_h", "norm_h", "cos"))
+            cos = graph.edata.pop("cos")
+            e = self.beta * cos
+            graph.edata["p"] = edge_softmax(graph, e)
+            graph.update_all(fn.u_mul_e("h", "p", "m"), fn.sum("m", "h"))
+            return graph.dstdata.pop("h")

python/dgl/nn/pytorch/conv/appnpconv.py

@@ -0,0 +1,123 @@
+"""Torch Module for APPNPConv"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from .graphconv import EdgeWeightNorm
+
+
+class APPNPConv(nn.Module):
+    r"""Approximate Personalized Propagation of Neural Predictions layer from `Predict then
+    Propagate: Graph Neural Networks meet Personalized PageRank
+    <https://arxiv.org/pdf/1810.05997.pdf>`__
+
+    .. math::
+        H^{0} &= X
+
+        H^{l+1} &= (1-\alpha)\left(\tilde{D}^{-1/2}
+        \tilde{A} \tilde{D}^{-1/2} H^{l}\right) + \alpha H^{0}
+
+    where :math:`\tilde{A}` is :math:`A` + :math:`I`.
+
+    Parameters
+    ----------
+    k : int
+        The number of iterations :math:`K`.
+    alpha : float
+        The teleport probability :math:`\alpha`.
+    edge_drop : float, optional
+        The dropout rate on edges that controls the
+        messages received by each node. Default: ``0``.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import APPNPConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = APPNPConv(k=3, alpha=0.5)
+    >>> res = conv(g, feat)
+    >>> print(res)
+    tensor([[0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536,
+            0.8536],
+            [0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268,
+            0.9268],
+            [0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634,
+            0.9634],
+            [0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268,
+            0.9268],
+            [0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634,
+            0.9634],
+            [0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,
+            0.5000]])
+    """
+
+    def __init__(self, k, alpha, edge_drop=0.0):
+        super(APPNPConv, self).__init__()
+        self._k = k
+        self._alpha = alpha
+        self.edge_drop = nn.Dropout(edge_drop)
+
+    def forward(self, graph, feat, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute APPNP layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, *)`. :math:`N` is the
+            number of nodes, and :math:`*` could be of any shape.
+        edge_weight: torch.Tensor, optional
+            edge_weight to use in the message passing process. This is equivalent to
+            using weighted adjacency matrix in the equation above, and
+            :math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
+            is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, *)` where :math:`*`
+            should be the same as input shape.
+        """
+        with graph.local_scope():
+            if edge_weight is None:
+                src_norm = th.pow(
+                    graph.out_degrees().to(feat).clamp(min=1), -0.5
+                )
+                shp = src_norm.shape + (1,) * (feat.dim() - 1)
+                src_norm = th.reshape(src_norm, shp).to(feat.device)
+                dst_norm = th.pow(
+                    graph.in_degrees().to(feat).clamp(min=1), -0.5
+                )
+                shp = dst_norm.shape + (1,) * (feat.dim() - 1)
+                dst_norm = th.reshape(dst_norm, shp).to(feat.device)
+            else:
+                edge_weight = EdgeWeightNorm("both")(graph, edge_weight)
+            feat_0 = feat
+            for _ in range(self._k):
+                # normalization by src node
+                if edge_weight is None:
+                    feat = feat * src_norm
+                graph.ndata["h"] = feat
+                w = (
+                    th.ones(graph.num_edges(), 1)
+                    if edge_weight is None
+                    else edge_weight
+                )
+                graph.edata["w"] = self.edge_drop(w).to(feat.device)
+                graph.update_all(fn.u_mul_e("h", "w", "m"), fn.sum("m", "h"))
+                feat = graph.ndata.pop("h")
+                # normalization by dst node
+                if edge_weight is None:
+                    feat = feat * dst_norm
+                feat = (1 - self._alpha) * feat + self._alpha * feat_0
+            return feat

python/dgl/nn/pytorch/conv/atomicconv.py

@@ -0,0 +1,301 @@
+"""Torch Module for Atomic Convolution Layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import numpy as np
+import torch as th
+import torch.nn as nn
+
+
+class RadialPooling(nn.Module):
+    r"""Radial pooling from `Atomic Convolutional Networks for
+    Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__
+
+    We denote the distance between atom :math:`i` and :math:`j` by :math:`r_{ij}`.
+
+    A radial pooling layer transforms distances with radial filters. For radial filter
+    indexed by :math:`k`, it projects edge distances with
+
+    .. math::
+        h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2)
+
+    If :math:`r_{ij} < c_k`,
+
+    .. math::
+        f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1),
+
+    else,
+
+    .. math::
+        f_{ij}^{k} = 0.
+
+    Finally,
+
+    .. math::
+        e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k}
+
+    Parameters
+    ----------
+    interaction_cutoffs : float32 tensor of shape (K)
+        :math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs
+        and two atoms are considered as connected if the distance between them is smaller than
+        the cutoffs. K for the number of radial filters.
+    rbf_kernel_means : float32 tensor of shape (K)
+        :math:`r_k` in the equations above. K for the number of radial filters.
+    rbf_kernel_scaling : float32 tensor of shape (K)
+        :math:`\gamma_k` in the equations above. K for the number of radial filters.
+    """
+
+    def __init__(
+        self, interaction_cutoffs, rbf_kernel_means, rbf_kernel_scaling
+    ):
+        super(RadialPooling, self).__init__()
+
+        self.interaction_cutoffs = nn.Parameter(
+            interaction_cutoffs.reshape(-1, 1, 1), requires_grad=True
+        )
+        self.rbf_kernel_means = nn.Parameter(
+            rbf_kernel_means.reshape(-1, 1, 1), requires_grad=True
+        )
+        self.rbf_kernel_scaling = nn.Parameter(
+            rbf_kernel_scaling.reshape(-1, 1, 1), requires_grad=True
+        )
+
+    def forward(self, distances):
+        """
+
+        Description
+        -----------
+        Apply the layer to transform edge distances.
+
+        Parameters
+        ----------
+        distances : Float32 tensor of shape (E, 1)
+            Distance between end nodes of edges. E for the number of edges.
+
+        Returns
+        -------
+        Float32 tensor of shape (K, E, 1)
+            Transformed edge distances. K for the number of radial filters.
+        """
+        scaled_euclidean_distance = (
+            -self.rbf_kernel_scaling * (distances - self.rbf_kernel_means) ** 2
+        )  # (K, E, 1)
+        rbf_kernel_results = th.exp(scaled_euclidean_distance)  # (K, E, 1)
+
+        cos_values = 0.5 * (
+            th.cos(np.pi * distances / self.interaction_cutoffs) + 1
+        )  # (K, E, 1)
+        cutoff_values = th.where(
+            distances <= self.interaction_cutoffs,
+            cos_values,
+            th.zeros_like(cos_values),
+        )  # (K, E, 1)
+
+        # Note that there appears to be an inconsistency between the paper and
+        # DeepChem's implementation. In the paper, the scaled_euclidean_distance first
+        # gets multiplied by cutoff_values, followed by exponentiation. Here we follow
+        # the practice of DeepChem.
+        return rbf_kernel_results * cutoff_values
+
+
+def msg_func(edges):
+    """
+
+    Description
+    -----------
+    Send messages along edges.
+
+    Parameters
+    ----------
+    edges : EdgeBatch
+        A batch of edges.
+
+    Returns
+    -------
+    dict mapping 'm' to Float32 tensor of shape (E, K * T)
+        Messages computed. E for the number of edges, K for the number of
+        radial filters and T for the number of features to use
+        (types of atomic number in the paper).
+    """
+    return {
+        "m": th.einsum("ij,ik->ijk", edges.src["hv"], edges.data["he"]).view(
+            len(edges), -1
+        )
+    }
+
+
+def reduce_func(nodes):
+    """
+
+    Description
+    -----------
+    Collect messages and update node representations.
+
+    Parameters
+    ----------
+    nodes : NodeBatch
+        A batch of nodes.
+
+    Returns
+    -------
+    dict mapping 'hv_new' to Float32 tensor of shape (V, K * T)
+        Updated node representations. V for the number of nodes, K for the number of
+        radial filters and T for the number of features to use
+        (types of atomic number in the paper).
+    """
+    return {"hv_new": nodes.mailbox["m"].sum(1)}
+
+
+class AtomicConv(nn.Module):
+    r"""Atomic Convolution Layer from `Atomic Convolutional Networks for
+    Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__
+
+    Denoting the type of atom :math:`i` by :math:`z_i` and the distance between atom
+    :math:`i` and :math:`j` by :math:`r_{ij}`.
+
+    **Distance Transformation**
+
+    An atomic convolution layer first transforms distances with radial filters and
+    then perform a pooling operation.
+
+    For radial filter indexed by :math:`k`, it projects edge distances with
+
+    .. math::
+        h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2)
+
+    If :math:`r_{ij} < c_k`,
+
+    .. math::
+        f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1),
+
+    else,
+
+    .. math::
+        f_{ij}^{k} = 0.
+
+    Finally,
+
+    .. math::
+        e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k}
+
+    **Aggregation**
+
+    For each type :math:`t`, each atom collects distance information from all neighbor atoms
+    of type :math:`t`:
+
+    .. math::
+        p_{i, t}^{k} = \sum_{j\in N(i)} e_{ij}^{k} * 1(z_j == t)
+
+    Then concatenate the results for all RBF kernels and atom types.
+
+    Parameters
+    ----------
+    interaction_cutoffs : float32 tensor of shape (K)
+        :math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs
+        and two atoms are considered as connected if the distance between them is smaller than
+        the cutoffs. K for the number of radial filters.
+    rbf_kernel_means : float32 tensor of shape (K)
+        :math:`r_k` in the equations above. K for the number of radial filters.
+    rbf_kernel_scaling : float32 tensor of shape (K)
+        :math:`\gamma_k` in the equations above. K for the number of radial filters.
+    features_to_use : None or float tensor of shape (T)
+        In the original paper, these are atomic numbers to consider, representing the types
+        of atoms. T for the number of types of atomic numbers. Default to None.
+
+    Note
+    ----
+
+    * This convolution operation is designed for molecular graphs in Chemistry, but it might
+      be possible to extend it to more general graphs.
+
+    * There seems to be an inconsistency about the definition of :math:`e_{ij}^{k}` in the
+      paper and the author's implementation. We follow the author's implementation. In the
+      paper, :math:`e_{ij}^{k}` was defined as
+      :math:`\exp(-\gamma_{k}|r_{ij}-r_{k}|^2 * f_{ij}^{k})`.
+
+    * :math:`\gamma_{k}`, :math:`r_k` and :math:`c_k` are all learnable.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import AtomicConv
+
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 1)
+    >>> edist = th.ones(6, 1)
+    >>> interaction_cutoffs = th.ones(3).float() * 2
+    >>> rbf_kernel_means = th.ones(3).float()
+    >>> rbf_kernel_scaling = th.ones(3).float()
+    >>> conv = AtomicConv(interaction_cutoffs, rbf_kernel_means, rbf_kernel_scaling)
+    >>> res = conv(g, feat, edist)
+    >>> res
+    tensor([[0.5000, 0.5000, 0.5000],
+                [0.5000, 0.5000, 0.5000],
+                [0.5000, 0.5000, 0.5000],
+                [1.0000, 1.0000, 1.0000],
+                [0.5000, 0.5000, 0.5000],
+                [0.0000, 0.0000, 0.0000]], grad_fn=<ViewBackward>)
+    """
+
+    def __init__(
+        self,
+        interaction_cutoffs,
+        rbf_kernel_means,
+        rbf_kernel_scaling,
+        features_to_use=None,
+    ):
+        super(AtomicConv, self).__init__()
+
+        self.radial_pooling = RadialPooling(
+            interaction_cutoffs=interaction_cutoffs,
+            rbf_kernel_means=rbf_kernel_means,
+            rbf_kernel_scaling=rbf_kernel_scaling,
+        )
+        if features_to_use is None:
+            self.num_channels = 1
+            self.features_to_use = None
+        else:
+            self.num_channels = len(features_to_use)
+            self.features_to_use = nn.Parameter(
+                features_to_use, requires_grad=False
+            )
+
+    def forward(self, graph, feat, distances):
+        """
+
+        Description
+        -----------
+        Apply the atomic convolution layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            Topology based on which message passing is performed.
+        feat : Float32 tensor of shape :math:`(V, 1)`
+            Initial node features, which are atomic numbers in the paper.
+            :math:`V` for the number of nodes.
+        distances : Float32 tensor of shape :math:`(E, 1)`
+            Distance between end nodes of edges. E for the number of edges.
+
+        Returns
+        -------
+        Float32 tensor of shape :math:`(V, K * T)`
+            Updated node representations. :math:`V` for the number of nodes, :math:`K` for the
+            number of radial filters, and :math:`T` for the number of types of atomic numbers.
+        """
+        with graph.local_scope():
+            radial_pooled_values = self.radial_pooling(distances).to(
+                feat
+            )  # (K, E, 1)
+            if self.features_to_use is not None:
+                feat = (feat == self.features_to_use).to(feat)  # (V, T)
+            graph.ndata["hv"] = feat
+            graph.edata["he"] = radial_pooled_values.transpose(1, 0).squeeze(
+                -1
+            )  # (E, K)
+            graph.update_all(msg_func, reduce_func)
+
+            return graph.ndata["hv_new"].view(
+                graph.num_nodes(), -1
+            )  # (V, K * T)

python/dgl/nn/pytorch/conv/cfconv.py

@@ -0,0 +1,145 @@
+"""Torch modules for interaction blocks in SchNet"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import numpy as np
+import torch.nn as nn
+
+from .... import function as fn
+
+
+class ShiftedSoftplus(nn.Module):
+    r"""Applies the element-wise function:
+
+    .. math::
+        \text{SSP}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x)) - \log(\text{shift})
+
+    Attributes
+    ----------
+    beta : int
+        :math:`\beta` value for the mathematical formulation. Default to 1.
+    shift : int
+        :math:`\text{shift}` value for the mathematical formulation. Default to 2.
+    """
+
+    def __init__(self, beta=1, shift=2, threshold=20):
+        super(ShiftedSoftplus, self).__init__()
+
+        self.shift = shift
+        self.softplus = nn.Softplus(beta=beta, threshold=threshold)
+
+    def forward(self, inputs):
+        """
+
+        Description
+        -----------
+        Applies the activation function.
+
+        Parameters
+        ----------
+        inputs : float32 tensor of shape (N, *)
+            * denotes any number of additional dimensions.
+
+        Returns
+        -------
+        float32 tensor of shape (N, *)
+            Result of applying the activation function to the input.
+        """
+        return self.softplus(inputs) - np.log(float(self.shift))
+
+
+class CFConv(nn.Module):
+    r"""CFConv from `SchNet: A continuous-filter convolutional neural network for
+    modeling quantum interactions <https://arxiv.org/abs/1706.08566>`__
+
+    It combines node and edge features in message passing and updates node representations.
+
+    .. math::
+        h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} h_j^{l} \circ W^{(l)}e_ij
+
+    where :math:`\circ` represents element-wise multiplication and for :math:`\text{SPP}` :
+
+    .. math::
+        \text{SSP}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x)) - \log(\text{shift})
+
+    Parameters
+    ----------
+    node_in_feats : int
+        Size for the input node features :math:`h_j^{(l)}`.
+    edge_in_feats : int
+        Size for the input edge features :math:`e_ij`.
+    hidden_feats : int
+        Size for the hidden representations.
+    out_feats : int
+        Size for the output representations :math:`h_j^{(l+1)}`.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import CFConv
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> nfeat = th.ones(6, 10)
+    >>> efeat = th.ones(6, 5)
+    >>> conv = CFConv(10, 5, 3, 2)
+    >>> res = conv(g, nfeat, efeat)
+    >>> res
+    tensor([[-0.1209, -0.2289],
+            [-0.1209, -0.2289],
+            [-0.1209, -0.2289],
+            [-0.1135, -0.2338],
+            [-0.1209, -0.2289],
+            [-0.1283, -0.2240]], grad_fn=<SubBackward0>)
+    """
+
+    def __init__(self, node_in_feats, edge_in_feats, hidden_feats, out_feats):
+        super(CFConv, self).__init__()
+
+        self.project_edge = nn.Sequential(
+            nn.Linear(edge_in_feats, hidden_feats),
+            ShiftedSoftplus(),
+            nn.Linear(hidden_feats, hidden_feats),
+            ShiftedSoftplus(),
+        )
+        self.project_node = nn.Linear(node_in_feats, hidden_feats)
+        self.project_out = nn.Sequential(
+            nn.Linear(hidden_feats, out_feats), ShiftedSoftplus()
+        )
+
+    def forward(self, g, node_feats, edge_feats):
+        """
+
+        Description
+        -----------
+        Performs message passing and updates node representations.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        node_feats : torch.Tensor or pair of torch.Tensor
+            The input node features. If a torch.Tensor is given, it represents the input
+            node feature of shape :math:`(N, D_{in})` where :math:`D_{in}` is size of
+            input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, which is the case for bipartite graph,
+            the pair must contain two tensors of shape :math:`(N_{src}, D_{in_{src}})` and
+            :math:`(N_{dst}, D_{in_{dst}})` separately for the source and destination nodes.
+
+        edge_feats : torch.Tensor
+            The input edge feature of shape :math:`(E, edge_in_feats)`
+            where :math:`E` is the number of edges.
+
+        Returns
+        -------
+        torch.Tensor
+            The output node feature of shape :math:`(N_{out}, out_feats)`
+            where :math:`N_{out}` is the number of destination nodes.
+        """
+        with g.local_scope():
+            if isinstance(node_feats, tuple):
+                node_feats_src, _ = node_feats
+            else:
+                node_feats_src = node_feats
+            g.srcdata["hv"] = self.project_node(node_feats_src)
+            g.edata["he"] = self.project_edge(edge_feats)
+            g.update_all(fn.u_mul_e("hv", "he", "m"), fn.sum("m", "h"))
+            return self.project_out(g.dstdata["h"])

python/dgl/nn/pytorch/conv/chebconv.py

@@ -0,0 +1,152 @@
+"""Torch Module for Chebyshev Spectral Graph Convolution layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+import torch.nn.functional as F
+from torch import nn
+
+from .... import broadcast_nodes, function as fn
+from ....base import dgl_warning
+
+
+class ChebConv(nn.Module):
+    r"""Chebyshev Spectral Graph Convolution layer from `Convolutional
+    Neural Networks on Graphs with Fast Localized Spectral Filtering
+    <https://arxiv.org/pdf/1606.09375.pdf>`__
+
+    .. math::
+        h_i^{l+1} &= \sum_{k=0}^{K-1} W^{k, l}z_i^{k, l}
+
+        Z^{0, l} &= H^{l}
+
+        Z^{1, l} &= \tilde{L} \cdot H^{l}
+
+        Z^{k, l} &= 2 \cdot \tilde{L} \cdot Z^{k-1, l} - Z^{k-2, l}
+
+        \tilde{L} &= 2\left(I - \tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2}\right)/\lambda_{max} - I
+
+    where :math:`\tilde{A}` is :math:`A` + :math:`I`, :math:`W` is learnable weight.
+
+
+    Parameters
+    ----------
+    in_feats: int
+        Dimension of input features; i.e, the number of dimensions of :math:`h_i^{(l)}`.
+    out_feats: int
+        Dimension of output features :math:`h_i^{(l+1)}`.
+    k : int
+        Chebyshev filter size :math:`K`.
+    activation : function, optional
+        Activation function. Default ``ReLu``.
+    bias : bool, optional
+        If True, adds a learnable bias to the output. Default: ``True``.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import ChebConv
+    >>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = ChebConv(10, 2, 2)
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[ 0.6163, -0.1809],
+            [ 0.6163, -0.1809],
+            [ 0.6163, -0.1809],
+            [ 0.9698, -1.5053],
+            [ 0.3664,  0.7556],
+            [-0.2370,  3.0164]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(self, in_feats, out_feats, k, activation=F.relu, bias=True):
+        super(ChebConv, self).__init__()
+        self._k = k
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self.activation = activation
+        self.linear = nn.Linear(k * in_feats, out_feats, bias)
+
+    def forward(self, graph, feat, lambda_max=None):
+        r"""Compute ChebNet layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
+            is size of input feature, :math:`N` is the number of nodes.
+        lambda_max : list or tensor or None, optional.
+            A list(tensor) with length :math:`B`, stores the largest eigenvalue
+            of the normalized laplacian of each individual graph in ``graph``,
+            where :math:`B` is the batch size of the input graph. Default: None.
+
+            If None, this method would set the default value to 2.
+            One can use :func:`dgl.laplacian_lambda_max` to compute this value.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is size of output feature.
+        """
+
+        def unnLaplacian(feat, D_invsqrt, graph):
+            """Operation Feat * D^-1/2 A D^-1/2"""
+            graph.ndata["h"] = feat * D_invsqrt
+            graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
+            return graph.ndata.pop("h") * D_invsqrt
+
+        with graph.local_scope():
+            D_invsqrt = th.pow(
+                graph.in_degrees().to(feat).clamp(min=1), -0.5
+            ).unsqueeze(-1)
+
+            if lambda_max is None:
+                dgl_warning(
+                    "lambda_max is not provided, using default value of 2.  "
+                    "Please use dgl.laplacian_lambda_max to compute the eigenvalues."
+                )
+                lambda_max = [2] * graph.batch_size
+
+            if isinstance(lambda_max, list):
+                lambda_max = th.Tensor(lambda_max).to(feat)
+            if lambda_max.dim() == 1:
+                lambda_max = lambda_max.unsqueeze(-1)  # (B,) to (B, 1)
+
+            # broadcast from (B, 1) to (N, 1)
+            lambda_max = broadcast_nodes(graph, lambda_max)
+            re_norm = 2.0 / lambda_max
+
+            # X_0 is the raw feature, Xt is the list of X_0, X_1, ... X_t
+            X_0 = feat
+            Xt = [X_0]
+
+            # X_1(f)
+            if self._k > 1:
+                h = unnLaplacian(X_0, D_invsqrt, graph)
+                X_1 = -re_norm * h + X_0 * (re_norm - 1)
+                # Append X_1 to Xt
+                Xt.append(X_1)
+
+            # Xi(x), i = 2...k
+            for _ in range(2, self._k):
+                h = unnLaplacian(X_1, D_invsqrt, graph)
+                X_i = -2 * re_norm * h + X_1 * 2 * (re_norm - 1) - X_0
+                # Add X_1 to Xt
+                Xt.append(X_i)
+                X_1, X_0 = X_i, X_1
+
+            # Create the concatenation
+            Xt = th.cat(Xt, dim=1)
+
+            # linear projection
+            h = self.linear(Xt)
+
+            # activation
+            if self.activation:
+                h = self.activation(h)
+
+        return h

python/dgl/nn/pytorch/conv/cugraph_base.py

@@ -0,0 +1,57 @@
+"""An abstract base class for cugraph-ops nn module."""
+import torch
+from torch import nn
+
+
+class CuGraphBaseConv(nn.Module):
+    r"""An abstract base class for cugraph-ops nn module."""
+
+    def __init__(self):
+        super().__init__()
+        self._cached_offsets_fg = None
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        raise NotImplementedError
+
+    def forward(self, *args):
+        r"""Runs the forward pass of the module."""
+        raise NotImplementedError
+
+    def pad_offsets(self, offsets: torch.Tensor, size: int) -> torch.Tensor:
+        r"""Pad zero-in-degree nodes to the end of offsets to reach size.
+
+        cugraph-ops often provides two variants of aggregation functions for a
+        specific model: one intended for sampled-graph use cases, one for
+        full-graph ones. The former is in general more performant, however, it
+        only works when the sample size (the max of in-degrees) is small (<200),
+        due to the limit of GPU shared memory. For graphs with a larger max
+        in-degree, we need to fall back to the full-graph option, which requires
+        to convert a DGL block to a full graph. With the csc-representation,
+        this is equivalent to pad zero-in-degree nodes to the end of the offsets
+        array (also called indptr or colptr).
+
+        Parameters
+        ----------
+        offsets :
+            The (monotonically increasing) index pointer array in a CSC-format
+            graph.
+        size : int
+            The length of offsets after padding.
+
+        Returns
+        -------
+        torch.Tensor
+            The augmented offsets array.
+        """
+        if self._cached_offsets_fg is None:
+            self._cached_offsets_fg = torch.empty(
+                size, dtype=offsets.dtype, device=offsets.device
+            )
+        elif self._cached_offsets_fg.numel() < size:
+            self._cached_offsets_fg.resize_(size)
+
+        self._cached_offsets_fg[: offsets.numel()] = offsets
+        self._cached_offsets_fg[offsets.numel() : size] = offsets[-1]
+
+        return self._cached_offsets_fg[:size]

python/dgl/nn/pytorch/conv/cugraph_gatconv.py

@@ -0,0 +1,213 @@
+"""Torch Module for graph attention network layer using the aggregation
+primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+
+import torch
+from torch import nn
+
+from .cugraph_base import CuGraphBaseConv
+
+try:
+    from pylibcugraphops.pytorch import SampledCSC, StaticCSC
+    from pylibcugraphops.pytorch.operators import mha_gat_n2n as GATConvAgg
+
+    HAS_PYLIBCUGRAPHOPS = True
+except ImportError:
+    HAS_PYLIBCUGRAPHOPS = False
+
+
+class CuGraphGATConv(CuGraphBaseConv):
+    r"""Graph attention layer from `Graph Attention Networks
+    <https://arxiv.org/pdf/1710.10903.pdf>`__, with the sparse aggregation
+    accelerated by cugraph-ops.
+
+    See :class:`dgl.nn.pytorch.conv.GATConv` for mathematical model.
+
+    This module depends on :code:`pylibcugraphops` package, which can be
+    installed via :code:`conda install -c nvidia pylibcugraphops=23.04`.
+    :code:`pylibcugraphops` 23.04 requires python 3.8.x or 3.10.x.
+
+    .. note::
+        This is an **experimental** feature.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size.
+    out_feats : int
+        Output feature size.
+    num_heads : int
+        Number of heads in Multi-Head Attention.
+    feat_drop : float, optional
+        Dropout rate on feature. Defaults: ``0``.
+    negative_slope : float, optional
+        LeakyReLU angle of negative slope. Defaults: ``0.2``.
+    residual : bool, optional
+        If True, use residual connection. Defaults: ``False``.
+    activation : callable activation function/layer or None, optional.
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+    bias : bool, optional
+        If True, learns a bias term. Defaults: ``True``.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import torch
+    >>> from dgl.nn import CuGraphGATConv
+    >>> device = 'cuda'
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = torch.ones(6, 10).to(device)
+    >>> conv = CuGraphGATConv(10, 2, num_heads=3).to(device)
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[[ 0.2340,  1.9226],
+            [ 1.6477, -1.9986],
+            [ 1.1138, -1.9302]],
+            [[ 0.2340,  1.9226],
+            [ 1.6477, -1.9986],
+            [ 1.1138, -1.9302]],
+            [[ 0.2340,  1.9226],
+            [ 1.6477, -1.9986],
+            [ 1.1138, -1.9302]],
+            [[ 0.2340,  1.9226],
+            [ 1.6477, -1.9986],
+            [ 1.1138, -1.9302]],
+            [[ 0.2340,  1.9226],
+            [ 1.6477, -1.9986],
+            [ 1.1138, -1.9302]],
+            [[ 0.2340,  1.9226],
+            [ 1.6477, -1.9986],
+            [ 1.1138, -1.9302]]], device='cuda:0', grad_fn=<ViewBackward0>)
+    """
+    MAX_IN_DEGREE_MFG = 200
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        num_heads,
+        feat_drop=0.0,
+        negative_slope=0.2,
+        residual=False,
+        activation=None,
+        bias=True,
+    ):
+        if HAS_PYLIBCUGRAPHOPS is False:
+            raise ModuleNotFoundError(
+                f"{self.__class__.__name__} requires pylibcugraphops=23.04. "
+                f"Install via `conda install -c nvidia 'pylibcugraphops=23.04'`."
+                f"pylibcugraphops requires Python 3.8 or 3.10."
+            )
+        super().__init__()
+        self.in_feats = in_feats
+        self.out_feats = out_feats
+        self.num_heads = num_heads
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.negative_slope = negative_slope
+        self.activation = activation
+
+        self.fc = nn.Linear(in_feats, out_feats * num_heads, bias=False)
+        self.attn_weights = nn.Parameter(
+            torch.Tensor(2 * num_heads * out_feats)
+        )
+
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(num_heads * out_feats))
+        else:
+            self.register_buffer("bias", None)
+
+        if residual:
+            if in_feats == out_feats * num_heads:
+                self.res_fc = nn.Identity()
+            else:
+                self.res_fc = nn.Linear(
+                    in_feats, out_feats * num_heads, bias=False
+                )
+        else:
+            self.register_buffer("res_fc", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Reinitialize learnable parameters."""
+
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_normal_(self.fc.weight, gain=gain)
+        nn.init.xavier_normal_(
+            self.attn_weights.view(2, self.num_heads, self.out_feats), gain=gain
+        )
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+        if isinstance(self.res_fc, nn.Linear):
+            self.res_fc.reset_parameters()
+
+    def forward(self, g, feat, max_in_degree=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            Input features of shape :math:`(N, D_{in})`.
+        max_in_degree : int
+            Maximum in-degree of destination nodes. It is only effective when
+            :attr:`g` is a :class:`DGLBlock`, i.e., bipartite graph. When
+            :attr:`g` is generated from a neighbor sampler, the value should be
+            set to the corresponding :attr:`fanout`. If not given,
+            :attr:`max_in_degree` will be calculated on-the-fly.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, H, D_{out})` where
+            :math:`H` is the number of heads, and :math:`D_{out}` is size of
+            output feature.
+        """
+        offsets, indices, _ = g.adj_tensors("csc")
+
+        if g.is_block:
+            if max_in_degree is None:
+                max_in_degree = g.in_degrees().max().item()
+
+            if max_in_degree < self.MAX_IN_DEGREE_MFG:
+                _graph = SampledCSC(
+                    offsets,
+                    indices,
+                    max_in_degree,
+                    g.num_src_nodes(),
+                )
+            else:
+                offsets_fg = self.pad_offsets(offsets, g.num_src_nodes() + 1)
+                _graph = StaticCSC(offsets_fg, indices)
+        else:
+            _graph = StaticCSC(offsets, indices)
+
+        feat = self.feat_drop(feat)
+        feat_transformed = self.fc(feat)
+        out = GATConvAgg(
+            feat_transformed,
+            self.attn_weights,
+            _graph,
+            self.num_heads,
+            "LeakyReLU",
+            self.negative_slope,
+            concat_heads=True,
+        )[: g.num_dst_nodes()].view(-1, self.num_heads, self.out_feats)
+
+        feat_dst = feat[: g.num_dst_nodes()]
+        if self.res_fc is not None:
+            out = out + self.res_fc(feat_dst).view(
+                -1, self.num_heads, self.out_feats
+            )
+
+        if self.bias is not None:
+            out = out + self.bias.view(-1, self.num_heads, self.out_feats)
+
+        if self.activation is not None:
+            out = self.activation(out)
+
+        return out

python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py

@@ -0,0 +1,228 @@
+"""Torch Module for Relational graph convolution layer using the aggregation
+primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch
+from torch import nn
+
+from .cugraph_base import CuGraphBaseConv
+
+try:
+    from pylibcugraphops.pytorch import HeteroCSC
+    from pylibcugraphops.pytorch.operators import (
+        agg_hg_basis_n2n_post as RelGraphConvAgg,
+    )
+
+    HAS_PYLIBCUGRAPHOPS = True
+except ImportError:
+    HAS_PYLIBCUGRAPHOPS = False
+
+
+class CuGraphRelGraphConv(CuGraphBaseConv):
+    r"""An accelerated relational graph convolution layer from `Modeling
+    Relational Data with Graph Convolutional Networks
+    <https://arxiv.org/abs/1703.06103>`__ that leverages the highly-optimized
+    aggregation primitives in cugraph-ops.
+
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    This module depends on :code:`pylibcugraphops` package, which can be
+    installed via :code:`conda install -c nvidia pylibcugraphops=23.04`.
+    :code:`pylibcugraphops` 23.04 requires python 3.8.x or 3.10.x.
+
+    .. note::
+        This is an **experimental** feature.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. It comes into effect when a regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``.
+    apply_norm : bool, optional
+        True to normalize aggregation output by the in-degree of the destination
+        node per edge type, i.e. :math:`|\mathcal{N}^r_i|`. Default: ``True``.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import torch
+    >>> from dgl.nn import CuGraphRelGraphConv
+    ...
+    >>> device = 'cuda'
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
+    >>> feat = torch.ones(6, 10).to(device)
+    >>> conv = CuGraphRelGraphConv(
+    ...     10, 2, 3, regularizer='basis', num_bases=2).to(device)
+    >>> etype = torch.tensor([0,1,2,0,1,2]).to(device)
+    >>> res = conv(g, feat, etype)
+    >>> res
+    tensor([[-1.7774, -2.0184],
+            [-1.4335, -2.3758],
+            [-1.7774, -2.0184],
+            [-0.4698, -3.0876],
+            [-1.4335, -2.3758],
+            [-1.4331, -2.3295]], device='cuda:0', grad_fn=<AddBackward0>)
+    """
+    MAX_IN_DEGREE_MFG = 500
+
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        self_loop=True,
+        dropout=0.0,
+        apply_norm=False,
+    ):
+        if HAS_PYLIBCUGRAPHOPS is False:
+            raise ModuleNotFoundError(
+                f"{self.__class__.__name__} requires pylibcugraphops=23.04. "
+                f"Install via `conda install -c nvidia 'pylibcugraphops=23.04'`."
+                f"pylibcugraphops requires Python 3.8 or 3.10."
+            )
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.apply_norm = apply_norm
+        self.dropout = nn.Dropout(dropout)
+
+        dim_self_loop = 1 if self_loop else 0
+        self.self_loop = self_loop
+        if regularizer is None:
+            self.W = nn.Parameter(
+                torch.Tensor(num_rels + dim_self_loop, in_feat, out_feat)
+            )
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(
+                torch.Tensor(num_bases + dim_self_loop, in_feat, out_feat)
+            )
+            self.coeff = nn.Parameter(torch.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis' or None, but got "
+                f"'{regularizer}'."
+            )
+        self.regularizer = regularizer
+
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_feat))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Reinitialize learnable parameters."""
+        bound = 1 / math.sqrt(self.in_feat)
+        end = -1 if self.self_loop else None
+        nn.init.uniform_(self.W[:end], -bound, bound)
+        if self.regularizer == "basis":
+            nn.init.xavier_uniform_(
+                self.coeff, gain=nn.init.calculate_gain("relu")
+            )
+        if self.self_loop:
+            nn.init.xavier_uniform_(self.W[-1], nn.init.calculate_gain("relu"))
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def forward(self, g, feat, etypes, max_in_degree=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor
+            A 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
+            Note that cugraph-ops only accepts edge type tensors in int32,
+            so any input of other integer types will be casted into int32,
+            thus introducing some overhead. Pass in int32 tensors directly
+            for best performance.
+        max_in_degree : int, optional
+            Maximum in-degree of destination nodes. It is only effective when
+            :attr:`g` is a :class:`DGLBlock`, i.e., bipartite graph. When
+            :attr:`g` is generated from a neighbor sampler, the value should be
+            set to the corresponding :attr:`fanout`. If not given,
+            :attr:`max_in_degree` will be calculated on-the-fly.
+
+        Returns
+        -------
+        torch.Tensor
+            New node features. Shape: :math:`(|V|, D_{out})`.
+        """
+        offsets, indices, edge_ids = g.adj_tensors("csc")
+        edge_types_perm = etypes[edge_ids.long()].int()
+
+        if g.is_block:
+            if max_in_degree is None:
+                max_in_degree = g.in_degrees().max().item()
+
+            if max_in_degree < self.MAX_IN_DEGREE_MFG:
+                _graph = HeteroCSC(
+                    offsets,
+                    indices,
+                    edge_types_perm,
+                    g.num_src_nodes(),
+                    self.num_rels,
+                )
+            else:
+                offsets_fg = self.pad_offsets(offsets, g.num_src_nodes() + 1)
+                _graph = HeteroCSC(
+                    offsets_fg,
+                    indices,
+                    edge_types_perm,
+                    g.num_src_nodes(),
+                    self.num_rels,
+                )
+        else:
+            _graph = HeteroCSC(
+                offsets,
+                indices,
+                edge_types_perm,
+                g.num_src_nodes(),
+                self.num_rels,
+            )
+
+        h = RelGraphConvAgg(
+            feat,
+            self.coeff,
+            _graph,
+            concat_own=self.self_loop,
+            norm_by_out_degree=self.apply_norm,
+        )[: g.num_dst_nodes()]
+        h = h @ self.W.view(-1, self.out_feat)
+        if self.bias is not None:
+            h = h + self.bias
+        h = self.dropout(h)
+
+        return h

python/dgl/nn/pytorch/conv/cugraph_sageconv.py

@@ -0,0 +1,148 @@
+"""Torch Module for GraphSAGE layer using the aggregation primitives in
+cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+
+from torch import nn
+
+from .cugraph_base import CuGraphBaseConv
+
+try:
+    from pylibcugraphops.pytorch import SampledCSC, StaticCSC
+    from pylibcugraphops.pytorch.operators import agg_concat_n2n as SAGEConvAgg
+
+    HAS_PYLIBCUGRAPHOPS = True
+except ImportError:
+    HAS_PYLIBCUGRAPHOPS = False
+
+
+class CuGraphSAGEConv(CuGraphBaseConv):
+    r"""An accelerated GraphSAGE layer from `Inductive Representation Learning
+    on Large Graphs <https://arxiv.org/pdf/1706.02216.pdf>`__ that leverages the
+    highly-optimized aggregation primitives in cugraph-ops:
+
+    .. math::
+        h_{\mathcal{N}(i)}^{(l+1)} &= \mathrm{aggregate}
+        \left(\{h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)
+
+        h_{i}^{(l+1)} &= W \cdot \mathrm{concat}
+        (h_{i}^{l}, h_{\mathcal{N}(i)}^{(l+1)})
+
+    This module depends on :code:`pylibcugraphops` package, which can be
+    installed via :code:`conda install -c nvidia pylibcugraphops=23.04`.
+    :code:`pylibcugraphops` 23.04 requires python 3.8.x or 3.10.x.
+
+    .. note::
+        This is an **experimental** feature.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size.
+    out_feats : int
+        Output feature size.
+    aggregator_type : str
+        Aggregator type to use (``mean``, ``sum``, ``min``, ``max``).
+    feat_drop : float
+        Dropout rate on features, default: ``0``.
+    bias : bool
+        If True, adds a learnable bias to the output. Default: ``True``.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import torch
+    >>> from dgl.nn import CuGraphSAGEConv
+    >>> device = 'cuda'
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = torch.ones(6, 10).to(device)
+    >>> conv = CuGraphSAGEConv(10, 2, 'mean').to(device)
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[-1.1690,  0.1952],
+            [-1.1690,  0.1952],
+            [-1.1690,  0.1952],
+            [-1.1690,  0.1952],
+            [-1.1690,  0.1952],
+            [-1.1690,  0.1952]], device='cuda:0', grad_fn=<AddmmBackward0>)
+    """
+    MAX_IN_DEGREE_MFG = 500
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        aggregator_type="mean",
+        feat_drop=0.0,
+        bias=True,
+    ):
+        if HAS_PYLIBCUGRAPHOPS is False:
+            raise ModuleNotFoundError(
+                f"{self.__class__.__name__} requires pylibcugraphops=23.04. "
+                f"Install via `conda install -c nvidia 'pylibcugraphops=23.04'`."
+                f"pylibcugraphops requires Python 3.8 or 3.10."
+            )
+
+        valid_aggr_types = {"max", "min", "mean", "sum"}
+        if aggregator_type not in valid_aggr_types:
+            raise ValueError(
+                f"Invalid aggregator_type. Must be one of {valid_aggr_types}. "
+                f"But got '{aggregator_type}' instead."
+            )
+
+        super().__init__()
+        self.in_feats = in_feats
+        self.out_feats = out_feats
+        self.aggr = aggregator_type
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.linear = nn.Linear(2 * in_feats, out_feats, bias=bias)
+
+    def reset_parameters(self):
+        r"""Reinitialize learnable parameters."""
+        self.linear.reset_parameters()
+
+    def forward(self, g, feat, max_in_degree=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            Node features. Shape: :math:`(N, D_{in})`.
+        max_in_degree : int
+            Maximum in-degree of destination nodes. It is only effective when
+            :attr:`g` is a :class:`DGLBlock`, i.e., bipartite graph. When
+            :attr:`g` is generated from a neighbor sampler, the value should be
+            set to the corresponding :attr:`fanout`. If not given,
+            :attr:`max_in_degree` will be calculated on-the-fly.
+
+        Returns
+        -------
+        torch.Tensor
+            Output node features. Shape: :math:`(N, D_{out})`.
+        """
+        offsets, indices, _ = g.adj_tensors("csc")
+
+        if g.is_block:
+            if max_in_degree is None:
+                max_in_degree = g.in_degrees().max().item()
+
+            if max_in_degree < self.MAX_IN_DEGREE_MFG:
+                _graph = SampledCSC(
+                    offsets,
+                    indices,
+                    max_in_degree,
+                    g.num_src_nodes(),
+                )
+            else:
+                offsets_fg = self.pad_offsets(offsets, g.num_src_nodes() + 1)
+                _graph = StaticCSC(offsets_fg, indices)
+        else:
+            _graph = StaticCSC(offsets, indices)
+
+        feat = self.feat_drop(feat)
+        h = SAGEConvAgg(feat, _graph, self.aggr)[: g.num_dst_nodes()]
+        h = self.linear(h)
+
+        return h

python/dgl/nn/pytorch/conv/densechebconv.py

@@ -0,0 +1,125 @@
+"""Torch Module for DenseChebConv"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+from torch.nn import init
+
+
+class DenseChebConv(nn.Module):
+    r"""Chebyshev Spectral Graph Convolution layer from `Convolutional
+    Neural Networks on Graphs with Fast Localized Spectral Filtering
+    <https://arxiv.org/pdf/1606.09375.pdf>`__
+
+    We recommend to use this module when applying ChebConv on dense graphs.
+
+    Parameters
+    ----------
+    in_feats: int
+        Dimension of input features :math:`h_i^{(l)}`.
+    out_feats: int
+        Dimension of output features :math:`h_i^{(l+1)}`.
+    k : int
+        Chebyshev filter size.
+    activation : function, optional
+        Activation function, default is ReLu.
+    bias : bool, optional
+        If True, adds a learnable bias to the output. Default: ``True``.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import DenseChebConv
+    >>>
+    >>> feat = th.ones(6, 10)
+    >>> adj = th.tensor([[0., 0., 1., 0., 0., 0.],
+    ...         [1., 0., 0., 0., 0., 0.],
+    ...         [0., 1., 0., 0., 0., 0.],
+    ...         [0., 0., 1., 0., 0., 1.],
+    ...         [0., 0., 0., 1., 0., 0.],
+    ...         [0., 0., 0., 0., 0., 0.]])
+    >>> conv = DenseChebConv(10, 2, 2)
+    >>> res = conv(adj, feat)
+    >>> res
+    tensor([[-3.3516, -2.4797],
+            [-3.3516, -2.4797],
+            [-3.3516, -2.4797],
+            [-4.5192, -3.0835],
+            [-2.5259, -2.0527],
+            [-0.5327, -1.0219]], grad_fn=<AddBackward0>)
+
+    See also
+    --------
+    `ChebConv <https://docs.dgl.ai/api/python/nn.pytorch.html#chebconv>`__
+    """
+
+    def __init__(self, in_feats, out_feats, k, bias=True):
+        super(DenseChebConv, self).__init__()
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self._k = k
+        self.W = nn.Parameter(th.Tensor(k, in_feats, out_feats))
+        if bias:
+            self.bias = nn.Parameter(th.Tensor(out_feats))
+        else:
+            self.register_buffer("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """Reinitialize learnable parameters."""
+        if self.bias is not None:
+            init.zeros_(self.bias)
+        for i in range(self._k):
+            init.xavier_normal_(self.W[i], init.calculate_gain("relu"))
+
+    def forward(self, adj, feat, lambda_max=None):
+        r"""Compute (Dense) Chebyshev Spectral Graph Convolution layer
+
+        Parameters
+        ----------
+        adj : torch.Tensor
+            The adjacency matrix of the graph to apply Graph Convolution on,
+            should be of shape :math:`(N, N)`, where a row represents the destination
+            and a column represents the source.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
+            is size of input feature, :math:`N` is the number of nodes.
+        lambda_max : float or None, optional
+            A float value indicates the largest eigenvalue of given graph.
+            Default: None.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is size of output feature.
+        """
+        A = adj.to(feat)
+        num_nodes = A.shape[0]
+
+        in_degree = 1 / A.sum(dim=1).clamp(min=1).sqrt()
+        D_invsqrt = th.diag(in_degree)
+        I = th.eye(num_nodes).to(A)
+        L = I - D_invsqrt @ A @ D_invsqrt
+
+        if lambda_max is None:
+            lambda_ = th.eig(L)[0][:, 0]
+            lambda_max = lambda_.max()
+
+        L_hat = 2 * L / lambda_max - I
+        Z = [th.eye(num_nodes).to(A)]
+        for i in range(1, self._k):
+            if i == 1:
+                Z.append(L_hat)
+            else:
+                Z.append(2 * L_hat @ Z[-1] - Z[-2])
+
+        Zs = th.stack(Z, 0)  # (k, n, n)
+
+        Zh = Zs @ feat.unsqueeze(0) @ self.W
+        Zh = Zh.sum(0)
+
+        if self.bias is not None:
+            Zh = Zh + self.bias
+        return Zh

python/dgl/nn/pytorch/conv/densegraphconv.py

@@ -0,0 +1,145 @@
+"""Torch Module for DenseGraphConv"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+from torch.nn import init
+
+
+class DenseGraphConv(nn.Module):
+    """Graph Convolutional layer from `Semi-Supervised Classification with Graph
+    Convolutional Networks <https://arxiv.org/abs/1609.02907>`__
+
+    We recommend user to use this module when applying graph convolution on
+    dense graphs.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
+    out_feats : int
+        Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
+    norm : str, optional
+        How to apply the normalizer. If is `'right'`, divide the aggregated messages
+        by each node's in-degrees, which is equivalent to averaging the received messages.
+        If is `'none'`, no normalization is applied. Default is `'both'`,
+        where the :math:`c_{ij}` in the paper is applied.
+    bias : bool, optional
+        If True, adds a learnable bias to the output. Default: ``True``.
+    activation : callable activation function/layer or None, optional
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+
+    Notes
+    -----
+    Zero in-degree nodes will lead to all-zero output. A common practice
+    to avoid this is to add a self-loop for each node in the graph,
+    which can be achieved by setting the diagonal of the adjacency matrix to be 1.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import DenseGraphConv
+    >>>
+    >>> feat = th.ones(6, 10)
+    >>> adj = th.tensor([[0., 0., 1., 0., 0., 0.],
+    ...         [1., 0., 0., 0., 0., 0.],
+    ...         [0., 1., 0., 0., 0., 0.],
+    ...         [0., 0., 1., 0., 0., 1.],
+    ...         [0., 0., 0., 1., 0., 0.],
+    ...         [0., 0., 0., 0., 0., 0.]])
+    >>> conv = DenseGraphConv(10, 2)
+    >>> res = conv(adj, feat)
+    >>> res
+    tensor([[0.2159, 1.9027],
+            [0.3053, 2.6908],
+            [0.3053, 2.6908],
+            [0.3685, 3.2481],
+            [0.3053, 2.6908],
+            [0.0000, 0.0000]], grad_fn=<AddBackward0>)
+
+    See also
+    --------
+    `GraphConv <https://docs.dgl.ai/api/python/nn.pytorch.html#graphconv>`__
+    """
+
+    def __init__(
+        self, in_feats, out_feats, norm="both", bias=True, activation=None
+    ):
+        super(DenseGraphConv, self).__init__()
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self._norm = norm
+        self.weight = nn.Parameter(th.Tensor(in_feats, out_feats))
+        if bias:
+            self.bias = nn.Parameter(th.Tensor(out_feats))
+        else:
+            self.register_buffer("bias", None)
+
+        self.reset_parameters()
+        self._activation = activation
+
+    def reset_parameters(self):
+        """Reinitialize learnable parameters."""
+        init.xavier_uniform_(self.weight)
+        if self.bias is not None:
+            init.zeros_(self.bias)
+
+    def forward(self, adj, feat):
+        r"""Compute (Dense) Graph Convolution layer.
+
+        Parameters
+        ----------
+        adj : torch.Tensor
+            The adjacency matrix of the graph to apply Graph Convolution on, when
+            applied to a unidirectional bipartite graph, ``adj`` should be of shape
+            should be of shape :math:`(N_{out}, N_{in})`; when applied to a homo
+            graph, ``adj`` should be of shape :math:`(N, N)`. In both cases,
+            a row represents a destination node while a column represents a source
+            node.
+        feat : torch.Tensor
+            The input feature.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is size of output feature.
+        """
+        adj = adj.to(feat)
+        src_degrees = adj.sum(dim=0).clamp(min=1)
+        dst_degrees = adj.sum(dim=1).clamp(min=1)
+        feat_src = feat
+
+        if self._norm == "both":
+            norm_src = th.pow(src_degrees, -0.5)
+            shp = norm_src.shape + (1,) * (feat.dim() - 1)
+            norm_src = th.reshape(norm_src, shp).to(feat.device)
+            feat_src = feat_src * norm_src
+
+        if self._in_feats > self._out_feats:
+            # mult W first to reduce the feature size for aggregation.
+            feat_src = th.matmul(feat_src, self.weight)
+            rst = adj @ feat_src
+        else:
+            # aggregate first then mult W
+            rst = adj @ feat_src
+            rst = th.matmul(rst, self.weight)
+
+        if self._norm != "none":
+            if self._norm == "both":
+                norm_dst = th.pow(dst_degrees, -0.5)
+            else:  # right
+                norm_dst = 1.0 / dst_degrees
+            shp = norm_dst.shape + (1,) * (feat.dim() - 1)
+            norm_dst = th.reshape(norm_dst, shp).to(feat.device)
+            rst = rst * norm_dst
+
+        if self.bias is not None:
+            rst = rst + self.bias
+
+        if self._activation is not None:
+            rst = self._activation(rst)
+
+        return rst

python/dgl/nn/pytorch/conv/densesageconv.py

@@ -0,0 +1,138 @@
+"""Torch Module for DenseSAGEConv"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+from torch import nn
+
+from ....utils import check_eq_shape
+
+
+class DenseSAGEConv(nn.Module):
+    """GraphSAGE layer from `Inductive Representation Learning on Large Graphs
+    <https://arxiv.org/abs/1706.02216>`__
+
+    We recommend to use this module when appying GraphSAGE on dense graphs.
+
+    Note that we only support gcn aggregator in DenseSAGEConv.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
+    out_feats : int
+        Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
+    feat_drop : float, optional
+        Dropout rate on features. Default: 0.
+    bias : bool
+        If True, adds a learnable bias to the output. Default: ``True``.
+    norm : callable activation function/layer or None, optional
+        If not None, applies normalization to the updated node features.
+    activation : callable activation function/layer or None, optional
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import DenseSAGEConv
+    >>>
+    >>> feat = th.ones(6, 10)
+    >>> adj = th.tensor([[0., 0., 1., 0., 0., 0.],
+    ...         [1., 0., 0., 0., 0., 0.],
+    ...         [0., 1., 0., 0., 0., 0.],
+    ...         [0., 0., 1., 0., 0., 1.],
+    ...         [0., 0., 0., 1., 0., 0.],
+    ...         [0., 0., 0., 0., 0., 0.]])
+    >>> conv = DenseSAGEConv(10, 2)
+    >>> res = conv(adj, feat)
+    >>> res
+    tensor([[1.0401, 2.1008],
+            [1.0401, 2.1008],
+            [1.0401, 2.1008],
+            [1.0401, 2.1008],
+            [1.0401, 2.1008],
+            [1.0401, 2.1008]], grad_fn=<AddmmBackward>)
+
+    See also
+    --------
+    `SAGEConv <https://docs.dgl.ai/api/python/nn.pytorch.html#sageconv>`__
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        feat_drop=0.0,
+        bias=True,
+        norm=None,
+        activation=None,
+    ):
+        super(DenseSAGEConv, self).__init__()
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self._norm = norm
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.activation = activation
+        self.fc = nn.Linear(in_feats, out_feats, bias=bias)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Notes
+        -----
+        The linear weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
+        """
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_uniform_(self.fc.weight, gain=gain)
+
+    def forward(self, adj, feat):
+        r"""
+
+        Description
+        -----------
+        Compute (Dense) Graph SAGE layer.
+
+        Parameters
+        ----------
+        adj : torch.Tensor
+            The adjacency matrix of the graph to apply SAGE Convolution on, when
+            applied to a unidirectional bipartite graph, ``adj`` should be of shape
+            should be of shape :math:`(N_{out}, N_{in})`; when applied to a homo
+            graph, ``adj`` should be of shape :math:`(N, N)`. In both cases,
+            a row represents a destination node while a column represents a source
+            node.
+        feat : torch.Tensor or a pair of torch.Tensor
+            If a torch.Tensor is given, the input feature of shape :math:`(N, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, D_{in})` and :math:`(N_{out}, D_{in})`.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is size of output feature.
+        """
+        check_eq_shape(feat)
+        if isinstance(feat, tuple):
+            feat_src = self.feat_drop(feat[0])
+            feat_dst = self.feat_drop(feat[1])
+        else:
+            feat_src = feat_dst = self.feat_drop(feat)
+        adj = adj.to(feat_src)
+        in_degrees = adj.sum(dim=1, keepdim=True)
+        h_neigh = (adj @ feat_src + feat_dst) / (in_degrees + 1)
+        rst = self.fc(h_neigh)
+        # activation
+        if self.activation is not None:
+            rst = self.activation(rst)
+        # normalization
+        if self._norm is not None:
+            rst = self._norm(rst)
+
+        return rst

python/dgl/nn/pytorch/conv/dgnconv.py

@@ -0,0 +1,265 @@
+"""Torch Module for Directional Graph Networks Convolution Layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from .pnaconv import AGGREGATORS, PNAConv, PNAConvTower, SCALERS
+
+
+def aggregate_dir_av(h, eig_s, eig_d, eig_idx):
+    """directional average aggregation"""
+    h_mod = torch.mul(
+        h,
+        (
+            torch.abs(eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx])
+            / (
+                torch.sum(
+                    torch.abs(eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx]),
+                    keepdim=True,
+                    dim=1,
+                )
+                + 1e-30
+            )
+        ).unsqueeze(-1),
+    )
+    return torch.sum(h_mod, dim=1)
+
+
+def aggregate_dir_dx(h, eig_s, eig_d, h_in, eig_idx):
+    """directional derivative aggregation"""
+    eig_w = (
+        (eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx])
+        / (
+            torch.sum(
+                torch.abs(eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx]),
+                keepdim=True,
+                dim=1,
+            )
+            + 1e-30
+        )
+    ).unsqueeze(-1)
+    h_mod = torch.mul(h, eig_w)
+    return torch.abs(torch.sum(h_mod, dim=1) - torch.sum(eig_w, dim=1) * h_in)
+
+
+for k in range(1, 4):
+    AGGREGATORS[f"dir{k}-av"] = partial(aggregate_dir_av, eig_idx=k - 1)
+    AGGREGATORS[f"dir{k}-dx"] = partial(aggregate_dir_dx, eig_idx=k - 1)
+
+
+class DGNConvTower(PNAConvTower):
+    """A single DGN tower with modified reduce function"""
+
+    def message(self, edges):
+        """message function for DGN layer"""
+        if self.edge_feat_size > 0:
+            f = torch.cat(
+                [edges.src["h"], edges.dst["h"], edges.data["a"]], dim=-1
+            )
+        else:
+            f = torch.cat([edges.src["h"], edges.dst["h"]], dim=-1)
+        return {
+            "msg": self.M(f),
+            "eig_s": edges.src["eig"],
+            "eig_d": edges.dst["eig"],
+        }
+
+    def reduce_func(self, nodes):
+        """reduce function for DGN layer"""
+        h_in = nodes.data["h"]
+        eig_s = nodes.mailbox["eig_s"]
+        eig_d = nodes.mailbox["eig_d"]
+        msg = nodes.mailbox["msg"]
+        degree = msg.size(1)
+
+        h = []
+        for agg in self.aggregators:
+            if agg.startswith("dir"):
+                if agg.endswith("av"):
+                    h.append(AGGREGATORS[agg](msg, eig_s, eig_d))
+                else:
+                    h.append(AGGREGATORS[agg](msg, eig_s, eig_d, h_in))
+            else:
+                h.append(AGGREGATORS[agg](msg))
+        h = torch.cat(h, dim=1)
+        h = torch.cat(
+            [
+                SCALERS[scaler](h, D=degree, delta=self.delta)
+                if scaler != "identity"
+                else h
+                for scaler in self.scalers
+            ],
+            dim=1,
+        )
+        return {"h_neigh": h}
+
+
+class DGNConv(PNAConv):
+    r"""Directional Graph Network Layer from `Directional Graph Networks
+    <https://arxiv.org/abs/2010.02863>`__
+
+    DGN introduces two special directional aggregators according to the vector field
+    :math:`F`, which is defined as the gradient of the low-frequency eigenvectors of graph
+    laplacian.
+
+    The directional average aggregator is defined as
+    :math:`h_i' = \sum_{j\in\mathcal{N}(i)}\frac{|F_{i,j}|\cdot h_j}{||F_{i,:}||_1+\epsilon}`
+
+    The directional derivative aggregator is defined as
+    :math:`h_i' = \sum_{j\in\mathcal{N}(i)}\frac{F_{i,j}\cdot h_j}{||F_{i,:}||_1+\epsilon}
+    -h_i\cdot\sum_{j\in\mathcal{N}(i)}\frac{F_{i,j}}{||F_{i,:}||_1+\epsilon}`
+
+    :math:`\epsilon` is the infinitesimal to keep the computation numerically stable.
+
+    Parameters
+    ----------
+    in_size : int
+        Input feature size; i.e. the size of :math:`h_i^l`.
+    out_size : int
+        Output feature size; i.e. the size of :math:`h_i^{l+1}`.
+    aggregators : list of str
+        List of aggregation function names(each aggregator specifies a way to aggregate
+        messages from neighbours), selected from:
+
+        * ``mean``: the mean of neighbour messages
+
+        * ``max``: the maximum of neighbour messages
+
+        * ``min``: the minimum of neighbour messages
+
+        * ``std``: the standard deviation of neighbour messages
+
+        * ``var``: the variance of neighbour messages
+
+        * ``sum``: the sum of neighbour messages
+
+        * ``moment3``, ``moment4``, ``moment5``: the normalized moments aggregation
+        :math:`(E[(X-E[X])^n])^{1/n}`
+
+        * ``dir{k}-av``: directional average aggregation with directions defined by the k-th
+        smallest eigenvectors. k can be selected from 1, 2, 3.
+
+        * ``dir{k}-dx``: directional derivative aggregation with directions defined by the k-th
+        smallest eigenvectors. k can be selected from 1, 2, 3.
+
+        Note that using directional aggregation requires the LaplacianPE transform on the input
+        graph for eigenvector computation (the PE size must be >= k above).
+    scalers: list of str
+        List of scaler function names, selected from:
+
+        * ``identity``: no scaling
+
+        * ``amplification``: multiply the aggregated message by :math:`\log(d+1)/\delta`,
+        where :math:`d` is the in-degree of the node.
+
+        * ``attenuation``: multiply the aggregated message by :math:`\delta/\log(d+1)`
+    delta: float
+        The in-degree-related normalization factor computed over the training set, used by scalers
+        for normalization. :math:`E[\log(d+1)]`, where :math:`d` is the in-degree for each node
+        in the training set.
+    dropout: float, optional
+        The dropout ratio. Default: 0.0.
+    num_towers: int, optional
+        The number of towers used. Default: 1. Note that in_size and out_size must be divisible
+        by num_towers.
+    edge_feat_size: int, optional
+        The edge feature size. Default: 0.
+    residual : bool, optional
+        The bool flag that determines whether to add a residual connection for the
+        output. Default: True. If in_size and out_size of the DGN conv layer are not
+        the same, this flag will be set as False forcibly.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import torch as th
+    >>> from dgl.nn import DGNConv
+    >>> from dgl import LaplacianPE
+    >>>
+    >>> # DGN requires precomputed eigenvectors, with 'eig' as feature name.
+    >>> transform = LaplacianPE(k=3, feat_name='eig')
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = transform(g)
+    >>> eig = g.ndata['eig']
+    >>> feat = th.ones(6, 10)
+    >>> conv = DGNConv(10, 10, ['dir1-av', 'dir1-dx', 'sum'], ['identity', 'amplification'], 2.5)
+    >>> ret = conv(g, feat, eig_vec=eig)
+    """
+
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        aggregators,
+        scalers,
+        delta,
+        dropout=0.0,
+        num_towers=1,
+        edge_feat_size=0,
+        residual=True,
+    ):
+        super(DGNConv, self).__init__(
+            in_size,
+            out_size,
+            aggregators,
+            scalers,
+            delta,
+            dropout,
+            num_towers,
+            edge_feat_size,
+            residual,
+        )
+
+        self.towers = nn.ModuleList(
+            [
+                DGNConvTower(
+                    self.tower_in_size,
+                    self.tower_out_size,
+                    aggregators,
+                    scalers,
+                    delta,
+                    dropout=dropout,
+                    edge_feat_size=edge_feat_size,
+                )
+                for _ in range(num_towers)
+            ]
+        )
+
+        self.use_eig_vec = False
+        for aggr in aggregators:
+            if aggr.startswith("dir"):
+                self.use_eig_vec = True
+                break
+
+    def forward(self, graph, node_feat, edge_feat=None, eig_vec=None):
+        r"""
+        Description
+        -----------
+        Compute DGN layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        node_feat : torch.Tensor
+            The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
+            nodes, and :math:`h_n` must be the same as in_size.
+        edge_feat : torch.Tensor, optional
+            The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
+            edges, and :math:`h_e` must be the same as edge_feat_size.
+        eig_vec : torch.Tensor, optional
+            K smallest non-trivial eigenvectors of Graph Laplacian of shape :math:`(N, K)`.
+            It is only required when :attr:`aggregators` contains directional aggregators.
+
+        Returns
+        -------
+        torch.Tensor
+            The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
+            should be the same as out_size.
+        """
+        with graph.local_scope():
+            if self.use_eig_vec:
+                graph.ndata["eig"] = eig_vec
+            return super().forward(graph, node_feat, edge_feat)

python/dgl/nn/pytorch/conv/dotgatconv.py

@@ -0,0 +1,242 @@
+"""Torch modules for graph attention networks(GAT)."""
+# pylint: disable= no-member, arguments-differ, invalid-name
+from torch import nn
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import expand_as_pair
+from ...functional import edge_softmax
+
+
+class DotGatConv(nn.Module):
+    r"""Apply dot product version of self attention in `Graph Attention Network
+    <https://arxiv.org/pdf/1710.10903.pdf>`__
+
+        .. math::
+            h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{i, j} h_j^{(l)}
+
+        where :math:`\alpha_{ij}` is the attention score bewteen node :math:`i` and node :math:`j`:
+
+        .. math::
+            \alpha_{i, j} &= \mathrm{softmax_i}(e_{ij}^{l})
+
+            e_{ij}^{l} &= ({W_i^{(l)} h_i^{(l)}})^T \cdot {W_j^{(l)} h_j^{(l)}}
+
+        where :math:`W_i` and :math:`W_j` transform node :math:`i`'s and node :math:`j`'s
+        features into the same dimension, so that when compute note features' similarity,
+        it can use dot-product.
+
+    Parameters
+    ----------
+    in_feats : int, or pair of ints
+        Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
+        DotGatConv can be applied on homogeneous graph and unidirectional
+        `bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
+        If the layer is to be applied to a unidirectional bipartite graph, ``in_feats``
+        specifies the input feature size on both the source and destination nodes.  If
+        a scalar is given, the source and destination node feature size would take the
+        same value.
+    out_feats : int
+        Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
+    num_heads : int
+        Number of head in Multi-Head Attention
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Default: ``False``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import DotGatConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> dotgatconv = DotGatConv(10, 2, num_heads=3)
+    >>> res = dotgatconv(g, feat)
+    >>> res
+    tensor([[[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]]], grad_fn=<BinaryReduceBackward>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
+    >>> u_feat = th.tensor(np.random.rand(2, 5).astype(np.float32))
+    >>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
+    >>> dotgatconv = DotGatConv((5,10), 2, 3)
+    >>> res = dotgatconv(g, (u_feat, v_feat))
+    >>> res
+    tensor([[[-0.6066,  1.0268],
+            [-0.5945, -0.4801],
+            [ 0.1594,  0.3825]],
+            [[ 0.0268,  1.0783],
+            [ 0.5041, -1.3025],
+            [ 0.6568,  0.7048]],
+            [[-0.2688,  1.0543],
+            [-0.0315, -0.9016],
+            [ 0.3943,  0.5347]],
+            [[-0.6066,  1.0268],
+            [-0.5945, -0.4801],
+            [ 0.1594,  0.3825]]], grad_fn=<BinaryReduceBackward>)
+    """
+
+    def __init__(
+        self, in_feats, out_feats, num_heads, allow_zero_in_degree=False
+    ):
+        super(DotGatConv, self).__init__()
+        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
+        self._out_feats = out_feats
+        self._allow_zero_in_degree = allow_zero_in_degree
+        self._num_heads = num_heads
+
+        if isinstance(in_feats, tuple):
+            self.fc_src = nn.Linear(
+                self._in_src_feats,
+                self._out_feats * self._num_heads,
+                bias=False,
+            )
+            self.fc_dst = nn.Linear(
+                self._in_dst_feats,
+                self._out_feats * self._num_heads,
+                bias=False,
+            )
+        else:
+            self.fc = nn.Linear(
+                self._in_src_feats,
+                self._out_feats * self._num_heads,
+                bias=False,
+            )
+
+    def forward(self, graph, feat, get_attention=False):
+        r"""
+
+        Description
+        -----------
+        Apply dot product version of self attention in GCN.
+
+        Parameters
+        ----------
+        graph: DGLGraph or bi_partities graph
+            The graph
+        feat: torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, the input feature of shape :math:`(N, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
+        get_attention : bool, optional
+            Whether to return the attention values. Default to False.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}` is size
+            of output feature.
+        torch.Tensor, optional
+            The attention values of shape :math:`(E, 1)`, where :math:`E` is the number of
+            edges. This is returned only when :attr:`get_attention` is ``True``.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+        """
+
+        graph = graph.local_var()
+
+        if not self._allow_zero_in_degree:
+            if (graph.in_degrees() == 0).any():
+                raise DGLError(
+                    "There are 0-in-degree nodes in the graph, "
+                    "output for those nodes will be invalid. "
+                    "This is harmful for some applications, "
+                    "causing silent performance regression. "
+                    "Adding self-loop on the input graph by "
+                    "calling `g = dgl.add_self_loop(g)` will resolve "
+                    "the issue. Setting ``allow_zero_in_degree`` "
+                    "to be `True` when constructing this module will "
+                    "suppress the check and let the code run."
+                )
+
+        # check if feat is a tuple
+        if isinstance(feat, tuple):
+            h_src = feat[0]
+            h_dst = feat[1]
+            feat_src = self.fc_src(h_src).view(
+                -1, self._num_heads, self._out_feats
+            )
+            feat_dst = self.fc_dst(h_dst).view(
+                -1, self._num_heads, self._out_feats
+            )
+        else:
+            h_src = feat
+            feat_src = feat_dst = self.fc(h_src).view(
+                -1, self._num_heads, self._out_feats
+            )
+            if graph.is_block:
+                feat_dst = feat_src[: graph.number_of_dst_nodes()]
+
+        # Assign features to nodes
+        graph.srcdata.update({"ft": feat_src})
+        graph.dstdata.update({"ft": feat_dst})
+
+        # Step 1. dot product
+        graph.apply_edges(fn.u_dot_v("ft", "ft", "a"))
+
+        # Step 2. edge softmax to compute attention scores
+        graph.edata["sa"] = edge_softmax(
+            graph, graph.edata["a"] / self._out_feats**0.5
+        )
+
+        # Step 3. Broadcast softmax value to each edge, and aggregate dst node
+        graph.update_all(
+            fn.u_mul_e("ft", "sa", "attn"), fn.sum("attn", "agg_u")
+        )
+
+        # output results to the destination nodes
+        rst = graph.dstdata["agg_u"]
+
+        if get_attention:
+            return rst, graph.edata["sa"]
+        else:
+            return rst

python/dgl/nn/pytorch/conv/edgeconv.py

@@ -0,0 +1,201 @@
+"""Torch Module for EdgeConv Layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+from torch import nn
+
+from .... import function as fn
+
+from ....base import DGLError
+from ....utils import expand_as_pair
+
+
+class EdgeConv(nn.Module):
+    r"""EdgeConv layer from `Dynamic Graph CNN for Learning on Point Clouds
+    <https://arxiv.org/pdf/1801.07829>`__
+
+    It can be described as follows:
+
+    .. math::
+       h_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} (
+       \Theta \cdot (h_j^{(l)} - h_i^{(l)}) + \Phi \cdot h_i^{(l)})
+
+    where :math:`\mathcal{N}(i)` is the neighbor of :math:`i`.
+    :math:`\Theta` and :math:`\Phi` are linear layers.
+
+    .. note::
+
+       The original formulation includes a ReLU inside the maximum operator.
+       This is equivalent to first applying a maximum operator then applying
+       the ReLU.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
+    out_feat : int
+        Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
+    batch_norm : bool
+        Whether to include batch normalization on messages. Default: ``False``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Default: ``False``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import EdgeConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> conv = EdgeConv(10, 2)
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[-0.2347,  0.5849],
+            [-0.2347,  0.5849],
+            [-0.2347,  0.5849],
+            [-0.2347,  0.5849],
+            [-0.2347,  0.5849],
+            [-0.2347,  0.5849]], grad_fn=<CopyReduceBackward>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
+    >>> u_fea = th.rand(2, 5)
+    >>> v_fea = th.rand(4, 5)
+    >>> conv = EdgeConv(5, 2, 3)
+    >>> res = conv(g, (u_fea, v_fea))
+    >>> res
+    tensor([[ 1.6375,  0.2085],
+            [-1.1925, -1.2852],
+            [ 0.2101,  1.3466],
+            [ 0.2342, -0.9868]], grad_fn=<CopyReduceBackward>)
+    """
+
+    def __init__(
+        self, in_feat, out_feat, batch_norm=False, allow_zero_in_degree=False
+    ):
+        super(EdgeConv, self).__init__()
+        self.batch_norm = batch_norm
+        self._allow_zero_in_degree = allow_zero_in_degree
+
+        self.theta = nn.Linear(in_feat, out_feat)
+        self.phi = nn.Linear(in_feat, out_feat)
+
+        if batch_norm:
+            self.bn = nn.BatchNorm1d(out_feat)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, g, feat):
+        """
+
+        Description
+        -----------
+        Forward computation
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : Tensor or pair of tensors
+            :math:`(N, D)` where :math:`N` is the number of nodes and
+            :math:`D` is the number of feature dimensions.
+
+            If a pair of tensors is given, the graph must be a uni-bipartite graph
+            with only one edge type, and the two tensors must have the same
+            dimensionality on all except the first axis.
+
+        Returns
+        -------
+        torch.Tensor
+            New node features.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+        """
+        with g.local_scope():
+            if not self._allow_zero_in_degree:
+                if (g.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            h_src, h_dst = expand_as_pair(feat, g)
+            g.srcdata["x"] = h_src
+            g.dstdata["x"] = h_dst
+            g.apply_edges(fn.v_sub_u("x", "x", "theta"))
+            g.edata["theta"] = self.theta(g.edata["theta"])
+            g.dstdata["phi"] = self.phi(g.dstdata["x"])
+            if not self.batch_norm:
+                g.update_all(fn.e_add_v("theta", "phi", "e"), fn.max("e", "x"))
+            else:
+                g.apply_edges(fn.e_add_v("theta", "phi", "e"))
+                # Although the official implementation includes a per-edge
+                # batch norm within EdgeConv, I choose to replace it with a
+                # global batch norm for a number of reasons:
+                #
+                # (1) When the point clouds within each batch do not have the
+                #     same number of points, batch norm would not work.
+                #
+                # (2) Even if the point clouds always have the same number of
+                #     points, the points may as well be shuffled even with the
+                #     same (type of) object (and the official implementation
+                #     *does* shuffle the points of the same example for each
+                #     epoch).
+                #
+                #     For example, the first point of a point cloud of an
+                #     airplane does not always necessarily reside at its nose.
+                #
+                #     In this case, the learned statistics of each position
+                #     by batch norm is not as meaningful as those learned from
+                #     images.
+                g.edata["e"] = self.bn(g.edata["e"])
+                g.update_all(fn.copy_e("e", "e"), fn.max("e", "x"))
+            return g.dstdata["x"]

python/dgl/nn/pytorch/conv/edgegatconv.py

@@ -0,0 +1,390 @@
+"""Torch modules for graph attention networks(GAT)."""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import expand_as_pair
+from ...functional import edge_softmax
+
+# pylint: enable=W0235
+class EdgeGATConv(nn.Module):
+    r"""Graph attention layer with edge features from `SCENE
+    <https://arxiv.org/pdf/2301.03512.pdf>`__
+
+    .. math::
+
+        \mathbf{v}_i^\prime = \mathbf{\Theta}_\mathrm{s} \cdot \mathbf{v}_i +
+        \sum\limits_{j \in \mathcal{N}(v_i)} \alpha_{j, i} \left( \mathbf{\Theta}_\mathrm{n}
+        \cdot \mathbf{v}_j + \mathbf{\Theta}_\mathrm{e} \cdot \mathbf{e}_{j,i} \right)
+
+    where :math:`\mathbf{\Theta}` is used to denote learnable weight matrices
+    for the transformation of features of the node to update (s=self),
+    neighboring nodes (n=neighbor) and edge features (e=edge).
+    Attention weights are obtained by
+
+    .. math::
+
+        \alpha_{j, i} = \mathrm{softmax}_i \Big( \mathrm{LeakyReLU} \big( \mathbf{a}^T
+        [ \mathbf{\Theta}_\mathrm{n} \cdot \mathbf{v}_i || \mathbf{\Theta}_\mathrm{n}
+        \cdot \mathbf{v}_j || \mathbf{\Theta}_\mathrm{e} \cdot \mathbf{e}_{j,i} ] \big) \Big)
+
+    with :math:`\mathbf{a}` corresponding to a learnable vector.
+    :math:`\mathrm{softmax_i}` stands for the normalization by all incoming edges of node :math:`i`.
+
+    Parameters
+    ----------
+    in_feats : int, or pair of ints
+        Input feature size; i.e, the number of dimensions of :math:`\mathbf{v}_i`.
+        GATConv can be applied on homogeneous graph and unidirectional
+        `bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
+        If the layer is to be applied to a unidirectional bipartite graph, ``in_feats``
+        specifies the input feature size on both the source and destination nodes.  If
+        a scalar is given, the source and destination node feature size would take the
+        same value.
+    edge_feats: int
+        Edge feature size; i.e., the number of dimensions of :math:\mathbf{e}_{j,i}`.
+    out_feats : int
+        Output feature size; i.e, the number of dimensions of :math:`\mathbf{v}_i^\prime`.
+    num_heads : int
+        Number of heads in Multi-Head Attention.
+    feat_drop : float, optional
+        Dropout rate on feature. Defaults: ``0``.
+    attn_drop : float, optional
+        Dropout rate on attention weight. Defaults: ``0``.
+    negative_slope : float, optional
+        LeakyReLU angle of negative slope. Defaults: ``0.2``.
+    residual : bool, optional
+        If True, use residual connection. Defaults: ``False``.
+    activation : callable activation function/layer or None, optional.
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Defaults: ``False``.
+    bias : bool, optional
+        If True, learns a bias term. Defaults: ``True``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    ----------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import EdgeGATConv
+
+    >>> # Case 1: Homogeneous graph.
+    >>> num_nodes, num_edges = 8, 30
+    >>> # Generate a graph.
+    >>> graph = dgl.rand_graph(num_nodes,num_edges)
+    >>> node_feats = th.rand((num_nodes, 20))
+    >>> edge_feats = th.rand((num_edges, 12))
+    >>> edge_gat = EdgeGATConv(
+    ...     in_feats=20,
+    ...     edge_feats=12,
+    ...     out_feats=15,
+    ...     num_heads=3,
+    ... )
+    >>> # Forward pass.
+    >>> new_node_feats = edge_gat(graph, node_feats, edge_feats)
+    >>> new_node_feats.shape
+    torch.Size([8, 3, 15]) torch.Size([30, 3, 10])
+
+    >>> # Case 2: Unidirectional bipartite graph.
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('A', 'r', 'B'): (u, v)})
+    >>> u_feat = th.tensor(np.random.rand(2, 25).astype(np.float32))
+    >>> v_feat = th.tensor(np.random.rand(4, 30).astype(np.float32))
+    >>> nfeats = (u_feat,v_feat)
+    >>> efeats = th.tensor(np.random.rand(5, 15).astype(np.float32))
+    >>> in_feats = (25,30)
+    >>> edge_feats = 15
+    >>> out_feats = 10
+    >>> num_heads = 3
+    >>> egat_model =  EdgeGATConv(
+    ...     in_feats,
+    ...     edge_feats,
+    ...     out_feats,
+    ...     num_heads,
+    ... )
+    >>> # Forward pass.
+    >>> new_node_feats, attention_weights = egat_model(g, nfeats, efeats, get_attention=True)
+    >>> new_node_feats.shape, attention_weights.shape
+    (torch.Size([4, 3, 10]), torch.Size([5, 3, 1]))
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        edge_feats,
+        out_feats,
+        num_heads,
+        feat_drop=0.0,
+        attn_drop=0.0,
+        negative_slope=0.2,
+        residual=True,
+        activation=None,
+        allow_zero_in_degree=False,
+        bias=True,
+    ):
+        super(EdgeGATConv, self).__init__()
+        self._num_heads = num_heads
+        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
+        self._out_feats = out_feats
+        self._allow_zero_in_degree = allow_zero_in_degree
+        if isinstance(in_feats, tuple):
+            self.fc_src = nn.Linear(
+                self._in_src_feats, out_feats * num_heads, bias=False
+            )
+            self.fc_dst = nn.Linear(
+                self._in_dst_feats, out_feats * num_heads, bias=False
+            )
+        else:
+            self.fc = nn.Linear(
+                self._in_src_feats, out_feats * num_heads, bias=False
+            )
+        self.attn_l = nn.Parameter(
+            th.FloatTensor(size=(1, num_heads, out_feats))
+        )
+        self.attn_r = nn.Parameter(
+            th.FloatTensor(size=(1, num_heads, out_feats))
+        )
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.leaky_relu = nn.LeakyReLU(negative_slope)
+        if bias:
+            self.bias = nn.Parameter(
+                th.FloatTensor(size=(num_heads * out_feats,))
+            )
+        else:
+            self.register_buffer("bias", None)
+        if residual:
+            self.res_fc = nn.Linear(
+                self._in_dst_feats, num_heads * out_feats, bias=False
+            )
+        else:
+            self.register_buffer("res_fc", None)
+
+        self._edge_feats = edge_feats
+        self.fc_edge = nn.Linear(edge_feats, out_feats * num_heads, bias=False)
+        self.attn_edge = nn.Parameter(
+            th.FloatTensor(size=(1, num_heads, out_feats))
+        )
+
+        self.reset_parameters()
+        self.activation = activation
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The fc weights :math:`\mathbf{\Theta}` are and the
+        attention weights are using xavier initialization method.
+        """
+        gain = nn.init.calculate_gain("relu")
+        if hasattr(self, "fc"):
+            nn.init.xavier_normal_(self.fc.weight, gain=gain)
+        else:
+            nn.init.xavier_normal_(self.fc_src.weight, gain=gain)
+            nn.init.xavier_normal_(self.fc_dst.weight, gain=gain)
+        nn.init.xavier_normal_(self.attn_l, gain=gain)
+        nn.init.xavier_normal_(self.attn_r, gain=gain)
+
+        nn.init.xavier_normal_(self.fc_edge.weight, gain=gain)
+        nn.init.xavier_normal_(self.attn_edge, gain=gain)
+        if self.bias is not None:
+            nn.init.constant_(self.bias, 0)
+        if isinstance(self.res_fc, nn.Linear):
+            nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, edge_feat, get_attention=False):
+        r"""
+
+        Description
+        -----------
+        Compute graph attention network layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, the input feature of shape :math:`(N, *, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, *, D_{in_{src}})` and :math:`(N_{out}, *, D_{in_{dst}})`.
+        edge_feat : torch.Tensor
+            The input edge feature of shape :math:`(E, D_{in_{edge}})`,
+            where :math:`E` is the number of edges and :math:`D_{in_{edge}}`
+            the size of the edge features.
+        get_attention : bool, optional
+            Whether to return the attention values. Default to False.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, *, H, D_{out})` where :math:`H`
+            is the number of heads, and :math:`D_{out}` is size of output feature.
+        torch.Tensor, optional
+            The attention values of shape :math:`(E, *, H, 1)`. This is returned only
+            when :attr:`get_attention` is ``True``.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            if isinstance(feat, tuple):
+                src_prefix_shape = feat[0].shape[:-1]
+                dst_prefix_shape = feat[1].shape[:-1]
+                h_src = self.feat_drop(feat[0])
+                h_dst = self.feat_drop(feat[1])
+                if not hasattr(self, "fc_src"):
+                    feat_src = self.fc(h_src).view(
+                        *src_prefix_shape, self._num_heads, self._out_feats
+                    )
+                    feat_dst = self.fc(h_dst).view(
+                        *dst_prefix_shape, self._num_heads, self._out_feats
+                    )
+                else:
+                    feat_src = self.fc_src(h_src).view(
+                        *src_prefix_shape, self._num_heads, self._out_feats
+                    )
+                    feat_dst = self.fc_dst(h_dst).view(
+                        *dst_prefix_shape, self._num_heads, self._out_feats
+                    )
+            else:
+                src_prefix_shape = dst_prefix_shape = feat.shape[:-1]
+                h_src = h_dst = self.feat_drop(feat)
+                feat_src = feat_dst = self.fc(h_src).view(
+                    *src_prefix_shape, self._num_heads, self._out_feats
+                )
+                if graph.is_block:
+                    feat_dst = feat_src[: graph.number_of_dst_nodes()]
+                    h_dst = h_dst[: graph.number_of_dst_nodes()]
+                    dst_prefix_shape = (
+                        graph.number_of_dst_nodes(),
+                    ) + dst_prefix_shape[1:]
+
+            # Linearly tranform the edge features.
+            n_edges = edge_feat.shape[:-1]
+            feat_edge = self.fc_edge(edge_feat).view(
+                *n_edges, self._num_heads, self._out_feats
+            )
+
+            # Add edge features to graph.
+            graph.edata["ft_edge"] = feat_edge
+
+            el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(-1)
+            er = (feat_dst * self.attn_r).sum(dim=-1).unsqueeze(-1)
+
+            # Calculate scalar for each edge.
+            ee = (feat_edge * self.attn_edge).sum(dim=-1).unsqueeze(-1)
+            graph.edata["ee"] = ee
+
+            graph.srcdata.update({"ft": feat_src, "el": el})
+            graph.dstdata.update({"er": er})
+            # Compute edge attention, el and er are a_l Wh_i and a_r Wh_j respectively.
+            graph.apply_edges(fn.u_add_v("el", "er", "e_tmp"))
+
+            # e_tmp combines attention weights of source and destination node.
+            # Add the attention weight of the edge.
+            graph.edata["e"] = graph.edata["e_tmp"] + graph.edata["ee"]
+
+            # Create new edges features that combine the
+            # features of the source node and the edge features.
+            graph.apply_edges(fn.u_add_e("ft", "ft_edge", "ft_combined"))
+
+            e = self.leaky_relu(graph.edata.pop("e"))
+            # Compute softmax.
+            graph.edata["a"] = self.attn_drop(edge_softmax(graph, e))
+
+            # For each edge, element-wise multiply the combined features with
+            # the attention coefficient.
+            graph.edata["m_combined"] = (
+                graph.edata["ft_combined"] * graph.edata["a"]
+            )
+
+            # First copy the edge features and then sum them up.
+            graph.update_all(fn.copy_e("m_combined", "m"), fn.sum("m", "ft"))
+
+            rst = graph.dstdata["ft"]
+            # Residual.
+            if self.res_fc is not None:
+                # Use -1 rather than self._num_heads to handle broadcasting.
+                if h_dst.numel() != 0:
+                    resval = self.res_fc(h_dst).view(
+                        *dst_prefix_shape, -1, self._out_feats
+                    )
+                    rst = rst + resval
+            # Bias.
+            if self.bias is not None:
+                rst = rst + self.bias.view(
+                    *((1,) * len(dst_prefix_shape)),
+                    self._num_heads,
+                    self._out_feats
+                )
+            # Activation.
+            if self.activation:
+                rst = self.activation(rst)
+
+            if get_attention:
+                return rst, graph.edata["a"]
+            else:
+                return rst

python/dgl/nn/pytorch/conv/egatconv.py

@@ -0,0 +1,260 @@
+"""Torch modules for graph attention networks with fully valuable edges (EGAT)."""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+from torch.nn import init
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import expand_as_pair
+from ...functional import edge_softmax
+
+
+# pylint: enable=W0235
+class EGATConv(nn.Module):
+    r"""Graph attention layer that handles edge features from `Rossmann-Toolbox
+    <https://pubmed.ncbi.nlm.nih.gov/34571541/>`__ (see supplementary data)
+
+    The difference lies in how unnormalized attention scores :math:`e_{ij}` are obtained:
+
+    .. math::
+        e_{ij} &= \vec{F} (f_{ij}^{\prime})
+
+        f_{ij}^{\prime} &= \mathrm{LeakyReLU}\left(A [ h_{i} \| f_{ij} \| h_{j}]\right)
+
+    where :math:`f_{ij}^{\prime}` are edge features, :math:`\mathrm{A}` is weight matrix and
+    :math:`\vec{F}` is weight vector. After that, resulting node features
+    :math:`h_{i}^{\prime}` are updated in the same way as in regular GAT.
+
+    Parameters
+    ----------
+    in_node_feats : int, or pair of ints
+        Input feature size; i.e, the number of dimensions of :math:`h_{i}`.
+        EGATConv can be applied on homogeneous graph and unidirectional
+        `bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
+        If the layer is to be applied to a unidirectional bipartite graph, ``in_feats``
+        specifies the input feature size on both the source and destination nodes.  If
+        a scalar is given, the source and destination node feature size would take the
+        same value.
+    in_edge_feats : int
+        Input edge feature size :math:`f_{ij}`.
+    out_node_feats : int
+        Output node feature size.
+    out_edge_feats : int
+        Output edge feature size :math:`f_{ij}^{\prime}`.
+    num_heads : int
+        Number of attention heads.
+    bias : bool, optional
+        If True, add bias term to :math:`f_{ij}^{\prime}`. Defaults: ``True``.
+
+    Examples
+    ----------
+    >>> import dgl
+    >>> import torch as th
+    >>> from dgl.nn import EGATConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> num_nodes, num_edges = 8, 30
+    >>> # generate a graph
+    >>> graph = dgl.rand_graph(num_nodes,num_edges)
+    >>> node_feats = th.rand((num_nodes, 20))
+    >>> edge_feats = th.rand((num_edges, 12))
+    >>> egat = EGATConv(in_node_feats=20,
+    ...                 in_edge_feats=12,
+    ...                 out_node_feats=15,
+    ...                 out_edge_feats=10,
+    ...                 num_heads=3)
+    >>> #forward pass
+    >>> new_node_feats, new_edge_feats = egat(graph, node_feats, edge_feats)
+    >>> new_node_feats.shape, new_edge_feats.shape
+    torch.Size([8, 3, 15]) torch.Size([30, 3, 10])
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('A', 'r', 'B'): (u, v)})
+    >>> u_feat = th.tensor(np.random.rand(2, 25).astype(np.float32))
+    >>> v_feat = th.tensor(np.random.rand(4, 30).astype(np.float32))
+    >>> nfeats = (u_feat,v_feat)
+    >>> efeats = th.tensor(np.random.rand(5, 15).astype(np.float32))
+    >>> in_node_feats = (25,30)
+    >>> in_edge_feats = 15
+    >>> out_node_feats = 10
+    >>> out_edge_feats = 5
+    >>> num_heads = 3
+    >>> egat_model =  EGATConv(in_node_feats,
+    ...                        in_edge_feats,
+    ...                        out_node_feats,
+    ...                        out_edge_feats,
+    ...                        num_heads,
+    ...                        bias=True)
+    >>> #forward pass
+    >>> new_node_feats,
+    >>> new_edge_feats,
+    >>> attentions = egat_model(g, nfeats, efeats, get_attention=True)
+    >>> new_node_feats.shape, new_edge_feats.shape, attentions.shape
+    (torch.Size([4, 3, 10]), torch.Size([5, 3, 5]), torch.Size([5, 3, 1]))
+    """
+
+    def __init__(
+        self,
+        in_node_feats,
+        in_edge_feats,
+        out_node_feats,
+        out_edge_feats,
+        num_heads,
+        bias=True,
+    ):
+        super().__init__()
+        self._num_heads = num_heads
+        self._in_src_node_feats, self._in_dst_node_feats = expand_as_pair(
+            in_node_feats
+        )
+        self._out_node_feats = out_node_feats
+        self._out_edge_feats = out_edge_feats
+        if isinstance(in_node_feats, tuple):
+            self.fc_node_src = nn.Linear(
+                self._in_src_node_feats, out_node_feats * num_heads, bias=False
+            )
+            self.fc_ni = nn.Linear(
+                self._in_src_node_feats, out_edge_feats * num_heads, bias=False
+            )
+            self.fc_nj = nn.Linear(
+                self._in_dst_node_feats, out_edge_feats * num_heads, bias=False
+            )
+        else:
+            self.fc_node_src = nn.Linear(
+                self._in_src_node_feats, out_node_feats * num_heads, bias=False
+            )
+            self.fc_ni = nn.Linear(
+                self._in_src_node_feats, out_edge_feats * num_heads, bias=False
+            )
+            self.fc_nj = nn.Linear(
+                self._in_src_node_feats, out_edge_feats * num_heads, bias=False
+            )
+
+        self.fc_fij = nn.Linear(
+            in_edge_feats, out_edge_feats * num_heads, bias=False
+        )
+        self.attn = nn.Parameter(
+            th.FloatTensor(size=(1, num_heads, out_edge_feats))
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                th.FloatTensor(size=(num_heads * out_edge_feats,))
+            )
+        else:
+            self.register_buffer("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """
+        Reinitialize learnable parameters.
+        """
+        gain = init.calculate_gain("relu")
+        init.xavier_normal_(self.fc_node_src.weight, gain=gain)
+        init.xavier_normal_(self.fc_ni.weight, gain=gain)
+        init.xavier_normal_(self.fc_fij.weight, gain=gain)
+        init.xavier_normal_(self.fc_nj.weight, gain=gain)
+        init.xavier_normal_(self.attn, gain=gain)
+        init.constant_(self.bias, 0)
+
+    def forward(
+        self, graph, nfeats, efeats, edge_weight=None, get_attention=False
+    ):
+        r"""
+        Compute new node and edge features.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        nfeat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, the input feature of shape :math:`(N, D_{in})`
+            where:
+                :math:`D_{in}` is size of input node feature,
+                :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+                :math:`(N_{in}, D_{in_{src}})` and
+                :math:`(N_{out}, D_{in_{dst}})`.
+        efeats: torch.Tensor
+             The input edge feature of shape :math:`(E, F_{in})`
+             where:
+                 :math:`F_{in}` is size of input node feature,
+                 :math:`E` is the number of edges.
+        edge_weight : torch.Tensor, optional
+            A 1D tensor of edge weight values.  Shape: :math:`(|E|,)`.
+        get_attention : bool, optional
+                Whether to return the attention values. Default to False.
+
+        Returns
+        -------
+        pair of torch.Tensor
+            node output features followed by edge output features.
+            The node output feature is of shape :math:`(N, H, D_{out})`
+            The edge output feature is of shape :math:`(F, H, F_{out})`
+            where:
+                :math:`H` is the number of heads,
+                :math:`D_{out}` is size of output node feature,
+                :math:`F_{out}` is size of output edge feature.
+        torch.Tensor, optional
+            The attention values of shape :math:`(E, H, 1)`.
+            This is returned only when :attr:`get_attention` is ``True``.
+        """
+
+        with graph.local_scope():
+            if (graph.in_degrees() == 0).any():
+                raise DGLError(
+                    "There are 0-in-degree nodes in the graph, "
+                    "output for those nodes will be invalid. "
+                    "This is harmful for some applications, "
+                    "causing silent performance regression. "
+                    "Adding self-loop on the input graph by "
+                    "calling `g = dgl.add_self_loop(g)` will resolve "
+                    "the issue."
+                )
+
+            # calc edge attention
+            # same trick way as in dgl.nn.pytorch.GATConv, but also includes edge feats
+            # https://github.com/dmlc/dgl/blob/master/python/dgl/nn/pytorch/conv/gatconv.py
+            if isinstance(nfeats, tuple):
+                nfeats_src, nfeats_dst = nfeats
+            else:
+                nfeats_src = nfeats_dst = nfeats
+
+            f_ni = self.fc_ni(nfeats_src)
+            f_nj = self.fc_nj(nfeats_dst)
+            f_fij = self.fc_fij(efeats)
+
+            graph.srcdata.update({"f_ni": f_ni})
+            graph.dstdata.update({"f_nj": f_nj})
+            # add ni, nj factors
+            graph.apply_edges(fn.u_add_v("f_ni", "f_nj", "f_tmp"))
+            # add fij to node factor
+            f_out = graph.edata.pop("f_tmp") + f_fij
+            if self.bias is not None:
+                f_out = f_out + self.bias
+            f_out = nn.functional.leaky_relu(f_out)
+            f_out = f_out.view(-1, self._num_heads, self._out_edge_feats)
+            # compute attention factor
+            e = (f_out * self.attn).sum(dim=-1).unsqueeze(-1)
+            graph.edata["a"] = edge_softmax(graph, e)
+            if edge_weight is not None:
+                graph.edata["a"] = graph.edata["a"] * edge_weight.tile(
+                    1, self._num_heads, 1
+                ).transpose(0, 2)
+            graph.srcdata["h_out"] = self.fc_node_src(nfeats_src).view(
+                -1, self._num_heads, self._out_node_feats
+            )
+            # calc weighted sum
+            graph.update_all(
+                fn.u_mul_e("h_out", "a", "m"), fn.sum("m", "h_out")
+            )
+
+            h_out = graph.dstdata["h_out"].view(
+                -1, self._num_heads, self._out_node_feats
+            )
+            if get_attention:
+                return h_out, f_out, graph.edata.pop("a")
+            else:
+                return h_out, f_out

python/dgl/nn/pytorch/conv/egnnconv.py

@@ -0,0 +1,163 @@
+"""Torch Module for E(n) Equivariant Graph Convolutional Layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch
+import torch.nn as nn
+
+from .... import function as fn
+
+
+class EGNNConv(nn.Module):
+    r"""Equivariant Graph Convolutional Layer from `E(n) Equivariant Graph
+    Neural Networks <https://arxiv.org/abs/2102.09844>`__
+
+    .. math::
+
+        m_{ij}=\phi_e(h_i^l, h_j^l, ||x_i^l-x_j^l||^2, a_{ij})
+
+        x_i^{l+1} = x_i^l + C\sum_{j\in\mathcal{N}(i)}(x_i^l-x_j^l)\phi_x(m_{ij})
+
+        m_i = \sum_{j\in\mathcal{N}(i)} m_{ij}
+
+        h_i^{l+1} = \phi_h(h_i^l, m_i)
+
+    where :math:`h_i`, :math:`x_i`, :math:`a_{ij}` are node features, coordinate
+    features, and edge features respectively. :math:`\phi_e`, :math:`\phi_h`, and
+    :math:`\phi_x` are two-layer MLPs. :math:`C` is a constant for normalization,
+    computed as :math:`1/|\mathcal{N}(i)|`.
+
+    Parameters
+    ----------
+    in_size : int
+        Input feature size; i.e. the size of :math:`h_i^l`.
+    hidden_size : int
+        Hidden feature size; i.e. the size of hidden layer in the two-layer MLPs in
+        :math:`\phi_e, \phi_x, \phi_h`.
+    out_size : int
+        Output feature size; i.e. the size of :math:`h_i^{l+1}`.
+    edge_feat_size : int, optional
+        Edge feature size; i.e. the size of :math:`a_{ij}`. Default: 0.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import torch as th
+    >>> from dgl.nn import EGNNConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> node_feat, coord_feat, edge_feat = th.ones(6, 10), th.ones(6, 3), th.ones(6, 2)
+    >>> conv = EGNNConv(10, 10, 10, 2)
+    >>> h, x = conv(g, node_feat, coord_feat, edge_feat)
+    """
+
+    def __init__(self, in_size, hidden_size, out_size, edge_feat_size=0):
+        super(EGNNConv, self).__init__()
+
+        self.in_size = in_size
+        self.hidden_size = hidden_size
+        self.out_size = out_size
+        self.edge_feat_size = edge_feat_size
+        act_fn = nn.SiLU()
+
+        # \phi_e
+        self.edge_mlp = nn.Sequential(
+            # +1 for the radial feature: ||x_i - x_j||^2
+            nn.Linear(in_size * 2 + edge_feat_size + 1, hidden_size),
+            act_fn,
+            nn.Linear(hidden_size, hidden_size),
+            act_fn,
+        )
+
+        # \phi_h
+        self.node_mlp = nn.Sequential(
+            nn.Linear(in_size + hidden_size, hidden_size),
+            act_fn,
+            nn.Linear(hidden_size, out_size),
+        )
+
+        # \phi_x
+        self.coord_mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            act_fn,
+            nn.Linear(hidden_size, 1, bias=False),
+        )
+
+    def message(self, edges):
+        """message function for EGNN"""
+        # concat features for edge mlp
+        if self.edge_feat_size > 0:
+            f = torch.cat(
+                [
+                    edges.src["h"],
+                    edges.dst["h"],
+                    edges.data["radial"],
+                    edges.data["a"],
+                ],
+                dim=-1,
+            )
+        else:
+            f = torch.cat(
+                [edges.src["h"], edges.dst["h"], edges.data["radial"]], dim=-1
+            )
+
+        msg_h = self.edge_mlp(f)
+        msg_x = self.coord_mlp(msg_h) * edges.data["x_diff"]
+
+        return {"msg_x": msg_x, "msg_h": msg_h}
+
+    def forward(self, graph, node_feat, coord_feat, edge_feat=None):
+        r"""
+        Description
+        -----------
+        Compute EGNN layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        node_feat : torch.Tensor
+            The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
+            nodes, and :math:`h_n` must be the same as in_size.
+        coord_feat : torch.Tensor
+            The coordinate feature of shape :math:`(N, h_x)`. :math:`N` is the
+            number of nodes, and :math:`h_x` can be any positive integer.
+        edge_feat : torch.Tensor, optional
+            The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
+            edges, and :math:`h_e` must be the same as edge_feat_size.
+
+        Returns
+        -------
+        node_feat_out : torch.Tensor
+            The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
+            is the same as out_size.
+        coord_feat_out: torch.Tensor
+            The output coordinate feature of shape :math:`(N, h_x)` where :math:`h_x`
+            is the same as the input coordinate feature dimension.
+        """
+        with graph.local_scope():
+            # node feature
+            graph.ndata["h"] = node_feat
+            # coordinate feature
+            graph.ndata["x"] = coord_feat
+            # edge feature
+            if self.edge_feat_size > 0:
+                assert edge_feat is not None, "Edge features must be provided."
+                graph.edata["a"] = edge_feat
+            # get coordinate diff & radial features
+            graph.apply_edges(fn.u_sub_v("x", "x", "x_diff"))
+            graph.edata["radial"] = (
+                graph.edata["x_diff"].square().sum(dim=1).unsqueeze(-1)
+            )
+            # normalize coordinate difference
+            graph.edata["x_diff"] = graph.edata["x_diff"] / (
+                graph.edata["radial"].sqrt() + 1e-30
+            )
+            graph.apply_edges(self.message)
+            graph.update_all(fn.copy_e("msg_x", "m"), fn.mean("m", "x_neigh"))
+            graph.update_all(fn.copy_e("msg_h", "m"), fn.sum("m", "h_neigh"))
+
+            h_neigh, x_neigh = graph.ndata["h_neigh"], graph.ndata["x_neigh"]
+
+            h = self.node_mlp(torch.cat([node_feat, h_neigh], dim=-1))
+            x = coord_feat + x_neigh
+
+            return h, x

python/dgl/nn/pytorch/conv/gatconv.py

@@ -0,0 +1,370 @@
+"""Torch modules for graph attention networks(GAT)."""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import expand_as_pair
+from ...functional import edge_softmax
+from ..utils import Identity
+
+
+# pylint: enable=W0235
+class GATConv(nn.Module):
+    r"""Graph attention layer from `Graph Attention Network
+    <https://arxiv.org/pdf/1710.10903.pdf>`__
+
+    .. math::
+        h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{i,j} W^{(l)} h_j^{(l)}
+
+    where :math:`\alpha_{ij}` is the attention score bewteen node :math:`i` and
+    node :math:`j`:
+
+    .. math::
+        \alpha_{ij}^{l} &= \mathrm{softmax_i} (e_{ij}^{l})
+
+        e_{ij}^{l} &= \mathrm{LeakyReLU}\left(\vec{a}^T [W h_{i} \| W h_{j}]\right)
+
+    Parameters
+    ----------
+    in_feats : int, or pair of ints
+        Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
+        GATConv can be applied on homogeneous graph and unidirectional
+        `bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
+        If the layer is to be applied to a unidirectional bipartite graph, ``in_feats``
+        specifies the input feature size on both the source and destination nodes.  If
+        a scalar is given, the source and destination node feature size would take the
+        same value.
+    out_feats : int
+        Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
+    num_heads : int
+        Number of heads in Multi-Head Attention.
+    feat_drop : float, optional
+        Dropout rate on feature. Defaults: ``0``.
+    attn_drop : float, optional
+        Dropout rate on attention weight. Defaults: ``0``.
+    negative_slope : float, optional
+        LeakyReLU angle of negative slope. Defaults: ``0.2``.
+    residual : bool, optional
+        If True, use residual connection. Defaults: ``False``.
+    activation : callable activation function/layer or None, optional.
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Defaults: ``False``.
+    bias : bool, optional
+        If True, learns a bias term. Defaults: ``True``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import GATConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> gatconv = GATConv(10, 2, num_heads=3)
+    >>> res = gatconv(g, feat)
+    >>> res
+    tensor([[[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]],
+            [[ 3.4570,  1.8634],
+            [ 1.3805, -0.0762],
+            [ 1.0390, -1.1479]]], grad_fn=<BinaryReduceBackward>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('A', 'r', 'B'): (u, v)})
+    >>> u_feat = th.tensor(np.random.rand(2, 5).astype(np.float32))
+    >>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
+    >>> gatconv = GATConv((5,10), 2, 3)
+    >>> res = gatconv(g, (u_feat, v_feat))
+    >>> res
+    tensor([[[-0.6066,  1.0268],
+            [-0.5945, -0.4801],
+            [ 0.1594,  0.3825]],
+            [[ 0.0268,  1.0783],
+            [ 0.5041, -1.3025],
+            [ 0.6568,  0.7048]],
+            [[-0.2688,  1.0543],
+            [-0.0315, -0.9016],
+            [ 0.3943,  0.5347]],
+            [[-0.6066,  1.0268],
+            [-0.5945, -0.4801],
+            [ 0.1594,  0.3825]]], grad_fn=<BinaryReduceBackward>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        num_heads,
+        feat_drop=0.0,
+        attn_drop=0.0,
+        negative_slope=0.2,
+        residual=False,
+        activation=None,
+        allow_zero_in_degree=False,
+        bias=True,
+    ):
+        super(GATConv, self).__init__()
+        self._num_heads = num_heads
+        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
+        self._out_feats = out_feats
+        self._allow_zero_in_degree = allow_zero_in_degree
+        if isinstance(in_feats, tuple):
+            self.fc_src = nn.Linear(
+                self._in_src_feats, out_feats * num_heads, bias=False
+            )
+            self.fc_dst = nn.Linear(
+                self._in_dst_feats, out_feats * num_heads, bias=False
+            )
+        else:
+            self.fc = nn.Linear(
+                self._in_src_feats, out_feats * num_heads, bias=False
+            )
+        self.attn_l = nn.Parameter(
+            th.FloatTensor(size=(1, num_heads, out_feats))
+        )
+        self.attn_r = nn.Parameter(
+            th.FloatTensor(size=(1, num_heads, out_feats))
+        )
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.leaky_relu = nn.LeakyReLU(negative_slope)
+
+        self.has_linear_res = False
+        self.has_explicit_bias = False
+        if residual:
+            if self._in_dst_feats != out_feats * num_heads:
+                self.res_fc = nn.Linear(
+                    self._in_dst_feats, num_heads * out_feats, bias=bias
+                )
+                self.has_linear_res = True
+            else:
+                self.res_fc = Identity()
+        else:
+            self.register_buffer("res_fc", None)
+
+        if bias and not self.has_linear_res:
+            self.bias = nn.Parameter(
+                th.FloatTensor(size=(num_heads * out_feats,))
+            )
+            self.has_explicit_bias = True
+        else:
+            self.register_buffer("bias", None)
+
+        self.reset_parameters()
+        self.activation = activation
+
+    def reset_parameters(self):
+        """
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The fc weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
+        The attention weights are using xavier initialization method.
+        """
+        gain = nn.init.calculate_gain("relu")
+        if hasattr(self, "fc"):
+            nn.init.xavier_normal_(self.fc.weight, gain=gain)
+        else:
+            nn.init.xavier_normal_(self.fc_src.weight, gain=gain)
+            nn.init.xavier_normal_(self.fc_dst.weight, gain=gain)
+        nn.init.xavier_normal_(self.attn_l, gain=gain)
+        nn.init.xavier_normal_(self.attn_r, gain=gain)
+        if self.has_explicit_bias:
+            nn.init.constant_(self.bias, 0)
+        if isinstance(self.res_fc, nn.Linear):
+            nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
+            if self.res_fc.bias is not None:
+                nn.init.constant_(self.res_fc.bias, 0)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, edge_weight=None, get_attention=False):
+        r"""
+
+        Description
+        -----------
+        Compute graph attention network layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, the input feature of shape :math:`(N, *, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, *, D_{in_{src}})` and :math:`(N_{out}, *, D_{in_{dst}})`.
+        edge_weight : torch.Tensor, optional
+            A 1D tensor of edge weight values.  Shape: :math:`(|E|,)`.
+        get_attention : bool, optional
+            Whether to return the attention values. Default to False.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, *, H, D_{out})` where :math:`H`
+            is the number of heads, and :math:`D_{out}` is size of output feature.
+        torch.Tensor, optional
+            The attention values of shape :math:`(E, *, H, 1)`, where :math:`E` is the number of
+            edges. This is returned only when :attr:`get_attention` is ``True``.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            if isinstance(feat, tuple):
+                src_prefix_shape = feat[0].shape[:-1]
+                dst_prefix_shape = feat[1].shape[:-1]
+                h_src = self.feat_drop(feat[0])
+                h_dst = self.feat_drop(feat[1])
+                if not hasattr(self, "fc_src"):
+                    feat_src = self.fc(h_src).view(
+                        *src_prefix_shape, self._num_heads, self._out_feats
+                    )
+                    feat_dst = self.fc(h_dst).view(
+                        *dst_prefix_shape, self._num_heads, self._out_feats
+                    )
+                else:
+                    feat_src = self.fc_src(h_src).view(
+                        *src_prefix_shape, self._num_heads, self._out_feats
+                    )
+                    feat_dst = self.fc_dst(h_dst).view(
+                        *dst_prefix_shape, self._num_heads, self._out_feats
+                    )
+            else:
+                src_prefix_shape = dst_prefix_shape = feat.shape[:-1]
+                h_src = h_dst = self.feat_drop(feat)
+                feat_src = feat_dst = self.fc(h_src).view(
+                    *src_prefix_shape, self._num_heads, self._out_feats
+                )
+                if graph.is_block:
+                    feat_dst = feat_src[: graph.number_of_dst_nodes()]
+                    h_dst = h_dst[: graph.number_of_dst_nodes()]
+                    dst_prefix_shape = (
+                        graph.number_of_dst_nodes(),
+                    ) + dst_prefix_shape[1:]
+            # NOTE: GAT paper uses "first concatenation then linear projection"
+            # to compute attention scores, while ours is "first projection then
+            # addition", the two approaches are mathematically equivalent:
+            # We decompose the weight vector a mentioned in the paper into
+            # [a_l || a_r], then
+            # a^T [Wh_i || Wh_j] = a_l Wh_i + a_r Wh_j
+            # Our implementation is much efficient because we do not need to
+            # save [Wh_i || Wh_j] on edges, which is not memory-efficient. Plus,
+            # addition could be optimized with DGL's built-in function u_add_v,
+            # which further speeds up computation and saves memory footprint.
+            el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(-1)
+            er = (feat_dst * self.attn_r).sum(dim=-1).unsqueeze(-1)
+            graph.srcdata.update({"ft": feat_src, "el": el})
+            graph.dstdata.update({"er": er})
+            # compute edge attention, el and er are a_l Wh_i and a_r Wh_j respectively.
+            graph.apply_edges(fn.u_add_v("el", "er", "e"))
+            e = self.leaky_relu(graph.edata.pop("e"))
+            # compute softmax
+            graph.edata["a"] = self.attn_drop(edge_softmax(graph, e))
+            if edge_weight is not None:
+                graph.edata["a"] = graph.edata["a"] * edge_weight.tile(
+                    1, self._num_heads, 1
+                ).transpose(0, 2)
+            # message passing
+            graph.update_all(fn.u_mul_e("ft", "a", "m"), fn.sum("m", "ft"))
+            rst = graph.dstdata["ft"]
+            # residual
+            if self.res_fc is not None:
+                # Use -1 rather than self._num_heads to handle broadcasting
+                if h_dst.numel() != 0:
+                    resval = self.res_fc(h_dst).view(
+                        *dst_prefix_shape, -1, self._out_feats
+                    )
+                    rst = rst + resval
+            # bias
+            if self.has_explicit_bias:
+                rst = rst + self.bias.view(
+                    *((1,) * len(dst_prefix_shape)),
+                    self._num_heads,
+                    self._out_feats
+                )
+            # activation
+            if self.activation:
+                rst = self.activation(rst)
+
+            if get_attention:
+                return rst, graph.edata["a"]
+            else:
+                return rst

python/dgl/nn/pytorch/conv/gatedgcnconv.py

@@ -0,0 +1,173 @@
+"""Torch Module for GatedGCN layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name, cell-var-from-loop
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .... import function as fn
+
+
+class GatedGCNConv(nn.Module):
+    r"""Gated graph convolutional layer from `Benchmarking Graph Neural Networks
+    <https://arxiv.org/abs/2003.00982>`__
+
+    .. math::
+        e_{ij}^{l+1}=D^l h_{i}^{l}+E^l h_{j}^{l}+C^l e_{ij}^{l}
+
+        norm_{ij}=\Sigma_{j\in N_{i}} \sigma\left(e_{ij}^{l+1}\right)+\varepsilon
+
+        \hat{e}_{ij}^{l+1}=\sigma(e_{ij}^{l+1}) / norm_{ij}
+
+        h_{i}^{l+1}=A^l h_{i}^{l}+\Sigma_{j \in N_{i}} \hat{e}_{ij}^{l+1} \odot B^l h_{j}^{l}
+
+    where :math:`h_{i}^{l}` is node :math:`i` feature of layer :math:`l`,
+    :math:`e_{ij}^{l}` is edge :math:`ij` feature of layer :math:`l`,
+    :math:`\sigma` is sigmoid function, :math:`\varepsilon` is a small fixed constant
+    for numerical stability, :math:`A^l, B^l, C^l, D^l, E^l` are linear layers.
+
+    Parameters
+    ----------
+    input_feats : int
+        Input feature size; i.e, the number of dimensions of :math:`h_{i}^{l}`.
+    edge_feats: int
+        Edge feature size; i.e., the number of dimensions of :math:`e_{ij}^{l}`.
+    output_feats : int
+        Output feature size; i.e., the number of dimensions of :math:`h_{i}^{l+1}`.
+    dropout : float, optional
+        Dropout rate on node and edge feature. Default: ``0``.
+    batch_norm : bool, optional
+        Whether to include batch normalization on node and edge feature. Default: ``True``.
+    residual : bool, optional
+        Whether to include residual connections. Default: ``True``.
+    activation : callable activation function/layer or None, optional
+        If not None, apply an activation function to the updated node features.
+        Default: ``F.relu``.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import torch as th
+    >>> import torch.nn.functional as F
+    >>> from dgl.nn import GatedGCNConv
+
+    >>> num_nodes, num_edges = 8, 30
+    >>> graph = dgl.rand_graph(num_nodes,num_edges)
+    >>> node_feats = th.rand(num_nodes, 20)
+    >>> edge_feats = th.rand(num_edges, 12)
+    >>> gatedGCN = GatedGCNConv(20, 12, 20)
+    >>> new_node_feats, new_edge_feats = gatedGCN(graph, node_feats, edge_feats)
+    >>> new_node_feats.shape, new_edge_feats.shape
+    (torch.Size([8, 20]), torch.Size([30, 20]))
+
+    """
+
+    def __init__(
+        self,
+        input_feats,
+        edge_feats,
+        output_feats,
+        dropout=0,
+        batch_norm=True,
+        residual=True,
+        activation=F.relu,
+    ):
+        super(GatedGCNConv, self).__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.batch_norm = batch_norm
+        self.residual = residual
+
+        if input_feats != output_feats or edge_feats != output_feats:
+            self.residual = False
+
+        # Linearly transform the node features.
+        self.A = nn.Linear(input_feats, output_feats, bias=True)
+        self.B = nn.Linear(input_feats, output_feats, bias=True)
+        self.D = nn.Linear(input_feats, output_feats, bias=True)
+        self.E = nn.Linear(input_feats, output_feats, bias=True)
+
+        # Linearly transform the edge features.
+        self.C = nn.Linear(edge_feats, output_feats, bias=True)
+
+        # Batch normalization on the node/edge features.
+        self.bn_node = nn.BatchNorm1d(output_feats)
+        self.bn_edge = nn.BatchNorm1d(output_feats)
+
+        self.activation = activation
+
+    def forward(self, graph, feat, edge_feat):
+        """
+
+        Description
+        -----------
+        Compute gated graph convolution layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, D_{in})` where :math:`N`
+            is the number of nodes of the graph and :math:`D_{in}` is the
+            input feature size.
+        edge_feat : torch.Tensor
+            The input edge feature of shape :math:`(E, D_{edge})`,
+            where :math:`E` is the number of edges and :math:`D_{edge}`
+            is the size of the edge features.
+
+        Returns
+        -------
+        torch.Tensor
+            The output node feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is the output feature size.
+        torch.Tensor
+            The output edge feature of shape :math:`(E, D_{out})` where :math:`D_{out}`
+            is the output feature size.
+        """
+        with graph.local_scope():
+            # For residual connection
+            h_in = feat
+            e_in = edge_feat
+
+            graph.ndata["Ah"] = self.A(feat)
+            graph.ndata["Bh"] = self.B(feat)
+            graph.ndata["Dh"] = self.D(feat)
+            graph.ndata["Eh"] = self.E(feat)
+            graph.edata["Ce"] = self.C(edge_feat)
+
+            graph.apply_edges(fn.u_add_v("Dh", "Eh", "DEh"))
+
+            # Get edge feature
+            graph.edata["e"] = graph.edata["DEh"] + graph.edata["Ce"]
+            graph.edata["sigma"] = torch.sigmoid(graph.edata["e"])
+
+            graph.update_all(
+                fn.u_mul_e("Bh", "sigma", "m"), fn.sum("m", "sum_sigma_h")
+            )
+            graph.update_all(fn.copy_e("sigma", "m"), fn.sum("m", "sum_sigma"))
+            graph.ndata["h"] = graph.ndata["Ah"] + graph.ndata[
+                "sum_sigma_h"
+            ] / (graph.ndata["sum_sigma"] + 1e-6)
+
+            # Result of graph convolution.
+            feat = graph.ndata["h"]
+            edge_feat = graph.edata["e"]
+
+            # Batch normalization.
+            if self.batch_norm:
+                feat = self.bn_node(feat)
+                edge_feat = self.bn_edge(edge_feat)
+
+            # Non-linear activation.
+            if self.activation:
+                feat = self.activation(feat)
+                edge_feat = self.activation(edge_feat)
+
+            # Residual connection.
+            if self.residual:
+                feat = h_in + feat
+                edge_feat = e_in + edge_feat
+
+            feat = self.dropout(feat)
+            edge_feat = self.dropout(edge_feat)
+
+            return feat, edge_feat

python/dgl/nn/pytorch/conv/gatedgraphconv.py

@@ -0,0 +1,173 @@
+"""Torch Module for Gated Graph Convolution layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name, cell-var-from-loop
+import torch as th
+from torch import nn
+from torch.nn import init
+
+from .... import function as fn
+
+
+class GatedGraphConv(nn.Module):
+    r"""Gated Graph Convolution layer from `Gated Graph Sequence
+    Neural Networks <https://arxiv.org/pdf/1511.05493.pdf>`__
+
+    .. math::
+        h_{i}^{0} &= [ x_i \| \mathbf{0} ]
+
+        a_{i}^{t} &= \sum_{j\in\mathcal{N}(i)} W_{e_{ij}} h_{j}^{t}
+
+        h_{i}^{t+1} &= \mathrm{GRU}(a_{i}^{t}, h_{i}^{t})
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size; i.e, the number of dimensions of :math:`x_i`.
+    out_feats : int
+        Output feature size; i.e., the number of dimensions of :math:`h_i^{(t+1)}`.
+    n_steps : int
+        Number of recurrent steps; i.e, the :math:`t` in the above formula.
+    n_etypes : int
+        Number of edge types.
+    bias : bool
+        If True, adds a learnable bias to the output. Default: ``True``.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import GatedGraphConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = GatedGraphConv(10, 10, 2, 3)
+    >>> etype = th.tensor([0,1,2,0,1,2])
+    >>> res = conv(g, feat, etype)
+    >>> res
+    tensor([[ 0.4652,  0.4458,  0.5169,  0.4126,  0.4847,  0.2303,  0.2757,  0.7721,
+            0.0523,  0.0857],
+            [ 0.0832,  0.1388, -0.5643,  0.7053, -0.2524, -0.3847,  0.7587,  0.8245,
+            0.9315,  0.4063],
+            [ 0.6340,  0.4096,  0.7692,  0.2125,  0.2106,  0.4542, -0.0580,  0.3364,
+            -0.1376,  0.4948],
+            [ 0.5551,  0.7946,  0.6220,  0.8058,  0.5711,  0.3063, -0.5454,  0.2272,
+            -0.6931, -0.1607],
+            [ 0.2644,  0.2469, -0.6143,  0.6008, -0.1516, -0.3781,  0.5878,  0.7993,
+            0.9241,  0.1835],
+            [ 0.6393,  0.3447,  0.3893,  0.4279,  0.3342,  0.3809,  0.0406,  0.5030,
+            0.1342,  0.0425]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(self, in_feats, out_feats, n_steps, n_etypes, bias=True):
+        super(GatedGraphConv, self).__init__()
+        assert in_feats <= out_feats, "out_feats must be not less than in_feats"
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self._n_steps = n_steps
+        self._n_etypes = n_etypes
+        self.linears = nn.ModuleList(
+            [nn.Linear(out_feats, out_feats) for _ in range(n_etypes)]
+        )
+        self.gru = nn.GRUCell(out_feats, out_feats, bias=bias)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The model parameters are initialized using Glorot uniform initialization
+        and the bias is initialized to be zero.
+        """
+        gain = init.calculate_gain("relu")
+        self.gru.reset_parameters()
+        for linear in self.linears:
+            init.xavier_normal_(linear.weight, gain=gain)
+            init.zeros_(linear.bias)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, etypes=None):
+        """
+
+        Description
+        -----------
+        Compute Gated Graph Convolution layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, D_{in})` where :math:`N`
+            is the number of nodes of the graph and :math:`D_{in}` is the
+            input feature size.
+        etypes : torch.LongTensor, or None
+            The edge type tensor of shape :math:`(E,)` where :math:`E` is
+            the number of edges of the graph. When there's only one edge type,
+            this argument can be skipped
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is the output feature size.
+        """
+        with graph.local_scope():
+            assert graph.is_homogeneous, (
+                "not a homogeneous graph; convert it with to_homogeneous "
+                "and pass in the edge type as argument"
+            )
+            if self._n_etypes != 1:
+                assert (
+                    etypes.min() >= 0 and etypes.max() < self._n_etypes
+                ), "edge type indices out of range [0, {})".format(
+                    self._n_etypes
+                )
+
+            zero_pad = feat.new_zeros(
+                (feat.shape[0], self._out_feats - feat.shape[1])
+            )
+            feat = th.cat([feat, zero_pad], -1)
+
+            for _ in range(self._n_steps):
+                if self._n_etypes == 1 and etypes is None:
+                    # Fast path when graph has only one edge type
+                    graph.ndata["h"] = self.linears[0](feat)
+                    graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "a"))
+                    a = graph.ndata.pop("a")  # (N, D)
+                else:
+                    graph.ndata["h"] = feat
+                    for i in range(self._n_etypes):
+                        eids = (
+                            th.nonzero(etypes == i, as_tuple=False)
+                            .view(-1)
+                            .type(graph.idtype)
+                        )
+                        if len(eids) > 0:
+                            graph.apply_edges(
+                                lambda edges: {
+                                    "W_e*h": self.linears[i](edges.src["h"])
+                                },
+                                eids,
+                            )
+                    graph.update_all(fn.copy_e("W_e*h", "m"), fn.sum("m", "a"))
+                    a = graph.ndata.pop("a")  # (N, D)
+                feat = self.gru(a, feat)
+            return feat

python/dgl/nn/pytorch/conv/gatv2conv.py

@@ -0,0 +1,335 @@
+"""Torch modules for graph attention networks v2 (GATv2)."""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import expand_as_pair
+from ...functional import edge_softmax
+from ..utils import Identity
+
+
+# pylint: enable=W0235
+class GATv2Conv(nn.Module):
+    r"""GATv2 from `How Attentive are Graph Attention Networks?
+    <https://arxiv.org/pdf/2105.14491.pdf>`__
+
+    .. math::
+        h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{ij}^{(l)} W^{(l)}_{right} h_j^{(l)}
+
+    where :math:`\alpha_{ij}` is the attention score bewteen node :math:`i` and
+    node :math:`j`:
+
+    .. math::
+        \alpha_{ij}^{(l)} &= \mathrm{softmax_i} (e_{ij}^{(l)})
+
+        e_{ij}^{(l)} &= {\vec{a}^T}^{(l)}\mathrm{LeakyReLU}\left(
+            W^{(l)}_{left} h_{i} + W^{(l)}_{right} h_{j}\right)
+
+    Parameters
+    ----------
+    in_feats : int, or pair of ints
+        Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
+        If the layer is to be applied to a unidirectional bipartite graph, `in_feats`
+        specifies the input feature size on both the source and destination nodes.
+        If a scalar is given, the source and destination node feature size
+        would take the same value.
+    out_feats : int
+        Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
+    num_heads : int
+        Number of heads in Multi-Head Attention.
+    feat_drop : float, optional
+        Dropout rate on feature. Defaults: ``0``.
+    attn_drop : float, optional
+        Dropout rate on attention weight. Defaults: ``0``.
+    negative_slope : float, optional
+        LeakyReLU angle of negative slope. Defaults: ``0.2``.
+    residual : bool, optional
+        If True, use residual connection. Defaults: ``False``.
+    activation : callable activation function/layer or None, optional.
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Defaults: ``False``.
+    bias : bool, optional
+        If set to :obj:`False`, the layer will not learn
+        an additive bias. (default: :obj:`True`)
+    share_weights : bool, optional
+        If set to :obj:`True`, the same matrix for :math:`W_{left}` and :math:`W_{right}` in
+        the above equations, will be applied to the source and the target node of every edge.
+        (default: :obj:`False`)
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be applied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import GATv2Conv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> gatv2conv = GATv2Conv(10, 2, num_heads=3)
+    >>> res = gatv2conv(g, feat)
+    >>> res
+    tensor([[[ 1.9599,  1.0239],
+            [ 3.2015, -0.5512],
+            [ 2.3700, -2.2182]],
+            [[ 1.9599,  1.0239],
+            [ 3.2015, -0.5512],
+            [ 2.3700, -2.2182]],
+            [[ 1.9599,  1.0239],
+            [ 3.2015, -0.5512],
+            [ 2.3700, -2.2182]],
+            [[ 1.9599,  1.0239],
+            [ 3.2015, -0.5512],
+            [ 2.3700, -2.2182]],
+            [[ 1.9599,  1.0239],
+            [ 3.2015, -0.5512],
+            [ 2.3700, -2.2182]],
+            [[ 1.9599,  1.0239],
+            [ 3.2015, -0.5512],
+            [ 2.3700, -2.2182]]], grad_fn=<GSpMMBackward>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('A', 'r', 'B'): (u, v)})
+    >>> u_feat = th.tensor(np.random.rand(2, 5).astype(np.float32))
+    >>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
+    >>> gatv2conv = GATv2Conv((5,10), 2, 3)
+    >>> res = gatv2conv(g, (u_feat, v_feat))
+    >>> res
+    tensor([[[-0.0935, -0.4273],
+            [-1.1850,  0.1123],
+            [-0.2002,  0.1155]],
+            [[ 0.1908, -1.2095],
+            [-0.0129,  0.6408],
+            [-0.8135,  0.1157]],
+            [[ 0.0596, -0.8487],
+            [-0.5421,  0.4022],
+            [-0.4805,  0.1156]],
+            [[-0.0935, -0.4273],
+            [-1.1850,  0.1123],
+            [-0.2002,  0.1155]]], grad_fn=<GSpMMBackward>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        num_heads,
+        feat_drop=0.0,
+        attn_drop=0.0,
+        negative_slope=0.2,
+        residual=False,
+        activation=None,
+        allow_zero_in_degree=False,
+        bias=True,
+        share_weights=False,
+    ):
+        super(GATv2Conv, self).__init__()
+        self._num_heads = num_heads
+        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
+        self._out_feats = out_feats
+        self._allow_zero_in_degree = allow_zero_in_degree
+        if isinstance(in_feats, tuple):
+            self.fc_src = nn.Linear(
+                self._in_src_feats, out_feats * num_heads, bias=bias
+            )
+            self.fc_dst = nn.Linear(
+                self._in_dst_feats, out_feats * num_heads, bias=bias
+            )
+        else:
+            self.fc_src = nn.Linear(
+                self._in_src_feats, out_feats * num_heads, bias=bias
+            )
+            if share_weights:
+                self.fc_dst = self.fc_src
+            else:
+                self.fc_dst = nn.Linear(
+                    self._in_src_feats, out_feats * num_heads, bias=bias
+                )
+        self.attn = nn.Parameter(th.FloatTensor(size=(1, num_heads, out_feats)))
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.leaky_relu = nn.LeakyReLU(negative_slope)
+        if residual:
+            if self._in_dst_feats != out_feats * num_heads:
+                self.res_fc = nn.Linear(
+                    self._in_dst_feats, num_heads * out_feats, bias=bias
+                )
+            else:
+                self.res_fc = Identity()
+        else:
+            self.register_buffer("res_fc", None)
+        self.activation = activation
+        self.share_weights = share_weights
+        self.bias = bias
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The fc weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
+        The attention weights are using xavier initialization method.
+        """
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_normal_(self.fc_src.weight, gain=gain)
+        if self.bias:
+            nn.init.constant_(self.fc_src.bias, 0)
+        if not self.share_weights:
+            nn.init.xavier_normal_(self.fc_dst.weight, gain=gain)
+            if self.bias:
+                nn.init.constant_(self.fc_dst.bias, 0)
+        nn.init.xavier_normal_(self.attn, gain=gain)
+        if isinstance(self.res_fc, nn.Linear):
+            nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
+            if self.bias:
+                nn.init.constant_(self.res_fc.bias, 0)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, get_attention=False):
+        r"""
+        Description
+        -----------
+        Compute graph attention network layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, the input feature of shape :math:`(N, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
+        get_attention : bool, optional
+            Whether to return the attention values. Default to False.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, H, D_{out})` where :math:`H`
+            is the number of heads, and :math:`D_{out}` is size of output feature.
+        torch.Tensor, optional
+            The attention values of shape :math:`(E, H, 1)`, where :math:`E` is the number of
+            edges. This is returned only when :attr:`get_attention` is ``True``.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            if isinstance(feat, tuple):
+                h_src = self.feat_drop(feat[0])
+                h_dst = self.feat_drop(feat[1])
+                feat_src = self.fc_src(h_src).view(
+                    -1, self._num_heads, self._out_feats
+                )
+                feat_dst = self.fc_dst(h_dst).view(
+                    -1, self._num_heads, self._out_feats
+                )
+            else:
+                h_src = h_dst = self.feat_drop(feat)
+                feat_src = self.fc_src(h_src).view(
+                    -1, self._num_heads, self._out_feats
+                )
+                if self.share_weights:
+                    feat_dst = feat_src
+                else:
+                    feat_dst = self.fc_dst(h_dst).view(
+                        -1, self._num_heads, self._out_feats
+                    )
+                if graph.is_block:
+                    feat_dst = feat_dst[: graph.number_of_dst_nodes()]
+                    h_dst = h_dst[: graph.number_of_dst_nodes()]
+            graph.srcdata.update(
+                {"el": feat_src}
+            )  # (num_src_edge, num_heads, out_dim)
+            graph.dstdata.update({"er": feat_dst})
+            graph.apply_edges(fn.u_add_v("el", "er", "e"))
+            e = self.leaky_relu(
+                graph.edata.pop("e")
+            )  # (num_src_edge, num_heads, out_dim)
+            e = (
+                (e * self.attn).sum(dim=-1).unsqueeze(dim=2)
+            )  # (num_edge, num_heads, 1)
+            # compute softmax
+            graph.edata["a"] = self.attn_drop(
+                edge_softmax(graph, e)
+            )  # (num_edge, num_heads)
+            # message passing
+            graph.update_all(fn.u_mul_e("el", "a", "m"), fn.sum("m", "ft"))
+            rst = graph.dstdata["ft"]
+            # residual
+            if self.res_fc is not None:
+                if h_dst.numel() != 0:
+                    resval = self.res_fc(h_dst).view(
+                        h_dst.shape[0], -1, self._out_feats
+                    )
+                    rst = rst + resval
+            # activation
+            if self.activation:
+                rst = self.activation(rst)
+
+            if get_attention:
+                return rst, graph.edata["a"]
+            else:
+                return rst

python/dgl/nn/pytorch/conv/gcn2conv.py

@@ -0,0 +1,285 @@
+"""Torch Module for Graph Convolutional Network via Initial residual
+    and Identity mapping (GCNII) layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import math
+
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from ....base import DGLError
+from .graphconv import EdgeWeightNorm
+
+
+class GCN2Conv(nn.Module):
+    r"""Graph Convolutional Network via Initial residual
+    and Identity mapping (GCNII) from `Simple and Deep Graph Convolutional
+    Networks <https://arxiv.org/abs/2007.02133>`__
+
+    It is mathematically is defined as follows:
+
+    .. math::
+
+        \mathbf{h}^{(l+1)} =\left( (1 - \alpha)(\mathbf{D}^{-1/2} \mathbf{\hat{A}}
+        \mathbf{D}^{-1/2})\mathbf{h}^{(l)} + \alpha {\mathbf{h}^{(0)}} \right)
+        \left( (1 - \beta_l) \mathbf{I} + \beta_l \mathbf{W} \right)
+
+    where :math:`\mathbf{\hat{A}}` is the adjacency matrix with self-loops,
+    :math:`\mathbf{D}_{ii} = \sum_{j=0} \mathbf{A}_{ij}` is its diagonal degree matrix,
+    :math:`\mathbf{h}^{(0)}` is the initial node features,
+    :math:`\mathbf{h}^{(l)}` is the feature of layer :math:`l`,
+    :math:`\alpha` is the fraction of initial node features, and
+    :math:`\beta_l` is the hyperparameter to tune the strength of identity mapping.
+    It is defined by :math:`\beta_l = \log(\frac{\lambda}{l}+1)\approx\frac{\lambda}{l}`,
+    where :math:`\lambda` is a hyperparameter. :math:`\beta` ensures that the decay of
+    the weight matrix adaptively increases as we stack more layers.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
+    layer : int
+        the index of current layer.
+    alpha : float
+        The fraction of the initial input features. Default: ``0.1``
+    lambda_ : float
+        The hyperparameter to ensure the decay of the weight matrix
+        adaptively increases. Default: ``1``
+    project_initial_features : bool
+        Whether to share a weight matrix between initial features and
+        smoothed features. Default: ``True``
+    bias : bool, optional
+        If True, adds a learnable bias to the output. Default: ``True``.
+    activation : callable activation function/layer or None, optional
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Default: ``False``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import GCN2Conv
+
+    >>> # Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 3)
+    >>> g = dgl.add_self_loop(g)
+    >>> conv1 = GCN2Conv(3, layer=1, alpha=0.5, \
+    ...         project_initial_features=True, allow_zero_in_degree=True)
+    >>> conv2 = GCN2Conv(3, layer=2, alpha=0.5, \
+    ...         project_initial_features=True, allow_zero_in_degree=True)
+    >>> res = feat
+    >>> res = conv1(g, res, feat)
+    >>> res = conv2(g, res, feat)
+    >>> print(res)
+    tensor([[1.3803, 3.3191, 2.9572],
+            [1.3803, 3.3191, 2.9572],
+            [1.3803, 3.3191, 2.9572],
+            [1.4770, 3.8326, 3.2451],
+            [1.3623, 3.2102, 2.8679],
+            [1.3803, 3.3191, 2.9572]], grad_fn=<AddBackward0>)
+
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        layer,
+        alpha=0.1,
+        lambda_=1,
+        project_initial_features=True,
+        allow_zero_in_degree=False,
+        bias=True,
+        activation=None,
+    ):
+        super().__init__()
+
+        self._in_feats = in_feats
+        self._project_initial_features = project_initial_features
+
+        self.alpha = alpha
+        self.beta = math.log(lambda_ / layer + 1)
+
+        self._bias = bias
+        self._activation = activation
+        self._allow_zero_in_degree = allow_zero_in_degree
+
+        self.weight1 = nn.Parameter(th.Tensor(self._in_feats, self._in_feats))
+
+        if self._project_initial_features:
+            self.register_parameter("weight2", None)
+        else:
+            self.weight2 = nn.Parameter(
+                th.Tensor(self._in_feats, self._in_feats)
+            )
+
+        if self._bias:
+            self.bias = nn.Parameter(th.Tensor(self._in_feats))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        """
+        nn.init.normal_(self.weight1)
+        if not self._project_initial_features:
+            nn.init.normal_(self.weight2)
+        if self._bias:
+            nn.init.zeros_(self.bias)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, feat_0, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute graph convolution.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape
+            :math:`(N, D_{in})`
+            where :math:`D_{in}` is the size of input feature and :math:`N` is the number of nodes.
+        feat_0 : torch.Tensor
+            The initial feature of shape :math:`(N, D_{in})`
+        edge_weight: torch.Tensor, optional
+            edge_weight to use in the message passing process. This is equivalent to
+            using weighted adjacency matrix in the equation above, and
+            :math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
+            is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
+
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+
+        Note
+        ----
+        * Input shape: :math:`(N, *, \text{in_feats})` where * means any number of additional
+          dimensions, :math:`N` is the number of nodes.
+        * Output shape: :math:`(N, *, \text{out_feats})` where all but the last dimension are
+          the same shape as the input.
+        * Weight shape: :math:`(\text{in_feats}, \text{out_feats})`.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            # normalize  to get smoothed representation
+            if edge_weight is None:
+                degs = graph.in_degrees().to(feat).clamp(min=1)
+                norm = th.pow(degs, -0.5)
+                norm = norm.to(feat.device).unsqueeze(1)
+            else:
+                edge_weight = EdgeWeightNorm("both")(graph, edge_weight)
+
+            if edge_weight is None:
+                feat = feat * norm
+            graph.ndata["h"] = feat
+            msg_func = fn.copy_u("h", "m")
+            if edge_weight is not None:
+                graph.edata["_edge_weight"] = edge_weight
+                msg_func = fn.u_mul_e("h", "_edge_weight", "m")
+            graph.update_all(msg_func, fn.sum("m", "h"))
+            feat = graph.ndata.pop("h")
+            if edge_weight is None:
+                feat = feat * norm
+            # scale
+            feat = feat * (1 - self.alpha)
+
+            # initial residual connection to the first layer
+            feat_0 = feat_0[: feat.size(0)] * self.alpha
+            feat_sum = feat + feat_0
+
+            if self._project_initial_features:
+                feat_proj_sum = feat_sum @ self.weight1
+            else:
+                feat_proj_sum = feat @ self.weight1 + feat_0 @ self.weight2
+
+            rst = (1 - self.beta) * feat_sum + self.beta * feat_proj_sum
+
+            if self._bias:
+                rst = rst + self.bias
+
+            if self._activation is not None:
+                rst = self._activation(rst)
+
+            return rst
+
+    def extra_repr(self):
+        """Set the extra representation of the module,
+        which will come into effect when printing the model.
+        """
+        summary = "in={_in_feats}"
+        summary += ", alpha={alpha}, beta={beta}"
+        if "self._bias" in self.__dict__:
+            summary += ", bias={bias}"
+        if "self._activation" in self.__dict__:
+            summary += ", activation={_activation}"
+
+        return summary.format(**self.__dict__)

python/dgl/nn/pytorch/conv/ginconv.py

@@ -0,0 +1,158 @@
+"""Torch Module for Graph Isomorphism Network layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from ....utils import expand_as_pair
+
+
+class GINConv(nn.Module):
+    r"""Graph Isomorphism Network layer from `How Powerful are Graph
+    Neural Networks? <https://arxiv.org/pdf/1810.00826.pdf>`__
+
+    .. math::
+        h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} +
+        \mathrm{aggregate}\left(\left\{h_j^{l}, j\in\mathcal{N}(i)
+        \right\}\right)\right)
+
+    If a weight tensor on each edge is provided, the weighted graph convolution is defined as:
+
+    .. math::
+        h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} +
+        \mathrm{aggregate}\left(\left\{e_{ji} h_j^{l}, j\in\mathcal{N}(i)
+        \right\}\right)\right)
+
+    where :math:`e_{ji}` is the weight on the edge from node :math:`j` to node :math:`i`.
+    Please make sure that `e_{ji}` is broadcastable with `h_j^{l}`.
+
+    Parameters
+    ----------
+    apply_func : callable activation function/layer or None
+        If not None, apply this function to the updated node feature,
+        the :math:`f_\Theta` in the formula, default: None.
+    aggregator_type : str
+        Aggregator type to use (``sum``, ``max`` or ``mean``), default: 'sum'.
+    init_eps : float, optional
+        Initial :math:`\epsilon` value, default: ``0``.
+    learn_eps : bool, optional
+        If True, :math:`\epsilon` will be a learnable parameter. Default: ``False``.
+    activation : callable activation function/layer or None, optional
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import GINConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> lin = th.nn.Linear(10, 10)
+    >>> conv = GINConv(lin, 'max')
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[-0.4821,  0.0207, -0.7665,  0.5721, -0.4682, -0.2134, -0.5236,  1.2855,
+            0.8843, -0.8764],
+            [-0.4821,  0.0207, -0.7665,  0.5721, -0.4682, -0.2134, -0.5236,  1.2855,
+            0.8843, -0.8764],
+            [-0.4821,  0.0207, -0.7665,  0.5721, -0.4682, -0.2134, -0.5236,  1.2855,
+            0.8843, -0.8764],
+            [-0.4821,  0.0207, -0.7665,  0.5721, -0.4682, -0.2134, -0.5236,  1.2855,
+            0.8843, -0.8764],
+            [-0.4821,  0.0207, -0.7665,  0.5721, -0.4682, -0.2134, -0.5236,  1.2855,
+            0.8843, -0.8764],
+            [-0.1804,  0.0758, -0.5159,  0.3569, -0.1408, -0.1395, -0.2387,  0.7773,
+            0.5266, -0.4465]], grad_fn=<AddmmBackward>)
+
+    >>> # With activation
+    >>> from torch.nn.functional import relu
+    >>> conv = GINConv(lin, 'max', activation=relu)
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
+             0.0000],
+            [5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
+             0.0000],
+            [5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
+             0.0000],
+            [5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
+             0.0000],
+            [5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
+             0.0000],
+            [2.5011, 0.0000, 0.0089, 2.0541, 0.8262, 0.0000, 0.0000, 0.1371, 0.0000,
+             0.0000]], grad_fn=<ReluBackward0>)
+    """
+
+    def __init__(
+        self,
+        apply_func=None,
+        aggregator_type="sum",
+        init_eps=0,
+        learn_eps=False,
+        activation=None,
+    ):
+        super(GINConv, self).__init__()
+        self.apply_func = apply_func
+        self._aggregator_type = aggregator_type
+        self.activation = activation
+        if aggregator_type not in ("sum", "max", "mean"):
+            raise KeyError(
+                "Aggregator type {} not recognized.".format(aggregator_type)
+            )
+        # to specify whether eps is trainable or not.
+        if learn_eps:
+            self.eps = th.nn.Parameter(th.FloatTensor([init_eps]))
+        else:
+            self.register_buffer("eps", th.FloatTensor([init_eps]))
+
+    def forward(self, graph, feat, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute Graph Isomorphism Network layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, the input feature of shape :math:`(N, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, D_{in})` and :math:`(N_{out}, D_{in})`.
+            If ``apply_func`` is not None, :math:`D_{in}` should
+            fit the input dimensionality requirement of ``apply_func``.
+        edge_weight : torch.Tensor, optional
+            Optional tensor on the edge. If given, the convolution will weight
+            with regard to the message.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where
+            :math:`D_{out}` is the output dimensionality of ``apply_func``.
+            If ``apply_func`` is None, :math:`D_{out}` should be the same
+            as input dimensionality.
+        """
+        _reducer = getattr(fn, self._aggregator_type)
+        with graph.local_scope():
+            aggregate_fn = fn.copy_u("h", "m")
+            if edge_weight is not None:
+                assert edge_weight.shape[0] == graph.num_edges()
+                graph.edata["_edge_weight"] = edge_weight
+                aggregate_fn = fn.u_mul_e("h", "_edge_weight", "m")
+
+            feat_src, feat_dst = expand_as_pair(feat, graph)
+            graph.srcdata["h"] = feat_src
+            graph.update_all(aggregate_fn, _reducer("m", "neigh"))
+            rst = (1 + self.eps) * feat_dst + graph.dstdata["neigh"]
+            if self.apply_func is not None:
+                rst = self.apply_func(rst)
+            # activation
+            if self.activation is not None:
+                rst = self.activation(rst)
+            return rst

python/dgl/nn/pytorch/conv/gineconv.py

@@ -0,0 +1,97 @@
+"""Torch Module for Graph Isomorphism Network layer variant with edge features"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+import torch.nn.functional as F
+from torch import nn
+
+from .... import function as fn
+from ....utils import expand_as_pair
+
+
+class GINEConv(nn.Module):
+    r"""Graph Isomorphism Network with Edge Features, introduced by
+    `Strategies for Pre-training Graph Neural Networks <https://arxiv.org/abs/1905.12265>`__
+
+    .. math::
+        h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} +
+        \sum_{j\in\mathcal{N}(i)}\mathrm{ReLU}(h_j^{l} + e_{j,i}^{l})\right)
+
+    where :math:`e_{j,i}^{l}` is the edge feature.
+
+    Parameters
+    ----------
+    apply_func : callable module or None
+        The :math:`f_\Theta` in the formula. If not None, it will be applied to
+        the updated node features. The default value is None.
+    init_eps : float, optional
+        Initial :math:`\epsilon` value, default: ``0``.
+    learn_eps : bool, optional
+        If True, :math:`\epsilon` will be a learnable parameter. Default: ``False``.
+
+    Examples
+    --------
+
+    >>> import dgl
+    >>> import torch
+    >>> import torch.nn as nn
+    >>> from dgl.nn import GINEConv
+
+    >>> g = dgl.graph(([0, 1, 2], [1, 1, 3]))
+    >>> in_feats = 10
+    >>> out_feats = 20
+    >>> nfeat = torch.randn(g.num_nodes(), in_feats)
+    >>> efeat = torch.randn(g.num_edges(), in_feats)
+    >>> conv = GINEConv(nn.Linear(in_feats, out_feats))
+    >>> res = conv(g, nfeat, efeat)
+    >>> print(res.shape)
+    torch.Size([4, 20])
+    """
+
+    def __init__(self, apply_func=None, init_eps=0, learn_eps=False):
+        super(GINEConv, self).__init__()
+        self.apply_func = apply_func
+        # to specify whether eps is trainable or not.
+        if learn_eps:
+            self.eps = nn.Parameter(th.FloatTensor([init_eps]))
+        else:
+            self.register_buffer("eps", th.FloatTensor([init_eps]))
+
+    def message(self, edges):
+        r"""User-defined Message Function"""
+        return {"m": F.relu(edges.src["hn"] + edges.data["he"])}
+
+    def forward(self, graph, node_feat, edge_feat):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        node_feat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, it is the input feature of shape :math:`(N, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, D_{in})` and :math:`(N_{out}, D_{in})`.
+            If ``apply_func`` is not None, :math:`D_{in}` should
+            fit the input feature size requirement of ``apply_func``.
+        edge_feat : torch.Tensor
+            Edge feature. It is a tensor of shape :math:`(E, D_{in})` where :math:`E`
+            is the number of edges.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where
+            :math:`D_{out}` is the output feature size of ``apply_func``.
+            If ``apply_func`` is None, :math:`D_{out}` should be the same
+            as :math:`D_{in}`.
+        """
+        with graph.local_scope():
+            feat_src, feat_dst = expand_as_pair(node_feat, graph)
+            graph.srcdata["hn"] = feat_src
+            graph.edata["he"] = edge_feat
+            graph.update_all(self.message, fn.sum("m", "neigh"))
+            rst = (1 + self.eps) * feat_dst + graph.dstdata["neigh"]
+            if self.apply_func is not None:
+                rst = self.apply_func(rst)
+            return rst

python/dgl/nn/pytorch/conv/gmmconv.py

@@ -0,0 +1,268 @@
+"""Torch Module for GMM Conv"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+from torch.nn import init
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import expand_as_pair
+from ..utils import Identity
+
+
+class GMMConv(nn.Module):
+    r"""Gaussian Mixture Model Convolution layer from `Geometric Deep
+    Learning on Graphs and Manifolds using Mixture Model CNNs
+    <https://arxiv.org/abs/1611.08402>`__
+
+    .. math::
+        u_{ij} &= f(x_i, x_j), x_j \in \mathcal{N}(i)
+
+        w_k(u) &= \exp\left(-\frac{1}{2}(u-\mu_k)^T \Sigma_k^{-1} (u - \mu_k)\right)
+
+        h_i^{l+1} &= \mathrm{aggregate}\left(\left\{\frac{1}{K}
+         \sum_{k}^{K} w_k(u_{ij}), \forall j\in \mathcal{N}(i)\right\}\right)
+
+    where :math:`u` denotes the pseudo-coordinates between a vertex and one of its neighbor,
+    computed using function :math:`f`, :math:`\Sigma_k^{-1}` and :math:`\mu_k` are
+    learnable parameters representing the covariance matrix and mean vector of a Gaussian kernel.
+
+    Parameters
+    ----------
+    in_feats : int
+        Number of input features; i.e., the number of dimensions of :math:`x_i`.
+    out_feats : int
+        Number of output features; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
+    dim : int
+        Dimensionality of pseudo-coordinte; i.e, the number of dimensions of :math:`u_{ij}`.
+    n_kernels : int
+        Number of kernels :math:`K`.
+    aggregator_type : str
+        Aggregator type (``sum``, ``mean``, ``max``). Default: ``sum``.
+    residual : bool
+        If True, use residual connection inside this layer. Default: ``False``.
+    bias : bool
+        If True, adds a learnable bias to the output. Default: ``True``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Default: ``False``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import GMMConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> conv = GMMConv(10, 2, 3, 2, 'mean')
+    >>> pseudo = th.ones(12, 3)
+    >>> res = conv(g, feat, pseudo)
+    >>> res
+    tensor([[-0.3462, -0.2654],
+            [-0.3462, -0.2654],
+            [-0.3462, -0.2654],
+            [-0.3462, -0.2654],
+            [-0.3462, -0.2654],
+            [-0.3462, -0.2654]], grad_fn=<AddBackward0>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
+    >>> u_fea = th.rand(2, 5)
+    >>> v_fea = th.rand(4, 10)
+    >>> pseudo = th.ones(5, 3)
+    >>> conv = GMMConv((10, 5), 2, 3, 2, 'mean')
+    >>> res = conv(g, (u_fea, v_fea), pseudo)
+    >>> res
+    tensor([[-0.1107, -0.1559],
+            [-0.1646, -0.2326],
+            [-0.1377, -0.1943],
+            [-0.1107, -0.1559]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        dim,
+        n_kernels,
+        aggregator_type="sum",
+        residual=False,
+        bias=True,
+        allow_zero_in_degree=False,
+    ):
+        super(GMMConv, self).__init__()
+        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
+        self._out_feats = out_feats
+        self._dim = dim
+        self._n_kernels = n_kernels
+        self._allow_zero_in_degree = allow_zero_in_degree
+        if aggregator_type == "sum":
+            self._reducer = fn.sum
+        elif aggregator_type == "mean":
+            self._reducer = fn.mean
+        elif aggregator_type == "max":
+            self._reducer = fn.max
+        else:
+            raise KeyError(
+                "Aggregator type {} not recognized.".format(aggregator_type)
+            )
+
+        self.mu = nn.Parameter(th.Tensor(n_kernels, dim))
+        self.inv_sigma = nn.Parameter(th.Tensor(n_kernels, dim))
+        self.fc = nn.Linear(
+            self._in_src_feats, n_kernels * out_feats, bias=False
+        )
+        if residual:
+            if self._in_dst_feats != out_feats:
+                self.res_fc = nn.Linear(
+                    self._in_dst_feats, out_feats, bias=False
+                )
+            else:
+                self.res_fc = Identity()
+        else:
+            self.register_buffer("res_fc", None)
+
+        if bias:
+            self.bias = nn.Parameter(th.Tensor(out_feats))
+        else:
+            self.register_buffer("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The fc parameters are initialized using Glorot uniform initialization
+        and the bias is initialized to be zero.
+        The mu weight is initialized using normal distribution and
+        inv_sigma is initialized with constant value 1.0.
+        """
+        gain = init.calculate_gain("relu")
+        init.xavier_normal_(self.fc.weight, gain=gain)
+        if isinstance(self.res_fc, nn.Linear):
+            init.xavier_normal_(self.res_fc.weight, gain=gain)
+        init.normal_(self.mu.data, 0, 0.1)
+        init.constant_(self.inv_sigma.data, 1)
+        if self.bias is not None:
+            init.zeros_(self.bias.data)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, pseudo):
+        """
+
+        Description
+        -----------
+        Compute Gaussian Mixture Model Convolution layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            If a single tensor is given, the input feature of shape :math:`(N, D_{in})` where
+            :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of tensors are given, the pair must contain two tensors of shape
+            :math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
+        pseudo : torch.Tensor
+            The pseudo coordinate tensor of shape :math:`(E, D_{u})` where
+            :math:`E` is the number of edges of the graph and :math:`D_{u}`
+            is the dimensionality of pseudo coordinate.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is the output feature size.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            feat_src, feat_dst = expand_as_pair(feat, graph)
+            graph.srcdata["h"] = self.fc(feat_src).view(
+                -1, self._n_kernels, self._out_feats
+            )
+            E = graph.num_edges()
+            # compute gaussian weight
+            gaussian = -0.5 * (
+                (
+                    pseudo.view(E, 1, self._dim)
+                    - self.mu.view(1, self._n_kernels, self._dim)
+                )
+                ** 2
+            )
+            gaussian = gaussian * (
+                self.inv_sigma.view(1, self._n_kernels, self._dim) ** 2
+            )
+            gaussian = th.exp(gaussian.sum(dim=-1, keepdim=True))  # (E, K, 1)
+            graph.edata["w"] = gaussian
+            graph.update_all(fn.u_mul_e("h", "w", "m"), self._reducer("m", "h"))
+            rst = graph.dstdata["h"].sum(1)
+            # residual connection
+            if self.res_fc is not None:
+                rst = rst + self.res_fc(feat_dst)
+            # bias
+            if self.bias is not None:
+                rst = rst + self.bias
+            return rst

python/dgl/nn/pytorch/conv/graphconv.py

@@ -0,0 +1,488 @@
+"""Torch modules for graph convolutions(GCN)."""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+from torch.nn import init
+
+from .... import function as fn
+from ....base import DGLError
+from ....convert import block_to_graph
+from ....heterograph import DGLBlock
+from ....transforms import reverse
+from ....utils import expand_as_pair
+
+
+class EdgeWeightNorm(nn.Module):
+    r"""This module normalizes positive scalar edge weights on a graph
+    following the form in `GCN <https://arxiv.org/abs/1609.02907>`__.
+
+    Mathematically, setting ``norm='both'`` yields the following normalization term:
+
+    .. math::
+      c_{ji} = (\sqrt{\sum_{k\in\mathcal{N}(j)}e_{jk}}\sqrt{\sum_{k\in\mathcal{N}(i)}e_{ki}})
+
+    And, setting ``norm='right'`` yields the following normalization term:
+
+    .. math::
+      c_{ji} = (\sum_{k\in\mathcal{N}(i)}e_{ki})
+
+    where :math:`e_{ji}` is the scalar weight on the edge from node :math:`j` to node :math:`i`.
+
+    The module returns the normalized weight :math:`e_{ji} / c_{ji}`.
+
+    Parameters
+    ----------
+    norm : str, optional
+        The normalizer as specified above. Default is `'both'`.
+    eps : float, optional
+        A small offset value in the denominator. Default is 0.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import EdgeWeightNorm, GraphConv
+
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> edge_weight = th.tensor([0.5, 0.6, 0.4, 0.7, 0.9, 0.1, 1, 1, 1, 1, 1, 1])
+    >>> norm = EdgeWeightNorm(norm='both')
+    >>> norm_edge_weight = norm(g, edge_weight)
+    >>> conv = GraphConv(10, 2, norm='none', weight=True, bias=True)
+    >>> res = conv(g, feat, edge_weight=norm_edge_weight)
+    >>> print(res)
+    tensor([[-1.1849, -0.7525],
+            [-1.3514, -0.8582],
+            [-1.2384, -0.7865],
+            [-1.9949, -1.2669],
+            [-1.3658, -0.8674],
+            [-0.8323, -0.5286]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(self, norm="both", eps=0.0):
+        super(EdgeWeightNorm, self).__init__()
+        self._norm = norm
+        self._eps = eps
+
+    def forward(self, graph, edge_weight):
+        r"""
+
+        Description
+        -----------
+        Compute normalized edge weight for the GCN model.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        edge_weight : torch.Tensor
+            Unnormalized scalar weights on the edges.
+            The shape is expected to be :math:`(|E|)`.
+
+        Returns
+        -------
+        torch.Tensor
+            The normalized edge weight.
+
+        Raises
+        ------
+        DGLError
+            Case 1:
+            The edge weight is multi-dimensional. Currently this module
+            only supports a scalar weight on each edge.
+
+            Case 2:
+            The edge weight has non-positive values with ``norm='both'``.
+            This will trigger square root and division by a non-positive number.
+        """
+        with graph.local_scope():
+            if isinstance(graph, DGLBlock):
+                graph = block_to_graph(graph)
+            if len(edge_weight.shape) > 1:
+                raise DGLError(
+                    "Currently the normalization is only defined "
+                    "on scalar edge weight. Please customize the "
+                    "normalization for your high-dimensional weights."
+                )
+            if self._norm == "both" and th.any(edge_weight <= 0).item():
+                raise DGLError(
+                    'Non-positive edge weight detected with `norm="both"`. '
+                    "This leads to square root of zero or negative values."
+                )
+
+            dev = graph.device
+            dtype = edge_weight.dtype
+            graph.srcdata["_src_out_w"] = th.ones(
+                graph.number_of_src_nodes(), dtype=dtype, device=dev
+            )
+            graph.dstdata["_dst_in_w"] = th.ones(
+                graph.number_of_dst_nodes(), dtype=dtype, device=dev
+            )
+            graph.edata["_edge_w"] = edge_weight
+
+            if self._norm == "both":
+                reversed_g = reverse(graph)
+                reversed_g.edata["_edge_w"] = edge_weight
+                reversed_g.update_all(
+                    fn.copy_e("_edge_w", "m"), fn.sum("m", "out_weight")
+                )
+                degs = reversed_g.dstdata["out_weight"] + self._eps
+                norm = th.pow(degs, -0.5)
+                graph.srcdata["_src_out_w"] = norm
+
+            if self._norm != "none":
+                graph.update_all(
+                    fn.copy_e("_edge_w", "m"), fn.sum("m", "in_weight")
+                )
+                degs = graph.dstdata["in_weight"] + self._eps
+                if self._norm == "both":
+                    norm = th.pow(degs, -0.5)
+                else:
+                    norm = 1.0 / degs
+                graph.dstdata["_dst_in_w"] = norm
+
+            graph.apply_edges(
+                lambda e: {
+                    "_norm_edge_weights": e.src["_src_out_w"]
+                    * e.dst["_dst_in_w"]
+                    * e.data["_edge_w"]
+                }
+            )
+            return graph.edata["_norm_edge_weights"]
+
+
+# pylint: disable=W0235
+class GraphConv(nn.Module):
+    r"""Graph convolutional layer from `Semi-Supervised Classification with Graph Convolutional
+    Networks <https://arxiv.org/abs/1609.02907>`__
+
+    Mathematically it is defined as follows:
+
+    .. math::
+      h_i^{(l+1)} = \sigma(b^{(l)} + \sum_{j\in\mathcal{N}(i)}\frac{1}{c_{ji}}h_j^{(l)}W^{(l)})
+
+    where :math:`\mathcal{N}(i)` is the set of neighbors of node :math:`i`,
+    :math:`c_{ji}` is the product of the square root of node degrees
+    (i.e.,  :math:`c_{ji} = \sqrt{|\mathcal{N}(j)|}\sqrt{|\mathcal{N}(i)|}`),
+    and :math:`\sigma` is an activation function.
+
+    If a weight tensor on each edge is provided, the weighted graph convolution is defined as:
+
+    .. math::
+      h_i^{(l+1)} = \sigma(b^{(l)} + \sum_{j\in\mathcal{N}(i)}\frac{e_{ji}}{c_{ji}}h_j^{(l)}W^{(l)})
+
+    where :math:`e_{ji}` is the scalar weight on the edge from node :math:`j` to node :math:`i`.
+    This is NOT equivalent to the weighted graph convolutional network formulation in the paper.
+
+    To customize the normalization term :math:`c_{ji}`, one can first set ``norm='none'`` for
+    the model, and send the pre-normalized :math:`e_{ji}` to the forward computation. We provide
+    :class:`~dgl.nn.pytorch.EdgeWeightNorm` to normalize scalar edge weight following the GCN paper.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
+    out_feats : int
+        Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
+    norm : str, optional
+        How to apply the normalizer.  Can be one of the following values:
+
+        * ``right``, to divide the aggregated messages by each node's in-degrees,
+          which is equivalent to averaging the received messages.
+
+        * ``none``, where no normalization is applied.
+
+        * ``both`` (default), where the messages are scaled with :math:`1/c_{ji}` above, equivalent
+          to symmetric normalization.
+
+        * ``left``, to divide the messages sent out from each node by its out-degrees,
+          equivalent to random walk normalization.
+    weight : bool, optional
+        If True, apply a linear layer. Otherwise, aggregating the messages
+        without a weight matrix.
+    bias : bool, optional
+        If True, adds a learnable bias to the output. Default: ``True``.
+    activation : callable activation function/layer or None, optional
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Default: ``False``.
+
+    Attributes
+    ----------
+    weight : torch.Tensor
+        The learnable weight tensor.
+    bias : torch.Tensor
+        The learnable bias tensor.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import GraphConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> conv = GraphConv(10, 2, norm='both', weight=True, bias=True)
+    >>> res = conv(g, feat)
+    >>> print(res)
+    tensor([[ 1.3326, -0.2797],
+            [ 1.4673, -0.3080],
+            [ 1.3326, -0.2797],
+            [ 1.6871, -0.3541],
+            [ 1.7711, -0.3717],
+            [ 1.0375, -0.2178]], grad_fn=<AddBackward0>)
+    >>> # allow_zero_in_degree example
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> conv = GraphConv(10, 2, norm='both', weight=True, bias=True, allow_zero_in_degree=True)
+    >>> res = conv(g, feat)
+    >>> print(res)
+    tensor([[-0.2473, -0.4631],
+            [-0.3497, -0.6549],
+            [-0.3497, -0.6549],
+            [-0.4221, -0.7905],
+            [-0.3497, -0.6549],
+            [ 0.0000,  0.0000]], grad_fn=<AddBackward0>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('_U', '_E', '_V') : (u, v)})
+    >>> u_fea = th.rand(2, 5)
+    >>> v_fea = th.rand(4, 5)
+    >>> conv = GraphConv(5, 2, norm='both', weight=True, bias=True)
+    >>> res = conv(g, (u_fea, v_fea))
+    >>> res
+    tensor([[-0.2994,  0.6106],
+            [-0.4482,  0.5540],
+            [-0.5287,  0.8235],
+            [-0.2994,  0.6106]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        norm="both",
+        weight=True,
+        bias=True,
+        activation=None,
+        allow_zero_in_degree=False,
+    ):
+        super(GraphConv, self).__init__()
+        if norm not in ("none", "both", "right", "left"):
+            raise DGLError(
+                'Invalid norm value. Must be either "none", "both", "right" or "left".'
+                ' But got "{}".'.format(norm)
+            )
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self._norm = norm
+        self._allow_zero_in_degree = allow_zero_in_degree
+
+        if weight:
+            self.weight = nn.Parameter(th.Tensor(in_feats, out_feats))
+        else:
+            self.register_parameter("weight", None)
+
+        if bias:
+            self.bias = nn.Parameter(th.Tensor(out_feats))
+        else:
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+        self._activation = activation
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The model parameters are initialized as in the
+        `original implementation <https://github.com/tkipf/gcn/blob/master/gcn/layers.py>`__
+        where the weight :math:`W^{(l)}` is initialized using Glorot uniform initialization
+        and the bias is initialized to be zero.
+
+        """
+        if self.weight is not None:
+            init.xavier_uniform_(self.weight)
+        if self.bias is not None:
+            init.zeros_(self.bias)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, weight=None, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute graph convolution.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, it represents the input feature of shape
+            :math:`(N, D_{in})`
+            where :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, which is the case for bipartite graph, the pair
+            must contain two tensors of shape :math:`(N_{in}, D_{in_{src}})` and
+            :math:`(N_{out}, D_{in_{dst}})`.
+        weight : torch.Tensor, optional
+            Optional external weight tensor.
+        edge_weight : torch.Tensor, optional
+            Optional tensor on the edge. If given, the convolution will weight
+            with regard to the message.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature
+
+        Raises
+        ------
+        DGLError
+            Case 1:
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+
+            Case 2:
+            External weight is provided while at the same time the module
+            has defined its own weight parameter.
+
+        Note
+        ----
+        * Input shape: :math:`(N, *, \text{in_feats})` where * means any number of additional
+          dimensions, :math:`N` is the number of nodes.
+        * Output shape: :math:`(N, *, \text{out_feats})` where all but the last dimension are
+          the same shape as the input.
+        * Weight shape: :math:`(\text{in_feats}, \text{out_feats})`.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+            aggregate_fn = fn.copy_u("h", "m")
+            if edge_weight is not None:
+                assert edge_weight.shape[0] == graph.num_edges()
+                graph.edata["_edge_weight"] = edge_weight
+                aggregate_fn = fn.u_mul_e("h", "_edge_weight", "m")
+
+            # (BarclayII) For RGCN on heterogeneous graphs we need to support GCN on bipartite.
+            feat_src, feat_dst = expand_as_pair(feat, graph)
+            if self._norm in ["left", "both"]:
+                degs = graph.out_degrees().to(feat_src).clamp(min=1)
+                if self._norm == "both":
+                    norm = th.pow(degs, -0.5)
+                else:
+                    norm = 1.0 / degs
+                shp = norm.shape + (1,) * (feat_src.dim() - 1)
+                norm = th.reshape(norm, shp)
+                feat_src = feat_src * norm
+
+            if weight is not None:
+                if self.weight is not None:
+                    raise DGLError(
+                        "External weight is provided while at the same time the"
+                        " module has defined its own weight parameter. Please"
+                        " create the module with flag weight=False."
+                    )
+            else:
+                weight = self.weight
+
+            if self._in_feats > self._out_feats:
+                # mult W first to reduce the feature size for aggregation.
+                if weight is not None:
+                    feat_src = th.matmul(feat_src, weight)
+                graph.srcdata["h"] = feat_src
+                graph.update_all(aggregate_fn, fn.sum(msg="m", out="h"))
+                rst = graph.dstdata["h"]
+            else:
+                # aggregate first then mult W
+                graph.srcdata["h"] = feat_src
+                graph.update_all(aggregate_fn, fn.sum(msg="m", out="h"))
+                rst = graph.dstdata["h"]
+                if weight is not None:
+                    rst = th.matmul(rst, weight)
+
+            if self._norm in ["right", "both"]:
+                degs = graph.in_degrees().to(feat_dst).clamp(min=1)
+                if self._norm == "both":
+                    norm = th.pow(degs, -0.5)
+                else:
+                    norm = 1.0 / degs
+                shp = norm.shape + (1,) * (feat_dst.dim() - 1)
+                norm = th.reshape(norm, shp)
+                rst = rst * norm
+
+            if self.bias is not None:
+                rst = rst + self.bias
+
+            if self._activation is not None:
+                rst = self._activation(rst)
+
+            return rst
+
+    def extra_repr(self):
+        """Set the extra representation of the module,
+        which will come into effect when printing the model.
+        """
+        summary = "in={_in_feats}, out={_out_feats}"
+        summary += ", normalization={_norm}"
+        if "_activation" in self.__dict__:
+            summary += ", activation={_activation}"
+        return summary.format(**self.__dict__)

python/dgl/nn/pytorch/conv/grouprevres.py

@@ -0,0 +1,257 @@
+"""Torch module for grouped reversible residual connections for GNNs"""
+# pylint: disable= no-member, arguments-differ, invalid-name, C0116, R1728
+from copy import deepcopy
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class InvertibleCheckpoint(torch.autograd.Function):
+    r"""Extension of torch.autograd"""
+
+    @staticmethod
+    def forward(ctx, fn, fn_inverse, num_inputs, *inputs_and_weights):
+        ctx.fn = fn
+        ctx.fn_inverse = fn_inverse
+        ctx.weights = inputs_and_weights[num_inputs:]
+        inputs = inputs_and_weights[:num_inputs]
+        ctx.input_requires_grad = []
+
+        with torch.no_grad():
+            # Make a detached copy, which shares the storage
+            x = []
+            for element in inputs:
+                if isinstance(element, torch.Tensor):
+                    x.append(element.detach())
+                    ctx.input_requires_grad.append(element.requires_grad)
+                else:
+                    x.append(element)
+                    ctx.input_requires_grad.append(None)
+            # Detach the output, which then allows discarding the intermediary results
+            outputs = ctx.fn(*x).detach_()
+
+        # clear memory of input node features
+        inputs[1].untyped_storage().resize_(0)
+
+        # store for backward pass
+        ctx.inputs = [inputs]
+        ctx.outputs = [outputs]
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        if not torch.autograd._is_checkpoint_valid():
+            raise RuntimeError(
+                "InvertibleCheckpoint is not compatible with .grad(), \
+                               please use .backward() if possible"
+            )
+        # retrieve input and output tensor nodes
+        if len(ctx.outputs) == 0:
+            raise RuntimeError(
+                "Trying to perform backward on the InvertibleCheckpoint \
+                               for more than once."
+            )
+        inputs = ctx.inputs.pop()
+        outputs = ctx.outputs.pop()
+
+        # reconstruct input node features
+        with torch.no_grad():
+            # inputs[0] is DGLGraph and inputs[1] is input node features
+            inputs_inverted = ctx.fn_inverse(
+                *((inputs[0], outputs) + inputs[2:])
+            )
+            # clear memory of outputs
+            outputs.untyped_storage().resize_(0)
+
+            x = inputs[1]
+            x.untyped_storage().resize_(int(np.prod(x.size())))
+            x.set_(inputs_inverted)
+
+        # compute gradients
+        with torch.set_grad_enabled(True):
+            detached_inputs = []
+            for i, element in enumerate(inputs):
+                if isinstance(element, torch.Tensor):
+                    element = element.detach()
+                    element.requires_grad = ctx.input_requires_grad[i]
+                detached_inputs.append(element)
+
+            detached_inputs = tuple(detached_inputs)
+            temp_output = ctx.fn(*detached_inputs)
+
+        filtered_detached_inputs = tuple(
+            filter(
+                lambda x: getattr(x, "requires_grad", False), detached_inputs
+            )
+        )
+        gradients = torch.autograd.grad(
+            outputs=(temp_output,),
+            inputs=filtered_detached_inputs + ctx.weights,
+            grad_outputs=grad_outputs,
+        )
+
+        input_gradients = []
+        i = 0
+        for rg in ctx.input_requires_grad:
+            if rg:
+                input_gradients.append(gradients[i])
+                i += 1
+            else:
+                input_gradients.append(None)
+
+        gradients = tuple(input_gradients) + gradients[-len(ctx.weights) :]
+
+        return (None, None, None) + gradients
+
+
+class GroupRevRes(nn.Module):
+    r"""Grouped reversible residual connections for GNNs, as introduced in
+    `Training Graph Neural Networks with 1000 Layers <https://arxiv.org/abs/2106.07476>`__
+
+    It uniformly partitions an input node feature :math:`X` into :math:`C` groups
+    :math:`X_1, X_2, \cdots, X_C` across the channel dimension. Besides, it makes
+    :math:`C` copies of the input GNN module :math:`f_{w1}, \cdots, f_{wC}`. In the
+    forward pass, each GNN module only takes the corresponding group of node features.
+
+    The output node representations :math:`X^{'}` are computed as follows.
+
+    .. math::
+
+        X_0^{'} = \sum_{i=2}^{C}X_i
+
+        X_i^{'} = f_{wi}(X_{i-1}^{'}, g, U) + X_i, i\in\{1,\cdots,C\}
+
+        X^{'} = X_1^{'} \, \Vert \, \ldots \, \Vert \, X_C^{'}
+
+    where :math:`g` is the input graph, :math:`U` is arbitrary additional input arguments like
+    edge features, and :math:`\, \Vert \,` is concatenation.
+
+    Parameters
+    ----------
+    gnn_module : nn.Module
+        GNN module for message passing. :attr:`GroupRevRes` will clone the module for
+        :attr:`groups`-1 number of times, yielding :attr:`groups` copies in total.
+        The input and output node representation size need to be the same. Its forward
+        function needs to take a DGLGraph and the associated input node features in order,
+        optionally followed by additional arguments like edge features.
+    groups : int, optional
+        The number of groups.
+
+    Examples
+    --------
+
+    >>> import dgl
+    >>> import torch
+    >>> import torch.nn as nn
+    >>> from dgl.nn import GraphConv, GroupRevRes
+
+    >>> class GNNLayer(nn.Module):
+    ...     def __init__(self, feats, dropout=0.2):
+    ...         super(GNNLayer, self).__init__()
+    ...         # Use BatchNorm and dropout to prevent gradient vanishing
+    ...         # In particular if you use a large number of GNN layers
+    ...         self.norm = nn.BatchNorm1d(feats)
+    ...         self.conv = GraphConv(feats, feats)
+    ...         self.dropout = nn.Dropout(dropout)
+    ...
+    ...     def forward(self, g, x):
+    ...         x = self.norm(x)
+    ...         x = self.dropout(x)
+    ...         return self.conv(g, x)
+
+    >>> num_nodes = 5
+    >>> num_edges = 20
+    >>> feats = 32
+    >>> groups = 2
+    >>> g = dgl.rand_graph(num_nodes, num_edges)
+    >>> x = torch.randn(num_nodes, feats)
+    >>> conv = GNNLayer(feats // groups)
+    >>> model = GroupRevRes(conv, groups)
+    >>> out = model(g, x)
+    """
+
+    def __init__(self, gnn_module, groups=2):
+        super(GroupRevRes, self).__init__()
+        self.gnn_modules = nn.ModuleList()
+        for i in range(groups):
+            if i == 0:
+                self.gnn_modules.append(gnn_module)
+            else:
+                self.gnn_modules.append(deepcopy(gnn_module))
+        self.groups = groups
+
+    def _forward(self, g, x, *args):
+        xs = torch.chunk(x, self.groups, dim=-1)
+
+        if len(args) == 0:
+            args_chunks = [()] * self.groups
+        else:
+            chunked_args = list(
+                map(lambda arg: torch.chunk(arg, self.groups, dim=-1), args)
+            )
+            args_chunks = list(zip(*chunked_args))
+        y_in = sum(xs[1:])
+
+        ys = []
+        for i in range(self.groups):
+            y_in = xs[i] + self.gnn_modules[i](g, y_in, *args_chunks[i])
+            ys.append(y_in)
+
+        out = torch.cat(ys, dim=-1)
+
+        return out
+
+    def _inverse(self, g, y, *args):
+        ys = torch.chunk(y, self.groups, dim=-1)
+
+        if len(args) == 0:
+            args_chunks = [()] * self.groups
+        else:
+            chunked_args = list(
+                map(lambda arg: torch.chunk(arg, self.groups, dim=-1), args)
+            )
+            args_chunks = list(zip(*chunked_args))
+
+        xs = []
+        for i in range(self.groups - 1, -1, -1):
+            if i != 0:
+                y_in = ys[i - 1]
+            else:
+                y_in = sum(xs)
+
+            x = ys[i] - self.gnn_modules[i](g, y_in, *args_chunks[i])
+            xs.append(x)
+
+        x = torch.cat(xs[::-1], dim=-1)
+
+        return x
+
+    def forward(self, g, x, *args):
+        r"""Apply the GNN module with grouped reversible residual connection.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        x : torch.Tensor
+            The input feature of shape :math:`(N, D_{in})`, where :math:`D_{in}` is size
+            of input feature, :math:`N` is the number of nodes.
+        args
+            Additional arguments to pass to :attr:`gnn_module`.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{in})`.
+        """
+        args = (g, x) + args
+        y = InvertibleCheckpoint.apply(
+            self._forward,
+            self._inverse,
+            len(args),
+            *(args + tuple([p for p in self.parameters() if p.requires_grad]))
+        )
+
+        return y

python/dgl/nn/pytorch/conv/hgtconv.py

@@ -0,0 +1,201 @@
+"""Heterogeneous Graph Transformer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import math
+
+import torch
+import torch.nn as nn
+
+from .... import function as fn
+from ..linear import TypedLinear
+from ..softmax import edge_softmax
+
+
+class HGTConv(nn.Module):
+    r"""Heterogeneous graph transformer convolution from `Heterogeneous Graph Transformer
+    <https://arxiv.org/abs/2003.01332>`__
+
+    Given a graph :math:`G(V, E)` and input node features :math:`H^{(l-1)}`,
+    it computes the new node features as follows:
+
+    Compute a multi-head attention score for each edge :math:`(s, e, t)` in the graph:
+
+    .. math::
+
+      Attention(s, e, t) = \text{Softmax}\left(||_{i\in[1,h]}ATT-head^i(s, e, t)\right) \\
+      ATT-head^i(s, e, t) = \left(K^i(s)W^{ATT}_{\phi(e)}Q^i(t)^{\top}\right)\cdot
+        \frac{\mu_{(\tau(s),\phi(e),\tau(t)}}{\sqrt{d}} \\
+      K^i(s) = \text{K-Linear}^i_{\tau(s)}(H^{(l-1)}[s]) \\
+      Q^i(t) = \text{Q-Linear}^i_{\tau(t)}(H^{(l-1)}[t]) \\
+
+    Compute the message to send on each edge :math:`(s, e, t)`:
+
+    .. math::
+
+      Message(s, e, t) = ||_{i\in[1, h]} MSG-head^i(s, e, t) \\
+      MSG-head^i(s, e, t) = \text{M-Linear}^i_{\tau(s)}(H^{(l-1)}[s])W^{MSG}_{\phi(e)} \\
+
+    Send messages to target nodes :math:`t` and aggregate:
+
+    .. math::
+
+      \tilde{H}^{(l)}[t] = \sum_{\forall s\in \mathcal{N}(t)}\left( Attention(s,e,t)
+      \cdot Message(s,e,t)\right)
+
+    Compute new node features:
+
+    .. math::
+
+      H^{(l)}[t]=\text{A-Linear}_{\tau(t)}(\sigma(\tilde(H)^{(l)}[t])) + H^{(l-1)}[t]
+
+    Parameters
+    ----------
+    in_size : int
+        Input node feature size.
+    head_size : int
+        Output head size. The output node feature size is ``head_size * num_heads``.
+    num_heads : int
+        Number of heads. The output node feature size is ``head_size * num_heads``.
+    num_ntypes : int
+        Number of node types.
+    num_etypes : int
+        Number of edge types.
+    dropout : optional, float
+        Dropout rate.
+    use_norm : optiona, bool
+        If true, apply a layer norm on the output node feature.
+
+    Examples
+    --------
+    """
+
+    def __init__(
+        self,
+        in_size,
+        head_size,
+        num_heads,
+        num_ntypes,
+        num_etypes,
+        dropout=0.2,
+        use_norm=False,
+    ):
+        super().__init__()
+        self.in_size = in_size
+        self.head_size = head_size
+        self.num_heads = num_heads
+        self.sqrt_d = math.sqrt(head_size)
+        self.use_norm = use_norm
+
+        self.linear_k = TypedLinear(in_size, head_size * num_heads, num_ntypes)
+        self.linear_q = TypedLinear(in_size, head_size * num_heads, num_ntypes)
+        self.linear_v = TypedLinear(in_size, head_size * num_heads, num_ntypes)
+        self.linear_a = TypedLinear(
+            head_size * num_heads, head_size * num_heads, num_ntypes
+        )
+
+        self.relation_pri = nn.ParameterList(
+            [nn.Parameter(torch.ones(num_etypes)) for i in range(num_heads)]
+        )
+        self.relation_att = nn.ModuleList(
+            [
+                TypedLinear(head_size, head_size, num_etypes)
+                for i in range(num_heads)
+            ]
+        )
+        self.relation_msg = nn.ModuleList(
+            [
+                TypedLinear(head_size, head_size, num_etypes)
+                for i in range(num_heads)
+            ]
+        )
+        self.skip = nn.Parameter(torch.ones(num_ntypes))
+        self.drop = nn.Dropout(dropout)
+        if use_norm:
+            self.norm = nn.LayerNorm(head_size * num_heads)
+        if in_size != head_size * num_heads:
+            self.residual_w = nn.Parameter(
+                torch.Tensor(in_size, head_size * num_heads)
+            )
+            nn.init.xavier_uniform_(self.residual_w)
+
+    def forward(self, g, x, ntype, etype, *, presorted=False):
+        """Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The input graph.
+        x : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        ntype : torch.Tensor
+            An 1D integer tensor of node types. Shape: :math:`(|V|,)`.
+        etype : torch.Tensor
+            An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
+        presorted : bool, optional
+            Whether *both* the nodes and the edges of the input graph have been sorted by
+            their types. Forward on pre-sorted graph may be faster. Graphs created by
+            :func:`~dgl.to_homogeneous` automatically satisfy the condition.
+            Also see :func:`~dgl.reorder_graph` for manually reordering the nodes and edges.
+
+        Returns
+        -------
+        torch.Tensor
+            New node features. Shape: :math:`(|V|, D_{head} * N_{head})`.
+        """
+        self.presorted = presorted
+        if g.is_block:
+            x_src = x
+            x_dst = x[: g.num_dst_nodes()]
+            srcntype = ntype
+            dstntype = ntype[: g.num_dst_nodes()]
+        else:
+            x_src = x
+            x_dst = x
+            srcntype = ntype
+            dstntype = ntype
+        with g.local_scope():
+            k = self.linear_k(x_src, srcntype, presorted).view(
+                -1, self.num_heads, self.head_size
+            )
+            q = self.linear_q(x_dst, dstntype, presorted).view(
+                -1, self.num_heads, self.head_size
+            )
+            v = self.linear_v(x_src, srcntype, presorted).view(
+                -1, self.num_heads, self.head_size
+            )
+            g.srcdata["k"] = k
+            g.dstdata["q"] = q
+            g.srcdata["v"] = v
+            g.edata["etype"] = etype
+            g.apply_edges(self.message)
+            g.edata["m"] = g.edata["m"] * edge_softmax(
+                g, g.edata["a"]
+            ).unsqueeze(-1)
+            g.update_all(fn.copy_e("m", "m"), fn.sum("m", "h"))
+            h = g.dstdata["h"].view(-1, self.num_heads * self.head_size)
+            # target-specific aggregation
+            h = self.drop(self.linear_a(h, dstntype, presorted))
+            alpha = torch.sigmoid(self.skip[dstntype]).unsqueeze(-1)
+            if x_dst.shape != h.shape:
+                h = h * alpha + (x_dst @ self.residual_w) * (1 - alpha)
+            else:
+                h = h * alpha + x_dst * (1 - alpha)
+            if self.use_norm:
+                h = self.norm(h)
+            return h
+
+    def message(self, edges):
+        """Message function."""
+        a, m = [], []
+        etype = edges.data["etype"]
+        k = torch.unbind(edges.src["k"], dim=1)
+        q = torch.unbind(edges.dst["q"], dim=1)
+        v = torch.unbind(edges.src["v"], dim=1)
+        for i in range(self.num_heads):
+            kw = self.relation_att[i](k[i], etype, self.presorted)  # (E, O)
+            a.append(
+                (kw * q[i]).sum(-1) * self.relation_pri[i][etype] / self.sqrt_d
+            )  # (E,)
+            m.append(
+                self.relation_msg[i](v[i], etype, self.presorted)
+            )  # (E, O)
+        return {"a": torch.stack(a, dim=1), "m": torch.stack(m, dim=1)}

python/dgl/nn/pytorch/conv/nnconv.py

@@ -0,0 +1,187 @@
+"""Torch Module for NNConv layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+from torch.nn import init
+
+from .... import function as fn
+from ....utils import expand_as_pair
+from ..utils import Identity
+
+
+class NNConv(nn.Module):
+    r"""Graph Convolution layer from `Neural Message Passing
+    for Quantum Chemistry <https://arxiv.org/pdf/1704.01212.pdf>`__
+
+    .. math::
+        h_{i}^{l+1} = h_{i}^{l} + \mathrm{aggregate}\left(\left\{
+        f_\Theta (e_{ij}) \cdot h_j^{l}, j\in \mathcal{N}(i) \right\}\right)
+
+    where :math:`e_{ij}` is the edge feature, :math:`f_\Theta` is a function
+    with learnable parameters.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
+        NNConv can be applied on homogeneous graph and unidirectional
+        `bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
+        If the layer is to be applied on a unidirectional bipartite graph, ``in_feats``
+        specifies the input feature size on both the source and destination nodes.  If
+        a scalar is given, the source and destination node feature size would take the
+        same value.
+    out_feats : int
+        Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
+    edge_func : callable activation function/layer
+        Maps each edge feature to a vector of shape
+        ``(in_feats * out_feats)`` as weight to compute
+        messages.
+        Also is the :math:`f_\Theta` in the formula.
+    aggregator_type : str
+        Aggregator type to use (``sum``, ``mean`` or ``max``).
+    residual : bool, optional
+        If True, use residual connection. Default: ``False``.
+    bias : bool, optional
+        If True, adds a learnable bias to the output. Default: ``True``.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import NNConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> lin = th.nn.Linear(5, 20)
+    >>> def edge_func(efeat):
+    ...     return lin(efeat)
+    >>> efeat = th.ones(6+6, 5)
+    >>> conv = NNConv(10, 2, edge_func, 'mean')
+    >>> res = conv(g, feat, efeat)
+    >>> res
+    tensor([[-1.5243, -0.2719],
+            [-1.5243, -0.2719],
+            [-1.5243, -0.2719],
+            [-1.5243, -0.2719],
+            [-1.5243, -0.2719],
+            [-1.5243, -0.2719]], grad_fn=<AddBackward0>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
+    >>> u_feat = th.tensor(np.random.rand(2, 10).astype(np.float32))
+    >>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
+    >>> conv = NNConv(10, 2, edge_func, 'mean')
+    >>> efeat = th.ones(5, 5)
+    >>> res = conv(g, (u_feat, v_feat), efeat)
+    >>> res
+    tensor([[-0.6568,  0.5042],
+            [ 0.9089, -0.5352],
+            [ 0.1261, -0.0155],
+            [-0.6568,  0.5042]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        edge_func,
+        aggregator_type="mean",
+        residual=False,
+        bias=True,
+    ):
+        super(NNConv, self).__init__()
+        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
+        self._out_feats = out_feats
+        self.edge_func = edge_func
+        if aggregator_type == "sum":
+            self.reducer = fn.sum
+        elif aggregator_type == "mean":
+            self.reducer = fn.mean
+        elif aggregator_type == "max":
+            self.reducer = fn.max
+        else:
+            raise KeyError(
+                "Aggregator type {} not recognized: ".format(aggregator_type)
+            )
+        self._aggre_type = aggregator_type
+        if residual:
+            if self._in_dst_feats != out_feats:
+                self.res_fc = nn.Linear(
+                    self._in_dst_feats, out_feats, bias=False
+                )
+            else:
+                self.res_fc = Identity()
+        else:
+            self.register_buffer("res_fc", None)
+        if bias:
+            self.bias = nn.Parameter(th.Tensor(out_feats))
+        else:
+            self.register_buffer("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The model parameters are initialized using Glorot uniform initialization
+        and the bias is initialized to be zero.
+        """
+        gain = init.calculate_gain("relu")
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+        if isinstance(self.res_fc, nn.Linear):
+            nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
+
+    def forward(self, graph, feat, efeat):
+        r"""Compute MPNN Graph Convolution layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor or pair of torch.Tensor
+            The input feature of shape :math:`(N, D_{in})` where :math:`N`
+            is the number of nodes of the graph and :math:`D_{in}` is the
+            input feature size.
+        efeat : torch.Tensor
+            The edge feature of shape :math:`(E, *)`, which should fit the input
+            shape requirement of ``edge_func``. :math:`E` is the number of edges
+            of the graph.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is the output feature size.
+        """
+        with graph.local_scope():
+            feat_src, feat_dst = expand_as_pair(feat, graph)
+
+            # (n, d_in, 1)
+            graph.srcdata["h"] = feat_src.unsqueeze(-1)
+            # (n, d_in, d_out)
+            graph.edata["w"] = self.edge_func(efeat).view(
+                -1, self._in_src_feats, self._out_feats
+            )
+            # (n, d_in, d_out)
+            graph.update_all(
+                fn.u_mul_e("h", "w", "m"), self.reducer("m", "neigh")
+            )
+            rst = graph.dstdata["neigh"].sum(dim=1)  # (n, d_out)
+            # residual connection
+            if self.res_fc is not None:
+                rst = rst + self.res_fc(feat_dst)
+            # bias
+            if self.bias is not None:
+                rst = rst + self.bias
+            return rst

python/dgl/nn/pytorch/conv/pnaconv.py

@@ -0,0 +1,347 @@
+"""Torch Module for Principal Neighbourhood Aggregation Convolution Layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+def aggregate_mean(h):
+    """mean aggregation"""
+    return torch.mean(h, dim=1)
+
+
+def aggregate_max(h):
+    """max aggregation"""
+    return torch.max(h, dim=1)[0]
+
+
+def aggregate_min(h):
+    """min aggregation"""
+    return torch.min(h, dim=1)[0]
+
+
+def aggregate_sum(h):
+    """sum aggregation"""
+    return torch.sum(h, dim=1)
+
+
+def aggregate_std(h):
+    """standard deviation aggregation"""
+    return torch.sqrt(aggregate_var(h) + 1e-30)
+
+
+def aggregate_var(h):
+    """variance aggregation"""
+    h_mean_squares = torch.mean(h * h, dim=1)
+    h_mean = torch.mean(h, dim=1)
+    var = torch.relu(h_mean_squares - h_mean * h_mean)
+    return var
+
+
+def _aggregate_moment(h, n):
+    """moment aggregation: for each node (E[(X-E[X])^n])^{1/n}"""
+    h_mean = torch.mean(h, dim=1, keepdim=True)
+    h_n = torch.mean(torch.pow(h - h_mean, n), dim=1)
+    rooted_h_n = torch.sign(h_n) * torch.pow(torch.abs(h_n) + 1e-30, 1.0 / n)
+    return rooted_h_n
+
+
+def aggregate_moment_3(h):
+    """moment aggregation with n=3"""
+    return _aggregate_moment(h, n=3)
+
+
+def aggregate_moment_4(h):
+    """moment aggregation with n=4"""
+    return _aggregate_moment(h, n=4)
+
+
+def aggregate_moment_5(h):
+    """moment aggregation with n=5"""
+    return _aggregate_moment(h, n=5)
+
+
+def scale_identity(h):
+    """identity scaling (no scaling operation)"""
+    return h
+
+
+def scale_amplification(h, D, delta):
+    """amplification scaling"""
+    return h * (np.log(D + 1) / delta)
+
+
+def scale_attenuation(h, D, delta):
+    """attenuation scaling"""
+    return h * (delta / np.log(D + 1))
+
+
+AGGREGATORS = {
+    "mean": aggregate_mean,
+    "sum": aggregate_sum,
+    "max": aggregate_max,
+    "min": aggregate_min,
+    "std": aggregate_std,
+    "var": aggregate_var,
+    "moment3": aggregate_moment_3,
+    "moment4": aggregate_moment_4,
+    "moment5": aggregate_moment_5,
+}
+SCALERS = {
+    "identity": scale_identity,
+    "amplification": scale_amplification,
+    "attenuation": scale_attenuation,
+}
+
+
+class PNAConvTower(nn.Module):
+    """A single PNA tower in PNA layers"""
+
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        aggregators,
+        scalers,
+        delta,
+        dropout=0.0,
+        edge_feat_size=0,
+    ):
+        super(PNAConvTower, self).__init__()
+        self.in_size = in_size
+        self.out_size = out_size
+        self.aggregators = aggregators
+        self.scalers = scalers
+        self.delta = delta
+        self.edge_feat_size = edge_feat_size
+
+        self.M = nn.Linear(2 * in_size + edge_feat_size, in_size)
+        self.U = nn.Linear(
+            (len(aggregators) * len(scalers) + 1) * in_size, out_size
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.batchnorm = nn.BatchNorm1d(out_size)
+
+    def reduce_func(self, nodes):
+        """reduce function for PNA layer:
+        tensordot of multiple aggregation and scaling operations"""
+        msg = nodes.mailbox["msg"]
+        degree = msg.size(1)
+        h = torch.cat(
+            [AGGREGATORS[agg](msg) for agg in self.aggregators], dim=1
+        )
+        h = torch.cat(
+            [
+                SCALERS[scaler](h, D=degree, delta=self.delta)
+                if scaler != "identity"
+                else h
+                for scaler in self.scalers
+            ],
+            dim=1,
+        )
+        return {"h_neigh": h}
+
+    def message(self, edges):
+        """message function for PNA layer"""
+        if self.edge_feat_size > 0:
+            f = torch.cat(
+                [edges.src["h"], edges.dst["h"], edges.data["a"]], dim=-1
+            )
+        else:
+            f = torch.cat([edges.src["h"], edges.dst["h"]], dim=-1)
+        return {"msg": self.M(f)}
+
+    def forward(self, graph, node_feat, edge_feat=None):
+        """compute the forward pass of a single tower in PNA convolution layer"""
+        # calculate graph normalization factors
+        snorm_n = torch.cat(
+            [
+                torch.ones(N, 1).to(node_feat) / N
+                for N in graph.batch_num_nodes()
+            ],
+            dim=0,
+        ).sqrt()
+        with graph.local_scope():
+            graph.ndata["h"] = node_feat
+            if self.edge_feat_size > 0:
+                assert edge_feat is not None, "Edge features must be provided."
+                graph.edata["a"] = edge_feat
+
+            graph.update_all(self.message, self.reduce_func)
+            h = self.U(torch.cat([node_feat, graph.ndata["h_neigh"]], dim=-1))
+            h = h * snorm_n
+            return self.dropout(self.batchnorm(h))
+
+
+class PNAConv(nn.Module):
+    r"""Principal Neighbourhood Aggregation Layer from `Principal Neighbourhood Aggregation
+    for Graph Nets <https://arxiv.org/abs/2004.05718>`__
+
+    A PNA layer is composed of multiple PNA towers. Each tower takes as input a split of the
+    input features, and computes the message passing as below.
+
+    .. math::
+        h_i^(l+1) = U(h_i^l, \oplus_{(i,j)\in E}M(h_i^l, e_{i,j}, h_j^l))
+
+    where :math:`h_i` and :math:`e_{i,j}` are node features and edge features, respectively.
+    :math:`M` and :math:`U` are MLPs, taking the concatenation of input for computing
+    output features. :math:`\oplus` represents the combination of various aggregators
+    and scalers. Aggregators aggregate messages from neighbours and scalers scale the
+    aggregated messages in different ways. :math:`\oplus` concatenates the output features
+    of each combination.
+
+    The output of multiple towers are concatenated and fed into a linear mixing layer for the
+    final output.
+
+    Parameters
+    ----------
+    in_size : int
+        Input feature size; i.e. the size of :math:`h_i^l`.
+    out_size : int
+        Output feature size; i.e. the size of :math:`h_i^{l+1}`.
+    aggregators : list of str
+        List of aggregation function names(each aggregator specifies a way to aggregate
+        messages from neighbours), selected from:
+
+        * ``mean``: the mean of neighbour messages
+
+        * ``max``: the maximum of neighbour messages
+
+        * ``min``: the minimum of neighbour messages
+
+        * ``std``: the standard deviation of neighbour messages
+
+        * ``var``: the variance of neighbour messages
+
+        * ``sum``: the sum of neighbour messages
+
+        * ``moment3``, ``moment4``, ``moment5``: the normalized moments aggregation
+        :math:`(E[(X-E[X])^n])^{1/n}`
+    scalers: list of str
+        List of scaler function names, selected from:
+
+        * ``identity``: no scaling
+
+        * ``amplification``: multiply the aggregated message by :math:`\log(d+1)/\delta`,
+        where :math:`d` is the degree of the node.
+
+        * ``attenuation``: multiply the aggregated message by :math:`\delta/\log(d+1)`
+    delta: float
+        The degree-related normalization factor computed over the training set, used by scalers
+        for normalization. :math:`E[\log(d+1)]`, where :math:`d` is the degree for each node
+        in the training set.
+    dropout: float, optional
+        The dropout ratio. Default: 0.0.
+    num_towers: int, optional
+        The number of towers used. Default: 1. Note that in_size and out_size must be divisible
+        by num_towers.
+    edge_feat_size: int, optional
+        The edge feature size. Default: 0.
+    residual : bool, optional
+        The bool flag that determines whether to add a residual connection for the
+        output. Default: True. If in_size and out_size of the PNA conv layer are not
+        the same, this flag will be set as False forcibly.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import torch as th
+    >>> from dgl.nn import PNAConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = PNAConv(10, 10, ['mean', 'max', 'sum'], ['identity', 'amplification'], 2.5)
+    >>> ret = conv(g, feat)
+    """
+
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        aggregators,
+        scalers,
+        delta,
+        dropout=0.0,
+        num_towers=1,
+        edge_feat_size=0,
+        residual=True,
+    ):
+        super(PNAConv, self).__init__()
+
+        self.in_size = in_size
+        self.out_size = out_size
+        assert (
+            in_size % num_towers == 0
+        ), "in_size must be divisible by num_towers"
+        assert (
+            out_size % num_towers == 0
+        ), "out_size must be divisible by num_towers"
+        self.tower_in_size = in_size // num_towers
+        self.tower_out_size = out_size // num_towers
+        self.edge_feat_size = edge_feat_size
+        self.residual = residual
+        if self.in_size != self.out_size:
+            self.residual = False
+
+        self.towers = nn.ModuleList(
+            [
+                PNAConvTower(
+                    self.tower_in_size,
+                    self.tower_out_size,
+                    aggregators,
+                    scalers,
+                    delta,
+                    dropout=dropout,
+                    edge_feat_size=edge_feat_size,
+                )
+                for _ in range(num_towers)
+            ]
+        )
+
+        self.mixing_layer = nn.Sequential(
+            nn.Linear(out_size, out_size), nn.LeakyReLU()
+        )
+
+    def forward(self, graph, node_feat, edge_feat=None):
+        r"""
+        Description
+        -----------
+        Compute PNA layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        node_feat : torch.Tensor
+            The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
+            nodes, and :math:`h_n` must be the same as in_size.
+        edge_feat : torch.Tensor, optional
+            The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
+            edges, and :math:`h_e` must be the same as edge_feat_size.
+
+        Returns
+        -------
+        torch.Tensor
+            The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
+            should be the same as out_size.
+        """
+        h_cat = torch.cat(
+            [
+                tower(
+                    graph,
+                    node_feat[
+                        :,
+                        ti * self.tower_in_size : (ti + 1) * self.tower_in_size,
+                    ],
+                    edge_feat,
+                )
+                for ti, tower in enumerate(self.towers)
+            ],
+            dim=1,
+        )
+        h_out = self.mixing_layer(h_cat)
+        # add residual connection
+        if self.residual:
+            h_out = h_out + node_feat
+
+        return h_out

python/dgl/nn/pytorch/conv/relgraphconv.py

@@ -0,0 +1,195 @@
+"""Torch Module for Relational graph convolution layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from ..linear import TypedLinear
+
+
+class RelGraphConv(nn.Module):
+    r"""Relational graph convolution layer from `Modeling Relational Data with Graph
+    Convolutional Networks <https://arxiv.org/abs/1703.06103>`__
+
+    It can be described in as below:
+
+    .. math::
+
+       h_i^{(l+1)} = \sigma(\sum_{r\in\mathcal{R}}
+       \sum_{j\in\mathcal{N}^r(i)}e_{j,i}W_r^{(l)}h_j^{(l)}+W_0^{(l)}h_i^{(l)})
+
+    where :math:`\mathcal{N}^r(i)` is the neighbor set of node :math:`i` w.r.t. relation
+    :math:`r`. :math:`e_{j,i}` is the normalizer. :math:`\sigma` is an activation
+    function. :math:`W_0` is the self-loop weight.
+
+    The basis regularization decomposes :math:`W_r` by:
+
+    .. math::
+
+       W_r^{(l)} = \sum_{b=1}^B a_{rb}^{(l)}V_b^{(l)}
+
+    where :math:`B` is the number of bases, :math:`V_b^{(l)}` are linearly combined
+    with coefficients :math:`a_{rb}^{(l)}`.
+
+    The block-diagonal-decomposition regularization decomposes :math:`W_r` into :math:`B`
+    number of block diagonal matrices. We refer :math:`B` as the number of bases.
+
+    The block regularization decomposes :math:`W_r` by:
+
+    .. math::
+
+       W_r^{(l)} = \oplus_{b=1}^B Q_{rb}^{(l)}
+
+    where :math:`B` is the number of bases, :math:`Q_{rb}^{(l)}` are block
+    bases with shape :math:`R^{(d^{(l+1)}/B)*(d^{l}/B)}`.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
+    out_feat : int
+        Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
+    num_rels : int
+        Number of relations.
+    regularizer : str, optional
+        Which weight regularizer to use ("basis", "bdd" or ``None``):
+
+         - "basis" is for basis-decomposition.
+         - "bdd" is for block-diagonal-decomposition.
+         - ``None`` applies no regularization.
+
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. It comes into effect when a regularizer is applied.
+        If ``None``, it uses number of relations (``num_rels``). Default: ``None``.
+        Note that ``in_feat`` and ``out_feat`` must be divisible by ``num_bases``
+        when applying "bdd" regularizer.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import RelGraphConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = RelGraphConv(10, 2, 3, regularizer='basis', num_bases=2)
+    >>> etype = th.tensor([0,1,2,0,1,2])
+    >>> res = conv(g, feat, etype)
+    >>> res
+    tensor([[ 0.3996, -2.3303],
+            [-0.4323, -0.1440],
+            [ 0.3996, -2.3303],
+            [ 2.1046, -2.8654],
+            [-0.4323, -0.1440],
+            [-0.1309, -1.0000]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        if regularizer is not None and num_bases is None:
+            num_bases = num_rels
+        self.linear_r = TypedLinear(
+            in_feat, out_feat, num_rels, regularizer, num_bases
+        )
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # TODO(minjie): consider remove those options in the future to make
+        #   the module only about graph convolution.
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def message(self, edges):
+        """Message function."""
+        m = self.linear_r(edges.src["h"], edges.data["etype"], self.presorted)
+        if "norm" in edges.data:
+            m = m * edges.data["norm"]
+        return {"m": m}
+
+    def forward(self, g, feat, etypes, norm=None, *, presorted=False):
+        """Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor or list[int]
+            An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
+        norm : torch.Tensor, optional
+            An 1D tensor of edge norm value.  Shape: :math:`(|E|,)`.
+        presorted : bool, optional
+            Whether the edges of the input graph have been sorted by their types.
+            Forward on pre-sorted graph may be faster. Graphs created
+            by :func:`~dgl.to_homogeneous` automatically satisfy the condition.
+            Also see :func:`~dgl.reorder_graph` for sorting edges manually.
+
+        Returns
+        -------
+        torch.Tensor
+            New node features. Shape: :math:`(|V|, D_{out})`.
+        """
+        self.presorted = presorted
+        with g.local_scope():
+            g.srcdata["h"] = feat
+            if norm is not None:
+                g.edata["norm"] = norm
+            g.edata["etype"] = etypes
+            # message passing
+            g.update_all(self.message, fn.sum("m", "h"))
+            # apply bias and activation
+            h = g.dstdata["h"]
+            if self.layer_norm:
+                h = self.layer_norm_weight(h)
+            if self.bias:
+                h = h + self.h_bias
+            if self.self_loop:
+                h = h + feat[: g.num_dst_nodes()] @ self.loop_weight
+            if self.activation:
+                h = self.activation(h)
+            h = self.dropout(h)
+            return h

python/dgl/nn/pytorch/conv/sageconv.py

@@ -0,0 +1,295 @@
+"""Torch Module for GraphSAGE layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .... import function as fn
+from ....base import DGLError
+from ....utils import check_eq_shape, expand_as_pair
+
+
+class SAGEConv(nn.Module):
+    r"""GraphSAGE layer from `Inductive Representation Learning on
+    Large Graphs <https://arxiv.org/pdf/1706.02216.pdf>`__
+
+    .. math::
+        h_{\mathcal{N}(i)}^{(l+1)} &= \mathrm{aggregate}
+        \left(\{h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)
+
+        h_{i}^{(l+1)} &= \sigma \left(W \cdot \mathrm{concat}
+        (h_{i}^{l}, h_{\mathcal{N}(i)}^{l+1}) \right)
+
+        h_{i}^{(l+1)} &= \mathrm{norm}(h_{i}^{(l+1)})
+
+    If a weight tensor on each edge is provided, the aggregation becomes:
+
+    .. math::
+        h_{\mathcal{N}(i)}^{(l+1)} = \mathrm{aggregate}
+        \left(\{e_{ji} h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)
+
+    where :math:`e_{ji}` is the scalar weight on the edge from node :math:`j` to node :math:`i`.
+    Please make sure that :math:`e_{ji}` is broadcastable with :math:`h_j^{l}`.
+
+    Parameters
+    ----------
+    in_feats : int, or pair of ints
+        Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
+
+        SAGEConv can be applied on homogeneous graph and unidirectional
+        `bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
+        If the layer applies on a unidirectional bipartite graph, ``in_feats``
+        specifies the input feature size on both the source and destination nodes.  If
+        a scalar is given, the source and destination node feature size would take the
+        same value.
+
+        If aggregator type is ``gcn``, the feature size of source and destination nodes
+        are required to be the same.
+    out_feats : int
+        Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
+    aggregator_type : str
+        Aggregator type to use (``mean``, ``gcn``, ``pool``, ``lstm``).
+    feat_drop : float
+        Dropout rate on features, default: ``0``.
+    bias : bool
+        If True, adds a learnable bias to the output. Default: ``True``.
+    norm : callable activation function/layer or None, optional
+        If not None, applies normalization to the updated node features.
+    activation : callable activation function/layer or None, optional
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+
+    Examples
+    --------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import SAGEConv
+
+    >>> # Case 1: Homogeneous graph
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> conv = SAGEConv(10, 2, 'pool')
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[-1.0888, -2.1099],
+            [-1.0888, -2.1099],
+            [-1.0888, -2.1099],
+            [-1.0888, -2.1099],
+            [-1.0888, -2.1099],
+            [-1.0888, -2.1099]], grad_fn=<AddBackward0>)
+
+    >>> # Case 2: Unidirectional bipartite graph
+    >>> u = [0, 1, 0, 0, 1]
+    >>> v = [0, 1, 2, 3, 2]
+    >>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
+    >>> u_fea = th.rand(2, 5)
+    >>> v_fea = th.rand(4, 10)
+    >>> conv = SAGEConv((5, 10), 2, 'mean')
+    >>> res = conv(g, (u_fea, v_fea))
+    >>> res
+    tensor([[ 0.3163,  3.1166],
+            [ 0.3866,  2.5398],
+            [ 0.5873,  1.6597],
+            [-0.2502,  2.8068]], grad_fn=<AddBackward0>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        aggregator_type,
+        feat_drop=0.0,
+        bias=True,
+        norm=None,
+        activation=None,
+    ):
+        super(SAGEConv, self).__init__()
+        valid_aggre_types = {"mean", "gcn", "pool", "lstm"}
+        if aggregator_type not in valid_aggre_types:
+            raise DGLError(
+                "Invalid aggregator_type. Must be one of {}. "
+                "But got {!r} instead.".format(
+                    valid_aggre_types, aggregator_type
+                )
+            )
+
+        self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
+        self._out_feats = out_feats
+        self._aggre_type = aggregator_type
+        self.norm = norm
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.activation = activation
+
+        # aggregator type: mean/pool/lstm/gcn
+        if aggregator_type == "pool":
+            self.fc_pool = nn.Linear(self._in_src_feats, self._in_src_feats)
+        if aggregator_type == "lstm":
+            self.lstm = nn.LSTM(
+                self._in_src_feats, self._in_src_feats, batch_first=True
+            )
+
+        self.fc_neigh = nn.Linear(self._in_src_feats, out_feats, bias=False)
+
+        if aggregator_type != "gcn":
+            self.fc_self = nn.Linear(self._in_dst_feats, out_feats, bias=bias)
+        elif bias:
+            self.bias = nn.parameter.Parameter(torch.zeros(self._out_feats))
+        else:
+            self.register_buffer("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The linear weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
+        The LSTM module is using xavier initialization method for its weights.
+        """
+        gain = nn.init.calculate_gain("relu")
+        if self._aggre_type == "pool":
+            nn.init.xavier_uniform_(self.fc_pool.weight, gain=gain)
+        if self._aggre_type == "lstm":
+            self.lstm.reset_parameters()
+        if self._aggre_type != "gcn":
+            nn.init.xavier_uniform_(self.fc_self.weight, gain=gain)
+        nn.init.xavier_uniform_(self.fc_neigh.weight, gain=gain)
+
+    def _lstm_reducer(self, nodes):
+        """LSTM reducer
+        NOTE(zihao): lstm reducer with default schedule (degree bucketing)
+        is slow, we could accelerate this with degree padding in the future.
+        """
+        m = nodes.mailbox["m"]  # (B, L, D)
+        batch_size = m.shape[0]
+        h = (
+            m.new_zeros((1, batch_size, self._in_src_feats)),
+            m.new_zeros((1, batch_size, self._in_src_feats)),
+        )
+        _, (rst, _) = self.lstm(m, h)
+        return {"neigh": rst.squeeze(0)}
+
+    def forward(self, graph, feat, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute GraphSAGE layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor or pair of torch.Tensor
+            If a torch.Tensor is given, it represents the input feature of shape
+            :math:`(N, D_{in})`
+            where :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
+            If a pair of torch.Tensor is given, the pair must contain two tensors of shape
+            :math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
+        edge_weight : torch.Tensor, optional
+            Optional tensor on the edge. If given, the convolution will weight
+            with regard to the message.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N_{dst}, D_{out})`
+            where :math:`N_{dst}` is the number of destination nodes in the input graph,
+            :math:`D_{out}` is the size of the output feature.
+        """
+        with graph.local_scope():
+            if isinstance(feat, tuple):
+                feat_src = self.feat_drop(feat[0])
+                feat_dst = self.feat_drop(feat[1])
+            else:
+                feat_src = feat_dst = self.feat_drop(feat)
+                if graph.is_block:
+                    feat_dst = feat_src[: graph.number_of_dst_nodes()]
+            msg_fn = fn.copy_u("h", "m")
+            if edge_weight is not None:
+                assert edge_weight.shape[0] == graph.num_edges()
+                graph.edata["_edge_weight"] = edge_weight
+                msg_fn = fn.u_mul_e("h", "_edge_weight", "m")
+
+            h_self = feat_dst
+
+            # Handle the case of graphs without edges
+            if graph.num_edges() == 0:
+                graph.dstdata["neigh"] = torch.zeros(
+                    feat_dst.shape[0], self._in_src_feats
+                ).to(feat_dst)
+
+            # Determine whether to apply linear transformation before message passing A(XW)
+            lin_before_mp = self._in_src_feats > self._out_feats
+
+            # Message Passing
+            if self._aggre_type == "mean":
+                graph.srcdata["h"] = (
+                    self.fc_neigh(feat_src) if lin_before_mp else feat_src
+                )
+                graph.update_all(msg_fn, fn.mean("m", "neigh"))
+                h_neigh = graph.dstdata["neigh"]
+                if not lin_before_mp:
+                    h_neigh = self.fc_neigh(h_neigh)
+            elif self._aggre_type == "gcn":
+                check_eq_shape(feat)
+                graph.srcdata["h"] = (
+                    self.fc_neigh(feat_src) if lin_before_mp else feat_src
+                )
+                if isinstance(feat, tuple):  # heterogeneous
+                    graph.dstdata["h"] = (
+                        self.fc_neigh(feat_dst) if lin_before_mp else feat_dst
+                    )
+                else:
+                    if graph.is_block:
+                        graph.dstdata["h"] = graph.srcdata["h"][
+                            : graph.num_dst_nodes()
+                        ]
+                    else:
+                        graph.dstdata["h"] = graph.srcdata["h"]
+                graph.update_all(msg_fn, fn.sum("m", "neigh"))
+                # divide in_degrees
+                degs = graph.in_degrees().to(feat_dst)
+                h_neigh = (graph.dstdata["neigh"] + graph.dstdata["h"]) / (
+                    degs.unsqueeze(-1) + 1
+                )
+                if not lin_before_mp:
+                    h_neigh = self.fc_neigh(h_neigh)
+            elif self._aggre_type == "pool":
+                graph.srcdata["h"] = F.relu(self.fc_pool(feat_src))
+                graph.update_all(msg_fn, fn.max("m", "neigh"))
+                h_neigh = self.fc_neigh(graph.dstdata["neigh"])
+            elif self._aggre_type == "lstm":
+                graph.srcdata["h"] = feat_src
+                graph.update_all(msg_fn, self._lstm_reducer)
+                h_neigh = self.fc_neigh(graph.dstdata["neigh"])
+            else:
+                raise KeyError(
+                    "Aggregator type {} not recognized.".format(
+                        self._aggre_type
+                    )
+                )
+
+            # GraphSAGE GCN does not require fc_self.
+            if self._aggre_type == "gcn":
+                rst = h_neigh
+                # add bias manually for GCN
+                if self.bias is not None:
+                    rst = rst + self.bias
+            else:
+                rst = self.fc_self(h_self) + h_neigh
+
+            # activation
+            if self.activation is not None:
+                rst = self.activation(rst)
+            # normalization
+            if self.norm is not None:
+                rst = self.norm(rst)
+            return rst

python/dgl/nn/pytorch/conv/sgconv.py

@@ -0,0 +1,218 @@
+"""Torch Module for Simplifying Graph Convolution layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from ....base import DGLError
+from .graphconv import EdgeWeightNorm
+
+
+class SGConv(nn.Module):
+    r"""SGC layer from `Simplifying Graph
+    Convolutional Networks <https://arxiv.org/pdf/1902.07153.pdf>`__
+
+    .. math::
+        H^{K} = (\tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2})^K X \Theta
+
+    where :math:`\tilde{A}` is :math:`A` + :math:`I`.
+    Thus the graph input is expected to have self-loop edges added.
+
+    Parameters
+    ----------
+    in_feats : int
+        Number of input features; i.e, the number of dimensions of :math:`X`.
+    out_feats : int
+        Number of output features; i.e, the number of dimensions of :math:`H^{K}`.
+    k : int
+        Number of hops :math:`K`. Defaults:``1``.
+    cached : bool
+        If True, the module would cache
+
+        .. math::
+            (\tilde{D}^{-\frac{1}{2}}\tilde{A}\tilde{D}^{-\frac{1}{2}})^K X\Theta
+
+        at the first forward call. This parameter should only be set to
+        ``True`` in Transductive Learning setting.
+    bias : bool
+        If True, adds a learnable bias to the output. Default: ``True``.
+    norm : callable activation function/layer or None, optional
+        If not None, applies normalization to the updated node features.  Default: ``False``.
+    allow_zero_in_degree : bool, optional
+        If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
+        since no message will be passed to those nodes. This is harmful for some applications
+        causing silent performance regression. This module will raise a DGLError if it detects
+        0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
+        and let the users handle it by themselves. Default: ``False``.
+
+    Note
+    ----
+    Zero in-degree nodes will lead to invalid output value. This is because no message
+    will be passed to those nodes, the aggregation function will be appied on empty input.
+    A common practice to avoid this is to add a self-loop for each node in the graph if
+    it is homogeneous, which can be achieved by:
+
+    >>> g = ... # a DGLGraph
+    >>> g = dgl.add_self_loop(g)
+
+    Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
+    since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
+    to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
+    A common practise to handle this is to filter out the nodes with zero-in-degree when use
+    after conv.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import SGConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> g = dgl.add_self_loop(g)
+    >>> feat = th.ones(6, 10)
+    >>> conv = SGConv(10, 2, k=2)
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[-1.9441, -0.9343],
+            [-1.9441, -0.9343],
+            [-1.9441, -0.9343],
+            [-2.7709, -1.3316],
+            [-1.9297, -0.9273],
+            [-1.9441, -0.9343]], grad_fn=<AddmmBackward>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        k=1,
+        cached=False,
+        bias=True,
+        norm=None,
+        allow_zero_in_degree=False,
+    ):
+        super(SGConv, self).__init__()
+        self.fc = nn.Linear(in_feats, out_feats, bias=bias)
+        self._cached = cached
+        self._cached_h = None
+        self._k = k
+        self.norm = norm
+        self._allow_zero_in_degree = allow_zero_in_degree
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The model parameters are initialized using xavier initialization
+        and the bias is initialized to be zero.
+        """
+        nn.init.xavier_uniform_(self.fc.weight)
+        if self.fc.bias is not None:
+            nn.init.zeros_(self.fc.bias)
+
+    def set_allow_zero_in_degree(self, set_value):
+        r"""
+
+        Description
+        -----------
+        Set allow_zero_in_degree flag.
+
+        Parameters
+        ----------
+        set_value : bool
+            The value to be set to the flag.
+        """
+        self._allow_zero_in_degree = set_value
+
+    def forward(self, graph, feat, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute Simplifying Graph Convolution layer.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
+            is size of input feature, :math:`N` is the number of nodes.
+        edge_weight: torch.Tensor, optional
+            edge_weight to use in the message passing process. This is equivalent to
+            using weighted adjacency matrix in the equation above, and
+            :math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
+            is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is size of output feature.
+
+        Raises
+        ------
+        DGLError
+            If there are 0-in-degree nodes in the input graph, it will raise DGLError
+            since no message will be passed to those nodes. This will cause invalid output.
+            The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
+
+        Note
+        ----
+        If ``cache`` is set to True, ``feat`` and ``graph`` should not change during
+        training, or you will get wrong results.
+        """
+        with graph.local_scope():
+            if not self._allow_zero_in_degree:
+                if (graph.in_degrees() == 0).any():
+                    raise DGLError(
+                        "There are 0-in-degree nodes in the graph, "
+                        "output for those nodes will be invalid. "
+                        "This is harmful for some applications, "
+                        "causing silent performance regression. "
+                        "Adding self-loop on the input graph by "
+                        "calling `g = dgl.add_self_loop(g)` will resolve "
+                        "the issue. Setting ``allow_zero_in_degree`` "
+                        "to be `True` when constructing this module will "
+                        "suppress the check and let the code run."
+                    )
+
+            msg_func = fn.copy_u("h", "m")
+            if edge_weight is not None:
+                graph.edata["_edge_weight"] = EdgeWeightNorm("both")(
+                    graph, edge_weight
+                )
+                msg_func = fn.u_mul_e("h", "_edge_weight", "m")
+
+            if self._cached_h is not None:
+                feat = self._cached_h
+            else:
+                if edge_weight is None:
+                    # compute normalization
+                    degs = graph.in_degrees().to(feat).clamp(min=1)
+                    norm = th.pow(degs, -0.5)
+                    norm = norm.to(feat.device).unsqueeze(1)
+                # compute (D^-1 A^k D)^k X
+                for _ in range(self._k):
+                    if edge_weight is None:
+                        feat = feat * norm
+                    graph.ndata["h"] = feat
+                    graph.update_all(msg_func, fn.sum("m", "h"))
+                    feat = graph.ndata.pop("h")
+                    if edge_weight is None:
+                        feat = feat * norm
+
+                if self.norm is not None:
+                    feat = self.norm(feat)
+
+                # cache feature
+                if self._cached:
+                    self._cached_h = feat
+            return self.fc(feat)

python/dgl/nn/pytorch/conv/tagconv.py

@@ -0,0 +1,150 @@
+"""Torch Module for Topology Adaptive Graph Convolutional layer"""
+# pylint: disable= no-member, arguments-differ, invalid-name
+import torch as th
+from torch import nn
+
+from .... import function as fn
+from .graphconv import EdgeWeightNorm
+
+
+class TAGConv(nn.Module):
+    r"""Topology Adaptive Graph Convolutional layer from `Topology
+    Adaptive Graph Convolutional Networks <https://arxiv.org/pdf/1710.10370.pdf>`__
+
+    .. math::
+        H^{K} = {\sum}_{k=0}^K (D^{-1/2} A D^{-1/2})^{k} X {\Theta}_{k},
+
+    where :math:`A` denotes the adjacency matrix,
+    :math:`D_{ii} = \sum_{j=0} A_{ij}` its diagonal degree matrix,
+    :math:`{\Theta}_{k}` denotes the linear weights to sum the results of different hops together.
+
+    Parameters
+    ----------
+    in_feats : int
+        Input feature size. i.e, the number of dimensions of :math:`X`.
+    out_feats : int
+        Output feature size.  i.e, the number of dimensions of :math:`H^{K}`.
+    k: int, optional
+        Number of hops :math:`K`. Default: ``2``.
+    bias: bool, optional
+        If True, adds a learnable bias to the output. Default: ``True``.
+    activation: callable activation function/layer or None, optional
+        If not None, applies an activation function to the updated node features.
+        Default: ``None``.
+
+    Attributes
+    ----------
+    lin : torch.Module
+        The learnable linear module.
+
+    Example
+    -------
+    >>> import dgl
+    >>> import numpy as np
+    >>> import torch as th
+    >>> from dgl.nn import TAGConv
+    >>>
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = TAGConv(10, 2, k=2)
+    >>> res = conv(g, feat)
+    >>> res
+    tensor([[ 0.5490, -1.6373],
+            [ 0.5490, -1.6373],
+            [ 0.5490, -1.6373],
+            [ 0.5513, -1.8208],
+            [ 0.5215, -1.6044],
+            [ 0.3304, -1.9927]], grad_fn=<AddmmBackward>)
+    """
+
+    def __init__(
+        self,
+        in_feats,
+        out_feats,
+        k=2,
+        bias=True,
+        activation=None,
+    ):
+        super(TAGConv, self).__init__()
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self._k = k
+        self._activation = activation
+        self.lin = nn.Linear(in_feats * (self._k + 1), out_feats, bias=bias)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""
+
+        Description
+        -----------
+        Reinitialize learnable parameters.
+
+        Note
+        ----
+        The model parameters are initialized using Glorot uniform initialization.
+        """
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_normal_(self.lin.weight, gain=gain)
+
+    def forward(self, graph, feat, edge_weight=None):
+        r"""
+
+        Description
+        -----------
+        Compute topology adaptive graph convolution.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
+            is size of input feature, :math:`N` is the number of nodes.
+        edge_weight: torch.Tensor, optional
+            edge_weight to use in the message passing process. This is equivalent to
+            using weighted adjacency matrix in the equation above, and
+            :math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
+            is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
+
+        Returns
+        -------
+        torch.Tensor
+            The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
+            is size of output feature.
+        """
+        with graph.local_scope():
+            assert graph.is_homogeneous, "Graph is not homogeneous"
+            if edge_weight is None:
+                norm = th.pow(graph.in_degrees().to(feat).clamp(min=1), -0.5)
+                shp = norm.shape + (1,) * (feat.dim() - 1)
+                norm = th.reshape(norm, shp).to(feat.device)
+
+            msg_func = fn.copy_u("h", "m")
+            if edge_weight is not None:
+                graph.edata["_edge_weight"] = EdgeWeightNorm("both")(
+                    graph, edge_weight
+                )
+                msg_func = fn.u_mul_e("h", "_edge_weight", "m")
+            # D-1/2 A D -1/2 X
+            fstack = [feat]
+            for _ in range(self._k):
+                if edge_weight is None:
+                    rst = fstack[-1] * norm
+                else:
+                    rst = fstack[-1]
+                graph.ndata["h"] = rst
+
+                graph.update_all(msg_func, fn.sum(msg="m", out="h"))
+                rst = graph.ndata["h"]
+                if edge_weight is None:
+                    rst = rst * norm
+                fstack.append(rst)
+
+            rst = self.lin(th.cat(fstack, dim=-1))
+
+            if self._activation is not None:
+                rst = self._activation(rst)
+
+            return rst

python/dgl/nn/pytorch/conv/twirlsconv.py

@@ -0,0 +1,699 @@
+"""Torch modules for TWIRLS"""
+# pylint: disable=invalid-name, useless-super-delegation, no-member
+
+import torch as tc
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .... import function as fn
+
+
+class TWIRLSConv(nn.Module):
+    r"""Convolution together with iteratively reweighting least squre from
+    `Graph Neural Networks Inspired by Classical Iterative Algorithms
+    <https://arxiv.org/pdf/2103.06064.pdf>`__
+
+    Parameters
+    ----------
+    input_d : int
+        Number of input features.
+    output_d : int
+        Number of output features.
+    hidden_d : int
+        Size of hidden layers.
+    prop_step : int
+        Number of propagation steps
+    num_mlp_before : int
+        Number of mlp layers before propagation. Default: ``1``.
+    num_mlp_after : int
+        Number of mlp layers after propagation.  Default: ``1``.
+    norm : str
+        The type of norm layers inside mlp layers. Can be ``'batch'``, ``'layer'`` or ``'none'``.
+        Default: ``'none'``
+    precond : str
+        If True, use pre conditioning and unormalized laplacian, else not use pre conditioning
+        and use normalized laplacian. Default: ``True``
+    alp : float
+        The :math:`\alpha` in paper. If equal to :math:`0`, will be automatically decided based
+        on other hyper prameters. Default: ``0``.
+    lam : float
+        The :math:`\lambda` in paper. Default: ``1``.
+    attention : bool
+        If ``True``, add an attention layer inside propagations. Default: ``False``.
+    tau : float
+        The :math:`\tau` in paper. Default: ``0.2``.
+    T : float
+        The :math:`T` in paper. If < 0, :math:`T` will be set to `\infty`. Default: ``-1``.
+    p : float
+        The :math:`p` in paper. Default: ``1``.
+    use_eta : bool
+        If ``True``, add a learnable weight on each dimension in attention. Default: ``False``.
+    attn_bef : bool
+        If ``True``, add another attention layer before propagation. Default: ``False``.
+    dropout : float
+        The dropout rate in mlp layers. Default: ``0.0``.
+    attn_dropout : float
+        The dropout rate of attention values. Default: ``0.0``.
+    inp_dropout : float
+        The dropout rate on input features. Default: ``0.0``.
+
+
+    Note
+    ----
+     ``add_self_loop`` will be automatically called before propagation.
+
+    Example
+    -------
+    >>> import dgl
+    >>> from dgl.nn import TWIRLSConv
+    >>> import torch as th
+
+    >>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
+    >>> feat = th.ones(6, 10)
+    >>> conv = TWIRLSConv(10, 2, 128, prop_step = 64)
+    >>> res = conv(g , feat)
+    >>> res.size()
+    torch.Size([6, 2])
+    """
+
+    def __init__(
+        self,
+        input_d,
+        output_d,
+        hidden_d,
+        prop_step,
+        num_mlp_before=1,
+        num_mlp_after=1,
+        norm="none",
+        precond=True,
+        alp=0,
+        lam=1,
+        attention=False,
+        tau=0.2,
+        T=-1,
+        p=1,
+        use_eta=False,
+        attn_bef=False,
+        dropout=0.0,
+        attn_dropout=0.0,
+        inp_dropout=0.0,
+    ):
+        super().__init__()
+        self.input_d = input_d
+        self.output_d = output_d
+        self.hidden_d = hidden_d
+        self.prop_step = prop_step
+        self.num_mlp_before = num_mlp_before
+        self.num_mlp_after = num_mlp_after
+        self.norm = norm
+        self.precond = precond
+        self.attention = attention
+        self.alp = alp
+        self.lam = lam
+        self.tau = tau
+        self.T = T
+        self.p = p
+        self.use_eta = use_eta
+        self.init_att = attn_bef
+        self.dropout = dropout
+        self.attn_dropout = attn_dropout
+        self.inp_dropout = inp_dropout
+
+        # ----- initialization of some variables -----
+        # where to put attention
+        self.attn_aft = prop_step // 2 if attention else -1
+
+        # whether we can cache unfolding result
+        self.cacheable = (
+            (not self.attention)
+            and self.num_mlp_before == 0
+            and self.inp_dropout <= 0
+        )
+        if self.cacheable:
+            self.cached_unfolding = None
+
+        # if only one layer, then no hidden size
+        self.size_bef_unf = self.hidden_d
+        self.size_aft_unf = self.hidden_d
+        if self.num_mlp_before == 0:
+            self.size_aft_unf = self.input_d  # as the input  of mlp_aft
+        if self.num_mlp_after == 0:
+            self.size_bef_unf = self.output_d  # as the output of mlp_bef
+
+        # ----- computational modules -----
+        self.mlp_bef = MLP(
+            self.input_d,
+            self.hidden_d,
+            self.size_bef_unf,
+            self.num_mlp_before,
+            self.dropout,
+            self.norm,
+            init_activate=False,
+        )
+
+        self.unfolding = TWIRLSUnfoldingAndAttention(
+            self.hidden_d,
+            self.alp,
+            self.lam,
+            self.prop_step,
+            self.attn_aft,
+            self.tau,
+            self.T,
+            self.p,
+            self.use_eta,
+            self.init_att,
+            self.attn_dropout,
+            self.precond,
+        )
+
+        # if there are really transformations before unfolding, then do init_activate in mlp_aft
+        self.mlp_aft = MLP(
+            self.size_aft_unf,
+            self.hidden_d,
+            self.output_d,
+            self.num_mlp_after,
+            self.dropout,
+            self.norm,
+            init_activate=(self.num_mlp_before > 0)
+            and (self.num_mlp_after > 0),
+        )
+
+    def forward(self, graph, feat):
+        r"""
+
+        Description
+        -----------
+        Run TWIRLS forward.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            The initial node features.
+        Returns
+        -------
+        torch.Tensor
+            The output feature
+
+        Note
+        ----
+        * Input shape: :math:`(N, \text{input_d})` where :math:`N` is the number of nodes.
+        * Output shape: :math:`(N, \text{output_d})`.
+        """
+
+        # ensure self loop
+        graph = graph.remove_self_loop()
+        graph = graph.add_self_loop()
+
+        x = feat
+
+        if self.cacheable:
+            # to cache unfolding result becase there is no paramaters before it
+            if self.cached_unfolding is None:
+                self.cached_unfolding = self.unfolding(graph, x)
+
+            x = self.cached_unfolding
+        else:
+            if self.inp_dropout > 0:
+                x = F.dropout(x, self.inp_dropout, training=self.training)
+            x = self.mlp_bef(x)
+            x = self.unfolding(graph, x)
+
+        x = self.mlp_aft(x)
+
+        return x
+
+
+class Propagate(nn.Module):
+    r"""
+
+    Description
+    -----------
+    The propagation method which is with pre-conditioning and reparameterizing. Correspond to
+    eq.28 in the paper.
+
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def _prop(self, graph, Y, lam):
+        """propagation part."""
+        Y = D_power_bias_X(graph, Y, -0.5, lam, 1 - lam)
+        Y = AX(graph, Y)
+        Y = D_power_bias_X(graph, Y, -0.5, lam, 1 - lam)
+
+        return Y
+
+    def forward(self, graph, Y, X, alp, lam):
+        r"""
+
+        Description
+        -----------
+        Propagation forward.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        Y : torch.Tensor
+            The feature under propagation. Corresponds to :math:`Z^{(k)}` in eq.28 in the paper.
+        X : torch.Tensor
+            The original feature. Corresponds to :math:`Z^{(0)}` in eq.28 in the paper.
+        alp : float
+            The step size. Corresponds to :math:`\alpha` in the paper.
+        lam : torch.Tensor
+            The coefficient of smoothing term. Corresponds to :math:`\lambda` in the paper.
+        Returns
+        -------
+        torch.Tensor
+            Propagated feature. :math:`Z^{(k+1)}` in eq.28 in the paper.
+        """
+
+        return (
+            (1 - alp) * Y
+            + alp * lam * self._prop(graph, Y, lam)
+            + alp * D_power_bias_X(graph, X, -1, lam, 1 - lam)
+        )
+
+
+class PropagateNoPrecond(nn.Module):
+    r"""
+
+    Description
+    -----------
+    The propagation method which is without pre-conditioning and reparameterizing and using
+    normalized laplacian.
+    Correspond to eq.30 in the paper.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, graph, Y, X, alp, lam):
+        r"""
+
+        Description
+        -----------
+        Propagation forward.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        Y : torch.Tensor
+            The feature under propagation. Corresponds to :math:`Y^{(k)}` in eq.30 in the paper.
+        X : torch.Tensor
+            The original feature. Corresponds to :math:`Y^{(0)}` in eq.30 in the paper.
+        alp : float
+            The step size. Corresponds to :math:`\alpha` in the paper.
+        lam : torch.Tensor
+            The coefficient of smoothing term. Corresponds to :math:`\lambda` in the paper.
+        Returns
+        -------
+        torch.Tensor
+            Propagated feature. :math:`Y^{(k+1)}` in eq.30 in the paper.
+        """
+
+        return (
+            (1 - alp * lam - alp) * Y
+            + alp * lam * normalized_AX(graph, Y)
+            + alp * X
+        )
+
+
+class Attention(nn.Module):
+    r"""
+
+    Description
+    -----------
+    The attention function. Correspond to :math:`s` in eq.27 the paper.
+
+    Parameters
+    ----------
+    tau : float
+        The lower thresholding parameter. Correspond to :math:`\tau` in the paper.
+    T : float
+        The upper thresholding parameter. Correspond to :math:`T` in the paper.
+    p : float
+        Correspond to :math:`\rho` in the paper..
+    attn_dropout : float
+        the dropout rate of attention value. Default: ``0.0``.
+
+    Returns
+    -------
+    torch.Tensor
+        The output feature
+    """
+
+    def __init__(self, tau, T, p, attn_dropout=0.0):
+        super().__init__()
+
+        self.tau = tau
+        self.T = T
+        self.p = p
+        self.attn_dropout = attn_dropout
+
+    def reweighting(self, graph):
+        """Compute graph edge weight. Would be stored in ``graph.edata['w']``"""
+
+        w = graph.edata["w"]
+
+        # It is not activation here but to ensure w > 0.
+        # w can be < 0 here because of some precision issue in dgl, which causes NaN afterwards.
+        w = F.relu(w) + 1e-7
+
+        w = tc.pow(w, 1 - 0.5 * self.p)
+
+        w[(w < self.tau)] = self.tau
+        if self.T > 0:
+            w[(w > self.T)] = float("inf")
+
+        w = 1 / w
+
+        # if not (w == w).all():
+        #     raise "nan occured!"
+
+        graph.edata["w"] = w + 1e-9  # avoid 0 degree
+
+    def forward(self, graph, Y, etas=None):
+        r"""
+
+        Description
+        -----------
+        Attention forward. Will update ``graph.edata['w']`` and ``graph.ndata['deg']``.
+
+        Parameters
+        ----------
+        graph : DGLGraph
+            The graph.
+        Y : torch.Tensor
+            The feature to compute attention.
+        etas : float
+            The weight of each dimension. If ``None``, then weight of each dimension is 1.
+            Default: ``None``.
+
+        Returns
+        -------
+        DGLGraph
+            The graph.
+        """
+
+        if etas is not None:
+            Y = Y * etas.view(-1)
+
+        # computing edge distance
+        graph.srcdata["h"] = Y
+        graph.srcdata["h_norm"] = (Y**2).sum(-1)
+        graph.apply_edges(fn.u_dot_v("h", "h", "dot_"))
+        graph.apply_edges(fn.u_add_v("h_norm", "h_norm", "norm_"))
+        graph.edata["dot_"] = graph.edata["dot_"].view(-1)
+        graph.edata["norm_"] = graph.edata["norm_"].view(-1)
+        graph.edata["w"] = graph.edata["norm_"] - 2 * graph.edata["dot_"]
+
+        # apply edge distance to get edge weight
+        self.reweighting(graph)
+
+        # update node degrees
+        graph.update_all(fn.copy_e("w", "m"), fn.sum("m", "deg"))
+        graph.ndata["deg"] = graph.ndata["deg"].view(-1)
+
+        # attention dropout. the implementation can ensure the degrees do not change in expectation.
+        # FIXME: consider if there is a better way
+        if self.attn_dropout > 0:
+            graph.edata["w"] = F.dropout(
+                graph.edata["w"], self.attn_dropout, training=self.training
+            )
+
+        return graph
+
+
+def normalized_AX(graph, X):
+    """Y = D^{-1/2}AD^{-1/2}X"""
+
+    Y = D_power_X(graph, X, -0.5)  # Y = D^{-1/2}X
+    Y = AX(graph, Y)  # Y = AD^{-1/2}X
+    Y = D_power_X(graph, Y, -0.5)  # Y = D^{-1/2}AD^{-1/2}X
+
+    return Y
+
+
+def AX(graph, X):
+    """Y = AX"""
+
+    graph.srcdata["h"] = X
+    graph.update_all(
+        fn.u_mul_e("h", "w", "m"),
+        fn.sum("m", "h"),
+    )
+    Y = graph.dstdata["h"]
+
+    return Y
+
+
+def D_power_X(graph, X, power):
+    """Y = D^{power}X"""
+
+    degs = graph.ndata["deg"]
+    norm = tc.pow(degs, power)
+    Y = X * norm.view(X.size(0), 1)
+    return Y
+
+
+def D_power_bias_X(graph, X, power, coeff, bias):
+    """Y = (coeff*D + bias*I)^{power} X"""
+    degs = graph.ndata["deg"]
+    degs = coeff * degs + bias
+    norm = tc.pow(degs, power)
+    Y = X * norm.view(X.size(0), 1)
+    return Y
+
+
+class TWIRLSUnfoldingAndAttention(nn.Module):
+    r"""
+
+    Description
+    -----------
+    Combine propagation and attention together.
+
+    Parameters
+    ----------
+    d : int
+        Size of graph feature.
+    alp : float
+        Step size. :math:`\alpha` in ther paper.
+    lam : int
+        Coefficient of graph smooth term. :math:`\lambda` in ther paper.
+    prop_step : int
+        Number of propagation steps
+    attn_aft : int
+        Where to put attention layer. i.e. number of propagation steps before attention.
+        If set to ``-1``, then no attention.
+    tau : float
+        The lower thresholding parameter. Correspond to :math:`\tau` in the paper.
+    T : float
+        The upper thresholding parameter. Correspond to :math:`T` in the paper.
+    p : float
+        Correspond to :math:`\rho` in the paper..
+    use_eta : bool
+        If `True`, learn a weight vector for each dimension when doing attention.
+    init_att : bool
+        If ``True``, add an extra attention layer before propagation.
+    attn_dropout : float
+        the dropout rate of attention value. Default: ``0.0``.
+    precond : bool
+        If ``True``, use pre-conditioned & reparameterized version propagation (eq.28), else use
+        normalized laplacian (eq.30).
+
+    Example
+    -------
+    >>> import dgl
+    >>> from dgl.nn import TWIRLSUnfoldingAndAttention
+    >>> import torch as th
+
+    >>> g = dgl.graph(([0, 1, 2, 3, 2, 5], [1, 2, 3, 4, 0, 3])).add_self_loop()
+    >>> feat = th.ones(6,5)
+    >>> prop = TWIRLSUnfoldingAndAttention(10, 1, 1, prop_step=3)
+    >>> res = prop(g,feat)
+    >>> res
+    tensor([[2.5000, 2.5000, 2.5000, 2.5000, 2.5000],
+            [2.5000, 2.5000, 2.5000, 2.5000, 2.5000],
+            [2.5000, 2.5000, 2.5000, 2.5000, 2.5000],
+            [3.7656, 3.7656, 3.7656, 3.7656, 3.7656],
+            [2.5217, 2.5217, 2.5217, 2.5217, 2.5217],
+            [4.0000, 4.0000, 4.0000, 4.0000, 4.0000]])
+
+    """
+
+    def __init__(
+        self,
+        d,
+        alp,
+        lam,
+        prop_step,
+        attn_aft=-1,
+        tau=0.2,
+        T=-1,
+        p=1,
+        use_eta=False,
+        init_att=False,
+        attn_dropout=0,
+        precond=True,
+    ):
+        super().__init__()
+
+        self.d = d
+        self.alp = alp if alp > 0 else 1 / (lam + 1)  # automatic set alpha
+        self.lam = lam
+        self.tau = tau
+        self.p = p
+        self.prop_step = prop_step
+        self.attn_aft = attn_aft
+        self.use_eta = use_eta
+        self.init_att = init_att
+
+        prop_method = Propagate if precond else PropagateNoPrecond
+        self.prop_layers = nn.ModuleList(
+            [prop_method() for _ in range(prop_step)]
+        )
+
+        self.init_attn = (
+            Attention(tau, T, p, attn_dropout) if self.init_att else None
+        )
+        self.attn_layer = (
+            Attention(tau, T, p, attn_dropout) if self.attn_aft >= 0 else None
+        )
+        self.etas = nn.Parameter(tc.ones(d)) if self.use_eta else None
+
+    def forward(self, g, X):
+        r"""
+
+        Description
+        -----------
+        Compute forward pass of propagation & attention.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        X : torch.Tensor
+            Init features.
+
+        Returns
+        -------
+        torch.Tensor
+            The graph.
+        """
+        Y = X
+
+        g.edata["w"] = tc.ones(g.num_edges(), 1, device=g.device)
+        g.ndata["deg"] = g.in_degrees().to(X)
+
+        if self.init_att:
+            g = self.init_attn(g, Y, self.etas)
+
+        for k, layer in enumerate(self.prop_layers):
+            # do unfolding
+            Y = layer(g, Y, X, self.alp, self.lam)
+
+            # do attention at certain layer
+            if k == self.attn_aft - 1:
+                g = self.attn_layer(g, Y, self.etas)
+
+        return Y
+
+
+class MLP(nn.Module):
+    r"""
+
+    Description
+    -----------
+    An MLP module.
+
+    Parameters
+    ----------
+    input_d : int
+        Number of input features.
+    output_d : int
+        Number of output features.
+    hidden_d : int
+        Size of hidden layers.
+    num_layers : int
+        Number of mlp layers.
+    dropout : float
+        The dropout rate in mlp layers.
+    norm : str
+        The type of norm layers inside mlp layers. Can be ``'batch'``, ``'layer'`` or ``'none'``.
+    init_activate : bool
+        If add a relu at the beginning.
+
+    """
+
+    def __init__(
+        self,
+        input_d,
+        hidden_d,
+        output_d,
+        num_layers,
+        dropout,
+        norm,
+        init_activate,
+    ):
+        super().__init__()
+
+        self.init_activate = init_activate
+        self.norm = norm
+        self.dropout = dropout
+
+        self.layers = nn.ModuleList([])
+
+        if num_layers == 1:
+            self.layers.append(nn.Linear(input_d, output_d))
+        elif num_layers > 1:
+            self.layers.append(nn.Linear(input_d, hidden_d))
+            for _ in range(num_layers - 2):
+                self.layers.append(nn.Linear(hidden_d, hidden_d))
+            self.layers.append(nn.Linear(hidden_d, output_d))
+
+        # how many norm layers we have
+        self.norm_cnt = num_layers - 1 + int(init_activate)
+        if norm == "batch":
+            self.norms = nn.ModuleList(
+                [nn.BatchNorm1d(hidden_d) for _ in range(self.norm_cnt)]
+            )
+        elif norm == "layer":
+            self.norms = nn.ModuleList(
+                [nn.LayerNorm(hidden_d) for _ in range(self.norm_cnt)]
+            )
+
+        self.reset_params()
+
+    def reset_params(self):
+        """reset mlp parameters using xavier_norm"""
+        for layer in self.layers:
+            nn.init.xavier_normal_(layer.weight.data)
+            nn.init.constant_(layer.bias.data, 0)
+
+    def activate(self, x):
+        """do normlaization and activation"""
+        if self.norm != "none":
+            x = self.norms[self.cur_norm_idx](x)  # use the last norm layer
+            self.cur_norm_idx += 1
+        x = F.relu(x)
+        x = F.dropout(x, self.dropout, training=self.training)
+        return x
+
+    def forward(self, x):
+        """The forward pass of mlp."""
+        self.cur_norm_idx = 0
+
+        if self.init_activate:
+            x = self.activate(x)
+
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+            if i != len(self.layers) - 1:  # do not activate in the last layer
+                x = self.activate(x)
+
+        return x