Update test condition to check bipartite graph in forward pass

pyg-team · Dec 10, 2024 · e4cdd0f · e4cdd0f
1 parent 2bc5232
commit e4cdd0f
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 19 deletions.
diff --git a/test/nn/conv/test_gat_conv.py b/test/nn/conv/test_gat_conv.py
@@ -267,10 +267,10 @@ def test_gat_norm_csc_error():
 
 def test_gat_conv_bipartite_error():
     x1 = torch.randn(4, 8)
+    x2 = torch.randn(2, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
 
     with pytest.raises(NotImplementedError,
                        match="not supported for bipartite message passing"):
-        conv = GATConv(8, 32, heads=2, normalize=True)
-        _ = conv(x1, adj1.t())
+        conv = GATConv((8, 16), 32, heads=2, normalize=True)
+        conv((x1, x2), edge_index)
diff --git a/test/nn/conv/test_gatv2_conv.py b/test/nn/conv/test_gatv2_conv.py
@@ -205,10 +205,10 @@ def test_gatv2_conv_with_edge_attr():
 
 def test_gat_conv_bipartite_error():
     x1 = torch.randn(4, 8)
+    x2 = torch.randn(2, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
-    adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
 
     with pytest.raises(NotImplementedError,
                        match="not supported for bipartite message passing"):
-        conv = GATv2Conv(8, 32, heads=2, normalize=True)
-        _ = conv(x1, adj1.t())
+        conv = GATv2Conv((8, 16), 32, heads=2, normalize=True)
+        conv((x1, x2), edge_index)
diff --git a/torch_geometric/nn/conv/gat_conv.py b/torch_geometric/nn/conv/gat_conv.py
@@ -502,12 +502,10 @@ def forward(  # noqa: F811
                         "'edge_index' in a 'SparseTensor' form")
 
         if self.normalize:
-            if isinstance(edge_index,
-                          SparseTensor) or is_torch_sparse_tensor(edge_index):
-                if edge_index.size(0) != edge_index.size(1):
-                    raise NotImplementedError(
-                        "The usage of 'normalize' is not supported "
-                        "for bipartite message passing.")
+            if not isinstance(self.in_channels, int):
+                raise NotImplementedError(
+                    "The usage of 'normalize' is not supported "
+                    "for bipartite message passing.")
 
             if isinstance(edge_index, Tensor):
                 edge_index, edge_attr = remove_self_loops(
@@ -520,7 +518,7 @@ def forward(  # noqa: F811
                                   size=size)
 
         if self.normalize:
-            num_nodes = None
+            num_nodes: Optional[int] = None
             if isinstance(edge_index, Tensor):
                 num_nodes = x_src.size(0)
                 if x_dst is not None:

diff --git a/torch_geometric/nn/conv/gatv2_conv.py b/torch_geometric/nn/conv/gatv2_conv.py
@@ -331,12 +331,10 @@ def forward(  # noqa: F811
                         "'edge_index' in a 'SparseTensor' form")
 
         if self.normalize:
-            if isinstance(edge_index,
-                          SparseTensor) or is_torch_sparse_tensor(edge_index):
-                if edge_index.size(0) != edge_index.size(1):
-                    raise NotImplementedError(
-                        "The usage of 'normalize' is not supported "
-                        "for bipartite message passing.")
+            if not isinstance(self.in_channels, int):
+                raise NotImplementedError(
+                    "The usage of 'normalize' is not supported "
+                    "for bipartite message passing.")
 
             if isinstance(edge_index, Tensor):
                 edge_index, edge_attr = remove_self_loops(