Repository: pyg-team/pytorch_geometric

Issue:
`HeteroData.subgraph()`
### 🚀 The feature, motivation and pitch

Similar to `Data.subgraph()`, there should exist a `HeteroData.subgraph()` method to compute subgraphs in a heterogeneous graph setting, *e.g.*, for obtaining inductive node splits. Here, `mask`/`index` should be of type `dict`, holding masks/indices for each/a subset of node types:
```python
hetero_data.subgraph({'paper': mask})
```

### Alternatives

_No response_

### Additional context

_No response_


Discussion:
Hi @rusty1s , I was about to open a new discussion, and just realized you are already on this. Just commenting to share my interest in this feature. Cheers. 
@michalisfrangos  are you interested in contributing this feature?
Pinging @mananshah99 and @sdulloor here who shared interest in contributing this feature as well.
It might be useful to implement a `utils.subgraph_bipartite(subset:Tuple[torch.Tensor,torch.Tensor],...)` or add support to `utils.subgraph` for bipartite graphs. I prefer adding a new function over modifying the existing one to make the code more clean.
That way `HeteroData.subgraph()` would make multiple calls to `subgraph_bipartite`. Something like
```
subgraph(node_mask_dict):
     ....
     for edge_type in self.edge_types:
           if edge_type[0] in node_mask_dict and..:
                   new_edge, _ , _ = utils.subgraph_bipartite((node_mask_dict[edge_type[0], node_mask_dict[edge_type[-1]))

```
WDYT?
Yes, this looks good to me. Although we overload a lot of functionality with bipartite graph support already (by passing tuples instead of single tensors), I agree that adding this directly to `subgraph` might makes the code overly complex. `bipartite_subgraph` is a good alternative that we do not even have to expose. 
how would this be different from just sampling a heterogeneous graph with large neighbourhoods to get different node types in the new sampled bipartite graph? 
Not sure I understand. Can you clarify? The `subgraph()` method might be useful to gather subgraphs prior to any training or sampling, e.g. for obtaining inductive subgraphs based on a pre-defined split.
so I have a transductive problem (for now) and for heterognn classification I am planning to just use the `HGTLoader` to get smaller batches for a list of nodes to train my model. Does that set up seem correct? I'm not sure how/if i should be using something like the `subgraph()` method, (whenever its implemented).
It depends on which data you want to train on. If you want to shrink the data prior to training, then `HeteroData.subgraph` would be applicable to create a smaller subgraph from your original graph. If you just want to operate on smaller batches during training, then you may want to adjust the `batch_size` argument of a loader.

Let me know if that makes sense to you.

Resolving patch:
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -12,7 +12,7 @@
 from torch_geometric.data.data import BaseData, Data, size_repr
 from torch_geometric.data.storage import BaseStorage, EdgeStorage, NodeStorage
 from torch_geometric.typing import EdgeType, NodeType, QueryType
-from torch_geometric.utils import is_undirected
+from torch_geometric.utils import bipartite_subgraph, is_undirected
 
 NodeOrEdgeType = Union[NodeType, EdgeType]
 NodeOrEdgeStorage = Union[NodeStorage, EdgeStorage]
@@ -445,6 +445,83 @@ def rename(self, name: NodeType, new_name: NodeType) -> 'HeteroData':
 
         return self
 
