Repository: pyg-team/pytorch_geometric
Pull request: Add `reduce` argument to `nn.pool.pool.pool_edge`
To specify edge_attr reduction scheme (issue #9107)

Problem statement:
Add edge_attr reduction argument to torch_geometric.nn.pool.pool.pool_edge
### 🚀 The feature, motivation and pitch

I am working on my own graph pooling scheme based on cluster vectors assigning nodes to the cluster they get pooled into. To pool my edges and edge_attributes, I am using `pool_edge` from `torch_geometric/nn/pool/pool.py`. However, I would want to be able to choose the reduce operation that will be used in `coalesce`. As of now, it is hardcoded to the default value `'sum'`.

### Alternatives

The only solution I see is to rewrite the `pool_edge` function identically and adding an argument `reduce` that will be passed to `coalesce`.

### Additional context

_No response_

Hints from issue discussion:
That's a good request. Do you mind sending a PR for that?

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/torch_geometric/nn/pool/pool.py b/torch_geometric/nn/pool/pool.py
index da076bc1cbef..22b5e134ee66 100644
--- a/torch_geometric/nn/pool/pool.py
+++ b/torch_geometric/nn/pool/pool.py
@@ -5,12 +5,14 @@
 from torch_geometric.utils import coalesce, remove_self_loops, scatter
 
 
-def pool_edge(cluster, edge_index, edge_attr: Optional[torch.Tensor] = None):
+def pool_edge(cluster, edge_index, edge_attr: Optional[torch.Tensor] = None,
+              reduce: Optional[str] = "add"):
     num_nodes = cluster.size(0)
     edge_index = cluster[edge_index.view(-1)].view(2, -1)
     edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
     if edge_index.numel() > 0:
-        edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes)
+        edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes,
+                                         reduce=reduce)
     return edge_index, edge_attr
 
 


Code review comments:
diff_hunk: @@ -5,12 +5,14 @@
 from torch_geometric.utils import coalesce, remove_self_loops, scatter
 
 
-def pool_edge(cluster, edge_index, edge_attr: Optional[torch.Tensor] = None):
+def pool_edge(cluster, edge_index, edge_attr: Optional[torch.Tensor] = None,
+              reduce: Optional[str] = "add"):
path: torch_geometric/nn/pool/pool.py
text: @user1:
```suggestion
              reduce: Optional[str] = 'sum', ):
```