Repository: pyg-team/pytorch_geometric

Issue:
TypeError: global_add_pool() missing 1 required positional argument: 'batch'
### 🐛 Describe the bug

The documentation of `global_add_pool ` states that the batch parameter is optional, however I get this error ```TypeError: global_add_pool() missing 1 required positional argument: 'batch'``` when I don't pass a value for the batch. 

Here is a snippet of code to reproduce the bug:
```Python
import torch
from torch_geometric.nn import global_add_pool
x = torch.zeros(4,5)
z = global_add_pool(x)
```

I think the origin of the issue is the absence of a default value for the batch parameter in the definition of `global_add_pool`.
And the same bug happens when using `global_mean_pool` or `global_max_pool`.

### Environment

* PyG version: 2.0.5
* PyTorch version: 1.11.0
* OS:  Fedora release 35
* Python version: 3.9
* CUDA/cuDNN version: 11.3
* How you installed PyTorch and PyG (`conda`, `pip`, source): Poetry



Discussion:
Yes, this is intended. I feel that `global_add_pool(x)` can lead to some serious bugs when not being careful, as such:
```python
def forward(x, edge_index, batch=None):
    x = global_ad__pool(x, batch)
```
is a good and more robust workaround. WDYT?
I agree with being careful with the `global_add_pool` and the workaround looks good.
I have two additional remarks:

- The current documentation of `global_add_pool` doesn't reflect that `batch` is  a mandatory parameter:
```global_add_pool(x: Tensor, batch: Optional[Tensor], size: Optional[int] = None) -> Tensor```
- When the batch is None it is likely because the node features matrix is a dense  node feature tensor $X \in \mathbb{R}^{B \times N_{max} \times F}$. In this case the dimension to reduce is dim=1, however this is not the behavior of `global_add_pool`:
```python
if batch is None:
    return x.sum(dim=0, keepdim=True)
```

I need to use `global_add_pool` after applying dense pooling layer such as diffpool. Although `torch.sum` can do the job, I prefer to have one function for global add pool to use with both sparse and dense batches. I think such function will help build modular pipelines such the one in GraphGym.
I think both suggestions are great. We could adjust the documentation to specify that batch is mendatory? And similarly for other global pooling functions too.

The ability to handle both sparse and dense batches sounds great too. As you said this seems to be easy with torch sum. Should we have a PR on this?
We should be able to do this by doing a dense aggregation in `dim=-2` instead of `dim=0`. Happy to accept a PR on this one.

Resolving patch:
diff --git a/torch_geometric/nn/glob/glob.py b/torch_geometric/nn/glob/glob.py
--- a/torch_geometric/nn/glob/glob.py
+++ b/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):


Test patch:
diff --git a/test/nn/glob/test_glob.py b/test/nn/glob/test_glob.py
--- a/test/nn/glob/test_glob.py
+++ b/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))