Repository: pyg-team/pytorch_geometric
Pull request #8477 (2023-11-29 09:03:59)

`EdgeIndex` conversions

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8477?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `5 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`57a4afc`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/57a4afc1091b427d1026c0a9637605878d9adf79?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.70% compared to head [(`bd367b4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8477?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.71%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8477?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/data/edge\_index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8477?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 91.22% | [5 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8477?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

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


```diff
@@           Coverage Diff           @@
##           master    #8477   +/-   ##
=======================================
  Coverage   87.70%   87.71%           
=======================================
  Files         478      478           
  Lines       29287    29333   +46     
=======================================
+ Hits        25687    25728   +41     
- Misses       3600     3605    +5     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8477?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 e50cb1942..8468444f7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,7 +10,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `is_torch_instance` to check against the original class of compiled models ([#8461](https://github.com/pyg-team/pytorch_geometric/pull/8461))
 - Added dense computation for `AddRandomWalkPE` ([#8431](https://github.com/pyg-team/pytorch_geometric/pull/8431))
 - Added a tutorial for point cloud processing ([#8015](https://github.com/pyg-team/pytorch_geometric/pull/8015))
-- Added `fsspec` as file system backend ([#8379](https://github.com/pyg-team/pytorch_geometric/pull/8379), [#8426](https://github.com/pyg-team/pytorch_geometric/pull/8426), [#8434](https://github.com/pyg-team/pytorch_geometric/pull/8434))
+- Added `fsspec` as file system backend ([#8379](https://github.com/pyg-team/pytorch_geometric/pull/8379), [#8426](https://github.com/pyg-team/pytorch_geometric/pull/8426), [#8434](https://github.com/pyg-team/pytorch_geometric/pull/8434), [#8474](https://github.com/pyg-team/pytorch_geometric/pull/8474))
 - Added support for floating-point average degree numbers in `FakeDataset` and `FakeHeteroDataset` ([#8404](https://github.com/pyg-team/pytorch_geometric/pull/8404))
 - Added support for device conversions of `InMemoryDataset` ([#8402](https://github.com/pyg-team/pytorch_geometric/pull/8402))
 - Added support for edge-level temporal sampling in `NeighborLoader` and `LinkNeighborLoader` ([#8372](https://github.com/pyg-team/pytorch_geometric/pull/8372))
diff --git a/examples/compile/gcn.py b/examples/compile/gcn.py
index 0f18af43c..35c828712 100644
--- a/examples/compile/gcn.py
+++ b/examples/compile/gcn.py
@@ -48,7 +48,7 @@ model = GCN(
 ).to(device)
 
 # Compile the model into an optimized version and enforce zero graph breaks:
-model = torch_geometric.compile(model, dynamic=False, fullgraph=True)
+model = torch_geometric.compile(model, dynamic=False)
 
 optimizer = torch.optim.Adam([
     dict(params=model.conv1.parameters(), weight_decay=5e-4),
diff --git a/examples/compile/gin.py b/examples/compile/gin.py
index ea1513876..c66dc741f 100644
--- a/examples/compile/gin.py
+++ b/examples/compile/gin.py
@@ -57,7 +57,7 @@ model = GIN(
 ).to(device)
 
 # Compile the model into an optimized version and enforce zero graph breaks:
-model = torch_geometric.compile(model, dynamic=True, fullgraph=True)
+model = torch_geometric.compile(model, dynamic=True)
 
 optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index 2a2bf38c7..5dd5bf286 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -5,6 +5,7 @@ from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
 from torch_geometric.utils import index_sort
 
 HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
@@ -167,47 +168,51 @@ class EdgeIndex(Tensor):
 
         return self
 
-    def fill_cache(self) -> 'EdgeIndex':
-        r"""Fills the cache with (meta)data information.
-        No-op in case :class:`EdgeIndex` is not sorted.
-        """
-        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."""
         return self._sparse_size
 
+    def get_sparse_size(self) -> Tuple[int, int]:
+        r"""The size of the underlying sparse matrix.
+        Automatically computed and cached when not explicitly set.
+        """
+        return (self.get_num_rows(), self.get_num_cols())
+
     @property
     def num_rows(self) -> Optional[int]:
         r"""The number of rows of the underlying sparse matrix."""
         return self._sparse_size[0]
 
+    def get_num_rows(self) -> int:
+        r"""The number of rows of the underlying sparse matrix.
+        Automatically computed and cached when not explicitly set.
+        """
+        if self.num_rows is None:
+            self._sparse_size = (
+                int(self[0].max()) + 1 if self.numel() > 0 else 0,
+                self.num_cols,
+            )
+
+        return self.num_rows
+
     @property
     def num_cols(self) -> Optional[int]:
         r"""The number of columns of the underlying sparse matrix."""
         return self._sparse_size[1]
 
+    def get_num_cols(self) -> int:
+        r"""The number of columns of the underlying sparse matrix.
+        Automatically computed and cached when not explicitly set.
+        """
+        if self.num_cols is None:
+            self._sparse_size = (
+                self.num_rows,
+                int(self[1].max()) + 1 if self.numel() > 0 else 0,
+            )
+
+        return self.num_cols
+
     @property
     def sort_order(self) -> Optional[str]:
         r"""The sort order of indices, either :obj:`"row"`, :obj:`"col"` or
@@ -215,6 +220,56 @@ class EdgeIndex(Tensor):
         """
         return None if self._sort_order is None else self._sort_order.value
 
+    def get_rowptr(self) -> Tensor:
+        r"""Returns the :obj:`rowptr` vector of :class:`EdgeIndex`, a
+        compressed representation of row indices in case :class:`EdgeIndex` is
+        sorted by rows.
+        """
+        if self._sort_order != SortOrder.ROW:
+            raise ValueError(
+                f"Cannot access 'rowptr' in '{self.__class__.__name__}' "
+                f"since it is not sorted by rows (got '{self.sort_order}')")
+
+        if self._rowptr is not None:
+            return self._rowptr
+
+        self._rowptr = torch._convert_indices_from_coo_to_csr(
+            self[0], self.get_num_rows(), out_int32=self.dtype != torch.int64)
+        self._rowptr = self._rowptr.to(self.dtype)
+
+        return self._rowptr
+
+    def get_colptr(self) -> Tensor:
+        r"""Returns the :obj:`colptr` vector of :class:`EdgeIndex`, a
+        compressed representation of column indices in case :class:`EdgeIndex`
+        is sorted by columns.
+        """
+        if self._sort_order != SortOrder.COL:
+            raise ValueError(
+                f"Cannot access 'colptr' in '{self.__class__.__name__}' "
+                f"since it is not sorted by columns (got '{self.sort_order}')")
+
+        if self._colptr is not None:
+            return self._colptr
+
+        self._colptr = torch._convert_indices_from_coo_to_csr(
+            self[1], self.get_num_cols(), out_int32=self.dtype != torch.int64)
+        self._colptr = self._colptr.to(self.dtype)
+
+        return self._colptr
+
+    def fill_cache(self) -> 'EdgeIndex':
+        r"""Fills the cache with (meta)data information.
+        No-op in case :class:`EdgeIndex` is not sorted.
+        """
+        self.get_num_rows()
+        self.get_num_cols()
+
+        if self._sort_order == SortOrder.ROW:
+            self.get_rowptr()
+        if self._sort_order == SortOrder.COL:
+            self.get_colptr()
+
     def as_tensor(self) -> Tensor:
         r"""Zero-copies the :class:`EdgeIndex` representation back to a
         :class:`torch.Tensor` representation.
@@ -504,3 +559,78 @@ def getitem(input: EdgeIndex, index: Any) -> Union[EdgeIndex, Tensor]:
             out._sort_order = input._sort_order
 
     return out
+
+
+def _get_value(
+    numel: int,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+    if (torch_geometric.typing.WITH_PT20
+            and not torch_geometric.typing.WITH_ARM):
+        return torch.ones(1, dtype=dtype, device=device).expand(numel)
+    return torch.ones(numel, dtype=dtype, device=device)
+
+
+# `to_sparse_coo()` uses `to_sparse(layout=None)` dispatch logic:
+def to_sparse_coo(tensor: EdgeIndex) -> Tensor:
+    out = torch.sparse_coo_tensor(
+        indices=tensor.as_tensor(),
+        values=_get_value(tensor.size(1), device=tensor.device),
+        size=tensor.get_sparse_size(),
+        device=tensor.device,
+    )
+
+    if tensor._sort_order == SortOrder.ROW:
+        out = out._coalesced_(True)
+
+    return out
+
+
+@implements(Tensor.to_sparse_csr)
+def to_sparse_csr(tensor: EdgeIndex) -> Tensor:
+    return torch.sparse_csr_tensor(
+        crow_indices=tensor.get_rowptr(),
+        col_indices=tensor[1],
+        values=_get_value(tensor.size(1), device=tensor.device),
+        size=tensor.get_sparse_size(),
+        device=tensor.device,
+    )
+
+
+if torch_geometric.typing.WITH_PT112:
+
+    @implements(Tensor.to_sparse_csc)
+    def to_sparse_csc(tensor: EdgeIndex) -> Tensor:
+        return torch.sparse_csc_tensor(
+            ccol_indices=tensor.get_colptr(),
+            row_indices=tensor[0],
+            values=_get_value(tensor.size(1), device=tensor.device),
+            size=tensor.get_sparse_size(),
+            device=tensor.device,
+        )
+
+
+if torch_geometric.typing.WITH_PT20:
+
+    @implements(Tensor.to_sparse)
+    def to_sparse(
+        tensor: EdgeIndex,
+        *,
+        layout: Optional[torch.layout] = None,
+    ) -> Tensor:
+
+        if layout is None or layout == torch.sparse_coo:
+            return to_sparse_coo(tensor)
+        if layout == torch.sparse_csr:
+            return tensor.to_sparse_csr()
+        if torch_geometric.typing.WITH_PT112 and layout == torch.sparse_csc:
+            return tensor.to_sparse_csc()
+
+        raise ValueError(f"Unexpected tensor layout (got '{layout}')")
+
+else:
+
+    @implements(Tensor.to_sparse)
+    def to_sparse(tensor: EdgeIndex) -> Tensor:
+        return to_sparse_coo(tensor)
diff --git a/torch_geometric/io/fs.py b/torch_geometric/io/fs.py
index be7c91299..46ea357da 100644
--- a/torch_geometric/io/fs.py
+++ b/torch_geometric/io/fs.py
@@ -93,6 +93,7 @@ def cp(
     path2: str,
     extract: bool = False,
     log: bool = True,
+    use_cache: bool = True,
     clear_cache: bool = True,
 ):
     kwargs = {}
@@ -103,10 +104,11 @@ def cp(
         if log and 'pytest' not in sys.modules:
             print(f'Downloading {path1}', file=sys.stderr)
 
-        home_dir = torch_geometric.get_home_dir()
-        cache_dir = osp.join(home_dir, 'simplecache', uuid4().hex)
-        kwargs.setdefault('simplecache', dict(cache_storage=cache_dir))
-        path1 = f'simplecache::{path1}'
+        if use_cache:  # Cache seems to confuse the gcs filesystem.
+            home_dir = torch_geometric.get_home_dir()
+            cache_dir = osp.join(home_dir, 'simplecache', uuid4().hex)
+            kwargs.setdefault('simplecache', dict(cache_storage=cache_dir))
+            path1 = f'simplecache::{path1}'
 
     # Handle automatic extraction:
     if extract and path1.endswith('.tar.gz'):
@@ -147,7 +149,7 @@ def cp(
                         break
                     f_to.write(chunk)
 
-    if clear_cache and cache_dir is not None:
+    if use_cache and clear_cache and cache_dir is not None:
         rm(cache_dir)
 
 


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 27ff274e7..0b43ac138 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -33,12 +33,12 @@ def test_basic():
 def test_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 adj.sparse_size == (3, 3)
     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 adj.sparse_size == (3, 3)
     assert torch.equal(adj._colptr, torch.tensor([0, 1, 3, 4]))
 
 
@@ -160,14 +160,14 @@ def test_flip():
     out = adj.flip(0)
     assert isinstance(out, EdgeIndex)
     assert torch.equal(out, torch.tensor([[1, 0, 2, 1], [0, 1, 1, 2]]))
-    assert out.sparse_size == (None, 3)
+    assert out.sparse_size == (3, 3)
     assert out.sort_order == 'col'
     assert torch.equal(out._colptr, torch.tensor([0, 1, 3, 4]))
 
     out = adj.flip([0, 1])
     assert isinstance(out, EdgeIndex)
     assert torch.equal(out, torch.tensor([[1, 2, 0, 1], [2, 1, 1, 0]]))
-    assert out.sparse_size == (None, 3)
+    assert out.sparse_size == (3, 3)
     assert out.sort_order is None
     assert out._colptr is None
 
@@ -222,6 +222,58 @@ def test_getitem():
     assert not isinstance(out, EdgeIndex)
 
 
+def test_to_sparse_coo():
+    adj = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]])
+    if torch_geometric.typing.WITH_PT20:
+        out = adj.to_sparse(layout=torch.sparse_coo)
+    else:
+        out = adj.to_sparse()
+    assert isinstance(out, Tensor)
+    assert out.layout == torch.sparse_coo
+    assert out.size() == (3, 3)
+    assert torch.equal(adj, out._indices())
+
+    # Test clunky dispatch logic for `to_sparse_coo()`:
+    adj = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]])
+    out = adj.to_sparse_coo()
+    assert isinstance(out, Tensor)
+    assert out.layout == torch.sparse_coo
+    assert out.size() == (3, 3)
+    assert torch.equal(adj, out._indices())
+
+
+def test_to_sparse_csr():
+    with pytest.raises(ValueError, match="not sorted by rows"):
+        EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]]).to_sparse_csr()
+
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+    if torch_geometric.typing.WITH_PT20:
+        out = adj.to_sparse(layout=torch.sparse_csr)
+    else:
+        out = adj.to_sparse_csr()
+    assert isinstance(out, Tensor)
+    assert out.layout == torch.sparse_csr
+    assert out.size() == (3, 3)
+    assert torch.equal(adj._rowptr, out.crow_indices())
+    assert torch.equal(adj[1], out.col_indices())
+
+
+def test_to_sparse_csc():
+    with pytest.raises(ValueError, match="not sorted by columns"):
+        EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]]).to_sparse_csc()
+
+    adj = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]], sort_order='col')
+    if torch_geometric.typing.WITH_PT20:
+        out = adj.to_sparse(layout=torch.sparse_csc)
+    else:
+        out = adj.to_sparse_csc()
+    assert isinstance(out, Tensor)
+    assert out.layout == torch.sparse_csc
+    assert out.size() == (3, 3)
+    assert torch.equal(adj._colptr, out.ccol_indices())
+    assert torch.equal(adj[0], out.row_indices())
+
+
 def test_save_and_load(tmp_path):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     adj.fill_cache()