Repository: pyg-team/pytorch_geometric
Pull request #8447 (2023-11-27 09:31:10)

Properly handle data type conversion in `EdgeIndex`

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

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8447?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/8447?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvZWRnZV9pbmRleC5weQ==) | 91.66% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8447?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    #8447   +/-   ##
=======================================
  Coverage   87.64%   87.64%           
=======================================
  Files         477      477           
  Lines       29161    29163    +2     
=======================================
+ Hits        25557    25559    +2     
  Misses       3604     3604           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8447?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 54328efa0..b6e733e49 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -1,12 +1,25 @@
 import functools
 from enum import Enum
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Set, Tuple, Union
 
 import torch
 from torch import Tensor
 
 HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
 
+SUPPORTED_DTYPES: Set[torch.dtype] = {
+    torch.uint8,
+    torch.int8,
+    torch.int16,
+    torch.int32,
+    torch.int64,
+}
+
+
+class SortOrder(Enum):
+    ROW = 'row'
+    COL = 'col'
+
 
 def implements(torch_function: Callable):
     r"""Registers a :pytorch:`PyTorch` function override."""
@@ -18,10 +31,11 @@ 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_valid_dtype(tensor: Tensor):
+    if tensor.dtype not in SUPPORTED_DTYPES:
+        raise ValueError(f"'EdgeIndex' holds an unsupported data type "
+                         f"(got '{tensor.dtype}', but expected one of "
+                         f"{SUPPORTED_DTYPES})")
 
 
 def assert_two_dimensional(tensor: Tensor):
@@ -33,11 +47,6 @@ def assert_two_dimensional(tensor: Tensor):
                          f"[2, *] (got {list(tensor.size())})")
 
 
-class SortOrder(Enum):
-    ROW = 'row'
-    COL = 'col'
-
-
 class EdgeIndex(Tensor):
     # See "https://pytorch.org/docs/stable/notes/extending.html"
     # for a basic tutorial on how to subclass `torch.Tensor`.
@@ -74,7 +83,7 @@ class EdgeIndex(Tensor):
                             f"{set(kwargs.keys())})")
 
         assert isinstance(data, Tensor)
-        assert_integral_type(data)
+        assert_valid_dtype(data)
         assert_two_dimensional(data)
 
         out = super().__new__(cls, data)
@@ -90,7 +99,7 @@ class EdgeIndex(Tensor):
         * that :class:`EdgeIndex` only holds valid entries.
         * that the sort order is correctly set.
         """
-        assert_integral_type(self)
+        assert_valid_dtype(self)
         assert_two_dimensional(self)
 
         if self.numel() > 0 and self.min() < 0:
@@ -223,7 +232,10 @@ def apply_(
 @implements(Tensor.to)
 def to(tensor: EdgeIndex, *args, **kwargs) -> EdgeIndex:
     out = apply_(tensor, Tensor.to, *args, **kwargs)
-    assert_integral_type(out)
+
+    if out.dtype not in SUPPORTED_DTYPES:
+        out = out.as_tensor()
+
     return out
 
 


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index bc463eccf..8f1e452af 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -1,5 +1,4 @@
 import torch
-from torch import Tensor
 
 from torch_geometric.data.edge_index import EdgeIndex
 from torch_geometric.testing import onlyCUDA, withCUDA
@@ -13,9 +12,9 @@ def test_edge_index():
     assert adj.sparse_size == (3, 3)
     assert adj.sort_order is None
 
-    assert isinstance(adj.as_tensor(), Tensor)
+    assert not isinstance(adj.as_tensor(), EdgeIndex)
 
-    assert isinstance(adj + 1, Tensor)
+    assert not isinstance(adj + 1, EdgeIndex)
 
 
 def test_edge_index_fill_cache():
@@ -38,6 +37,10 @@ def test_edge_index_to(device):
     assert isinstance(out, EdgeIndex)
     assert out.dtype == torch.int
 
+    out = adj.to(torch.float)
+    assert not isinstance(out, EdgeIndex)
+    assert out.dtype == torch.float
+
     out = adj.to(device)
     assert isinstance(out, EdgeIndex)
     assert out.device == device
@@ -68,4 +71,4 @@ def test_edge_index_cat():
 
     out = torch.cat([adj1, adj2], dim=0)
     assert out.size() == (4, 4)
-    assert isinstance(out, Tensor)
+    assert not isinstance(out, EdgeIndex)