Repository: pyg-team/pytorch_geometric
Pull request #7948 (2023-08-29 06:57:44)

Add `group_argsort` implementation

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7948?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7948](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7948?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (aaa4a76) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/ccf6706c9f63911667ba88647889ea0638f54019?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (ccf6706) will **decrease** coverage by `0.74%`.
> Report is 1 commits behind head on master.
> The diff coverage is `94.73%`.

> :exclamation: Current head aaa4a76 differs from pull request most recent head 48435ee. Consider uploading reports for the commit 48435ee to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7948      +/-   ##
==========================================
- Coverage   90.21%   89.48%   -0.74%     
==========================================
  Files         459      459              
  Lines       26922    26936      +14     
==========================================
- Hits        24289    24103     -186     
- Misses       2633     2833     +200     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7948?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/scatter.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3NjYXR0ZXIucHk=) | `88.63% <94.44%> (-3.96%)` | :arrow_down: |
| [torch\_geometric/utils/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7948?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |

... and [31 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7948/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

: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 3f210d30e..cf1b6dbcd 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 `group_argsort` implementation ([#7948](https://github.com/pyg-team/pytorch_geometric/pull/7948))
 - Added `CachedLoader` implementation ([#7896](https://github.com/pyg-team/pytorch_geometric/pull/7896), [#7897](https://github.com/pyg-team/pytorch_geometric/pull/7897))
 - Added possibility to run training benchmarks on XPU device ([#7925](https://github.com/pyg-team/pytorch_geometric/pull/7925))
 - Added `utils.ppr` for personalized PageRank computation ([#7917](https://github.com/pyg-team/pytorch_geometric/pull/7917))
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
index 46cc541ae..30c648a61 100644
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -1,6 +1,6 @@
 import copy
 
-from .scatter import scatter
+from .scatter import scatter, group_argsort
 from .segment import segment
 from .sort import index_sort
 from .degree import degree
@@ -56,6 +56,7 @@ from .train_test_split_edges import train_test_split_edges
 
 __all__ = [
     'scatter',
+    'group_argsort',
     'segment',
     'index_sort',
     'degree',
diff --git a/torch_geometric/utils/scatter.py b/torch_geometric/utils/scatter.py
index 7a87aa742..e1c314b15 100644
--- a/torch_geometric/utils/scatter.py
+++ b/torch_geometric/utils/scatter.py
@@ -170,8 +170,9 @@ def scatter_argmax(src: Tensor, index: Tensor, dim: int = 0,
         return out[1]
 
     # Only implemented under certain conditions for now :(
-    assert dim == 0
     assert src.dim() == 1 and index.dim() == 1
+    assert dim == 0 or dim == -1
+    assert src.numel() == index.numel()
 
     if dim_size is None:
         dim_size = index.max() + 1 if index.numel() > 0 else 0
@@ -191,3 +192,65 @@ def scatter_argmax(src: Tensor, index: Tensor, dim: int = 0,
     out[index[nonzero]] = nonzero
 
     return out
+
+
+def group_argsort(
+    src: Tensor,
+    index: Tensor,
+    dim: int = 0,
+    num_groups: Optional[int] = None,
+    descending: bool = False,
+    return_consecutive: bool = False,
+    stable: bool = False,
+) -> Tensor:
+    r"""Returns the indices that sort the tensor :obj:`src` along a given
+    dimension in ascending order by value.
+    In contrast to :meth:`torch.argsort`, sorting is performed in groups
+    according to the values in :obj:`index`.
+
+    Args:
+        src (torch.Tensor): The source tensor.
+        index (torch.Tensor): The index tensor.
+        dim (int, optional): The dimension along which to index.
+            (default: :obj:`0`)
+        num_groups (int, optional): The number of groups.
+            (default: :obj:`None`)
+        descending (bool, optional): Controls the sorting order (ascending or
+            descending). (default: :obj:`False`)
+        return_consecutive (bool, optional): If set to :obj:`True`, will not
+            offset the output to start from :obj:`0` for each group.
+            (default: :obj:`False`)
+        stable (bool, optional): Controls the relative order of equivalent
+            elements. (default: :obj:`False`)
+    """
+    # Only implemented under certain conditions for now :(
+    assert src.dim() == 1 and index.dim() == 1
+    assert dim == 0 or dim == -1
+    assert src.numel() == index.numel() and src.numel() > 0
+
+    # Normalize `src` to range [0, 1]:
+    src = src - src.min()
+    src = src / src.max()
+
+    # Compute `grouped_argsort`:
+    src = src - 2 * index if descending else src + 2 * index
+    if torch_geometric.typing.WITH_PT113:
+        perm = src.argsort(descending=descending, stable=stable)
+    else:
+        perm = src.argsort(descending=descending)
+        if stable:
+            warnings.warn("Ignoring option `stable=True` in 'group_argsort' "
+                          "since it requires PyTorch >= 1.13.0")
+    out = torch.empty_like(index)
+    out[perm] = torch.arange(index.numel(), device=index.device)
+
+    if return_consecutive:
+        return out
+
+    # Compute cumulative sum of number of entries with the same index:
+    count = scatter(torch.ones_like(index), index, dim=dim,
+                    dim_size=num_groups, reduce='sum')
+    ptr = count.new_zeros(count.numel() + 1)
+    torch.cumsum(count, dim=0, out=ptr[1:])
+
+    return out - ptr[index]


Test patch:
diff --git a/test/utils/test_scatter.py b/test/utils/test_scatter.py
index 44e132c1c..658a7ac80 100644
--- a/test/utils/test_scatter.py
+++ b/test/utils/test_scatter.py
@@ -5,7 +5,7 @@ import torch
 
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import disableExtensions, withCUDA, withPackage
-from torch_geometric.utils import scatter
+from torch_geometric.utils import group_argsort, scatter
 from torch_geometric.utils.scatter import scatter_argmax
 
 
@@ -74,6 +74,26 @@ def test_scatter_any(device):
             assert float(out[i, j]) in src[2 * i:2 * i + 2, j].tolist()
 
 
+@withCUDA
+@pytest.mark.parametrize('num_groups', [4])
+@pytest.mark.parametrize('descending', [False, True])
+def test_group_argsort(num_groups, descending, device):
+    src = torch.randn(20, device=device)
+    index = torch.randint(0, num_groups, (20, ), device=device)
+
+    out = group_argsort(src, index, 0, num_groups, descending=descending)
+
+    expected = torch.empty_like(index)
+    for i in range(num_groups):
+        mask = index == i
+        tmp = src[mask].argsort(descending=descending)
+        perm = torch.empty_like(tmp)
+        perm[tmp] = torch.arange(tmp.numel(), device=device)
+        expected[mask] = perm
+
+    assert torch.equal(out, expected)
+
+
 @withCUDA
 @disableExtensions
 def test_scatter_argmax(device):