pyg-nightly 2.7.0.dev20250215__tar.gz → 2.7.0.dev20250217__tar.gz

{pyg_nightly-2.7.0.dev20250215 → pyg_nightly-2.7.0.dev20250217}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pyg-nightly
-Version: 2.7.0.dev20250215
+Version: 2.7.0.dev20250217
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>
@@ -22,6 +22,7 @@ Requires-Dist: psutil>=5.8.0
 Requires-Dist: pyparsing
 Requires-Dist: requests
 Requires-Dist: tqdm
+Requires-Dist: xxhash
 Requires-Dist: matplotlib ; extra == "benchmark"
 Requires-Dist: networkx ; extra == "benchmark"
 Requires-Dist: pandas ; extra == "benchmark"

{pyg_nightly-2.7.0.dev20250215 → pyg_nightly-2.7.0.dev20250217}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250215"
+version="2.7.0.dev20250217"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]
@@ -37,6 +37,7 @@ dependencies=[
     "pyparsing",
     "requests",
     "tqdm",
+    "xxhash",
 ]
 
 [project.optional-dependencies]

{pyg_nightly-2.7.0.dev20250215 → pyg_nightly-2.7.0.dev20250217}/torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250215'
+__version__ = '2.7.0.dev20250217'
 
 __all__ = [
     'Index',

{pyg_nightly-2.7.0.dev20250215 → pyg_nightly-2.7.0.dev20250217}/torch_geometric/hash_tensor.py

@@ -2,16 +2,54 @@ from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
 
 import torch
 import torch.utils._pytree as pytree
+import xxhash
 from torch import Tensor
 
+import torch_geometric.typing
 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:
+        device = device or torch.get_default_device()
+        # On GPU, we default to int32 for faster 'CUDAHashMap' implementation:
+        if device.type == 'cuda':
+            key = torch.tensor(
+                [xxhash.xxh32(x).intdigest() & 0x7FFFFFFF for x in key],
+                dtype=torch.int32, device=device)
+        else:
+            key = torch.tensor([
+                xxhash.xxh64(x).intdigest() & 0x7FFFFFFFFFFFFFFF for x 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]
+    _min_key: Tensor
+    _max_key: Tensor
 
     @staticmethod
     def __new__(
@@ -24,18 +62,11 @@ class HashTensor(Tensor):
     ) -> '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 "
@@ -81,6 +112,34 @@ class HashTensor(Tensor):
 
         out._value = value
 
+        out._min_key = key.min() if key.numel() > 0 else key.new_zeros(())
+        out._max_key = key.max() if key.numel() > 0 else key.new_zeros(())
+
+        _range = out._max_key - out._min_key
+        # TODO Expose fixed threshold as argument.
+        if (key.dtype in {torch.uint8, torch.int16} or _range <= 1_000_000
+                or _range <= 2 * key.numel()):
+            out._map = torch.full(
+                size=(_range + 2, ),
+                fill_value=-1,
+                dtype=torch.int64,
+                device=device,
+            )
+            out._map[(key - (out._min_key - 1)).long()] = torch.arange(
+                key.numel(),
+                dtype=out._map.dtype,
+                device=out._map.device,
+            )
+        elif torch_geometric.typing.WITH_CUDA_HASH_MAP and key.is_cuda:
+            # TODO Convert int64 to int32.
+            out._map = CUDAHashMap(key, 0.5)
+        elif torch_geometric.typing.WITH_CPU_HASH_MAP and key.is_cpu:
+            out._map = CPUHashMap(key, -1)
+        else:
+            # TODO Expose pandas fallback.
+            # warnings.warn()
+            raise NotImplementedError
+
         return out
 
     def as_tensor(self) -> Tensor:

{pyg_nightly-2.7.0.dev20250215 → pyg_nightly-2.7.0.dev20250217}/torch_geometric/typing.py

@@ -69,10 +69,16 @@ try:
         pyg_lib.sampler.neighbor_sample).parameters)
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = ('edge_weight' in inspect.signature(
         pyg_lib.sampler.neighbor_sample).parameters)
-    # WITH_CPU_HASH_MAP = hasattr(torch.classes.pyg, 'CPUHashMap')
-    # WITH_CUDA_HASH_MAP = hasattr(torch.classes.pyg, 'CUDAHashMap')
-    WITH_CPU_HASH_MAP = False
-    WITH_CUDA_HASH_MAP = False
+    try:
+        torch.classes.pyg.CPUHashMap
+        WITH_CPU_HASH_MAP = True
+    except Exception:
+        WITH_CPU_HASH_MAP = False
+    try:
+        torch.classes.pyg.CUDAHashMap
+        WITH_CUDA_HASH_MAP = True
+    except Exception:
+        WITH_CUDA_HASH_MAP = False
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
         warnings.warn(f"An issue occurred while importing 'pyg-lib'. "