Repository: pyg-team/pytorch_geometric
Pull request #9281 (2024-05-04 13:11:43)

Add `Index.sort()`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9281?dropdown=coverage&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 `95.23810%` with `2 lines` in your changes are missing coverage. Please review.
> Project coverage is 86.89%. Comparing base [(`73d3402`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/73d340275a4e8569da77a5b81add548efc6b88b9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`e4e9bfe`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9281?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).
> Report is 4 commits behind head on master.

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

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9281?dropdown=coverage&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/index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9281?src=pr&el=tree&filepath=torch_geometric%2Findex.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2luZGV4LnB5) | 95.23% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9281?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    #9281      +/-   ##
==========================================
+ Coverage   86.79%   86.89%   +0.09%     
==========================================
  Files         471      460      -11     
  Lines       30376    30154     -222     
==========================================
- Hits        26366    26203     -163     
+ Misses       4010     3951      -59     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9281?dropdown=coverage&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 5dedb1449..871c3bba5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added `torch_geometric.Index` ([#9276](https://github.com/pyg-team/pytorch_geometric/pull/9276), [#9277](https://github.com/pyg-team/pytorch_geometric/pull/9277), [#9278](https://github.com/pyg-team/pytorch_geometric/pull/9278), [#9279](https://github.com/pyg-team/pytorch_geometric/pull/9279), [#9280](https://github.com/pyg-team/pytorch_geometric/pull/9280))
+- Added `torch_geometric.Index` ([#9276](https://github.com/pyg-team/pytorch_geometric/pull/9276), [#9277](https://github.com/pyg-team/pytorch_geometric/pull/9277), [#9278](https://github.com/pyg-team/pytorch_geometric/pull/9278), [#9279](https://github.com/pyg-team/pytorch_geometric/pull/9279), [#9280](https://github.com/pyg-team/pytorch_geometric/pull/9280), [#9281](https://github.com/pyg-team/pytorch_geometric/pull/9281))
 - Added support for PyTorch 2.3 ([#9240](https://github.com/pyg-team/pytorch_geometric/pull/9240))
 - Added support for `EdgeIndex` in `message_and_aggregate` ([#9131](https://github.com/pyg-team/pytorch_geometric/pull/9131))
 - Added `CornellTemporalHyperGraphDataset` ([#9090](https://github.com/pyg-team/pytorch_geometric/pull/9090))
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index c27e19e98..aff331a4c 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -414,6 +414,48 @@ def _to_copy(
     )
 
 
-# def sort(self) -> None:
-#     # TODO MOVE BEHIND TORCH DISPATCH
-#     raise NotImplementedError
+@implements(aten.sort.default)
+def _sort(
+    tensor: Index,
+    dim: int = -1,
+    descending: bool = False,
+) -> Tuple[Index, Tensor]:
+
+    if tensor.is_sorted and not descending:
+        return tensor, torch.arange(tensor._data.numel(),
+                                    device=tensor._data.device)
+
+    data, perm = aten.sort.default(tensor._data, dim, descending)
+
+    out = Index(data)
+    out._dim_size = tensor._dim_size
+
+    if not descending:
+        out._is_sorted = True
+
+    return out, perm
+
+
+@implements(aten.sort.stable)
+def _sort_stable(
+    tensor: Index,
+    *,
+    stable: bool = False,
+    dim: int = -1,
+    descending: bool = False,
+) -> Tuple[Index, Tensor]:
+
+    if tensor.is_sorted and not descending:
+        return tensor, torch.arange(tensor._data.numel(),
+                                    device=tensor._data.device)
+
+    data, perm = aten.sort.stable(tensor._data, stable=stable, dim=dim,
+                                  descending=descending)
+
+    out = Index(data)
+    out._dim_size = tensor._dim_size
+
+    if not descending:
+        out._is_sorted = True
+
+    return out, perm


Test patch:
diff --git a/test/test_index.py b/test/test_index.py
index 20b828adf..d57eed322 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -216,3 +216,55 @@ def test_contiguous(dtype, device):
     out = index.contiguous()
     assert isinstance(out, Index)
     assert out.is_contiguous
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_sort(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([1, 0, 2, 1], dim_size=3, **kwargs)
+
+    index, _ = index.sort()
+    assert isinstance(index, Index)
+    assert index.equal(tensor([0, 1, 1, 2], device=device))
+    assert index.dim_size == 3
+    assert index.is_sorted
+
+    out, perm = index.sort()
+    assert isinstance(out, Index)
+    assert out._data.data_ptr() == index._data.data_ptr()
+    assert perm.equal(tensor([0, 1, 2, 3], device=device))
+    assert out.dim_size == 3
+
+    index, _ = index.sort(descending=True)
+    assert isinstance(index, Index)
+    assert index.equal(tensor([2, 1, 1, 0], device=device))
+    assert index.dim_size == 3
+    assert not index.is_sorted
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_sort_stable(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([1, 0, 2, 1], dim_size=3, **kwargs)
+
+    index, perm = index.sort(stable=True)
+    assert isinstance(index, Index)
+    assert index.equal(tensor([0, 1, 1, 2], device=device))
+    assert perm.equal(tensor([1, 0, 3, 2], device=device))
+    assert index.dim_size == 3
+    assert index.is_sorted
+
+    out, perm = index.sort(stable=True)
+    assert isinstance(out, Index)
+    assert out._data.data_ptr() == index._data.data_ptr()
+    assert perm.equal(tensor([0, 1, 2, 3], device=device))
+    assert out.dim_size == 3
+
+    index, perm = index.sort(descending=True, stable=True)
+    assert isinstance(index, Index)
+    assert index.equal(tensor([2, 1, 1, 0], device=device))
+    assert perm.equal(tensor([3, 1, 2, 0], device=device))
+    assert index.dim_size == 3
+    assert not index.is_sorted