Repository: pyg-team/pytorch_geometric
Pull request #8442 (2023-11-26 20:31:27)

Introduce `HANDLED_FUNCTIONS` to `EdgeIndex`

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

> :exclamation: Current head 768dc63 differs from pull request most recent head c169646. Consider uploading reports for the commit c169646 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8442?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/8442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 72.88% | [16 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8442?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    #8442      +/-   ##
==========================================
- Coverage   87.65%   87.62%   -0.04%     
==========================================
  Files         477      477              
  Lines       29071    29121      +50     
==========================================
+ Hits        25481    25516      +35     
- Misses       3590     3605      +15     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8442?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 c6ef6b6cb..eb9b9cc09 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -1,9 +1,22 @@
+import functools
 from enum import Enum
-from typing import Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
 
+HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
+
+
+def implements(torch_function: Callable):
+    r"""Registers a :pytorch:`PyTorch` function override."""
+    @functools.wraps(torch_function)
+    def decorator(my_function: Callable):
+        HANDLED_FUNCTIONS[torch_function] = my_function
+        return my_function
+
+    return decorator
+
 
 class SortOrder(Enum):
     ROW = 'row'
@@ -11,20 +24,28 @@ class SortOrder(Enum):
 
 
 class EdgeIndex(Tensor):
+    # See "https://pytorch.org/docs/stable/notes/extending.html"
+    # for a basic tutorial on how to subclass `torch.Tensor`.
+
+    # The size of the underlying sparse matrix:
     _sparse_size: Tuple[Optional[int], Optional[int]] = (None, None)
+
+    # Whether the `edge_index` represented is non-sorted (`None`), or sorted
+    # based on row or column values.
     _sort_order: Optional[SortOrder] = None
 
-    _rowptr: Optional[Tensor] = None
-    _colptr: Optional[Tensor] = None
-    _csr2csc: Optional[Tensor] = None
-    _csc2csr: Optional[Tensor] = None
+    # An additional data cache:
+    _rowptr: Optional[Tensor] = None  # The CSR `rowptr` in case sorted by row.
+    _colptr: Optional[Tensor] = None  # The CSC `colptr` in case sorted by col.
+    _csr2csc: Optional[Tensor] = None  # Permutation from CSR to CSC.
+    _csc2csr: Optional[Tensor] = None  # Permutation from CSC to CSR.
 
     def __new__(
         cls,
         data,
         *args,
         sparse_size: Tuple[Optional[int], Optional[int]] = (None, None),
-        sort_order: Optional[SortOrder] = None,
+        sort_order: Optional[Union[str, SortOrder]] = None,
         **kwargs,
     ):
         if not isinstance(data, Tensor):
@@ -57,21 +78,116 @@ class EdgeIndex(Tensor):
 
         return out
 
+    def validate(self) -> 'EdgeIndex':
+        r"""Validates the :class:`EdgeIndex` representation, i.e., it ensures
+        * that :class:`EdgeIndex` only holds valid entries.
+        * that the sort order is correctly set.
+        """
+        if self.numel() > 0 and self.min() < 0:
+            raise ValueError(f"'{self.__class__.__name__}' contains negative "
+                             f"indices (got {int(self.min())})")
+
+        if (self.numel() > 0 and self.num_rows is not None
+                and self[0].max() >= self.num_rows):
+            raise ValueError(f"'{self.__class__.__name__}' contains larger "
+                             f"indices than its number of rows "
+                             f"(got {int(self[0].max())}, but expected values "
+                             f"smaller than {self.num_rows})")
+
+        if (self.numel() > 0 and self.num_cols is not None
+                and self[1].max() >= self.num_cols):
+            raise ValueError(f"'{self.__class__.__name__}' contains larger "
+                             f"indices than its number of columns "
+                             f"(got {int(self[1].max())}, but expected values "
+                             f"smaller than {self.num_cols})")
+
+        if self._sort_order == SortOrder.ROW and (self[0].diff() < 0).any():
+            raise ValueError(f"'{self.__class__.__name__}' is not sorted by "
+                             f"row indices")
+
+        if self._sort_order == SortOrder.COL and (self[1].diff() < 0).any():
+            raise ValueError(f"'{self.__class__.__name__}' is not sorted by "
+                             f"column indices")
+
+        return self
+
     @property
     def sparse_size(self) -> Tuple[Optional[int], Optional[int]]:
+        r"""The size of the underlying sparse matrix."""
         return self._sparse_size
 
     @property
-    def sort_order(self) -> Optional[SortOrder]:
-        return self._sort_order
+    def num_rows(self) -> Optional[int]:
+        r"""The number of rows of the underlying sparse matrix."""
+        return self._sparse_size[0]
 
-    def validate(self) -> bool:
-        raise NotImplementedError  # TODO
+    @property
+    def num_cols(self) -> Optional[int]:
+        r"""The number of columns of the underlying sparse matrix."""
+        return self._sparse_size[1]
 
-    def to_tensor(self) -> Tensor:
-        return torch.as_tensor(self)  # TODO
+    @property
+    def sort_order(self) -> Optional[str]:
+        r"""The sort order of indices, either :obj:`"row"`, :obj:`"col"` or
+        :obj:`None`.
+        """
+        return None if self._sort_order is None else self._sort_order.value
+
+    def as_tensor(self) -> Tensor:
+        r"""Zero-copies the :class:`EdgeIndex` representation back to a
+        :class:`torch.Tensor` representation.
+        """
+        return self.as_subclass(Tensor)
 
     @classmethod
     def __torch_function__(cls, func, types, args=(), kwargs=None):
-        # TODO Disallow in-place operations?
-        return super().__torch_function__(func, types, args, kwargs)
+        # `EdgeIndex` should be treated as a regular PyTorch tensor for all
+        # standard PyTorch functionalities. However,
+        # * some of its metadata can be transferred to new functions, e.g.,
+        #   `torch.cat(dim=1)` can inherit the sparse matrix size, or
+        #   `torch.narrow(dim=1)` can inherit cached pointers.
+        # * not all operations lead to valid `EdgeIndex` tensors again, e.g.,
+        #   `torch.sum()` does not yield a `EdgeIndex` as its output, or
+        #   `torch.cat(dim=0) violates the [2, *] shape assumption.
+
+        # To account for this, we hold a number of `HANDLED_FUNCTIONS` that
+        # implement specific functions for valid `EdgeIndex` routines.
+        if func in HANDLED_FUNCTIONS:
+            return HANDLED_FUNCTIONS[func](*args, **(kwargs or {}))
+
+        # For all other PyTorch functions, we return a vanilla PyTorch tensor.
+        types = tuple(torch.Tensor if issubclass(t, cls) else t for t in types)
+        return torch.Tensor.__torch_function__(func, types, args, kwargs)
+
+
+@implements(torch.cat)
+def cat(
+    tensors: List[Union[EdgeIndex, Tensor]],
+    dim: int = 0,
+) -> Union[EdgeIndex, Tensor]:
+
+    out = Tensor.__torch_function__(torch.cat, (Tensor, ), (tensors, dim))
+
+    if dim != 1 and dim != -1:  # No valid `EdgeIndex` anymore.
+        return out
+
+    out = out.as_subclass(EdgeIndex)
+
+    # Post-process `sparse_size`:
+    num_rows = 0
+    for tensor in tensors:
+        if not isinstance(tensor, EdgeIndex) or tensor.num_rows is None:
+            num_rows = None
+            break
+        num_rows = max(num_rows, tensor.num_rows)
+
+    num_cols = 0
+    for tensor in tensors:
+        if not isinstance(tensor, EdgeIndex) or tensor.num_cols is None:
+            num_cols = None
+            break
+        num_cols = max(num_cols, tensor.num_cols)
+
+    out._sparse_size = (num_rows, num_cols)
+
+    return out


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 425a6ead4..792605150 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -1,19 +1,30 @@
 import torch
+from torch import Tensor
 
 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)
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(3, 3))
+    assert isinstance(adj, EdgeIndex)
+    assert isinstance(adj.as_tensor(), Tensor)
+    assert adj.sparse_size == (3, 3)
     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
+    assert isinstance(out, Tensor)
+
 
-    out = torch.cat([adj, adj], dim=1)
+def test_edge_index_cat():
+    adj1 = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(3, 3))
+    adj2 = EdgeIndex([[1, 2, 2, 3], [2, 1, 3, 2]], sparse_size=(4, 4))
+
+    out = torch.cat([adj1, adj2], dim=1)
+    assert out.size() == (2, 8)
     assert isinstance(out, EdgeIndex)
-    assert out.sparse_size == (None, None)
+    assert out.sparse_size == (4, 4)
     assert out.sort_order is None
+
+    out = torch.cat([adj1, adj2], dim=0)
+    assert out.size() == (4, 4)
+    assert isinstance(out, Tensor)