Repository: pyg-team/pytorch_geometric
Pull request #9598 (2024-08-16 18:11:35)

Destroy process group in `graphgym` test
When `test_graphgym.py::test_train` is executed, a process group is created.  If this group is not properly destroyed, it can interfere with subsequent tests, causing failures such as:
```
FAILED ../../pytorch_geometric/test/metrics/test_link_pred_metric.py::test_precision[1-32-3000-1000-100] - RuntimeError: No backend type associated with device type cpu
FAILED ../../pytorch_geometric/test/metrics/test_link_pred_metric.py::test_precision[10-32-3000-1000-100] - RuntimeError: No backend type associated with device type cpu
FAILED ../../pytorch_geometric/test/metrics/test_link_pred_metric.py::test_precision[100-32-3000-1000-100] - RuntimeError: No backend type associated with device type cpu
FAILED ../../pytorch_geometric/test/metrics/test_link_pred_metric.py::test_recall - RuntimeError: No backend type associated with device type cpu
FAILED ../../pytorch_geometric/test/metrics/test_link_pred_metric.py::test_f1 - RuntimeError: No backend type associated with device type cpu
FAILED ../../pytorch_geometric/test/metrics/test_link_pred_metric.py::test_map - RuntimeError: No backend type associated with device type cpu
FAILED ../../pytorch_geometric/test/metrics/test_link_pred_metric.py::test_ndcg - RuntimeError: No backend type associated with device type CPU
```
A similar bug was recently fixed by [PR#9534](https://github.com/pyg-team/pytorch_geometric/pull/9534). 


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f22c4893b..ecc716aba 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -65,6 +65,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `VirtualNode` transform for empty edge indices ([#9605](https://github.com/pyg-team/pytorch_geometric/pull/9605))
 - Fixed an issue where import order in the multi-GPU `cugraph` example could cause an `rmm` error ([#9577](https://github.com/pyg-team/pytorch_geometric/pull/9577))
 - Made the output of the single-GPU `cugraph` example more readable ([#9577](https://github.com/pyg-team/pytorch_geometric/pull/9577))
 - Fixed `load_state_dict` behavior with lazy parameters in `HeteroDictLinear` ([#9493](https://github.com/pyg-team/pytorch_geometric/pull/9493))
diff --git a/torch_geometric/transforms/virtual_node.py b/torch_geometric/transforms/virtual_node.py
index 3f7a59170..3bb8297de 100644
--- a/torch_geometric/transforms/virtual_node.py
+++ b/torch_geometric/transforms/virtual_node.py
@@ -37,7 +37,8 @@ class VirtualNode(BaseTransform):
         col = torch.cat([col, full, arange], dim=0)
         edge_index = torch.stack([row, col], dim=0)
 
-        new_type = edge_type.new_full((num_nodes, ), int(edge_type.max()) + 1)
+        num_edge_types = int(edge_type.max()) if edge_type.numel() > 0 else 0
+        new_type = edge_type.new_full((num_nodes, ), num_edge_types + 1)
         edge_type = torch.cat([edge_type, new_type, new_type + 1], dim=0)
 
         old_data = copy.copy(data)


Test patch:
diff --git a/test/data/lightning/test_datamodule.py b/test/data/lightning/test_datamodule.py
index 279782272..cb86d810d 100644
--- a/test/data/lightning/test_datamodule.py
+++ b/test/data/lightning/test_datamodule.py
@@ -290,7 +290,8 @@ class LinearHeteroNodeModule(LightningModule):
 def preserve_context():
     num_threads = torch.get_num_threads()
     yield
-    torch.distributed.destroy_process_group()
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
     torch.set_num_threads(num_threads)
 
 
diff --git a/test/graphgym/test_graphgym.py b/test/graphgym/test_graphgym.py
index 321c31bdf..894bc3dce 100644
--- a/test/graphgym/test_graphgym.py
+++ b/test/graphgym/test_graphgym.py
@@ -163,10 +163,17 @@ def test_graphgym_module(tmp_path):
     assert isinstance(outputs["loss"], torch.Tensor)
 
 
+@pytest.fixture
+def destroy_process_group():
+    yield
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
+
+
 @onlyOnline
 @onlyLinux
 @withPackage('yacs', 'pytorch_lightning')
-def test_train(tmp_path, capfd):
+def test_train(destroy_process_group, tmp_path, capfd):
     warnings.filterwarnings('ignore', ".*does not have many workers.*")
 
     import pytorch_lightning as pl
diff --git a/test/transforms/test_virtual_node.py b/test/transforms/test_virtual_node.py
index 7111c771a..35517eb76 100644
--- a/test/transforms/test_virtual_node.py
+++ b/test/transforms/test_virtual_node.py
@@ -36,3 +36,17 @@ def test_virtual_node():
     assert data.num_nodes == 5
 
     assert data.edge_type.tolist() == [0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
+
+    data = Data(x=x, edge_index=torch.empty(2, 0, dtype=torch.long))
+    data = VirtualNode()(data)
+    assert len(data) == 3
+
+    assert data.x.size() == (5, 16)
+    assert torch.allclose(data.x[:4], x)
+    assert data.x[4:].abs().sum() == 0
+
+    assert data.edge_index.tolist() == [
+        [0, 1, 2, 3, 4, 4, 4, 4],
+        [4, 4, 4, 4, 0, 1, 2, 3],
+    ]
+    assert data.edge_type.tolist() == [1, 1, 1, 1, 2, 2, 2, 2]