+    def subgraph(self, subset_dict: Dict[NodeType, Tensor]) -> 'HeteroData':
+        r"""Returns the induced subgraph containing the node types and
+        corresponding nodes in :obj:`subset_dict`.
+
+        .. code-block:: python
+
+            data = HeteroData()
+            data['paper'].x = ...
+            data['author'].x = ...
+            data['conference'].x = ...
+            data['paper', 'cites', 'paper'].edge_index = ...
+            data['author', 'paper'].edge_index = ...
+            data['paper', 'conference'].edge_index = ...
+            print(data)
+            >>> HeteroData(
+                paper={ x=[10, 16] },
+                author={ x=[5, 32] },
+                conference={ x=[5, 8] },
+                (paper, cites, paper)={ edge_index=[2, 50] },
+                (author, to, paper)={ edge_index=[2, 30] },
+                (paper, to, conference)={ edge_index=[2, 25] }
+            )
+
+            subset_dict = {
+                'paper': torch.tensor([3, 4, 5, 6]),
+                'author': torch.tensor([0, 2]),
+            }
+
+            print(data.subgraph(subset_dict))
+            >>> HeteroData(
+                paper={ x=[4, 16] },
+                author={ x=[2, 32] },
+                (paper, cites, paper)={ edge_index=[2, 24] },
+                (author, to, paper)={ edge_index=[2, 5] }
+            )
+
+        Args:
+            subset_dict (Dict[str, LongTensor or BoolTensor]): A dictonary
+                holding the nodes to keep for each node type.
+        """
+        data = self.__class__(self._global_store)
+
+        for node_type, subset in subset_dict.items():
+            for key, value in self[node_type].items():
+                if key == '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:
+                    data[node_type][key] = value
+
+        for edge_type in self.edge_types:
+            src, _, dst = edge_type
+            if src not in subset_dict or dst not in subset_dict:
+                continue
+
+            edge_index, _, edge_mask = bipartite_subgraph(
+                (subset_dict[src], subset_dict[dst]),
+                self[edge_type].edge_index,
+                relabel_nodes=True,
+                size=(self[src].num_nodes, self[dst].num_nodes),
+                return_edge_mask=True,
+            )
+
+            for key, value in self[edge_type].items():
+                if key == 'edge_index':
+                    data[edge_type].edge_index = edge_index
+                elif self[edge_type].is_edge_attr(key):
+                    data[edge_type][key] = value[edge_mask]
+                else:
+                    data[edge_type][key] = value
+
+        return data
+
     def to_homogeneous(self, node_attrs: Optional[List[str]] = None,
                        edge_attrs: Optional[List[str]] = None,
                        add_node_type: bool = True,


Test patch:
diff --git a/test/data/test_hetero_data.py b/test/data/test_hetero_data.py
--- a/test/data/test_hetero_data.py
+++ b/test/data/test_hetero_data.py
@@ -6,15 +6,20 @@
 
 x_paper = torch.randn(10, 16)
 x_author = torch.randn(5, 32)
+x_conference = torch.randn(5, 8)
 
 idx_paper = torch.randint(x_paper.size(0), (100, ), dtype=torch.long)
 idx_author = torch.randint(x_author.size(0), (100, ), dtype=torch.long)
+idx_conference = torch.randint(x_conference.size(0), (100, ), dtype=torch.long)
 
 edge_index_paper_paper = torch.stack([idx_paper[:50], idx_paper[:50]], dim=0)
 edge_index_paper_author = torch.stack([idx_paper[:30], idx_author[:30]], dim=0)
-edge_index_author_paper = torch.stack([idx_paper[:30], idx_author[:30]], dim=0)
+edge_index_author_paper = torch.stack([idx_author[:30], idx_paper[:30]], dim=0)
+edge_index_paper_conference = torch.stack(
+    [idx_paper[:25], idx_conference[:25]], dim=0)
 
 edge_attr_paper_paper = torch.randn(edge_index_paper_paper.size(1), 8)
+edge_attr_author_paper = torch.randn(edge_index_author_paper.size(1), 8)
 
 
 def get_edge_index(num_src_nodes, num_dst_nodes, num_edges):
@@ -159,6 +164,57 @@ def test_hetero_data_rename():
     assert edge_index.tolist() == edge_index_paper_paper.tolist()
 
 
+def test_hetero_data_subgraph():
+    data = HeteroData()
+    data.num_node_types = 3
+    data['paper'].x = x_paper
+    data['paper'].name = 'paper'
+    data['paper'].num_nodes = x_paper.size(0)
+    data['author'].x = x_author
+    data['author'].num_nodes = x_author.size(0)
+    data['conference'].x = x_conference
+    data['conference'].num_nodes = x_conference.size(0)
+    data['paper', 'paper'].edge_index = edge_index_paper_paper
+    data['paper', 'paper'].edge_attr = edge_attr_paper_paper
+    data['paper', 'paper'].name = 'cites'
+    data['author', 'paper'].edge_index = edge_index_author_paper
+    data['paper', 'author'].edge_index = edge_index_paper_author
+    data['paper', 'conference'].edge_index = edge_index_paper_conference
+
+    subset = {
+        'paper': torch.randperm(x_paper.size(0))[:4],
+        'author': torch.randperm(x_author.size(0))[:2]
+    }
+
+    out = data.subgraph(subset)
+
+    assert out.num_node_types == data.num_node_types
+    assert out.node_types == ['paper', 'author']
+
+    assert len(out['paper']) == 3
+    assert torch.allclose(out['paper'].x, data['paper'].x[subset['paper']])
+    assert out['paper'].name == 'paper'
+    assert out['paper'].num_nodes == 4
+    assert len(out['author']) == 2
+    assert torch.allclose(out['author'].x, data['author'].x[subset['author']])
+    assert out['author'].num_nodes == 2
+
+    assert out.edge_types == [
+        ('paper', 'to', 'paper'),
+        ('author', 'to', 'paper'),
+        ('paper', 'to', 'author'),
+    ]
+
+    assert len(out['paper', 'paper']) == 3
+    assert out['paper', 'paper'].edge_index is not None
+    assert out['paper', 'paper'].edge_attr is not None
+    assert out['paper', 'paper'].name == 'cites'
+    assert len(out['paper', 'author']) == 1
+    assert out['paper', 'author'].edge_index is not None
+    assert len(out['author', 'paper']) == 1
+    assert out['author', 'paper'].edge_index is not None
+
+
 def test_copy_hetero_data():
     data = HeteroData()
     data['paper'].x = x_paper