Support for dense aggregations in global_*_pool (#4827)

* global dense pool

* update

Author: rusty1s

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 support for dense aggregations in `global_*_pool` ([#4827](https://github.com/pyg-team/pytorch_geometric/pull/4827))
 - Added Python version requirement ([#4825](https://github.com/pyg-team/pytorch_geometric/pull/4825))
 - Added TorchScript support to `JumpingKnowledge` module ([#4805](https://github.com/pyg-team/pytorch_geometric/pull/4805))
 - Added a `max_sample` argument to `AddMetaPaths` in order to tackle very dense metapath edges ([#4750](https://github.com/pyg-team/pytorch_geometric/pull/4750))

test/nn/glob/test_glob.py

@@ -65,3 +65,8 @@ def test_permuted_global_pool():
     assert out.size() == (2, 4)
     assert torch.allclose(out[0], px1.max(dim=0)[0])
     assert torch.allclose(out[1], px2.max(dim=0)[0])
+
+
+def test_dense_global_pool():
+    x = torch.randn(3, 16, 32)
+    assert torch.allclose(global_add_pool(x, None), x.sum(dim=1))

torch_geometric/nn/glob/glob.py

@@ -24,9 +24,9 @@ def global_add_pool(x: Tensor, batch: Optional[Tensor],
             Automatically calculated if not given. (default: :obj:`None`)
     """
     if batch is None:
-        return x.sum(dim=0, keepdim=True)
+        return x.sum(dim=-2, keepdim=x.dim() == 2)
     size = int(batch.max().item() + 1) if size is None else size
-    return scatter(x, batch, dim=0, dim_size=size, reduce='add')
+    return scatter(x, batch, dim=-2, dim_size=size, reduce='add')
 
 
 def global_mean_pool(x: Tensor, batch: Optional[Tensor],
@@ -48,9 +48,9 @@ def global_mean_pool(x: Tensor, batch: Optional[Tensor],
             Automatically calculated if not given. (default: :obj:`None`)
     """
     if batch is None:
-        return x.mean(dim=0, keepdim=True)
+        return x.mean(dim=-2, keepdim=x.dim() == 2)
     size = int(batch.max().item() + 1) if size is None else size
-    return scatter(x, batch, dim=0, dim_size=size, reduce='mean')
+    return scatter(x, batch, dim=-2, dim_size=size, reduce='mean')
 
 
 def global_max_pool(x: Tensor, batch: Optional[Tensor],
@@ -72,9 +72,9 @@ def global_max_pool(x: Tensor, batch: Optional[Tensor],
             Automatically calculated if not given. (default: :obj:`None`)
     """
     if batch is None:
-        return x.max(dim=0, keepdim=True)[0]
+        return x.max(dim=-2, keepdim=x.dim() == 2)[0]
     size = int(batch.max().item() + 1) if size is None else size
-    return scatter(x, batch, dim=0, dim_size=size, reduce='max')
+    return scatter(x, batch, dim=-2, dim_size=size, reduce='max')
 
 
 class GlobalPooling(torch.nn.Module):