Repository: pyg-team/pytorch_geometric
Pull request #9291 (2024-05-04 18:42:41)

Support `EdgeIndex.sparse_narrow` for non-sorted edge indices


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f671f2f27..31baea56f 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
 
+- Support `EdgeIndex.sparse_narrow` for non-sorted edge indices ([#9291](https://github.com/pyg-team/pytorch_geometric/pull/9291))
 - 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), [#9284](https://github.com/pyg-team/pytorch_geometric/pull/9284), [#9285](https://github.com/pyg-team/pytorch_geometric/pull/9285), [#9286](https://github.com/pyg-team/pytorch_geometric/pull/9286), [#9287](https://github.com/pyg-team/pytorch_geometric/pull/9287), [#9288](https://github.com/pyg-team/pytorch_geometric/pull/9288), [#9289](https://github.com/pyg-team/pytorch_geometric/pull/9289))
 - 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))
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index b849122df..0dab2e362 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1041,64 +1041,84 @@ class EdgeIndex(Tensor):
                              f"(got {start})")
 
         if dim == 0:
-            (rowptr, col), _ = self.get_csr()
-            rowptr = rowptr.narrow(0, start, length + 1)
+            if self.is_sorted_by_row:
+                (rowptr, col), _ = self.get_csr()
+                rowptr = rowptr.narrow(0, start, length + 1)
+
+                if rowptr.numel() < 2:
+                    row, col = self._data[0, :0], self._data[1, :0]
+                    rowptr = None
+                    num_rows = 0
+                else:
+                    col = col[rowptr[0]:rowptr[-1]]
+                    rowptr = rowptr - rowptr[0]
+                    num_rows = rowptr.numel() - 1
+
+                    row = torch.arange(
+                        num_rows,
+                        dtype=col.dtype,
+                        device=col.device,
+                    ).repeat_interleave(
+                        rowptr.diff(),
+                        output_size=col.numel(),
+                    )
 
