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

pyg_nightly-2.7.0.dev20250219.dist-info/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.dist-info → pyg_nightly-2.7.0.dev20250219.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pyg-nightly
-Version: 2.7.0.dev20250217
+Version: 2.7.0.dev20250219
 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.dist-info → pyg_nightly-2.7.0.dev20250219.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=LZOz8NuqxM3Gqi_c9iZqnbd58Dk5w1hMcWGTBRlN8jM,1978
+torch_geometric/__init__.py,sha256=P-W8gnhQltNcKfeLyPOC4diB9zVIHQ5yi0IXFmjYzd4,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=N4_LPbg9F35g_e9Wy2WgkLMA9kJPnKiYEgTn_4YrypU,5976
+torch_geometric/hash_tensor.py,sha256=T4fRmS6TD-j7PKa9LMralSbI6naqA4ctW8AnjWsREzw,9615
 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,7 @@ 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.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,,
+pyg_nightly-2.7.0.dev20250219.dist-info/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250219.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
+pyg_nightly-2.7.0.dev20250219.dist-info/METADATA,sha256=vvByDyHVHwRDOLKEnsUSA1Dh2obTtvRZvzGks5iU8XA,62999
+pyg_nightly-2.7.0.dev20250219.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.dev20250217'
+__version__ = '2.7.0.dev20250219'
 
 __all__ = [
     'Index',

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,43 +184,21 @@ 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
 
+    # Methods #################################################################
+
     def as_tensor(self) -> Tensor:
         r"""Zero-copies the :class:`HashTensor` representation back to a
         :class:`torch.Tensor` representation.
@@ -150,6 +207,11 @@ class HashTensor(Tensor):
             return self._value
         return torch.arange(self.size(0), dtype=self.dtype, device=self.device)
 
+    # PyTorch/Python builtins #################################################
+
+    # 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 +231,67 @@ 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._from_data(
+        _map,
+        value,
+        min_key,
+        max_key,
+        num_keys=tensor.size(0),
+        dtype=dtype or tensor.dtype,
+    )
+
+
+@implements(aten.unsqueeze.default)
+def _unsqueeze(tensor: HashTensor, dim: int) -> HashTensor:
+    if dim == 0 or dim == -(tensor.dim() + 1):
+        raise IndexError(f"Cannot unsqueeze '{tensor.__class__.__name__}' in "
+                         f"the first dimension")
+
+    return tensor._from_data(
+        tensor._map,
+        aten.unsqueeze.default(tensor.as_tensor(), dim),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )