Repository: pyg-team/pytorch_geometric
Pull request #7649 (2023-06-27 08:05:16)

Add `Data.sort()` and better documentation on sort requirements in `Aggregation`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d13eb7fa3..be9a97918 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 `Data.sort()` and `HeteroData.sort()` functionalities ([#7649](https://github.com/pyg-team/pytorch_geometric/pull/7649))
 - Added `torch.nested_tensor` support in `Data` and `Batch` ([#7643](https://github.com/pyg-team/pytorch_geometric/pull/7643), [#7647](https://github.com/pyg-team/pytorch_geometric/pull/7647))
 - Added `interval` argument to `Cartesian` and `LocalCartesian` transformations ([#7533](https://github.com/pyg-team/pytorch_geometric/pull/7533), [#7614](https://github.com/pyg-team/pytorch_geometric/pull/7614))
 - Added a `LightGCN` example on the `AmazonBook` dataset ([7603](https://github.com/pyg-team/pytorch_geometric/pull/7603))
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index a9fa1f237..35bdaf37e 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -194,11 +194,6 @@ class BaseData:
         r"""Returns all edge-level tensor attribute names."""
         return list(set(chain(*[s.edge_attrs() for s in self.edge_stores])))
 
-    def is_coalesced(self) -> bool:
-        r"""Returns :obj:`True` if edge indices :obj:`edge_index` are sorted
-        and do not contain duplicate entries."""
-        return all([store.is_coalesced() for store in self.edge_stores])
-
     def generate_ids(self):
         r"""Generates and sets :obj:`n_id` and :obj:`e_id` attributes to assign
         each node and edge to a continuously ascending and unique ID."""
@@ -207,12 +202,42 @@ class BaseData:
         for store in self.edge_stores:
             store.e_id = torch.arange(store.num_edges)
 
-    def coalesce(self):
+    def is_sorted(self, sort_by_row: bool = True) -> bool:
+        r"""Returns :obj:`True` if edge indices :obj:`edge_index` are sorted.
+
+        Args:
+            sort_by_row (bool, optional): If set to :obj:`False`, will require
+                column-wise order/by destination node of :obj:`edge_index`.
+                (default: :obj:`True`)
+        """
+        return all(
+            [store.is_sorted(sort_by_row) for store in self.edge_stores])
+
+    def sort(self, sort_by_row: bool = True) -> 'Data':
+        r"""Sorts edge indices :obj:`edge_index`.
+
+        Args:
+            sort_by_row (bool, optional): If set to :obj:`False`, will sort
+                :obj:`edge_index` column-wise order/by destination node.
+                (default: :obj:`True`)
+        """
+        out = copy.copy(self)
+        for store in out.edge_stores:
+            store.sort(sort_by_row)
+        return out
+
+    def is_coalesced(self) -> bool:
+        r"""Returns :obj:`True` if edge indices :obj:`edge_index` are sorted
+        and do not contain duplicate entries."""
+        return all([store.is_coalesced() for store in self.edge_stores])
+
+    def coalesce(self) -> 'Data':
         r"""Sorts and removes duplicated entries from edge indices
         :obj:`edge_index`."""
-        for store in self.edge_stores:
+        out = copy.copy(self)
+        for store in out.edge_stores:
             store.coalesce()
-        return self
+        return out
 
     def has_isolated_nodes(self) -> bool:
         r"""Returns :obj:`True` if the graph contains isolated nodes."""
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index b46b0484a..716d00d92 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -28,6 +28,7 @@ from torch_geometric.utils import (
     contains_isolated_nodes,
     is_torch_sparse_tensor,
     is_undirected,
+    sort_edge_index,
 )
 
 N_KEYS = {'x', 'feat', 'pos', 'batch', 'node_type', 'n_id'}
@@ -496,6 +497,23 @@ class EdgeStorage(BaseStorage):
     def edge_attrs(self) -> List[str]:
         return [key for key in self.keys() if self.is_edge_attr(key)]
 
+    def is_sorted(self, sort_by_row: bool = True) -> bool:
+        if 'edge_index' in self:
+            index = self.edge_index[0] if sort_by_row else self.edge_index[1]
+            return bool(torch.all(index[:-1] <= index[1:]))
+        return True
+
+    def sort(self, sort_by_row: bool = True) -> 'EdgeStorage':
+        if 'edge_index' in self:
+            edge_attrs = self.edge_attrs()
+            edge_attrs.remove('edge_index')
+            edge_feats = [self[edge_attr] for edge_attr in edge_attrs]
+            self.edge_index, edge_feats = sort_edge_index(
+                self.edge_index, edge_feats, sort_by_row=sort_by_row)
+            for key, edge_feat in zip(edge_attrs, edge_feats):
+                self[key] = edge_feat
+        return self
+
     def is_coalesced(self) -> bool:
         for value in self.values('adj', 'adj_t'):
             return value.is_coalesced()
@@ -510,7 +528,7 @@ class EdgeStorage(BaseStorage):
 
         return True
 
-    def coalesce(self, reduce: str = 'sum'):
+    def coalesce(self, reduce: str = 'sum') -> 'EdgeStorage':
         for key, value in self.items('adj', 'adj_t'):
             self[key] = value.coalesce(reduce)
 
diff --git a/torch_geometric/nn/aggr/base.py b/torch_geometric/nn/aggr/base.py
index 593cce8a5..58be08b3d 100644
--- a/torch_geometric/nn/aggr/base.py
+++ b/torch_geometric/nn/aggr/base.py
@@ -140,7 +140,11 @@ class Aggregation(torch.nn.Module):
     def assert_sorted_index(self, index: Optional[Tensor]):
         if index is not None and not torch.all(index[:-1] <= index[1:]):
             raise ValueError("Can not perform aggregation since the 'index' "
-                             "tensor is not sorted")
+                             "tensor is not sorted. Specifically, if you use "
+                             "this aggregation as part of 'MessagePassing`, "
+                             "ensure that 'edge_index' is sorted by "
+                             "destination nodes, e.g., by calling "
+                             "`data.sort(sort_by_row=False)`")
 
     def assert_two_dimensional_input(self, x: Tensor, dim: int):
         if x.dim() != 2:
diff --git a/torch_geometric/nn/aggr/gmt.py b/torch_geometric/nn/aggr/gmt.py
index 3f81785b1..b18bb8d6d 100644
--- a/torch_geometric/nn/aggr/gmt.py
+++ b/torch_geometric/nn/aggr/gmt.py
@@ -22,6 +22,15 @@ class GraphMultisetTransformer(Aggregation):
     and finally pools the representative elements via attention-based pooling
     into a single cluster.
 
+    .. note::
+
+        :class:`GraphMultisetTransformer` requires sorted indices :obj:`index`
+        as input. Specifically, if you use this aggregation as part of
+        :class:`~torch_geometric.nn.conv.MessagePassing`, ensure that
+        :obj:`edge_index` is sorted by destination nodes, either by manually
+        sorting edge indices via :meth:`~torch_geometric.utils.sort_edge_index`
+        or by calling :meth:`~torch_geometric.data.Data.sort`.
+
     Args:
         channels (int): Size of each input sample.
         k (int): Number of :math:`k` representative nodes after pooling.
diff --git a/torch_geometric/nn/aggr/gru.py b/torch_geometric/nn/aggr/gru.py
index c731d6025..29abbbb2e 100644
--- a/torch_geometric/nn/aggr/gru.py
+++ b/torch_geometric/nn/aggr/gru.py
@@ -12,7 +12,17 @@ class GRUAggregation(Aggregation):
     interpreted as a sequence, as described in the `"Graph Neural Networks
     with Adaptive Readouts" <https://arxiv.org/abs/2211.04952>`_ paper.
 
+    .. note::
+
+        :class:`MLPAggregation` requires sorted indices :obj:`index` as input.
+        Specifically, if you use this aggregation as part of
+        :class:`~torch_geometric.nn.conv.MessagePassing`, ensure that
+        :obj:`edge_index` is sorted by destination nodes, either by manually
+        sorting edge indices via :meth:`~torch_geometric.utils.sort_edge_index`
+        or by calling :meth:`~torch_geometric.data.Data.sort`.
+
     .. warning::
+
         :class:`GRUAggregation` is not a permutation-invariant operator.
 
     Args:
diff --git a/torch_geometric/nn/aggr/lstm.py b/torch_geometric/nn/aggr/lstm.py
index 25da0b030..6cab8b55f 100644
--- a/torch_geometric/nn/aggr/lstm.py
+++ b/torch_geometric/nn/aggr/lstm.py
@@ -12,7 +12,17 @@ class LSTMAggregation(Aggregation):
     interpreted as a sequence, as described in the `"Inductive Representation
     Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper.
 
+    .. note::
+
+        :class:`LSTMAggregation` requires sorted indices :obj:`index` as input.
+        Specifically, if you use this aggregation as part of
+        :class:`~torch_geometric.nn.conv.MessagePassing`, ensure that
+        :obj:`edge_index` is sorted by destination nodes, either by manually
+        sorting edge indices via :meth:`~torch_geometric.utils.sort_edge_index`
+        or by calling :meth:`~torch_geometric.data.Data.sort`.
+
     .. warning::
+
         :class:`LSTMAggregation` is not a permutation-invariant operator.
 
     Args:
diff --git a/torch_geometric/nn/aggr/mlp.py b/torch_geometric/nn/aggr/mlp.py
index bfbd34c46..8410b0e70 100644
--- a/torch_geometric/nn/aggr/mlp.py
+++ b/torch_geometric/nn/aggr/mlp.py
@@ -11,7 +11,17 @@ class MLPAggregation(Aggregation):
     a Multi-Layer Perceptron (MLP), as described in the `"Graph Neural Networks
     with Adaptive Readouts" <https://arxiv.org/abs/2211.04952>`_ paper.
 
+    .. note::
+
+        :class:`GRUAggregation` requires sorted indices :obj:`index` as input.
+        Specifically, if you use this aggregation as part of
+        :class:`~torch_geometric.nn.conv.MessagePassing`, ensure that
+        :obj:`edge_index` is sorted by destination nodes, either by manually
+        sorting edge indices via :meth:`~torch_geometric.utils.sort_edge_index`
+        or by calling :meth:`~torch_geometric.data.Data.sort`.
+
     .. warning::
+
         :class:`MLPAggregation` is not a permutation-invariant operator.
 
     Args:
diff --git a/torch_geometric/nn/aggr/set2set.py b/torch_geometric/nn/aggr/set2set.py
index 50bc7c4b2..0e281e966 100644
--- a/torch_geometric/nn/aggr/set2set.py
+++ b/torch_geometric/nn/aggr/set2set.py
@@ -44,7 +44,6 @@ class Set2Set(Aggregation):
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2) -> Tensor:
 
-        # TODO Currently, `to_dense_batch` can only operate on `index`:
         self.assert_index_present(index)
         self.assert_two_dimensional_input(x, dim)
 
diff --git a/torch_geometric/nn/aggr/set_transformer.py b/torch_geometric/nn/aggr/set_transformer.py
index 01b837966..72afdd07d 100644
--- a/torch_geometric/nn/aggr/set_transformer.py
+++ b/torch_geometric/nn/aggr/set_transformer.py
@@ -17,6 +17,15 @@ class SetTransformerAggregation(Aggregation):
     the `"Graph Neural Networks with Adaptive Readouts"
     <https://arxiv.org/abs/2211.04952>`_ paper.
 
+    .. note::
+
+        :class:`SetTransformerAggregation` requires sorted indices :obj:`index`
+        as input. Specifically, if you use this aggregation as part of
+        :class:`~torch_geometric.nn.conv.MessagePassing`, ensure that
+        :obj:`edge_index` is sorted by destination nodes, either by manually
+        sorting edge indices via :meth:`~torch_geometric.utils.sort_edge_index`
+        or by calling :meth:`~torch_geometric.data.Data.sort`.
+
     Args:
         channels (int): Size of each input sample.
         num_seed_points (int, optional): Number of seed points.
diff --git a/torch_geometric/nn/aggr/sort.py b/torch_geometric/nn/aggr/sort.py
index ba7c216ff..79ee9a5ab 100644
--- a/torch_geometric/nn/aggr/sort.py
+++ b/torch_geometric/nn/aggr/sort.py
@@ -14,6 +14,15 @@ class SortAggregation(Aggregation):
     where node features are sorted in descending order based on their last
     feature channel. The first :math:`k` nodes form the output of the layer.
 
+    .. note::
+
+        :class:`SortAggregation` requires sorted indices :obj:`index` as input.
+        Specifically, if you use this aggregation as part of
+        :class:`~torch_geometric.nn.conv.MessagePassing`, ensure that
+        :obj:`edge_index` is sorted by destination nodes, either by manually
+        sorting edge indices via :meth:`~torch_geometric.utils.sort_edge_index`
+        or by calling :meth:`~torch_geometric.data.Data.sort`.
+
     Args:
         k (int): The number of nodes to hold for each graph.
     """
diff --git a/torch_geometric/utils/sort_edge_index.py b/torch_geometric/utils/sort_edge_index.py
index 3b7a6fe27..d09d0e230 100644
--- a/torch_geometric/utils/sort_edge_index.py
+++ b/torch_geometric/utils/sort_edge_index.py
@@ -45,7 +45,8 @@ def sort_edge_index(  # noqa
         num_nodes (int, optional): The number of nodes, *i.e.*
             :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
         sort_by_row (bool, optional): If set to :obj:`False`, will sort
-            :obj:`edge_index` column-wise.
+            :obj:`edge_index` column-wise/by destination node.
+            (default: :obj:`True`)
 
     :rtype: :class:`LongTensor` if :attr:`edge_attr` is not passed, else
         (:class:`LongTensor`, :obj:`Optional[Tensor]` or :obj:`List[Tensor]]`)


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index 144c94fa1..4b4fd3ee1 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -284,6 +284,36 @@ def test_copy_data():
     assert data.x.tolist() == out.x.tolist()
 
 
+def test_data_sort():
+    x = torch.randn(4, 16)
+    edge_index = torch.tensor([[0, 0, 0, 2, 1, 3], [1, 2, 3, 0, 0, 0]])
+    edge_attr = torch.randn(6, 8)
+
+    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
+    assert not data.is_sorted(sort_by_row=True)
+    assert not data.is_sorted(sort_by_row=False)
+
+    out = data.sort(sort_by_row=True)
+    assert out.is_sorted(sort_by_row=True)
+    assert not out.is_sorted(sort_by_row=False)
+    assert torch.equal(out.x, data.x)
+    assert out.edge_index.tolist() == [[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]]
+    assert torch.equal(
+        out.edge_attr,
+        data.edge_attr[torch.tensor([0, 1, 2, 4, 3, 5])],
+    )
+
+    out = data.sort(sort_by_row=False)
+    assert not out.is_sorted(sort_by_row=True)
+    assert out.is_sorted(sort_by_row=False)
+    assert torch.equal(out.x, data.x)
+    assert out.edge_index.tolist() == [[1, 2, 3, 0, 0, 0], [0, 0, 0, 1, 2, 3]]
+    assert torch.equal(
+        out.edge_attr,
+        data.edge_attr[torch.tensor([4, 3, 5, 0, 1, 2])],
+    )
+
+
 def test_debug_data():
     torch_geometric.set_debug(True)