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

{pyg_nightly-2.7.0.dev20250216.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.dev20250216
+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>

{pyg_nightly-2.7.0.dev20250216.dist-info → pyg_nightly-2.7.0.dev20250217.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=QlYg3d7vRWC0k_OzPCmAanyTUBiTvikPuzM9feT0REk,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=VBuz9n16ouSk2u4DifpCIW5MXiuoH5UGO4rPj-Celjw,4418
+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
@@ -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.dev20250216.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
-pyg_nightly-2.7.0.dev20250216.dist-info/METADATA,sha256=nrJ478w_yYVSpIyqlfR5lxdYTu6NOIzCv_atkMY5w5w,62999
-pyg_nightly-2.7.0.dev20250216.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.dev20250216'
+__version__ = '2.7.0.dev20250217'
 
 __all__ = [
     'Index',

torch_geometric/hash_tensor.py

@@ -5,6 +5,7 @@ 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] = {}
@@ -18,9 +19,16 @@ def as_key_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)
+        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)
@@ -40,6 +48,8 @@ def as_key_tensor(
 class HashTensor(Tensor):
     _map: Union[Tensor, CPUHashMap, CUDAHashMap]
     _value: Optional[Tensor]
+    _min_key: Tensor
+    _max_key: Tensor
 
     @staticmethod
     def __new__(
@@ -102,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: