Add resize method to EdgeIndex (#8983)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: rusty1s <matthias.fey@tu-dortmund.de>

CHANGELOG.md

@@ -7,9 +7,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `EdgeIndex.resize_` functionality ([#8983](https://github.com/pyg-team/pytorch_geometric/pull/8983))
+- Added approximate `faiss`-based KNN-search ([#8952](https://github.com/pyg-team/pytorch_geometric/pull/8952))
+
 ### Changed
 
-- Added approximate `faiss`-based KNN-search ([#8952](https://github.com/pyg-team/pytorch_geometric/pull/8952))
 - Breaking Change: Added support for `EdgeIndex` in `cugraph` GNN layers ([#8938](https://github.com/pyg-team/pytorch_geometric/pull/8937))
 - Added the `dim` arg to `torch.cross` calls ([#8918](https://github.com/pyg-team/pytorch_geometric/pull/8918))
 

test/test_edge_index.py

@@ -1023,6 +1023,37 @@ def test_sparse_narrow(device):
     assert out._indptr is None
 
 
+@withCUDA
+def test_resize(device):
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], device=device)
+
+    out = adj.sort_by('row')[0].fill_cache_()
+    assert out.sparse_size() == (3, 3)
+    assert out._indptr.equal(tensor([0, 1, 3, 4], device=device))
+    assert out._T_indptr.equal(tensor([0, 1, 3, 4], device=device))
+    out = out.resize_((4, 5))
+    assert out.sparse_size() == (4, 5)
+    assert out._indptr.equal(tensor([0, 1, 3, 4, 4], device=device))
+    assert out._T_indptr.equal(tensor([0, 1, 3, 4, 4, 4], device=device))
+    out = out.resize_((3, 3))
+    assert out.sparse_size() == (3, 3)
+    assert out._indptr is None
+    assert out._T_indptr is None
+
+    out = adj.sort_by('col')[0].fill_cache_()
+    assert out.sparse_size() == (3, 3)
+    assert out._indptr.equal(tensor([0, 1, 3, 4], device=device))
+    assert out._T_indptr.equal(tensor([0, 1, 3, 4], device=device))
+    out = out.resize_((4, 5))
+    assert out.sparse_size() == (4, 5)
+    assert out._indptr.equal(tensor([0, 1, 3, 4, 4, 4], device=device))
+    assert out._T_indptr.equal(tensor([0, 1, 3, 4, 4], device=device))
+    out = out.resize_((3, 3))
+    assert out.sparse_size() == (3, 3)
+    assert out._indptr is None
+    assert out._T_indptr is None
+
+
 @withCUDA
 @pytest.mark.parametrize('dtype', DTYPES)
 def test_save_and_load(dtype, device, tmp_path):

torch_geometric/edge_index.py

@@ -516,6 +516,43 @@ def get_sparse_size(
 
         return torch.Size((self.get_sparse_size(0), self.get_sparse_size(1)))
 
+    def resize_(self, sparse_size: Tuple[int, int]) -> 'EdgeIndex':
+        r"""Assigns or re-assigns the size of the underlying sparse matrix.
+
+        Args:
+            sparse_size (tuple[int, int]): The size of the sparse matrix.
+        """
+        num_rows, num_cols = sparse_size
+
+        if self.is_undirected and num_rows != num_cols:
+            raise ValueError(f"'EdgeIndex' is undirected but received a "
+                             f"non-symmetric size (got {list(sparse_size)})")
+
+        def _modify_ptr(ptr: Optional[Tensor], size: int) -> Optional[Tensor]:
+            if ptr is None:
+                return None
+
+            if ptr.numel() - 1 == size:
+                return ptr
+
+            if ptr.numel() - 1 > size:
+                return None
+
+            val = ptr.new_full((size - ptr.numel() + 1, ), fill_value=ptr[-1])
+            return torch.cat([ptr, val], dim=0)
+
+        if self.is_sorted_by_row:
+            self._indptr = _modify_ptr(self._indptr, num_rows)
+            self._T_indptr = _modify_ptr(self._T_indptr, num_cols)
+
+        if self.is_sorted_by_col:
+            self._indptr = _modify_ptr(self._indptr, num_cols)
+            self._T_indptr = _modify_ptr(self._T_indptr, num_rows)
+
+        self._sparse_size = sparse_size
+
+        return self
+
     def get_num_rows(self) -> int:
         r"""The number of rows of the underlying sparse matrix.
         Automatically computed and cached when not explicitly set.