Repository: pyg-team/pytorch_geometric
Pull request #8441 (2023-11-26 13:25:32)

Introduce `EdgeIndex` sub-class
This is a currently hidden class that inherits from `torch.Tensor`. Its main goal is to store additional meta-information to an `edge_index` tensor, such as its underlying sparse size, and whether it is sorted or not.

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

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8441?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/8441?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 81.81% | [8 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8441?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    #8441      +/-   ##
==========================================
- Coverage   87.65%   87.65%   -0.01%     
==========================================
  Files         476      477       +1     
  Lines       29027    29071      +44     
==========================================
+ Hits        25445    25481      +36     
- Misses       3582     3590       +8     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8441?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
new file mode 100644
index 000000000..c6ef6b6cb
--- /dev/null
+++ b/torch_geometric/data/edge_index.py
@@ -0,0 +1,77 @@
+from enum import Enum
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+
+class SortOrder(Enum):
+    ROW = 'row'
+    COL = 'col'
+
+
+class EdgeIndex(Tensor):
+    _sparse_size: Tuple[Optional[int], Optional[int]] = (None, None)
+    _sort_order: Optional[SortOrder] = None
+
+    _rowptr: Optional[Tensor] = None
+    _colptr: Optional[Tensor] = None
+    _csr2csc: Optional[Tensor] = None
+    _csc2csr: Optional[Tensor] = None
+
+    def __new__(
+        cls,
+        data,
+        *args,
+        sparse_size: Tuple[Optional[int], Optional[int]] = (None, None),
+        sort_order: Optional[SortOrder] = None,
+        **kwargs,
+    ):
+        if not isinstance(data, Tensor):
+            data = torch.tensor(data, *args, **kwargs)
+        elif len(args) > 0:
+            raise TypeError(
+                f"new() received an invalid combination of arguments - got "
+                f"(Tensor, {', '.join(str(type(arg)) for arg in args)})")
+        elif len(kwargs) > 0:
+            raise TypeError(f"new() received invalid keyword arguments - got "
+                            f"{set(kwargs.keys())})")
+
+        assert isinstance(data, Tensor)
+
+        if data.is_floating_point():
+            raise ValueError(f"'{cls.__name__}' needs to be of integral type "
+                             f"(got '{data.dtype}')")
+        if data.dim() != 2:
+            raise ValueError(f"'{cls.__name__}' needs to be two-dimensional "
+                             f"(got {data.dim()} dimensions)")
+        if data.size(0) != 2:
+            raise ValueError(f"'{cls.__name__}' needs to have a shape of "
+                             f"[2, *] (got {list(data.size())})")
+
+        out = super().__new__(cls, data)
+
+        # Attach metadata:
+        out._sparse_size = sparse_size
+        out._sort_order = None if sort_order is None else SortOrder(sort_order)
+
+        return out
+
+    @property
+    def sparse_size(self) -> Tuple[Optional[int], Optional[int]]:
+        return self._sparse_size
+
+    @property
+    def sort_order(self) -> Optional[SortOrder]:
+        return self._sort_order
+
+    def validate(self) -> bool:
+        raise NotImplementedError  # TODO
+
+    def to_tensor(self) -> Tensor:
+        return torch.as_tensor(self)  # TODO
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        # TODO Disallow in-place operations?
+        return super().__torch_function__(func, types, args, kwargs)


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
new file mode 100644
index 000000000..425a6ead4
--- /dev/null
+++ b/test/data/test_edge_index.py
@@ -0,0 +1,19 @@
+import torch
+
+from torch_geometric.data.edge_index import EdgeIndex
+
+
+def test_edge_index():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(10, 2))
+    assert adj.sparse_size == (10, 2)
+    assert adj.sort_order is None
+
+    out = adj + 1
+    assert isinstance(out, EdgeIndex)
+    assert out.sparse_size == (None, None)
+    assert out.sort_order is None
+
+    out = torch.cat([adj, adj], dim=1)
+    assert isinstance(out, EdgeIndex)
+    assert out.sparse_size == (None, None)
+    assert out.sort_order is None