Repository: pyg-team/pytorch_geometric
Pull request #10073 (2025-02-25 09:35:02)

Add support for `numpy()`/`tolist()` in `Index`/`EdgeIndex` 


Patch:
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index df9a5a686..f56c85142 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -17,6 +17,7 @@ from typing import (
     overload,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 from torch import Tensor
@@ -1246,6 +1247,14 @@ class EdgeIndex(Tensor):
         return torch._tensor_str._add_suffixes(prefix + tensor_str, suffixes,
                                                indent, force_newline=False)
 
+    def tolist(self) -> List[Any]:
+        """"""  # noqa: D419
+        return self._data.tolist()
+
+    def numpy(self, *, force: bool = False) -> np.ndarray:
+        """"""  # noqa: D419
+        return self._data.numpy(force=force)
+
     # Helpers #################################################################
 
     def _shallow_copy(self) -> 'EdgeIndex':
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index 3a9a89fc7..811b82cc9 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -12,6 +12,7 @@ from typing import (
     Union,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 from torch import Tensor
@@ -410,6 +411,14 @@ class Index(Tensor):
         return torch._tensor_str._add_suffixes(prefix + tensor_str, suffixes,
                                                indent, force_newline=False)
 
+    def tolist(self) -> List[Any]:
+        """"""  # noqa: D419
+        return self._data.tolist()
+
+    def numpy(self, *, force: bool = False) -> np.ndarray:
+        """"""  # noqa: D419
+        return self._data.numpy(force=force)
+
     # Helpers #################################################################
 
     def _shallow_copy(self) -> 'Index':


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 39ecbe2a7..dcb9599e4 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -2,6 +2,7 @@ import os.path as osp
 import warnings
 from typing import List, Optional
 
+import numpy as np
 import pytest
 import torch
 from torch import Tensor, tensor
@@ -1213,15 +1214,15 @@ def test_sparse_resize(dtype, device):
 
 
 def test_tolist():
-    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
-    with pytest.raises(RuntimeError, match="supported for tensor subclasses"):
-        adj.tolist()
+    data = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    adj = EdgeIndex(data)
+    assert adj.tolist() == data.tolist()
 
 
 def test_numpy():
-    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
-    with pytest.raises(RuntimeError, match="supported for tensor subclasses"):
-        adj.numpy()
+    data = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    adj = EdgeIndex(data)
+    assert np.array_equal(adj.numpy(), data.numpy())
 
 
 @withCUDA
diff --git a/test/test_index.py b/test/test_index.py
index af921e343..85e0af7a9 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -1,6 +1,7 @@
 import os.path as osp
 from typing import List
 
+import numpy as np
 import pytest
 import torch
 from torch import tensor
@@ -510,15 +511,15 @@ def test_sub(dtype, device):
 
 
 def test_to_list():
-    index = Index([0, 1, 1, 2])
-    with pytest.raises(RuntimeError, match="supported for tensor subclasses"):
-        index.tolist()
+    data = torch.tensor([0, 1, 1, 2])
+    index = Index(data)
+    assert index.tolist() == data.tolist()
 
 
 def test_numpy():
-    index = Index([0, 1, 1, 2])
-    with pytest.raises(RuntimeError, match="supported for tensor subclasses"):
-        index.numpy()
+    data = torch.tensor([0, 1, 1, 2])
+    index = Index(data)
+    assert np.array_equal(index.numpy(), data.numpy())
 
 
 @withCUDA