-            if rowptr.numel() < 2:
-                row, col = self._data[0, :0], self._data[1, :0]
-                rowptr = None
-                num_rows = 0
-            else:
-                col = col[rowptr[0]:rowptr[-1]]
-                rowptr = rowptr - rowptr[0]
-                num_rows = rowptr.numel() - 1
-
-                row = torch.arange(
-                    num_rows,
-                    dtype=col.dtype,
-                    device=col.device,
-                ).repeat_interleave(
-                    rowptr.diff(),
-                    output_size=col.numel(),
+                edge_index = EdgeIndex(
+                    torch.stack([row, col], dim=0),
+                    sparse_size=(num_rows, self.sparse_size(1)),
+                    sort_order='row',
                 )
+                edge_index._indptr = rowptr
+                return edge_index
 
-            edge_index = EdgeIndex(
-                torch.stack([row, col], dim=0),
-                sparse_size=(num_rows, self.sparse_size(1)),
-                sort_order='row',
-            )
-            edge_index._indptr = rowptr
-            return edge_index
+            else:
+                mask = self._data[0] >= start
+                mask &= self._data[0] < (start + length)
+                offset = torch.tensor([[start], [0]], device=self.device)
+                edge_index = self[:, mask].sub_(offset)  # type: ignore
+                edge_index._sparse_size = (length, edge_index._sparse_size[1])
+                return edge_index
 
-        else:  # dim == 0:
-            (colptr, row), _ = self.get_csc()
-            colptr = colptr.narrow(0, start, length + 1)
+        else:
+            assert dim == 1
 
-            if colptr.numel() < 2:
-                row, col = self._data[0, :0], self._data[1, :0]
-                colptr = None
-                num_cols = 0
-            else:
-                row = row[colptr[0]:colptr[-1]]
-                colptr = colptr - colptr[0]
-                num_cols = colptr.numel() - 1
-
-                col = torch.arange(
-                    num_cols,
-                    dtype=row.dtype,
-                    device=row.device,
-                ).repeat_interleave(
-                    colptr.diff(),
-                    output_size=row.numel(),
+            if self.is_sorted_by_col:
+                (colptr, row), _ = self.get_csc()
+                colptr = colptr.narrow(0, start, length + 1)
+
+                if colptr.numel() < 2:
+                    row, col = self._data[0, :0], self._data[1, :0]
+                    colptr = None
+                    num_cols = 0
+                else:
+                    row = row[colptr[0]:colptr[-1]]
+                    colptr = colptr - colptr[0]
+                    num_cols = colptr.numel() - 1
+
+                    col = torch.arange(
+                        num_cols,
+                        dtype=row.dtype,
+                        device=row.device,
+                    ).repeat_interleave(
+                        colptr.diff(),
+                        output_size=row.numel(),
+                    )
+
+                edge_index = EdgeIndex(
+                    torch.stack([row, col], dim=0),
+                    sparse_size=(self.sparse_size(0), num_cols),
+                    sort_order='col',
                 )
+                edge_index._indptr = colptr
+                return edge_index
 
-            edge_index = EdgeIndex(
-                torch.stack([row, col], dim=0),
-                sparse_size=(self.sparse_size(0), num_cols),
-                sort_order='col',
-            )
-            edge_index._indptr = colptr
-            return edge_index
+            else:
+                mask = self._data[1] >= start
+                mask &= self._data[1] < (start + length)
+                offset = torch.tensor([[0], [start]], device=self.device)
+                edge_index = self[:, mask].sub_(offset)  # type: ignore
+                edge_index._sparse_size = (edge_index._sparse_size[0], length)
+                return edge_index
 
     def to_vector(self) -> Tensor:
         r"""Converts :class:`EdgeIndex` into a one-dimensional index


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 9b6358508..a16f695d3 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -148,9 +148,9 @@ def test_validate():
     with pytest.raises(ValueError, match="received a non-symmetric size"):
         EdgeIndex([[0, 1], [1, 0]], is_undirected=True, sparse_size=(2, 3))
     with pytest.raises(TypeError, match="invalid combination of arguments"):
-        EdgeIndex(torch.tensor([[0, 1], [1, 0]]), torch.long)
+        EdgeIndex(tensor([[0, 1], [1, 0]]), torch.long)
     with pytest.raises(TypeError, match="invalid keyword arguments"):
-        EdgeIndex(torch.tensor([[0, 1], [1, 0]]), dtype=torch.long)
+        EdgeIndex(tensor([[0, 1], [1, 0]]), dtype=torch.long)
     with pytest.raises(ValueError, match="contains negative indices"):
         EdgeIndex([[-1, 0], [0, 1]]).validate()
     with pytest.raises(ValueError, match="than its number of rows"):
@@ -696,19 +696,19 @@ def test_add(dtype, device, is_undirected):
     assert out.is_undirected == is_undirected
     assert out.sparse_size() == (7, 7)
 
-    out = adj + torch.tensor([2], dtype=dtype, device=device)
+    out = adj + tensor([2], dtype=dtype, device=device)
     assert isinstance(out, EdgeIndex)
     assert out.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
     assert out.is_undirected == is_undirected
     assert out.sparse_size() == (5, 5)
 
-    out = adj + torch.tensor([[2], [1]], dtype=dtype, device=device)
+    out = adj + tensor([[2], [1]], dtype=dtype, device=device)
     assert isinstance(out, EdgeIndex)
     assert out.equal(tensor([[2, 3, 3, 4], [2, 1, 3, 2]], device=device))
     assert not out.is_undirected
     assert out.sparse_size() == (5, 4)
 
-    out = adj + torch.tensor([[2], [2]], dtype=dtype, device=device)
+    out = adj + tensor([[2], [2]], dtype=dtype, device=device)
     assert isinstance(out, EdgeIndex)
     assert out.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
     assert out.is_undirected == is_undirected
@@ -743,19 +743,19 @@ def test_sub(dtype, device, is_undirected):
     assert out.is_undirected == is_undirected
     assert out.sparse_size() == (3, 3)
 
-    out = adj - torch.tensor([2], dtype=dtype, device=device)
+    out = adj - tensor([2], dtype=dtype, device=device)
     assert isinstance(out, EdgeIndex)
     assert out.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
     assert out.is_undirected == is_undirected
     assert out.sparse_size() == (5, 5)
 
-    out = adj - torch.tensor([[2], [1]], dtype=dtype, device=device)
+    out = adj - tensor([[2], [1]], dtype=dtype, device=device)
     assert isinstance(out, EdgeIndex)
     assert out.equal(tensor([[2, 3, 3, 4], [4, 3, 5, 4]], device=device))
     assert not out.is_undirected
     assert out.sparse_size() == (5, 6)
 
-    out = adj - torch.tensor([[2], [2]], dtype=dtype, device=device)
+    out = adj - tensor([[2], [2]], dtype=dtype, device=device)
     assert isinstance(out, EdgeIndex)
     assert out.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
     assert out.is_undirected == is_undirected
@@ -1013,32 +1013,62 @@ def test_matmul(without_extensions, device):
 
 @withCUDA
 @pytest.mark.parametrize('dtype', DTYPES)
-def test_sparse_narrow(dtype, device):
+def test_sparse_row_narrow(dtype, device):
     kwargs = dict(dtype=dtype, device=device)
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row', **kwargs)
 
     out = adj.sparse_narrow(dim=0, start=1, length=1)
-    assert out.equal(torch.tensor([[0, 0], [0, 2]], device=device))
+    assert out.equal(tensor([[0, 0], [0, 2]], device=device))
     assert out.sparse_size() == (1, None)
     assert out.sort_order == 'row'
-    assert out._indptr.equal(torch.tensor([0, 2], device=device))
+    assert out._indptr.equal(tensor([0, 2], device=device))
 
     out = adj.sparse_narrow(dim=0, start=2, length=0)
-    assert out.equal(torch.tensor([[], []], device=device))
+    assert out.equal(tensor([[], []], device=device))
     assert out.sparse_size() == (0, None)
     assert out.sort_order == 'row'
     assert out._indptr is None
 
     out = adj.sparse_narrow(dim=1, start=1, length=1)
-    assert (out.equal(torch.tensor([[0, 2], [0, 0]], device=device))
-            or out.equal(torch.tensor([[2, 0], [0, 0]], device=device)))
+    assert out.equal(tensor([[0, 2], [0, 0]], device=device))
     assert out.sparse_size() == (3, 1)
-    assert out.sort_order == 'col'
-    assert out._indptr.equal(torch.tensor([0, 2], device=device))
+    assert out.sort_order == 'row'
+    assert out._indptr is None
 
     out = adj.sparse_narrow(dim=1, start=2, length=0)
-    assert out.equal(torch.tensor([[], []], device=device))
+    assert out.equal(tensor([[], []], device=device))
     assert out.sparse_size() == (3, 0)
+    assert out.sort_order == 'row'
+    assert out._indptr is None
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_sparse_col_narrow(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    adj = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]], sort_order='col', **kwargs)
+
+    out = adj.sparse_narrow(dim=1, start=1, length=1)
+    assert out.equal(tensor([[0, 2], [0, 0]], device=device))
+    assert out.sparse_size() == (None, 1)
+    assert out.sort_order == 'col'
+    assert out._indptr.equal(tensor([0, 2], device=device))
+
+    out = adj.sparse_narrow(dim=1, start=2, length=0)
+    assert out.equal(tensor([[], []], device=device))
+    assert out.sparse_size() == (None, 0)
+    assert out.sort_order == 'col'
+    assert out._indptr is None
+
+    out = adj.sparse_narrow(dim=0, start=1, length=1)
+    assert out.equal(tensor([[0, 0], [0, 2]], device=device))
+    assert out.sparse_size() == (1, 3)
+    assert out.sort_order == 'col'
+    assert out._indptr is None
+
+    out = adj.sparse_narrow(dim=0, start=2, length=0)
+    assert out.equal(tensor([[], []], device=device))
+    assert out.sparse_size() == (0, 3)
     assert out.sort_order == 'col'
     assert out._indptr is None
 
@@ -1099,7 +1129,7 @@ def test_numpy():
 @pytest.mark.parametrize('dtype', DTYPES)
 def test_global_mapping(device, dtype):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], device=device, dtype=dtype)
-    n_id = torch.tensor([10, 20, 30], device=device, dtype=dtype)
+    n_id = tensor([10, 20, 30], device=device, dtype=dtype)
 
     expected = tensor([[10, 20, 20, 30], [20, 10, 30, 20]], device=device)
     out = n_id[adj]
diff --git a/test/test_index.py b/test/test_index.py
index 1f8e4fcab..f5e8820b1 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -67,9 +67,9 @@ def test_validate():
     with pytest.raises(ValueError, match="needs to be one-dimensional"):
         Index([[0], [1]])
     with pytest.raises(TypeError, match="invalid combination of arguments"):
-        Index(torch.tensor([0, 1]), torch.long)
+        Index(tensor([0, 1]), torch.long)
     with pytest.raises(TypeError, match="invalid keyword arguments"):
-        Index(torch.tensor([0, 1]), dtype=torch.long)
+        Index(tensor([0, 1]), dtype=torch.long)
     with pytest.raises(ValueError, match="contains negative indices"):
         Index([-1, 0]).validate()
     with pytest.raises(ValueError, match="than its registered size"):
@@ -370,7 +370,7 @@ def test_getitem(dtype, device):
     assert out.dim_size == 3
     assert out.is_sorted
 
-    out = index[torch.tensor([False, True, False, True], device=device)]
+    out = index[tensor([False, True, False, True], device=device)]
     assert isinstance(out, Index)
     assert out.equal(tensor([1, 2], device=device))
     assert out.dim_size == 3
@@ -431,7 +431,7 @@ def test_add(dtype, device):
     assert out.dim_size == 7
     assert out.is_sorted
 
-    out = index + torch.tensor([2], dtype=dtype, device=device)
+    out = index + tensor([2], dtype=dtype, device=device)
     assert isinstance(out, Index)
     assert out.equal(tensor([2, 3, 3, 4], device=device))
     assert out.dim_size == 5
@@ -465,7 +465,7 @@ def test_sub(dtype, device):
     assert out.dim_size == 3
     assert out.is_sorted
 
-    out = index - torch.tensor([2], dtype=dtype, device=device)
+    out = index - tensor([2], dtype=dtype, device=device)
     assert isinstance(out, Index)
     assert out.equal(tensor([2, 3, 3, 4], device=device))
     assert out.dim_size == 5