Repository: pyg-team/pytorch_geometric
Pull request #7299 (2023-05-05 05:15:42)

rebase master branch
rebase master branch


Patch:
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
new file mode 100644
index 000000000..ac674f04e
--- /dev/null
+++ b/.github/CODEOWNERS
@@ -0,0 +1,34 @@
+# About code owners
+# https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
+#
+# Code owners recieve review requests when parts of code they "own" change.
+
+*.py @rusty1s @wsad1
+
+/torch_geometric/nn/ @EdisonLeeeee
+
+/torch_geometric/explain/ @RexYing @RBendias @dufourc1
+
+/torch_geometric/data/ @mananshah99
+
+/torch_geometric/loader/ @mananshah99 @pyg-team/intel-team
+
+/torch_geometric/sampler/ @mananshah99 @pyg-team/intel-team
+
+/torch_geometric/transforms/ @wsad1
+
+/torch_geometric/utils/ @wsad1
+
+/torch_geometric/datasets/ @wsad1
+
+/torch_geometric/graphgym/ @JiaxuanYou
+
+/test/ @wsad1
+
+/docs/ @rusty1s
+
+/graphgym/ @JiaxuanYou
+
+/examples/ @wsad1
+
+/benchmark/ @pyg-team/intel-team
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 91cde2ba5..73052b80b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,16 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for tuples as keys in `ModuleDict`/`ParameterDict` ([#7294](https://github.com/pyg-team/pytorch_geometric/pull/7294))
+- Added `NodePropertySplit` transform for creating node-level splits using structural node properties ([#6894](https://github.com/pyg-team/pytorch_geometric/pull/6894))
+- Added an option to preserve directed graphs in `CitationFull` datasets ([#7275](https://github.com/pyg-team/pytorch_geometric/pull/7275))
+- Added support for `torch.sparse.Tensor` in `DataLoader` ([#7252](https://github.com/pyg-team/pytorch_geometric/pull/7252))
+- Added `save` and `load` methods to `InMemoryDataset` ([#7250](https://github.com/pyg-team/pytorch_geometric/pull/7250))
+- Added an example for heterogeneous GNN explanation via `CaptumExplainer` ([#7096](https://github.com/pyg-team/pytorch_geometric/pull/7096))
+- Added `visualize_feature_importance` functionality to `HeteroExplanation` ([#7096](https://github.com/pyg-team/pytorch_geometric/pull/7096))
+- Added a `AddRemainingSelfLoops` transform ([#7192](https://github.com/pyg-team/pytorch_geometric/pull/7192))
+- Added `optimizer_resolver` ([#7209](https://github.com/pyg-team/pytorch_geometric/pull/7209))
+- Added `type_ptr` argument to `HeteroLayerNorm` ([#7208](https://github.com/pyg-team/pytorch_geometric/pull/7208))
 - Added an option to benchmark scripts to write PyTorch profiler results to CSV ([#7114](https://github.com/pyg-team/pytorch_geometric/pull/7114))
 - Added subgraph type sampling option with bidirectional edge support ([#7199](https://github.com/pyg-team/pytorch_geometric/pull/7199), [#7200](https://github.com/pyg-team/pytorch_geometric/pull/7200))
 - Added support for `"any"`-reductions in `scatter` ([#7198](https://github.com/pyg-team/pytorch_geometric/pull/7198))
@@ -24,6 +34,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Updated examples to use `NeighborLoader` instead of `NeighborSampler` ([#7152](https://github.com/pyg-team/pytorch_geometric/pull/7152))
+- Fixed `HGTConv` utility function `_construct_src_node_feat` ([#7194](https://github.com/pyg-team/pytorch_geometric/pull/7194))
+- Extend dataset summary to create stats for each node/edge type ([#7203](https://github.com/pyg-team/pytorch_geometric/pull/7203))
+- Added an optional `batch_size` argument to `avg_pool_x` and `max_pool_x` ([#7216](https://github.com/pyg-team/pytorch_geometric/pull/7216))
 - Fixed `subgraph` on unordered inputs ([#7187](https://github.com/pyg-team/pytorch_geometric/pull/7187))
 - Allow missing node types in `HeteroDictLinear` ([#7185](https://github.com/pyg-team/pytorch_geometric/pull/7185))
 - Optimized `from_networkx` memory footprint by reducing unnecessary copies ([#7119](https://github.com/pyg-team/pytorch_geometric/pull/7119))
@@ -37,6 +51,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed a bug in `FastHGTConv` that computed values via parameters used to compute the keys ([#7050](https://github.com/pyg-team/pytorch_geometric/pull/7050))
 - Accelerated sparse tensor conversion routines ([#7042](https://github.com/pyg-team/pytorch_geometric/pull/7042), [#7043](https://github.com/pyg-team/pytorch_geometric/pull/7043))
 - Change `torch_sparse.SparseTensor` logic to utilize `torch.sparse_csr` instead ([#7041](https://github.com/pyg-team/pytorch_geometric/pull/7041))
+- Added an optional `batch_size` and `max_num_nodes` arguments to `MemPooling` layer ([#7239](https://github.com/pyg-team/pytorch_geometric/pull/7239))
 
 ### Removed
 
diff --git a/benchmark/kernel/datasets.py b/benchmark/kernel/datasets.py
index e7accdb55..0033db2e5 100644
--- a/benchmark/kernel/datasets.py
+++ b/benchmark/kernel/datasets.py
@@ -7,7 +7,7 @@ from torch_geometric.datasets import TUDataset
 from torch_geometric.utils import degree
 
 
-class NormalizedDegree(object):
+class NormalizedDegree:
     def __init__(self, mean, std):
         self.mean = mean
         self.std = std
diff --git a/benchmark/runtime/dgl/hidden.py b/benchmark/runtime/dgl/hidden.py
index d38590e8c..69e19bd6e 100644
--- a/benchmark/runtime/dgl/hidden.py
+++ b/benchmark/runtime/dgl/hidden.py
@@ -5,7 +5,7 @@ import warnings
 warnings.filterwarnings('ignore')
 
 
-class HiddenPrint(object):
+class HiddenPrint:
     def __enter__(self):
         self._original_stdout = sys.stdout
         sys.stdout = open(os.devnull, 'w')
diff --git a/docs/source/advanced/graphgym.rst b/docs/source/advanced/graphgym.rst
index 22d879ebd..344898fab 100644
--- a/docs/source/advanced/graphgym.rst
+++ b/docs/source/advanced/graphgym.rst
@@ -107,12 +107,7 @@ To use GraphGym, you need to clone :pyg:`PyG` from :github:`GitHub`, then change
        bash run_batch.sh # run a batch of experiments
 
 #. **Run GraphGym with CPU backend:**
-   GraphGym supports CPU backend as well -- you only need to add the line :obj:`device: cpu` to the :obj:`*.yaml` file.
-   Here we provide an example:
-
-   .. code-block:: bash
-
-       bash run_single_cpu.sh # run a single experiment using CPU backend
+   GraphGym supports CPU backend as well -- you only need to add the line :obj:`accelerator: cpu` to the :obj:`*.yaml` file.
 
 In-Depth Usage
 --------------
diff --git a/docs/source/get_started/colabs.rst b/docs/source/get_started/colabs.rst
index 46a20acbf..9ad7726d2 100644
--- a/docs/source/get_started/colabs.rst
+++ b/docs/source/get_started/colabs.rst
@@ -16,6 +16,7 @@ We have prepared a list of :colab:`Colab` notebooks that practically introduces
 8. `Node Classification Instrumented with <https://colab.research.google.com/github/wandb/examples/blob/master/colabs/pyg/8_Node_Classification_(with_W&B).ipynb>`__ :wandb:`null` `Weights&Biases <https://colab.research.google.com/github/wandb/examples/blob/master/colabs/pyg/8_Node_Classification_(with_W&B).ipynb>`__
 9. `Graph Classification Instrumented with <https://colab.research.google.com/github/wandb/examples/blob/pyg/graph-classification/colabs/pyg/Graph_Classification_with_PyG_and_W%26B.ipynb>`__ :wandb:`null` `Weights&Biases <https://colab.research.google.com/github/wandb/examples/blob/pyg/graph-classification/colabs/pyg/Graph_Classification_with_PyG_and_W%26B.ipynb>`__
 10. `Link Prediction on MovieLens <https://colab.research.google.com/drive/1xpzn1Nvai1ygd_P5Yambc_oe4VBPK_ZT?usp=sharing>`__
+11. `Link Regression on MovieLens <https://colab.research.google.com/drive/1N3LvAO0AXV4kBPbTMX866OwJM9YS6Ji2?usp=sharing>`__
 
 All :colab:`Colab` notebooks are released under the MIT license.
 
diff --git a/docs/source/tutorial/load_csv.rst b/docs/source/tutorial/load_csv.rst
index 57281859c..22ff4402a 100644
--- a/docs/source/tutorial/load_csv.rst
+++ b/docs/source/tutorial/load_csv.rst
@@ -118,7 +118,7 @@ For this, we make use of the excellent `sentence-transformers <https://www.sbert
 
 .. code-block:: python
 
-    class SequenceEncoder(object):
+    class SequenceEncoder:
         def __init__(self, model_name='all-MiniLM-L6-v2', device=None):
             self.device = device
             self.model = SentenceTransformer(model_name, device=device)
@@ -137,7 +137,7 @@ For this, we first need to find all existing genres present in the data, create
 
 .. code-block:: python
 
-    class GenresEncoder(object):
+    class GenresEncoder:
         def __init__(self, sep='|'):
             self.sep = sep
 
@@ -218,7 +218,7 @@ For example, for loading the ratings from the :obj:`rating` column in :obj:`rati
 
 .. code-block:: python
 
-    class IdentityEncoder(object):
+    class IdentityEncoder:
         def __init__(self, dtype=None):
             self.dtype = dtype
 
diff --git a/examples/attentive_fp.py b/examples/attentive_fp.py
index ced66593f..7bd9baee5 100644
--- a/examples/attentive_fp.py
+++ b/examples/attentive_fp.py
@@ -10,7 +10,7 @@ from torch_geometric.loader import DataLoader
 from torch_geometric.nn.models import AttentiveFP
 
 
-class GenFeatures(object):
+class GenFeatures:
     def __init__(self):
         self.symbols = [
             'B', 'C', 'N', 'O', 'F', 'Si', 'P', 'S', 'Cl', 'As', 'Se', 'Br',
diff --git a/examples/colors_topk_pool.py b/examples/colors_topk_pool.py
index 7fb17fc55..ce9d91e2a 100644
--- a/examples/colors_topk_pool.py
+++ b/examples/colors_topk_pool.py
@@ -12,7 +12,7 @@ from torch_geometric.nn import GINConv, TopKPooling, global_add_pool
 from torch_geometric.utils import scatter
 
 
-class HandleNodeAttention(object):
+class HandleNodeAttention:
     def __call__(self, data):
         data.attn = torch.softmax(data.x[:, 0], dim=0)
         data.x = data.x[:, 1:]
diff --git a/examples/cora.py b/examples/cora.py
index 76f6142bd..d30a0200f 100644
--- a/examples/cora.py
+++ b/examples/cora.py
@@ -6,6 +6,10 @@ import torch.nn.functional as F
 import torch_geometric.transforms as T
 from torch_geometric.datasets import Planetoid
 from torch_geometric.nn import SplineConv
+from torch_geometric.typing import WITH_TORCH_SPLINE_CONV
+
+if not WITH_TORCH_SPLINE_CONV:
+    quit("This example requires 'torch-spline-conv'")
 
 dataset = 'Cora'
 transform = T.Compose([
diff --git a/examples/explain/captum_explainer_hetero_link.py b/examples/explain/captum_explainer_hetero_link.py
new file mode 100644
index 000000000..3523d07ba
--- /dev/null
+++ b/examples/explain/captum_explainer_hetero_link.py
@@ -0,0 +1,120 @@
+import os.path as osp
+
+import torch
+import torch.nn.functional as F
+from torch.nn import Linear
+
+import torch_geometric.transforms as T
+from torch_geometric.datasets import MovieLens
+from torch_geometric.explain import CaptumExplainer, Explainer
+from torch_geometric.nn import SAGEConv, to_hetero
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data/MovieLens')
+dataset = MovieLens(path, model_name='all-MiniLM-L6-v2')
+data = dataset[0].to(device)
+
+# Add user node features for message passing:
+data['user'].x = torch.eye(data['user'].num_nodes, device=device)
+del data['user'].num_nodes
+
+# Add a reverse ('movie', 'rev_rates', 'user') relation for message passing:
+data = T.ToUndirected()(data)
+data['user', 'movie'].edge_label = data['user',
+                                        'movie'].edge_label.to(torch.float)
+del data['movie', 'rev_rates', 'user'].edge_label  # Remove "reverse" label.
+
+# Perform a link-level split into training, validation, and test edges:
+data, _, _ = T.RandomLinkSplit(
+    num_val=0.1,
+    num_test=0.1,
+    neg_sampling_ratio=0.0,
+    edge_types=[('user', 'rates', 'movie')],
+    rev_edge_types=[('movie', 'rev_rates', 'user')],
+)(data)
+
+
+class GNNEncoder(torch.nn.Module):
+    def __init__(self, hidden_channels, out_channels):
+        super().__init__()
+        self.conv1 = SAGEConv((-1, -1), hidden_channels)
+        self.conv2 = SAGEConv((-1, -1), out_channels)
+
+    def forward(self, x, edge_index):
+        x = self.conv1(x, edge_index).relu()
+        x = self.conv2(x, edge_index)
+        return x
+
+
+class EdgeDecoder(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
+        self.lin2 = Linear(hidden_channels, 1)
+
+    def forward(self, z_dict, edge_label_index):
+        row, col = edge_label_index
+        z = torch.cat([z_dict['user'][row], z_dict['movie'][col]], dim=-1)
+
+        z = self.lin1(z).relu()
+        z = self.lin2(z)
+        return z.view(-1)
+
+
+class Model(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
+        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
+        self.decoder = EdgeDecoder(hidden_channels)
+
+    def forward(self, x_dict, edge_index_dict, edge_label_index):
+        z_dict = self.encoder(x_dict, edge_index_dict)
+        return self.decoder(z_dict, edge_label_index)
+
+
+model = Model(hidden_channels=32).to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+for epoch in range(1, 10):
+    model.train()
+    optimizer.zero_grad()
+    pred = model(
+        data.x_dict,
+        data.edge_index_dict,
+        data['user', 'movie'].edge_label_index,
+    )
+    loss = F.mse_loss(pred, data['user', 'movie'].edge_label)
+    loss.backward()
+    optimizer.step()
+
+explainer = Explainer(
+    model=model,
+    algorithm=CaptumExplainer('IntegratedGradients'),
+    explanation_type='model',
+    model_config=dict(
+        mode='regression',
+        task_level='edge',
+        return_type='raw',
+    ),
+    node_mask_type='attributes',
+    edge_mask_type='object',
+    threshold_config=dict(
+        threshold_type='topk',
+        value=200,
+    ),
+)
+
+index = torch.tensor([2, 10])  # Explain edge labels with index 2 and 10.
+explanation = explainer(
+    data.x_dict,
+    data.edge_index_dict,
+    index=index,
+    edge_label_index=data['user', 'movie'].edge_label_index,
+)
+print(f'Generated explanations in {explanation.available_explanations}')
+
+path = 'feature_importance.png'
+explanation.visualize_feature_importance(path, top_k=10)
+print(f"Feature importance plot has been saved to '{path}'")
diff --git a/examples/faust.py b/examples/faust.py
index 0d63e2049..899c6a196 100644
--- a/examples/faust.py
+++ b/examples/faust.py
@@ -7,6 +7,10 @@ import torch_geometric.transforms as T
 from torch_geometric.datasets import FAUST
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import SplineConv
+from torch_geometric.typing import WITH_TORCH_SPLINE_CONV
+
+if not WITH_TORCH_SPLINE_CONV:
+    quit("This example requires 'torch-spline-conv'")
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'FAUST')
 pre_transform = T.Compose([T.FaceToEdge(), T.Constant(value=1)])
diff --git a/examples/gcn2_ppi.py b/examples/gcn2_ppi.py
index 87fcbb4ce..a22cc5ffa 100644
--- a/examples/gcn2_ppi.py
+++ b/examples/gcn2_ppi.py
@@ -9,10 +9,6 @@ import torch_geometric.transforms as T
 from torch_geometric.datasets import PPI
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import GCN2Conv
-from torch_geometric.typing import WITH_TORCH_SPARSE
-
-if not WITH_TORCH_SPARSE:
-    quit("This example requires 'torch-sparse'")
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'GCN2_PPI')
 pre_transform = T.Compose([T.GCNNorm(), T.ToSparseTensor()])
diff --git a/examples/hetero/hetero_link_pred.py b/examples/hetero/hetero_link_pred.py
index d4bdcf8e5..b77393c8b 100644
--- a/examples/hetero/hetero_link_pred.py
+++ b/examples/hetero/hetero_link_pred.py
@@ -91,12 +91,6 @@ class Model(torch.nn.Module):
 
 
 model = Model(hidden_channels=32).to(device)
-
-# Due to lazy initialization, we need to run one model step so the number
-# of parameters can be inferred:
-with torch.no_grad():
-    model.encoder(train_data.x_dict, train_data.edge_index_dict)
-
 optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 
 
diff --git a/examples/hetero/load_csv.py b/examples/hetero/load_csv.py
index edc4d72ed..c6057bb6f 100644
--- a/examples/hetero/load_csv.py
+++ b/examples/hetero/load_csv.py
@@ -42,7 +42,7 @@ def load_edge_csv(path, src_index_col, src_mapping, dst_index_col, dst_mapping,
     return edge_index, edge_attr
 
 
-class SequenceEncoder(object):
+class SequenceEncoder:
     # The 'SequenceEncoder' encodes raw column strings into embeddings.
     def __init__(self, model_name='all-MiniLM-L6-v2', device=None):
         self.device = device
@@ -55,7 +55,7 @@ class SequenceEncoder(object):
         return x.cpu()
 
 
-class GenresEncoder(object):
+class GenresEncoder:
     # The 'GenreEncoder' splits the raw column strings by 'sep' and converts
     # individual elements to categorical labels.
     def __init__(self, sep='|'):
@@ -72,7 +72,7 @@ class GenresEncoder(object):
         return x
 
 
-class IdentityEncoder(object):
+class IdentityEncoder:
     # The 'IdentityEncoder' takes the raw column values and converts them to
     # PyTorch tensors.
     def __init__(self, dtype=None):
diff --git a/examples/infomax_inductive.py b/examples/infomax_inductive.py
index 74a5354dd..73e96e137 100644
--- a/examples/infomax_inductive.py
+++ b/examples/infomax_inductive.py
@@ -5,20 +5,18 @@ import torch.nn as nn
 from tqdm import tqdm
 
 from torch_geometric.datasets import Reddit
-from torch_geometric.loader import NeighborSampler
+from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import DeepGraphInfomax, SAGEConv
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Reddit')
 dataset = Reddit(path)
-data = dataset[0]
-
-train_loader = NeighborSampler(data.edge_index, node_idx=None,
-                               sizes=[10, 10, 25], batch_size=256,
-                               shuffle=True, num_workers=12)
+data = dataset[0].to(device, 'x', 'edge_index')
 
-test_loader = NeighborSampler(data.edge_index, node_idx=None,
-                              sizes=[10, 10, 25], batch_size=256,
-                              shuffle=False, num_workers=12)
+train_loader = NeighborLoader(data, num_neighbors=[10, 10, 25], batch_size=256,
+                              shuffle=True, num_workers=12)
+test_loader = NeighborLoader(data, num_neighbors=[10, 10, 25], batch_size=256,
+                             num_workers=12)
 
 
 class Encoder(nn.Module):
@@ -37,19 +35,17 @@ class Encoder(nn.Module):
             nn.PReLU(hidden_channels)
         ])
 
-    def forward(self, x, adjs):
-        for i, (edge_index, _, size) in enumerate(adjs):
-            x_target = x[:size[1]]  # Target nodes are always placed first.
-            x = self.convs[i]((x, x_target), edge_index)
-            x = self.activations[i](x)
-        return x
+    def forward(self, x, edge_index, batch_size):
+        for conv, act in zip(self.convs, self.activations):
+            x = conv(x, edge_index)
+            x = act(x)
+        return x[:batch_size]
 
 
-def corruption(x, edge_index):
-    return x[torch.randperm(x.size(0))], edge_index
+def corruption(x, edge_index, batch_size):
+    return x[torch.randperm(x.size(0))], edge_index, batch_size
 
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = DeepGraphInfomax(
     hidden_channels=512, encoder=Encoder(dataset.num_features, 512),
     summary=lambda z, *args, **kwargs: torch.sigmoid(z.mean(dim=0)),
@@ -58,20 +54,15 @@ model = DeepGraphInfomax(
 model = model.to(device)
 optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
 
-x, y = data.x.to(device), data.y.to(device)
-
 
 def train(epoch):
     model.train()
 
     total_loss = total_examples = 0
-    for batch_size, n_id, adjs in tqdm(train_loader,
-                                       desc=f'Epoch {epoch:02d}'):
-        # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
-        adjs = [adj.to(device) for adj in adjs]
-
+    for batch in tqdm(train_loader, desc=f'Epoch {epoch:02d}'):
         optimizer.zero_grad()
-        pos_z, neg_z, summary = model(x[n_id], adjs)
+        pos_z, neg_z, summary = model(batch.x, batch.edge_index,
+                                      batch.batch_size)
         loss = model.loss(pos_z, neg_z, summary)
         loss.backward()
         optimizer.step()
@@ -86,13 +77,13 @@ def test():
     model.eval()
 
     zs = []
-    for i, (batch_size, n_id, adjs) in enumerate(test_loader):
-        adjs = [adj.to(device) for adj in adjs]
-        zs.append(model(x[n_id], adjs)[0])
+    for batch in tqdm(test_loader, desc='Evaluating'):
+        pos_z, _, _ = model(batch.x, batch.edge_index, batch.batch_size)
+        zs.append(pos_z.cpu())
     z = torch.cat(zs, dim=0)
     train_val_mask = data.train_mask | data.val_mask
-    acc = model.test(z[train_val_mask], y[train_val_mask], z[data.test_mask],
-                     y[data.test_mask], max_iter=10000)
+    acc = model.test(z[train_val_mask], data.y[train_val_mask],
+                     z[data.test_mask], data.y[data.test_mask], max_iter=10000)
     return acc
 
 
diff --git a/examples/mnist_graclus.py b/examples/mnist_graclus.py
index 2c1a85677..17e026f97 100644
--- a/examples/mnist_graclus.py
+++ b/examples/mnist_graclus.py
@@ -13,8 +13,14 @@ from torch_geometric.nn import (
     max_pool,
     max_pool_x,
 )
+from torch_geometric.typing import WITH_TORCH_CLUSTER, WITH_TORCH_SPLINE_CONV
 from torch_geometric.utils import normalized_cut
 
+if not WITH_TORCH_CLUSTER:
+    quit("This example requires 'torch-cluster'")
+if not WITH_TORCH_SPLINE_CONV:
+    quit("This example requires 'torch-spline-conv'")
+
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST')
 transform = T.Cartesian(cat=False)
 train_dataset = MNISTSuperpixels(path, True, transform=transform)
diff --git a/examples/mnist_voxel_grid.py b/examples/mnist_voxel_grid.py
index dc2a91959..62bad025c 100644
--- a/examples/mnist_voxel_grid.py
+++ b/examples/mnist_voxel_grid.py
@@ -7,6 +7,10 @@ import torch_geometric.transforms as T
 from torch_geometric.datasets import MNISTSuperpixels
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import SplineConv, max_pool, max_pool_x, voxel_grid
+from torch_geometric.typing import WITH_TORCH_SPLINE_CONV
+
+if not WITH_TORCH_SPLINE_CONV:
+    quit("This example requires 'torch-spline-conv'")
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'MNIST')
 transform = T.Cartesian(cat=False)
diff --git a/examples/multi_gpu/data_parallel.py b/examples/multi_gpu/data_parallel.py
index 42eb40fea..e78fbafd1 100644
--- a/examples/multi_gpu/data_parallel.py
+++ b/examples/multi_gpu/data_parallel.py
@@ -2,11 +2,18 @@ import os.path as osp
 
 import torch
 import torch.nn.functional as F
+from torch.nn import Linear, ReLU, Sequential
 
 import torch_geometric.transforms as T
 from torch_geometric.datasets import MNISTSuperpixels
 from torch_geometric.loader import DataListLoader
-from torch_geometric.nn import DataParallel, SplineConv, global_mean_pool
+from torch_geometric.nn import (
+    DataParallel,
+    NNConv,
+    SplineConv,
+    global_mean_pool,
+)
+from torch_geometric.typing import WITH_TORCH_SPLINE_CONV
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '../../data', 'MNIST')
 dataset = MNISTSuperpixels(path, transform=T.Cartesian()).shuffle()
@@ -16,8 +23,18 @@ loader = DataListLoader(dataset, batch_size=1024, shuffle=True)
 class Net(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        self.conv1 = SplineConv(dataset.num_features, 32, dim=2, kernel_size=5)
-        self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
+        if WITH_TORCH_SPLINE_CONV:
+            self.conv1 = SplineConv(dataset.num_features, 32, dim=2,
+                                    kernel_size=5)
+            self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
+        else:
+            nn1 = Sequential(Linear(2, 25), ReLU(),
+                             Linear(25, dataset.num_features * 32))
+            self.conv1 = NNConv(dataset.num_features, 32, nn1, aggr='mean')
+
+            nn2 = Sequential(Linear(2, 25), ReLU(), Linear(25, 32 * 64))
+            self.conv2 = NNConv(32, 64, nn2, aggr='mean')
+
         self.lin1 = torch.nn.Linear(64, 128)
         self.lin2 = torch.nn.Linear(128, dataset.num_classes)
 
diff --git a/examples/node2vec.py b/examples/node2vec.py
index e54dabc88..7c2d4ee35 100644
--- a/examples/node2vec.py
+++ b/examples/node2vec.py
@@ -7,10 +7,6 @@ from sklearn.manifold import TSNE
 
 from torch_geometric.datasets import Planetoid
 from torch_geometric.nn import Node2Vec
-from torch_geometric.typing import WITH_TORCH_CLUSTER
-
-if not WITH_TORCH_CLUSTER:
-    quit("This example requires 'torch-cluster'")
 
 
 def main():
@@ -20,9 +16,17 @@ def main():
     data = dataset[0]
 
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    model = Node2Vec(data.edge_index, embedding_dim=128, walk_length=20,
-                     context_size=10, walks_per_node=10,
-                     num_negative_samples=1, p=1, q=1, sparse=True).to(device)
+    model = Node2Vec(
+        data.edge_index,
+        embedding_dim=128,
+        walk_length=20,
+        context_size=10,
+        walks_per_node=10,
+        num_negative_samples=1,
+        p=1,
+        q=1,
+        sparse=True,
+    ).to(device)
 
     num_workers = 0 if sys.platform.startswith('win') else 4
     loader = model.loader(batch_size=128, shuffle=True,
@@ -52,7 +56,7 @@ def main():
     for epoch in range(1, 101):
         loss = train()
         acc = test()
-        print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Acc: {acc:.4f}')
+        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Acc: {acc:.4f}')
 
     @torch.no_grad()
     def plot_points(colors):
diff --git a/examples/ogbn_products_gat.py b/examples/ogbn_products_gat.py
index 4688e8b61..93a596178 100644
--- a/examples/ogbn_products_gat.py
+++ b/examples/ogbn_products_gat.py
@@ -8,22 +8,33 @@ from ogb.nodeproppred import Evaluator, PygNodePropPredDataset
 from torch.nn import Linear as Lin
 from tqdm import tqdm
 
-from torch_geometric.loader import NeighborSampler
+from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import GATConv
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'products')
 dataset = PygNodePropPredDataset('ogbn-products', root)
 split_idx = dataset.get_idx_split()
 evaluator = Evaluator(name='ogbn-products')
-data = dataset[0]
-
-train_idx = split_idx['train']
-train_loader = NeighborSampler(data.edge_index, node_idx=train_idx,
-                               sizes=[10, 10, 10], batch_size=512,
-                               shuffle=True, num_workers=12)
-subgraph_loader = NeighborSampler(data.edge_index, node_idx=None, sizes=[-1],
-                                  batch_size=1024, shuffle=False,
-                                  num_workers=12)
+data = dataset[0].to(device, 'x', 'y')
+
+train_loader = NeighborLoader(
+    data,
+    input_nodes=split_idx['train'],
+    num_neighbors=[10, 10, 5],
+    batch_size=512,
+    shuffle=True,
+    num_workers=12,
+    persistent_workers=True,
+)
+subgraph_loader = NeighborLoader(
+    data,
+    input_nodes=None,
+    num_neighbors=[-1],
+    batch_size=2048,
+    num_workers=12,
+    persistent_workers=True,
+)
 
 
 class GAT(torch.nn.Module):
@@ -56,21 +67,13 @@ class GAT(torch.nn.Module):
         for skip in self.skips:
             skip.reset_parameters()
 
-    def forward(self, x, adjs):
-        # `train_loader` computes the k-hop neighborhood of a batch of nodes,
-        # and returns, for each layer, a bipartite graph object, holding the
-        # bipartite edges `edge_index`, the index `e_id` of the original edges,
-        # and the size/shape `size` of the bipartite graph.
-        # Target nodes are also included in the source nodes so that one can
-        # easily apply skip-connections or add self-loops.
-        for i, (edge_index, _, size) in enumerate(adjs):
-            x_target = x[:size[1]]  # Target nodes are always placed first.
-            x = self.convs[i]((x, x_target), edge_index)
-            x = x + self.skips[i](x_target)
+    def forward(self, x, edge_index):
+        for i, (conv, skip) in enumerate(zip(self.convs, self.skips)):
+            x = conv(x, edge_index) + skip(x)
             if i != self.num_layers - 1:
                 x = F.elu(x)
                 x = F.dropout(x, p=0.5, training=self.training)
-        return x.log_softmax(dim=-1)
+        return x
 
     def inference(self, x_all):
         pbar = tqdm(total=x_all.size(0) * self.num_layers)
@@ -79,22 +82,18 @@ class GAT(torch.nn.Module):
         # Compute representations of nodes layer by layer, using *all*
         # available edges. This leads to faster computation in contrast to
         # immediately computing the final representations of each batch.
-        total_edges = 0
         for i in range(self.num_layers):
             xs = []
-            for batch_size, n_id, adj in subgraph_loader:
-                edge_index, _, size = adj.to(device)
-                total_edges += edge_index.size(1)
-                x = x_all[n_id].to(device)
-                x_target = x[:size[1]]
-                x = self.convs[i]((x, x_target), edge_index)
-                x = x + self.skips[i](x_target)
-
+            for batch in subgraph_loader:
+                x = x_all[batch.n_id].to(device)
+                edge_index = batch.edge_index.to(device)
+                x = self.convs[i](x, edge_index) + self.skips[i](x)
+                x = x[:batch.batch_size]
                 if i != self.num_layers - 1:
                     x = F.elu(x)
                 xs.append(x.cpu())
 
-                pbar.update(batch_size)
+                pbar.update(batch.batch_size)
 
             x_all = torch.cat(xs, dim=0)
 
@@ -103,40 +102,33 @@ class GAT(torch.nn.Module):
         return x_all
 
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = GAT(dataset.num_features, 128, dataset.num_classes, num_layers=3,
-            heads=4)
-model = model.to(device)
-
-x = data.x.to(device)
-y = data.y.squeeze().to(device)
+            heads=4).to(device)
 
 
 def train(epoch):
     model.train()
 
-    pbar = tqdm(total=train_idx.size(0))
+    pbar = tqdm(total=split_idx['train'].size(0))
     pbar.set_description(f'Epoch {epoch:02d}')
 
     total_loss = total_correct = 0
-    for batch_size, n_id, adjs in train_loader:
-        # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
-        adjs = [adj.to(device) for adj in adjs]
-
+    for batch in train_loader:
         optimizer.zero_grad()
-        out = model(x[n_id], adjs)
-        loss = F.nll_loss(out, y[n_id[:batch_size]])
+        out = model(batch.x, batch.edge_index.to(device))[:batch.batch_size]
+        y = batch.y[:batch.batch_size].squeeze()
+        loss = F.cross_entropy(out, y)
         loss.backward()
         optimizer.step()
 
         total_loss += float(loss)
-        total_correct += int(out.argmax(dim=-1).eq(y[n_id[:batch_size]]).sum())
-        pbar.update(batch_size)
+        total_correct += int(out.argmax(dim=-1).eq(y).sum())
+        pbar.update(batch.batch_size)
 
     pbar.close()
 
     loss = total_loss / len(train_loader)
-    approx_acc = total_correct / train_idx.size(0)
+    approx_acc = total_correct / split_idx['train'].size(0)
 
     return loss, approx_acc
 
@@ -145,9 +137,9 @@ def train(epoch):
 def test():
     model.eval()
 
-    out = model.inference(x)
+    out = model.inference(data.x)
 
-    y_true = y.cpu().unsqueeze(-1)
+    y_true = data.y.cpu()
     y_pred = out.argmax(dim=-1, keepdim=True)
 
     train_acc = evaluator.eval({
@@ -168,9 +160,7 @@ def test():
 
 test_accs = []
 for run in range(1, 11):
-    print('')
-    print(f'Run {run:02d}:')
-    print('')
+    print(f'\nRun {run:02d}:\n')
 
     model.reset_parameters()
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
diff --git a/examples/ogbn_products_sage.py b/examples/ogbn_products_sage.py
index 0ff2406d6..7dbb96b48 100644
--- a/examples/ogbn_products_sage.py
+++ b/examples/ogbn_products_sage.py
@@ -7,22 +7,33 @@ import torch.nn.functional as F
 from ogb.nodeproppred import Evaluator, PygNodePropPredDataset
 from tqdm import tqdm
 
-from torch_geometric.loader import NeighborSampler
+from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import SAGEConv
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'products')
 dataset = PygNodePropPredDataset('ogbn-products', root)
 split_idx = dataset.get_idx_split()
 evaluator = Evaluator(name='ogbn-products')
-data = dataset[0]
-
-train_idx = split_idx['train']
-train_loader = NeighborSampler(data.edge_index, node_idx=train_idx,
-                               sizes=[15, 10, 5], batch_size=1024,
-                               shuffle=True, num_workers=12)
-subgraph_loader = NeighborSampler(data.edge_index, node_idx=None, sizes=[-1],
-                                  batch_size=4096, shuffle=False,
-                                  num_workers=12)
+data = dataset[0].to(device, 'x', 'y')
+
+train_loader = NeighborLoader(
+    data,
+    input_nodes=split_idx['train'],
+    num_neighbors=[15, 10, 5],
+    batch_size=1024,
+    shuffle=True,
+    num_workers=12,
+    persistent_workers=True,
+)
+subgraph_loader = NeighborLoader(
+    data,
+    input_nodes=None,
+    num_neighbors=[-1],
+    batch_size=4096,
+    num_workers=12,
+    persistent_workers=True,
+)
 
 
 class SAGE(torch.nn.Module):
@@ -41,20 +52,13 @@ class SAGE(torch.nn.Module):
         for conv in self.convs:
             conv.reset_parameters()
 
-    def forward(self, x, adjs):
-        # `train_loader` computes the k-hop neighborhood of a batch of nodes,
-        # and returns, for each layer, a bipartite graph object, holding the
-        # bipartite edges `edge_index`, the index `e_id` of the original edges,
-        # and the size/shape `size` of the bipartite graph.
-        # Target nodes are also included in the source nodes so that one can
-        # easily apply skip-connections or add self-loops.
-        for i, (edge_index, _, size) in enumerate(adjs):
-            x_target = x[:size[1]]  # Target nodes are always placed first.
-            x = self.convs[i]((x, x_target), edge_index)
+    def forward(self, x, edge_index):
+        for i, conv in enumerate(self.convs):
+            x = conv(x, edge_index)
             if i != self.num_layers - 1:
-                x = F.relu(x)
+                x = x.relu()
                 x = F.dropout(x, p=0.5, training=self.training)
-        return x.log_softmax(dim=-1)
+        return x
 
     def inference(self, x_all):
         pbar = tqdm(total=x_all.size(0) * self.num_layers)
@@ -63,20 +67,18 @@ class SAGE(torch.nn.Module):
         # Compute representations of nodes layer by layer, using *all*
         # available edges. This leads to faster computation in contrast to
         # immediately computing the final representations of each batch.
-        total_edges = 0
         for i in range(self.num_layers):
             xs = []
-            for batch_size, n_id, adj in subgraph_loader:
-                edge_index, _, size = adj.to(device)
-                total_edges += edge_index.size(1)
-                x = x_all[n_id].to(device)
-                x_target = x[:size[1]]
-                x = self.convs[i]((x, x_target), edge_index)
+            for batch in subgraph_loader:
+                x = x_all[batch.n_id].to(device)
+                edge_index = batch.edge_index.to(device)
+                x = self.convs[i](x, edge_index)
+                x = x[:batch.batch_size]
                 if i != self.num_layers - 1:
-                    x = F.relu(x)
+                    x = x.relu()
                 xs.append(x.cpu())
 
-                pbar.update(batch_size)
+                pbar.update(batch.batch_size)
 
             x_all = torch.cat(xs, dim=0)
 
@@ -85,39 +87,33 @@ class SAGE(torch.nn.Module):
         return x_all
 
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = SAGE(dataset.num_features, 256, dataset.num_classes, num_layers=3)
 model = model.to(device)
 
-x = data.x.to(device)
-y = data.y.squeeze().to(device)
-
 
 def train(epoch):
     model.train()
 
-    pbar = tqdm(total=train_idx.size(0))
+    pbar = tqdm(total=split_idx['train'].size(0))
     pbar.set_description(f'Epoch {epoch:02d}')
 
     total_loss = total_correct = 0
-    for batch_size, n_id, adjs in train_loader:
-        # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
-        adjs = [adj.to(device) for adj in adjs]
-
+    for batch in train_loader:
         optimizer.zero_grad()
-        out = model(x[n_id], adjs)
-        loss = F.nll_loss(out, y[n_id[:batch_size]])
+        out = model(batch.x, batch.edge_index.to(device))[:batch.batch_size]
+        y = batch.y[:batch.batch_size].squeeze()
+        loss = F.cross_entropy(out, y)
         loss.backward()
         optimizer.step()
 
         total_loss += float(loss)
-        total_correct += int(out.argmax(dim=-1).eq(y[n_id[:batch_size]]).sum())
-        pbar.update(batch_size)
+        total_correct += int(out.argmax(dim=-1).eq(y).sum())
+        pbar.update(batch.batch_size)
 
     pbar.close()
 
     loss = total_loss / len(train_loader)
-    approx_acc = total_correct / train_idx.size(0)
+    approx_acc = total_correct / split_idx['train'].size(0)
 
     return loss, approx_acc
 
@@ -126,9 +122,9 @@ def train(epoch):
 def test():
     model.eval()
 
-    out = model.inference(x)
+    out = model.inference(data.x)
 
-    y_true = y.cpu().unsqueeze(-1)
+    y_true = data.y.cpu()
     y_pred = out.argmax(dim=-1, keepdim=True)
 
     train_acc = evaluator.eval({
@@ -149,9 +145,7 @@ def test():
 
 test_accs = []
 for run in range(1, 11):
-    print('')
-    print(f'Run {run:02d}:')
-    print('')
+    print(f'\nRun {run:02d}:\n')
 
     model.reset_parameters()
     optimizer = torch.optim.Adam(model.parameters(), lr=0.003)
diff --git a/examples/point_transformer_classification.py b/examples/point_transformer_classification.py
index 8556249c4..9a22e6a2c 100644
--- a/examples/point_transformer_classification.py
+++ b/examples/point_transformer_classification.py
@@ -3,14 +3,24 @@ import os.path as osp
 import torch
 import torch.nn.functional as F
 from torch.nn import Linear as Lin
-from torch_cluster import fps, knn_graph
 
 import torch_geometric.transforms as T
 from torch_geometric.datasets import ModelNet
 from torch_geometric.loader import DataLoader
-from torch_geometric.nn import MLP, PointTransformerConv, global_mean_pool, knn
+from torch_geometric.nn import (
+    MLP,
+    PointTransformerConv,
+    fps,
+    global_mean_pool,
+    knn,
+    knn_graph,
+)
+from torch_geometric.typing import WITH_TORCH_CLUSTER
 from torch_geometric.utils import scatter
 
+if not WITH_TORCH_CLUSTER:
+    quit("This example requires 'torch-cluster'")
+
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data/ModelNet10')
 pre_transform, transform = T.NormalizeScale(), T.SamplePoints(1024)
 train_dataset = ModelNet(path, '10', True, transform, pre_transform)
diff --git a/examples/point_transformer_segmentation.py b/examples/point_transformer_segmentation.py
index 1d0026863..feeb247fe 100644
--- a/examples/point_transformer_segmentation.py
+++ b/examples/point_transformer_segmentation.py
@@ -3,15 +3,18 @@ import os.path as osp
 import torch
 import torch.nn.functional as F
 from point_transformer_classification import TransformerBlock, TransitionDown
-from torch_cluster import knn_graph
 from torchmetrics.functional import jaccard_index
 
 import torch_geometric.transforms as T
 from torch_geometric.datasets import ShapeNet
 from torch_geometric.loader import DataLoader
-from torch_geometric.nn import MLP, knn_interpolate
+from torch_geometric.nn import MLP, knn_graph, knn_interpolate
+from torch_geometric.typing import WITH_TORCH_CLUSTER
 from torch_geometric.utils import scatter
 
+if not WITH_TORCH_CLUSTER:
+    quit("This example requires 'torch-cluster'")
+
 category = 'Airplane'  # Pass in `None` to train on all categories.
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ShapeNet')
 transform = T.Compose([
diff --git a/examples/pointnet2_classification.py b/examples/pointnet2_classification.py
index 4bae3b395..fa1975491 100644
--- a/examples/pointnet2_classification.py
+++ b/examples/pointnet2_classification.py
@@ -7,6 +7,10 @@ import torch_geometric.transforms as T
 from torch_geometric.datasets import ModelNet
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import MLP, PointNetConv, fps, global_max_pool, radius
+from torch_geometric.typing import WITH_TORCH_CLUSTER
+
+if not WITH_TORCH_CLUSTER:
+    quit("This example requires 'torch-cluster'")
 
 
 class SAModule(torch.nn.Module):
diff --git a/examples/pointnet2_segmentation.py b/examples/pointnet2_segmentation.py
index 1c575b403..5ced25b06 100644
--- a/examples/pointnet2_segmentation.py
+++ b/examples/pointnet2_segmentation.py
@@ -9,8 +9,12 @@ import torch_geometric.transforms as T
 from torch_geometric.datasets import ShapeNet
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import MLP, knn_interpolate
+from torch_geometric.typing import WITH_TORCH_CLUSTER
 from torch_geometric.utils import scatter
 
+if not WITH_TORCH_CLUSTER:
+    quit("This example requires 'torch-cluster'")
+
 category = 'Airplane'  # Pass in `None` to train on all categories.
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ShapeNet')
 transform = T.Compose([
diff --git a/examples/proteins_diff_pool.py b/examples/proteins_diff_pool.py
index 99d0f33e1..e021b1d36 100644
--- a/examples/proteins_diff_pool.py
+++ b/examples/proteins_diff_pool.py
@@ -12,7 +12,7 @@ from torch_geometric.nn import DenseSAGEConv, dense_diff_pool
 max_nodes = 150
 
 
-class MyFilter(object):
+class MyFilter:
     def __call__(self, data):
         return data.num_nodes <= max_nodes
 
diff --git a/examples/qm9_nn_conv.py b/examples/qm9_nn_conv.py
index 72f971c59..55998d8bd 100644
--- a/examples/qm9_nn_conv.py
+++ b/examples/qm9_nn_conv.py
@@ -14,14 +14,14 @@ target = 0
 dim = 64
 
 
-class MyTransform(object):
+class MyTransform:
     def __call__(self, data):
         # Specify target.
         data.y = data.y[:, target]
         return data
 
 
-class Complete(object):
+class Complete:
     def __call__(self, data):
         device = data.edge_index.device
 
diff --git a/examples/qm9_pretrained_schnet.py b/examples/qm9_pretrained_schnet.py
index da91f1db0..61d6c31bf 100644
--- a/examples/qm9_pretrained_schnet.py
+++ b/examples/qm9_pretrained_schnet.py
@@ -14,7 +14,7 @@ parser.add_argument('--cutoff', type=float, default=10.0,
 args = parser.parse_args()
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'QM9')
-dataset = QM9(osp.join())
+dataset = QM9(path)
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
diff --git a/examples/randlanet_classification.py b/examples/randlanet_classification.py
index 9d9031e05..f65407059 100644
--- a/examples/randlanet_classification.py
+++ b/examples/randlanet_classification.py
@@ -20,8 +20,12 @@ from torch_geometric.nn.aggr import MaxAggregation
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.pool import knn_graph
 from torch_geometric.nn.pool.decimation import decimation_indices
+from torch_geometric.typing import WITH_TORCH_CLUSTER
 from torch_geometric.utils import softmax
 
+if not WITH_TORCH_CLUSTER:
+    quit("This example requires 'torch-cluster'")
+
 # Default activation and batch norm parameters used by RandLA-Net:
 lrelu02_kwargs = {'negative_slope': 0.2}
 bn099_kwargs = {'momentum': 0.01, 'eps': 1e-6}
diff --git a/examples/randlanet_segmentation.py b/examples/randlanet_segmentation.py
index cc5af0733..39602da37 100644
--- a/examples/randlanet_segmentation.py
+++ b/examples/randlanet_segmentation.py
@@ -16,8 +16,12 @@ import torch_geometric.transforms as T
 from torch_geometric.datasets import ShapeNet
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import knn_interpolate
+from torch_geometric.typing import WITH_TORCH_CLUSTER
 from torch_geometric.utils import scatter
 
+if not WITH_TORCH_CLUSTER:
+    quit("This example requires 'torch-cluster'")
+
 category = 'Airplane'  # Pass in `None` to train on all categories.
 category_num_classes = 4  # 4 for Airplane - see ShapeNet for details
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ShapeNet')
diff --git a/examples/triangles_sag_pool.py b/examples/triangles_sag_pool.py
index 4ed7e181c..fc70a2196 100644
--- a/examples/triangles_sag_pool.py
+++ b/examples/triangles_sag_pool.py
@@ -13,7 +13,7 @@ from torch_geometric.nn import GCNConv, GINConv, SAGPooling, global_max_pool
 from torch_geometric.utils import scatter
 
 
-class HandleNodeAttention(object):
+class HandleNodeAttention:
     def __call__(self, data):
         data.attn = torch.softmax(data.x, dim=0).flatten()
         data.x = None
diff --git a/readthedocs.yml b/readthedocs.yml
index aae9ffd43..35620da86 100644
--- a/readthedocs.yml
+++ b/readthedocs.yml
@@ -1,11 +1,11 @@
 version: 2
 
 build:
-   image: latest
+   os: ubuntu-22.04
+   tools:
+      python: "3.8"
 
 python:
-   version: 3.8
-   system_packages: true
    install:
       - requirements: docs/requirements.txt
       - method: pip
diff --git a/torch_geometric/data/batch.py b/torch_geometric/data/batch.py
index b24ff9b32..cf5e1f6c0 100644
--- a/torch_geometric/data/batch.py
+++ b/torch_geometric/data/batch.py
@@ -48,7 +48,7 @@ class DynamicInheritance(type):
         return super(DynamicInheritance, new_cls).__call__(*args, **kwargs)
 
 
-class DynamicInheritanceGetter(object):
+class DynamicInheritanceGetter:
     def __call__(self, cls, base_cls):
         return cls(_base_cls=base_cls)
 
diff --git a/torch_geometric/data/collate.py b/torch_geometric/data/collate.py
index 3d76bc194..a9896916d 100644
--- a/torch_geometric/data/collate.py
+++ b/torch_geometric/data/collate.py
@@ -9,6 +9,8 @@ import torch_geometric.typing
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage, NodeStorage
 from torch_geometric.typing import SparseTensor, torch_sparse
+from torch_geometric.utils import is_sparse, is_torch_sparse_tensor
+from torch_geometric.utils.sparse import cat
 
 
 def collate(
@@ -122,7 +124,7 @@ def _collate(
 
     elem = values[0]
 
-    if isinstance(elem, Tensor):
+    if isinstance(elem, Tensor) and not is_sparse(elem):
         # Concatenate a list of `torch.Tensor` along the `cat_dim`.
         # NOTE: We need to take care of incrementing elements appropriately.
         key = str(key)
@@ -160,7 +162,7 @@ def _collate(
         value = torch.cat(values, dim=cat_dim or 0, out=out)
         return value, slices, incs
 
-    elif isinstance(elem, SparseTensor) and increment:
+    elif is_sparse(elem) and increment:
         # Concatenate a list of `SparseTensor` along the `cat_dim`.
         # NOTE: `cat_dim` may return a tuple to allow for diagonal stacking.
         key = str(key)
@@ -168,7 +170,10 @@ def _collate(
         cat_dims = (cat_dim, ) if isinstance(cat_dim, int) else cat_dim
         repeats = [[value.size(dim) for dim in cat_dims] for value in values]
         slices = cumsum(repeats)
-        value = torch_sparse.cat(values, dim=cat_dim)
+        if is_torch_sparse_tensor(elem):
+            value = cat(values, dim=cat_dim)
+        else:
+            value = torch_sparse.cat(values, dim=cat_dim)
         return value, slices, None
 
     elif isinstance(elem, (int, float)):
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 98ba68889..e620a5520 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -37,10 +37,10 @@ from torch_geometric.typing import (
     OptTensor,
     SparseTensor,
 )
-from torch_geometric.utils import select, subgraph
+from torch_geometric.utils import is_sparse, select, subgraph
 
 
-class BaseData(object):
+class BaseData:
     def __getattr__(self, key: str) -> Any:
         raise NotImplementedError
 
@@ -518,7 +518,7 @@ class Data(BaseData, FeatureStore, GraphStore):
         return self
 
     def __cat_dim__(self, key: str, value: Any, *args, **kwargs) -> Any:
-        if isinstance(value, SparseTensor) and 'adj' in key:
+        if is_sparse(value) and 'adj' in key:
             return (0, 1)
         elif 'index' in key or key == 'face':
             return -1
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index 7a9144e51..abca8fd99 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -125,6 +125,23 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
             else:
                 setattr(self, key, value)
 
+    @classmethod
+    def from_dict(cls, mapping: Dict[str, Any]) -> 'HeteroData':
+        r"""Creates a :class:`~torch_geometric.data.HeteroData` object from a
+        Python dictionary."""
+        out = cls()
+        for key, value in mapping.items():
+            if key == '_global_store':
+                out.__dict__['_global_store'] = BaseStorage(
+                    _parent=out, **value)
+            elif isinstance(key, str):
+                out._node_store_dict[key] = NodeStorage(
+                    _parent=out, _key=key, **value)
+            else:
+                out._edge_store_dict[key] = EdgeStorage(
+                    _parent=out, _key=key, **value)
+        return out
+
     def __getattr__(self, key: str) -> Any:
         # `data.*_dict` => Link to node and edge stores.
         # `data.*` => Link to the `_global_store`.
@@ -258,11 +275,12 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         return list(self._edge_store_dict.items())
 
     def to_dict(self) -> Dict[str, Any]:
-        out = self._global_store.to_dict()
+        out_dict: Dict[str, Any] = {}
+        out_dict['_global_store'] = self._global_store.to_dict()
         for key, store in chain(self._node_store_dict.items(),
                                 self._edge_store_dict.items()):
-            out[key] = store.to_dict()
-        return out
+            out_dict[key] = store.to_dict()
+        return out_dict
 
     def to_namedtuple(self) -> NamedTuple:
         field_names = list(self._global_store.keys())
diff --git a/torch_geometric/data/in_memory_dataset.py b/torch_geometric/data/in_memory_dataset.py
index 1e1b3ee19..9d54cf1b7 100644
--- a/torch_geometric/data/in_memory_dataset.py
+++ b/torch_geometric/data/in_memory_dataset.py
@@ -2,12 +2,24 @@ import copy
 import warnings
 from abc import ABC
 from collections.abc import Mapping, Sequence
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
-
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
+
+import torch
 from torch import Tensor
 
 from torch_geometric.data import Batch, Data
 from torch_geometric.data.collate import collate
+from torch_geometric.data.data import BaseData
 from torch_geometric.data.dataset import Dataset, IndexType
 from torch_geometric.data.separate import separate
 
@@ -93,6 +105,19 @@ class InMemoryDataset(Dataset, ABC):
 
         return data
 
+    @classmethod
+    def save(cls, data_list: List[BaseData], path: str):
+        r"""Saves a list of data objects to the file path :obj:`path`."""
+        data, slices = cls.collate(data_list)
+        torch.save((data.to_dict(), slices), path)
+
+    def load(self, path: str, data_cls: Type[BaseData] = Data):
+        r"""Loads the dataset from the file path :obj:`path`."""
+        data, self.slices = torch.load(path)
+        if isinstance(data, dict):  # Backward compatibility.
+            data = data_cls.from_dict(data)
+        self.data = data
+
     @staticmethod
     def collate(
             data_list: List[Data]) -> Tuple[Data, Optional[Dict[str, Tensor]]]:
diff --git a/torch_geometric/data/separate.py b/torch_geometric/data/separate.py
index 5fcb98b92..c16e94cbe 100644
--- a/torch_geometric/data/separate.py
+++ b/torch_geometric/data/separate.py
@@ -65,7 +65,8 @@ def _separate(
         start, end = int(slices[idx]), int(slices[idx + 1])
         value = narrow(value, cat_dim or 0, start, end - start)
         value = value.squeeze(0) if cat_dim is None else value
-        if decrement and (incs.dim() > 1 or int(incs[idx]) != 0):
+        if (decrement and incs is not None
+                and (incs.dim() > 1 or int(incs[idx]) != 0)):
             value = value - incs[idx].to(value.device)
         return value
 
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index 41a68c699..f43cf1254 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -191,7 +191,13 @@ class BaseStorage(MutableMapping):
 
     def to_dict(self) -> Dict[str, Any]:
         r"""Returns a dictionary of stored key/value pairs."""
-        return copy.copy(self._mapping)
+        out_dict = copy.copy(self._mapping)
+        # Needed to preserve individual `num_nodes` attributes when calling
+        # `BaseData.collate`.
+        # TODO (matthias) Try to make this more generic.
+        if '_num_nodes' in self.__dict__:
+            out_dict['_num_nodes'] = self.__dict__['_num_nodes']
+        return out_dict
 
     def to_namedtuple(self) -> NamedTuple:
         r"""Returns a :obj:`NamedTuple` of stored key/value pairs."""
diff --git a/torch_geometric/data/summary.py b/torch_geometric/data/summary.py
index c8570a319..4e1aa0eba 100644
--- a/torch_geometric/data/summary.py
+++ b/torch_geometric/data/summary.py
@@ -1,10 +1,11 @@
+from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import torch
 from tqdm import tqdm
 
-from torch_geometric.data import Dataset
+from torch_geometric.data import Dataset, HeteroData
 
 
 @dataclass
@@ -40,12 +41,15 @@ class Summary:
     num_graphs: int
     num_nodes: Stats
     num_edges: Stats
+    num_nodes_per_type: Optional[Dict[str, Stats]] = None
+    num_edges_per_type: Optional[Dict[str, Stats]] = None
 
     @classmethod
     def from_dataset(
         cls,
         dataset: Dataset,
         progress_bar: Optional[bool] = None,
+        per_type: bool = True,
     ):
         r"""Creates a summary of a :class:`~torch_geometric.data.Dataset`
         object.
@@ -56,6 +60,9 @@ class Summary:
                 progress bar during stats computation. If set to :obj:`None`,
                 will automatically decide whether to show a progress bar based
                 on dataset size. (default: :obj:`None`)
+            per_type (bool, optional). If set to :obj:`True`, will separate
+                statistics per node and edge type (only applicable in
+                heterogeneous graph datasets). (default: :obj:`True`)
         """
         name = dataset.__class__.__name__
 
@@ -65,28 +72,86 @@ class Summary:
         if progress_bar:
             dataset = tqdm(dataset)
 
-        num_nodes_list, num_edges_list = [], []
+        num_nodes, num_edges = [], []
+        num_nodes_per_type = defaultdict(list)
+        num_edges_per_type = defaultdict(list)
+
         for data in dataset:
-            num_nodes_list.append(data.num_nodes)
-            num_edges_list.append(data.num_edges)
+            num_nodes.append(data.num_nodes)
+            num_edges.append(data.num_edges)
+
+            if per_type and isinstance(data, HeteroData):
+                for node_type in data.node_types:
+                    num_nodes_per_type[node_type].append(
+                        data[node_type].num_nodes)
+                for edge_type in data.edge_types:
+                    num_edges_per_type[edge_type].append(
+                        data[edge_type].num_edges)
+
+        if len(num_nodes_per_type) > 0:
+            num_nodes_per_type = {
+                node_type: Stats.from_data(num_nodes_list)
+                for node_type, num_nodes_list in num_nodes_per_type.items()
+            }
+        else:
+            num_nodes_per_type = None
+
+        if len(num_edges_per_type) > 0:
+            num_edges_per_type = {
+                edge_type: Stats.from_data(num_edges_list)
+                for edge_type, num_edges_list in num_edges_per_type.items()
+            }
+        else:
+            num_edges_per_type = None
 
         return cls(
             name=name,
             num_graphs=len(dataset),
-            num_nodes=Stats.from_data(num_nodes_list),
-            num_edges=Stats.from_data(num_edges_list),
+            num_nodes=Stats.from_data(num_nodes),
+            num_edges=Stats.from_data(num_edges),
+            num_nodes_per_type=num_nodes_per_type,
+            num_edges_per_type=num_edges_per_type,
         )
 
     def __repr__(self) -> str:
         from tabulate import tabulate
 
-        prefix = f'{self.name} (#graphs={self.num_graphs}):\n'
+        body = f'{self.name} (#graphs={self.num_graphs}):\n'
 
         content = [['', '#nodes', '#edges']]
         stats = [self.num_nodes, self.num_edges]
         for field in Stats.__dataclass_fields__:
             row = [field] + [f'{getattr(s, field):.1f}' for s in stats]
             content.append(row)
-        body = tabulate(content, headers='firstrow', tablefmt='psql')
-
-        return prefix + body
+        body += tabulate(content, headers='firstrow', tablefmt='psql')
+
+        if self.num_nodes_per_type is not None:
+            content = [['']]
+            content[0] += list(self.num_nodes_per_type.keys())
+
+            for field in Stats.__dataclass_fields__:
+                row = [field] + [
+                    f'{getattr(s, field):.1f}'
+                    for s in self.num_nodes_per_type.values()
+                ]
+                content.append(row)
+            body += "\nNumber of nodes per node type:\n"
+            body += tabulate(content, headers='firstrow', tablefmt='psql')
+
+        if self.num_edges_per_type is not None:
+            content = [['']]
+            content[0] += [
+                f"({', '.join(edge_type)})"
+                for edge_type in self.num_edges_per_type.keys()
+            ]
+
+            for field in Stats.__dataclass_fields__:
+                row = [field] + [
+                    f'{getattr(s, field):.1f}'
+                    for s in self.num_edges_per_type.values()
+                ]
+                content.append(row)
+            body += "\nNumber of edges per edge type:\n"
+            body += tabulate(content, headers='firstrow', tablefmt='psql')
+
+        return body
diff --git a/torch_geometric/data/view.py b/torch_geometric/data/view.py
index 75e1b2c02..d03b96723 100644
--- a/torch_geometric/data/view.py
+++ b/torch_geometric/data/view.py
@@ -2,7 +2,7 @@ from collections.abc import Mapping
 from typing import Iterable, List
 
 
-class MappingView(object):
+class MappingView:
     def __init__(self, mapping: Mapping, *args: List[str]):
         self._mapping = mapping
         self._args = args
diff --git a/torch_geometric/datasets/amazon.py b/torch_geometric/datasets/amazon.py
index 6b3ebc3a3..176c8d08a 100644
--- a/torch_geometric/datasets/amazon.py
+++ b/torch_geometric/datasets/amazon.py
@@ -54,9 +54,13 @@ class Amazon(InMemoryDataset):
 
     url = 'https://github.com/shchur/gnn-benchmark/raw/master/data/npz/'
 
-    def __init__(self, root: str, name: str,
-                 transform: Optional[Callable] = None,
-                 pre_transform: Optional[Callable] = None):
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+    ):
         self.name = name.lower()
         assert self.name in ['computers', 'photo']
         super().__init__(root, transform, pre_transform)
@@ -82,7 +86,7 @@ class Amazon(InMemoryDataset):
         download_url(self.url + self.raw_file_names, self.raw_dir)
 
     def process(self):
-        data = read_npz(self.raw_paths[0])
+        data = read_npz(self.raw_paths[0], to_undirected=True)
         data = data if self.pre_transform is None else self.pre_transform(data)
         data, slices = self.collate([data])
         torch.save((data, slices), self.processed_paths[0])
diff --git a/torch_geometric/datasets/citation_full.py b/torch_geometric/datasets/citation_full.py
index c3ec045e4..70c2d6f51 100644
--- a/torch_geometric/datasets/citation_full.py
+++ b/torch_geometric/datasets/citation_full.py
@@ -27,6 +27,8 @@ class CitationFull(InMemoryDataset):
             an :obj:`torch_geometric.data.Data` object and returns a
             transformed version. The data object will be transformed before
             being saved to disk. (default: :obj:`None`)
+        to_undirected (bool, optional): Whether the original graph is
+            converted to an undirected one. (default: :obj:`True`)
 
     **STATS:**
 
@@ -68,10 +70,16 @@ class CitationFull(InMemoryDataset):
 
     url = 'https://github.com/abojchevski/graph2gauss/raw/master/data/{}.npz'
 
-    def __init__(self, root: str, name: str,
-                 transform: Optional[Callable] = None,
-                 pre_transform: Optional[Callable] = None):
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        to_undirected: bool = True,
+    ):
         self.name = name.lower()
+        self.to_undirected = to_undirected
         assert self.name in ['cora', 'cora_ml', 'citeseer', 'dblp', 'pubmed']
         super().__init__(root, transform, pre_transform)
         self.data, self.slices = torch.load(self.processed_paths[0])
@@ -96,7 +104,7 @@ class CitationFull(InMemoryDataset):
         download_url(self.url.format(self.name), self.raw_dir)
 
     def process(self):
-        data = read_npz(self.raw_paths[0])
+        data = read_npz(self.raw_paths[0], to_undirected=self.to_undirected)
         data = data if self.pre_transform is None else self.pre_transform(data)
         data, slices = self.collate([data])
         torch.save((data, slices), self.processed_paths[0])
diff --git a/torch_geometric/datasets/coauthor.py b/torch_geometric/datasets/coauthor.py
index e6f015b60..39db63e10 100644
--- a/torch_geometric/datasets/coauthor.py
+++ b/torch_geometric/datasets/coauthor.py
@@ -53,9 +53,13 @@ class Coauthor(InMemoryDataset):
 
     url = 'https://github.com/shchur/gnn-benchmark/raw/master/data/npz/'
 
-    def __init__(self, root: str, name: str,
-                 transform: Optional[Callable] = None,
-                 pre_transform: Optional[Callable] = None):
+    def __init__(
+        self,
+        root: str,
+        name: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+    ):
         assert name.lower() in ['cs', 'physics']
         self.name = 'CS' if name.lower() == 'cs' else 'Physics'
         super().__init__(root, transform, pre_transform)
@@ -81,7 +85,7 @@ class Coauthor(InMemoryDataset):
         download_url(self.url + self.raw_file_names, self.raw_dir)
 
     def process(self):
-        data = read_npz(self.raw_paths[0])
+        data = read_npz(self.raw_paths[0], to_undirected=True)
         data = data if self.pre_transform is None else self.pre_transform(data)
         data, slices = self.collate([data])
         torch.save((data, slices), self.processed_paths[0])
diff --git a/torch_geometric/datasets/entities.py b/torch_geometric/datasets/entities.py
index 52fcc0981..210d27864 100644
--- a/torch_geometric/datasets/entities.py
+++ b/torch_geometric/datasets/entities.py
@@ -204,7 +204,7 @@ class Entities(InMemoryDataset):
         return f'{self.name.upper()}{self.__class__.__name__}()'
 
 
-class hide_stdout(object):
+class hide_stdout:
     def __enter__(self):
         self.level = logging.getLogger().level
         logging.getLogger().setLevel(logging.ERROR)
diff --git a/torch_geometric/explain/explanation.py b/torch_geometric/explain/explanation.py
index 71fe5073b..5c2ec4118 100644
--- a/torch_geometric/explain/explanation.py
+++ b/torch_geometric/explain/explanation.py
@@ -200,21 +200,18 @@ class Explanation(Data, ExplanationMixin):
         feat_labels: Optional[List[str]] = None,
         top_k: Optional[int] = None,
     ):
-        r"""Creates a bar plot of the node features importance by summing up
-        :attr:`self.node_mask` across all nodes.
+        r"""Creates a bar plot of the node feature importances by summing up
+        the node mask across all nodes.
 
         Args:
             path (str, optional): The path to where the plot is saved.
                 If set to :obj:`None`, will visualize the plot on-the-fly.
                 (default: :obj:`None`)
-            feat_labels (List[str], optional): Optional labels for features.
+            feat_labels (List[str], optional): The labels of features.
                 (default :obj:`None`)
             top_k (int, optional): Top k features to plot. If :obj:`None`
                 plots all features. (default: :obj:`None`)
         """
-        import matplotlib.pyplot as plt
-        import pandas as pd
-
         node_mask = self.get('node_mask')
         if node_mask is None:
             raise ValueError(f"The attribute 'node_mask' is not available "
@@ -225,44 +222,12 @@ class Explanation(Data, ExplanationMixin):
                              f"object-level 'node_mask' "
                              f"(got shape {node_mask.size()})")
 
-        feat_importance = node_mask.sum(dim=0).cpu().numpy()
-
         if feat_labels is None:
-            feat_labels = range(feat_importance.shape[0])
-
-        if len(feat_labels) != feat_importance.shape[0]:
-            raise ValueError(f"The '{self.__class__.__name__}' object holds "
-                             f"{feat_importance.numel()} features, but "
-                             f"only {len(feat_labels)} were passed")
-
-        df = pd.DataFrame({'feat_importance': feat_importance},
-                          index=feat_labels)
-        df = df.sort_values("feat_importance", ascending=False)
-        df = df.round(decimals=3)
-
-        if top_k is not None:
-            df = df.head(top_k)
-            title = f"Feature importance for top {len(df)} features"
-        else:
-            title = f"Feature importance for {len(df)} features"
-
-        ax = df.plot(
-            kind='barh',
-            figsize=(10, 7),
-            title=title,
-            ylabel='Feature label',
-            xlim=[0, float(feat_importance.max()) + 0.3],
-            legend=False,
-        )
-        plt.gca().invert_yaxis()
-        ax.bar_label(container=ax.containers[0], label_type='edge')
+            feat_labels = range(node_mask.size(1))
 
-        if path is not None:
-            plt.savefig(path)
-        else:
-            plt.show()
+        score = node_mask.sum(dim=0)
 
-        plt.close()
+        return _visualize_score(score, feat_labels, path, top_k)
 
     def visualize_graph(self, path: Optional[str] = None,
                         backend: Optional[str] = None):
@@ -343,3 +308,93 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
                     out[edge_type][key] = value[edge_mask]
 
         return out.subgraph(node_mask_dict)
+
+    def visualize_feature_importance(
+        self,
+        path: Optional[str] = None,
+        feat_labels: Optional[Dict[NodeType, List[str]]] = None,
+        top_k: Optional[int] = None,
+    ):
+        r"""Creates a bar plot of the node feature importances by summing up
+        node masks across all nodes for each node type.
+
+        Args:
+            path (str, optional): The path to where the plot is saved.
+                If set to :obj:`None`, will visualize the plot on-the-fly.
+                (default: :obj:`None`)
+            feat_labels (Dict[NodeType, List[str]], optional): The labels of
+                features for each node type. (default :obj:`None`)
+            top_k (int, optional): Top k features to plot. If :obj:`None`
+                plots all features. (default: :obj:`None`)
+        """
+        node_mask_dict = self.node_mask_dict
+        if len(node_mask_dict) == 0:
+            raise ValueError(f"The attribute 'node_mask' is not available "
+                             f"in '{self.__class__.__name__}' "
+                             f"(got {self.available_explanations})")
+        for node_mask in node_mask_dict.values():
+            if node_mask.dim() != 2 or node_mask.size(1) <= 1:
+                raise ValueError(f"Cannot compute feature importance for "
+                                 f"object-level 'node_mask' "
+                                 f"(got shape {node_mask_dict.size()})")
+
+        if feat_labels is None:
+            feat_labels = {}
+            for node_type, node_mask in node_mask_dict.items():
+                feat_labels[node_type] = range(node_mask.size(1))
+
+        score = torch.cat(
+            [node_mask.sum(dim=0) for node_mask in node_mask_dict.values()],
+            dim=0)
+
+        all_feat_labels = []
+        for node_type in node_mask_dict.keys():
+            all_feat_labels += [
+                f'{node_type}#{label}' for label in feat_labels[node_type]
+            ]
+
+        return _visualize_score(score, all_feat_labels, path, top_k)
+
+
+def _visualize_score(
+    score: torch.Tensor,
+    labels: List[str],
+    path: Optional[str] = None,
+    top_k: Optional[int] = None,
+):
+    import matplotlib.pyplot as plt
+    import pandas as pd
+
+    if len(labels) != score.numel():
+        raise ValueError(f"The number of labels (got {len(labels)}) must "
+                         f"match the number of scores (got {score.numel()})")
+
+    score = score.cpu().numpy()
+
+    df = pd.DataFrame({'score': score}, index=labels)
+    df = df.sort_values('score', ascending=False)
+    df = df.round(decimals=3)
+
+    if top_k is not None:
+        df = df.head(top_k)
+        title = f"Feature importance for top {len(df)} features"
+    else:
+        title = f"Feature importance for {len(df)} features"
+
+    ax = df.plot(
+        kind='barh',
+        figsize=(10, 7),
+        title=title,
+        ylabel='Feature label',
+        xlim=[0, float(df['score'].max()) + 0.3],
+        legend=False,
+    )
+    plt.gca().invert_yaxis()
+    ax.bar_label(container=ax.containers[0], label_type='edge')
+
+    if path is not None:
+        plt.savefig(path)
+    else:
+        plt.show()
+
+    plt.close()
diff --git a/torch_geometric/graphgym/logger.py b/torch_geometric/graphgym/logger.py
index ac753b63e..38c04ecef 100644
--- a/torch_geometric/graphgym/logger.py
+++ b/torch_geometric/graphgym/logger.py
@@ -35,7 +35,7 @@ def set_printing():
     logging.basicConfig(**logging_cfg)
 
 
-class Logger(object):
+class Logger:
     def __init__(self, name='train', task_type=None):
         self.name = name
         self.task_type = task_type
diff --git a/torch_geometric/graphgym/model_builder.py b/torch_geometric/graphgym/model_builder.py
index 90e7237fa..246ccfde5 100644
--- a/torch_geometric/graphgym/model_builder.py
+++ b/torch_geometric/graphgym/model_builder.py
@@ -74,7 +74,7 @@ def create_model(to_device=True, dim_in=None, dim_out=None) -> GraphGymModule:
     dim_in = cfg.share.dim_in if dim_in is None else dim_in
     dim_out = cfg.share.dim_out if dim_out is None else dim_out
     # binary classification, output dim = 1
-    if 'classification' in cfg.dataset.task_type and dim_out == 2:
+    if 'classification' == cfg.dataset.task_type and dim_out == 2:
         dim_out = 1
 
     model = GraphGymModule(dim_in, dim_out, cfg)
diff --git a/torch_geometric/io/npz.py b/torch_geometric/io/npz.py
index 075b91900..3e7ffd07f 100644
--- a/torch_geometric/io/npz.py
+++ b/torch_geometric/io/npz.py
@@ -1,17 +1,20 @@
+from typing import Any, Dict
+
 import numpy as np
 import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import Data
-from torch_geometric.utils import remove_self_loops, to_undirected
+from torch_geometric.utils import remove_self_loops
+from torch_geometric.utils import to_undirected as to_undirected_fn
 
 
-def read_npz(path):
+def read_npz(path: str, to_undirected: bool = True) -> Data:
     with np.load(path) as f:
         return parse_npz(f)
 
 
-def parse_npz(f):
+def parse_npz(f: Dict[str, Any], to_undirected: bool = True) -> Data:
     x = sp.csr_matrix((f['attr_data'], f['attr_indices'], f['attr_indptr']),
                       f['attr_shape']).todense()
     x = torch.from_numpy(x).to(torch.float)
@@ -23,7 +26,8 @@ def parse_npz(f):
     col = torch.from_numpy(adj.col).to(torch.long)
     edge_index = torch.stack([row, col], dim=0)
     edge_index, _ = remove_self_loops(edge_index)
-    edge_index = to_undirected(edge_index, num_nodes=x.size(0))
+    if to_undirected:
+        edge_index = to_undirected_fn(edge_index, num_nodes=x.size(0))
 
     y = torch.from_numpy(f['labels']).to(torch.long)
 
diff --git a/torch_geometric/loader/utils.py b/torch_geometric/loader/utils.py
index 1a2eeb855..00eca9c70 100644
--- a/torch_geometric/loader/utils.py
+++ b/torch_geometric/loader/utils.py
@@ -7,6 +7,7 @@ import numpy as np
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
 from torch_geometric.data import (
     Data,
     FeatureStore,
@@ -41,7 +42,11 @@ def index_select(value: FeatureTensorType, index: Tensor,
             size = list(value.shape)
             size[dim] = index.numel()
             numel = math.prod(size)
-            storage = value.storage()._new_shared(numel)
+            if torch_geometric.typing.WITH_PT2:
+                storage = value.untyped_storage()._new_shared(
+                    numel * value.element_size())
+            else:
+                storage = value.storage()._new_shared(numel)
             out = value.new(storage).view(size)
 
         return torch.index_select(value, dim, index, out=out)
diff --git a/torch_geometric/nn/conv/gcn_conv.py b/torch_geometric/nn/conv/gcn_conv.py
index cc61dfc0b..7ad729ae9 100644
--- a/torch_geometric/nn/conv/gcn_conv.py
+++ b/torch_geometric/nn/conv/gcn_conv.py
@@ -124,7 +124,7 @@ class GCNConv(MessagePassing):
 
     .. math::
         \mathbf{x}^{\prime}_i = \mathbf{\Theta}^{\top} \sum_{j \in
-        \mathcal{N}(v) \cup \{ i \}} \frac{e_{j,i}}{\sqrt{\hat{d}_j
+        \mathcal{N}(i) \cup \{ i \}} \frac{e_{j,i}}{\sqrt{\hat{d}_j
         \hat{d}_i}} \mathbf{x}_j
 
     with :math:`\hat{d}_i = 1 + \sum_{j \in \mathcal{N}(i)} e_{j,i}`, where
diff --git a/torch_geometric/nn/conv/hgt_conv.py b/torch_geometric/nn/conv/hgt_conv.py
index a97f2121a..97071cbe4 100644
--- a/torch_geometric/nn/conv/hgt_conv.py
+++ b/torch_geometric/nn/conv/hgt_conv.py
@@ -25,11 +25,6 @@ class HGTConv(MessagePassing):
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
         hetero/hgt_dblp.py>`_.
 
-    .. note::
-
-        For a faster alternative, use :class:`FastHGTConv` which does not
-        iterate over individual node and edge types.
-
     Args:
         in_channels (int or Dict[str, int]): Size of each input sample of every
             node type, or :obj:`-1` to derive the size from the first input(s)
@@ -71,6 +66,10 @@ class HGTConv(MessagePassing):
         self.heads = heads
         self.node_types = metadata[0]
         self.edge_types = metadata[1]
+        self.edge_types_map = {
+            edge_type: i
+            for i, edge_type in enumerate(metadata[1])
+        }
 
         self.dst_node_types = set([key[-1] for key in self.edge_types])
 
@@ -83,10 +82,10 @@ class HGTConv(MessagePassing):
         dim = out_channels // heads
         num_types = heads * len(self.edge_types)
 
-        self.k_rel = HeteroLinear(dim, dim, num_types, is_sorted=True,
-                                  bias=False)
-        self.v_rel = HeteroLinear(dim, dim, num_types, is_sorted=True,
-                                  bias=False)
+        self.k_rel = HeteroLinear(dim, dim, num_types, bias=False,
+                                  is_sorted=True)
+        self.v_rel = HeteroLinear(dim, dim, num_types, bias=False,
+                                  is_sorted=True)
 
         self.skip = ParameterDict({
             node_type: Parameter(torch.Tensor(1))
@@ -121,36 +120,40 @@ class HGTConv(MessagePassing):
         return torch.cat(outs, dim=0), offset
 
     def _construct_src_node_feat(
-        self,
-        k_dict: Dict[str, Tensor],
-        v_dict: Dict[str, Tensor],
+        self, k_dict: Dict[str, Tensor], v_dict: Dict[str, Tensor],
+        edge_index_dict: Dict[EdgeType, Adj]
     ) -> Tuple[Tensor, Tensor, Dict[EdgeType, int]]:
         """Constructs the source node representations."""
-        count = 0
         cumsum = 0
+        num_edge_types = len(self.edge_types)
         H, D = self.heads, self.out_channels // self.heads
 
         # Flatten into a single tensor with shape [num_edge_types * heads, D]:
         ks: List[Tensor] = []
         vs: List[Tensor] = []
-        type_list: List[int] = []
+        type_list: List[Tensor] = []
         offset: Dict[EdgeType] = {}
-        for edge_type in self.edge_types:
-            src, _, _ = edge_type
-
-            ks.append(k_dict[src].reshape(-1, D))
-            vs.append(v_dict[src].reshape(-1, D))
-
+        for edge_type in edge_index_dict.keys():
+            src = edge_type[0]
             N = k_dict[src].size(0)
-            for _ in range(H):
-                type_list.append(torch.full((N, ), count, dtype=torch.long))
-                count += 1
             offset[edge_type] = cumsum
             cumsum += N
 
-        type_vec = torch.cat(type_list, dim=0)
-        k = self.k_rel(torch.cat(ks, dim=0), type_vec).view(-1, H, D)
-        v = self.v_rel(torch.cat(vs, dim=0), type_vec).view(-1, H, D)
+            # construct type_vec for curr edge_type with shape [H, D]
+            edge_type_offset = self.edge_types_map[edge_type]
+            type_vec = torch.arange(H, dtype=torch.long).view(-1, 1).repeat(
+                1, N) * num_edge_types + edge_type_offset
+
+            type_list.append(type_vec)
+            ks.append(k_dict[src])
+            vs.append(v_dict[src])
+
+        ks = torch.cat(ks, dim=0).transpose(0, 1).reshape(-1, D)
+        vs = torch.cat(vs, dim=0).transpose(0, 1).reshape(-1, D)
+        type_vec = torch.cat(type_list, dim=1).flatten()
+
+        k = self.k_rel(ks, type_vec).view(H, -1, D).transpose(0, 1)
+        v = self.v_rel(vs, type_vec).view(H, -1, D).transpose(0, 1)
 
         return k, v, offset
 
@@ -184,12 +187,14 @@ class HGTConv(MessagePassing):
         # Compute K, Q, V over node types:
         kqv_dict = self.kqv_lin(x_dict)
         for key, val in kqv_dict.items():
-            k_dict[key] = val[:, :F].view(-1, H, D)
-            q_dict[key] = val[:, F:2 * F].view(-1, H, D)
-            v_dict[key] = val[:, 2 * F:].view(-1, H, D)
+            k, q, v = torch.tensor_split(val, 3, dim=1)
+            k_dict[key] = k.view(-1, H, D)
+            q_dict[key] = q.view(-1, H, D)
+            v_dict[key] = v.view(-1, H, D)
 
         q, dst_offset = self._cat(q_dict)
-        k, v, src_offset = self._construct_src_node_feat(k_dict, v_dict)
+        k, v, src_offset = self._construct_src_node_feat(
+            k_dict, v_dict, edge_index_dict)
 
         edge_index, edge_attr = construct_bipartite_edge_index(
             edge_index_dict, src_offset, dst_offset, edge_attr_dict=self.p_rel)
@@ -200,7 +205,8 @@ class HGTConv(MessagePassing):
         # Reconstruct output node embeddings dict:
         for node_type, start_offset in dst_offset.items():
             end_offset = start_offset + q_dict[node_type].size(0)
-            out_dict[node_type] = out[start_offset:end_offset]
+            if node_type in self.dst_node_types:
+                out_dict[node_type] = out[start_offset:end_offset]
 
         # Transform output node embeddings:
         a_dict = self.out_lin({
@@ -210,11 +216,7 @@ class HGTConv(MessagePassing):
 
         # Iterate over node types:
         for node_type, out in out_dict.items():
-            if node_type not in self.dst_node_types:
-                out_dict[node_type] = None
-                continue
-            else:
-                out = a_dict[node_type]
+            out = a_dict[node_type]
 
             if out.size(-1) == x_dict[node_type].size(-1):
                 alpha = self.skip[node_type].sigmoid()
diff --git a/torch_geometric/nn/conv/pdn_conv.py b/torch_geometric/nn/conv/pdn_conv.py
index bc2740596..bdb00f42e 100644
--- a/torch_geometric/nn/conv/pdn_conv.py
+++ b/torch_geometric/nn/conv/pdn_conv.py
@@ -15,7 +15,7 @@ class PDNConv(MessagePassing):
     <https://arxiv.org/pdf/2010.12878.pdf>`_ paper
 
     .. math::
-        \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(v) \cup
+        \mathbf{x}^{\prime}_i = \sum_{j \in \mathcal{N}(i) \cup
         \{i\}}f_{\Theta}(\textbf{e}_{(j,i)}) \cdot f_{\Omega}(\mathbf{x}_{j})
 
     where :math:`z_{i,j}` denotes the edge feature vector from source node
diff --git a/torch_geometric/nn/conv/utils/inspector.py b/torch_geometric/nn/conv/utils/inspector.py
index 538db7988..e5d42b77a 100644
--- a/torch_geometric/nn/conv/utils/inspector.py
+++ b/torch_geometric/nn/conv/utils/inspector.py
@@ -6,7 +6,7 @@ from typing import Any, Callable, Dict, List, Optional, Set
 from torch_geometric.nn.conv.utils.typing import parse_types
 
 
-class Inspector(object):
+class Inspector:
     def __init__(self, base_class: Any):
         self.base_class: Any = base_class
         self.params: Dict[str, Dict[str, Any]] = {}
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index e042abb9e..b13c474bd 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -387,7 +387,7 @@ class HeteroDictLinear(torch.nn.Module):
                     biases.append(lin.bias)
             biases = None if biases[0] is None else biases
             outs = pyg_lib.ops.grouped_matmul(xs, weights, biases)
-            for key, out in zip(self.lins.keys(), outs):
+            for key, out in zip(x_dict.keys(), outs):
                 if key in x_dict:
                     out_dict[key] = out
         else:
diff --git a/torch_geometric/nn/fx.py b/torch_geometric/nn/fx.py
index faf3f6314..573e3f2bc 100644
--- a/torch_geometric/nn/fx.py
+++ b/torch_geometric/nn/fx.py
@@ -10,7 +10,7 @@ except (ImportError, ModuleNotFoundError, AttributeError):
     GraphModule, Graph, Node = 'GraphModule', 'Graph', 'Node'
 
 
-class Transformer(object):
+class Transformer:
     r"""A :class:`Transformer` executes an FX graph node-by-node, applies
     transformations to each node, and produces a new :class:`torch.nn.Module`.
     It exposes a :func:`transform` method that returns the transformed
diff --git a/torch_geometric/nn/models/label_prop.py b/torch_geometric/nn/models/label_prop.py
index 41645db3f..810c5fdb5 100644
--- a/torch_geometric/nn/models/label_prop.py
+++ b/torch_geometric/nn/models/label_prop.py
@@ -10,8 +10,8 @@ from torch_geometric.utils import one_hot, spmm
 
 
 class LabelPropagation(MessagePassing):
-    r"""The label propagation operator from the `"Learning from Labeled and
-    Unlabeled Data with Label Propagation"
+    r"""The label propagation operator, firstly introduced in the
+    `"Learning from Labeled and Unlabeled Data with Label Propagation"
     <http://mlg.eng.cam.ac.uk/zoubin/papers/CMU-CALD-02-107.pdf>`_ paper
 
     .. math::
@@ -19,6 +19,9 @@ class LabelPropagation(MessagePassing):
         \mathbf{D}^{-1/2} \mathbf{Y} + (1 - \alpha) \mathbf{Y},
 
     where unlabeled data is inferred by labeled data via propagation.
+    This concrete implementation here is derived from the `"Combining Label
+    Propagation And Simple Models Out-performs Graph Neural Networks"
+    <https://arxiv.org/abs/2010.13993>`_ paper.
 
     .. note::
 
diff --git a/torch_geometric/nn/models/mlp.py b/torch_geometric/nn/models/mlp.py
index 402a6c087..d2be2475e 100644
--- a/torch_geometric/nn/models/mlp.py
+++ b/torch_geometric/nn/models/mlp.py
@@ -194,9 +194,14 @@ class MLP(torch.nn.Module):
         Args:
             x (torch.Tensor): The source tensor.
             return_emb (bool, optional): If set to :obj:`True`, will
-                additionally return the embeddings before execution of to the
+                additionally return the embeddings before execution of the
                 final output layer. (default: :obj:`False`)
         """
+        # `return_emb` is annotated here as `NoneType` to be compatible with
+        # TorchScript, which does not support different return types based on
+        # the value of an input argument.
+        emb: Optional[Tensor] = None
+
         for i, (lin, norm) in enumerate(zip(self.lins, self.norms)):
             x = lin(x)
             if self.act is not None and self.act_first:
@@ -205,7 +210,8 @@ class MLP(torch.nn.Module):
             if self.act is not None and not self.act_first:
                 x = self.act(x)
             x = F.dropout(x, p=self.dropout[i], training=self.training)
-            emb = x
+            if isinstance(return_emb, bool) and return_emb is True:
+                emb = x
 
         if self.plain_last:
             x = self.lins[-1](x)
diff --git a/torch_geometric/nn/models/node2vec.py b/torch_geometric/nn/models/node2vec.py
index ae91b4a49..37ece5297 100644
--- a/torch_geometric/nn/models/node2vec.py
+++ b/torch_geometric/nn/models/node2vec.py
@@ -5,16 +5,11 @@ from torch import Tensor
 from torch.nn import Embedding
 from torch.utils.data import DataLoader
 
+from torch_geometric.typing import WITH_PYG_LIB, WITH_TORCH_CLUSTER
 from torch_geometric.utils import sort_edge_index
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 from torch_geometric.utils.sparse import index2ptr
 
-try:
-    import torch_cluster  # noqa
-    random_walk = torch.ops.torch_cluster.random_walk
-except ImportError:
-    random_walk = None
-
 
 class Node2Vec(torch.nn.Module):
     r"""The Node2Vec model from the
@@ -63,8 +58,18 @@ class Node2Vec(torch.nn.Module):
     ):
         super().__init__()
 
-        if random_walk is None:
-            raise ImportError('`Node2Vec` requires `torch-cluster`.')
+        if WITH_PYG_LIB and p == 1.0 and q == 1.0:
+            self.random_walk_fn = torch.ops.pyg.random_walk
+        elif WITH_TORCH_CLUSTER:
+            self.random_walk_fn = torch.ops.torch_cluster.random_walk
+        else:
+            if p == 1.0 and q == 1.0:
+                raise ImportError(f"'{self.__class__.__name__}' "
+                                  f"requires either the 'pyg-lib' or "
+                                  f"'torch-cluster' package")
+            else:
+                raise ImportError(f"'{self.__class__.__name__}' "
+                                  f"requires the 'torch-cluster' package")
 
         self.num_nodes = maybe_num_nodes(edge_index, num_nodes)
 
@@ -103,8 +108,8 @@ class Node2Vec(torch.nn.Module):
     @torch.jit.export
     def pos_sample(self, batch: Tensor) -> Tensor:
         batch = batch.repeat(self.walks_per_node)
-        rw = random_walk(self.rowptr, self.col, batch, self.walk_length,
-                         self.p, self.q)
+        rw = self.random_walk_fn(self.rowptr, self.col, batch,
+                                 self.walk_length, self.p, self.q)
         if not isinstance(rw, Tensor):
             rw = rw[0]
 
diff --git a/torch_geometric/nn/models/re_net.py b/torch_geometric/nn/models/re_net.py
index d15f173a0..cb4ec8e2b 100644
--- a/torch_geometric/nn/models/re_net.py
+++ b/torch_geometric/nn/models/re_net.py
@@ -94,7 +94,7 @@ class RENet(torch.nn.Module):
         of a :class:`torch_geometric.datasets.icews.EventDataset` with
         :math:`k` denoting the sequence length :obj:`seq_len`.
         """
-        class PreTransform(object):
+        class PreTransform:
             def __init__(self, seq_len: int):
                 self.seq_len = seq_len
                 self.inc = 5000
diff --git a/torch_geometric/nn/models/tgn.py b/torch_geometric/nn/models/tgn.py
index 9026fb4e8..d224f9b06 100644
--- a/torch_geometric/nn/models/tgn.py
+++ b/torch_geometric/nn/models/tgn.py
@@ -220,7 +220,7 @@ class TimeEncoder(torch.nn.Module):
         return self.lin(t.view(-1, 1)).cos()
 
 
-class LastNeighborLoader(object):
+class LastNeighborLoader:
     def __init__(self, num_nodes: int, size: int, device=None):
         self.size = size
 
diff --git a/torch_geometric/nn/module_dict.py b/torch_geometric/nn/module_dict.py
index 09c3cab8c..122e03a30 100644
--- a/torch_geometric/nn/module_dict.py
+++ b/torch_geometric/nn/module_dict.py
@@ -1,15 +1,20 @@
-from typing import Iterable, Mapping, Optional, Tuple
+from typing import Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Module
 
+Key = Union[str, Tuple[str, ...]]
+
 
 # `torch.nn.ModuleDict` doesn't allow `.` to be used in key names.
 # This `ModuleDict` will support it by converting the `.` to `#` in the
 # internal representation and converts it back to `.` in the external
-# representation.
+# representation. It also allows passing tuples as keys.
 class ModuleDict(torch.nn.ModuleDict):
-    def __init__(self, modules: Optional[Mapping[str, Module]] = None):
+    def __init__(
+        self,
+        modules: Optional[Mapping[Union[str, Tuple[str, ...]], Module]] = None,
+    ):
         if modules is not None:  # Replace the keys in modules:
             modules = {
                 self.to_internal_key(key): module
@@ -18,28 +23,34 @@ class ModuleDict(torch.nn.ModuleDict):
         super().__init__(modules)
 
     @staticmethod
-    def to_internal_key(key: str) -> str:
-        return key.replace(".", "#")
+    def to_internal_key(key: Key) -> str:
+        if isinstance(key, tuple):
+            assert len(key) > 1
+            key = f"<{'___'.join(key)}>"
+        assert isinstance(key, str)
+        return key.replace('.', '#')
 
     @staticmethod
-    def to_external_key(key: str) -> str:
-        return key.replace("#", ".")
+    def to_external_key(key: str) -> Key:
+        key = key.replace('#', '.')
+        if key.startswith('<') and key.endswith('>') and '___' in key:
+            key = tuple(key[1:-1].split('___'))
+        return key
 
-    def __getitem__(self, key: str) -> Module:
+    def __getitem__(self, key: Key) -> Module:
         return super().__getitem__(self.to_internal_key(key))
 
-    def __setitem__(self, key: str, module: Module):
+    def __setitem__(self, key: Key, module: Module):
         return super().__setitem__(self.to_internal_key(key), module)
 
-    def __delitem__(self, key: str):
+    def __delitem__(self, key: Key):
         return super().__delitem__(self.to_internal_key(key))
 
-    def __contains__(self, key: str) -> bool:
+    def __contains__(self, key: Key) -> bool:
         return super().__contains__(self.to_internal_key(key))
 
-    def keys(self) -> Iterable[str]:
+    def keys(self) -> Iterable[Key]:
         return [self.to_external_key(key) for key in super().keys()]
 
-    def items(self) -> Iterable[Tuple[str, Module]]:
-        return [(self.to_external_key(key), value)
-                for key, value in super().items()]
+    def items(self) -> Iterable[Tuple[Key, Module]]:
+        return [(self.to_external_key(k), v) for k, v in super().items()]
diff --git a/torch_geometric/nn/norm/layer_norm.py b/torch_geometric/nn/norm/layer_norm.py
index c40d3cc9c..fdaaeeb76 100644
--- a/torch_geometric/nn/norm/layer_norm.py
+++ b/torch_geometric/nn/norm/layer_norm.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import List, Optional, Union
 
 import torch
 import torch.nn.functional as F
@@ -143,6 +143,7 @@ class HeteroLayerNorm(torch.nn.Module):
         mode: str = 'node',
     ):
         super().__init__()
+        assert mode == 'node'
 
         self.in_channels = in_channels
         self.num_types = num_types
@@ -164,16 +165,41 @@ class HeteroLayerNorm(torch.nn.Module):
             torch.nn.init.ones_(self.weight)
             torch.nn.init.zeros_(self.bias)
 
-    def forward(self, x: Tensor, type_vec: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        type_vec: OptTensor = None,
+        type_ptr: Optional[Union[Tensor, List[int]]] = None,
+    ) -> Tensor:
         r"""
+        .. note::
+            Either :obj:`type_vec` or :obj:`type_ptr` needs to be specified.
+            In general, relying on :obj:`type_ptr` is more efficient in case
+            the input tensor is sorted by types.
+
         Args:
             x (torch.Tensor): The input features.
-            type_vec (torch.Tensor): A vector that maps each entry to a type.
+            type_vec (torch.Tensor, optional): A vector that maps each entry to
+                a type. (default: :obj:`None`)
+            type_ptr (torch.Tensor or List[int]): A vector denoting the
+                boundaries of types. (default: :obj:`None`)
         """
+        if type_vec is None and type_ptr is None:
+            raise ValueError("Either 'type_vec' or 'type_ptr' must be given")
+
         out = F.layer_norm(x, (self.in_channels, ), None, None, self.eps)
 
         if self.affine:
-            out = out * self.weight[type_vec] + self.bias[type_vec]
+            # TODO Revisit this logic completely as it performs worse than just
+            # operating on a dictionary of tensors
+            # (especially the `type_vec` code path)
+            if type_ptr is not None:
+                h = torch.empty_like(out)
+                for i, (s, e) in enumerate(zip(type_ptr[:-1], type_ptr[1:])):
+                    h[s:e] = out[s:e] * self.weight[i] + self.bias[i]
+                out = h
+            else:
+                out = out * self.weight[type_vec] + self.bias[type_vec]
 
         return out
 
diff --git a/torch_geometric/nn/parameter_dict.py b/torch_geometric/nn/parameter_dict.py
index 1d8ebc59b..99d887bd1 100644
--- a/torch_geometric/nn/parameter_dict.py
+++ b/torch_geometric/nn/parameter_dict.py
@@ -1,42 +1,57 @@
-from typing import Iterable, Mapping, Optional
+from typing import Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Parameter
 
+Key = Union[str, Tuple[str, ...]]
+
 
 # `torch.nn.ParameterDict` doesn't allow `.` to be used in key names.
 # This `ParameterDict` will support it by converting the `.` to `#` in the
 # internal representation and converts it back to `.` in the external
-# representation.
+# representation. It also allows passing tuples as keys.
 class ParameterDict(torch.nn.ParameterDict):
-    def __init__(self, parameters: Optional[Mapping[str, Parameter]] = None):
+    def __init__(
+        self,
+        parameters: Optional[Mapping[Key, Parameter]] = None,
+    ):
         # Replace the keys in modules.
         if parameters:
             parameters = {
-                self.to_internal_key(key): module
-                for key, module in parameters.items()
+                self.to_internal_key(key): parameter
+                for key, parameter in parameters.items()
             }
         super().__init__(parameters)
 
     @staticmethod
-    def to_internal_key(key: str) -> str:
-        return key.replace(".", "#")
+    def to_internal_key(key: str) -> Key:
+        if isinstance(key, tuple):
+            assert len(key) > 1
+            key = f"<{'___'.join(key)}>"
+        assert isinstance(key, str)
+        return key.replace('.', '#')
 
     @staticmethod
-    def to_external_key(key: str) -> str:
-        return key.replace("#", ".")
+    def to_external_key(key: str) -> Key:
+        key = key.replace('#', '.')
+        if key.startswith('<') and key.endswith('>') and '___' in key:
+            key = tuple(key[1:-1].split('___'))
+        return key
 
-    def __getitem__(self, key: str) -> Parameter:
+    def __getitem__(self, key: Key) -> Parameter:
         return super().__getitem__(self.to_internal_key(key))
 
-    def __setitem__(self, key: str, parameter: Parameter):
+    def __setitem__(self, key: Key, parameter: Parameter):
         return super().__setitem__(self.to_internal_key(key), parameter)
 
-    def __delitem__(self, key: str):
+    def __delitem__(self, key: Key):
         return super().__delitem__(self.to_internal_key(key))
 
-    def __contains__(self, key: str) -> bool:
+    def __contains__(self, key: Key) -> bool:
         return super().__contains__(self.to_internal_key(key))
 
-    def keys(self) -> Iterable[str]:
+    def keys(self) -> Iterable[Key]:
         return [self.to_external_key(key) for key in super().keys()]
+
+    def items(self) -> Iterable[Tuple[Key, Parameter]]:
+        return [(self.to_external_key(k), v) for k, v in super().items()]
diff --git a/torch_geometric/nn/pool/avg_pool.py b/torch_geometric/nn/pool/avg_pool.py
index 48ea071dd..487ab4139 100644
--- a/torch_geometric/nn/pool/avg_pool.py
+++ b/torch_geometric/nn/pool/avg_pool.py
@@ -20,6 +20,7 @@ def avg_pool_x(
     cluster: Tensor,
     x: Tensor,
     batch: Tensor,
+    batch_size: Optional[int] = None,
     size: Optional[int] = None,
 ) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Average pools node features according to the clustering defined in
@@ -34,6 +35,8 @@ def avg_pool_x(
         batch (torch.Tensor): The batch vector
             :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
             node to a specific example.
+        batch_size (int, optional): The number of examples :math:`B`.
+            Automatically calculated if not given. (default: :obj:`None`)
         size (int, optional): The maximum number of clusters in a single
             example. (default: :obj:`None`)
 
@@ -41,7 +44,8 @@ def avg_pool_x(
         :obj:`None`, else :class:`torch.Tensor`
     """
     if size is not None:
-        batch_size = int(batch.max().item()) + 1
+        if batch_size is None:
+            batch_size = int(batch.max().item()) + 1
         return _avg_pool_x(cluster, x, batch_size * size), None
 
     cluster, perm = consecutive_cluster(cluster)
diff --git a/torch_geometric/nn/pool/max_pool.py b/torch_geometric/nn/pool/max_pool.py
index cf849a8cb..d099455ec 100644
--- a/torch_geometric/nn/pool/max_pool.py
+++ b/torch_geometric/nn/pool/max_pool.py
@@ -20,6 +20,7 @@ def max_pool_x(
     cluster: Tensor,
     x: Tensor,
     batch: Tensor,
+    batch_size: Optional[int] = None,
     size: Optional[int] = None,
 ) -> Tuple[Tensor, Optional[Tensor]]:
     r"""Max-Pools node features according to the clustering defined in
@@ -33,6 +34,8 @@ def max_pool_x(
         batch (torch.Tensor): The batch vector
             :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
             node to a specific example.
+        batch_size (int, optional): The number of examples :math:`B`.
+            Automatically calculated if not given. (default: :obj:`None`)
         size (int, optional): The maximum number of clusters in a single
             example. This property is useful to obtain a batch-wise dense
             representation, *e.g.* for applying FC layers, but should only be
@@ -43,7 +46,8 @@ def max_pool_x(
         :obj:`None`, else :class:`torch.Tensor`
     """
     if size is not None:
-        batch_size = int(batch.max().item()) + 1
+        if batch_size is None:
+            batch_size = int(batch.max().item()) + 1
         return _max_pool_x(cluster, x, batch_size * size), None
 
     cluster, perm = consecutive_cluster(cluster)
diff --git a/torch_geometric/nn/pool/mem_pool.py b/torch_geometric/nn/pool/mem_pool.py
index 446916b60..00ea93808 100644
--- a/torch_geometric/nn/pool/mem_pool.py
+++ b/torch_geometric/nn/pool/mem_pool.py
@@ -78,8 +78,14 @@ class MemPooling(torch.nn.Module):
         loss = KLDivLoss(reduction='batchmean', log_target=False)
         return loss(S.clamp(EPS).log(), P.clamp(EPS))
 
-    def forward(self, x: Tensor, batch: Optional[Tensor] = None,
-                mask: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
+    def forward(
+        self,
+        x: Tensor,
+        batch: Optional[Tensor] = None,
+        mask: Optional[Tensor] = None,
+        max_num_nodes: Optional[int] = None,
+        batch_size: Optional[int] = None,
+    ) -> Tuple[Tensor, Tensor]:
         r"""
         Args:
             x (torch.Tensor): The node feature tensor of shape
@@ -97,9 +103,15 @@ class MemPooling(torch.nn.Module):
                 node features of shape
                 :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`.
                 (default: :obj:`None`)
+            max_num_nodes (int, optional): The size of the :math:`B` node
+                dimension. Automatically calculated if not given.
+                (default: :obj:`None`)
+            batch_size (int, optional): The number of examples :math:`B`.
+                Automatically calculated if not given. (default: :obj:`None`)
         """
         if x.dim() <= 2:
-            x, mask = to_dense_batch(x, batch)
+            x, mask = to_dense_batch(x, batch, max_num_nodes=max_num_nodes,
+                                     batch_size=batch_size)
         elif mask is None:
             mask = x.new_ones((x.size(0), x.size(1)), dtype=torch.bool)
 
diff --git a/torch_geometric/nn/resolver.py b/torch_geometric/nn/resolver.py
index ce6a94958..099ec3ff7 100644
--- a/torch_geometric/nn/resolver.py
+++ b/torch_geometric/nn/resolver.py
@@ -28,7 +28,6 @@ def swish(x: Tensor) -> Tensor:
 
 
 def activation_resolver(query: Union[Any, str] = 'relu', *args, **kwargs):
-    import torch
     base_cls = torch.nn.Module
     base_cls_repr = 'Act'
     acts = [
@@ -47,8 +46,6 @@ def activation_resolver(query: Union[Any, str] = 'relu', *args, **kwargs):
 
 
 def normalization_resolver(query: Union[Any, str], *args, **kwargs):
-    import torch
-
     import torch_geometric.nn.norm as norm
     base_cls = torch.nn.Module
     base_cls_repr = 'Norm'
@@ -77,6 +74,18 @@ def aggregation_resolver(query: Union[Any, str], *args, **kwargs):
     return resolver(aggrs, aggr_dict, query, base_cls, None, *args, **kwargs)
 
 
+# Optimizer Resolver ##########################################################
+
+
+def optimizer_resolver(query: Union[Any, str], *args, **kwargs):
+    base_cls = Optimizer
+    optimizers = [
+        optimizer for optimizer in vars(torch.optim).values()
+        if isinstance(optimizer, type) and issubclass(optimizer, base_cls)
+    ]
+    return resolver(optimizers, {}, query, base_cls, None, *args, **kwargs)
+
+
 # Learning Rate Scheduler Resolver ############################################
 
 
diff --git a/torch_geometric/nn/to_hetero_transformer.py b/torch_geometric/nn/to_hetero_transformer.py
index 219c05cdf..e6c4b3a6e 100644
--- a/torch_geometric/nn/to_hetero_transformer.py
+++ b/torch_geometric/nn/to_hetero_transformer.py
@@ -401,7 +401,11 @@ class ToHeteroTransformer(Transformer):
                         self.find_by_name(f'{value.name}__{key2str(key[-1])}'),
                     )
                 else:
-                    raise NotImplementedError
+                    raise ValueError(f"Cannot generate a graph node '{node}' "
+                                     f"for type '{key}' since it does not "
+                                     f"exist. Please make sure that all "
+                                     f"node types get updated during message "
+                                     f"passing.")
             elif isinstance(value, dict):
                 return {k: _recurse(v) for k, v in value.items()}
             elif isinstance(value, list):
diff --git a/torch_geometric/profile/benchmark.py b/torch_geometric/profile/benchmark.py
index f71ef9c18..60e6b2b8e 100644
--- a/torch_geometric/profile/benchmark.py
+++ b/torch_geometric/profile/benchmark.py
@@ -1,5 +1,5 @@
 import time
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -7,9 +7,22 @@ from torch import Tensor
 from torch_geometric.utils import is_torch_sparse_tensor
 
 
+def require_grad(x: Any, requires_grad: bool = True) -> Any:
+    if (isinstance(x, Tensor) and x.is_floating_point()
+            and not is_torch_sparse_tensor(x)):
+        return x.detach().requires_grad_(requires_grad)
+    elif isinstance(x, list):
+        return [require_grad(v, requires_grad) for v in x]
+    elif isinstance(x, tuple):
+        return tuple(require_grad(v, requires_grad) for v in x)
+    elif isinstance(x, dict):
+        return {k: require_grad(v, requires_grad) for k, v in x.items()}
+    return x
+
+
 def benchmark(
     funcs: List[Callable],
-    args: Tuple[Any],
+    args: Union[Tuple[Any], List[Tuple[Any]]],
     num_steps: int,
     func_names: Optional[List[str]] = None,
     num_warmups: int = 10,
@@ -20,7 +33,9 @@ def benchmark(
 
     Args:
         funcs ([Callable]): The list of functions to benchmark.
-        args ((Any, )): The arguments to pass to the functions.
+        args ((Any, ) or [(Any, )]): The arguments to pass to the functions.
+            Can be a list of arguments for each function in :obj:`funcs` in
+            case their headers differ.
         num_steps (int): The number of steps to run the benchmark.
         func_names ([str], optional): The names of the functions. If not given,
             will try to infer the name from the function itself.
@@ -47,15 +62,14 @@ def benchmark(
         raise ValueError(f"Length of 'funcs' (got {len(funcs)}) and "
                          f"'func_names' (got {len(func_names)}) must be equal")
 
+    # Zero-copy `args` for each function (if necessary):
+    args_list = [args] * len(funcs) if isinstance(args, tuple) else args
+
     ts: List[List[str]] = []
-    for func, name in zip(funcs, func_names):
+    for func, args, name in zip(funcs, args_list, func_names):
         t_forward = t_backward = 0
         for i in range(num_warmups + num_steps):
-            args = [
-                arg.detach().requires_grad_(backward)
-                if isinstance(arg, Tensor) and arg.is_floating_point()
-                and not is_torch_sparse_tensor(arg) else arg for arg in args
-            ]
+            args = require_grad(args, backward)
 
             if torch.cuda.is_available():
                 torch.cuda.synchronize()
@@ -69,6 +83,13 @@ def benchmark(
                 t_forward += time.perf_counter() - t_start
 
             if backward:
+                # TODO Generalize this logic. This is also a bit unfair as the
+                # concatenation leads to incorrectly measured backward speeds.
+                if isinstance(out, (tuple, list)):
+                    out = torch.cat(out, dim=0)
+                elif isinstance(out, dict):
+                    out = torch.cat(list(out.values()), dim=0)
+
                 out_grad = torch.randn_like(out)
                 t_start = time.perf_counter()
 
diff --git a/torch_geometric/resolver.py b/torch_geometric/resolver.py
index 6b642d5bc..d10c36ab4 100644
--- a/torch_geometric/resolver.py
+++ b/torch_geometric/resolver.py
@@ -22,9 +22,7 @@ def resolver(classes: List[Any], class_dict: Dict[str, Any],
         if query_repr == key_repr:
             if inspect.isclass(cls):
                 obj = cls(*args, **kwargs)
-                assert callable(obj)
                 return obj
-            assert callable(cls)
             return cls
 
     for cls in classes:
@@ -32,9 +30,7 @@ def resolver(classes: List[Any], class_dict: Dict[str, Any],
         if query_repr in [cls_repr, cls_repr.replace(base_cls_repr, '')]:
             if inspect.isclass(cls):
                 obj = cls(*args, **kwargs)
-                assert callable(obj)
                 return obj
-            assert callable(cls)
             return cls
 
     choices = set(cls.__name__ for cls in classes) | set(class_dict.keys())
diff --git a/torch_geometric/transforms/__init__.py b/torch_geometric/transforms/__init__.py
index f56db80f8..aef8b7913 100644
--- a/torch_geometric/transforms/__init__.py
+++ b/torch_geometric/transforms/__init__.py
@@ -10,6 +10,7 @@ from .svd_feature_reduction import SVDFeatureReduction
 from .remove_training_classes import RemoveTrainingClasses
 from .random_node_split import RandomNodeSplit
 from .random_link_split import RandomLinkSplit
+from .node_property_split import NodePropertySplit
 from .mask import IndexToMask, MaskToIndex
 from .pad import Pad
 
@@ -18,6 +19,7 @@ from .one_hot_degree import OneHotDegree
 from .target_indegree import TargetIndegree
 from .local_degree_profile import LocalDegreeProfile
 from .add_self_loops import AddSelfLoops
+from .add_remaining_self_loops import AddRemainingSelfLoops
 from .remove_isolated_nodes import RemoveIsolatedNodes
 from .remove_duplicated_edges import RemoveDuplicatedEdges
 from .knn_graph import KNNGraph
@@ -71,6 +73,7 @@ general_transforms = [
     'RemoveTrainingClasses',
     'RandomNodeSplit',
     'RandomLinkSplit',
+    'NodePropertySplit',
     'IndexToMask',
     'MaskToIndex',
     'Pad',
@@ -82,6 +85,7 @@ graph_transforms = [
     'TargetIndegree',
     'LocalDegreeProfile',
     'AddSelfLoops',
+    'AddRemainingSelfLoops',
     'RemoveIsolatedNodes',
     'RemoveDuplicatedEdges',
     'KNNGraph',
diff --git a/torch_geometric/transforms/add_remaining_self_loops.py b/torch_geometric/transforms/add_remaining_self_loops.py
new file mode 100644
index 000000000..74a1470a0
--- /dev/null
+++ b/torch_geometric/transforms/add_remaining_self_loops.py
@@ -0,0 +1,49 @@
+from typing import Optional, Union
+
+from torch import Tensor
+
+from torch_geometric.data import Data, HeteroData
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+from torch_geometric.utils import add_remaining_self_loops
+
+
+@functional_transform('add_remaining_self_loops')
+class AddRemainingSelfLoops(BaseTransform):
+    r"""Adds remaining self-loops to the given homogeneous or heterogeneous
+    graph (functional name: :obj:`add_remaining_self_loops`).
+
+    Args:
+        attr (str, optional): The name of the attribute of edge weights
+            or multi-dimensional edge features to pass to
+            :meth:`torch_geometric.utils.add_remaining_self_loops`.
+            (default: :obj:`"edge_weight"`)
+        fill_value (float or Tensor or str, optional): The way to generate
+            edge features of self-loops (in case :obj:`attr != None`).
+            If given as :obj:`float` or :class:`torch.Tensor`, edge features of
+            self-loops will be directly given by :obj:`fill_value`.
+            If given as :obj:`str`, edge features of self-loops are computed by
+            aggregating all features of edges that point to the specific node,
+            according to a reduce operation. (:obj:`"add"`, :obj:`"mean"`,
+            :obj:`"min"`, :obj:`"max"`, :obj:`"mul"`). (default: :obj:`1.`)
+    """
+    def __init__(self, attr: Optional[str] = 'edge_weight',
+                 fill_value: Union[float, Tensor, str] = 1.0):
+        self.attr = attr
+        self.fill_value = fill_value
+
+    def __call__(
+        self,
+        data: Union[Data, HeteroData],
+    ) -> Union[Data, HeteroData]:
+        for store in data.edge_stores:
+            if store.is_bipartite() or 'edge_index' not in store:
+                continue
+
+            store.edge_index, edge_weight = add_remaining_self_loops(
+                store.edge_index, getattr(store, self.attr, None),
+                fill_value=self.fill_value, num_nodes=store.size(0))
+
+            setattr(store, self.attr, edge_weight)
+
+        return data
diff --git a/torch_geometric/transforms/node_property_split.py b/torch_geometric/transforms/node_property_split.py
new file mode 100644
index 000000000..df4a313a5
--- /dev/null
+++ b/torch_geometric/transforms/node_property_split.py
@@ -0,0 +1,165 @@
+from typing import Any, Dict, List
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import Data
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+from torch_geometric.utils import to_networkx
+
+
+@functional_transform('node_property_split')
+class NodePropertySplit(BaseTransform):
+    r"""Creates a node-level split with distributional shift based on a given
+    node property, as proposed in the `"Evaluating Robustness and Uncertainty
+    of Graph Models Under Structural Distributional Shifts"
+    <https://arxiv.org/abs/2302.13875v1>`__ paper
+    (functional name: :obj:`node_property_split`).
+
+    It splits the nodes in a given graph into five non-intersecting parts
+    based on their structural properties.
+    This can be used for transductive node prediction tasks with distributional
+    shifts.
+    It considers the in-distribution (ID) and out-of-distribution (OOD) subsets
+    of nodes.
+    The ID subset includes training, validation and testing parts, while
+    the OOD subset includes validation and testing parts.
+    As a result, it creates five associated node mask vectors for each graph,
+    three which are for the ID nodes (:obj:`id_train_mask`,
+    :obj:`id_val_mask`, :obj:`id_test_mask`), and two which are for the OOD
+    nodes (:obj:`ood_val_mask`, :obj:`ood_test_mask`).
+
+    This class implements three particular strategies for inducing
+    distributional shifts in a graph — based on **popularity**, **locality**
+    or **density**.
+
+    Args:
+        property_name (str): The name of the node property to be used
+            (:obj:`"popularity"`, :obj:`"locality"`, :obj:`"density"`).
+        ratios ([float]): A list of five ratio values for ID training,
+            ID validation, ID test, OOD validation and OOD test parts.
+            The values must sum to :obj:`1.0`.
+        ascending (bool, optional): Whether to sort nodes in ascending order
+            of the node property, so that nodes with greater values of the
+            property are considered to be OOD (default: :obj:`True`)
+
+    Example:
+
+    .. code-block:: python
+
+        from torch_geometric.transforms import NodePropertySplit
+        from torch_geometric.datasets.graph_generator import ERGraph
+
+        data = ERGraph(num_nodes=1000, edge_prob=0.4)()
+
+        property_name = 'popularity'
+        ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
+        tranaform = NodePropertySplit(property_name, ratios)
+
+        data = transform(data)
+    """
+    def __init__(
+        self,
+        property_name: str,
+        ratios: List[float],
+        ascending: bool = True,
+    ):
+        if property_name not in {'popularity', 'locality', 'density'}:
+            raise ValueError(f"Unexpected 'property_name' "
+                             f"(got '{property_name}')")
+
+        if len(ratios) != 5:
+            raise ValueError(f"'ratios' must contain 5 values "
+                             f"(got {len(ratios)})")
+
+        if sum(ratios) != 1.0:
+            raise ValueError(f"'ratios' must sum to 1.0 (got {sum(ratios)})")
+
+        self.property_name = property_name
+        self.compute_fn = _property_name_to_compute_fn[property_name]
+        self.ratios = ratios
+        self.ascending = ascending
+
+    def __call__(self, data: Data) -> Data:
+
+        G = to_networkx(data, to_undirected=True, remove_self_loops=True)
+        property_values = self.compute_fn(G, self.ascending)
+        mask_dict = self._mask_nodes_by_property(property_values, self.ratios)
+
+        for key, mask in mask_dict.items():
+            data[key] = mask
+
+        return data
+
+    @staticmethod
+    def _compute_popularity_property(G: Any, ascending: bool = True) -> Tensor:
+        import networkx.algorithms as A
+
+        property_values = torch.tensor(list(A.pagerank(G).values()))
+        property_values *= -1 if ascending else 1
+        return property_values
+
+    @staticmethod
+    def _compute_locality_property(G: Any, ascending: bool = True) -> Tensor:
+        import networkx.algorithms as A
+
+        pagerank_values = torch.tensor(list(A.pagerank(G).values()))
+
+        num_nodes = G.number_of_nodes()
+        personalization = dict(zip(range(num_nodes), [0.0] * num_nodes))
+        personalization[int(pagerank_values.argmax())] = 1.0
+
+        property_values = torch.tensor(
+            list(A.pagerank(G, personalization=personalization).values()))
+        property_values *= -1 if ascending else 1
+        return property_values
+
+    @staticmethod
+    def _compute_density_property(G: Any, ascending: bool = True) -> Tensor:
+        import networkx.algorithms as A
+
+        property_values = torch.tensor(list(A.clustering(G).values()))
+        property_values *= -1 if ascending else 1
+        return property_values
+
+    @staticmethod
+    def _mask_nodes_by_property(
+        property_values: Tensor,
+        ratios: List[float],
+    ) -> Dict[str, Tensor]:
+
+        num_nodes = property_values.size(0)
+        sizes = (num_nodes * torch.tensor(ratios)).round().long()
+        sizes[-1] -= sizes.sum() - num_nodes
+
+        perm = torch.randperm(num_nodes)
+        id_size = int(sizes[:3].sum())
+        perm = perm[property_values[perm].argsort()]
+        perm[:id_size] = perm[:id_size][torch.randperm(id_size)]
+
+        node_splits = perm.split(sizes.tolist())
+        names = [
+            'id_train_mask',
+            'id_val_mask',
+            'id_test_mask',
+            'ood_val_mask',
+            'ood_test_mask',
+        ]
+
+        split_masks = {}
+        for name, node_split in zip(names, node_splits):
+            split_mask = torch.zeros(num_nodes, dtype=torch.bool)
+            split_mask[node_split] = True
+            split_masks[name] = split_mask
+        return split_masks
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.property_name})'
+
+
+_property_name_to_compute_fn = {
+    'popularity': NodePropertySplit._compute_popularity_property,
+    'locality': NodePropertySplit._compute_locality_property,
+    'density': NodePropertySplit._compute_density_property,
+}
diff --git a/torch_geometric/transforms/normalize_features.py b/torch_geometric/transforms/normalize_features.py
index 5b7a0fa8b..57ca32376 100644
--- a/torch_geometric/transforms/normalize_features.py
+++ b/torch_geometric/transforms/normalize_features.py
@@ -23,7 +23,8 @@ class NormalizeFeatures(BaseTransform):
     ) -> Union[Data, HeteroData]:
         for store in data.stores:
             for key, value in store.items(*self.attrs):
-                value = value - value.min()
-                value.div_(value.sum(dim=-1, keepdim=True).clamp_(min=1.))
-                store[key] = value
+                if value.numel() > 0:
+                    value = value - value.min()
+                    value.div_(value.sum(dim=-1, keepdim=True).clamp_(min=1.))
+                    store[key] = value
         return data
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index a05fd2bd0..372739b93 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -42,6 +42,16 @@ except (ImportError, OSError) as e:
                       f"Disabling its usage. Stacktrace: {e}")
     WITH_TORCH_CLUSTER = False
 
+try:
+    import torch_spline_conv  # noqa
+    WITH_TORCH_SPLINE_CONV = True
+except (ImportError, OSError) as e:
+    if isinstance(e, OSError):
+        warnings.warn(
+            f"An issue occurred while importing 'torch-spline-conv'. "
+            f"Disabling its usage. Stacktrace: {e}")
+    WITH_TORCH_SPLINE_CONV = False
+
 try:
     import torch_sparse  # noqa
     from torch_sparse import SparseStorage, SparseTensor
@@ -53,7 +63,21 @@ except (ImportError, OSError) as e:
     WITH_TORCH_SPARSE = False
 
     class SparseStorage:
-        def __init__(*args, **kwargs):
+        def __init__(
+            self,
+            row: Optional[Tensor] = None,
+            rowptr: Optional[Tensor] = None,
+            col: Optional[Tensor] = None,
+            value: Optional[Tensor] = None,
+            sparse_sizes: Optional[Tuple[Optional[int], Optional[int]]] = None,
+            rowcount: Optional[Tensor] = None,
+            colptr: Optional[Tensor] = None,
+            colcount: Optional[Tensor] = None,
+            csr2csc: Optional[Tensor] = None,
+            csc2csr: Optional[Tensor] = None,
+            is_sorted: bool = False,
+            trust_data: bool = False,
+        ):
             raise ImportError("'SparseStorage' requires 'torch-sparse'")
 
     class SparseTensor:
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
index f1a602a65..552a83c36 100644
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -27,7 +27,8 @@ from .to_dense_adj import to_dense_adj
 from .nested import to_nested_tensor, from_nested_tensor
 from .sparse import (dense_to_sparse, is_sparse, is_torch_sparse_tensor,
                      to_torch_coo_tensor, to_torch_csr_tensor,
-                     to_torch_csc_tensor, to_edge_index)
+                     to_torch_csc_tensor, to_torch_sparse_tensor,
+                     to_edge_index)
 from .spmm import spmm
 from .unbatch import unbatch, unbatch_edge_index
 from .one_hot import one_hot
@@ -99,6 +100,7 @@ __all__ = [
     'to_torch_coo_tensor',
     'to_torch_csr_tensor',
     'to_torch_csc_tensor',
+    'to_torch_sparse_tensor',
     'to_edge_index',
     'spmm',
     'unbatch',
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
index d05d91730..ada6d2e6a 100644
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -239,6 +239,44 @@ def to_torch_csc_tensor(
     return adj
 
 
+def to_torch_sparse_tensor(
+    edge_index: Tensor,
+    edge_attr: Optional[Tensor] = None,
+    size: Optional[Union[int, Tuple[int, int]]] = None,
+    is_coalesced: bool = False,
+    layout: torch.layout = torch.sparse_coo,
+):
+    r"""Converts a sparse adjacency matrix defined by edge indices and edge
+    attributes to a :class:`torch.sparse.Tensor` with custom :obj:`layout`.
+    See :meth:`~torch_geometric.utils.to_edge_index` for the reverse operation.
+
+    Args:
+        edge_index (LongTensor): The edge indices.
+        edge_attr (Tensor, optional): The edge attributes.
+            (default: :obj:`None`)
+        size (int or (int, int), optional): The size of the sparse matrix.
+            If given as an integer, will create a quadratic sparse matrix.
+            If set to :obj:`None`, will infer a quadratic sparse matrix based
+            on :obj:`edge_index.max() + 1`. (default: :obj:`None`)
+        is_coalesced (bool): If set to :obj:`True`, will assume that
+            :obj:`edge_index` is already coalesced and thus avoids expensive
+            computation. (default: :obj:`False`)
+        layout (torch.layout, optional): The layout of the output sparse tensor
+            (:obj:`torch.sparse_coo`, :obj:`torch.sparse_csr`,
+            :obj:`torch.sparse_csc`). (default: :obj:`torch.sparse_coo`)
+
+    :rtype: :class:`torch.sparse.Tensor`
+    """
+    if layout == torch.sparse_coo:
+        return to_torch_coo_tensor(edge_index, edge_attr, size, is_coalesced)
+    if layout == torch.sparse_csr:
+        return to_torch_csr_tensor(edge_index, edge_attr, size, is_coalesced)
+    if layout == torch.sparse_csc:
+        return to_torch_csc_tensor(edge_index, edge_attr, size, is_coalesced)
+
+    raise ValueError(f"Unexpected sparse tensor layout (got '{layout}')")
+
+
 def to_edge_index(adj: Union[Tensor, SparseTensor]) -> Tuple[Tensor, Tensor]:
     r"""Converts a :class:`torch.sparse.Tensor` or a
     :class:`torch_sparse.SparseTensor` to edge indices and edge attributes.
@@ -341,3 +379,42 @@ def ptr2index(ptr: Tensor) -> Tensor:
 def index2ptr(index: Tensor, size: int) -> Tensor:
     return torch._convert_indices_from_coo_to_csr(
         index, size, out_int32=index.dtype == torch.int32)
+
+
+def cat(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
+    # TODO (matthias) We can make this more efficient by directly operating on
+    # the individual sparse tensor layouts.
+    assert dim in {0, 1, (0, 1)}
+
+    size = [0, 0]
+    edge_indices = []
+    edge_attrs = []
+    for tensor in tensors:
+        assert is_torch_sparse_tensor(tensor)
+        edge_index, edge_attr = to_edge_index(tensor)
+        edge_index = edge_index.clone()
+
+        if dim == 0:
+            edge_index[0] += size[0]
+            size[0] += tensor.size(0)
+            size[1] = max(size[1], tensor.size(1))
+        elif dim == 1:
+            edge_index[1] += size[1]
+            size[0] = max(size[0], tensor.size(0))
+            size[1] += tensor.size(1)
+        else:
+            edge_index[0] += size[0]
+            edge_index[1] += size[1]
+            size[0] += tensor.size(0)
+            size[1] += tensor.size(1)
+
+        edge_indices.append(edge_index)
+        edge_attrs.append(edge_attr)
+
+    return to_torch_sparse_tensor(
+        edge_index=torch.cat(edge_indices, dim=1),
+        edge_attr=torch.cat(edge_attrs, dim=0),
+        size=size,
+        is_coalesced=dim == (0, 1),
+        layout=tensors[0].layout,
+    )


Test patch:
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
index 76dc187b6..ca5f46324 100644
--- a/test/data/test_batch.py
+++ b/test/data/test_batch.py
@@ -1,12 +1,14 @@
 import os.path as osp
 
 import numpy as np
+import pytest
 import torch
 
 import torch_geometric
 from torch_geometric.data import Batch, Data, HeteroData
 from torch_geometric.testing import get_random_edge_index, withPackage
 from torch_geometric.typing import SparseTensor
+from torch_geometric.utils import to_edge_index, to_torch_sparse_tensor
 
 
 def test_batch_basic():
@@ -466,12 +468,10 @@ def test_nested_follow_batch():
     d4 = Data(xs=[tr(4, 3), tr(16, 4), tr(1, 2)], a={"aa": tr(8, 3)},
               x=tr(8, 5))
 
-    # Dataset
     data_list = [d1, d2, d3, d4]
 
     batch = Batch.from_data_list(data_list, follow_batch=['xs', 'a'])
 
-    # assert shapes
     assert batch.xs[0].shape == (19, 3)
     assert batch.xs[1].shape == (56, 4)
     assert batch.xs[2].shape == (7, 2)
@@ -480,7 +480,6 @@ def test_nested_follow_batch():
     assert len(batch.xs_batch) == 3
     assert len(batch.a_batch) == 1
 
-    # assert _batch
     assert batch.xs_batch[0].tolist() == \
            [0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3]
     assert batch.xs_batch[1].tolist() == \
@@ -490,3 +489,31 @@ def test_nested_follow_batch():
 
     assert batch.a_batch['aa'].tolist() == \
            [0] * 11 + [1] * 2 + [2] * 4 + [3] * 8
+
+
+@withPackage('torch>=2.0.0')
+@pytest.mark.parametrize('layout', [
+    torch.sparse_coo,
+    torch.sparse_csr,
+    torch.sparse_csc,
+])
+def test_torch_sparse_batch(layout):
+    x_dense = torch.randn(3, 4)
+    x = x_dense.to_sparse(layout=layout)
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    edge_attr = torch.rand(4)
+    adj = to_torch_sparse_tensor(edge_index, edge_attr, layout=layout)
+
+    data = Data(x=x, adj=adj)
+
+    batch = Batch.from_data_list([data, data])
+
+    assert batch.x.size() == (6, 4)
+    assert batch.x.layout == layout
+    assert torch.equal(batch.x.to_dense(), torch.cat([x_dense, x_dense], 0))
+
+    assert batch.adj.size() == (6, 6)
+    assert batch.adj.layout == layout
+    out = to_edge_index(batch.adj.to_sparse(layout=torch.sparse_csr))
+    assert torch.equal(out[0], torch.cat([edge_index, edge_index + 3], 1))
+    assert torch.equal(out[1], torch.cat([edge_attr, edge_attr], 0))
diff --git a/test/data/test_dataset.py b/test/data/test_dataset.py
index 067a055a1..d82cb09bd 100644
--- a/test/data/test_dataset.py
+++ b/test/data/test_dataset.py
@@ -10,10 +10,24 @@ from torch_geometric.typing import SparseTensor
 
 class MyTestDataset(InMemoryDataset):
     def __init__(self, data_list, transform=None):
-        super().__init__('/tmp/MyTestDataset', transform=transform)
+        super().__init__(None, transform=transform)
         self.data, self.slices = self.collate(data_list)
 
 
+class MyStoredTestDataset(InMemoryDataset):
+    def __init__(self, root, data_list, transform=None):
+        self.data_list = data_list
+        super().__init__(root, transform=transform)
+        self.load(self.processed_paths[0], data_cls=data_list[0].__class__)
+
+    @property
+    def processed_file_names(self) -> str:
+        return 'data.pt'
+
+    def process(self):
+        self.save(self.data_list, self.processed_paths[0])
+
+
 def test_in_memory_dataset():
     x1 = torch.Tensor([[1], [1], [1]])
     x2 = torch.Tensor([[2], [2], [2]])
@@ -57,6 +71,54 @@ def test_in_memory_dataset():
     assert torch.equal(dataset[1:].x, x2)
 
 
+def test_stored_in_memory_dataset(tmp_path):
+    x1 = torch.Tensor([[1], [1], [1]])
+    x2 = torch.Tensor([[2], [2], [2], [2]])
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    data1 = Data(x1, edge_index, num_nodes=3, test_int=1, test_str='1')
+    data2 = Data(x2, edge_index, num_nodes=4, test_int=2, test_str='2')
+
+    dataset = MyStoredTestDataset(tmp_path, [data1, data2])
+    assert dataset._data.num_nodes == 7
+    assert dataset._data._num_nodes == [3, 4]
+
+    assert torch.equal(dataset[0].x, x1)
+    assert torch.equal(dataset[0].edge_index, edge_index)
+    assert dataset[0].num_nodes == 3
+    assert torch.equal(dataset[0].test_int, torch.tensor([1]))
+    assert dataset[0].test_str == '1'
+
+    assert torch.equal(dataset[1].x, x2)
+    assert torch.equal(dataset[1].edge_index, edge_index)
+    assert dataset[1].num_nodes == 4
+    assert torch.equal(dataset[1].test_int, torch.tensor([2]))
+    assert dataset[1].test_str == '2'
+
+
+def test_stored_hetero_in_memory_dataset(tmp_path):
+    x1 = torch.Tensor([[1], [1], [1]])
+    x2 = torch.Tensor([[2], [2], [2], [2]])
+
+    data1 = HeteroData()
+    data1['paper'].x = x1
+    data1['paper'].num_nodes = 3
+
+    data2 = HeteroData()
+    data2['paper'].x = x2
+    data2['paper'].num_nodes = 4
+
+    dataset = MyStoredTestDataset(tmp_path, [data1, data2])
+    assert dataset._data['paper'].num_nodes == 7
+    assert dataset._data['paper']._num_nodes == [3, 4]
+
+    assert torch.equal(dataset[0]['paper'].x, x1)
+    assert dataset[0]['paper'].num_nodes == 3
+
+    assert torch.equal(dataset[1]['paper'].x, x2)
+    assert dataset[1]['paper'].num_nodes == 4
+
+
 def test_in_memory_num_classes():
     dataset = MyTestDataset([Data(), Data()])
     assert dataset.num_classes == 0
diff --git a/test/data/test_dataset_summary.py b/test/data/test_dataset_summary.py
index 4ebee67df..5eed8708d 100644
--- a/test/data/test_dataset_summary.py
+++ b/test/data/test_dataset_summary.py
@@ -1,10 +1,21 @@
 import torch
+from torch import Tensor
 
-from torch_geometric.data.summary import Summary
+from torch_geometric.data.summary import Stats, Summary
 from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
 from torch_geometric.testing import withPackage
 
 
+def check_stats(stats: Stats, expected: Tensor):
+    assert stats.mean == float(expected.mean())
+    assert stats.std == float(expected.std())
+    assert stats.min == float(expected.min())
+    assert stats.quantile25 == float(expected.quantile(0.25))
+    assert stats.median == float(expected.median())
+    assert stats.quantile75 == float(expected.quantile(0.75))
+    assert stats.max == float(expected.max())
+
+
 def test_dataset_summary():
     dataset = FakeDataset(num_graphs=10)
     num_nodes = torch.Tensor([data.num_nodes for data in dataset])
@@ -15,27 +26,24 @@ def test_dataset_summary():
     assert summary.name == 'FakeDataset'
     assert summary.num_graphs == 10
 
-    assert summary.num_nodes.mean == num_nodes.mean().item()
-    assert summary.num_nodes.std == num_nodes.std().item()
-    assert summary.num_nodes.min == num_nodes.min().item()
-    assert summary.num_nodes.quantile25 == num_nodes.quantile(0.25).item()
-    assert summary.num_nodes.median == num_nodes.median().item()
-    assert summary.num_nodes.quantile75 == num_nodes.quantile(0.75).item()
-    assert summary.num_nodes.max == num_nodes.max().item()
+    check_stats(summary.num_nodes, num_nodes)
+    check_stats(summary.num_edges, num_edges)
+
+
+@withPackage('tabulate')
+def test_dataset_summary_representation():
+    dataset = FakeDataset(num_graphs=10)
+
+    summary1 = Summary.from_dataset(dataset, per_type=False)
+    summary2 = Summary.from_dataset(dataset, per_type=True)
 
-    assert summary.num_edges.mean == num_edges.mean().item()
-    assert summary.num_edges.std == num_edges.std().item()
-    assert summary.num_edges.min == num_edges.min().item()
-    assert summary.num_edges.quantile25 == num_edges.quantile(0.25).item()
-    assert summary.num_edges.median == num_edges.median().item()
-    assert summary.num_edges.quantile75 == num_edges.quantile(0.75).item()
-    assert summary.num_edges.max == num_edges.max().item()
+    assert str(summary1) == str(summary2)
 
 
 @withPackage('tabulate')
 def test_dataset_summary_hetero():
     dataset1 = FakeHeteroDataset(num_graphs=10)
-    summary1 = Summary.from_dataset(dataset1)
+    summary1 = Summary.from_dataset(dataset1, per_type=False)
 
     dataset2 = [data.to_homogeneous() for data in dataset1]
     summary2 = Summary.from_dataset(dataset2)
@@ -43,3 +51,48 @@ def test_dataset_summary_hetero():
 
     assert summary1 == summary2
     assert str(summary1) == str(summary2)
+
+
+@withPackage('tabulate')
+def test_dataset_summary_hetero_representation_length():
+    dataset = FakeHeteroDataset(num_graphs=10)
+    summary = Summary.from_dataset(dataset)
+    num_lines = len(str(summary).splitlines())
+
+    stats_len = len(Stats.__dataclass_fields__)
+    len_header_and_border = 5
+    num_tables = 3  # general, stats per node type, stats per edge type
+
+    assert num_lines == num_tables * (stats_len + len_header_and_border)
+
+
+def test_dataset_summary_hetero_per_type_check():
+    dataset = FakeHeteroDataset(num_graphs=10)
+    exp_num_nodes = torch.Tensor([data.num_nodes for data in dataset])
+    exp_num_edges = torch.Tensor([data.num_edges for data in dataset])
+
+    summary = dataset.get_summary()
+
+    assert summary.name == 'FakeHeteroDataset'
+    assert summary.num_graphs == 10
+
+    check_stats(summary.num_nodes, exp_num_nodes)
+    check_stats(summary.num_edges, exp_num_edges)
+
+    num_nodes_per_type = {}
+    for node_type in dataset.node_types:
+        num_nodes_per_type[node_type] = torch.Tensor(
+            [data[node_type].num_nodes for data in dataset])
+
+    assert len(summary.num_nodes_per_type) == len(dataset.node_types)
+    for node_type, stats in summary.num_nodes_per_type.items():
+        check_stats(stats, num_nodes_per_type[node_type])
+
+    num_edges_per_type = {}
+    for edge_type in dataset.edge_types:
+        num_edges_per_type[edge_type] = torch.Tensor(
+            [data[edge_type].num_edges for data in dataset])
+
+    assert len(summary.num_edges_per_type) == len(dataset.edge_types)
+    for edge_type, stats in summary.num_edges_per_type.items():
+        check_stats(stats, num_edges_per_type[edge_type])
diff --git a/test/data/test_hetero_data.py b/test/data/test_hetero_data.py
index 1ea3c79bb..e0be1dd69 100644
--- a/test/data/test_hetero_data.py
+++ b/test/data/test_hetero_data.py
@@ -90,6 +90,20 @@ def test_init_hetero_data():
     assert len(data.edge_items()) == 3
 
 
+def test_hetero_data_to_from_dict():
+    data = HeteroData()
+    data.global_id = '1'
+    data['v1'].x = torch.randn(5, 16)
+    data['v2'].y = torch.randn(4, 16)
+    data['v1', 'v2'].edge_index = torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]])
+
+    out = HeteroData.from_dict(data.to_dict())
+    assert out.global_id == data.global_id
+    assert torch.equal(out['v1'].x, data['v1'].x)
+    assert torch.equal(out['v2'].y, data['v2'].y)
+    assert torch.equal(out['v1', 'v2'].edge_index, data['v1', 'v2'].edge_index)
+
+
 def test_hetero_data_functions():
     data = HeteroData()
     data['paper'].x = x_paper
diff --git a/test/explain/test_hetero_explanation.py b/test/explain/test_hetero_explanation.py
index 838eda2f6..f492407ca 100644
--- a/test/explain/test_hetero_explanation.py
+++ b/test/explain/test_hetero_explanation.py
@@ -1,3 +1,4 @@
+import os.path as osp
 from typing import Optional, Union
 
 import pytest
@@ -6,6 +7,7 @@ import torch
 from torch_geometric.data import HeteroData
 from torch_geometric.explain import HeteroExplanation
 from torch_geometric.explain.config import MaskType
+from torch_geometric.testing import withPackage
 
 
 def create_random_explanation(
@@ -115,3 +117,27 @@ def test_edge_mask():
     assert out['paper', 'author'].edge_mask.size() == (1, )
     assert torch.equal(out['paper', 'author'].edge_index,
                        torch.tensor([[1], [1]]))
+
+
+@withPackage('matplotlib')
+@pytest.mark.parametrize('top_k', [2, None])
+@pytest.mark.parametrize('node_mask_type', [None, 'attributes'])
+def test_visualize_feature_importance(
+    top_k,
+    node_mask_type,
+    tmp_path,
+    hetero_data,
+):
+    explanation = create_random_explanation(
+        hetero_data,
+        node_mask_type=node_mask_type,
+    )
+
+    path = osp.join(tmp_path, 'feature_importance.png')
+
+    if node_mask_type is None:
+        with pytest.raises(ValueError, match="node_mask' is not"):
+            explanation.visualize_feature_importance(path, top_k=top_k)
+    else:
+        explanation.visualize_feature_importance(path, top_k=top_k)
+        assert osp.exists(path)
diff --git a/test/nn/conv/test_hgt_conv.py b/test/nn/conv/test_hgt_conv.py
index 02bb7c8e4..a41a755a9 100644
--- a/test/nn/conv/test_hgt_conv.py
+++ b/test/nn/conv/test_hgt_conv.py
@@ -193,7 +193,45 @@ def test_hgt_conv_missing_dst_node_type():
     out_dict = conv(data.x_dict, data.edge_index_dict)
     assert out_dict['author'].size() == (4, 64)
     assert out_dict['paper'].size() == (6, 64)
-    assert out_dict['university'] is None
+    assert 'university' not in out_dict
+
+
+def test_hgt_conv_missing_input_node_type():
+    data = HeteroData()
+    data['author'].x = torch.randn(4, 16)
+    data['paper'].x = torch.randn(6, 32)
+    data['author', 'writes',
+         'paper'].edge_index = get_random_edge_index(4, 6, 20)
+
+    # Some nodes from metadata are missing in data.
+    # This might happen while using NeighborLoader.
+    metadata = (['author', 'paper',
+                 'university'], [('author', 'writes', 'paper')])
+    conv = HGTConv(-1, 64, metadata, heads=1)
+
+    out_dict = conv(data.x_dict, data.edge_index_dict)
+    assert out_dict['paper'].size() == (6, 64)
+    assert 'university' not in out_dict
+
+
+def test_hgt_conv_missing_edge_type():
+    data = HeteroData()
+    data['author'].x = torch.randn(4, 16)
+    data['paper'].x = torch.randn(6, 32)
+    data['university'].x = torch.randn(10, 32)
+
+    data['author', 'writes',
+         'paper'].edge_index = get_random_edge_index(4, 6, 20)
+
+    metadata = (['author', 'paper',
+                 'university'], [('author', 'writes', 'paper'),
+                                 ('university', 'employs', 'author')])
+    conv = HGTConv(-1, 64, metadata, heads=1)
+
+    out_dict = conv(data.x_dict, data.edge_index_dict)
+    assert out_dict['author'].size() == (4, 64)
+    assert out_dict['paper'].size() == (6, 64)
+    assert 'university' not in out_dict
 
 
 if __name__ == '__main__':
diff --git a/test/nn/dense/test_linear.py b/test/nn/dense/test_linear.py
index 15919831d..cfa20fc40 100644
--- a/test/nn/dense/test_linear.py
+++ b/test/nn/dense/test_linear.py
@@ -1,13 +1,18 @@
 import copy
+import warnings
+from typing import List
 
 import pytest
 import torch
+from torch import Tensor
 from torch.nn import Linear as PTLinear
 from torch.nn.parameter import UninitializedParameter
 
 import torch_geometric.typing
 from torch_geometric.nn import HeteroDictLinear, HeteroLinear, Linear
+from torch_geometric.profile import benchmark
 from torch_geometric.testing import withCUDA, withPackage
+from torch_geometric.typing import pyg_lib
 
 weight_inits = ['glorot', 'kaiming_uniform', None]
 bias_inits = ['zeros', None]
@@ -216,3 +221,84 @@ def test_hetero_linear_sort(type_vec, device):
         node_type = int(type_vec[i])
         expected = x[i] @ lin.weight[node_type] + lin.bias[node_type]
         assert torch.allclose(out[i], expected, atol=1e-3)
+
+
+if __name__ == '__main__':
+    import argparse
+    try:
+        import dgl
+        WITH_DLG = True
+    except:  # noqa
+        WITH_DGL = False
+
+    warnings.filterwarnings('ignore', '.*API of nested tensors.*')
+    warnings.filterwarnings('ignore', '.*TypedStorage is deprecated.*')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--backward', action='store_true')
+    args = parser.parse_args()
+
+    torch.manual_seed(12345)
+
+    def get_xs(mean: float, std: float, num_types: int,
+               channels: int) -> List[Tensor]:
+        num_nodes_list = torch.normal(
+            mean=torch.tensor([mean] * num_types, dtype=torch.float),
+            std=torch.tensor([std] * num_types, dtype=torch.float),
+        ).round().to(torch.long).tolist()
+
+        return [
+            torch.randn(num_nodes, channels, device=args.device)
+            for num_nodes in num_nodes_list
+        ]
+
+    def sequential(xs: List[Tensor], weights: List[Tensor]) -> List[Tensor]:
+        return [x @ weight for x, weight in zip(xs, weights)]
+
+    def nested(xs: List[Tensor], weights: List[Tensor]) -> List[Tensor]:
+        x = torch.nested.nested_tensor(xs)
+        weight = torch.nested.nested_tensor(weights)
+        return list(torch.matmul(x, weight).unbind(0))
+
+    def grouped(x: Tensor, ptr: Tensor, weight: Tensor) -> Tensor:
+        return pyg_lib.ops.segment_matmul(x, ptr, weight)
+
+    def padded(x: Tensor, weight: Tensor) -> Tensor:
+        return torch.matmul(x, weight)
+
+    def dgl_mm(x: Tensor, count: Tensor, weight: Tensor) -> Tensor:
+        return dgl.ops.segment_mm(x, weight, count)
+
+    num_nodes, channels = 1_000_000, 64
+
+    for num_types in [3, 5, 10, 50, 100, 200, 500, 1000]:
+        print(f'Number of types: {num_types}')
+        mean = num_nodes // num_types
+        std = mean // 4
+
+        xs = get_xs(mean, std, num_types, channels)
+        count = torch.tensor([x.size(0) for x in xs])
+        ptr = torch.tensor([0] + [x.size(0) for x in xs]).cumsum(0)
+        x = torch.cat(xs, dim=0)
+        padded_x = torch.nested.nested_tensor(xs).to_padded_tensor(padding=0.0)
+        weight = torch.randn(num_types, channels, channels, device=args.device)
+        weights = list(weight.unbind(0))
+
+        funcs = [sequential, grouped, padded]
+        func_names = ['Sequential', 'Grouped', 'Padded']
+        args_list = [(xs, weights), (x, ptr, weight), (padded_x, weight)]
+
+        if WITH_DGL:
+            funcs.append(dgl_mm)
+            func_names.append('DGL')
+            args_list.append((x, count, weight))
+
+        benchmark(
+            funcs=funcs,
+            func_names=func_names,
+            args=args_list,
+            num_steps=50 if args.device == 'cpu' else 500,
+            num_warmups=10 if args.device == 'cpu' else 100,
+            backward=args.backward,
+        )
diff --git a/test/nn/models/test_mlp.py b/test/nn/models/test_mlp.py
index cb28a89ac..d830a8ed7 100644
--- a/test/nn/models/test_mlp.py
+++ b/test/nn/models/test_mlp.py
@@ -2,7 +2,6 @@ import pytest
 import torch
 
 from torch_geometric.nn import MLP
-from torch_geometric.testing import is_full_test
 
 
 @pytest.mark.parametrize('norm', ['batch_norm', None])
@@ -22,9 +21,8 @@ def test_mlp(norm, act_first, plain_last):
     out = mlp(x)
     assert out.size() == (4, 64)
 
-    if is_full_test():
-        jit = torch.jit.script(mlp)
-        assert torch.allclose(jit(x), out)
+    jit = torch.jit.script(mlp)
+    assert torch.allclose(jit(x), out)
 
     torch.manual_seed(12345)
     mlp = MLP(
@@ -39,6 +37,20 @@ def test_mlp(norm, act_first, plain_last):
     assert torch.allclose(mlp(x), out)
 
 
+def test_mlp_return_emb():
+    x = torch.randn(4, 16)
+
+    mlp = MLP([16, 32, 1])
+
+    out, emb = mlp(x, return_emb=True)
+    assert out.size() == (4, 1)
+    assert emb.size() == (4, 32)
+
+    out, emb = mlp(x, return_emb=False)
+    assert out.size() == (4, 1)
+    assert emb is None
+
+
 @pytest.mark.parametrize('plain_last', [False, True])
 def test_fine_grained_mlp(plain_last):
     mlp = MLP(
diff --git a/test/nn/models/test_node2vec.py b/test/nn/models/test_node2vec.py
index 77d195656..0ce285ed5 100644
--- a/test/nn/models/test_node2vec.py
+++ b/test/nn/models/test_node2vec.py
@@ -1,3 +1,4 @@
+import pytest
 import torch
 
 from torch_geometric.nn import Node2Vec
@@ -5,12 +6,21 @@ from torch_geometric.testing import is_full_test, withCUDA, withPackage
 
 
 @withCUDA
+@withPackage('pyg_lib')
 @withPackage('torch_cluster')
-def test_node2vec(device):
+@pytest.mark.parametrize('p', [1.0])
+@pytest.mark.parametrize('q', [1.0, 0.5])
+def test_node2vec(device, p, q):
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], device=device)
 
-    model = Node2Vec(edge_index, embedding_dim=16, walk_length=2,
-                     context_size=2).to(device)
+    model = Node2Vec(
+        edge_index,
+        embedding_dim=16,
+        walk_length=2,
+        context_size=2,
+        p=p,
+        q=q,
+    ).to(device)
     assert str(model) == 'Node2Vec(3, 16)'
 
     assert model(torch.arange(3, device=device)).size() == (3, 16)
diff --git a/test/nn/norm/test_layer_norm.py b/test/nn/norm/test_layer_norm.py
index 3c17c1670..ba3ac977f 100644
--- a/test/nn/norm/test_layer_norm.py
+++ b/test/nn/norm/test_layer_norm.py
@@ -31,26 +31,39 @@ def test_layer_norm(device, affine, mode):
 @pytest.mark.parametrize('affine', [False, True])
 def test_hetero_layer_norm(device, affine):
     x = torch.randn((100, 16), device=device)
+    expected = LayerNorm(16, affine=affine, mode='node').to(device)(x)
 
     # Test single type:
-    norm = LayerNorm(16, affine=affine, mode='node').to(device)
-    expected = norm(x)
-
     type_vec = torch.zeros(100, dtype=torch.long, device=device)
+    type_ptr = [0, 100]
+
     norm = HeteroLayerNorm(16, num_types=1, affine=affine).to(device)
     assert str(norm) == 'HeteroLayerNorm(16, num_types=1)'
 
     out = norm(x, type_vec)
     assert out.size() == (100, 16)
-    assert torch.allclose(out, expected)
+    assert torch.allclose(out, expected, atol=1e-3)
+    assert torch.allclose(norm(out, type_ptr=type_ptr), expected, atol=1e-3)
+
+    mean = out.mean(dim=-1)
+    std = out.std(unbiased=False, dim=-1)
+    assert torch.allclose(mean, torch.zeros_like(mean), atol=1e-2)
+    assert torch.allclose(std, torch.ones_like(std), atol=1e-2)
 
     # Test multiple types:
-    type_vec = torch.randint(5, (100, ), device=device)
+    type_vec = torch.arange(5, device=device)
+    type_vec = type_vec.view(-1, 1).repeat(1, 20).view(-1)
+    type_ptr = [0, 20, 40, 60, 80, 100]
+
     norm = HeteroLayerNorm(16, num_types=5, affine=affine).to(device)
+    assert str(norm) == 'HeteroLayerNorm(16, num_types=5)'
+
     out = norm(x, type_vec)
     assert out.size() == (100, 16)
+    assert torch.allclose(out, expected, atol=1e-3)
+    assert torch.allclose(norm(out, type_ptr=type_ptr), expected, atol=1e-3)
 
     mean = out.mean(dim=-1)
     std = out.std(unbiased=False, dim=-1)
-    assert torch.allclose(mean, torch.zeros_like(mean), atol=1e-5)
-    assert torch.allclose(std, torch.ones_like(std), atol=1e-5)
+    assert torch.allclose(mean, torch.zeros_like(mean), atol=1e-2)
+    assert torch.allclose(std, torch.ones_like(std), atol=1e-2)
diff --git a/test/nn/pool/test_avg_pool.py b/test/nn/pool/test_avg_pool.py
index c1e0f6dd8..15c8113df 100644
--- a/test/nn/pool/test_avg_pool.py
+++ b/test/nn/pool/test_avg_pool.py
@@ -23,11 +23,18 @@ def test_avg_pool_x():
     out, _ = avg_pool_x(cluster, x, batch, size=2)
     assert out.tolist() == [[3, 4], [5, 6], [10, 11], [0, 0]]
 
+    batch_size = int(batch.max().item()) + 1
+    out2, _ = avg_pool_x(cluster, x, batch, batch_size=batch_size, size=2)
+    assert torch.equal(out, out2)
+
     if is_full_test():
         jit = torch.jit.script(avg_pool_x)
         out, _ = jit(cluster, x, batch, size=2)
         assert out.tolist() == [[3, 4], [5, 6], [10, 11], [0, 0]]
 
+        out2, _ = jit(cluster, x, batch, batch_size=batch_size, size=2)
+        assert torch.equal(out, out2)
+
 
 def test_avg_pool():
     cluster = torch.tensor([0, 1, 0, 1, 2, 2])
diff --git a/test/nn/pool/test_max_pool.py b/test/nn/pool/test_max_pool.py
index a6c4921fb..db58de3d9 100644
--- a/test/nn/pool/test_max_pool.py
+++ b/test/nn/pool/test_max_pool.py
@@ -23,11 +23,18 @@ def test_max_pool_x():
     out, _ = max_pool_x(cluster, x, batch, size=2)
     assert out.tolist() == [[5, 6], [7, 8], [11, 12], [0, 0]]
 
+    batch_size = int(batch.max().item()) + 1
+    out2, _ = max_pool_x(cluster, x, batch, batch_size=batch_size, size=2)
+    assert torch.equal(out, out2)
+
     if is_full_test():
         jit = torch.jit.script(max_pool_x)
         out, _ = jit(cluster, x, batch, size=2)
         assert out.tolist() == [[5, 6], [7, 8], [11, 12], [0, 0]]
 
+        out2, _ = jit(cluster, x, batch, batch_size=batch_size, size=2)
+        assert torch.equal(out, out2)
+
 
 def test_max_pool():
     cluster = torch.tensor([0, 1, 0, 1, 2, 2])
diff --git a/test/nn/test_module_dict.py b/test/nn/test_module_dict.py
index 16fced0bc..6dbfc9fc6 100644
--- a/test/nn/test_module_dict.py
+++ b/test/nn/test_module_dict.py
@@ -7,9 +7,13 @@ def test_internal_external_key_conversion():
     assert ModuleDict.to_internal_key('a.b') == 'a#b'
     assert ModuleDict.to_internal_key('ab') == 'ab'
     assert ModuleDict.to_internal_key('a.b.c') == 'a#b#c'
+    assert ModuleDict.to_internal_key(('a', 'b')) == '<a___b>'
+    assert ModuleDict.to_internal_key(('a.b', 'c')) == '<a#b___c>'
 
     assert ModuleDict.to_external_key('a#b') == 'a.b'
     assert ModuleDict.to_external_key('a#b#c') == 'a.b.c'
+    assert ModuleDict.to_external_key('<a___b>') == ('a', 'b')
+    assert ModuleDict.to_external_key('<a#b___c>') == ('a.b', 'c')
 
 
 def test_dot_syntax_keys():
@@ -26,5 +30,24 @@ def test_dot_syntax_keys():
     for key in expected_keys:
         assert key in module_dict
 
+    assert 'model.lin2' in module_dict
     del module_dict['model.lin2']
     assert 'model.lin2' not in module_dict
+
+
+def test_tuple_keys():
+    module_dict = ModuleDict({
+        ('a', 'b'): torch.nn.Linear(16, 16),
+        ('a.b', 'c'): torch.nn.Linear(8, 8),
+    })
+
+    expected_keys = {('a', 'b'), ('a.b', 'c')}
+    assert set(module_dict.keys()) == expected_keys
+    assert set([key for key, _ in module_dict.items()]) == expected_keys
+
+    for key in expected_keys:
+        assert key in module_dict
+
+    assert ('a', 'b') in module_dict
+    del module_dict['a', 'b']
+    assert ('a', 'b') not in module_dict
diff --git a/test/nn/test_parameter_dict.py b/test/nn/test_parameter_dict.py
index 34cde8c54..29d44f926 100644
--- a/test/nn/test_parameter_dict.py
+++ b/test/nn/test_parameter_dict.py
@@ -1,33 +1,55 @@
-from typing import Mapping
-
 import torch
-from torch.nn import Parameter
 
 from torch_geometric.nn.parameter_dict import ParameterDict
 
 
 def test_internal_external_key_conversion():
-    assert ParameterDict.to_internal_key("a.b") == "a#b"
-    assert ParameterDict.to_internal_key("ab") == "ab"
-    assert ParameterDict.to_internal_key("a.b.c") == "a#b#c"
+    assert ParameterDict.to_internal_key('a.b') == 'a#b'
+    assert ParameterDict.to_internal_key('ab') == 'ab'
+    assert ParameterDict.to_internal_key('a.b.c') == 'a#b#c'
+    assert ParameterDict.to_internal_key(('a', 'b')) == '<a___b>'
+    assert ParameterDict.to_internal_key(('a.b', 'c')) == '<a#b___c>'
 
-    assert ParameterDict.to_external_key("a#b") == "a.b"
-    assert ParameterDict.to_external_key("a#b#c") == "a.b.c"
+    assert ParameterDict.to_external_key('a#b') == 'a.b'
+    assert ParameterDict.to_external_key('a#b#c') == 'a.b.c'
+    assert ParameterDict.to_external_key('<a___b>') == ('a', 'b')
+    assert ParameterDict.to_external_key('<a#b___c>') == ('a.b', 'c')
 
 
 def test_dot_syntax_keys():
-    parameters: Mapping[str, Parameter] = {
-        "param1": Parameter(torch.Tensor(16, 16)),
-        "model.param2": Parameter(torch.Tensor(8, 8)),
-        "model.sub_model.param3": Parameter(torch.Tensor(4, 4)),
+    parameter_dict = {
+        'param1': torch.nn.Parameter(torch.Tensor(16, 16)),
+        'model.param2': torch.nn.Parameter(torch.Tensor(8, 8)),
+        'model.sub_model.param3': torch.nn.Parameter(torch.Tensor(4, 4)),
+    }
+    parameter_dict = ParameterDict(parameter_dict)
+
+    expected_keys = {'param1', 'model.param2', 'model.sub_model.param3'}
+    assert set(parameter_dict.keys()) == expected_keys
+    assert set([key for key, _ in parameter_dict.items()]) == expected_keys
+
+    for key in expected_keys:
+        assert key in parameter_dict
+
+    assert 'model.param2' in parameter_dict
+    del parameter_dict['model.param2']
+    assert 'model.param2' not in parameter_dict
+
+
+def test_tuple_keys():
+    parameter_dict = {
+        ('a', 'b'): torch.nn.Parameter(torch.Tensor(16, 16)),
+        ('a.b', 'c'): torch.nn.Parameter(torch.Tensor(8, 8)),
     }
-    parameter_dict = ParameterDict(parameters)
+    parameter_dict = ParameterDict(parameter_dict)
 
-    expected_keys = {"param1", "model.param2", "model.sub_model.param3"}
+    expected_keys = {('a', 'b'), ('a.b', 'c')}
     assert set(parameter_dict.keys()) == expected_keys
+    assert set([key for key, _ in parameter_dict.items()]) == expected_keys
 
     for key in expected_keys:
         assert key in parameter_dict
 
-    del parameter_dict["model.param2"]
-    assert "model.param2" not in parameter_dict
+    assert ('a', 'b') in parameter_dict
+    del parameter_dict['a', 'b']
+    assert ('a', 'b') not in parameter_dict
diff --git a/test/nn/test_resolver.py b/test/nn/test_resolver.py
index 46ffd4ed8..cc3dac69e 100644
--- a/test/nn/test_resolver.py
+++ b/test/nn/test_resolver.py
@@ -8,6 +8,7 @@ from torch_geometric.nn.resolver import (
     aggregation_resolver,
     lr_scheduler_resolver,
     normalization_resolver,
+    optimizer_resolver,
 )
 
 
@@ -58,6 +59,22 @@ def test_normalization_resolver(norm_tuple):
                       norm_module)
 
 
+def test_optimizer_resolver():
+    params = [torch.nn.Parameter(torch.Tensor(1))]
+
+    assert isinstance(optimizer_resolver(torch.optim.SGD(params, lr=0.01)),
+                      torch.optim.SGD)
+    assert isinstance(optimizer_resolver(torch.optim.Adam(params)),
+                      torch.optim.Adam)
+    assert isinstance(optimizer_resolver(torch.optim.Rprop(params)),
+                      torch.optim.Rprop)
+
+    assert isinstance(optimizer_resolver('sgd', params, lr=0.01),
+                      torch.optim.SGD)
+    assert isinstance(optimizer_resolver('adam', params), torch.optim.Adam)
+    assert isinstance(optimizer_resolver('rprop', params), torch.optim.Rprop)
+
+
 @pytest.mark.parametrize('scheduler_args', [
     ('constant_with_warmup', LambdaLR),
     ('linear_with_warmup', LambdaLR),
diff --git a/test/profile/test_profile_utils.py b/test/profile/test_profile_utils.py
index c9215be68..4f38aa21a 100644
--- a/test/profile/test_profile_utils.py
+++ b/test/profile/test_profile_utils.py
@@ -24,7 +24,7 @@ def test_count_parameters():
 
 def test_get_model_size():
     model_size = get_model_size(Linear(32, 128, bias=False))
-    assert model_size >= 32 * 128 * 4 and model_size < 32 * 128 * 4 + 1000
+    assert model_size >= 32 * 128 * 4 and model_size < 32 * 128 * 4 + 2000
 
 
 def test_get_data_size():
diff --git a/test/transforms/test_add_remaining_self_loops.py b/test/transforms/test_add_remaining_self_loops.py
new file mode 100644
index 000000000..1de56f5a3
--- /dev/null
+++ b/test/transforms/test_add_remaining_self_loops.py
@@ -0,0 +1,72 @@
+import torch
+
+from torch_geometric.data import Data, HeteroData
+from torch_geometric.transforms import AddRemainingSelfLoops
+
+
+def test_add_remaining_self_loops():
+    assert str(AddRemainingSelfLoops()) == 'AddRemainingSelfLoops()'
+
+    assert len(AddRemainingSelfLoops()(Data())) == 0
+
+    # No self-loops in `edge_index`.
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    edge_weight = torch.tensor([1, 2, 3, 4])
+    edge_attr = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
+
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert len(data) == 2
+    assert data.edge_index.tolist() == [[0, 1, 1, 2, 0, 1, 2],
+                                        [1, 0, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+
+    # Single self-loop in `edge_index`.
+    edge_index = torch.tensor([[0, 0, 1, 2], [1, 0, 2, 1]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert len(data) == 2
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [1, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+
+    data = Data(edge_index=edge_index, edge_weight=edge_weight, num_nodes=3)
+    data = AddRemainingSelfLoops(attr='edge_weight', fill_value=5)(data)
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [1, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+    assert data.edge_weight.tolist() == [1, 3, 4, 2, 5, 5]
+
+    data = Data(edge_index=edge_index, edge_attr=edge_attr, num_nodes=3)
+    data = AddRemainingSelfLoops(attr='edge_attr', fill_value='add')(data)
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [1, 2, 1, 0, 1, 2]]
+    assert data.num_nodes == 3
+    assert data.edge_attr.tolist() == [[1, 2], [5, 6], [7, 8], [3, 4], [8, 10],
+                                       [5, 6]]
+
+
+def test_add_remaining_self_loops_all_loops_exist():
+    # All self-loops already exist in the data object.
+    edge_index = torch.tensor([[0, 1, 2], [0, 1, 2]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert data.edge_index.tolist() == edge_index.tolist()
+
+    # All self-loops already exist in the data object, some of them appear
+    # multiple times.
+    edge_index = torch.tensor([[0, 0, 1, 1, 2], [0, 0, 1, 1, 2]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddRemainingSelfLoops()(data)
+    assert data.edge_index.tolist() == [[0, 1, 2], [0, 1, 2]]
+
+
+def test_hetero_add_remaining_self_loops():
+    edge_index = torch.tensor([[0, 0, 1, 2], [1, 0, 2, 1]])
+
+    data = HeteroData()
+    data['v'].num_nodes = 3
+    data['w'].num_nodes = 3
+    data['v', 'v'].edge_index = edge_index
+    data['v', 'w'].edge_index = edge_index
+    data = AddRemainingSelfLoops()(data)
+    assert data['v', 'v'].edge_index.tolist() == [[0, 1, 2, 0, 1, 2],
+                                                  [1, 2, 1, 0, 1, 2]]
+    assert data['v', 'w'].edge_index.tolist() == edge_index.tolist()
diff --git a/test/transforms/test_add_self_loops.py b/test/transforms/test_add_self_loops.py
index 45455bcda..38a4ed0ff 100644
--- a/test/transforms/test_add_self_loops.py
+++ b/test/transforms/test_add_self_loops.py
@@ -7,12 +7,12 @@ from torch_geometric.transforms import AddSelfLoops
 def test_add_self_loops():
     assert str(AddSelfLoops()) == 'AddSelfLoops()'
 
+    assert len(AddSelfLoops()(Data())) == 0
+
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
     edge_weight = torch.tensor([1, 2, 3, 4])
     edge_attr = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
 
-    assert len(AddSelfLoops()(Data())) == 0
-
     data = Data(edge_index=edge_index, num_nodes=3)
     data = AddSelfLoops()(data)
     assert len(data) == 2
@@ -36,6 +36,14 @@ def test_add_self_loops():
                                        [8, 10], [5, 6]]
 
 
+def test_add_self_loops_with_existing_self_loops():
+    edge_index = torch.tensor([[0, 1, 2], [0, 1, 2]])
+    data = Data(edge_index=edge_index, num_nodes=3)
+    data = AddSelfLoops()(data)
+    assert data.edge_index.tolist() == [[0, 1, 2, 0, 1, 2], [0, 1, 2, 0, 1, 2]]
+    assert data.num_nodes == 3
+
+
 def test_hetero_add_self_loops():
     edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
 
diff --git a/test/transforms/test_node_property_split.py b/test/transforms/test_node_property_split.py
new file mode 100644
index 000000000..27686fbc3
--- /dev/null
+++ b/test/transforms/test_node_property_split.py
@@ -0,0 +1,39 @@
+import pytest
+import torch
+
+from torch_geometric.datasets import graph_generator
+from torch_geometric.testing import withPackage
+from torch_geometric.transforms import NodePropertySplit
+
+
+@withPackage('networkx')
+@pytest.mark.parametrize('property_name', [
+    'popularity',
+    'locality',
+    'density',
+])
+def test_node_property_split(property_name):
+    ratios = [0.3, 0.1, 0.1, 0.2, 0.3]
+
+    transform = NodePropertySplit(property_name, ratios)
+    assert str(transform) == f'NodePropertySplit({property_name})'
+
+    data = graph_generator.ERGraph(num_nodes=100, edge_prob=0.4)()
+    data = transform(data)
+
+    node_ids = []
+    for name, ratio in zip([
+            'id_train_mask',
+            'id_val_mask',
+            'id_test_mask',
+            'ood_val_mask',
+            'ood_test_mask',
+    ], ratios):
+        assert data[name].dtype == torch.bool
+        assert data[name].size() == (100, )
+        assert int(data[name].sum()) == 100 * ratio
+        node_ids.extend(data[name].nonzero().view(-1).tolist())
+
+    # Check that masks are non-intersecting and cover all nodes:
+    node_ids = torch.tensor(node_ids)
+    assert node_ids.numel() == torch.unique(node_ids).numel() == 100
diff --git a/test/transforms/test_random_link_split.py b/test/transforms/test_random_link_split.py
index 5c6fe5df4..764f7b4b0 100644
--- a/test/transforms/test_random_link_split.py
+++ b/test/transforms/test_random_link_split.py
@@ -2,7 +2,7 @@ import pytest
 import torch
 
 from torch_geometric.data import Data, HeteroData
-from torch_geometric.testing import get_random_edge_index
+from torch_geometric.testing import get_random_edge_index, onlyFullTest
 from torch_geometric.transforms import RandomLinkSplit
 from torch_geometric.utils import is_undirected, to_undirected
 
@@ -288,3 +288,32 @@ def test_random_link_split_non_contiguous():
     train_data, val_data, test_data = transform(data)
     assert train_data['p', 'p'].num_edges == 60
     assert train_data['p', 'p'].edge_index.is_contiguous()
+
+
+@onlyFullTest
+def test_random_link_split_on_dataset(get_dataset):
+    dataset = get_dataset(name='MUTAG')
+
+    dataset.transform = RandomLinkSplit(
+        num_val=0.1,
+        num_test=0.1,
+        disjoint_train_ratio=0.3,
+        add_negative_train_samples=False,
+    )
+
+    train_dataset, val_dataset, test_dataset = zip(*dataset)
+    assert len(train_dataset) == len(dataset)
+    assert len(val_dataset) == len(dataset)
+    assert len(test_dataset) == len(dataset)
+
+    assert isinstance(train_dataset[0], Data)
+    assert train_dataset[0].edge_label.min() == 1.0
+    assert train_dataset[0].edge_label.max() == 1.0
+
+    assert isinstance(val_dataset[0], Data)
+    assert val_dataset[0].edge_label.min() == 0.0
+    assert val_dataset[0].edge_label.max() == 1.0
+
+    assert isinstance(test_dataset[0], Data)
+    assert test_dataset[0].edge_label.min() == 0.0
+    assert test_dataset[0].edge_label.max() == 1.0