Support variety of keys in HashTensor

pyproject.toml

@@ -37,6 +37,7 @@ dependencies=[
     "pyparsing",
     "requests",
     "tqdm",
+    "xxhash",
 ]
 
 [project.optional-dependencies]

test/test_hash_tensor.py

@@ -28,3 +28,8 @@ def test_basic(dtype, device):
     value = torch.randn(key.size(0), 2, device=device)
 
     HashTensor(key, value)
+
+
+@withCUDA
+def test_string_key(device):
+    HashTensor(['1', '2', '3'], device=device)

torch_geometric/hash_tensor.py

@@ -2,13 +2,41 @@
 
 import torch
 import torch.utils._pytree as pytree
+import xxhash
 from torch import Tensor
 
 from torch_geometric.typing import CPUHashMap, CUDAHashMap
 
 HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
 
 
+def as_key_tensor(
+    key: Any,
+    *,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+    try:
+        key = torch.as_tensor(key, device=device)
+    except Exception:
+        key = torch.tensor([
+            xxhash.xxh64(item).intdigest() & 0x7FFFFFFFFFFFFFFF for item in key
+        ], dtype=torch.int64, device=device)
+
+    if key.element_size() == 1:
+        key = key.view(torch.uint8)
+    elif key.element_size() == 2:
+        key = key.view(torch.int16)
+    elif key.element_size() == 4:
+        key = key.view(torch.int32)
+    elif key.element_size() == 8:
+        key = key.view(torch.int64)
+    else:
+        raise ValueError(f"Received invalid dtype '{key.dtype}' with "
+                         f"{key.element_size()} bytes")
+
+    return key
+
+
 class HashTensor(Tensor):
     _map: Union[Tensor, CPUHashMap, CUDAHashMap]
     _value: Optional[Tensor]
@@ -24,18 +52,11 @@ def __new__(
     ) -> 'HashTensor':
 
         if value is not None:
-            if not isinstance(value, Tensor):
-                value = torch.tensor(value, dtype=dtype, device=device)
-            else:
-                value = value.to(device, dtype)
-
-        # TODO Add numpy support
-        if not isinstance(key, Tensor):
-            device = value.device if value is not None else device
-            key = torch.tensor(key, device=device)
-        else:
-            key = key.to(device)
-            device = key.device
+            value = torch.as_tensor(value, dtype=dtype, device=device)
+            device = value.device
+
+        key = as_key_tensor(key, device=device)
+        device = key.device
 
         if key.dim() != 1:
             raise ValueError(f"'key' data in '{cls.__name__}' needs to be "