Repository: pyg-team/pytorch_geometric
Pull request #9029 (2024-03-08 00:22:22)

Add `group_cat` functionality
Concatenates the given sequence of `tensors` in the given dimension `dim`. Differ from `torch.cat`, values along the concatenating dimension are grouped according to the indicies defined in the `index` tensors. 

All tensors must either have the same shape (except in the concatenating dimension) or be empty and all index tensors must in the same sequence as in `tensors`.

Examples:
```
>>> x1 = torch.tensor([[0.2716, 0.4233, 0.2658, 0.8284],
...                    [0.3166, 0.0142, 0.1700, 0.2944],
...                    [0.2371, 0.3839, 0.7193, 0.2954],
...                    [0.4100, 0.0012, 0.5114, 0.3353]])
>>> x2 = torch.tensor([[0.3752, 0.5782, 0.7105, 0.4002],
...                    [0.7757, 0.5999, 0.7898, 0.0753]])

>>> x1_index = torch.LongTensor([0,0,1,2])
>>> x2_index = torch.LongTensor([0,2])

>>> scatter_concat([x1,x2], [x1_index,x2_index], 0)
tensor([[0.2716, 0.4233, 0.2658, 0.8284],
              [0.3166, 0.0142, 0.1700, 0.2944],
              [0.3752, 0.5782, 0.7105, 0.4002],
              [0.2371, 0.3839, 0.7193, 0.2954],
              [0.4100, 0.0012, 0.5114, 0.3353],
              [0.7757, 0.5999, 0.7898, 0.0753]])
```

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9029?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: Patch coverage is `88.88889%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 89.24%. Comparing base [(`af0f5f4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/af0f5f44d4dfa8accd898df3e8523d06b67940b0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`4780b12`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9029?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).

> :exclamation: Current head 4780b12 differs from pull request most recent head 9528425. Consider uploading reports for the commit 9528425 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9029?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/utils/\_scatter.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9029?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL19zY2F0dGVyLnB5) | 87.50% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9029?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #9029      +/-   ##
==========================================
- Coverage   89.86%   89.24%   -0.63%     
==========================================
  Files         468      468              
  Lines       30039    30041       +2     
==========================================
- Hits        26995    26809     -186     
- Misses       3044     3232     +188     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9029?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?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 1ede8d9a4..2bea8d940 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_cat` functionality ([#9029](https://github.com/pyg-team/pytorch_geometric/pull/9029))
 - Added support for `EdgeIndex` in `spmm` ([#9026](https://github.com/pyg-team/pytorch_geometric/pull/9026))
 - Added option to pre-allocate memory in GPU-based `ApproxKNN` ([#9046](https://github.com/pyg-team/pytorch_geometric/pull/9046))
 - Added support for `EdgeIndex` in `MessagePassing` ([#9007](https://github.com/pyg-team/pytorch_geometric/pull/9007))
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
index e00d169e4..1cf66aa33 100644
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -2,7 +2,7 @@ r"""Utility package."""
 
 import copy
 
-from ._scatter import scatter, group_argsort
+from ._scatter import scatter, group_argsort, group_cat
 from ._segment import segment
 from ._index_sort import index_sort
 from .functions import cumsum
@@ -60,6 +60,7 @@ from ._train_test_split_edges import train_test_split_edges
 __all__ = [
     'scatter',
     'group_argsort',
+    'group_cat',
     'segment',
     'index_sort',
     'cumsum',
diff --git a/torch_geometric/utils/_scatter.py b/torch_geometric/utils/_scatter.py
index 9f738b8b7..f04bd40c4 100644
--- a/torch_geometric/utils/_scatter.py
+++ b/torch_geometric/utils/_scatter.py
@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Sequence, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -283,3 +283,47 @@ def group_argsort(
     ptr = cumsum(count)
 
     return out - ptr[index]
+
+
+def group_cat(
+    tensors: Sequence[Tensor],
+    index: Sequence[Tensor],
+    dim: int = 0,
+    return_index: bool = False,
+) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+    r"""Concatenates the given sequence of tensors :obj:`tensors` in the given
+    dimension :obj:`dim`.
+    Different from :meth:`torch.cat`, values along the concatenating dimension
+    are grouped according to the indicies defined in the :obj:`index` tensors.
+    All tensors must have the same shape (except in the concatenating
+    dimension).
+
+    Args:
+        tensors ([Tensor]): Sequence of tensors.
+        index ([Tensor]): Sequence of index tensors.
+        dim (int, optional): The dimension along which the tensors are
+            concatenated. (default: :obj:`0`)
+        return_index (bool, optional): If set to :obj:`True`, will return the
+            new index tensor. (default: :obj:`False`)
+
+    Example:
+        >>> x1 = torch.tensor([[0.2716, 0.4233],
+        ...                    [0.3166, 0.0142],
+        ...                    [0.2371, 0.3839],
+        ...                    [0.4100, 0.0012]])
+        >>> x2 = torch.tensor([[0.3752, 0.5782],
+        ...                    [0.7757, 0.5999]])
+        >>> index1 = torch.tensor([0, 0, 1, 2])
+        >>> index2 = torch.tensor([0, 2])
+        >>> scatter_concat([x1,x2], [index1, index2], dim=0)
+        tensor([[0.2716, 0.4233],
+                [0.3166, 0.0142],
+                [0.3752, 0.5782],
+                [0.2371, 0.3839],
+                [0.4100, 0.0012],
+                [0.7757, 0.5999]])
+    """
+    assert len(tensors) == len(index)
+    index, perm = torch.cat(index).sort(stable=True)
+    out = torch.cat(tensors, dim=0)[perm]
+    return (out, index) if return_index else out


Test patch:
diff --git a/test/utils/test_scatter.py b/test/utils/test_scatter.py
index c9a32fa03..462309234 100644
--- a/test/utils/test_scatter.py
+++ b/test/utils/test_scatter.py
@@ -6,7 +6,7 @@ import torch
 import torch_geometric.typing
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import withCUDA, withPackage
-from torch_geometric.utils import group_argsort, scatter
+from torch_geometric.utils import group_argsort, group_cat, scatter
 from torch_geometric.utils._scatter import scatter_argmax
 
 
@@ -111,6 +111,25 @@ def test_scatter_argmax(device):
     assert argmax.tolist() == [3, 5, 1, 4, 5, 5]
 
 
+@withCUDA
+def test_group_cat(device):
+    x1 = torch.randn(4, 4, device=device)
+    x2 = torch.randn(2, 4, device=device)
+    index1 = torch.tensor([0, 0, 1, 2], device=device)
+    index2 = torch.tensor([0, 2], device=device)
+
+    expected = torch.cat([x1[:2], x2[:1], x1[2:4], x2[1:]], dim=0)
+
+    out, index = group_cat(
+        [x1, x2],
+        [index1, index2],
+        dim=0,
+        return_index=True,
+    )
+    assert torch.equal(out, expected)
+    assert index.tolist() == [0, 0, 0, 1, 2, 2]
+
+
 if __name__ == '__main__':
     # Insights on GPU:
     # ================