Repository: pyg-team/pytorch_geometric
Pull request #8445 (2023-11-27 06:31:55)

`EdgeIndex` device transfers

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8445?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `11 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`77e0b90`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/77e0b9073a98f66774136bc036dd3b600f70496a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.95% compared to head [(`6be6e21`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8445?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 87.60%.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8445?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/8445?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 71.05% | [11 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8445?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    #8445      +/-   ##
==========================================
- Coverage   87.95%   87.60%   -0.36%     
==========================================
  Files         477      477              
  Lines       29119    29150      +31     
==========================================
- Hits        25613    25538      -75     
- Misses       3506     3612     +106     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8445?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 eb9b9cc09..bc625627a 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -18,6 +18,21 @@ def implements(torch_function: Callable):
     return decorator
 
 
+def assert_integral_type(tensor: Tensor):
+    if tensor.is_floating_point():
+        raise ValueError(f"'EdgeIndex' needs to be of integral type "
+                         f"(got '{tensor.dtype}')")
+
+
+def assert_two_dimensional(tensor: Tensor):
+    if tensor.dim() != 2:
+        raise ValueError(f"'EdgeIndex' needs to be two-dimensional "
+                         f"(got {tensor.dim()} dimensions)")
+    if tensor.size(0) != 2:
+        raise ValueError(f"'EdgeIndex' needs to have a shape of "
+                         f"[2, *] (got {list(tensor.size())})")
+
+
 class SortOrder(Enum):
     ROW = 'row'
     COL = 'col'
@@ -59,16 +74,8 @@ class EdgeIndex(Tensor):
                             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())})")
+        assert_integral_type(data)
+        assert_two_dimensional(data)
 
         out = super().__new__(cls, data)
 
@@ -83,6 +90,9 @@ class EdgeIndex(Tensor):
         * that :class:`EdgeIndex` only holds valid entries.
         * that the sort order is correctly set.
         """
+        assert_integral_type(self)
+        assert_two_dimensional(self)
+
         if self.numel() > 0 and self.min() < 0:
             raise ValueError(f"'{self.__class__.__name__}' contains negative "
                              f"indices (got {int(self.min())})")
@@ -157,7 +167,51 @@ class EdgeIndex(Tensor):
 
         # 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)
+        return Tensor.__torch_function__(func, types, args, kwargs)
+
+
+def apply_(
+    tensor: EdgeIndex,
+    fn: Callable,
+    *args,
+    **kwargs,
+) -> EdgeIndex:
+
+    out = Tensor.__torch_function__(fn, (Tensor, ), (tensor, ) + args, kwargs)
+    out = out.as_subclass(EdgeIndex)
+
+    # Copy metadata:
+    out._sparse_size = out._sparse_size
+    out._sort_order = out._sort_order
+
+    # Convert metadata:
+    if tensor._rowptr is not None:
+        out._rowptr = fn(tensor._rowptr, *args, **kwargs)
+    if tensor._colptr is not None:
+        out._colptr = fn(tensor._colptr, *args, **kwargs)
+    if tensor._csr2csc is not None:
+        out._csr2csc = fn(tensor._csr2csc, *args, **kwargs)
+    if tensor._csc2csr is not None:
+        out._csc2csr = fn(tensor._csc2csr, *args, **kwargs)
+
+    return out
+
+
+@implements(Tensor.to)
+def to(tensor: EdgeIndex, *args, **kwargs) -> EdgeIndex:
+    out = apply_(tensor, Tensor.to, *args, **kwargs)
+    assert_integral_type(out)
+    return out
+
+
+@implements(Tensor.cpu)
+def cpu(tensor: EdgeIndex, *args, **kwargs) -> EdgeIndex:
+    return apply_(tensor, Tensor.cpu, *args, **kwargs)
+
+
+@implements(Tensor.cuda)
+def cuda(tensor: EdgeIndex, *args, **kwargs) -> EdgeIndex:
+    return apply_(tensor, Tensor.cuda, *args, **kwargs)
 
 
 @implements(torch.cat)


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 792605150..85a6bd0a2 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -2,19 +2,49 @@ import torch
 from torch import Tensor
 
 from torch_geometric.data.edge_index import EdgeIndex
+from torch_geometric.testing import onlyCUDA, withCUDA
 
 
 def test_edge_index():
     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 str(adj) == ('EdgeIndex([[0, 1, 1, 2],\n'
+                        '           [1, 0, 2, 1]])')
     assert adj.sparse_size == (3, 3)
     assert adj.sort_order is None
 
+    assert isinstance(adj.as_tensor(), Tensor)
+
     out = adj + 1
     assert isinstance(out, Tensor)
 
 
+@withCUDA
+def test_edge_index_to(device):
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    out = adj.to(torch.int)
+    assert isinstance(out, EdgeIndex)
+    assert out.dtype == torch.int
+
+    out = adj.to(device)
+    assert isinstance(out, EdgeIndex)
+    assert out.device == device
+
+
+@onlyCUDA
+def test_edge_index_cpu_cuda():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    out = adj.cuda()
+    assert isinstance(out, EdgeIndex)
+    assert out.is_cuda
+
+    out = out.cpu()
+    assert isinstance(out, EdgeIndex)
+    assert not out.is_cuda
+
+
 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))