chore: import upstream snapshot with attribution

examples/sparse/appnp.py

@@ -0,0 +1,114 @@
+"""
+[Predict then Propagate: Graph Neural Networks meet Personalized PageRank]
+(https://arxiv.org/abs/1810.05997)
+"""
+
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+class APPNP(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        hidden_size=64,
+        dropout=0.1,
+        num_hops=10,
+        alpha=0.1,
+    ):
+        super().__init__()
+
+        self.f_theta = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(in_size, hidden_size),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size, out_size),
+        )
+        self.num_hops = num_hops
+        self.A_dropout = nn.Dropout(dropout)
+        self.alpha = alpha
+
+    def forward(self, A_hat, X):
+        Z_0 = Z = self.f_theta(X)
+        for _ in range(self.num_hops):
+            A_drop = dglsp.val_like(A_hat, self.A_dropout(A_hat.val))
+            Z = (1 - self.alpha) * A_drop @ Z + self.alpha * Z_0
+        return Z
+
+
+def evaluate(g, pred):
+    label = g.ndata["label"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(model, g, A_hat, X):
+    label = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    optimizer = Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
+
+    for epoch in range(50):
+        # Forward.
+        model.train()
+        logits = model(A_hat, X)
+
+        # Compute loss with nodes in training set.
+        loss = F.cross_entropy(logits[train_mask], label[train_mask])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute prediction.
+        model.eval()
+        logits = model(A_hat, X)
+        pred = logits.argmax(dim=1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(g, pred)
+        print(
+            f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}, test"
+            f" acc: {test_acc:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+
+    # Create the sparse adjacency matrix A.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    # Calculate the symmetrically normalized adjacency matrix.
+    I = dglsp.identity(A.shape, device=dev)
+    A_hat = A + I
+    D_hat = dglsp.diag(A_hat.sum(dim=1)) ** -0.5
+    A_hat = D_hat @ A_hat @ D_hat
+
+    # Create APPNP model.
+    X = g.ndata["feat"]
+    in_size = X.shape[1]
+    out_size = dataset.num_classes
+    model = APPNP(in_size, out_size).to(dev)
+
+    # Kick off training.
+    train(model, g, A_hat, X)

examples/sparse/c_and_s.py

@@ -0,0 +1,129 @@
+"""
+[Combining Label Propagation and Simple Models Out-performs
+Graph Neural Networks](https://arxiv.org/abs/2010.13993)
+"""
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+###############################################################################
+# (HIGHLIGHT) Compute Label Propagation with Sparse Matrix API
+###############################################################################
+@torch.no_grad()
+def label_propagation(A_hat, label, num_layers=20, alpha=0.9):
+    Y = label
+    for _ in range(num_layers):
+        Y = alpha * A_hat @ Y + (1 - alpha) * label
+        Y = Y.clamp_(0.0, 1.0)
+    return Y
+
+
+def correct(A_hat, label, soft_label, mask):
+    # Compute error.
+    error = torch.zeros_like(soft_label)
+    error[mask] = label[mask] - soft_label[mask]
+
+    # Smooth error.
+    smoothed_error = label_propagation(A_hat, error)
+
+    # Autoscale.
+    sigma = error[mask].abs()
+    sigma = sigma.sum() / sigma.shape[0]
+    scale = sigma / smoothed_error.abs().sum(dim=1, keepdim=True)
+    scale[scale.isinf() | (scale > 1000)] = 1.0
+
+    # Correct.
+    result = soft_label + scale * smoothed_error
+    return result
+
+
+def smooth(A_hat, label, soft_label, mask):
+    soft_label[mask] = label[mask].float()
+    return label_propagation(A_hat, soft_label)
+
+
+def evaluate(g, pred):
+    label = g.ndata["label"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(base_model, g, X):
+    label = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+
+    optimizer = Adam(base_model.parameters(), lr=0.01)
+
+    for epoch in range(10):
+        # Forward.
+        base_model.train()
+        logits = base_model(X)
+
+        # Compute loss with nodes in training set.
+        loss = F.cross_entropy(logits[train_mask], label[train_mask])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute prediction.
+        base_model.eval()
+        logits = base_model(X)
+        pred = logits.argmax(dim=1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(g, pred)
+        print(
+            f"Base model, In epoch {epoch}, loss: {loss:.3f}, "
+            f"val acc: {val_acc:.3f}, test acc: {test_acc:.3f}"
+        )
+    return logits
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+
+    # Create the sparse adjacency matrix A.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    # Calculate the symmetrically normalized adjacency matrix.
+    I = dglsp.identity(A.shape, device=dev)
+    A_hat = A + I
+    D_hat = dglsp.diag(A_hat.sum(dim=1)) ** -0.5
+    A_hat = D_hat @ A_hat @ D_hat
+
+    # Create models.
+    X = g.ndata["feat"]
+    in_size = X.shape[1]
+    out_size = dataset.num_classes
+    base_model = nn.Linear(in_size, out_size).to(dev)
+
+    # Stage1: Train the base model.
+    logits = train(base_model, g, X)
+
+    # Stage2: Correct and Smooth.
+    soft_label = F.softmax(logits, dim=1)
+    label = F.one_hot(g.ndata["label"])
+    soft_label = correct(A_hat, label, soft_label, g.ndata["train_mask"])
+    soft_label = smooth(A_hat, label, soft_label, g.ndata["train_mask"])
+    pred = soft_label.argmax(dim=1)
+    val_acc, test_acc = evaluate(g, pred)
+    print(f"val acc: {val_acc:.3f}, test acc: {test_acc:.3f}")

examples/sparse/gat.py

@@ -0,0 +1,141 @@
+"""
+[Graph Attention Networks]
+(https://arxiv.org/abs/1710.10903)
+"""
+
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+class GATConv(nn.Module):
+    def __init__(self, in_size, out_size, num_heads, dropout):
+        super().__init__()
+
+        self.out_size = out_size
+        self.num_heads = num_heads
+
+        self.dropout = nn.Dropout(dropout)
+        self.W = nn.Linear(in_size, out_size * num_heads)
+        self.a_l = nn.Parameter(torch.zeros(1, out_size, num_heads))
+        self.a_r = nn.Parameter(torch.zeros(1, out_size, num_heads))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_normal_(self.W.weight, gain=gain)
+        nn.init.xavier_normal_(self.a_l, gain=gain)
+        nn.init.xavier_normal_(self.a_r, gain=gain)
+
+    ###########################################################################
+    # (HIGHLIGHT) Take the advantage of DGL sparse APIs to implement
+    # multihead attention.
+    ###########################################################################
+    def forward(self, A_hat, Z):
+        Z = self.dropout(Z)
+        Z = self.W(Z).view(Z.shape[0], self.out_size, self.num_heads)
+
+        # a^T [Wh_i || Wh_j] = a_l Wh_i + a_r Wh_j
+        e_l = (Z * self.a_l).sum(dim=1)
+        e_r = (Z * self.a_r).sum(dim=1)
+        e = e_l[A_hat.row] + e_r[A_hat.col]
+
+        a = F.leaky_relu(e)
+        A_atten = dglsp.val_like(A_hat, a).softmax()
+        a_drop = self.dropout(A_atten.val)
+        A_atten = dglsp.val_like(A_atten, a_drop)
+        return dglsp.bspmm(A_atten, Z)
+
+
+class GAT(nn.Module):
+    def __init__(
+        self, in_size, out_size, hidden_size=8, num_heads=8, dropout=0.6
+    ):
+        super().__init__()
+
+        self.in_conv = GATConv(
+            in_size, hidden_size, num_heads=num_heads, dropout=dropout
+        )
+        self.out_conv = GATConv(
+            hidden_size * num_heads, out_size, num_heads=1, dropout=dropout
+        )
+
+    def forward(self, A_hat, X):
+        # Flatten the head and feature dimension.
+        Z = F.elu(self.in_conv(A_hat, X)).flatten(1)
+        # Average over the head dimension.
+        Z = self.out_conv(A_hat, Z).mean(-1)
+        return Z
+
+
+def evaluate(g, pred):
+    label = g.ndata["label"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(model, g, A_hat, X):
+    label = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    optimizer = Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
+
+    for epoch in range(50):
+        # Forward.
+        model.train()
+        logits = model(A_hat, X)
+
+        # Compute loss with nodes in training set.
+        loss = F.cross_entropy(logits[train_mask], label[train_mask])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute prediction.
+        model.eval()
+        logits = model(A_hat, X)
+        pred = logits.argmax(dim=1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(g, pred)
+        print(
+            f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}, test"
+            f" acc: {test_acc:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+
+    # Create the sparse adjacency matrix A.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    # Add self-loops.
+    I = dglsp.identity(A.shape, device=dev)
+    A_hat = A + I
+
+    # Create GAT model.
+    X = g.ndata["feat"]
+    in_size = X.shape[1]
+    out_size = dataset.num_classes
+    model = GAT(in_size, out_size).to(dev)
+
+    # Kick off training.
+    train(model, g, A_hat, X)

examples/sparse/gcn.py

@@ -0,0 +1,107 @@
+"""
+[Semi-Supervised Classification with Graph Convolutional Networks]
+(https://arxiv.org/abs/1609.02907)
+"""
+
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+class GCN(nn.Module):
+    def __init__(self, in_size, out_size, hidden_size=16):
+        super().__init__()
+
+        # Two-layer GCN.
+        self.W1 = nn.Linear(in_size, hidden_size)
+        self.W2 = nn.Linear(hidden_size, out_size)
+
+    ############################################################################
+    # (HIGHLIGHT) Take the advantage of DGL sparse APIs to implement the GCN
+    # forward process.
+    ############################################################################
+    def forward(self, A_norm, X):
+        X = A_norm @ self.W1(X)
+        X = F.relu(X)
+        X = A_norm @ self.W2(X)
+        return X
+
+
+def evaluate(g, pred):
+    label = g.ndata["label"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(model, g, A_norm, X):
+    label = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    optimizer = Adam(model.parameters(), lr=1e-2, weight_decay=5e-4)
+    loss_fcn = nn.CrossEntropyLoss()
+
+    for epoch in range(200):
+        model.train()
+
+        # Forward.
+        logits = model(A_norm, X)
+
+        # Compute loss with nodes in the training set.
+        loss = loss_fcn(logits[train_mask], label[train_mask])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute prediction.
+        pred = logits.argmax(dim=1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(g, pred)
+        if epoch % 20 == 0:
+            print(
+                f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}"
+                f", test acc: {test_acc:.3f}"
+            )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+    num_classes = dataset.num_classes
+    X = g.ndata["feat"]
+
+    # Create the adjacency matrix of graph.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    ############################################################################
+    # (HIGHLIGHT) Compute the symmetrically normalized adjacency matrix with
+    # Sparse Matrix API
+    ############################################################################
+    I = dglsp.identity(A.shape, device=dev)
+    A_hat = A + I
+    D_hat = dglsp.diag(A_hat.sum(1)) ** -0.5
+    A_norm = D_hat @ A_hat @ D_hat
+
+    # Create model.
+    in_size = X.shape[1]
+    out_size = num_classes
+    model = GCN(in_size, out_size).to(dev)
+
+    # Kick off training.
+    train(model, g, A_norm, X)

examples/sparse/gcnii.py

@@ -0,0 +1,167 @@
+"""
+[Simple and Deep Graph Convolutional Networks]
+(https://arxiv.org/abs/2007.02133)
+"""
+
+import math
+
+import dgl.sparse as dglsp
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+class GCNIIConvolution(nn.Module):
+    def __init__(self, in_size, out_size):
+        super().__init__()
+        self.out_size = out_size
+        self.weight = nn.Linear(in_size, out_size, bias=False)
+
+    ############################################################################
+    # (HIGHLIGHT) Take the advantage of DGL sparse APIs to implement the GCNII
+    # forward process.
+    ############################################################################
+    def forward(self, A_norm, H, H0, lamda, alpha, l):
+        beta = math.log(lamda / l + 1)
+
+        # Multiply a sparse matrix by a dense matrix.
+        H = A_norm @ H
+        H = (1 - alpha) * H + alpha * H0
+        H = (1 - beta) * H + beta * self.weight(H)
+        return H
+
+
+class GCNII(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        hidden_size,
+        n_layers,
+        lamda,
+        alpha,
+        dropout=0.5,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.n_layers = n_layers
+        self.lamda = lamda
+        self.alpha = alpha
+
+        # The GCNII model.
+        self.layers = nn.ModuleList()
+        self.layers.append(nn.Linear(in_size, hidden_size))
+        for _ in range(n_layers):
+            self.layers.append(GCNIIConvolution(hidden_size, hidden_size))
+        self.layers.append(nn.Linear(hidden_size, out_size))
+
+        self.activation = nn.ReLU()
+        self.dropout = dropout
+
+    def forward(self, A_norm, feature):
+        H = feature
+        H = F.dropout(H, self.dropout, training=self.training)
+        H = self.layers[0](H)
+        H = self.activation(H)
+        H0 = H
+
+        # The GCNII convolution forward.
+        for i, conv in enumerate(self.layers[1:-1]):
+            H = F.dropout(H, self.dropout, training=self.training)
+            H = conv(A_norm, H, H0, self.lamda, self.alpha, i + 1)
+            H = self.activation(H)
+
+        H = F.dropout(H, self.dropout, training=self.training)
+        H = self.layers[-1](H)
+
+        return H
+
+
+def evaluate(model, A_norm, H, label, val_mask, test_mask):
+    model.eval()
+    logits = model(A_norm, H)
+    pred = logits.argmax(dim=1)
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(model, g, A_norm, H):
+    label = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+    optimizer = Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
+
+    loss_fcn = nn.CrossEntropyLoss()
+
+    for epoch in range(100):
+        model.train()
+        optimizer.zero_grad()
+
+        # Forward.
+        logits = model(A_norm, H)
+
+        # Compute loss with nodes in the training set.
+        loss = loss_fcn(logits[train_mask], label[train_mask])
+
+        # Backward.
+        loss.backward()
+        optimizer.step()
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(
+            model, A_norm, H, label, val_mask, test_mask
+        )
+        if epoch % 5 == 0:
+            print(
+                f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}"
+                f", test acc: {test_acc:.3f}"
+            )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+    num_classes = dataset.num_classes
+    H = g.ndata["feat"]
+
+    # Create the adjacency matrix of graph.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    ############################################################################
+    # (HIGHLIGHT) Compute the symmetrically normalized adjacency matrix with
+    # Sparse Matrix API
+    ############################################################################
+    I = dglsp.identity(A.shape, device=dev)
+    A_hat = A + I
+    D_hat = dglsp.diag(A_hat.sum(1)) ** -0.5
+    A_norm = D_hat @ A_hat @ D_hat
+
+    # Create model.
+    in_size = H.shape[1]
+    out_size = num_classes
+    model = GCNII(
+        in_size,
+        out_size,
+        hidden_size=64,
+        n_layers=64,
+        lamda=0.5,
+        alpha=0.2,
+        dropout=0.5,
+    ).to(dev)
+
+    # Kick off training.
+    train(model, g, A_norm, H)

examples/sparse/graph_transformer.py

@@ -0,0 +1,191 @@
+"""
+[A Generalization of Transformer Networks to Graphs]
+(https://arxiv.org/abs/2012.09699)
+"""
+
+import dgl
+import dgl.nn as dglnn
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+from dgl.data import AsGraphPredDataset
+from dgl.dataloading import GraphDataLoader
+from ogb.graphproppred import collate_dgl, DglGraphPropPredDataset, Evaluator
+from ogb.graphproppred.mol_encoder import AtomEncoder
+from tqdm import tqdm
+
+
+class SparseMHA(nn.Module):
+    """Sparse Multi-head Attention Module"""
+
+    def __init__(self, hidden_size=80, num_heads=8):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.scaling = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size)
+        self.k_proj = nn.Linear(hidden_size, hidden_size)
+        self.v_proj = nn.Linear(hidden_size, hidden_size)
+        self.out_proj = nn.Linear(hidden_size, hidden_size)
+
+    def forward(self, A, h):
+        N = len(h)
+        q = self.q_proj(h).reshape(N, self.head_dim, self.num_heads)
+        q *= self.scaling
+        k = self.k_proj(h).reshape(N, self.head_dim, self.num_heads)
+        v = self.v_proj(h).reshape(N, self.head_dim, self.num_heads)
+
+        ######################################################################
+        # (HIGHLIGHT) Compute the multi-head attention with Sparse Matrix API
+        ######################################################################
+        attn = dglsp.bsddmm(A, q, k.transpose(1, 0))  # [N, N, nh]
+        attn = attn.softmax()
+        out = dglsp.bspmm(attn, v)
+
+        return self.out_proj(out.reshape(N, -1))
+
+
+class GTLayer(nn.Module):
+    """Graph Transformer Layer"""
+
+    def __init__(self, hidden_size=80, num_heads=8):
+        super().__init__()
+        self.MHA = SparseMHA(hidden_size=hidden_size, num_heads=num_heads)
+        self.batchnorm1 = nn.BatchNorm1d(hidden_size)
+        self.batchnorm2 = nn.BatchNorm1d(hidden_size)
+        self.FFN1 = nn.Linear(hidden_size, hidden_size * 2)
+        self.FFN2 = nn.Linear(hidden_size * 2, hidden_size)
+
+    def forward(self, A, h):
+        h1 = h
+        h = self.MHA(A, h)
+        h = self.batchnorm1(h + h1)
+
+        h2 = h
+        h = self.FFN2(F.relu(self.FFN1(h)))
+        h = h2 + h
+
+        return self.batchnorm2(h)
+
+
+class GTModel(nn.Module):
+    def __init__(
+        self,
+        out_size,
+        hidden_size=80,
+        pos_enc_size=2,
+        num_layers=8,
+        num_heads=8,
+    ):
+        super().__init__()
+        self.atom_encoder = AtomEncoder(hidden_size)
+        self.pos_linear = nn.Linear(pos_enc_size, hidden_size)
+        self.layers = nn.ModuleList(
+            [GTLayer(hidden_size, num_heads) for _ in range(num_layers)]
+        )
+        self.pooler = dglnn.SumPooling()
+        self.predictor = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.ReLU(),
+            nn.Linear(hidden_size // 2, hidden_size // 4),
+            nn.ReLU(),
+            nn.Linear(hidden_size // 4, out_size),
+        )
+
+    def forward(self, g, X, pos_enc):
+        indices = torch.stack(g.edges())
+        N = g.num_nodes()
+        A = dglsp.spmatrix(indices, shape=(N, N))
+        h = self.atom_encoder(X) + self.pos_linear(pos_enc)
+        for layer in self.layers:
+            h = layer(A, h)
+        h = self.pooler(g, h)
+
+        return self.predictor(h)
+
+
+@torch.no_grad()
+def evaluate(model, dataloader, evaluator, device):
+    model.eval()
+    y_true = []
+    y_pred = []
+    for batched_g, labels in dataloader:
+        batched_g, labels = batched_g.to(device), labels.to(device)
+        y_hat = model(batched_g, batched_g.ndata["feat"], batched_g.ndata["PE"])
+        y_true.append(labels.view(y_hat.shape).detach().cpu())
+        y_pred.append(y_hat.detach().cpu())
+    y_true = torch.cat(y_true, dim=0).numpy()
+    y_pred = torch.cat(y_pred, dim=0).numpy()
+    input_dict = {"y_true": y_true, "y_pred": y_pred}
+    return evaluator.eval(input_dict)["rocauc"]
+
+
+def train(model, dataset, evaluator, device):
+    train_dataloader = GraphDataLoader(
+        dataset[dataset.train_idx],
+        batch_size=256,
+        shuffle=True,
+        collate_fn=collate_dgl,
+    )
+    valid_dataloader = GraphDataLoader(
+        dataset[dataset.val_idx], batch_size=256, collate_fn=collate_dgl
+    )
+    test_dataloader = GraphDataLoader(
+        dataset[dataset.test_idx], batch_size=256, collate_fn=collate_dgl
+    )
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    num_epochs = 50
+    scheduler = optim.lr_scheduler.StepLR(
+        optimizer, step_size=num_epochs, gamma=0.5
+    )
+    loss_fcn = nn.BCEWithLogitsLoss()
+
+    for epoch in range(num_epochs):
+        model.train()
+        total_loss = 0.0
+        for batched_g, labels in train_dataloader:
+            batched_g, labels = batched_g.to(device), labels.to(device)
+            logits = model(
+                batched_g, batched_g.ndata["feat"], batched_g.ndata["PE"]
+            )
+            loss = loss_fcn(logits, labels.float())
+            total_loss += loss.item()
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+        scheduler.step()
+        avg_loss = total_loss / len(train_dataloader)
+        val_metric = evaluate(model, valid_dataloader, evaluator, device)
+        test_metric = evaluate(model, test_dataloader, evaluator, device)
+        print(
+            f"Epoch: {epoch:03d}, Loss: {avg_loss:.4f}, "
+            f"Val: {val_metric:.4f}, Test: {test_metric:.4f}"
+        )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # load dataset
+    pos_enc_size = 8
+    dataset = AsGraphPredDataset(
+        DglGraphPropPredDataset("ogbg-molhiv", "./data/OGB")
+    )
+    evaluator = Evaluator("ogbg-molhiv")
+    # laplacian positional encoding
+    for g, _ in tqdm(dataset, desc="Computing Laplacian PE"):
+        g.ndata["PE"] = dgl.lap_pe(g, k=pos_enc_size, padding=True)
+
+    # Create model.
+    out_size = dataset.num_tasks
+    model = GTModel(out_size=out_size, pos_enc_size=pos_enc_size).to(dev)
+
+    # Kick off training.
+    train(model, dataset, evaluator, dev)

examples/sparse/han.py

@@ -0,0 +1,188 @@
+"""
+[Heterogeneous Graph Attention Network]
+(https://arxiv.org/abs/1903.07293)
+"""
+
+import pickle
+
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dgl.data.utils import _get_dgl_url, download, get_download_dir
+from torch.optim import Adam
+
+
+class GATConv(nn.Module):
+    def __init__(self, in_size, out_size, num_heads, dropout):
+        super().__init__()
+
+        self.out_size = out_size
+        self.num_heads = num_heads
+
+        self.dropout = nn.Dropout(dropout)
+        self.W = nn.Linear(in_size, out_size * num_heads)
+        self.a_l = nn.Parameter(torch.zeros(1, out_size, num_heads))
+        self.a_r = nn.Parameter(torch.zeros(1, out_size, num_heads))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_normal_(self.W.weight, gain=gain)
+        nn.init.xavier_normal_(self.a_l, gain=gain)
+        nn.init.xavier_normal_(self.a_r, gain=gain)
+
+    ###########################################################################
+    # (HIGHLIGHT) Take the advantage of DGL sparse APIs to implement
+    # multihead attention.
+    ###########################################################################
+    def forward(self, A_hat, Z):
+        Z = self.dropout(Z)
+        Z = self.W(Z).view(Z.shape[0], self.out_size, self.num_heads)
+
+        # a^T [Wh_i || Wh_j] = a_l Wh_i + a_r Wh_j
+        e_l = (Z * self.a_l).sum(dim=1)
+        e_r = (Z * self.a_r).sum(dim=1)
+        e = e_l[A_hat.row] + e_r[A_hat.col]
+
+        a = F.leaky_relu(e)
+        A_atten = dglsp.val_like(A_hat, a).softmax()
+        a_drop = self.dropout(A_atten.val)
+        A_atten = dglsp.val_like(A_atten, a_drop)
+        return dglsp.bspmm(A_atten, Z)
+
+
+class SemanticAttention(nn.Module):
+    def __init__(self, in_size, hidden_size=128):
+        super().__init__()
+
+        self.project = nn.Sequential(
+            nn.Linear(in_size, hidden_size),
+            nn.Tanh(),
+            nn.Linear(hidden_size, 1, bias=False),
+        )
+
+    def forward(self, z):
+        w = self.project(z).mean(0)
+        beta = torch.softmax(w, dim=0)
+        beta = beta.expand((z.shape[0],) + beta.shape)
+
+        return (beta * z).sum(1)
+
+
+class HAN(nn.Module):
+    def __init__(
+        self,
+        num_meta_paths,
+        in_size,
+        out_size,
+        hidden_size=8,
+        num_heads=8,
+        dropout=0.6,
+    ):
+        super().__init__()
+
+        self.gat_layers = nn.ModuleList()
+        for _ in range(num_meta_paths):
+            self.gat_layers.append(
+                GATConv(in_size, hidden_size, num_heads, dropout)
+            )
+
+        in_size = hidden_size * num_heads
+        self.semantic_attention = SemanticAttention(in_size)
+        self.predict = nn.Linear(in_size, out_size)
+
+    def forward(self, A_list, X):
+        meta_path_Z_list = []
+        for i, A in enumerate(A_list):
+            meta_path_Z_list.append(self.gat_layers[i](A, X).flatten(1))
+
+        # (num_nodes, num_meta_paths, hidden_size * num_heads)
+        meta_path_Z = torch.stack(meta_path_Z_list, dim=1)
+
+        Z = self.semantic_attention(meta_path_Z)
+        Z = self.predict(Z)
+
+        return Z
+
+
+def evaluate(label, val_idx, test_idx, pred):
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_idx] == label[val_idx]).float().mean()
+    test_acc = (pred[test_idx] == label[test_idx]).float().mean()
+    return val_acc, test_acc
+
+
+def train(model, data, A_list, X, label):
+    dev = X.device
+    train_idx = torch.from_numpy(data["train_idx"]).long().squeeze(0).to(dev)
+    val_idx = torch.from_numpy(data["val_idx"]).long().squeeze(0).to(dev)
+    test_idx = torch.from_numpy(data["test_idx"]).long().squeeze(0).to(dev)
+    optimizer = Adam(model.parameters(), lr=0.005, weight_decay=0.001)
+
+    for epoch in range(70):
+        # Forward.
+        model.train()
+        logits = model(A_list, X)
+
+        # Compute loss with nodes in training set.
+        loss = F.cross_entropy(logits[train_idx], label[train_idx])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute prediction.
+        model.eval()
+        logits = model(A_list, X)
+        pred = logits.argmax(dim=1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(label, val_idx, test_idx, pred)
+        print(
+            f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}, test"
+            f" acc: {test_acc:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # (TODO): Move the logic to a built-in dataset.
+    # Load the data.
+    url = "dataset/ACM3025.pkl"
+    data_path = get_download_dir() + "/ACM3025.pkl"
+    download(_get_dgl_url(url), path=data_path)
+
+    with open(data_path, "rb") as f:
+        data = pickle.load(f)
+
+    # Create sparse adjacency matrices corresponding to two meta paths.
+    # Self-loops already added.
+    PAP_dst, PAP_src = data["PAP"].nonzero()
+    PAP_indices = torch.stack(
+        [torch.from_numpy(PAP_src).long(), torch.from_numpy(PAP_dst).long()]
+    ).to(dev)
+    PAP_A = dglsp.spmatrix(PAP_indices)
+
+    PLP_dst, PLP_src = data["PLP"].nonzero()
+    PLP_indices = torch.stack(
+        [torch.from_numpy(PLP_src).long(), torch.from_numpy(PLP_src).long()]
+    ).to(dev)
+    PLP_A = dglsp.spmatrix(PLP_indices)
+    A_list = [PAP_A, PLP_A]
+
+    # Create HAN model.
+    X = torch.from_numpy(data["feature"].todense()).float().to(dev)
+    label = torch.from_numpy(data["label"].todense())
+    out_size = label.shape[1]
+    label = label.nonzero()[:, 1].to(dev)
+    in_size = X.shape[1]
+    model = HAN(len(A_list), in_size, out_size).to(dev)
+
+    # Kick off training.
+    train(model, data, A_list, X, label)

examples/sparse/hetero-rgcn.py

@@ -0,0 +1,280 @@
+"""
+Modeling Relational Data with Graph Convolutional Networks
+Paper: https://arxiv.org/abs/1703.06103
+Reference Code: https://github.com/tkipf/relational-gcn
+
+This script trains and tests a Hetero Relational Graph Convolutional Networks 
+(Hetero-RGCN) model based on the information of a full graph.
+
+This flowchart describes the main functional sequence of the provided example.
+main
+│
+├───> Load and preprocess full dataset
+│
+├───> Instantiate Hetero-RGCN model
+│
+├───> train
+│     │
+│     └───> Training loop
+│           │
+│           └───> Hetero-RGCN.forward
+└───> test
+      │
+      └───> Evaluate the model
+"""
+import argparse
+import time
+
+import dgl
+import dgl.sparse as dglsp
+
+import numpy as np
+
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dgl.data.rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
+
+
+class RelGraphEmbed(nn.Module):
+    r"""Embedding layer for featureless heterograph."""
+
+    def __init__(
+        self,
+        ntype_num,
+        embed_size,
+    ):
+        super(RelGraphEmbed, self).__init__()
+        self.embed_size = embed_size
+        self.dropout = nn.Dropout(0.0)
+
+        # Create weight embeddings for each node for each relation.
+        self.embeds = nn.ParameterDict()
+        for ntype, num_nodes in ntype_num.items():
+            embed = nn.Parameter(th.Tensor(num_nodes, self.embed_size))
+            nn.init.xavier_uniform_(embed, gain=nn.init.calculate_gain("relu"))
+            self.embeds[ntype] = embed
+
+    def forward(self):
+        return self.embeds
+
+
+class HeteroRelationalGraphConv(nn.Module):
+    r"""HeteroRelational graph convolution layer.
+
+    Parameters
+    ----------
+    in_size : int
+        Input feature size.
+    out_size : int
+        Output feature size.
+    relation_names : list[str]
+        Relation names.
+    """
+
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        relation_names,
+        activation=None,
+    ):
+        super(HeteroRelationalGraphConv, self).__init__()
+        self.in_size = in_size
+        self.out_size = out_size
+        self.relation_names = relation_names
+        self.activation = activation
+
+        ########################################################################
+        # (HIGHLIGHT) HeteroGraphConv is a graph convolution operator over
+        # heterogeneous graphs. A dictionary is passed where the key is the
+        # relation name and the value is the insatnce of conv layer.
+        ########################################################################
+        self.W = nn.ModuleDict(
+            {str(rel): nn.Linear(in_size, out_size) for rel in relation_names}
+        )
+
+        self.dropout = nn.Dropout(0.0)
+
+    def forward(self, A, inputs):
+        """Forward computation
+
+        Parameters
+        ----------
+        A : Hetero Sparse Matrix
+            Input graph.
+        inputs : dict[str, torch.Tensor]
+            Node feature for each node type.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            New node features for each node type.
+        """
+        hs = {}
+        for rel in A:
+            src_type, edge_type, dst_type = rel
+            if dst_type not in hs:
+                hs[dst_type] = th.zeros(
+                    inputs[dst_type].shape[0], self.out_size
+                )
+            ####################################################################
+            # (HIGHLIGHT) Sparse library use hetero sparse matrix to present
+            # heterogeneous graphs. A dictionary is passed where the key is
+            # the tuple of (source node type, edge type, destination node type)
+            # and the value is the sparse matrix contructed from the key on
+            # global graph. The convolution operation is the multiplication of
+            # sparse matrix and convolutional layer.
+            ####################################################################
+            hs[dst_type] = hs[dst_type] + (
+                A[rel].T @ self.W[str(edge_type)](inputs[src_type])
+            )
+            if self.activation:
+                hs[dst_type] = self.activation(hs[dst_type])
+            hs[dst_type] = self.dropout(hs[dst_type])
+
+        return hs
+
+
+class EntityClassify(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        relation_names,
+        embed_layer,
+    ):
+        super(EntityClassify, self).__init__()
+        self.in_size = in_size
+        self.out_size = out_size
+        self.relation_names = relation_names
+        self.relation_names.sort()
+        self.embed_layer = embed_layer
+
+        self.layers = nn.ModuleList()
+        # Input to hidden.
+        self.layers.append(
+            HeteroRelationalGraphConv(
+                self.in_size,
+                self.in_size,
+                self.relation_names,
+                activation=F.relu,
+            )
+        )
+        # Hidden to output.
+        self.layers.append(
+            HeteroRelationalGraphConv(
+                self.in_size,
+                self.out_size,
+                self.relation_names,
+            )
+        )
+
+    def forward(self, A):
+        h = self.embed_layer()
+        for layer in self.layers:
+            h = layer(A, h)
+        return h
+
+
+def main(args):
+    # Load graph data.
+    if args.dataset == "aifb":
+        dataset = AIFBDataset()
+    elif args.dataset == "bgs":
+        dataset = BGSDataset()
+    else:
+        raise ValueError()
+
+    g = dataset[0]
+    category = dataset.predict_category
+    num_classes = dataset.num_classes
+    train_mask = g.nodes[category].data.pop("train_mask")
+    test_mask = g.nodes[category].data.pop("test_mask")
+    train_idx = th.nonzero(train_mask, as_tuple=False).squeeze()
+    test_idx = th.nonzero(test_mask, as_tuple=False).squeeze()
+    labels = g.nodes[category].data.pop("labels")
+
+    # Split dataset into train, validate, test.
+    val_idx = train_idx[: len(train_idx) // 5]
+    train_idx = train_idx[len(train_idx) // 5 :]
+
+    embed_layer = RelGraphEmbed(
+        {ntype: g.num_nodes(ntype) for ntype in g.ntypes}, 16
+    )
+
+    # Create model.
+    model = EntityClassify(
+        16,
+        num_classes,
+        list(set(g.etypes)),
+        embed_layer,
+    )
+
+    # Optimizer.
+    optimizer = th.optim.Adam(model.parameters(), lr=1e-2, weight_decay=0)
+
+    # Construct hetero sparse matrix.
+    A = {}
+    for stype, etype, dtype in g.canonical_etypes:
+        eg = g[stype, etype, dtype]
+        indices = th.stack(eg.edges("uv"))
+        A[(stype, etype, dtype)] = dglsp.spmatrix(
+            indices, shape=(g.num_nodes(stype), g.num_nodes(dtype))
+        )
+        ###########################################################
+        # (HIGHLIGHT) Compute the normalized adjacency matrix with
+        # Sparse Matrix API
+        ###########################################################
+        D1_hat = dglsp.diag(A[(stype, etype, dtype)].sum(1)) ** -0.5
+        D2_hat = dglsp.diag(A[(stype, etype, dtype)].sum(0)) ** -0.5
+        A[(stype, etype, dtype)] = D1_hat @ A[(stype, etype, dtype)] @ D2_hat
+
+    # Training loop.
+    print("start training...")
+    model.train()
+    for epoch in range(10):
+        optimizer.zero_grad()
+        logits = model(A)[category]
+        loss = F.cross_entropy(logits[train_idx], labels[train_idx])
+        loss.backward()
+        optimizer.step()
+
+        train_acc = th.sum(
+            logits[train_idx].argmax(dim=1) == labels[train_idx]
+        ).item() / len(train_idx)
+        val_loss = F.cross_entropy(logits[val_idx], labels[val_idx])
+        val_acc = th.sum(
+            logits[val_idx].argmax(dim=1) == labels[val_idx]
+        ).item() / len(val_idx)
+        print(
+            f"Epoch {epoch:05d} | Train Acc: {train_acc:.4f} | "
+            f"Train Loss: {loss.item():.4f} | Valid Acc: {val_acc:.4f} | "
+            f"Valid loss: {val_loss.item():.4f} "
+        )
+    print()
+
+    model.eval()
+    logits = model.forward(A)[category]
+    test_loss = F.cross_entropy(logits[test_idx], labels[test_idx])
+    test_acc = th.sum(
+        logits[test_idx].argmax(dim=1) == labels[test_idx]
+    ).item() / len(test_idx)
+    print(
+        "Test Acc: {:.4f} | Test loss: {:.4f}".format(
+            test_acc, test_loss.item()
+        )
+    )
+    print()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="RGCN")
+    parser.add_argument(
+        "-d", "--dataset", type=str, required=True, help="dataset to use"
+    )
+
+    args = parser.parse_args()
+    print(args)
+    main(args)

examples/sparse/hgnn.py

@@ -0,0 +1,109 @@
+"""
+Hypergraph Neural Networks (https://arxiv.org/pdf/1809.09401.pdf)
+"""
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+from dgl.data import CoraGraphDataset
+from torchmetrics.functional import accuracy
+
+
+class HGNN(nn.Module):
+    def __init__(self, H, in_size, out_size, hidden_dims=16):
+        super().__init__()
+
+        self.Theta1 = nn.Linear(in_size, hidden_dims)
+        self.Theta2 = nn.Linear(hidden_dims, out_size)
+        self.dropout = nn.Dropout(0.5)
+
+        ###########################################################
+        # (HIGHLIGHT) Compute the Laplacian with Sparse Matrix API
+        ###########################################################
+        d_V = H.sum(1)  # node degree
+        d_E = H.sum(0)  # edge degree
+        n_edges = d_E.shape[0]
+        D_V_invsqrt = dglsp.diag(d_V**-0.5)  # D_V ** (-1/2)
+        D_E_inv = dglsp.diag(d_E**-1)  # D_E ** (-1)
+        W = dglsp.identity((n_edges, n_edges))
+        self.laplacian = D_V_invsqrt @ H @ W @ D_E_inv @ H.T @ D_V_invsqrt
+
+    def forward(self, X):
+        X = self.laplacian @ self.Theta1(self.dropout(X))
+        X = F.relu(X)
+        X = self.laplacian @ self.Theta2(self.dropout(X))
+        return X
+
+
+def train(model, optimizer, X, Y, train_mask):
+    model.train()
+    Y_hat = model(X)
+    loss = F.cross_entropy(Y_hat[train_mask], Y[train_mask])
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+
+def evaluate(model, X, Y, val_mask, test_mask, num_classes):
+    model.eval()
+    Y_hat = model(X)
+    val_acc = accuracy(
+        Y_hat[val_mask], Y[val_mask], task="multiclass", num_classes=num_classes
+    )
+    test_acc = accuracy(
+        Y_hat[test_mask],
+        Y[test_mask],
+        task="multiclass",
+        num_classes=num_classes,
+    )
+    return val_acc, test_acc
+
+
+def load_data():
+    dataset = CoraGraphDataset()
+
+    graph = dataset[0]
+    # The paper created a hypergraph from the original graph. For each node in
+    # the original graph, a hyperedge in the hypergraph is created to connect
+    # its neighbors and itself. In this case, the incidence matrix of the
+    # hypergraph is the same as the adjacency matrix of the original graph (with
+    # self-loops).
+    # We follow the paper and assume that the rows of the incidence matrix
+    # are for nodes and the columns are for edges.
+    indices = torch.stack(graph.edges())
+    H = dglsp.spmatrix(indices)
+    H = H + dglsp.identity(H.shape)
+
+    X = graph.ndata["feat"]
+    Y = graph.ndata["label"]
+    train_mask = graph.ndata["train_mask"]
+    val_mask = graph.ndata["val_mask"]
+    test_mask = graph.ndata["test_mask"]
+    return H, X, Y, dataset.num_classes, train_mask, val_mask, test_mask
+
+
+def main():
+    H, X, Y, num_classes, train_mask, val_mask, test_mask = load_data()
+    model = HGNN(H, X.shape[1], num_classes)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+    with tqdm.trange(500) as tq:
+        for epoch in tq:
+            train(model, optimizer, X, Y, train_mask)
+            val_acc, test_acc = evaluate(
+                model, X, Y, val_mask, test_mask, num_classes
+            )
+            tq.set_postfix(
+                {
+                    "Val acc": f"{val_acc:.5f}",
+                    "Test acc": f"{test_acc:.5f}",
+                },
+                refresh=False,
+            )
+
+    print(f"Test acc: {test_acc:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/sparse/hypergraphatt.py

@@ -0,0 +1,144 @@
+"""
+Hypergraph Convolution and Hypergraph Attention
+(https://arxiv.org/pdf/1901.08150.pdf).
+"""
+import argparse
+
+import dgl.sparse as dglsp
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import tqdm
+from dgl.data import CoraGraphDataset
+from torchmetrics.functional import accuracy
+
+
+def hypergraph_laplacian(H):
+    ###########################################################
+    # (HIGHLIGHT) Compute the Laplacian with Sparse Matrix API
+    ###########################################################
+    d_V = H.sum(1)  # node degree
+    d_E = H.sum(0)  # edge degree
+    n_edges = d_E.shape[0]
+    D_V_invsqrt = dglsp.diag(d_V**-0.5)  # D_V ** (-1/2)
+    D_E_inv = dglsp.diag(d_E**-1)  # D_E ** (-1)
+    W = dglsp.identity((n_edges, n_edges))
+    return D_V_invsqrt @ H @ W @ D_E_inv @ H.T @ D_V_invsqrt
+
+
+class HypergraphAttention(nn.Module):
+    """Hypergraph Attention module as in the paper
+    `Hypergraph Convolution and Hypergraph Attention
+    <https://arxiv.org/pdf/1901.08150.pdf>`_.
+    """
+
+    def __init__(self, in_size, out_size):
+        super().__init__()
+
+        self.P = nn.Linear(in_size, out_size)
+        self.a = nn.Linear(2 * out_size, 1)
+
+    def forward(self, H, X, X_edges):
+        Z = self.P(X)
+        Z_edges = self.P(X_edges)
+        sim = self.a(torch.cat([Z[H.row], Z_edges[H.col]], 1))
+        sim = F.leaky_relu(sim, 0.2).squeeze(1)
+        # Reassign the hypergraph new weights.
+        H_att = dglsp.val_like(H, sim)
+        H_att = H_att.softmax()
+        return hypergraph_laplacian(H_att) @ Z
+
+
+class Net(nn.Module):
+    def __init__(self, in_size, out_size, hidden_size=16):
+        super().__init__()
+
+        self.layer1 = HypergraphAttention(in_size, hidden_size)
+        self.layer2 = HypergraphAttention(hidden_size, out_size)
+
+    def forward(self, H, X):
+        Z = self.layer1(H, X, X)
+        Z = F.elu(Z)
+        Z = self.layer2(H, Z, Z)
+        return Z
+
+
+def train(model, optimizer, H, X, Y, train_mask):
+    model.train()
+    Y_hat = model(H, X)
+    loss = F.cross_entropy(Y_hat[train_mask], Y[train_mask])
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+    return loss.item()
+
+
+def evaluate(model, H, X, Y, val_mask, test_mask, num_classes):
+    model.eval()
+    Y_hat = model(H, X)
+    val_acc = accuracy(
+        Y_hat[val_mask], Y[val_mask], task="multiclass", num_classes=num_classes
+    )
+    test_acc = accuracy(
+        Y_hat[test_mask],
+        Y[test_mask],
+        task="multiclass",
+        num_classes=num_classes,
+    )
+    return val_acc, test_acc
+
+
+def load_data():
+    dataset = CoraGraphDataset()
+
+    graph = dataset[0]
+    # The paper created a hypergraph from the original graph. For each node in
+    # the original graph, a hyperedge in the hypergraph is created to connect
+    # its neighbors and itself. In this case, the incidence matrix of the
+    # hypergraph is the same as the adjacency matrix of the original graph (with
+    # self-loops).
+    # We follow the paper and assume that the rows of the incidence matrix
+    # are for nodes and the columns are for edges.
+    indices = torch.stack(graph.edges())
+    H = dglsp.spmatrix(indices)
+    H = H + dglsp.identity(H.shape)
+
+    X = graph.ndata["feat"]
+    Y = graph.ndata["label"]
+    train_mask = graph.ndata["train_mask"]
+    val_mask = graph.ndata["val_mask"]
+    test_mask = graph.ndata["test_mask"]
+    return H, X, Y, dataset.num_classes, train_mask, val_mask, test_mask
+
+
+def main(args):
+    H, X, Y, num_classes, train_mask, val_mask, test_mask = load_data()
+    model = Net(X.shape[1], num_classes)
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+    with tqdm.trange(args.epochs) as tq:
+        for epoch in tq:
+            loss = train(model, optimizer, H, X, Y, train_mask)
+            val_acc, test_acc = evaluate(
+                model, H, X, Y, val_mask, test_mask, num_classes
+            )
+            tq.set_postfix(
+                {
+                    "Loss": f"{loss:.5f}",
+                    "Val acc": f"{val_acc:.5f}",
+                    "Test acc": f"{test_acc:.5f}",
+                },
+                refresh=False,
+            )
+
+    print(f"Test acc: {test_acc:.3f}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Hypergraph Attention Example")
+    parser.add_argument(
+        "--epochs", type=int, default=500, help="Number of training epochs."
+    )
+    args = parser.parse_args()
+    main(args)

examples/sparse/pagerank.py

@@ -0,0 +1,31 @@
+import dgl.sparse as dglsp
+import networkx as nx
+import torch
+
+N = 100
+DAMP = 0.85
+K = 10
+
+
+def pagerank(A):
+    D = A.sum(0)
+    V = torch.ones(N) / N
+    for _ in range(K):
+        ########################################################################
+        # (HIGHLIGHT) Take the advantage of DGL sparse APIs to calculate the
+        # page rank.
+        ########################################################################
+        V = (1 - DAMP) / N + DAMP * A @ (V / D)
+    return V
+
+
+if __name__ == "__main__":
+    g = nx.erdos_renyi_graph(N, 0.05, seed=10086)
+
+    # Create the adjacency matrix of graph.
+    edges = list(g.to_directed().edges())
+    indices = torch.tensor(edges).transpose(0, 1)
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    V = pagerank(A)
+    print(V)

examples/sparse/sampling/graphsage.py

@@ -0,0 +1,233 @@
+"""
+This script demonstrate how to use dgl sparse library to sample on graph and 
+train model. It trains and tests a GraphSAGE model using the sparse sample and 
+compact operators to sample submatrix from the whole matrix.
+
+This flowchart describes the main functional sequence of the provided example.
+main
+│
+├───> Load and preprocess full dataset
+│
+├───> Instantiate SAGE model
+│
+├───> train
+│     │
+│     └───> Training loop
+│           │
+│           ├───> Sample submatrix
+│           │
+│           └───> SAGE.forward
+└───> test
+      │
+      ├───> Sample submatrix
+      │
+      └───> Evaluate the model
+"""
+import argparse
+
+import dgl.sparse as dglsp
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchmetrics.functional as MF
+from dgl.data import AsNodePredDataset
+from ogb.nodeproppred import DglNodePropPredDataset
+
+
+class SAGEConv(nn.Module):
+    r"""GraphSAGE layer from `Inductive Representation Learning on
+    Large Graphs <https://arxiv.org/pdf/1706.02216.pdf>`__
+    """
+
+    def __init__(
+        self,
+        in_size,
+        out_size,
+    ):
+        super(SAGEConv, self).__init__()
+        self._in_src_feats, self._in_dst_feats = in_size, in_size
+        self._out_size = out_size
+
+        self.fc_neigh = nn.Linear(self._in_src_feats, out_size, bias=False)
+        self.fc_self = nn.Linear(self._in_dst_feats, out_size, bias=True)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_uniform_(self.fc_self.weight, gain=gain)
+        nn.init.xavier_uniform_(self.fc_neigh.weight, gain=gain)
+
+    def forward(self, A, feat):
+        feat_src = feat
+        feat_dst = feat[: A.shape[1]]
+
+        # Aggregator type: mean.
+        srcdata = self.fc_neigh(feat_src)
+        # Divided by degree.
+        D_hat = dglsp.diag(A.sum(0)) ** -1
+        A_div = A @ D_hat
+        # Conv neighbors.
+        dstdata = A_div.T @ srcdata
+
+        rst = self.fc_self(feat_dst) + dstdata
+        return rst
+
+
+class SAGE(nn.Module):
+    def __init__(self, in_size, hid_size, out_size):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        # Three-layer GraphSAGE-gcn.
+        self.layers.append(SAGEConv(in_size, hid_size))
+        self.layers.append(SAGEConv(hid_size, hid_size))
+        self.layers.append(SAGEConv(hid_size, out_size))
+        self.dropout = nn.Dropout(0.5)
+        self.hid_size = hid_size
+        self.out_size = out_size
+
+    def forward(self, sampled_matrices, x):
+        hidden_x = x
+        for layer_idx, (layer, sampled_matrix) in enumerate(
+            zip(self.layers, sampled_matrices)
+        ):
+            hidden_x = layer(sampled_matrix, hidden_x)
+            if layer_idx != len(self.layers) - 1:
+                hidden_x = F.relu(hidden_x)
+                hidden_x = self.dropout(hidden_x)
+        return hidden_x
+
+
+def multilayer_sample(A, fanouts, seeds, ndata):
+    sampled_matrices = []
+    src = seeds
+
+    #####################################################################
+    # (HIGHLIGHT) Using the sparse sample operator to preform random
+    # sampling on the neighboring nodes of the seeds nodes. The sparse
+    # compact operator is then employed to compact and relabel the sampled
+    # matrix, resulting in the sampled matrix and the relabel index.
+    #####################################################################
+
+    for fanout in fanouts:
+        # Sample neighbors.
+        sampled_matrix = A.sample(1, fanout, ids=src).coalesce()
+        # Compact the sampled matrix.
+        compacted_mat, row_ids = sampled_matrix.compact(0)
+        sampled_matrices.insert(0, compacted_mat)
+        src = row_ids
+
+    x = ndata["feat"][src]
+    y = ndata["label"][seeds]
+    return sampled_matrices, x, y
+
+
+def evaluate(model, A, dataloader, ndata, num_classes):
+    model.eval()
+    ys = []
+    y_hats = []
+    fanouts = [10, 10, 10]
+    for it, seeds in enumerate(dataloader):
+        with torch.no_grad():
+            sampled_matrices, x, y = multilayer_sample(A, fanouts, seeds, ndata)
+            ys.append(y)
+            y_hats.append(model(sampled_matrices, x))
+
+    return MF.accuracy(
+        torch.cat(y_hats),
+        torch.cat(ys),
+        task="multiclass",
+        num_classes=num_classes,
+    )
+
+
+def validate(device, A, ndata, dataset, model, batch_size):
+    inf_id = dataset.test_idx.to(device)
+    inf_dataloader = torch.utils.data.DataLoader(inf_id, batch_size=batch_size)
+    acc = evaluate(model, A, inf_dataloader, ndata, dataset.num_classes)
+    return acc
+
+
+def train(device, A, ndata, dataset, model):
+    # Create sampler & dataloader.
+    train_idx = dataset.train_idx.to(device)
+    val_idx = dataset.val_idx.to(device)
+
+    train_dataloader = torch.utils.data.DataLoader(
+        train_idx, batch_size=1024, shuffle=True
+    )
+    val_dataloader = torch.utils.data.DataLoader(val_idx, batch_size=1024)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
+
+    fanouts = [10, 10, 10]
+    for epoch in range(10):
+        model.train()
+        total_loss = 0
+        for it, seeds in enumerate(train_dataloader):
+            sampled_matrices, x, y = multilayer_sample(A, fanouts, seeds, ndata)
+            y_hat = model(sampled_matrices, x)
+            loss = F.cross_entropy(y_hat, y)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+
+        acc = evaluate(model, A, val_dataloader, ndata, dataset.num_classes)
+        print(
+            "Epoch {:05d} | Loss {:.4f} | Accuracy {:.4f} ".format(
+                epoch, total_loss / (it + 1), acc.item()
+            )
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GraphSAGE")
+    parser.add_argument(
+        "--mode",
+        default="gpu",
+        choices=["cpu", "gpu"],
+        help="Training mode. 'cpu' for CPU training, 'gpu' for GPU training.",
+    )
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        args.mode = "cpu"
+    print(f"Training in {args.mode} mode.")
+
+    #####################################################################
+    # (HIGHLIGHT) This example implements a graphSAGE algorithm by sparse
+    # operators, which involves sampling a subgraph from a full graph and
+    # conducting training.
+    #
+    # First, the whole graph is loaded onto the CPU or GPU and transformed
+    # to sparse matrix. To obtain the training subgraph, it samples three
+    # submatrices by seed nodes, which contains their randomly sampled
+    # 1-hop, 2-hop, and 3-hop neighbors. Then, the features of the
+    # subgraph are input to the network for training.
+    #####################################################################
+
+    # Load and preprocess dataset.
+    print("Loading data")
+    device = torch.device("cpu" if args.mode == "cpu" else "cuda")
+    dataset = AsNodePredDataset(DglNodePropPredDataset("ogbn-products"))
+    g = dataset[0]
+    g = g.to(device)
+
+    # Create GraphSAGE model.
+    in_size = g.ndata["feat"].shape[1]
+    out_size = dataset.num_classes
+    model = SAGE(in_size, 256, out_size).to(device)
+
+    # Create sparse.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    # Model training.
+    print("Training...")
+    train(device, A, g.ndata, dataset, model)
+
+    # Test the model.
+    print("Testing...")
+    acc = validate(device, A, g.ndata, dataset, model, batch_size=4096)
+    print(f"Test accuracy {acc:.4f}")

examples/sparse/sampling/ladies.py

@@ -0,0 +1,256 @@
+"""
+This script demonstrates how to use dgl sparse library to sample on graph and 
+train model. It trains and tests a LADIES model using the sparse power and 
+sp_broadcast_v operators to sample submatrix from the whole matrix.
+
+This flowchart describes the main functional sequence of the provided example.
+main
+│
+├───> Load and preprocess full dataset
+│
+├───> Instantiate LADIES model
+│
+├───> train
+│     │
+│     └───> Training loop
+│           │
+│           ├───> Sample submatrix
+│           │
+│           └───> LADIES.forward
+└───> test
+      │
+      ├───> Sample submatrix
+      │
+      └───> Evaluate the model
+"""
+import argparse
+
+import dgl.sparse as dglsp
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchmetrics.functional as MF
+from dgl.data import AsNodePredDataset
+from dgl.sparse import sp_broadcast_v
+from ogb.nodeproppred import DglNodePropPredDataset
+
+
+class SAGEConv(nn.Module):
+    r"""LADIES layer from `Layer-Dependent Importance Sampling
+    for Training Deep and Large Graph Convolutional Networks
+    <https://arxiv.org/abs/1911.07323.pdf>`__"""
+
+    def __init__(
+        self,
+        in_size,
+        out_size,
+    ):
+        super(SAGEConv, self).__init__()
+        self._in_src_feats, self._in_dst_feats = in_size, in_size
+        self._out_size = out_size
+
+        self.fc_neigh = nn.Linear(self._in_src_feats, out_size, bias=False)
+        self.fc_self = nn.Linear(self._in_dst_feats, out_size, bias=True)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        gain = nn.init.calculate_gain("relu")
+        nn.init.xavier_uniform_(self.fc_self.weight, gain=gain)
+        nn.init.xavier_uniform_(self.fc_neigh.weight, gain=gain)
+
+    def forward(self, A, feat):
+        feat_src = feat
+        feat_dst = feat[: A.shape[1]]
+
+        # Aggregator type: mean.
+        srcdata = self.fc_neigh(feat_src)
+        # Divided by degree.
+        D_hat = dglsp.diag(A.sum(0)) ** -1
+        A_div = A @ D_hat
+        # Conv neighbors.
+        dstdata = A_div.T @ srcdata
+
+        rst = self.fc_self(feat_dst) + dstdata
+        return rst
+
+
+class LADIES(nn.Module):
+    def __init__(self, in_size, hid_size, out_size):
+        super().__init__()
+        self.layers = nn.ModuleList()
+        # Three-layer LADIES.
+        self.layers.append(SAGEConv(in_size, hid_size))
+        self.layers.append(SAGEConv(hid_size, hid_size))
+        self.layers.append(SAGEConv(hid_size, out_size))
+
+        self.dropout = nn.Dropout(0.5)
+        self.hid_size = hid_size
+        self.out_size = out_size
+
+    def forward(self, sampled_matrices, x):
+        hidden_x = x
+        for layer_idx, (layer, sampled_matrix) in enumerate(
+            zip(self.layers, sampled_matrices)
+        ):
+            hidden_x = layer(sampled_matrix, hidden_x)
+            if layer_idx != len(self.layers) - 1:
+                hidden_x = F.relu(hidden_x)
+                hidden_x = self.dropout(hidden_x)
+        return hidden_x
+
+
+def multilayer_sample(A, fanouts, seeds, ndata):
+    sampled_matrices = []
+    src = seeds
+
+    #########################################################################
+    # (HIGHLIGHT) Using the sparse sample operator to preform LADIES sampling
+    # algorithm from the neighboring nodes of the seeds nodes.
+    # The sparse sp_power operator is applied to compute sample probability,
+    # and sp_broadcast_v is then employed to normalize weight by performing
+    # division operations on column.
+    #########################################################################
+
+    for fanout in fanouts:
+        # Sample neighbors.
+        sub_A = A.index_select(1, src)
+        # Compute probability weight.
+        row_probs = (sub_A**2).sum(1)
+        row_probs = row_probs / row_probs.sum(0)
+        # Layer-wise sample nodes.
+        row_ids = torch.multinomial(row_probs, fanout, replacement=False)
+        # Add self-loop.
+        row_ids = torch.cat((row_ids, src), 0).unique()
+        sampled_matrix = sub_A.index_select(0, row_ids)
+        # Normalize edge weights.
+        div_matirx = sp_broadcast_v(
+            sampled_matrix, row_probs[row_ids].reshape(-1, 1), "truediv"
+        )
+        div_matirx = sp_broadcast_v(div_matirx, div_matirx.sum(0), "truediv")
+
+        # Save the sampled matrix.
+        sampled_matrices.insert(0, div_matirx)
+        src = row_ids
+
+    x = ndata["feat"][src]
+    y = ndata["label"][seeds]
+    return sampled_matrices, x, y
+
+
+def evaluate(model, A, dataloader, ndata, num_classes):
+    model.eval()
+    ys = []
+    y_hats = []
+    fanouts = [4000, 4000, 4000]
+    for seeds in dataloader:
+        with torch.no_grad():
+            sampled_matrices, x, y = multilayer_sample(A, fanouts, seeds, ndata)
+            ys.append(y)
+            y_hats.append(model(sampled_matrices, x))
+
+    return MF.accuracy(
+        torch.cat(y_hats),
+        torch.cat(ys),
+        task="multiclass",
+        num_classes=num_classes,
+    )
+
+
+def validate(device, A, ndata, dataset, model, batch_size):
+    inf_id = dataset.test_idx.to(device)
+    inf_dataloader = torch.utils.data.DataLoader(inf_id, batch_size=batch_size)
+    acc = evaluate(model, A, inf_dataloader, ndata, dataset.num_classes)
+    return acc
+
+
+def train(device, A, ndata, dataset, model):
+    # Create sampler & dataloader.
+    train_idx = dataset.train_idx.to(device)
+    val_idx = dataset.val_idx.to(device)
+
+    train_dataloader = torch.utils.data.DataLoader(
+        train_idx, batch_size=1024, shuffle=True
+    )
+    val_dataloader = torch.utils.data.DataLoader(val_idx, batch_size=1024)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=5e-4)
+
+    fanouts = [4000, 4000, 4000]
+    for epoch in range(20):
+        model.train()
+        total_loss = 0
+        for it, seeds in enumerate(train_dataloader):
+            sampled_matrices, x, y = multilayer_sample(A, fanouts, seeds, ndata)
+            y_hat = model(sampled_matrices, x)
+            loss = F.cross_entropy(y_hat, y)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+
+        acc = evaluate(model, A, val_dataloader, ndata, dataset.num_classes)
+        print(
+            "Epoch {:05d} | Loss {:.4f} | Accuracy {:.4f} ".format(
+                epoch, total_loss / (it + 1), acc.item()
+            )
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LADIESConv")
+    parser.add_argument(
+        "--mode",
+        default="gpu",
+        choices=["cpu", "gpu"],
+        help="Training mode. 'cpu' for CPU training, 'gpu' for GPU training.",
+    )
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        args.mode = "cpu"
+    print(f"Training in {args.mode} mode.")
+
+    #####################################################################
+    # (HIGHLIGHT) This example implements a LADIES algorithm by sparse
+    # operators, which involves sampling a subgraph from a full graph and
+    # conducting training.
+    #
+    # First, the whole graph is loaded onto the CPU or GPU and transformed
+    # to sparse matrix. To obtain the training subgraph, it samples three
+    # submatrices by seed nodes, which contains their layer-wise sampled
+    # 1-hop, 2-hop, and 3-hop neighbors. Then, the features of the
+    # subgraph are input to the network for training.
+    #####################################################################
+
+    # Load and preprocess dataset.
+    print("Loading data")
+    device = torch.device("cpu" if args.mode == "cpu" else "cuda")
+    dataset = AsNodePredDataset(DglNodePropPredDataset("ogbn-products"))
+    g = dataset[0]
+
+    # Create LADIES model.
+    in_size = g.ndata["feat"].shape[1]
+    out_size = dataset.num_classes
+    model = LADIES(in_size, 256, out_size).to(device)
+
+    # Create sparse.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N)).coalesce()
+    I = dglsp.identity(A.shape)
+
+    # Initialize laplacian matrix.
+    A_hat = A + I
+    D_hat = dglsp.diag(A_hat.sum(1)) ** -0.5
+    A_norm = D_hat @ A_hat @ D_hat
+    A_norm = A_norm.to(device)
+    g = g.to(device)
+
+    # Model training.
+    print("Training...")
+    train(device, A_norm, g.ndata, dataset, model)
+
+    # Test the model.
+    print("Testing...")
+    acc = validate(device, A_norm, g.ndata, dataset, model, batch_size=2048)
+    print(f"Test accuracy {acc:.4f}")

examples/sparse/sgc.py

@@ -0,0 +1,94 @@
+"""
+[Simplifying Graph Convolutional Networks]
+(https://arxiv.org/abs/1902.07153)
+"""
+
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+################################################################################
+# (HIGHLIGHT) Take the advantage of DGL sparse APIs to implement the feature
+# pre-computation.
+################################################################################
+def pre_compute(A, X, k):
+    for _ in range(k):
+        X = A @ X
+    return X
+
+
+def evaluate(g, pred):
+    label = g.ndata["label"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(model, g, X_sgc):
+    label = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    optimizer = Adam(model.parameters(), lr=2e-1, weight_decay=5e-6)
+
+    for epoch in range(20):
+        # Forward.
+        logits = model(X_sgc)
+
+        # Compute loss with nodes in the training set.
+        loss = F.cross_entropy(logits[train_mask], label[train_mask])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute prediction.
+        pred = logits.argmax(dim=1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(g, pred)
+        print(
+            f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}, test"
+            f" acc: {test_acc:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+
+    # Create the sparse adjacency matrix A
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    # Calculate the symmetrically normalized adjacency matrix.
+    I = dglsp.identity(A.shape, device=dev)
+    A_hat = A + I
+    D_hat = dglsp.diag(A_hat.sum(dim=1)) ** -0.5
+    A_hat = D_hat @ A_hat @ D_hat
+
+    # 2-hop diffusion.
+    k = 2
+    X = g.ndata["feat"]
+    X_sgc = pre_compute(A_hat, X, k)
+
+    # Create model.
+    in_size = X.shape[1]
+    out_size = dataset.num_classes
+    model = nn.Linear(in_size, out_size).to(dev)
+
+    # Kick off training.
+    train(model, g, X_sgc)

examples/sparse/sign.py

@@ -0,0 +1,127 @@
+"""
+[SIGN: Scalable Inception Graph Neural Networks]
+(https://arxiv.org/abs/2004.11198)
+
+This example shows a simplified version of SIGN: a precomputed 2-hops diffusion
+operator on top of symmetrically normalized adjacency matrix A_hat.
+"""
+
+import dgl.sparse as dglsp
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+################################################################################
+# (HIGHLIGHT) Take the advantage of DGL sparse APIs to implement the feature
+# diffusion in SIGN laconically.
+################################################################################
+def sign_diffusion(A, X, r):
+    # Perform the r-hop diffusion operation.
+    X_sign = [X]
+    for _ in range(r):
+        X = A @ X
+        X_sign.append(X)
+    return X_sign
+
+
+class SIGN(nn.Module):
+    def __init__(self, in_size, out_size, r, hidden_size=256):
+        super().__init__()
+        # Note that theta and omega refer to the learnable matrices in the
+        # original paper correspondingly. The variable r refers to subscript to
+        # theta.
+        self.theta = nn.ModuleList(
+            [nn.Linear(in_size, hidden_size) for _ in range(r + 1)]
+        )
+        self.omega = nn.Linear(hidden_size * (r + 1), out_size)
+
+    def forward(self, X_sign):
+        results = []
+        for i in range(len(X_sign)):
+            results.append(self.theta[i](X_sign[i]))
+        Z = F.relu(torch.cat(results, dim=1))
+        return self.omega(Z)
+
+
+def evaluate(g, pred):
+    label = g.ndata["label"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(model, g, X_sign):
+    label = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    optimizer = Adam(model.parameters(), lr=3e-3)
+
+    for epoch in range(10):
+        # Switch the model to training mode.
+        model.train()
+
+        # Forward.
+        logits = model(X_sign)
+
+        # Compute loss with nodes in training set.
+        loss = F.cross_entropy(logits[train_mask], label[train_mask])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Switch the model to evaluating mode.
+        model.eval()
+
+        # Compute prediction.
+        logits = model(X_sign)
+        pred = logits.argmax(1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(g, pred)
+        print(
+            f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}, test"
+            f" acc: {test_acc:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+
+    # Create the sparse adjacency matrix A (note that W was used as the notation
+    # for adjacency matrix in the original paper).
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    # Calculate the symmetrically normalized adjacency matrix.
+    I = dglsp.identity(A.shape, device=dev)
+    A_hat = A + I
+    D_hat = dglsp.diag(A_hat.sum(dim=1)) ** -0.5
+    A_hat = D_hat @ A_hat @ D_hat
+
+    # 2-hop diffusion.
+    r = 2
+    X = g.ndata["feat"]
+    X_sign = sign_diffusion(A_hat, X, r)
+
+    # Create SIGN model.
+    in_size = X.shape[1]
+    out_size = dataset.num_classes
+    model = SIGN(in_size, out_size, r).to(dev)
+
+    # Kick off training.
+    train(model, g, X_sign)

examples/sparse/twirls.py

@@ -0,0 +1,200 @@
+"""
+[Graph Neural Networks Inspired by Classical Iterative Algorithms]
+(https://arxiv.org/pdf/2103.06064.pdf)
+
+This example shows a simplified version of the TWIRLS model proposed
+in the paper. It implements two variants. One is the basic iterative
+graph diffusion algorithm. The other is an advanced implementation
+with attention.
+"""
+
+import argparse
+
+import dgl.sparse as dglsp
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.data import CoraGraphDataset
+from torch.optim import Adam
+
+
+class MLP(nn.Module):
+    def __init__(self, in_size, hidden_size):
+        super().__init__()
+        self.linear_1 = nn.Linear(in_size, hidden_size)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size)
+        self.dropout = nn.Dropout(0.8)
+
+    def forward(self, X):
+        H = self.linear_1(X)
+        H = F.relu(H)
+        H = self.dropout(H)
+        H = self.linear_2(H)
+        return H
+
+
+################################################################################
+# (HIGHLIGHT) Use DGL sparse API to implement the iterative graph diffusion
+# algorithm.
+################################################################################
+class TWIRLS(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        hidden_size=128,
+        num_steps=16,
+        lam=1.0,
+        alpha=0.5,
+    ):
+        super().__init__()
+        self.num_steps = num_steps
+        self.lam = lam
+        self.alpha = alpha
+        self.mlp = MLP(in_size, hidden_size)
+        self.linear_out = nn.Linear(hidden_size, out_size)
+
+    def forward(self, A, X):
+        # Compute Y = Y0 = f(X; W) using a two-layer MLP.
+        Y = Y0 = self.mlp(X)
+
+        # Compute diagonal matrix D_tild.
+        I = dglsp.identity(A.shape, device=A.device)
+        D_tild = self.lam * dglsp.diag(A.sum(1)) + I
+
+        # Iteratively compute new Y by equation (6) in the paper.
+        for k in range(self.num_steps):
+            Y_hat = self.lam * A @ Y + Y0
+            # The inverse of a diagonal matrix inverses its diagonal values.
+            Y = (1 - self.alpha) * Y + self.alpha * (D_tild**-1) @ Y_hat
+
+        # Apply a linear layer on the final output.
+        return self.linear_out(Y)
+
+
+################################################################################
+# (HIGHLIGHT) Implementation of the advanced TWIRLS model with attention
+# to show the usage of differentiable weighted sparse matrix.
+################################################################################
+class TWIRLSWithAttention(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        out_size,
+        hidden_size=128,
+        num_steps=16,
+        lam=1.0,
+        alpha=0.5,
+    ):
+        super().__init__()
+        self.num_steps = num_steps
+        self.lam = lam
+        self.alpha = alpha
+        self.mlp = MLP(in_size, hidden_size)
+        self.linear_out = nn.Linear(hidden_size, out_size)
+
+    def forward(self, A, X):
+        # Compute Y = Y0 = f(X; W) using a two-layer MLP.
+        Y = Y0 = self.mlp(X)
+
+        # Compute diagonal matrix D_tild.
+        I = dglsp.identity(A.shape, device=A.device)
+        D_tild = self.lam * dglsp.diag(A.sum(1)) + I
+
+        # Conduct half of the diffusion steps.
+        for k in range(self.num_steps // 2):
+            Y_hat = self.lam * A @ Y + Y0
+            Y = (1 - self.alpha) * Y + self.alpha * (D_tild**-1) @ Y_hat
+
+        # Calculate attention weight by equation (25) in the paper.
+        Y_i = Y[A.row]
+        Y_j = Y[A.col]
+        norm_ij = torch.linalg.vector_norm(Y_i - Y_j, dim=1)
+        # Bound the attention value within [0.0, 1.0).
+        gamma_ij = torch.clamp(0.5 / (norm_ij + 1e-7), min=0.0, max=1.0)
+        # Create a new adjacency matrix with the new weight.
+        A = dglsp.val_like(A, gamma_ij)
+        # Recompute D_tild.
+        D_tild = self.lam * dglsp.diag(A.sum(1)) + I
+
+        # Conduct the other half of the diffusion steps.
+        for k in range(self.num_steps // 2):
+            Y_hat = self.lam * A @ Y + Y0
+            Y = (1 - self.alpha) * Y + self.alpha * (D_tild**-1) @ Y_hat
+
+        # Apply a linear layer on the final output.
+        return self.linear_out(Y)
+
+
+def evaluate(g, pred):
+    model.eval()
+    label = g.ndata["label"]
+    val_mask = g.ndata["val_mask"]
+    test_mask = g.ndata["test_mask"]
+
+    # Compute accuracy on validation/test set.
+    val_acc = (pred[val_mask] == label[val_mask]).float().mean()
+    test_acc = (pred[test_mask] == label[test_mask]).float().mean()
+    return val_acc, test_acc
+
+
+def train(g, model, A, X):
+    labels = g.ndata["label"]
+    train_mask = g.ndata["train_mask"]
+    optimizer = Adam(model.parameters(), lr=5e-4)
+
+    for epoch in range(300):
+        model.train()
+        # Forward.
+        logits = model(A, X)
+
+        # Compute loss with nodes in training set.
+        loss = F.cross_entropy(logits[train_mask], labels[train_mask])
+
+        # Backward.
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # Compute prediction.
+        pred = logits.argmax(1)
+
+        # Evaluate the prediction.
+        val_acc, test_acc = evaluate(g, pred)
+        print(
+            f"In epoch {epoch}, loss: {loss:.3f}, val acc: {val_acc:.3f}, test"
+            f" acc: {test_acc:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("TWIRLS example in DGL Sparse.")
+    parser.add_argument(
+        "--attention", action="store_true", help="Use TWIRLS with attention."
+    )
+    args = parser.parse_args()
+    # If CUDA is available, use GPU to accelerate the training, use CPU
+    # otherwise.
+    dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    # Load graph from the existing dataset.
+    dataset = CoraGraphDataset()
+    g = dataset[0].to(dev)
+    X = g.ndata["feat"]
+
+    # Create the sparse adjacency matrix A.
+    indices = torch.stack(g.edges())
+    N = g.num_nodes()
+    A = dglsp.spmatrix(indices, shape=(N, N))
+
+    # Create the TWIRLS model.
+    in_size = X.shape[1]
+    out_size = dataset.num_classes
+    if args.attention:
+        model = TWIRLSWithAttention(in_size, out_size).to(dev)
+    else:
+        model = TWIRLS(in_size, out_size).to(dev)
+
+    # Kick off training.
+    train(g, model, A, X)