Repository: pyg-team/pytorch_geometric
Pull request #7611 (2023-06-19 15:46:38)

Hypergraph data object
Hello. I don't know the definition of a subgraph for a hypergraph (not mentioned in papers). When getting a subgraph, a hyperedge can be cut. For example, if a hyperedge connects vertices [1, 2, 3], but only [1, 3] are selected for the subgraph, then we have the option to: 1) remove this hyperedge; 2) change the hyperedge and make it only connects [1, 3]. So, in this version, I choose the second option.

CodeTal: > Thanks for the great work @CodeTal. Will look at your code in detail and leave more comments. For now, could you add tests for the `HyperGraphData` and all the new methods. Also add docs.

Sure.
CodeTal: Hi. I found that the modified `HyperGraphData` class should include a different `to_heterogeneous` function. Since the format of `edge_index` in HyperGraphData has changed (from (2, edge_num) to (2, edge_num*nodes_per_edge)), then the format of `edge_type` should change. Do you need to implement this function in this version? If so, I was thinking about how to express the `edge_type` of `HyperGraphData`. We can either assign each node-edge pair an `edge_type` or just assign the unique edge numbers an `edge_type`.
wsad1: Sorry for the late reply @CodeTal. For this PR we can skip the implementation of `to_heterogenous` and raise a `NotImplemntedError`. Lets implement this in a follow up PR, since reviewing smaller PRs is easier.
CodeTal: Hello! In this PR, I have added test cases to test `subgraph()` and `__inc__()`. I would like to add some docs to these code pieces, but I am not sure what are the docs. Do you mean a post on the [PyG documentation](https://pytorch-geometric.readthedocs.io/en/latest/), or a README file under the directory of the source code?
CodeTal: @wsad1 Thanks for your careful review and advice! I have added documentation comments, as well as fixing the old code. 
wsad1: @CodeTal I updated the the tests for `subgraph`  by adding `edge_attr`. This is when I realised there is an inconsistency in the way we define edge attributes. For hyper edges should an edge attribute be of shape `[num_edges, num_attrs]` ie. a single attribute per hyper edge OR of shape `[num_edges*num_nodes_per_edge, num_attrs]` ie. an attribute for each connection of a hyper edge. I think it should be the latter. Given that `hyper_subgraph` should return a masks of shape `num_edges` to filter out, and not a mask of shape  ` num_edges*num_nodes_per_edge`.
I'll try to make this update today. But if I am not able to please go ahead with it.
CodeTal: > @CodeTal thanks for the great work.

And thank you for your careful review and enlightening update! Sorry for not following the issue last week, I was busy with other things on last week and had little spare time.
mehular0ra: @CodeTal  I really liked your implementation. Just a question I had. Will this Hypergraph data object support multimodality? By that what I mean is: if let's say Incidence Matrix has a size of not [n,n], but [n, m] where m = k*n. k=1, 2, 3,...
That is, will your implementation support when there more number of hyperedges than the nodes?

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f469324a1..1fbfe0e1b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `HyperGraphData` to support hypergraphs ([#7611](https://github.com/pyg-team/pytorch_geometric/pull/7611))
 - Added the `PCQM4Mv2` dataset as a reference implementation for `OnDiskDataset` ([#8102](https://github.com/pyg-team/pytorch_geometric/pull/8102)
 - Added `module_headers` property to `nn.Sequential` models ([#8093](https://github.com/pyg-team/pytorch_geometric/pull/8093)
 - Added `OnDiskDataset` interface with data loader support ([#8066](https://github.com/pyg-team/pytorch_geometric/pull/8066), [#8088](https://github.com/pyg-team/pytorch_geometric/pull/8088), [#8092](https://github.com/pyg-team/pytorch_geometric/pull/8092), [#8106](https://github.com/pyg-team/pytorch_geometric/pull/8106))
diff --git a/torch_geometric/data/hypergraph_data.py b/torch_geometric/data/hypergraph_data.py
new file mode 100644
index 000000000..100ae37b6
--- /dev/null
+++ b/torch_geometric/data/hypergraph_data.py
@@ -0,0 +1,206 @@
+import copy
+import warnings
+from typing import Any, List, Optional
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import Data
+from torch_geometric.typing import EdgeType, NodeType, OptTensor
+from torch_geometric.utils import hyper_subgraph, select
+
+
+class HyperGraphData(Data):
+    r"""A data object describing a hypergraph.
+    The data object can hold node-level, link-level and graph-level attributes.
+    This object differs from a standard :obj:`~torch_geometric.data.Data`
+    object by having hyperedges, i.e. edges that connect more
+    than two nodes. For example, in the hypergraph scenario
+    :math:`\mathcal{G} = (\mathcal{V}, \mathcal{E})` with
+    :math:`\mathcal{V} = \{ 0, 1, 2, 3, 4 \}` and
+    :math:`\mathcal{E} = \{ \{ 0, 1, 2 \}, \{ 1, 2, 3, 4 \} \}`, the
+    hyperedge index :obj:`edge_index` is represented as:
+
+    .. code-block:: python
+
+        # hyper graph with two hyperedges
+        # connecting 3 and 4 nodes, respectively
+        edge_index = torch.tensor([
+            [0, 1, 2, 1, 2, 3, 4],
+            [0, 0, 0, 1, 1, 1, 1],
+        ])
+
+    Args:
+        x (torch.Tensor, optional): Node feature matrix with shape
+            :obj:`[num_nodes, num_node_features]`. (default: :obj:`None`)
+        edge_index (LongTensor, optional): Hyperedge tensor
+            with shape :obj:`[2, num_edges*num_nodes_per_edge]`.
+            Where `edge_index[1]` denotes the hyperedge index and
+            `edge_index[0]` denotes the node indicies that are connected
+            by the hyperedge. (default: :obj:`None`)
+            (default: :obj:`None`)
+        edge_attr (torch.Tensor, optional): Edge feature matrix with shape
+            :obj:`[num_edges, num_edge_features]`.
+            (default: :obj:`None`)
+        y (torch.Tensor, optional): Graph-level or node-level ground-truth
+            labels with arbitrary shape. (default: :obj:`None`)
+        pos (torch.Tensor, optional): Node position matrix with shape
+            :obj:`[num_nodes, num_dimensions]`. (default: :obj:`None`)
+        **kwargs (optional): Additional attributes.
+    """
+    def __init__(self, x: OptTensor = None, edge_index: OptTensor = None,
+                 edge_attr: OptTensor = None, y: OptTensor = None,
+                 pos: OptTensor = None, **kwargs):
+        super().__init__(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y,
+                         pos=pos, **kwargs)
+
+    @property
+    def num_edges(self) -> int:
+        r"""Returns the number of hyperedges in the hypergraph.
+        """
+        if self.edge_index is None:
+            return 0
+        return max(self.edge_index[1]) + 1
+
+    @property
+    def num_nodes(self) -> int:
+        num_nodes = super().num_nodes
+
+        # For a hyper graph, the `edge_index[1]`
+        # does not contain node indicies. Therefore,
+        # the below code is to prevent the `num_nodes`
+        # being estimated as the number of hyperedges.
+        if (self.edge_index is not None and num_nodes == self.num_edges):
+            return max(self.edge_index[0]) + 1
+        return num_nodes
+
+    def is_edge_attr(self, key: str) -> bool:
+        val = super().is_edge_attr(key)
+        if not val and self.edge_index is not None:
+            return key in self and self[key].size(0) == self.num_edges
+
+    def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
+        if key == 'edge_index':
+            return torch.tensor([[self.num_nodes], [self.num_edges]])
+        else:
+            return super().__inc__(key, value, *args, **kwargs)
+
+    def subgraph(self, subset: Tensor) -> 'HyperGraphData':
+        r"""Returns the induced subgraph given by the node indices
+        :obj:`subset`.
+
+        .. note::
+
+            If only a subset of a hyperedge's nodes are to be
+            selected in the subgraph, the hyperedge will remain in the
+            subgraph, but only the selected nodes will be connected by
+            the hyperedge. Hyperedges that only connects one node in the
+            subgraph will be removed.
+
+        Examples:
+            >>> x = torch.randn(4, 16)
+            >>> edge_index = torch.tensor([
+            ...     [0, 1, 0, 2, 1, 1, 2, 4],
+            ...     [0, 0, 1, 1, 1, 2, 2, 2]
+            >>> ])
+            >>> data = HyperGraphData(x = x, edge_index = edge_index)
+            >>> subset = torch.tensor([1, 2, 4])
+            >>> subgraph = data.subgraph(subset)
+            >>> subgraph.edge_index
+            tensor([[2, 1, 1, 2, 4],
+            [0, 0, 1, 1, 1]])
+
+        Args:
+            subset (LongTensor or BoolTensor): The nodes to keep.
+        """
+        out = hyper_subgraph(subset, self.edge_index, relabel_nodes=True,
+                             num_nodes=self.num_nodes, return_edge_mask=True)
+        edge_index, _, edge_mask = out
+
+        data = copy.copy(self)
+
+        for key, value in self:
+            if key == 'edge_index':
+                data.edge_index = edge_index
+            elif key == '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)
+            elif self.is_edge_attr(key):
+                cat_dim = self.__cat_dim__(key, value)
+                data[key] = select(value, edge_mask, dim=cat_dim)
+
+        return data
+
+    def edge_subgraph(self, subset: Tensor) -> 'Data':
+        raise NotImplementedError
+
+    def to_heterogeneous(
+        self,
+        node_type: Optional[Tensor] = None,
+        edge_type: Optional[Tensor] = None,
+        node_type_names: Optional[List[NodeType]] = None,
+        edge_type_names: Optional[List[EdgeType]] = None,
+    ):
+        raise NotImplementedError
+
+    def has_isolated_nodes(self) -> bool:
+        if self.edge_index is None:
+            return False
+        return torch.unique(self.edge_index[0]).size(0) < self.num_nodes
+
+    def is_directed(self) -> bool:
+        raise NotImplementedError
+
+    def is_undirected(self) -> bool:
+        raise NotImplementedError
+
+    def has_self_loops(self) -> bool:
+        raise NotImplementedError
+
+    def validate(self, raise_on_error: bool = True) -> bool:
+        r"""Validates the correctness of the data."""
+        cls_name = self.__class__.__name__
+        status = True
+
+        num_nodes = self.num_nodes
+        if num_nodes is None:
+            status = False
+            warn_or_raise(f"'num_nodes' is undefined in '{cls_name}'",
+                          raise_on_error)
+
+        if 'edge_index' in self:
+            if self.edge_index.dim() != 2 or self.edge_index.size(0) != 2:
+                status = False
+                warn_or_raise(
+                    f"'edge_index' needs to be of shape [2, num_edges] in "
+                    f"'{cls_name}' (found {self.edge_index.size()})",
+                    raise_on_error)
+
+        if 'edge_index' in self and self.edge_index.numel() > 0:
+            if self.edge_index.min() < 0:
+                status = False
+                warn_or_raise(
+                    f"'edge_index' contains negative indices in "
+                    f"'{cls_name}' (found {int(self.edge_index.min())})",
+                    raise_on_error)
+
+            if num_nodes is not None and self.edge_index[0].max() >= num_nodes:
+                status = False
+                warn_or_raise(
+                    f"'edge_index' contains larger indices than the number "
+                    f"of nodes ({num_nodes}) in '{cls_name}' "
+                    f"(found {int(self.edge_index.max())})", raise_on_error)
+
+        return status
+
+
+def warn_or_raise(msg: str, raise_on_error: bool = True):
+    if raise_on_error:
+        raise ValueError(msg)
+    else:
+        warnings.warn(msg)
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
index faa52fea5..2f3ff7c0a 100644
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -15,7 +15,7 @@ from .loop import (contains_self_loops, remove_self_loops,
                    add_remaining_self_loops, get_self_loop_attr)
 from .isolated import contains_isolated_nodes, remove_isolated_nodes
 from .subgraph import (get_num_hops, subgraph, k_hop_subgraph,
-                       bipartite_subgraph)
+                       bipartite_subgraph, hyper_subgraph)
 from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
 from .homophily import homophily
 from .assortativity import assortativity
@@ -80,6 +80,7 @@ __all__ = [
     'subgraph',
     'bipartite_subgraph',
     'k_hop_subgraph',
+    'hyper_subgraph',
     'dropout_node',
     'dropout_edge',
     'dropout_path',
diff --git a/torch_geometric/utils/subgraph.py b/torch_geometric/utils/subgraph.py
index 89c37b9c4..6353281f1 100644
--- a/torch_geometric/utils/subgraph.py
+++ b/torch_geometric/utils/subgraph.py
@@ -326,3 +326,108 @@ def k_hop_subgraph(
         edge_index = node_idx[edge_index]
 
     return subset, edge_index, inv, edge_mask
+
+
+def hyper_subgraph(
+    subset: Union[Tensor, List[int]],
+    edge_index: Tensor,
+    edge_attr: OptTensor = None,
+    relabel_nodes: bool = False,
+    num_nodes: Optional[int] = None,
+    return_edge_mask: bool = False,
+) -> Union[Tuple[Tensor, OptTensor], Tuple[Tensor, OptTensor, OptTensor]]:
+    r"""Returns the induced subgraph of the hyper graph of
+        :obj:`(edge_index, edge_attr)` containing the nodes in :obj:`subset`.
+
+        Args:
+            subset (LongTensor, BoolTensor or [int]): The nodes to keep.
+            edge_index (LongTensor): Hyperedge tensor
+                with shape :obj:`[2, num_edges*num_nodes_per_edge]`.
+                Where `edge_index[1]` denotes the hyperedge index and
+                `edge_index[0]` denotes the node indicies that are connected
+                by the hyperedge.
+            edge_attr (Tensor, optional): Edge weights or multi-dimensional
+                edge features of shape :obj:`[num_edges,-1]`.
+                (default: :obj:`None`)
+            relabel_nodes (bool, optional): If set to :obj:`True`, the
+                resulting :obj:`edge_index` will be relabeled to hold
+                consecutive indices
+                starting from zero. (default: :obj:`False`)
+            num_nodes (int, optional): The number of nodes, *i.e.*
+                :obj:`max_val + 1` of :attr:`edge_index`.
+                (default: :obj:`None`)
+            return_edge_mask (bool, optional): If set to :obj:`True`, will
+                return the edge mask of shape :obj:`num_edges`
+                to filter out additional edge features.
+                (default: :obj:`False`)
+
+        :rtype: (:class:`LongTensor`, :class:`Tensor`)
+
+        Examples:
+
+            >>> edge_index = torch.tensor([[0, 1, 2, 1, 2, 3, 0, 2, 3],
+            ...                            [0, 0, 0, 1, 1, 1, 2, 2, 2]])
+            >>> edge_attr = torch.tensor([3, 2, 6])
+            >>> subset = torch.tensor([0, 3])
+            >>> subgraph(subset, edge_index, edge_attr)
+            (tensor([[0, 3],
+                    [0, 0]]),
+            tensor([ 6.]))
+
+            >>> subgraph(subset, edge_index, edge_attr, return_edge_mask=True)
+            (tensor([[0, 3],
+                    [0, 0]]),
+            tensor([ 6.]))
+            tensor([False, False, True])
+
+        """
+
+    device = edge_index.device
+
+    if isinstance(subset, (list, tuple)):
+        subset = torch.tensor(subset, dtype=torch.long, device=device)
+
+    if subset.dtype != torch.bool:
+        num_nodes = maybe_num_nodes(edge_index, num_nodes)
+        node_mask = index_to_mask(subset, size=num_nodes)
+    else:
+        num_nodes = subset.size(0)
+        node_mask = subset
+
+    # Mask all connections that contain a node not in the subset
+    hyper_edge_connection_mask = node_mask[
+        edge_index[0]]  # num_edges*num_nodes_per_edge
+
+    # Mask hyperedges that contain one or less nodes from the subset
+    num_edges = edge_index[1].max() + 1
+    edge_mask = torch.scatter_add(
+        torch.zeros(num_edges, dtype=torch.long,
+                    device=device), 0, edge_index[1],
+        hyper_edge_connection_mask.to(dtype=torch.long)) > 1  # num_edges
+
+    # Mask connections if hyperedge contains one or less nodes from the subset
+    # or is connected to a node not in the subset
+    hyper_edge_connection_mask = hyper_edge_connection_mask & edge_mask[
+        edge_index[1]]
+
+    edge_index = edge_index[:, hyper_edge_connection_mask]
+    edge_attr = edge_attr[edge_mask] if edge_attr is not None else None
+
+    # Relabel edges
+    edge_idx = torch.zeros(edge_mask.size(0), dtype=torch.long, device=device)
+    edge_idx[edge_mask] = torch.arange(edge_mask.sum().item(), device=device)
+    edge_index = torch.cat(
+        [edge_index[0].unsqueeze(0), edge_idx[edge_index[1]].unsqueeze(0)], 0)
+
+    if relabel_nodes:
+        node_idx = torch.zeros(node_mask.size(0), dtype=torch.long,
+                               device=device)
+        node_idx[subset] = torch.arange(node_mask.sum().item(), device=device)
+        edge_index = torch.cat(
+            [node_idx[edge_index[0]].unsqueeze(0), edge_index[1].unsqueeze(0)],
+            0)
+
+    if return_edge_mask:
+        return edge_index, edge_attr, edge_mask
+    else:
+        return edge_index, edge_attr


Test patch:
diff --git a/test/data/test_hypergraph_data.py b/test/data/test_hypergraph_data.py
new file mode 100644
index 000000000..e4484d683
--- /dev/null
+++ b/test/data/test_hypergraph_data.py
@@ -0,0 +1,171 @@
+import pytest
+import torch
+
+import torch_geometric
+from torch_geometric.data.hypergraph_data import HyperGraphData
+from torch_geometric.loader import DataLoader
+
+
+def test_hypergraph_data():
+    torch_geometric.set_debug(True)
+
+    x = torch.tensor([[1, 3, 5, 7], [2, 4, 6, 8], [7, 8, 9, 10]],
+                     dtype=torch.float).t()
+    edge_index = torch.tensor([[0, 1, 2, 1, 2, 3, 0, 2, 3],
+                               [0, 0, 0, 1, 1, 1, 2, 2, 2]])
+    data = HyperGraphData(x=x, edge_index=edge_index).to(torch.device('cpu'))
+    data.validate(raise_on_error=True)
+
+    assert data.num_nodes == 4
+    assert data.num_edges == 3
+
+    assert data.node_attrs() == ['x']
+    assert data.edge_attrs() == ['edge_index']
+
+    assert data.x.tolist() == x.tolist()
+    assert data['x'].tolist() == x.tolist()
+    assert data.get('x').tolist() == x.tolist()
+    assert data.get('y', 2) == 2
+    assert data.get('y', None) is None
+
+    assert sorted(data.keys()) == ['edge_index', 'x']
+    assert len(data) == 2
+    assert 'x' in data and 'edge_index' in data and 'pos' not in data
+
+    D = data.to_dict()
+    assert len(D) == 2
+    assert 'x' in D and 'edge_index' in D
+
+    D = data.to_namedtuple()
+    assert len(D) == 2
+    assert D.x is not None and D.edge_index is not None
+
+    assert data.__cat_dim__('x', data.x) == 0
+    assert data.__cat_dim__('edge_index', data.edge_index) == -1
+    assert data.__inc__('x', data.x) == 0
+    assert torch.equal(data.__inc__('edge_index', data.edge_index),
+                       torch.tensor([[data.num_nodes], [data.num_edges]]))
+    data_list = [data, data]
+    loader = DataLoader(data_list, batch_size=2)
+    batch = next(iter(loader))
+    batched_edge_index = batch.edge_index
+    assert batched_edge_index.tolist() == [[
+        0, 1, 2, 1, 2, 3, 0, 2, 3, 4, 5, 6, 5, 6, 7, 4, 6, 7
+    ], [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5]]
+
+    assert not data.x.is_contiguous()
+    data.contiguous()
+    assert data.x.is_contiguous()
+
+    assert not data.is_coalesced()
+    data = data.coalesce()
+    assert data.is_coalesced()
+
+    clone = data.clone()
+    assert clone != data
+    assert len(clone) == len(data)
+    assert clone.x.data_ptr() != data.x.data_ptr()
+    assert clone.x.tolist() == data.x.tolist()
+    assert clone.edge_index.data_ptr() != data.edge_index.data_ptr()
+    assert clone.edge_index.tolist() == data.edge_index.tolist()
+
+    data['x'] = x + 1
+    assert data.x.tolist() == (x + 1).tolist()
+
+    assert str(data) == 'HyperGraphData(x=[4, 3], edge_index=[2, 9])'
+
+    dictionary = {'x': data.x, 'edge_index': data.edge_index}
+    data = HyperGraphData.from_dict(dictionary)
+    assert sorted(data.keys()) == ['edge_index', 'x']
+
+    assert not data.has_isolated_nodes()
+    # assert not data.has_self_loops()
+    # assert data.is_undirected()
+    # assert not data.is_directed()
+
+    assert data.num_nodes == 4
+    assert data.num_edges == 3
+    with pytest.warns(UserWarning, match='deprecated'):
+        assert data.num_faces is None
+    assert data.num_node_features == 3
+    assert data.num_features == 3
+
+    data.edge_attr = torch.randn(data.num_edges, 2)
+    assert data.num_edge_features == 2
+    assert data.is_edge_attr('edge_attr')
+    data.edge_attr = None
+
+    data.x = None
+    with pytest.warns(UserWarning, match='Unable to accurately infer'):
+        assert data.num_nodes == 4
+
+    data.edge_index = None
+    with pytest.warns(UserWarning, match='Unable to accurately infer'):
+        assert data.num_nodes is None
+    assert data.num_edges == 0
+
+    data.num_nodes = 4
+    assert data.num_nodes == 4
+
+    data = HyperGraphData(x=x, attribute=x)
+    assert len(data) == 2
+    assert data.x.tolist() == x.tolist()
+    assert data.attribute.tolist() == x.tolist()
+
+    face = torch.tensor([[0, 1], [1, 2], [2, 3]])
+    data = HyperGraphData(num_nodes=4, face=face)
+    with pytest.warns(UserWarning, match='deprecated'):
+        assert data.num_faces == 2
+    assert data.num_nodes == 4
+
+    data = HyperGraphData(title='test')
+    assert str(data) == "HyperGraphData(title='test')"
+    assert data.num_node_features == 0
+    # assert data.num_edge_features == 0
+
+    key = value = 'test_value'
+    data[key] = value
+    assert data[key] == value
+    del data[value]
+    del data[value]  # Deleting unset attributes should work as well.
+
+    assert data.get(key) is None
+    assert data.get('title') == 'test'
+
+    torch_geometric.set_debug(False)
+
+
+def test_hypergraphdata_subgraph():
+    x = torch.arange(5)
+    y = torch.tensor([0.])
+    edge_index = torch.tensor([[0, 1, 3, 2, 4, 0, 3, 4, 2, 1, 2, 3],
+                               [0, 0, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3]])
+    edge_attr = torch.rand(4, 2)
+    data = HyperGraphData(x=x, y=y, edge_index=edge_index, edge_attr=edge_attr,
+                          num_nodes=5)
+
+    out = data.subgraph(torch.tensor([1, 2, 4]))
+    assert len(out) == 5
+    assert torch.equal(out.x, torch.tensor([1, 2, 4]))
+    assert torch.equal(out.y, data.y)
+    assert out.edge_index.tolist() == [[1, 2, 2, 1, 0, 1], [0, 0, 1, 1, 2, 2]]
+    assert torch.equal(out.edge_attr, edge_attr[[1, 2, 3]])
+    assert out.num_nodes == 3
+
+    # Test unordered selection:
+    out = data.subgraph(torch.tensor([3, 1, 2]))
+    assert len(out) == 5
+    assert torch.equal(out.x, torch.tensor([3, 1, 2]))
+    assert torch.equal(out.y, data.y)
+    assert out.edge_index.tolist() == [[0, 2, 0, 2, 1, 2, 0],
+                                       [0, 0, 1, 1, 2, 2, 2]]
+    assert torch.equal(out.edge_attr, edge_attr[[1, 2, 3]])
+    assert out.num_nodes == 3
+
+    out = data.subgraph(torch.tensor([False, False, False, True, True]))
+    assert len(out) == 5
+    assert torch.equal(out.x, torch.arange(3, 5))
+    assert torch.equal(out.y, data.y)
+    assert out.edge_index.tolist() == [[0, 1, 0, 1], [0, 0, 1, 1]]
+    assert torch.equal(out.edge_attr, edge_attr[[1, 2]])
+    assert out.num_nodes == 2