Repository: pyg-team/pytorch_geometric
Pull request #8640 (2023-12-19 17:26:44)

Inherit metadata for `EdgeIndex(edge_index)`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8640?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 [(`d41de53`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/d41de53b8d3e273acb6fd3d7747139b5534ae022?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.19% compared to head [(`a497289`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8640?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.20%.


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


```diff
@@           Coverage Diff           @@
##           master    #8640   +/-   ##
=======================================
  Coverage   89.19%   89.20%           
=======================================
  Files         481      481           
  Lines       30709    30711    +2     
=======================================
+ Hits        27392    27395    +3     
+ Misses       3317     3316    -1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8640?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 91646cdcf..139aff97b 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -207,7 +207,7 @@ class EdgeIndex(Tensor):
         cls: Type,
         data: Any,
         *args: Any,
-        sparse_size: Tuple[Optional[int], Optional[int]] = (None, None),
+        sparse_size: Optional[Tuple[Optional[int], Optional[int]]] = None,
         sort_order: Optional[Union[str, SortOrder]] = None,
         is_undirected: bool = False,
         **kwargs: Any,
@@ -227,6 +227,14 @@ class EdgeIndex(Tensor):
         assert_two_dimensional(data)
         assert_contiguous(data)
 
+        if isinstance(data, cls):  # If passed `EdgeIndex`, inherit metadata:
+            sparse_size = sparse_size or data.sparse_size()
+            sort_order = sort_order or data.sort_order
+            is_undirected = is_undirected or data.is_undirected
+
+        if sparse_size is None:
+            sparse_size = (None, None)
+
         if is_undirected:
             assert_symmetric(sparse_size)
             if sparse_size[0] is not None and sparse_size[1] is None:
@@ -244,6 +252,13 @@ class EdgeIndex(Tensor):
         out._sort_order = None if sort_order is None else SortOrder(sort_order)
         out._is_undirected = is_undirected
 
+        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
+            out._value = out._value
+
         return out
 
     # Validation ##############################################################


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 580c8404a..484a524b9 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -74,6 +74,22 @@ def test_basic(dtype, device):
     assert out.device == device
 
 
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+@pytest.mark.parametrize('is_undirected', IS_UNDIRECTED)
+def test_identity(dtype, device, is_undirected):
+    kwargs = dict(dtype=dtype, device=device, is_undirected=is_undirected)
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(3, 3), **kwargs)
+
+    out = EdgeIndex(adj)
+    assert out.data_ptr() == adj.data_ptr()
+    assert out.dtype == adj.dtype
+    assert out.device == adj.device
+    assert out.sparse_size() == adj.sparse_size()
+    assert out.sort_order == adj.sort_order
+    assert out.is_undirected == adj.is_undirected
+
+
 def test_set_tuple_item():
     tmp = (0, 1, 2)
     assert set_tuple_item(tmp, 0, 3) == (3, 1, 2)