Repository: pyg-team/pytorch_geometric
Pull request #8498 (2023-12-01 06:52:32)

`[EdgeIndex]` Move modified `torch__function__` methods to class level

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8498?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 [(`369c955`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/369c955171e145ad68565919a9d970a22214fb57?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.76% compared to head [(`4e08f36`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8498?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.39%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8498?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/8498?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 95.12% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8498?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    #8498      +/-   ##
==========================================
- Coverage   88.76%   88.39%   -0.38%     
==========================================
  Files         478      478              
  Lines       29479    29472       -7     
==========================================
- Hits        26168    26052     -116     
- Misses       3311     3420     +109     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8498?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 3bbfd61a2..354ff80eb 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -381,6 +381,82 @@ class EdgeIndex(Tensor):
 
         return torch.return_types.sort([out, perm])
 
+    def to_dense(
+        self,
+        value: Optional[Tensor] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Tensor:
+
+        # TODO Respect duplicate edges.
+
+        dtype = value.dtype if value is not None else dtype
+
+        size = self.get_sparse_size()
+        if value is not None and value.dim() > 1:
+            size = size + value.shape[1:]
+
+        out = torch.zeros(size, dtype=dtype, device=self.device)
+        out[self[0], self[1]] = value if value is not None else 1
+
+        return out
+
+    def to_sparse_coo(self, value: Optional[Tensor] = None) -> Tensor:
+        value = self._get_value() if value is None else value
+        out = torch.sparse_coo_tensor(
+            indices=self.as_tensor(),
+            values=value,
+            size=self.get_sparse_size(),
+            device=self.device,
+            requires_grad=value.requires_grad,
+        )
+
+        if self._sort_order == SortOrder.ROW:
+            out = out._coalesced_(True)
+
+        return out
+
+    def to_sparse_csr(self, value: Optional[Tensor] = None) -> Tensor:
+        value = self._get_value() if value is None else value
+        return torch.sparse_csr_tensor(
+            crow_indices=self.get_rowptr(),
+            col_indices=self[1],
+            values=value,
+            size=self.get_sparse_size(),
+            device=self.device,
+            requires_grad=value.requires_grad,
+        )
+
+    def to_sparse_csc(self, value: Optional[Tensor] = None) -> Tensor:
+        if not torch_geometric.typing.WITH_PT112:
+            raise NotImplementedError(
+                "'to_sparse_csc' not supported for PyTorch < 1.12")
+
+        value = self._get_value() if value is None else value
+        return torch.sparse_csc_tensor(
+            ccol_indices=self.get_colptr(),
+            row_indices=self[0],
+            values=value,
+            size=self.get_sparse_size(),
+            device=self.device,
+            requires_grad=value.requires_grad,
+        )
+
+    def to_sparse(
+        self,
+        *,
+        layout: Optional[torch.layout] = None,
+        value: Optional[Tensor] = None,
+    ) -> Tensor:
+
+        if layout is None or layout == torch.sparse_coo:
+            return self.to_sparse_coo(value)
+        if layout == torch.sparse_csr:
+            return self.to_sparse_csr(value)
+        if torch_geometric.typing.WITH_PT112 and layout == torch.sparse_csc:
+            return self.to_sparse_csc(value)
+
+        raise ValueError(f"Unexpected tensor layout (got '{layout}')")
+
     def to_sparse_tensor(
         self,
         value: Optional[Tensor] = None,
@@ -403,6 +479,15 @@ class EdgeIndex(Tensor):
             trust_data=True,
         )
 
+    def matmul(
+        self,
+        other: Union[Tensor, 'EdgeIndex'],
+        input_value: Optional[Tensor] = None,
+        other_value: Optional[Tensor] = None,
+        reduce: ReduceType = 'sum',
+    ) -> Union[Tensor, Tuple['EdgeIndex', Tensor]]:
+        return matmul(self, other, input_value, other_value, reduce)
+
     @classmethod
     def __torch_function__(cls, func, types, args=(), kwargs=None):
         # `EdgeIndex` should be treated as a regular PyTorch tensor for all
@@ -636,104 +721,6 @@ def getitem(input: EdgeIndex, index: Any) -> Union[EdgeIndex, Tensor]:
     return out
 
 
-@implements(Tensor.to_dense)
-def to_dense(
-    tensor: EdgeIndex,
-    dtype: Optional[torch.dtype] = None,
-    value: Optional[Tensor] = None,
-) -> Tensor:
-
-    # TODO Respect duplicate edges.
-
-    dtype = value.dtype if value is not None else dtype
-
-    size = tensor.get_sparse_size()
-    if value is not None and value.dim() > 1:
-        size = size + value.shape[1:]
-
-    out = torch.zeros(size, dtype=dtype, device=tensor.device)
-    out[tensor[0], tensor[1]] = value if value is not None else 1
-
-    return out
-
-
-# `to_sparse_coo()` uses `to_sparse(layout=None)` dispatch logic:
-def to_sparse_coo(tensor: EdgeIndex, value: Optional[Tensor] = None) -> Tensor:
-    out = torch.sparse_coo_tensor(
-        indices=tensor.as_tensor(),
-        values=tensor._get_value() if value is None else value,
-        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, value: Optional[Tensor] = None) -> Tensor:
-    return torch.sparse_csr_tensor(
-        crow_indices=tensor.get_rowptr(),
-        col_indices=tensor[1],
-        values=tensor._get_value() if value is None else value,
-        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,
-        value: Optional[Tensor] = None,
-    ) -> Tensor:
-        return torch.sparse_csc_tensor(
-            ccol_indices=tensor.get_colptr(),
-            row_indices=tensor[0],
-            values=tensor._get_value() if value is None else value,
-            size=tensor.get_sparse_size(),
-            device=tensor.device,
-        )
-
-else:
-
-    def to_sparse_csc(
-        tensor: EdgeIndex,
-        value: Optional[Tensor] = None,
-    ) -> Tensor:
-        raise NotImplementedError(
-            "'to_sparse_csc' not supported for PyTorch < 1.12")
-
-
-if torch_geometric.typing.WITH_PT20:
-
-    @implements(Tensor.to_sparse)
-    def to_sparse(
-        tensor: EdgeIndex,
-        *,
-        layout: Optional[torch.layout] = None,
-        value: Optional[Tensor] = None,
-    ) -> Tensor:
-
-        if layout is None or layout == torch.sparse_coo:
-            return to_sparse_coo(tensor, value)
-        if layout == torch.sparse_csr:
-            return to_sparse_csr(tensor, value)
-        if torch_geometric.typing.WITH_PT112 and layout == torch.sparse_csc:
-            return to_sparse_csc(tensor, value)
-
-        raise ValueError(f"Unexpected tensor layout (got '{layout}')")
-
-else:
-
-    @implements(Tensor.to_sparse)
-    def to_sparse(tensor: EdgeIndex, value: Optional[Tensor] = None) -> Tensor:
-        return to_sparse_coo(tensor, value)
-
-
 class SparseDenseMatmul(torch.autograd.Function):
     @staticmethod
     def forward(
@@ -767,7 +754,7 @@ class SparseDenseMatmul(torch.autograd.Function):
 
         if input_value is None:
             input_value = input._get_value()
-        adj = to_sparse_csr(input, input_value)
+        adj = input.to_sparse_csr(input_value)
 
         return adj @ other
 
@@ -816,8 +803,6 @@ class SparseDenseMatmul(torch.autograd.Function):
         return None, other_grad, None, None
 
 
-@implements(torch.matmul)
-@implements(Tensor.matmul)
 def matmul(
     input: EdgeIndex,
     other: Union[Tensor, EdgeIndex],
@@ -836,14 +821,14 @@ def matmul(
         return SparseDenseMatmul.apply(input, other, input_value, reduce)
 
     if input._sort_order == SortOrder.COL:
-        input = to_sparse_csc(input, input_value)
+        input = input.to_sparse_csc(input_value)
     else:
-        input = to_sparse_csr(input, input_value)
+        input = input.to_sparse_csr(input_value)
 
     if other._sort_order == SortOrder.COL:
-        other = to_sparse_csc(other, other_value)
+        other = other.to_sparse_csc(other_value)
     else:
-        other = to_sparse_csr(other, other_value)
+        other = other.to_sparse_csr(other_value)
 
     out = torch.matmul(input, other)
     assert out.layout == torch.sparse_csr
@@ -861,8 +846,17 @@ def matmul(
     return edge_index, out.values()
 
 
+@implements(torch.matmul)
+@implements(Tensor.matmul)
+def _matmul1(
+    input: EdgeIndex,
+    other: Union[Tensor, EdgeIndex],
+) -> Union[Tensor, Tuple[EdgeIndex, Tensor]]:
+    return matmul(input, other)
+
+
 @implements(torch.sparse.mm)
-def _matmul(
+def _matmul2(
     mat1: EdgeIndex,
     mat2: Union[Tensor, EdgeIndex],
     reduce: ReduceType = 'sum',


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index cc90375a3..7113d7041 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -7,7 +7,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric
-from torch_geometric.data.edge_index import EdgeIndex, matmul, to_dense
+from torch_geometric.data.edge_index import EdgeIndex
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import (
     disableExtensions,
@@ -262,14 +262,14 @@ def test_to_dense(dtype):
     assert out.tolist() == [[0, 1, 0], [1, 0, 1], [0, 1, 0]]
 
     value = torch.arange(1, 5, dtype=dtype or torch.float)
-    out = to_dense(adj, value=value)
+    out = adj.to_dense(value)
     assert isinstance(out, Tensor)
     assert out.size() == (3, 3)
     assert out.dtype == dtype or torch.float
     assert out.tolist() == [[0, 2, 0], [1, 0, 4], [0, 3, 0]]
 
     value = torch.arange(1, 5, dtype=dtype or torch.float).view(-1, 1)
-    out = to_dense(adj, value=value)
+    out = adj.to_dense(value)
     assert isinstance(out, Tensor)
     assert out.size() == (3, 3, 1)
     assert out.dtype == dtype or torch.float
@@ -357,19 +357,19 @@ def test_matmul_forward():
     assert out.is_sorted_by_row
     assert out._sparse_size == (3, 3)
     assert out._rowptr is not None
-    assert torch.allclose(to_dense(out, value=value), adj1_dense @ adj1_dense)
+    assert torch.allclose(out.to_dense(value), adj1_dense @ adj1_dense)
 
     out, value = adj1 @ adj2
     assert isinstance(out, EdgeIndex)
-    assert torch.allclose(to_dense(out, value=value), adj1_dense @ adj2_dense)
+    assert torch.allclose(out.to_dense(value), adj1_dense @ adj2_dense)
 
     out, value = adj2 @ adj1
     assert isinstance(out, EdgeIndex)
-    assert torch.allclose(to_dense(out, value=value), adj2_dense @ adj1_dense)
+    assert torch.allclose(out.to_dense(value), adj2_dense @ adj1_dense)
 
     out, value = adj2 @ adj2
     assert isinstance(out, EdgeIndex)
-    assert torch.allclose(to_dense(out, value=value), adj2_dense @ adj2_dense)
+    assert torch.allclose(out.to_dense(value), adj2_dense @ adj2_dense)
 
 
 def test_matmul_input_value():
@@ -378,8 +378,8 @@ def test_matmul_input_value():
     x = torch.randn(3, 1)
     value = torch.randn(4)
 
-    out = matmul(adj, x, input_value=value)
-    assert torch.allclose(out, to_dense(adj, value=value) @ x)
+    out = adj.matmul(x, input_value=value)
+    assert torch.allclose(out, adj.to_dense(value) @ x)
 
 
 def test_matmul_backward():
@@ -388,12 +388,12 @@ def test_matmul_backward():
     x1 = torch.randn(3, 1, requires_grad=True)
     value = torch.randn(4)
 
-    out = matmul(adj, x1, input_value=value)
+    out = adj.matmul(x1, input_value=value)
     grad_out = torch.randn_like(out)
     out.backward(grad_out)
 
     x2 = x1.detach().requires_grad_()
-    dense_adj = to_dense(adj, value=value)
+    dense_adj = adj.to_dense(value)
     out = dense_adj @ x2
     out.backward(grad_out)