Repository: pyg-team/pytorch_geometric
Pull request #8397 (2023-11-17 12:12:23)

Fix issue with uninitialized module in `GAT+hetero+DDP` scenario
This PR fixes the error described in #8382.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8397?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`cf24b4b`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/cf24b4bcb4e825537ba08d8fc5f31073e2cd84c7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.73% compared to head [(`c1fb291`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8397?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.74%.

> :exclamation: Current head c1fb291 differs from pull request most recent head ce713f1. Consider uploading reports for the commit ce713f1 to get more accurate results


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8397   +/-   ##
=======================================
  Coverage   87.73%   87.74%           
=======================================
  Files         475      475           
  Lines       28890    28899    +9     
=======================================
+ Hits        25348    25357    +9     
  Misses       3542     3542           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8397?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).

DamianSzwichtenberg: > Hey, I pushed an alternative fix. I think this one should resolve your issues without any `is_unitialized` handling.

Yep, I confirm, it works. Thanks 😉 

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 196f79ba6..4ecc97adb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -37,6 +37,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Breaking Change: Properly initialize modules in `GATConv` depending on whether the input is bipartite or non-bipartite ([#8397](https://github.com/pyg-team/pytorch_geometric/pull/8397))
 - Fixed `input_id` computation in `NeighborLoader` in case a `mask` is given ([#8312](https://github.com/pyg-team/pytorch_geometric/pull/8312))
 - Respect current device when deep-copying `Linear` layers ([#8311](https://github.com/pyg-team/pytorch_geometric/pull/8311))
 - Fixed `Data.subgraph()`/`HeteroData.subgraph()` in case `edge_index` is not defined ([#8277](https://github.com/pyg-team/pytorch_geometric/pull/8277))
diff --git a/torch_geometric/nn/conv/gat_conv.py b/torch_geometric/nn/conv/gat_conv.py
index ad1d66359..3e4dcd69c 100644
--- a/torch_geometric/nn/conv/gat_conv.py
+++ b/torch_geometric/nn/conv/gat_conv.py
@@ -149,10 +149,10 @@ class GATConv(MessagePassing):
 
         # In case we are operating in bipartite graphs, we apply separate
         # transformations 'lin_src' and 'lin_dst' to source and target nodes:
+        self.lin = self.lin_src = self.lin_dst = None
         if isinstance(in_channels, int):
-            self.lin_src = Linear(in_channels, heads * out_channels,
-                                  bias=False, weight_initializer='glorot')
-            self.lin_dst = self.lin_src
+            self.lin = Linear(in_channels, heads * out_channels, bias=False,
+                              weight_initializer='glorot')
         else:
             self.lin_src = Linear(in_channels[0], heads * out_channels, False,
                                   weight_initializer='glorot')
@@ -182,8 +182,12 @@ class GATConv(MessagePassing):
 
     def reset_parameters(self):
         super().reset_parameters()
-        self.lin_src.reset_parameters()
-        self.lin_dst.reset_parameters()
+        if self.lin is not None:
+            self.lin.reset_parameters()
+        if self.lin_src is not None:
+            self.lin_src.reset_parameters()
+        if self.lin_dst is not None:
+            self.lin_dst.reset_parameters()
         if self.lin_edge is not None:
             self.lin_edge.reset_parameters()
         glorot(self.att_src)
@@ -232,13 +236,33 @@ class GATConv(MessagePassing):
         # transform source and target node features via separate weights:
         if isinstance(x, Tensor):
             assert x.dim() == 2, "Static graphs not supported in 'GATConv'"
-            x_src = x_dst = self.lin_src(x).view(-1, H, C)
+
+            if self.lin is not None:
+                x_src = x_dst = self.lin(x).view(-1, H, C)
+            else:
+                # If the module is initialized as bipartite, transform source
+                # and destination node features separately:
+                assert self.lin_src is not None and self.lin_dst is not None
+                x_src = self.lin_src(x).view(-1, H, C)
+                x_dst = self.lin_dst(x).view(-1, H, C)
+
         else:  # Tuple of source and target node features:
             x_src, x_dst = x
             assert x_src.dim() == 2, "Static graphs not supported in 'GATConv'"
-            x_src = self.lin_src(x_src).view(-1, H, C)
-            if x_dst is not None:
-                x_dst = self.lin_dst(x_dst).view(-1, H, C)
+
+            if self.lin is not None:
+                # If the module is initialized as non-bipartite, we expect that
+                # source and destination node features have the same shape and
+                # that they their transformations are shared:
+                x_src = self.lin(x_src).view(-1, H, C)
+                if x_dst is not None:
+                    x_dst = self.lin(x_dst).view(-1, H, C)
+            else:
+                assert self.lin_src is not None and self.lin_dst is not None
+
+                x_src = self.lin_src(x_src).view(-1, H, C)
+                if x_dst is not None:
+                    x_dst = self.lin_dst(x_dst).view(-1, H, C)
 
         x = (x_src, x_dst)
 


Test patch:
diff --git a/test/nn/conv/test_point_conv.py b/test/nn/conv/test_point_conv.py
index 84c0069ed..496b2442b 100644
--- a/test/nn/conv/test_point_conv.py
+++ b/test/nn/conv/test_point_conv.py
@@ -53,12 +53,14 @@ def test_point_net_conv():
     assert out.size() == (2, 32)
     assert torch.allclose(conv((x1, None), (pos1, pos2), edge_index), out)
     assert torch.allclose(conv(x1, (pos1, pos2), adj1.t()), out, atol=1e-6)
-    assert torch.allclose(conv((x1, None), (pos1, pos2), adj1.t()), out)
+    assert torch.allclose(conv((x1, None), (pos1, pos2), adj1.t()), out,
+                          atol=1e-6)
 
     if torch_geometric.typing.WITH_TORCH_SPARSE:
         adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 2))
         assert torch.allclose(conv(x1, (pos1, pos2), adj2.t()), out, atol=1e-6)
-        assert torch.allclose(conv((x1, None), (pos1, pos2), adj2.t()), out)
+        assert torch.allclose(conv((x1, None), (pos1, pos2), adj2.t()), out,
+                              atol=1e-6)
 
     if is_full_test():
         t = '(PairOptTensor, PairTensor, Tensor) -> Tensor'
@@ -68,4 +70,5 @@ def test_point_net_conv():
     if is_full_test() and torch_geometric.typing.WITH_TORCH_SPARSE:
         t = '(PairOptTensor, PairTensor, SparseTensor) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert torch.allclose(jit((x1, None), (pos1, pos2), adj2.t()), out)
+        assert torch.allclose(jit((x1, None), (pos1, pos2), adj2.t()), out,
+                              atol=1e-6)
diff --git a/test/nn/dense/test_dense_gat_conv.py b/test/nn/dense/test_dense_gat_conv.py
index d54ee28b4..7c2c17dc8 100644
--- a/test/nn/dense/test_dense_gat_conv.py
+++ b/test/nn/dense/test_dense_gat_conv.py
@@ -14,7 +14,7 @@ def test_dense_gat_conv(heads, concat):
     assert str(dense_conv) == f'DenseGATConv(16, 16, heads={heads})'
 
     # Ensure same weights and bias:
-    dense_conv.lin = sparse_conv.lin_src
+    dense_conv.lin = sparse_conv.lin
     dense_conv.att_src = sparse_conv.att_src
     dense_conv.att_dst = sparse_conv.att_dst
     dense_conv.bias = sparse_conv.bias