Repository: pyg-team/pytorch_geometric
Pull request #9279 (2024-05-04 11:21:28)

Add `torch_geometric.Index.dim_resize_()`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9279?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> 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 [(`908a965`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9279?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 2 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #9279      +/-   ##
==========================================
+ Coverage   86.79%   86.89%   +0.09%     
==========================================
  Files         471      471              
  Lines       30376    30444      +68     
==========================================
+ Hits        26366    26455      +89     
+ Misses       4010     3989      -21     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9279?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 901bb4c99..2d2e853b2 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))
+- 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))
 - 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/edge_index.py b/torch_geometric/edge_index.py
index aef377711..f7926443d 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -547,11 +547,8 @@ class EdgeIndex(Tensor):
             if ptr is None or size is None:
                 return None
 
-            if ptr.numel() - 1 == size:
-                return ptr
-
-            if ptr.numel() - 1 > size:
-                return None
+            if ptr.numel() - 1 >= size:
+                return ptr[:size + 1]
 
             fill_value = ptr.new_full(
                 (size - ptr.numel() + 1, ),
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index 328257b7d..32f5b2931 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -9,7 +9,6 @@ from typing import (
     Optional,
     Tuple,
     Type,
-    Union,
 )
 
 import torch
@@ -212,7 +211,23 @@ class Index(Tensor):
 
     def dim_resize_(self, dim_size: Optional[int]) -> 'Index':
         r"""TODO."""
-        raise NotImplementedError  # TODO
+        if self.is_sorted and self._indptr is not None:
+            if dim_size is None:
+                self._indptr = None
+
+            elif self._indptr.numel() - 1 >= dim_size:
+                self._indptr = self._indptr[:dim_size + 1]
+
+            else:
+                fill_value = self._indptr.new_full(
+                    (dim_size - self._indptr.numel() + 1, ),
+                    fill_value=self._indptr[-1],  # type: ignore
+                )
+                self._indptr = torch.cat([self._indptr, fill_value], dim=0)
+
+        self._dim_size = dim_size
+
+        return self
 
     @assert_sorted
     def get_indptr(self) -> Tensor:
@@ -249,10 +264,6 @@ class Index(Tensor):
         """
         return self._data
 
-    def dim_narrow(self, start: Union[int, Tensor], length: int) -> 'Index':
-        r"""TODO."""
-        raise NotImplementedError  # TODO
-
     # PyTorch/Python builtins #################################################
 
     def __tensor_flatten__(self) -> Tuple[List[str], Tuple[Any, ...]]:


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index e7f82fe92..596006c5f 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -995,12 +995,10 @@ def test_matmul(without_extensions, device):
 
 
 @withCUDA
-def test_sparse_narrow(device):
-    adj = EdgeIndex(
-        [[0, 1, 1, 2], [1, 0, 2, 1]],
-        device=device,
-        sort_order='row',
-    )
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_sparse_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))
@@ -1029,8 +1027,9 @@ def test_sparse_narrow(device):
 
 
 @withCUDA
-def test_sparse_resize(device):
-    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], device=device)
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_sparse_resize(dtype, device):
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=dtype, device=device)
 
     out = adj.sort_by('row')[0].fill_cache_()
     assert out.sparse_size() == (3, 3)
@@ -1042,6 +1041,10 @@ def test_sparse_resize(device):
     assert out._T_indptr.equal(tensor([0, 1, 3, 4, 4, 4], device=device))
     out = out.sparse_resize_(3, 3)
     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.sparse_resize_(None, None)
+    assert out.sparse_size() == (None, None)
     assert out._indptr is None
     assert out._T_indptr is None
 
@@ -1055,6 +1058,10 @@ def test_sparse_resize(device):
     assert out._T_indptr.equal(tensor([0, 1, 3, 4, 4], device=device))
     out = out.sparse_resize_(3, 3)
     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.sparse_resize_(None, None)
+    assert out.sparse_size() == (None, None)
     assert out._indptr is None
     assert out._T_indptr is None
 
diff --git a/test/test_index.py b/test/test_index.py
index 8b2b34cbb..28cdb09aa 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -88,3 +88,25 @@ def test_fill_cache_(dtype, device):
     index.validate().fill_cache_()
     assert index.dim_size == 3
     assert index._indptr is None
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_dim_resize(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([0, 1, 1, 2], is_sorted=True, **kwargs).fill_cache_()
+
+    assert index.dim_size == 3
+    assert index._indptr.equal(tensor([0, 1, 3, 4], device=device))
+
+    out = index.dim_resize_(4)
+    assert out.dim_size == 4
+    assert out._indptr.equal(tensor([0, 1, 3, 4, 4], device=device))
+
+    out = index.dim_resize_(3)
+    assert out.dim_size == 3
+    assert out._indptr.equal(tensor([0, 1, 3, 4], device=device))
+
+    out = index.dim_resize_(None)
+    assert out.dim_size is None
+    assert out._indptr is None