Repository: pyg-team/pytorch_geometric
Pull request #10066 (2025-02-24 13:27:06)

Fix `torch.jit.script` in combination with `HashTensor` monkey-patch 


Patch:
diff --git a/torch_geometric/hash_tensor.py b/torch_geometric/hash_tensor.py
index db894e99e..0a03b1e5a 100644
--- a/torch_geometric/hash_tensor.py
+++ b/torch_geometric/hash_tensor.py
@@ -12,6 +12,7 @@ from typing import (
     Union,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 import xxhash
@@ -303,6 +304,10 @@ class HashTensor(Tensor):
         """"""  # noqa: D419
         return self.as_tensor().tolist()
 
+    def numpy(self, *, force: bool = False) -> np.ndarray:
+        """"""  # noqa: D419
+        return self.as_tensor().numpy(force=force)
+
     def index_select(  # type: ignore
         self,
         dim: int,
@@ -555,22 +560,32 @@ _old_index_select = torch.index_select
 
 def _new_index_select(
     input: Tensor,
-    dim: int,
+    dim: Union[int, str],
     index: Tensor,
-    *,
     out: Optional[Tensor] = None,
 ) -> Tensor:
 
-    if dim < -input.dim() or dim >= input.dim():
+    if isinstance(dim, int) and (dim < -input.dim() or dim >= input.dim()):
         raise IndexError(f"Dimension out of range (expected to be in range of "
                          f"[{-input.dim()}, {input.dim()-1}], but got {dim})")
 
     # We convert any index tensor in the first dimension into a tensor. This
     # means that downstream handling (i.e. in `aten.index_select.default`)
     # needs to take this pre-conversion into account.
-    if isinstance(input, HashTensor) and (dim == 0 or dim == -input.dim()):
+    if (not torch.jit.is_scripting() and isinstance(input, HashTensor)
+            and isinstance(dim, int) and (dim == 0 or dim == -input.dim())):
         index = as_key_tensor(index, device=input.device)
-    return _old_index_select(input, dim, index, out=out)
+
+    if isinstance(dim, int):  # Type narrowing...
+        if out is None:
+            return _old_index_select(input, dim, index)
+        else:
+            return _old_index_select(input, dim, index, out=out)
+    else:
+        if out is None:
+            return _old_index_select(input, dim, index)
+        else:
+            return _old_index_select(input, dim, index, out=out)
 
 
 torch.index_select = _new_index_select  # type: ignore
@@ -603,20 +618,25 @@ _old_select = torch.select
 
 def _new_select(
     input: Tensor,
-    dim: int,
+    dim: Union[int, str],
     index: int,
 ) -> Tensor:
 
-    if dim < -input.dim() or dim >= input.dim():
+    if isinstance(dim, int) and (dim < -input.dim() or dim >= input.dim()):
         raise IndexError(f"Dimension out of range (expected to be in range of "
                          f"[{-input.dim()}, {input.dim()-1}], but got {dim})")
 
     # We convert any index in the first dimension into an integer. This means
     # that downstream handling (i.e. in `aten.select.int`) needs to take this
     # pre-conversion into account.
-    if isinstance(input, HashTensor) and (dim == 0 or dim == -input.dim()):
+    if (not torch.jit.is_scripting() and isinstance(input, HashTensor)
+            and isinstance(dim, int) and (dim == 0 or dim == -input.dim())):
         index = int(as_key_tensor([index]))
-    return _old_select(input, dim, index)
+
+    if isinstance(dim, int):  # Type narrowing...
+        return _old_select(input, dim, index)
+    else:
+        return _old_select(input, dim, index)
 
 
 torch.select = _new_select  # type: ignore


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index ed06d24d4..eb78b3246 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -1212,7 +1212,7 @@ def test_sparse_resize(dtype, device):
     assert out._T_indptr is None
 
 
-def test_to_list():
+def test_tolist():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
     with pytest.raises(RuntimeError, match="supported for tensor subclasses"):
         adj.tolist()
diff --git a/test/test_hash_tensor.py b/test/test_hash_tensor.py
index fbaaf7cf9..b77662eae 100644
--- a/test/test_hash_tensor.py
+++ b/test/test_hash_tensor.py
@@ -1,6 +1,7 @@
 import os.path as osp
 from typing import List
 
+import numpy as np
 import pytest
 import torch
 
@@ -440,6 +441,18 @@ def test_select(device):
     assert out.as_tensor().equal(value[:, 0])
 
 
+def test_tolist():
+    key = torch.tensor([2, 1, 0])
+    value = torch.randn(key.size(0), 2)
+    assert HashTensor(key, value).tolist() == value.tolist()
+
+
+def test_numpy():
+    key = torch.tensor([2, 1, 0])
+    value = torch.randn(key.size(0), 2)
+    assert np.allclose(HashTensor(key, value).numpy(), value.numpy())
+
+
 def _collate_fn(hash_tensors: List[HashTensor]) -> List[HashTensor]:
     return hash_tensors
 
diff --git a/test/transforms/test_to_superpixels.py b/test/transforms/test_to_superpixels.py
index 282a86c2b..0851452f0 100644
--- a/test/transforms/test_to_superpixels.py
+++ b/test/transforms/test_to_superpixels.py
@@ -15,7 +15,7 @@ resources = [
 
 
 @onlyOnline
-@withPackage('torchvision', 'skimage')
+@withPackage('torch>=2.0.0', 'torchvision', 'skimage')
 def test_to_superpixels(tmp_path):
     import torchvision.transforms as T
     from torchvision.datasets.mnist import (