Repository: pyg-team/pytorch_geometric
Pull request #8690 (2023-12-30 15:08:20)

Allow `EdgeIndex` creation from `torch.sparse` tensors

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8690?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`180a11d`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/180a11d14b06329f250a1e3042dbf870f8977252?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.49% compared to head [(`6675891`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8690?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.49%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8690?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/edge\_index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8690?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2VkZ2VfaW5kZXgucHk=) | 93.75% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8690?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    #8690   +/-   ##
=======================================
  Coverage   89.49%   89.49%           
=======================================
  Files         481      481           
  Lines       31047    31076   +29     
=======================================
+ Hits        27786    27813   +27     
- Misses       3261     3263    +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8690?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/edge_index.py b/torch_geometric/edge_index.py
index 46bd3e771..782f72933 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -113,6 +113,11 @@ def maybe_sub(
                  for v, o in zip(value, other))
 
 
+def ptr2index(ptr: Tensor, output_size: Optional[int] = None) -> Tensor:
+    index = torch.arange(ptr.numel() - 1, dtype=ptr.dtype, device=ptr.device)
+    return index.repeat_interleave(ptr.diff(), output_size=output_size)
+
+
 def assert_valid_dtype(tensor: Tensor) -> None:
     if tensor.dtype not in SUPPORTED_DTYPES:
         raise ValueError(f"'EdgeIndex' holds an unsupported data type "
@@ -270,15 +275,52 @@ class EdgeIndex(Tensor):
                             f"{set(kwargs.keys())})")
 
         assert isinstance(data, Tensor)
-        assert_valid_dtype(data)
-        assert_two_dimensional(data)
-        assert_contiguous(data)
+
+        indptr: Optional[Tensor] = None
 
         if isinstance(data, cls):  # If passed `EdgeIndex`, inherit metadata:
+            indptr = data._indptr
             sparse_size = sparse_size or data.sparse_size()
             sort_order = sort_order or data.sort_order
             is_undirected = is_undirected or data.is_undirected
 
+        # Convert `torch.sparse` tensors to `EdgeIndex` representation:
+        if data.layout == torch.sparse_coo:
+            sort_order = SortOrder.ROW
+            sparse_size = sparse_size or (data.size(0), data.size(1))
+            data = data.indices()
+
+        if data.layout == torch.sparse_csr:
+            indptr = data.crow_indices()
+            col = data.col_indices()
+
+            assert isinstance(indptr, Tensor)
+            row = ptr2index(indptr, output_size=col.numel())
+
+            sort_order = SortOrder.ROW
+            sparse_size = sparse_size or (data.size(0), data.size(1))
+            if sparse_size[0] is not None and sparse_size[0] != data.size(0):
+                indptr = None
+            data = torch.stack([row, col], dim=0)
+
+        if (torch_geometric.typing.WITH_PT112
+                and data.layout == torch.sparse_csc):
+            row = data.row_indices()
+            indptr = data.ccol_indices()
+
+            assert isinstance(indptr, Tensor)
+            col = ptr2index(indptr, output_size=row.numel())
+
+            sort_order = SortOrder.COL
+            sparse_size = sparse_size or (data.size(0), data.size(1))
+            if sparse_size[1] is not None and sparse_size[1] != data.size(1):
+                indptr = None
+            data = torch.stack([row, col], dim=0)
+
+        assert_valid_dtype(data)
+        assert_two_dimensional(data)
+        assert_contiguous(data)
+
         if sparse_size is None:
             sparse_size = (None, None)
 
@@ -298,9 +340,9 @@ class EdgeIndex(Tensor):
         out._sparse_size = sparse_size
         out._sort_order = None if sort_order is None else SortOrder(sort_order)
         out._is_undirected = is_undirected
+        out._indptr = indptr
 
         if isinstance(data, cls):  # If passed `EdgeIndex`, inherit metadata:
-            out._indptr = data._indptr
             out._T_perm = data._T_perm
             out._T_index = data._T_index
             out._T_indptr = data._T_indptr


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 7d7083df0..e8df4348e 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -48,15 +48,15 @@ def test_basic(dtype, device):
     adj.validate()
     assert isinstance(adj, EdgeIndex)
 
-    dtype_str = '' if dtype == torch.long else f', dtype={dtype}'
     if torch_geometric.typing.WITH_PT112:
-        assert str(adj) == (
-            f'EdgeIndex([[0, 1, 1, 2],\n'
-            f'           [1, 0, 2, 1]]{dtype_str}, sparse_size=(3, 3), nnz=4)')
+        assert str(adj).startswith('EdgeIndex([[0, 1, 1, 2],\n'
+                                   '           [1, 0, 2, 1]], ')
     else:
-        assert str(adj) == (
-            f'tensor([[0, 1, 1, 2],\n'
-            f'        [1, 0, 2, 1]]{dtype_str}, sparse_size=(3, 3), nnz=4)')
+        assert str(adj).startswith('tensor([[0, 1, 1, 2],\n'
+                                   '        [1, 0, 2, 1]], ')
+    assert str(adj).endswith('sparse_size=(3, 3), nnz=4)')
+    assert (f"device='{device}'" in str(adj)) == adj.is_cuda
+    assert (f'dtype={dtype}' in str(adj)) == (dtype != torch.long)
 
     assert adj.dtype == dtype
     assert adj.device == device
@@ -98,6 +98,31 @@ def test_identity(dtype, device, is_undirected):
     assert out.is_undirected == adj.is_undirected
 
 
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_sparse_tensor(dtype, device):
+    kwargs = dict(dtype=dtype, device=device, is_undirected=True)
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row', **kwargs)
+
+    out = EdgeIndex(adj.to_sparse_coo())
+    assert out.equal(adj)
+    assert out.sort_order == 'row'
+    assert out.sparse_size() == (3, 3)
+    assert out._indptr is None
+
+    out = EdgeIndex(adj.to_sparse_csr())
+    assert out.equal(adj)
+    assert out.sort_order == 'row'
+    assert out.sparse_size() == (3, 3)
+    assert out._indptr.equal(tensor([0, 1, 3, 4], device=device))
+
+    out = EdgeIndex(adj.to_sparse_csc())
+    assert out.equal(adj.sort_by('col')[0])
+    assert out.sort_order == 'col'
+    assert out.sparse_size() == (3, 3)
+    assert out._indptr.equal(tensor([0, 1, 3, 4], device=device))
+
+
 def test_set_tuple_item():
     tmp = (0, 1, 2)
     assert set_tuple_item(tmp, 0, 3) == (3, 1, 2)