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

pyg_nightly-2.7.0.dev20250218/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2023 PyG Team <team@pyg.org>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pyg-nightly
-Version: 2.7.0.dev20250217
+Version: 2.7.0.dev20250218
 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.dev20250217 → pyg_nightly-2.7.0.dev20250218}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250217"
+version="2.7.0.dev20250218"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]
@@ -18,6 +18,7 @@ keywords=[
     "graph-neural-networks",
     "graph-convolutional-networks",
 ]
+license = {file = "LICENSE"}
 classifiers=[
     "Development Status :: 5 - Production/Stable",
     "License :: OSI Approved :: MIT License",

{pyg_nightly-2.7.0.dev20250217 → pyg_nightly-2.7.0.dev20250218}/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.dev20250217'
+__version__ = '2.7.0.dev20250218'
 
 __all__ = [
     'Index',

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

@@ -1,3 +1,5 @@
+import functools
+import warnings
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
 
 import torch
@@ -8,9 +10,21 @@ from torch import Tensor
 import torch_geometric.typing
 from torch_geometric.typing import CPUHashMap, CUDAHashMap
 
+aten = torch.ops.aten
+
 HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
 
 
+def implements(torch_function: Callable) -> Callable:
+    r"""Registers a :pytorch:`PyTorch` function override."""
+    @functools.wraps(torch_function)
+    def decorator(my_function: Callable) -> Callable:
+        HANDLED_FUNCTIONS[torch_function] = my_function
+        return my_function
+
+    return decorator
+
+
 def as_key_tensor(
     key: Any,
     *,
@@ -20,15 +34,16 @@ def as_key_tensor(
         key = torch.as_tensor(key, device=device)
     except Exception:
         device = device or torch.get_default_device()
+        # TODO Convert int64 to int32.
         # 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 torch_geometric.typing.WITH_CUDA_HASH_MAP and device.type == 'cuda':
+            pass
+            # key = torch.tensor(
+            #     [xxhash.xxh32(x).intdigest() & 0x7FFFFFFF for x in key],
+            #     dtype=torch.int32, device=device)
+        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)
@@ -45,6 +60,28 @@ def as_key_tensor(
     return key
 
 
+def get_hash_map(key: Tensor) -> Union[CPUHashMap, CUDAHashMap]:
+    if torch_geometric.typing.WITH_CUDA_HASH_MAP and key.is_cuda:
+        # TODO Convert int64 to int32.
+        return CUDAHashMap(key, 0.5)
+
+    if key.is_cuda:
+        warnings.warn("Fallback to CPU-based mapping algorithm which may "
+                      "cause slowdowns and device synchronization. Please "
+                      "install 'pyg-lib' for an accelerated 'HashTensor' "
+                      "implementation.")
+
+    if torch_geometric.typing.WITH_CPU_HASH_MAP:
+        return CPUHashMap(key.cpu(), -1)
+
+    import pandas as pd
+
+    return pd.CategoricalDtype(
+        categories=key.cpu().numpy(),
+        ordered=True,
+    )
+
+
 class HashTensor(Tensor):
     _map: Union[Tensor, CPUHashMap, CUDAHashMap]
     _value: Optional[Tensor]
@@ -66,7 +103,6 @@ class HashTensor(Tensor):
             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 "
@@ -88,6 +124,49 @@ class HashTensor(Tensor):
                                  f"first dimension (got {key.size(0)} and "
                                  f"{value.size(0)})")
 
+        min_key = key.min() if key.numel() > 0 else key.new_zeros(())
+        max_key = key.max() if key.numel() > 0 else key.new_zeros(())
+
+        _range = max_key - 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()):
+            _map = torch.full(
+                size=(_range + 2, ),
+                fill_value=-1,
+                dtype=torch.int64,
+                device=device,
+            )
+            _map[(key - (min_key - 1)).long()] = torch.arange(
+                key.numel(),
+                dtype=_map.dtype,
+                device=_map.device,
+            )
+        else:
+            _map = get_hash_map(key)
+
+        return cls._from_data(
+            _map,
+            value,
+            min_key,
+            max_key,
+            num_keys=key.numel(),
+            dtype=dtype,
+        )
+
+    @classmethod
+    def _from_data(
+        cls,
+        _map: Union[Tensor, CPUHashMap, CUDAHashMap],
+        value: Optional[Tensor],
+        min_key: Tensor,
+        max_key: Tensor,
+        *,
+        num_keys: int,
+        dtype: Optional[torch.dtype],
+    ) -> 'HashTensor':
+
+        if value is not None:
             dtype = value.dtype
             size = value.size()
             stride = value.stride()
@@ -95,7 +174,7 @@ class HashTensor(Tensor):
             requires_grad = value.requires_grad
         else:
             dtype = dtype or torch.int64
-            size = torch.Size([key.numel()])
+            size = torch.Size([num_keys])
             stride = (1, )
             layout = torch.strided
             requires_grad = False
@@ -105,40 +184,16 @@ class HashTensor(Tensor):
             size=size,
             strides=stride,
             dtype=dtype,
-            device=device,
+            device=min_key.device,
             layout=layout,
             requires_grad=requires_grad,
         )
+        assert isinstance(out, HashTensor)
 
+        out._map = _map
         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
+        out._min_key = min_key
+        out._max_key = max_key
 
         return out
 
@@ -150,6 +205,9 @@ class HashTensor(Tensor):
             return self._value
         return torch.arange(self.size(0), dtype=self.dtype, device=self.device)
 
+    # Prevent auto-wrapping outputs back into the proper subclass type:
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
     @classmethod
     def __torch_dispatch__(
         cls: Type,
@@ -169,3 +227,51 @@ class HashTensor(Tensor):
             kwargs = pytree.tree_map_only(HashTensor, lambda x: x.as_tensor(),
                                           kwargs)
         return func(*args, **(kwargs or {}))
+
+
+@implements(aten._to_copy.default)
+def _to_copy(
+    tensor: HashTensor,
+    *,
+    dtype: Optional[torch.dtype] = None,
+    layout: Optional[torch.layout] = None,
+    device: Optional[torch.device] = None,
+    pin_memory: bool = False,
+    non_blocking: bool = False,
+    memory_format: Optional[torch.memory_format] = None,
+) -> HashTensor:
+
+    value = tensor._value
+    if value is not None:
+        value = aten._to_copy.default(
+            value,
+            dtype=dtype,
+            layout=layout,
+            device=device,
+            pin_memory=pin_memory,
+            non_blocking=non_blocking,
+            memory_format=memory_format,
+        )
+
+    min_key = aten._to_copy.default(tensor._min_key, device=device)
+    max_key = aten._to_copy.default(tensor._max_key, device=device)
+
+    _map = tensor._map
+    if isinstance(_map, Tensor):
+        _map = aten._to_copy.default(_map, device=device)
+    # Only convert `_map` in case `CUDAHashMap` exists - otherwise we use
+    # CPU-based mapping anyway and there is no need for a copy.
+    elif (torch_geometric.typing.WITH_CUDA_HASH_MAP and min_key.is_cuda
+          and tensor._min_key.device != min_key.device):
+        key = _map.keys()
+        key = aten._to_copy.default(key, device=device)
+        _map = get_hash_map(key)
+
+    return tensor.__class__._from_data(
+        _map,
+        value,
+        min_key,
+        max_key,
+        num_keys=tensor.size(0),
+        dtype=dtype or tensor.dtype,
+    )