Repository: pyg-team/pytorch_geometric
Pull request #8344 (2023-11-08 16:11:55)

Support `torch.compile` in `HeteroConv`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 881e376fe..e12e5242f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for `torch.compile` in `HeteroConv` ([#8344](https://github.com/pyg-team/pytorch_geometric/pull/8344))
 - Added support for weighted `sparse_cross_entropy` ([#8340](https://github.com/pyg-team/pytorch_geometric/pull/8340))
 - Added a multi GPU training benchmarks for XPU device ([#8288](https://github.com/pyg-team/pytorch_geometric/pull/8288))
 - Support MRR computation in `KGEModel.test()` ([#8298](https://github.com/pyg-team/pytorch_geometric/pull/8298))
diff --git a/torch_geometric/nn/conv/hetero_conv.py b/torch_geometric/nn/conv/hetero_conv.py
index 82cbe14d1..5b43dbc6d 100644
--- a/torch_geometric/nn/conv/hetero_conv.py
+++ b/torch_geometric/nn/conv/hetero_conv.py
@@ -1,5 +1,4 @@
 import warnings
-from collections import defaultdict
 from typing import Dict, List, Optional
 
 import torch
@@ -114,9 +113,9 @@ class HeteroConv(torch.nn.Module):
                 :meth:`~torch_geometric.nn.conv.HeteroConv.forward` via
                 :obj:`edge_attr_dict = { edge_type: edge_attr }`.
         """
-        out_dict = defaultdict(list)
+        out_dict: Dict[str, List[Tensor]] = {}
 
-        for edge_type in self.convs.keys():
+        for edge_type, conv in self.convs.items():
             src, rel, dst = edge_type
 
             has_edge_level_arg = False
@@ -156,8 +155,12 @@ class HeteroConv(torch.nn.Module):
             if not has_edge_level_arg:
                 continue
 
-            out = self.convs[edge_type](*args, **kwargs)
-            out_dict[dst].append(out)
+            out = conv(*args, **kwargs)
+
+            if dst not in out_dict:
+                out_dict[dst] = [out]
+            else:
+                out_dict[dst].append(out)
 
         for key, value in out_dict.items():
             out_dict[key] = group(value, self.aggr)


Test patch:
diff --git a/test/nn/conv/test_hetero_conv.py b/test/nn/conv/test_hetero_conv.py
index 7eb43f485..358fe8163 100644
--- a/test/nn/conv/test_hetero_conv.py
+++ b/test/nn/conv/test_hetero_conv.py
@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+import torch_geometric
 from torch_geometric.data import HeteroData
 from torch_geometric.nn import (
     GATConv,
@@ -11,7 +12,13 @@ from torch_geometric.nn import (
     MessagePassing,
     SAGEConv,
 )
-from torch_geometric.testing import get_random_edge_index
+from torch_geometric.testing import (
+    disableExtensions,
+    get_random_edge_index,
+    onlyLinux,
+    withCUDA,
+    withPackage,
+)
 
 
 @pytest.mark.parametrize('aggr', ['sum', 'mean', 'min', 'max', 'cat', None])
@@ -167,3 +174,34 @@ def test_hetero_conv_with_dot_syntax_node_types():
     assert len(out_dict) == 2
     assert out_dict['src.paper'].size() == (50, 64)
     assert out_dict['author'].size() == (30, 64)
+
+
+@withCUDA
+@onlyLinux
+@disableExtensions
+@withPackage('torch>=2.1.0')
+def test_compile_hetero_conv_graph_breaks(device):
+    import torch._dynamo as dynamo
+
+    data = HeteroData()
+    data['a'].x = torch.randn(50, 16, device=device)
+    data['b'].x = torch.randn(50, 16, device=device)
+    edge_index = get_random_edge_index(50, 50, 100, device=device)
+    data['a', 'to', 'b'].edge_index = edge_index
+    data['b', 'to', 'a'].edge_index = edge_index.flip([0])
+
+    conv = HeteroConv({
+        ('a', 'to', 'b'): SAGEConv(16, 32).jittable(),
+        ('b', 'to', 'a'): SAGEConv(16, 32).jittable(),
+    }).to(device)
+
+    explanation = dynamo.explain(conv)(data.x_dict, data.edge_index_dict)
+    assert explanation.graph_break_count == 0
+
+    compiled_conv = torch_geometric.compile(conv)
+
+    expected = conv(data.x_dict, data.edge_index_dict)
+    out = compiled_conv(data.x_dict, data.edge_index_dict)
+    assert len(out) == len(expected)
+    for key in expected.keys():
+        assert torch.allclose(out[key], expected[key])