Add check on add_self_loops in HeteroConv and to_hetero

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.0.5] - 2022-MM-DD
 ### Added
+- Added a check in `HeteroConv` and `to_hetero` to make sure `add_self_loops` is not true ([4647](https://github.com/pyg-team/pytorch_geometric/pull/4647))
 - Added `HeteroData.subgraph()` support ([#4635](https://github.com/pyg-team/pytorch_geometric/pull/4635))
 - Added the `AQSOL` dataset ([#4626](https://github.com/pyg-team/pytorch_geometric/pull/4626))
 - Added `HeteroData.node_items()` and `HeteroData.edge_items()` functionality ([#4644](https://github.com/pyg-team/pytorch_geometric/pull/4644))

test/nn/conv/test_hetero_conv.py

@@ -77,3 +77,16 @@ def test_hetero_conv_with_custom_conv():
     assert len(out) == 2
     assert out['paper'].size() == (50, 64)
     assert out['author'].size() == (30, 64)
+
+
+class MessagePassingLoops(MessagePassing):
+    def __init__(self, add_self_loops: bool = True):
+        super().__init__()
+        self.add_self_loops = add_self_loops
+
+
+def test_hetero_exception_self_loops():
+    model = MessagePassingLoops()
+    with pytest.raises(ValueError):
+        HeteroConv({("a", "to", "b"): model})
+    HeteroConv({("a", "to", "a"): model})

test/nn/test_to_hetero_transformer.py

@@ -1,5 +1,6 @@
 from typing import Tuple
 
+import pytest
 import torch
 from torch import Tensor
 from torch.nn import Linear, ReLU, Sequential
@@ -363,3 +364,19 @@ def test_graph_level_to_hetero():
     model = to_hetero(model, metadata, aggr='mean', debug=False)
     out = model(x_dict, edge_index_dict, batch_dict)
     assert out.size() == (1, 64)
+
+
+class MessagePassingLoops(MessagePassing):
+    def __init__(self, add_self_loops: bool = True):
+        super().__init__()
+        self.add_self_loops = add_self_loops
+
+    def forward(self):
+        pass
+
+
+def test_hetero_transformer_exception_self_loops():
+    model = MessagePassingLoops()
+    with pytest.raises(ValueError):
+        to_hetero(model, (['a'], (('a', 'to', 'b'), )), aggr='mean')
+    to_hetero(model, (['a'], (('a', 'to', 'a'), )), aggr='mean')

torch_geometric/nn/conv/hetero_conv.py

@@ -7,6 +7,7 @@
 
 from torch_geometric.nn.conv.hgt_conv import group
 from torch_geometric.typing import Adj, EdgeType, NodeType
+from torch_geometric.utils.hetero import check_add_self_loops
 
 
 class HeteroConv(Module):
@@ -46,6 +47,7 @@ class HeteroConv(Module):
     def __init__(self, convs: Dict[EdgeType, Module],
                  aggr: Optional[str] = "sum"):
         super().__init__()
+        self.validate_convs(convs)
 
         src_node_types = set([key[0] for key in convs.keys()])
         dst_node_types = set([key[-1] for key in convs.keys()])
@@ -59,6 +61,16 @@ def __init__(self, convs: Dict[EdgeType, Module],
         self.convs = ModuleDict({'__'.join(k): v for k, v in convs.items()})
         self.aggr = aggr
 
+    @staticmethod
+    def validate_convs(convs: Dict[EdgeType, Module]) -> None:
+        """
+        Checks to make sure that the convs provided for bipartite message
+        passing edges to not have an 'add_self_loops' argument that is set
+        to true. Raises :attr:`ValueError` if such a model exists.
+        """
+        for edge_type, module in convs.items():
+            check_add_self_loops(module, [edge_type])
+
     def reset_parameters(self):
         for conv in self.convs.values():
             conv.reset_parameters()

torch_geometric/nn/to_hetero_transformer.py

@@ -8,7 +8,10 @@
 
 from torch_geometric.nn.fx import Transformer
 from torch_geometric.typing import EdgeType, Metadata, NodeType
-from torch_geometric.utils.hetero import get_unused_node_types
+from torch_geometric.utils.hetero import (
+    check_add_self_loops,
+    get_unused_node_types,
+)
 
 try:
     from torch.fx import Graph, GraphModule, Node
@@ -132,6 +135,7 @@ def __init__(
         debug: bool = False,
     ):
         super().__init__(module, input_map, debug)
+        check_add_self_loops(module, metadata[1])
 
         unused_node_types = get_unused_node_types(*metadata)
         if len(unused_node_types) > 0:

torch_geometric/nn/to_hetero_with_bases_transformer.py

@@ -2,6 +2,7 @@
 import warnings
 from typing import Any, Dict, List, Optional, Union
 
+import pytest
 import torch
 from torch import Tensor
 from torch.nn import Module, Parameter
@@ -11,7 +12,10 @@
 from torch_geometric.nn.dense import Linear
 from torch_geometric.nn.fx import Transformer
 from torch_geometric.typing import EdgeType, Metadata, NodeType
-from torch_geometric.utils.hetero import get_unused_node_types
+from torch_geometric.utils.hetero import (
+    check_add_self_loops,
+    get_unused_node_types,
+)
 
 try:
     from torch.fx import Graph, GraphModule, Node
@@ -128,6 +132,7 @@ def forward(self, x, edge_index):
         debug (bool, optional): If set to :obj:`True`, will perform
             transformation in debug mode. (default: :obj:`False`)
     """
+
     transformer = ToHeteroWithBasesTransformer(module, metadata, num_bases,
                                                in_channels, input_map, debug)
     return transformer.transform()
@@ -144,6 +149,7 @@ def __init__(
         debug: bool = False,
     ):
         super().__init__(module, input_map, debug)
+        check_add_self_loops(module, metadata[1])
 
         unused_node_types = get_unused_node_types(*metadata)
         if len(unused_node_types) > 0:
@@ -547,3 +553,19 @@ def split_output(
 def key2str(key: Union[NodeType, EdgeType]) -> str:
     key = '__'.join(key) if isinstance(key, tuple) else key
     return key.replace(' ', '_').replace('-', '_').replace(':', '_')
+
+
+class MessagePassingLoops(MessagePassing):
+    def __init__(self, add_self_loops: bool = True):
+        super().__init__()
+        self.add_self_loops = add_self_loops
+
+    def forward(self):
+        pass
+
+
+def test_hetero_transformer_exception_self_loops():
+    model = MessagePassingLoops()
+    with pytest.raises(ValueError):
+        to_hetero_with_bases(model, (('a'), (('a', 'to', 'b'), )), 1)
+    to_hetero_with_bases(model, (('a'), (('a', 'to', 'a'), )), 1)

torch_geometric/utils/hetero.py

@@ -47,3 +47,15 @@ def get_unused_node_types(node_types: List[NodeType],
                           edge_types: List[EdgeType]) -> Set[NodeType]:
     dst_node_types = set(edge_type[-1] for edge_type in edge_types)
     return set(node_types) - set(dst_node_types)
+
+
+def check_add_self_loops(module: torch.nn.Module, edge_types: List[EdgeType]):
+    edge_types = [
+        edge_type for edge_type in edge_types if edge_type[0] != edge_type[-1]
+    ]
+    if len(edge_types) > 0 and hasattr(
+            module, "add_self_loops") and module.add_self_loops:
+        raise ValueError(
+            f"'add_self_loops' attribute set to 'True' on module {module} "
+            f"for use with edge type(s) '{edge_types}'. This will lead to "
+            f"incorrect message passing results.")