Merge pull request #483 from fdiehl/globalpool_batch_order

Allowing for permuted batch references for global pool operations

test/nn/glob/test_glob.py

@@ -1,4 +1,5 @@
 import torch
+import pytest
 from torch_geometric.nn import (global_add_pool, global_mean_pool,
                                 global_max_pool)
 
@@ -22,3 +23,28 @@ def test_global_pool():
     assert out.size() == (2, 4)
     assert out[0].tolist() == x[:4].max(dim=0)[0].tolist()
     assert out[1].tolist() == x[4:].max(dim=0)[0].tolist()
+
+
+def test_permuted_global_pool():
+    N_1, N_2 = 4, 6
+    x = torch.randn(N_1 + N_2, 4)
+    batch = torch.tensor([0 for _ in range(N_1)] + [1 for _ in range(N_2)])
+
+    permutation = torch.randperm(N_1 + N_2)
+    perm_x = x[permutation]
+    perm_batch = batch[permutation]
+
+    out = global_add_pool(perm_x, perm_batch)
+    assert out.size() == (2, 4)
+    assert out[0].tolist() == pytest.approx(x[:4].sum(dim=0).tolist())
+    assert out[1].tolist() == pytest.approx(x[4:].sum(dim=0).tolist())
+
+    out = global_mean_pool(perm_x, perm_batch)
+    assert out.size() == (2, 4)
+    assert out[0].tolist() == pytest.approx(x[:4].mean(dim=0).tolist())
+    assert out[1].tolist() == pytest.approx(x[4:].mean(dim=0).tolist())
+
+    out = global_max_pool(perm_x, perm_batch)
+    assert out.size() == (2, 4)
+    assert out[0].tolist() == pytest.approx(x[:4].max(dim=0)[0].tolist())
+    assert out[1].tolist() == pytest.approx(x[4:].max(dim=0)[0].tolist())

torch_geometric/nn/glob/glob.py

@@ -1,3 +1,4 @@
+import torch
 from torch_geometric.utils import scatter_
 
 
@@ -20,7 +21,7 @@ def global_add_pool(x, batch, size=None):
     :rtype: :class:`Tensor`
     """
 
-    size = batch[-1].item() + 1 if size is None else size
+    size = torch.max(batch) + 1 if size is None else size
     return scatter_('add', x, batch, dim_size=size)
 
 
@@ -43,7 +44,7 @@ def global_mean_pool(x, batch, size=None):
     :rtype: :class:`Tensor`
     """
 
-    size = batch[-1].item() + 1 if size is None else size
+    size = torch.max(batch) + 1 if size is None else size
     return scatter_('mean', x, batch, dim_size=size)
 
 
@@ -66,5 +67,5 @@ def global_max_pool(x, batch, size=None):
     :rtype: :class:`Tensor`
     """
 
-    size = batch[-1].item() + 1 if size is None else size
+    size = torch.max(batch) + 1 if size is None else size
     return scatter_('max', x, batch, dim_size=size)