Repository: pyg-team/pytorch_geometric
Pull request #8269 (2023-10-26 19:51:17)

Add alternative edge indices shape to `GCNConv` input description
<img width="752" alt="Screenshot 2023-10-26 at 1 05 22 PM" src="https://github.com/pyg-team/pytorch_geometric/assets/54658925/fa7d55b6-30ac-4036-baf3-1133e5c47257">
According to the implementation https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/gcn_conv.py#L49, sparse adjacency matrix of shape |V|*|V| can also be used as edge_indices.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8269?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8269](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8269?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (9ac32d3) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/3752d94447950e83e9d6e03c1332d0c031a2376c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (3752d94) will **not change** coverage.
> The diff coverage is `n/a`.

```diff
@@           Coverage Diff           @@
##           master    #8269   +/-   ##
=======================================
  Coverage   87.41%   87.41%           
=======================================
  Files         473      473           
  Lines       28635    28635           
=======================================
  Hits        25032    25032           
  Misses       3603     3603           
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8269?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/conv/gcn\_conv.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8269?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvZ2NuX2NvbnYucHk=) | `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 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/examples/kuzu/papers_100M/train.py b/examples/kuzu/papers_100M/train.py
index 9599c1c01..84c7484ab 100644
--- a/examples/kuzu/papers_100M/train.py
+++ b/examples/kuzu/papers_100M/train.py
@@ -68,10 +68,10 @@ class GraphSAGE(torch.nn.Module):
         self.norms = torch.nn.ModuleList()
 
         self.convs.append(SAGEConv(in_channels, hidden_channels))
-        self.bns.append(BatchNorm(hidden_channels))
+        self.norms.append(BatchNorm(hidden_channels))
         for i in range(1, num_layers):
-            self.layers.append(SAGEConv(hidden_channels, hidden_channels))
-            self.bns.append(BatchNorm(hidden_channels))
+            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
+            self.norms.append(BatchNorm(hidden_channels))
 
         self.mlp = MLP(
             in_channels=in_channels + num_layers * hidden_channels,
@@ -82,6 +82,8 @@ class GraphSAGE(torch.nn.Module):
             act='leaky_relu',
         )
 
+        self.dropout = dropout
+
     def forward(self, x, edge_index):
         x = F.dropout(x, p=self.dropout, training=self.training)
         xs = [x]
@@ -105,9 +107,12 @@ for epoch in range(1, NUM_EPOCHS + 1):
         batch_size = batch['paper'].batch_size
 
         optimizer.zero_grad()
-        out = model(batch.x, batch.edge_index)[:batch_size]
-        y = batch.y[:batch_size].long().view(-1)
-        loss = F.cross_entropy_loss(out, y)
+        out = model(
+            batch['paper'].x,
+            batch['paper', 'cites', 'paper'].edge_index,
+        )[:batch_size]
+        y = batch['paper'].y[:batch_size].long().view(-1)
+        loss = F.cross_entropy(out, y)
 
         loss.backward()
         optimizer.step()
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)
diff --git a/torch_geometric/nn/conv/gcn_conv.py b/torch_geometric/nn/conv/gcn_conv.py
index e1d330430..a18dff20d 100644
--- a/torch_geometric/nn/conv/gcn_conv.py
+++ b/torch_geometric/nn/conv/gcn_conv.py
@@ -159,7 +159,8 @@ class GCNConv(MessagePassing):
     Shapes:
         - **input:**
           node features :math:`(|\mathcal{V}|, F_{in})`,
-          edge indices :math:`(2, |\mathcal{E}|)`,
+          edge indices :math:`(2, |\mathcal{E}|)`
+          or sparse matrix :math:`(|\mathcal{V}|, |\mathcal{V}|)`,
           edge weights :math:`(|\mathcal{E}|)` *(optional)*
         - **output:** node features :math:`(|\mathcal{V}|, F_{out})`
     """


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 6c27d567d..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
@@ -702,7 +722,7 @@ def test_basic_graph_store():
     assert len(data.get_all_edge_attrs()) == 0
 
 
-def test_data_generate_ids():
+def test_generate_ids():
     data = HeteroData()
 
     data['paper'].x = torch.randn(100, 128)
@@ -718,3 +738,21 @@ def test_data_generate_ids():
     assert data['author'].n_id.tolist() == list(range(200))
     assert data['paper', 'author'].e_id.tolist() == list(range(300))
     assert data['author', 'paper'].e_id.tolist() == list(range(400))
+
+
+def test_invalid_keys():
+    data = HeteroData()
+
+    data['paper'].x = torch.randn(10, 128)
+    data['paper'].node_attrs = ['y']
+    data['paper', 'paper'].edge_index = get_random_edge_index(10, 10, 20)
+    data['paper', 'paper'].edge_attrs = ['edge_attr']
+
+    assert data['paper'].node_attrs() == ['x']
+    assert data['paper']['node_attrs'] == ['y']
+    assert data['paper', 'paper'].edge_attrs() == ['edge_index']
+    assert data['paper', 'paper']['edge_attrs'] == ['edge_attr']
+
+    out = data.to_homogeneous()
+    assert set(out.node_attrs()) == {'x', 'node_type'}
+    assert set(out.edge_attrs()) == {'edge_index', 'edge_type'}