Repository: pyg-team/pytorch_geometric
Pull request #7198 (2023-04-18 05:11:05)

Support for `scatter(..., reduce="any")`

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7198?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7198](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7198?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (5b84111) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/7c39ad55f155b18f3e74de14ce7c92a57395c055?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (7c39ad5) will **decrease** coverage by `0.03%`.
> The diff coverage is `33.33%`.

```diff
@@            Coverage Diff             @@
##           master    #7198      +/-   ##
==========================================
- Coverage   91.65%   91.62%   -0.03%     
==========================================
  Files         436      436              
  Lines       23933    23945      +12     
==========================================
+ Hits        21936    21940       +4     
- Misses       1997     2005       +8     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7198?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/coalesce.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7198?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2NvYWxlc2NlLnB5) | `100.00% <ø> (ø)` | |
| [torch\_geometric/utils/scatter.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7198?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3NjYXR0ZXIucHk=) | `52.17% <33.33%> (-20.56%)` | :arrow_down: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a99d47ed2..0c318546d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for `"any"`-reductions in `scatter` ([#7198](https://github.com/pyg-team/pytorch_geometric/pull/7198))
 - Added manual sampling interface to `NodeLoader` and `LinkLoader` ([#7197](https://github.com/pyg-team/pytorch_geometric/pull/7197))
 - Extending `torch.sparse` support ([#7155](https://github.com/pyg-team/pytorch_geometric/pull/7155))
 - Added edge weight support to `LightGCN` ([#7157](https://github.com/pyg-team/pytorch_geometric/pull/7157))
diff --git a/torch_geometric/utils/coalesce.py b/torch_geometric/utils/coalesce.py
index 80354b34c..92b54643e 100644
--- a/torch_geometric/utils/coalesce.py
+++ b/torch_geometric/utils/coalesce.py
@@ -50,7 +50,7 @@ def coalesce(
             :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
         reduce (str, optional): The reduce operation to use for merging edge
             features (:obj:`"add"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
-            :obj:`"mul"`). (default: :obj:`"add"`)
+            :obj:`"mul"`, :obj:`"any"`). (default: :obj:`"add"`)
         is_sorted (bool, optional): If set to :obj:`True`, will expect
             :obj:`edge_index` to be already sorted row-wise.
         sort_by_row (bool, optional): If set to :obj:`False`, will sort
diff --git a/torch_geometric/utils/scatter.py b/torch_geometric/utils/scatter.py
index 5ced98e72..de0457662 100644
--- a/torch_geometric/utils/scatter.py
+++ b/torch_geometric/utils/scatter.py
@@ -15,11 +15,6 @@ if has_pytorch112:  # pragma: no cover
 
     warnings.filterwarnings('ignore', '.*is in beta and the API may change.*')
 
-    def broadcast(src: Tensor, ref: Tensor, dim: int) -> Tensor:
-        size = [1] * ref.dim()
-        size[dim] = -1
-        return src.view(size).expand_as(ref)
-
     def scatter(src: Tensor, index: Tensor, dim: int = 0,
                 dim_size: Optional[int] = None, reduce: str = 'sum') -> Tensor:
         r"""Reduces all values from the :obj:`src` tensor at the indices
@@ -39,8 +34,8 @@ if has_pytorch112:  # pragma: no cover
                 minimal-sized output tensor according to
                 :obj:`index.max() + 1`. (default: :obj:`None`)
             reduce (str, optional): The reduce operation (:obj:`"sum"`,
-                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`).
-                (default: :obj:`"sum"`)
+                :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`,
+                :obj:`"any"`). (default: :obj:`"sum"`)
         """
         if index.dim() != 1:
             raise ValueError(f"The `index` argument must be one-dimensional "
@@ -68,6 +63,11 @@ if has_pytorch112:  # pragma: no cover
         size = list(src.size())
         size[dim] = dim_size
 
+        # For "any" reduction, we use regular `scatter_`:
+        if reduce == 'any':
+            index = broadcast(index, src, dim)
+            return src.new_zeros(size).scatter_(dim, index, src)
+
         # For "sum" and "mean" reduction, we make use of `scatter_add_`:
         if reduce == 'sum' or reduce == 'add':
             index = broadcast(index, src, dim)
@@ -145,7 +145,25 @@ else:  # pragma: no cover
                 :obj:`"mean"`, :obj:`"mul"`, :obj:`"min"` or :obj:`"max"`).
                 (default: :obj:`"sum"`)
         """
+        if reduce == 'any':
+            dim = src.dim() + dim if dim < 0 else dim
+
+            if dim_size is None:
+                dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
+
+            size = list(src.size())
+            size[dim] = dim_size
+
+            index = broadcast(index, src, dim)
+            return src.new_zeros(size).scatter_(dim, index, src)
+
         if not torch_geometric.typing.WITH_TORCH_SCATTER:
             raise ImportError("'scatter' requires the 'torch-scatter' package")
         return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
                                      reduce=reduce)
+
+
+def broadcast(src: Tensor, ref: Tensor, dim: int) -> Tensor:
+    size = [1] * ref.dim()
+    size[dim] = -1
+    return src.view(size).expand_as(ref)


Test patch:
diff --git a/test/utils/test_scatter.py b/test/utils/test_scatter.py
index b5b608725..735a5cee1 100644
--- a/test/utils/test_scatter.py
+++ b/test/utils/test_scatter.py
@@ -58,6 +58,18 @@ def test_scatter_backward(reduce, device):
     assert src.grad is not None
 
 
+@withCUDA
+def test_scatter_any(device):
+    src = torch.randn(6, 4, device=device)
+    index = torch.tensor([0, 0, 1, 1, 2, 2], device=device)
+
+    out = scatter(src, index, dim=0, reduce='any')
+
+    for i in range(3):
+        for j in range(4):
+            assert float(out[i, j]) in src[2 * i:2 * i + 2, j].tolist()
+
+
 if __name__ == '__main__':
     # Insights on GPU:
     # ================