Repository: pyg-team/pytorch_geometric
Pull request #8446 (2023-11-27 06:46:27)

`EdgeIndex` fill cache

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8446?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:
> Comparison is base [(`7b71b3e`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/7b71b3ec2f2b605ce067547f830799837e208e9f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.60% compared to head [(`725c6a8`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8446?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.64%.


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


```diff
@@            Coverage Diff             @@
##           master    #8446      +/-   ##
==========================================
+ Coverage   87.60%   87.64%   +0.03%     
==========================================
  Files         477      477              
  Lines       29150    29161      +11     
==========================================
+ Hits        25538    25557      +19     
+ Misses       3612     3604       -8     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8446?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/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index bc625627a..54328efa0 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -121,6 +121,29 @@ class EdgeIndex(Tensor):
 
         return self
 
+    def fill_cache(self) -> 'EdgeIndex':
+        if self._sort_order == SortOrder.ROW and self._rowptr is None:
+            if self.num_rows is None:
+                self._sparse_size = (
+                    int(self[0].max()) + 1 if self.numel() > 0 else 0,
+                    self.num_cols,
+                )
+
+            self._rowptr = torch._convert_indices_from_coo_to_csr(
+                self[0], self.num_rows, out_int32=self.dtype != torch.int64)
+            self._rowptr = self._rowptr.to(self.dtype)
+
+        if self._sort_order == SortOrder.COL and self._colptr is None:
+            if self.num_cols is None:
+                self._sparse_size = (
+                    self.num_rows,
+                    int(self[1].max()) + 1 if self.numel() > 0 else 0,
+                )
+
+            self._colptr = torch._convert_indices_from_coo_to_csr(
+                self[1], self.num_cols, out_int32=self.dtype != torch.int64)
+            self._colptr = self._colptr.to(self.dtype)
+
     @property
     def sparse_size(self) -> Tuple[Optional[int], Optional[int]]:
         r"""The size of the underlying sparse matrix."""


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 85a6bd0a2..bc463eccf 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -15,8 +15,19 @@ def test_edge_index():
 
     assert isinstance(adj.as_tensor(), Tensor)
 
-    out = adj + 1
-    assert isinstance(out, Tensor)
+    assert isinstance(adj + 1, Tensor)
+
+
+def test_edge_index_fill_cache():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+    adj.validate().fill_cache()
+    assert adj.sparse_size == (3, None)
+    assert torch.equal(adj._rowptr, torch.tensor([0, 1, 3, 4]))
+
+    adj = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]], sort_order='col')
+    adj.validate().fill_cache()
+    assert adj.sparse_size == (None, 3)
+    assert torch.equal(adj._colptr, torch.tensor([0, 1, 3, 4]))
 
 
 @withCUDA