Repository: pyg-team/pytorch_geometric
Pull request #7416 (2023-05-23 13:54:32)

Re-factor `ClusterLoader` + Integrate `pyg-lib` metis computation

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7416?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7416](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7416?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4764342) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/e2c4b76a064762dea278daca9ccea66f71baaede?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (e2c4b76) will **increase** coverage by `0.01%`.
> The diff coverage is `95.28%`.

```diff
@@            Coverage Diff             @@
##           master    #7416      +/-   ##
==========================================
+ Coverage   91.32%   91.33%   +0.01%     
==========================================
  Files         440      440              
  Lines       24576    24631      +55     
==========================================
+ Hits        22443    22496      +53     
- Misses       2133     2135       +2     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7416?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/loader/cluster.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7416?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2xvYWRlci9jbHVzdGVyLnB5) | `96.42% <95.23%> (-0.09%)` | :arrow_down: |
| [torch\_geometric/typing.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7416?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3R5cGluZy5weQ==) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 46163ff28..4d396ae2b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Re-factored `ClusterLoader` to integrate `pyg-lib` METIS routine ([#7416](https://github.com/pyg-team/pytorch_geometric/pull/7416))
 - Fixed an index-out-of-range bug in `QuantileAggregation` when `dim_size` is passed ([#7407](https://github.com/pyg-team/pytorch_geometric/pull/7407))
 - The `filter_per_worker` option will not get automatically inferred by default based on the device of the underlying data ([#7399](https://github.com/pyg-team/pytorch_geometric/pull/7399))
 - Fixed a bug in `LightGCN.recommendation_loss()` to only use the embeddings of the nodes involved in the current mini-batch ([#7384](https://github.com/pyg-team/pytorch_geometric/pull/7384))
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 846c105de..dcba87187 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -589,12 +589,6 @@ class Data(BaseData, FeatureStore, GraphStore):
         Args:
             subset (LongTensor or BoolTensor): The nodes to keep.
         """
-        if subset.dtype == torch.bool:
-            num_nodes = int(subset.sum())
-        else:
-            num_nodes = subset.size(0)
-            subset = torch.unique(subset, sorted=True)
-
         out = subgraph(subset, self.edge_index, relabel_nodes=True,
                        num_nodes=self.num_nodes, return_edge_mask=True)
         edge_index, _, edge_mask = out
@@ -605,7 +599,10 @@ class Data(BaseData, FeatureStore, GraphStore):
             if key == 'edge_index':
                 data.edge_index = edge_index
             elif key == 'num_nodes':
-                data.num_nodes = num_nodes
+                if subset.dtype == torch.bool:
+                    data.num_nodes = int(subset.sum())
+                else:
+                    data.num_nodes = subset.size(0)
             elif self.is_node_attr(key):
                 cat_dim = self.__cat_dim__(key, value)
                 data[key] = select(value, subset, dim=cat_dim)
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index 99c78b0f5..da19896d0 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -638,17 +638,12 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
         subset_dict = copy.copy(subset_dict)
 
         for node_type, subset in subset_dict.items():
-
-            if subset.dtype == torch.bool:
-                num_nodes = int(subset.sum())
-            else:
-                num_nodes = subset.size(0)
-                subset = torch.unique(subset, sorted=True)
-                subset_dict[node_type] = subset
-
             for key, value in self[node_type].items():
                 if key == 'num_nodes':
-                    data[node_type].num_nodes = num_nodes
+                    if subset.dtype == torch.bool:
+                        data[node_type].num_nodes = int(subset.sum())
+                    else:
+                        data[node_type].num_nodes = subset.size(0)
                 elif self[node_type].is_node_attr(key):
                     data[node_type][key] = value[subset]
                 else:
diff --git a/torch_geometric/loader/cluster.py b/torch_geometric/loader/cluster.py
index 23824f174..4b3cbb51a 100644
--- a/torch_geometric/loader/cluster.py
+++ b/torch_geometric/loader/cluster.py
@@ -1,13 +1,28 @@
 import copy
 import os.path as osp
 import sys
-from typing import Optional
+from dataclasses import dataclass
+from typing import List, Optional
 
 import torch
 import torch.utils.data
+from torch import Tensor
 
-from torch_geometric.typing import SparseTensor, torch_sparse
-from torch_geometric.utils import narrow, select
+import torch_geometric.typing
+from torch_geometric.data import Data
+from torch_geometric.typing import pyg_lib
+from torch_geometric.utils import index_sort, narrow, select, sort_edge_index
+from torch_geometric.utils.map import map_index
+from torch_geometric.utils.sparse import index2ptr, ptr2index
+
+
+@dataclass
+class Partition:
+    rowptr: Tensor
+    col: Tensor
+    partptr: Tensor
+    node_perm: Tensor
+    edge_perm: Tensor
 
 
 class ClusterData(torch.utils.data.Dataset):
@@ -44,82 +59,139 @@ class ClusterData(torch.utils.data.Dataset):
         assert data.edge_index is not None
 
         self.num_parts = num_parts
+        self.recursive = recursive
         self.keep_inter_cluster_edges = keep_inter_cluster_edges
 
         recursive_str = '_recursive' if recursive else ''
-        filename = f'partition_{num_parts}{recursive_str}.pt'
+        filename = f'metis_{num_parts}{recursive_str}.pt'
         path = osp.join(save_dir or '', filename)
         if save_dir is not None and osp.exists(path):
-            adj, partptr, perm = torch.load(path)
+            self.partition = torch.load(path)
         else:
             if log:  # pragma: no cover
                 print('Computing METIS partitioning...', file=sys.stderr)
 
-            N, E = data.num_nodes, data.num_edges
-            adj = SparseTensor(
-                row=data.edge_index[0], col=data.edge_index[1],
-                value=torch.arange(E, device=data.edge_index.device),
-                sparse_sizes=(N, N))
-            adj, partptr, perm = adj.partition(num_parts, recursive)
+            cluster = self._metis(data.edge_index, data.num_nodes)
+            self.partition = self._partition(data.edge_index, cluster)
 
             if save_dir is not None:
-                torch.save((adj, partptr, perm), path)
+                torch.save(self.partition, path)
 
             if log:  # pragma: no cover
                 print('Done!', file=sys.stderr)
 
-        self.data = self._permute_data(data, perm, adj)
-        self.partptr = partptr
-        self.perm = perm
-
-    def _permute_data(self, data, node_idx, adj):
+        self.data = self._permute_data(data, self.partition)
+
+    def _metis(self, edge_index: Tensor, num_nodes: int) -> Tensor:
+        # Computes a node-level partition assignment vector via METIS.
+
+        # Calculate CSR representation:
+        row, col = sort_edge_index(edge_index, num_nodes=num_nodes)
+        rowptr = index2ptr(row, size=num_nodes)
+
+        # Compute METIS partitioning:
+        if torch_geometric.typing.WITH_METIS:
+            return pyg_lib.partition.metis(
+                rowptr.cpu(),
+                col.cpu(),
+                self.num_parts,
+                recursive=self.recursive,
+            ).to(edge_index.device)
+
+        if torch_geometric.typing.WITH_TORCH_SPARSE:
+            return torch.ops.torch_sparse.partition(
+                rowptr.cpu(),
+                col.cpu(),
+                None,
+                self.num_parts,
+                self.recursive,
+            ).to(edge_index.device)
+
+        raise ImportError(f"'{self.__class__.__name__}' requires either "
+                          f"'pyg-lib' or 'torch-sparse'")
+
+    def _partition(self, edge_index: Tensor, cluster: Tensor) -> Partition:
+        # Computes node-level and edge-level permutations and permutes the edge
+        # connectivity accordingly:
+
+        # Sort `cluster` and compute boundaries `partptr`:
+        cluster, node_perm = index_sort(cluster, max_value=self.num_parts)
+        partptr = index2ptr(cluster, size=self.num_parts)
+
+        # Permute `edge_index` based on node permutation:
+        edge_perm = torch.arange(edge_index.size(1), device=edge_index.device)
+        arange = torch.empty_like(node_perm)
+        arange[node_perm] = torch.arange(cluster.numel(),
+                                         device=cluster.device)
+        edge_index = arange[edge_index]
+
+        # Compute final CSR representation:
+        (row, col), edge_perm = sort_edge_index(
+            edge_index,
+            edge_attr=edge_perm,
+            num_nodes=cluster.numel(),
+        )
+        rowptr = index2ptr(row, size=cluster.numel())
+
+        return Partition(rowptr, col, partptr, node_perm, edge_perm)
+
+    def _permute_data(self, data: Data, partition: Partition) -> Data:
+        # Permute node-level and edge-level attributes according to the
+        # calculated permutations in `Partition`:
         out = copy.copy(data)
         for key, value in data.items():
-            if data.is_node_attr(key):
+            if key == 'edge_index':
+                continue
+            elif data.is_node_attr(key):
                 cat_dim = data.__cat_dim__(key, value)
-                out[key] = select(value, node_idx, dim=cat_dim)
-
+                out[key] = select(value, partition.node_perm, dim=cat_dim)
+            elif data.is_edge_attr(key):
+                cat_dim = data.__cat_dim__(key, value)
+                out[key] = select(value, partition.edge_perm, dim=cat_dim)
         out.edge_index = None
-        out.adj = adj
 
         return out
 
-    def __len__(self):
-        return self.partptr.numel() - 1
-
-    def __getitem__(self, idx):
-        start = int(self.partptr[idx])
-        length = int(self.partptr[idx + 1]) - start
-
-        data = copy.copy(self.data)
-        adj, data.adj = data.adj, None
-
-        adj = adj.narrow(0, start, length)
+    def __len__(self) -> int:
+        return self.partition.partptr.numel() - 1
+
+    def __getitem__(self, idx: int) -> Data:
+        node_start = int(self.partition.partptr[idx])
+        node_end = int(self.partition.partptr[idx + 1])
+        node_length = node_end - node_start
+
+        rowptr = self.partition.rowptr[node_start:node_end + 1]
+        edge_start = int(rowptr[0])
+        edge_end = int(rowptr[-1])
+        edge_length = edge_end - edge_start
+        rowptr = rowptr - edge_start
+        row = ptr2index(rowptr)
+        col = self.partition.col[edge_start:edge_end]
         if not self.keep_inter_cluster_edges:
-            adj = adj.narrow(1, start, length)
+            edge_mask = (col >= node_start) & (col < node_end)
+            row = row[edge_mask]
+            col = col[edge_mask] - node_start
 
-        edge_idx = adj.storage.value()
+        out = copy.copy(self.data)
 
-        for key, value in data:
+        for key, value in self.data.items():
             if key == 'num_nodes':
-                data.num_nodes = length
+                out.num_nodes = node_length
             elif self.data.is_node_attr(key):
                 cat_dim = self.data.__cat_dim__(key, value)
-                data[key] = narrow(value, cat_dim, start, length)
+                out[key] = narrow(value, cat_dim, node_start, node_length)
             elif self.data.is_edge_attr(key):
                 cat_dim = self.data.__cat_dim__(key, value)
-                data[key] = select(value, edge_idx, dim=cat_dim)
+                out[key] = narrow(value, cat_dim, edge_start, edge_length)
+                if not self.keep_inter_cluster_edges:
+                    out[key] = out[key][edge_mask]
 
-        row, col, _ = adj.coo()
-        data.edge_index = torch.stack([row, col], dim=0)
+        out.edge_index = torch.stack([row, col], dim=0)
 
-        return data
+        return out
 
-    def __repr__(self):
-        return (f'{self.__class__.__name__}(\n'
-                f'  data={self.data},\n'
-                f'  num_parts={self.num_parts}\n'
-                f')')
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.num_parts})'
 
 
 class ClusterLoader(torch.utils.data.DataLoader):
@@ -149,36 +221,67 @@ class ClusterLoader(torch.utils.data.DataLoader):
     """
     def __init__(self, cluster_data, **kwargs):
         self.cluster_data = cluster_data
+        iterator = range(len(cluster_data))
+        super().__init__(iterator, collate_fn=self._collate, **kwargs)
 
-        super().__init__(range(len(cluster_data)), collate_fn=self._collate,
-                         **kwargs)
-
-    def _collate(self, batch):
+    def _collate(self, batch: List[int]) -> Data:
         if not isinstance(batch, torch.Tensor):
             batch = torch.tensor(batch)
 
-        start = self.cluster_data.partptr[batch].tolist()
-        end = self.cluster_data.partptr[batch + 1].tolist()
-        node_idx = torch.cat([torch.arange(s, e) for s, e in zip(start, end)])
-
-        data = copy.copy(self.cluster_data.data)
-
-        adj, data.adj = self.cluster_data.data.adj, None
-        adj = torch_sparse.cat(
-            [adj.narrow(0, s, e - s) for s, e in zip(start, end)], dim=0)
-        adj = adj.index_select(1, node_idx)
-        row, col, edge_idx = adj.coo()
-
-        for key, value in data:
+        global_rowptr = self.cluster_data.partition.rowptr
+        global_col = self.cluster_data.partition.col
+
+        # Get all node-level and edge-level start and end indices for the
+        # current mini-batch:
+        node_start = self.cluster_data.partition.partptr[batch]
+        node_end = self.cluster_data.partition.partptr[batch + 1]
+        edge_start = global_rowptr[node_start]
+        edge_end = global_rowptr[node_end]
+
+        # Iterate over each partition in the batch and calculate new edge
+        # connectivity. This is done by slicing the corresponding source and
+        # destination indices for each partition and adjusting their indices to
+        # start from zero:
+        rows, cols, nodes, cumsum = [], [], [], 0
+        for i in range(batch.numel()):
+            nodes.append(torch.arange(node_start[i], node_end[i]))
+            rowptr = global_rowptr[node_start[i]:node_end[i] + 1]
+            rowptr = rowptr - edge_start[i]
+            row = ptr2index(rowptr) + cumsum
+            col = global_col[edge_start[i]:edge_end[i]]
+            rows.append(row)
+            cols.append(col)
+            cumsum += rowptr.numel() - 1
+
+        node = torch.cat(nodes, dim=0)
+        row = torch.cat(rows, dim=0)
+        col = torch.cat(cols, dim=0)
+
+        # Map `col` vector to valid entries and remove any entries that do not
+        # connect two nodes within the same mini-batch:
+        col, edge_mask = map_index(col, node)
+        row = row[edge_mask]
+
+        out = copy.copy(self.cluster_data.data)
+
+        # Slice node-level and edge-level attributes according to its offsets:
+        for key, value in self.cluster_data.data.items():
             if key == 'num_nodes':
-                data.num_nodes = node_idx.numel()
+                out.num_nodes = cumsum
             elif self.cluster_data.data.is_node_attr(key):
                 cat_dim = self.cluster_data.data.__cat_dim__(key, value)
-                data[key] = select(value, node_idx, dim=cat_dim)
+                out[key] = torch.cat([
+                    narrow(out[key], cat_dim, s, e - s)
+                    for s, e in zip(node_start, node_end)
+                ], dim=cat_dim)
             elif self.cluster_data.data.is_edge_attr(key):
                 cat_dim = self.cluster_data.data.__cat_dim__(key, value)
-                data[key] = select(value, edge_idx, dim=cat_dim)
+                value = torch.cat([
+                    narrow(out[key], cat_dim, s, e - s)
+                    for s, e in zip(edge_start, edge_end)
+                ], dim=cat_dim)
+                out[key] = select(value, edge_mask, dim=cat_dim)
 
-        data.edge_index = torch.stack([row, col], dim=0)
+        out.edge_index = torch.stack([row, col], dim=0)
 
-        return data
+        return out
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index d25605ea7..3fd739f2c 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -13,6 +13,7 @@ try:
     WITH_GMM = WITH_PT2 and hasattr(pyg_lib.ops, 'grouped_matmul')
     WITH_SAMPLED_OP = hasattr(pyg_lib.ops, 'sampled_add')
     WITH_INDEX_SORT = hasattr(pyg_lib.ops, 'index_sort')
+    WITH_METIS = hasattr(pyg_lib, 'partition')
 except (ImportError, OSError) as e:
     if isinstance(e, OSError):
         warnings.warn(f"An issue occurred while importing 'pyg-lib'. "
@@ -22,6 +23,7 @@ except (ImportError, OSError) as e:
     WITH_GMM = False
     WITH_SAMPLED_OP = False
     WITH_INDEX_SORT = False
+    WITH_METIS = False
 
 try:
     import torch_scatter  # noqa
diff --git a/torch_geometric/utils/map.py b/torch_geometric/utils/map.py
new file mode 100644
index 000000000..cea56b0ed
--- /dev/null
+++ b/torch_geometric/utils/map.py
@@ -0,0 +1,50 @@
+from typing import Tuple
+
+import torch
+from torch import Tensor
+
+
+def map_index(src: Tensor, index: Tensor) -> Tuple[Tensor, Tensor]:
+    r"""Maps indices in :obj:`src` to the positional value of their
+    corresponding occurence in :obj:`index`.
+
+    Args:
+        src (torch.Tensor): The source tensor to map.
+        index (torch.Tensor): The index tensor that denotes the new mapping.
+
+    :rtype: (:class:`torch.Tensor`, :class:`torch.BoolTensor`)
+
+    Examples:
+
+        >>> src = torch.tensor([2, 0, 1, 0, 3])
+        >>> index = torch.tensor([3, 2, 0, 1])
+
+        >>> map_index(src, index)
+        (tensor([1, 2, 3, 2, 0]), tensor([True, True, True, True, True]))
+
+        >>> src = torch.tensor([2, 0, 1, 0, 3])
+        >>> index = torch.tensor([3, 2, 0])
+
+        >>> map_index(src, index)
+        (tensor([1, 2, 2, 0]), tensor([True, True, False, True, True]))
+    """
+    import pandas as pd
+
+    assert src.dim() == 1 and index.dim() == 1
+    assert not src.is_floating_point()
+    assert not index.is_floating_point()
+
+    arange = pd.RangeIndex(0, index.size(0))
+    df = pd.DataFrame(index=index.detach().cpu().numpy(), data={'out': arange})
+    ser = pd.Series(src.detach().cpu(), name='key')
+    result = df.merge(ser, how='right', left_index=True, right_on='key')
+    out = torch.from_numpy(result['out'].values).to(index.device)
+
+    if out.is_floating_point():
+        mask = torch.isnan(out).logical_not_()
+        out = out[mask].to(index.dtype)
+        return out, mask
+
+    out = out.to(index.dtype)
+    mask = torch.ones_like(out, dtype=torch.bool)
+    return out, mask


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index 5badd09eb..b38578c7f 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -208,12 +208,12 @@ def test_data_subgraph():
     assert torch.equal(out.edge_weight, edge_weight[torch.arange(2, 6)])
     assert out.num_nodes == 3
 
-    # test for unordered selection
+    # Test unordered selection:
     out = data.subgraph(torch.tensor([3, 1, 2]))
     assert len(out) == 5
-    assert torch.equal(out.x, torch.arange(1, 4))
+    assert torch.equal(out.x, torch.tensor([3, 1, 2]))
     assert torch.equal(out.y, data.y)
-    assert out.edge_index.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]]
+    assert out.edge_index.tolist() == [[1, 2, 2, 0], [2, 1, 0, 2]]
     assert torch.equal(out.edge_weight, edge_weight[torch.arange(2, 6)])
     assert out.num_nodes == 3
 
diff --git a/test/data/test_hetero_data.py b/test/data/test_hetero_data.py
index ddb3cd32b..5698ce7a0 100644
--- a/test/data/test_hetero_data.py
+++ b/test/data/test_hetero_data.py
@@ -214,8 +214,6 @@ def test_hetero_data_subgraph():
         'conf': torch.randperm(x_conference.size(0))[:2],
     }
 
-    subset_sorted = {key: torch.sort(idx)[0] for key, idx in subset.items()}
-
     out = data.subgraph(subset)
     out.validate(raise_on_error=True)
 
@@ -224,7 +222,7 @@ def test_hetero_data_subgraph():
 
     for key in out.node_types:
         assert len(out[key]) == len(data[key])
-        assert torch.allclose(out[key].x, data[key].x[subset_sorted[key]])
+        assert torch.allclose(out[key].x, data[key].x[subset[key]])
         assert out[key].num_nodes == subset[key].size(0)
         if key == 'paper':
             assert out['paper'].name == 'paper'
@@ -235,8 +233,8 @@ def test_hetero_data_subgraph():
     for key in out.node_types:
         node_mask[key] = torch.zeros((data[key].num_nodes, ), dtype=torch.bool)
         node_map[key] = torch.zeros((data[key].num_nodes, ), dtype=torch.long)
-        node_mask[key][subset_sorted[key]] = True
-        node_map[key][subset_sorted[key]] = torch.arange(subset[key].size(0))
+        node_mask[key][subset[key]] = True
+        node_map[key][subset[key]] = torch.arange(subset[key].size(0))
 
     edge_mask = {}  # for each edge type a mask of edges in the subgraph
     subgraph_edge_index = {
diff --git a/test/loader/test_cluster.py b/test/loader/test_cluster.py
index 884eac42e..c934eb52e 100644
--- a/test/loader/test_cluster.py
+++ b/test/loader/test_cluster.py
@@ -1,21 +1,22 @@
 import pytest
 import torch
 
+import torch_geometric.typing
 from torch_geometric.data import Data
 from torch_geometric.loader import ClusterData, ClusterLoader
 from torch_geometric.testing import onlyFullTest
-from torch_geometric.utils import to_dense_adj
+from torch_geometric.utils import sort_edge_index
 
 try:
     rowptr = torch.tensor([0, 1])
     col = torch.tensor([0])
     torch.ops.torch_sparse.partition(rowptr, col, None, 1, True)
-    with_metis = True
+    WITH_METIS = True
 except (AttributeError, RuntimeError):
-    with_metis = False
+    WITH_METIS = False or torch_geometric.typing.WITH_METIS
 
 
-@pytest.mark.skipif(not with_metis, reason='Not compiled with METIS support')
+@pytest.mark.skipif(not WITH_METIS, reason='Not compiled with METIS support')
 def test_cluster_gcn():
     adj = torch.tensor([
         [1, 1, 1, 0, 1, 0],
@@ -29,81 +30,86 @@ def test_cluster_gcn():
     x = torch.Tensor([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
     edge_index = adj.nonzero(as_tuple=False).t()
     edge_attr = torch.arange(edge_index.size(1))
-    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
+    n_id = torch.arange(6)
+    data = Data(x=x, n_id=n_id, edge_index=edge_index, edge_attr=edge_attr)
     data.num_nodes = 6
 
     cluster_data = ClusterData(data, num_parts=2, log=False)
 
-    assert cluster_data.partptr.tolist() == [0, 3, 6]
-    assert cluster_data.perm.tolist() == [0, 2, 4, 1, 3, 5]
-    assert cluster_data.data.x.tolist() == [
-        [0, 0],
-        [2, 2],
-        [4, 4],
-        [1, 1],
-        [3, 3],
-        [5, 5],
-    ]
-    assert cluster_data.data.adj.to_dense().tolist() == [
-        [0, 2, 3, 1, 0, 0],
-        [8, 9, 10, 0, 0, 0],
-        [14, 15, 16, 0, 0, 0],
-        [4, 0, 0, 5, 6, 7],
-        [0, 0, 0, 11, 12, 13],
-        [0, 0, 0, 17, 18, 19],
+    partition = cluster_data._partition(
+        edge_index, cluster=torch.tensor([0, 1, 0, 1, 0, 1]))
+    assert partition.partptr.tolist() == [0, 3, 6]
+    assert partition.node_perm.tolist() == [0, 2, 4, 1, 3, 5]
+    assert partition.edge_perm.tolist() == [
+        0, 2, 3, 1, 8, 9, 10, 14, 15, 16, 4, 5, 6, 7, 11, 12, 13, 17, 18, 19
     ]
 
-    data = cluster_data[0]
-    assert data.num_nodes == 3
-    assert data.x.tolist() == [[0, 0], [2, 2], [4, 4]]
-    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
-                                        [0, 1, 2, 0, 1, 2, 0, 1, 2]]
-    assert data.edge_attr.tolist() == [0, 2, 3, 8, 9, 10, 14, 15, 16]
+    assert cluster_data.partition.partptr.tolist() == [0, 3, 6]
+    assert torch.equal(
+        cluster_data.partition.node_perm.sort()[0],
+        torch.arange(data.num_nodes),
+    )
+    assert torch.equal(
+        cluster_data.partition.edge_perm.sort()[0],
+        torch.arange(data.num_edges),
+    )
 
-    data = cluster_data[1]
-    assert data.num_nodes == 3
-    assert data.x.tolist() == [[1, 1], [3, 3], [5, 5]]
-    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
-                                        [0, 1, 2, 0, 1, 2, 0, 1, 2]]
-    assert data.edge_attr.tolist() == [5, 6, 7, 11, 12, 13, 17, 18, 19]
+    out = cluster_data[0]
+    expected = data.subgraph(out.n_id)
+    out.validate()
+    assert out.num_nodes == 3
+    assert out.n_id.size() == (3, )
+    assert torch.equal(out.x, expected.x)
+    tmp = sort_edge_index(expected.edge_index, expected.edge_attr)
+    assert torch.equal(out.edge_index, tmp[0])
+    assert torch.equal(out.edge_attr, tmp[1])
+
+    out = cluster_data[1]
+    out.validate()
+    assert out.num_nodes == 3
+    assert out.n_id.size() == (3, )
+    expected = data.subgraph(out.n_id)
+    assert torch.equal(out.x, expected.x)
+    tmp = sort_edge_index(expected.edge_index, expected.edge_attr)
+    assert torch.equal(out.edge_index, tmp[0])
+    assert torch.equal(out.edge_attr, tmp[1])
 
     loader = ClusterLoader(cluster_data, batch_size=1)
     iterator = iter(loader)
 
-    data = next(iterator)
-    assert data.x.tolist() == [[0, 0], [2, 2], [4, 4]]
-    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
-                                        [0, 1, 2, 0, 1, 2, 0, 1, 2]]
-    assert data.edge_attr.tolist() == [0, 2, 3, 8, 9, 10, 14, 15, 16]
-
-    data = next(iterator)
-    assert data.x.tolist() == [[1, 1], [3, 3], [5, 5]]
-    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2],
-                                        [0, 1, 2, 0, 1, 2, 0, 1, 2]]
-    assert data.edge_attr.tolist() == [5, 6, 7, 11, 12, 13, 17, 18, 19]
+    out = next(iterator)
+    out.validate()
+    assert out.num_nodes == 3
+    assert out.n_id.size() == (3, )
+    expected = data.subgraph(out.n_id)
+    assert torch.equal(out.x, expected.x)
+    tmp = sort_edge_index(expected.edge_index, expected.edge_attr)
+    assert torch.equal(out.edge_index, tmp[0])
+    assert torch.equal(out.edge_attr, tmp[1])
+
+    out = next(iterator)
+    out.validate()
+    assert out.num_nodes == 3
+    assert out.n_id.size() == (3, )
+    expected = data.subgraph(out.n_id)
+    assert torch.equal(out.x, expected.x)
+    tmp = sort_edge_index(expected.edge_index, expected.edge_attr)
+    assert torch.equal(out.edge_index, tmp[0])
+    assert torch.equal(out.edge_attr, tmp[1])
 
     loader = ClusterLoader(cluster_data, batch_size=2, shuffle=False)
-    data = next(iter(loader))
-    assert data.num_nodes == 6
-    assert data.x.tolist() == [
-        [0, 0],
-        [2, 2],
-        [4, 4],
-        [1, 1],
-        [3, 3],
-        [5, 5],
-    ]
-    assert to_dense_adj(data.edge_index).squeeze().tolist() == [
-        [1, 1, 1, 1, 0, 0],
-        [1, 1, 1, 0, 0, 0],
-        [1, 1, 1, 0, 0, 0],
-        [1, 0, 0, 1, 1, 1],
-        [0, 0, 0, 1, 1, 1],
-        [0, 0, 0, 1, 1, 1],
-    ]
-
-
-@pytest.mark.skipif(not with_metis, reason='Not compiled with METIS support')
+    out = next(iter(loader))
+    out.validate()
+    assert out.num_nodes == 6
+    assert out.n_id.size() == (6, )
+    expected = data.subgraph(out.n_id)
+    assert torch.equal(out.x, expected.x)
+    tmp = sort_edge_index(expected.edge_index, expected.edge_attr)
+    assert torch.equal(out.edge_index, tmp[0])
+    assert torch.equal(out.edge_attr, tmp[1])
+
+
+@pytest.mark.skipif(not WITH_METIS, reason='Not compiled with METIS support')
 def test_keep_inter_cluster_edges():
     adj = torch.tensor([
         [1, 1, 1, 0, 1, 0],
@@ -124,29 +130,39 @@ def test_keep_inter_cluster_edges():
                                keep_inter_cluster_edges=True)
 
     data = cluster_data[0]
-    assert data.edge_index.tolist() == [[0, 0, 0, 0, 1, 1, 1, 2, 2, 2],
-                                        [0, 1, 2, 3, 0, 1, 2, 0, 1, 2]]
-    assert data.edge_attr.tolist() == [0, 2, 3, 1, 8, 9, 10, 14, 15, 16]
+    assert data.edge_index[0].min() == 0
+    assert data.edge_index[0].max() == 2
+    assert data.edge_index[1].min() == 0
+    assert data.edge_index[1].max() > 2
+    assert data.edge_index.size(1) == data.edge_attr.size(0)
 
     data = cluster_data[1]
-    assert data.edge_index.tolist() == [[0, 0, 0, 0, 1, 1, 1, 2, 2, 2],
-                                        [0, 3, 4, 5, 3, 4, 5, 3, 4, 5]]
-    assert data.edge_attr.tolist() == [4, 5, 6, 7, 11, 12, 13, 17, 18, 19]
+    assert data.edge_index[0].min() == 0
+    assert data.edge_index[0].max() == 2
+    assert data.edge_index[1].min() == 0
+    assert data.edge_index[1].max() > 2
+    assert data.edge_index.size(1) == data.edge_attr.size(0)
 
 
 @onlyFullTest
-@pytest.mark.skipif(not with_metis, reason='Not compiled with METIS support')
+@pytest.mark.skipif(not WITH_METIS, reason='Not compiled with METIS support')
 def test_cluster_gcn_correctness(get_dataset):
     dataset = get_dataset('Cora')
     data = dataset[0].clone()
     data.n_id = torch.arange(data.num_nodes)
-    cluster_data = ClusterData(data, num_parts=10)
+    cluster_data = ClusterData(data, num_parts=10, log=False)
     loader = ClusterLoader(cluster_data, batch_size=3, shuffle=False)
 
     for batch1 in loader:
+        batch1.validate()
         batch2 = data.subgraph(batch1.n_id)
         assert batch1.num_nodes == batch2.num_nodes
         assert batch1.num_edges == batch2.num_edges
+        assert torch.equal(batch1.x, batch2.x)
+        assert torch.equal(
+            batch1.edge_index,
+            sort_edge_index(batch2.edge_index),
+        )
 
 
 if __name__ == '__main__':
diff --git a/test/nn/conv/test_rgcn_conv.py b/test/nn/conv/test_rgcn_conv.py
index 791be8cec..1df7aaf65 100644
--- a/test/nn/conv/test_rgcn_conv.py
+++ b/test/nn/conv/test_rgcn_conv.py
@@ -83,7 +83,7 @@ def test_rgcn_conv(cls, conf, device):
     if is_full_test():
         t = '(OptTensor, Tensor, OptTensor) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert torch.allclose(jit(x1, edge_index, edge_type), out1)
+        assert torch.allclose(jit(x1, edge_index, edge_type), out1, atol=1e-3)
         if num_blocks is None:
             assert torch.allclose(jit(idx1, edge_index, edge_type), out2,
                                   atol=1e-3)
diff --git a/test/utils/test_map.py b/test/utils/test_map.py
new file mode 100644
index 000000000..923bb5b81
--- /dev/null
+++ b/test/utils/test_map.py
@@ -0,0 +1,21 @@
+import torch
+
+from torch_geometric.utils.map import map_index
+
+
+def test_map_index():
+    src = torch.tensor([2, 0, 1, 0, 3])
+    index = torch.tensor([3, 2, 0, 1])
+
+    out, mask = map_index(src, index)
+    assert out.tolist() == [1, 2, 3, 2, 0]
+    assert mask.tolist() == [True, True, True, True, True]
+
+
+def test_map_index_na():
+    src = torch.tensor([2, 0, 1, 0, 3])
+    index = torch.tensor([3, 2, 0])
+
+    out, mask = map_index(src, index)
+    assert out.tolist() == [1, 2, 2, 0]
+    assert mask.tolist() == [True, True, False, True, True]