Repository: pyg-team/pytorch_geometric

Issue:
PyTorch implementation of np.lexsort() in `sort_csc` util fucnction
### 🛠 Proposed Refactor

Hi,

I came upon this lexsort implementation in the `sort_csc` function in the following line:

https://github.com/pyg-team/pytorch_geometric/blob/9c79ce8570df194c15e4e6857c8b86d757dbbcb7/torch_geometric/sampler/utils.py#L27


And looks like an unofficial PyTorch implementation is proposed here:
https://discuss.pytorch.org/t/numpy-lexsort-equivalent-in-pytorch/47850/5?u=ebrahim.pichka

As below (with minor changes):

```py
def torch_lexsort(keys, dim=-1):
    if len(keys) < 2:
        raise ValueError(f"keys must be at least 2 sequences, but {len(keys)=}.")

    
    idx = keys[0].argsort(dim=dim, stable=True)
    for k in keys[1:]:
        idx = idx.gather(dim, k.gather(dim, idx).argsort(dim=dim, stable=True))
    
    return idx
```

I also ran some tests on the proposed implementation as follows:

```py
In [1]: import torch; import numpy as np

In [2]: N = 1000000

In [3]: a = np.random.rand(N); b = np.random.randint(N // 4, size=N)

In [4]: a_t = torch.tensor(a); b_t = torch.tensor(b)

In [5]: a_t_cu, b_t_cu = a_t.to(torch.device("cuda")), b_t.to(torch.device("cuda"))

In [6]: def torch_lexsort(keys, dim=-1):
   ...:     # defined as above

In [7]: %timeit -n 2 -r 20 np.lexsort([a, b])
302 ms ± 31.2 ms per loop (mean ± std. dev. of 20 runs, 2 loops each)

In [8]: %timeit -n 2 -r 20 torch_lexsort([a_t, b_t])
293 ms ± 35.1 ms per loop (mean ± std. dev. of 20 runs, 2 loops each)

In [9]: %timeit -n 20 -r 100 torch_lexsort([a_t_cu, b_t_cu])
The slowest run took 5.27 times longer than the fastest. This could mean that an intermediate result is being cached.
3.97 ms ± 334 µs per loop (mean ± std. dev. of 100 runs, 20 loops each)

In [10]: idx_np = np.lexsort([a, b]); idx_pt = torch_lexsort([a_t, b_t])

In [11]: (idx_np == idx_pt.numpy()).all()
Out[11]: True
```
Seems it helps with GPU support as well.

Thought it could replace the current implementation which detaches and recasts the tensor.

### Suggest a potential alternative/fix

Mentioned Above


Discussion:
Amazing. Do you wanna send a PR for this? Otherwise, I can go ahead and make the corresponding change as well.
@rusty1s I'd like to make a PR. 
Super :)

Resolving patch:
diff --git a/torch_geometric/sampler/utils.py b/torch_geometric/sampler/utils.py
--- a/torch_geometric/sampler/utils.py
+++ b/torch_geometric/sampler/utils.py
@@ -1,13 +1,12 @@
 from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union
 
-import numpy as np
 import torch
 from torch import Tensor
 
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.storage import EdgeStorage
 from torch_geometric.typing import NodeType, OptTensor
-from torch_geometric.utils import coalesce, index_sort
+from torch_geometric.utils import coalesce, index_sort, lexsort
 from torch_geometric.utils.sparse import index2ptr
 
 # Edge Layout Conversion ######################################################
@@ -22,14 +21,7 @@ def sort_csc(
         col, perm = index_sort(col)
         return row[perm], col, perm
     else:
-        # We use `np.lexsort` to sort based on multiple keys.
-        # TODO There does not seem to exist a PyTorch equivalent yet :(
-        perm = np.lexsort([
-            src_node_time[row].detach().cpu().numpy(),
-            col.detach().cpu().numpy()
-        ])
-        perm = torch.from_numpy(perm).to(col.device)
-
+        perm = lexsort([src_node_time[row], col])
         return row[perm], col[perm], perm
 
 
diff --git a/torch_geometric/utils/__init__.py b/torch_geometric/utils/__init__.py
--- a/torch_geometric/utils/__init__.py
+++ b/torch_geometric/utils/__init__.py
@@ -7,6 +7,7 @@
 from .softmax import softmax
 from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
 from .sort_edge_index import sort_edge_index
+from .lexsort import lexsort
 from .coalesce import coalesce
 from .undirected import is_undirected, to_undirected
 from .loop import (contains_self_loops, remove_self_loops,
@@ -63,6 +64,7 @@
     'dropout_path',
     'dropout_adj',
     'sort_edge_index',
+    'lexsort',
     'coalesce',
     'is_undirected',
     'to_undirected',
diff --git a/torch_geometric/utils/lexsort.py b/torch_geometric/utils/lexsort.py
new file mode 100644
--- /dev/null
+++ b/torch_geometric/utils/lexsort.py
@@ -0,0 +1,41 @@
+from typing import List
+
+import numpy as np
+import torch
+from torch import Tensor
+
+import torch_geometric.typing
+
+
+def lexsort(
+    keys: List[Tensor],
+    dim: int = -1,
+    descending: bool = False,
+) -> Tensor:
+    r"""Performs an indirect stable sort using a sequence of keys.
+
+    Given multiple sorting keys, returns an array of integer indices that
+    describe their sort order.
+    The last key in the sequence is used for the primary sort order, the
+    second-to-last key for the secondary sort order, and so on.
+
+    Args:
+        keys ([torch.Tensor]): The :math:`k` different columns to be sorted.
+            The last key is the primary sort key.
+        dim (int, optional): The dimension to sort along. (default: :obj:`-1`)
+        descending (bool, optional): Controls the sorting order (ascending or
+            descending). (default: :obj:`False`)
+    """
+    assert len(keys) >= 1
+
+    if not torch_geometric.typing.WITH_PT113:
+        keys = [k.neg() for k in keys] if descending else keys
+        out = np.lexsort([k.detach().cpu().numpy() for k in keys], axis=dim)
+        return torch.from_numpy(out).to(keys[0].device)
+
+    kwargs = dict(dim=dim, descending=descending, stable=True)
+    out = keys[0].argsort(**kwargs)
+    for k in keys[1:]:
+        out = out.gather(dim, k.gather(dim, out).argsort(**kwargs))
+
+    return out


Test patch:
diff --git a/test/utils/test_lexsort.py b/test/utils/test_lexsort.py
new file mode 100644
--- /dev/null
+++ b/test/utils/test_lexsort.py
@@ -0,0 +1,11 @@
+import numpy as np
+import torch
+
+from torch_geometric.utils import lexsort
+
+
+def test_lexsort():
+    keys = [torch.randn(100) for _ in range(3)]
+
+    expected = np.lexsort([key.numpy() for key in keys])
+    assert torch.equal(lexsort(keys), torch.from_numpy(expected))