pyg-nightly 2.7.0.dev20250215__py3-none-any.whl → 2.7.0.dev20250217__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250215.dist-info → pyg_nightly-2.7.0.dev20250217.dist-info}/METADATA

@@ -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.dist-info → pyg_nightly-2.7.0.dev20250217.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=WUHw4l8NRfqG5Jr-cDi2l7tNpLMNU6SgjrnbmA6eFhY,1978
+torch_geometric/__init__.py,sha256=LZOz8NuqxM3Gqi_c9iZqnbd58Dk5w1hMcWGTBRlN8jM,1978
 torch_geometric/_compile.py,sha256=f-WQeH4VLi5Hn9lrgztFUCSrN_FImjhQa6BxFzcYC38,1338
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -9,7 +9,7 @@ torch_geometric/deprecation.py,sha256=dWRymDIUkUVI2MeEmBG5WF4R6jObZeseSBV9G6FNfj
 torch_geometric/device.py,sha256=tU5-_lBNVbVHl_kUmWPwiG5mQ1pyapwMF4JkmtNN3MM,1224
 torch_geometric/edge_index.py,sha256=BsLh5tOZRjjSYDkjqOFAdBuvMaDg7EWaaLELYsUL0Z8,70048
 torch_geometric/experimental.py,sha256=JbtNNEXjFGI8hZ9raM6-qrZURP6Z5nlDK8QicZUIbz0,4756
-torch_geometric/hash_tensor.py,sha256=MRmjpF7dnH_IfC3HdT0LnjBa1-Q6zpPehPSCLxy2lLY,3912
+torch_geometric/hash_tensor.py,sha256=N4_LPbg9F35g_e9Wy2WgkLMA9kJPnKiYEgTn_4YrypU,5976
 torch_geometric/home.py,sha256=EV54B4Dmiv61GDbkCwtCfWGWJ4eFGwZ8s3KOgGjwYgY,790
 torch_geometric/index.py,sha256=9ChzWFCwj2slNcVBOgfV-wQn-KscJe_y7502w-Vf76w,24045
 torch_geometric/inspector.py,sha256=nKi5o4Mn6xsG0Ex1GudTEQt_EqnF9mcMqGtp7Shh9sQ,19336
@@ -19,7 +19,7 @@ torch_geometric/logging.py,sha256=HmHHLiCcM64k-6UYNOSfXPIeSGNAyiGGcn8cD8tlyuQ,85
 torch_geometric/resolver.py,sha256=fn-_6mCpI2xv7eDZnIFcYrHOn0IrwbkWFLDb9laQrWI,1270
 torch_geometric/seed.py,sha256=MJLbVwpb9i8mK3oi32sS__Cq-dRq_afTeoOL_HoA9ko,372
 torch_geometric/template.py,sha256=rqjDWgcSAgTCiV4bkOjWRPaO4PpUdC_RXigzxxBqAu8,1060
-torch_geometric/typing.py,sha256=-YM1o6csJsWsHpEzmnjIh82_3yq6m-1nwC2TjUMKRGk,15520
+torch_geometric/typing.py,sha256=mtSM6QhCsrohstnyvqMuxEIajCkhcvkQKOU4uVu-nDs,15596
 torch_geometric/warnings.py,sha256=t114CbkrmiqkXaavx5g7OO52dLdktf-U__B5QqYIQvI,413
 torch_geometric/contrib/__init__.py,sha256=0pWkmXfZtbdr-AKwlii5LTFggTEH-MCrSKpZxrtPlVs,352
 torch_geometric/contrib/datasets/__init__.py,sha256=lrGnWsEiJf5zsBRmshGZZFN_uYR2ezDjbj9n9nCpvtk,23
@@ -633,6 +633,6 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
 torch_geometric/visualization/graph.py,sha256=ZuLPL92yGRi7lxlqsUPwL_EVVXF7P2kMcveTtW79vpA,4784
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250215.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
-pyg_nightly-2.7.0.dev20250215.dist-info/METADATA,sha256=j3pySMauadU_aRqO78jwhWjPQhV5lapZkr3uvX3fffk,62977
-pyg_nightly-2.7.0.dev20250215.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250217.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
+pyg_nightly-2.7.0.dev20250217.dist-info/METADATA,sha256=ACH7q4dHQ_l_FbmuJwFJ_K6syGHcHVDwHD3WlqPs0lw,62999
+pyg_nightly-2.7.0.dev20250217.dist-info/RECORD,,

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',

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:

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'. "