Repository: pyg-team/pytorch_geometric
Pull request #8277 (2023-10-28 13:59:46)

Fix `subgraph` method in case `edge_index` is not defined
Fixes https://github.com/pyg-team/pytorch_geometric/issues/8272

codecov[bot]: We're currently processing your upload.  This comment will be updated when the results are available.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 29605db98..d72c4903f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `Data.subgraph()`/`HeteroData.subgraph()` in case `edge_index` is not defined ([#8277](https://github.com/pyg-team/pytorch_geometric/pull/8277))
 - Fixed empty edge handling in `MetaPath2Vec` ([#8248](https://github.com/pyg-team/pytorch_geometric/pull/8248))
 - Fixed `AttentionExplainer` usage within `AttentiveFP` ([#8244](https://github.com/pyg-team/pytorch_geometric/pull/8244))
 - Fixed `load_from_state_dict` in lazy `Linear` modules ([#8242](https://github.com/pyg-team/pytorch_geometric/pull/8242))
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 77d5e8610..add9903fe 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -630,9 +630,21 @@ class Data(BaseData, FeatureStore, GraphStore):
         Args:
             subset (LongTensor or BoolTensor): The nodes to keep.
         """
-        out = subgraph(subset, self.edge_index, relabel_nodes=True,
-                       num_nodes=self.num_nodes, return_edge_mask=True)
-        edge_index, _, edge_mask = out
+        if 'edge_index' in self:
+            edge_index, _, edge_mask = subgraph(
+                subset,
+                self.edge_index,
+                relabel_nodes=True,
+                num_nodes=self.num_nodes,
+                return_edge_mask=True,
+            )
+        else:
+            edge_index = None
+            edge_mask = torch.ones(
+                self.num_edges,
+                dtype=torch.bool,
+                device=subset.device,
+            )
 
         data = copy.copy(self)
 
@@ -843,31 +855,31 @@ class Data(BaseData, FeatureStore, GraphStore):
             yield key, value
 
     @property
-    def x(self) -> Any:
+    def x(self) -> Optional[Tensor]:
         return self['x'] if 'x' in self._store else None
 
     @property
-    def edge_index(self) -> Any:
+    def edge_index(self) -> Optional[Tensor]:
         return self['edge_index'] if 'edge_index' in self._store else None
 
     @property
-    def edge_weight(self) -> Any:
+    def edge_weight(self) -> Optional[Tensor]:
         return self['edge_weight'] if 'edge_weight' in self._store else None
 
     @property
-    def edge_attr(self) -> Any:
+    def edge_attr(self) -> Optional[Tensor]:
         return self['edge_attr'] if 'edge_attr' in self._store else None
 
     @property
-    def y(self) -> Any:
+    def y(self) -> Optional[Union[Tensor, int, float]]:
         return self['y'] if 'y' in self._store else None
 
     @property
-    def pos(self) -> Any:
+    def pos(self) -> Optional[Tensor]:
         return self['pos'] if 'pos' in self._store else None
 
     @property
-    def batch(self) -> Any:
+    def batch(self) -> Optional[Tensor]:
         return self['batch'] if 'batch' in self._store else None
 
     # Deprecated functions ####################################################
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index f28896e0c..72ccd9fc0 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -682,6 +682,9 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                     data[node_type][key] = value
 
         for edge_type in self.edge_types:
+            if 'edge_index' not in self[edge_type]:
+                continue
+
             src, _, dst = edge_type
 
             src_subset = subset_dict.get(src)


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index 299d42fc2..8ebc4e39a 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -268,6 +268,16 @@ def test_data_subgraph_with_list_field():
     assert out.x.tolist() == out.y == [1, 2, 3]
 
 
+def test_data_empty_subgraph():
+    data = Data(x=torch.arange(5), y=torch.tensor(0.0))
+
+    out = data.subgraph(torch.tensor([1, 2, 3]))
+    assert 'edge_index' not in out
+    assert torch.equal(out.x, torch.arange(1, 4))
+    assert torch.equal(out.y, data.y)
+    assert out.num_nodes == 3
+
+
 def test_copy_data():
     data = Data(x=torch.randn(20, 5))
 
diff --git a/test/data/test_hetero_data.py b/test/data/test_hetero_data.py
index 587dbf5d0..80ea0d868 100644
--- a/test/data/test_hetero_data.py
+++ b/test/data/test_hetero_data.py
@@ -332,6 +332,26 @@ def test_hetero_data_subgraph():
     assert out['paper', 'paper'].edge_attr.size() == (4, 8)
 
 
+def test_hetero_data_empty_subgraph():
+    data = HeteroData()
+    data.num_node_types = 3
+    data['paper'].x = torch.arange(5)
+    data['author'].x = torch.arange(5)
+    data['paper', 'author'].edge_weight = torch.arange(5)
+
+    out = data.subgraph(subset_dict={
+        'paper': torch.tensor([1, 2, 3]),
+        'author': torch.tensor([1, 2, 3]),
+    })
+
+    assert torch.equal(out['paper'].x, torch.arange(1, 4))
+    assert out['paper'].num_nodes == 3
+    assert torch.equal(out['author'].x, torch.arange(1, 4))
+    assert out['author'].num_nodes == 3
+    assert 'edge_index' not in out['paper', 'author']
+    assert torch.equal(out['paper', 'author'].edge_weight, torch.arange(5))
+
+
 def test_copy_hetero_data():
     data = HeteroData()
     data['paper'].x = x_paper