pyg-nightly 2.7.0.dev20250221__py3-none-any.whl → 2.7.0.dev20250223__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250221.dist-info → pyg_nightly-2.7.0.dev20250223.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=f29gN2VOOKFStDrxccL16W-vkm6S_5-CHkutK9E5NX4,1978
+torch_geometric/__init__.py,sha256=u79QBiX3vYzq10QGHDb7To3rzMR1UXXopOn8ptXDZ0A,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=aM5chfNpYWOU01nEkE2hBm_ZNlOwCjrZcYXjNtf6PMc,11836
+torch_geometric/hash_tensor.py,sha256=9Zg1KCebfN-xJE1dX2nGGYnK09snSyJkjaYVzCUOfkM,17278
 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
@@ -290,7 +290,7 @@ torch_geometric/loader/temporal_dataloader.py,sha256=AQ2QFeiXKbPp6I8sUeE8H7br-1_
 torch_geometric/loader/utils.py,sha256=f27mczQ7fEP2HpTsJGJxKS0slPu0j8zTba3jP8ViNck,14901
 torch_geometric/loader/zip_loader.py,sha256=3lt10fD15Rxm1WhWzypswGzCEwUz4h8OLCD1nE15yNg,3843
 torch_geometric/metrics/__init__.py,sha256=3krvDobW6vV5yHTjq2S2pmOXxNfysNG26muq7z48e94,699
-torch_geometric/metrics/link_pred.py,sha256=cz9GbvZthV2PAnVnxiZlksGr0VmTQOJGNuZ-OYYg04U,29667
+torch_geometric/metrics/link_pred.py,sha256=wGQG-Fl6BQYJMLZe_L_iIl4ixj6TWgLkkuHyMMraWBA,30480
 torch_geometric/nn/__init__.py,sha256=kQHHHUxFDht2ztD-XFQuv98TvC8MdodaFsIjAvltJBw,874
 torch_geometric/nn/data_parallel.py,sha256=lDAxRi83UNuzAQSj3eu9K2sQheOIU6wqR5elS6oDs90,4764
 torch_geometric/nn/encoding.py,sha256=QNjwWczYExZ1wRGBmpuqYbn6tB7NC4BU-DEgzjhcZqw,3115
@@ -513,10 +513,10 @@ torch_geometric/sampler/base.py,sha256=kT6hYM6losYta3pqLQlqiqboJiujLy6RlH8qM--U_
 torch_geometric/sampler/hgt_sampler.py,sha256=UAm8_wwzEcziKDJ8-TnfZh1705dXRsy_I5PKhZSDTK8,2721
 torch_geometric/sampler/neighbor_sampler.py,sha256=MAVphWqNf0-cwlHRvdiU8de86dBxwjm3Miam_6s1ep4,33971
 torch_geometric/sampler/utils.py,sha256=RJtasO6Q7Pp3oYEOWrbf2DEYuSfuKZOsF2I7-eJDnoA,5485
-torch_geometric/testing/__init__.py,sha256=QUTeYNkmibxFu08AlZGzAnMHfEoBp2kt9o65k0wmfmU,1249
+torch_geometric/testing/__init__.py,sha256=0mAGVWRrTBNsGV2YUkCu_FkyQ8JIcrYVw2LsdKgY9ak,1291
 torch_geometric/testing/asserts.py,sha256=DLC9HnBgFWuTIiQs2OalsQcXGhOVG-e6R99IWhkO32c,4606
 torch_geometric/testing/data.py,sha256=O1qo8FyNxt6RGf63Ys3eXBfa5RvYydeZLk74szrez3c,2604
-torch_geometric/testing/decorators.py,sha256=b0Xqpu-qdiElGo0cFG8cSu-Pqgce7NH8xcoI0NigWiM,8309
+torch_geometric/testing/decorators.py,sha256=j45wlxMB1-Pn3wPKBgDziqg6KkWJUb_fcwfUXzkL2mM,8677
 torch_geometric/testing/distributed.py,sha256=ZZCCXqiQC4-m1ExSjDZhS_a1qPXnHEwhJGTmACxNnVI,2227
 torch_geometric/testing/feature_store.py,sha256=J6JBIt2XK-t8yG8B4JzXp-aJcVl5jaCS1m2H7d6OUxs,2158
 torch_geometric/testing/graph_store.py,sha256=00B7QToCIspYmgN7svQKp1iU-qAzEtrt3VQRFxkHfuk,1044
@@ -633,7 +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.dev20250221.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250221.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
-pyg_nightly-2.7.0.dev20250221.dist-info/METADATA,sha256=LyCZFRktgqS8Zd9gDaXsIFQHAQGWPYgO7x7RDoG3yg8,63021
-pyg_nightly-2.7.0.dev20250221.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250223.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250223.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
+pyg_nightly-2.7.0.dev20250223.dist-info/METADATA,sha256=o3vW1MbKajweST33mDeCk-b1CKb5wGegFomLfUE_rOQ,63021
+pyg_nightly-2.7.0.dev20250223.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.dev20250221'
+__version__ = '2.7.0.dev20250223'
 
 __all__ = [
     'Index',

torch_geometric/hash_tensor.py

@@ -142,12 +142,12 @@ class HashTensor(Tensor):
         if (key.dtype in {torch.uint8, torch.int16} or _range <= 1_000_000
                 or _range <= 2 * key.numel()):
             _map = torch.full(
-                size=(_range + 2, ),
+                size=(_range + 3, ),
                 fill_value=-1,
                 dtype=torch.int64,
-                device=device,
+                device=key.device,
             )
-            _map[(key - (min_key - 1)).long()] = torch.arange(
+            _map[key.long() - (min_key.long() - 1)] = torch.arange(
                 key.numel(),
                 dtype=_map.dtype,
                 device=_map.device,
@@ -164,6 +164,8 @@ class HashTensor(Tensor):
             dtype=dtype,
         )
 
+    # Private Methods #########################################################
+
     @classmethod
     def _from_data(
         cls,
@@ -207,6 +209,20 @@ class HashTensor(Tensor):
 
         return out
 
+    @property
+    def _key(self) -> Tensor:
+        if isinstance(self._map, Tensor):
+            mask = self._map >= 0
+            key = mask.nonzero().view(-1) - 1
+            key = key[self._map[mask]]
+        elif (torch_geometric.typing.WITH_CUDA_HASH_MAP
+              or torch_geometric.typing.WITH_CPU_HASH_MAP):
+            key = self._map.keys().to(self.device)
+        else:
+            key = torch.from_numpy(self._map.categories.to_numpy())
+
+        return key.to(self.device)
+
     def _shallow_copy(self) -> 'HashTensor':
         return self._from_data(
             self._map,
@@ -217,6 +233,37 @@ class HashTensor(Tensor):
             dtype=self.dtype,
         )
 
+    def _get(self, query: Tensor) -> Tensor:
+        if isinstance(self._map, Tensor):
+            index = query.long() - (self._min_key.long() - 1)
+            index = self._map[index.clamp_(min=0, max=self._map.numel() - 1)]
+        elif torch_geometric.typing.WITH_CUDA_HASH_MAP and query.is_cuda:
+            index = self._map.get(query)
+        elif torch_geometric.typing.WITH_CPU_HASH_MAP:
+            index = self._map.get(query.cpu())
+        else:
+            import pandas as pd
+
+            ser = pd.Series(query.cpu().numpy(), dtype=self._map)
+            index = torch.from_numpy(ser.cat.codes.to_numpy()).to(torch.long)
+
+        index = index.to(self.device)
+
+        if self._value is None:
+            return index.to(self.dtype)
+
+        out = self._value[index]
+
+        mask = index != -1
+        mask = mask.view([-1] + [1] * (out.dim() - 1))
+        fill_value = float('NaN') if out.is_floating_point() else -1
+        if torch_geometric.typing.WITH_PT20:
+            other: Union[int, float, Tensor] = fill_value
+        else:
+            other = torch.full_like(out, fill_value)
+
+        return out.where(mask, other)
+
     # Methods #################################################################
 
     def as_tensor(self) -> Tensor:
@@ -252,6 +299,23 @@ class HashTensor(Tensor):
                                           kwargs)
         return func(*args, **(kwargs or {}))
 
+    def tolist(self) -> List[Any]:
+        return self.as_tensor().tolist()
+
+    def index_select(  # type: ignore
+        self,
+        dim: int,
+        index: Any,
+    ) -> Union['HashTensor', Tensor]:
+        return torch.index_select(self, dim, index)
+
+    def select(  # type: ignore
+        self,
+        dim: int,
+        index: Any,
+    ) -> Union['HashTensor', Tensor]:
+        return torch.select(self, dim, index)
+
 
 @implements(aten.alias.default)
 def _alias(tensor: HashTensor) -> HashTensor:
@@ -327,9 +391,13 @@ def _squeeze_default(tensor: HashTensor) -> HashTensor:
     if tensor._value is None:
         return tensor._shallow_copy()
 
+    value = tensor.as_tensor()
+    for d in range(tensor.dim() - 1, 0, -1):
+        value = value.squeeze(d)
+
     return tensor._from_data(
         tensor._map,
-        aten.squeeze.dims(tensor._value, list(range(1, tensor.dim()))),
+        value,
         tensor._min_key,
         tensor._max_key,
         num_keys=tensor.size(0),
@@ -355,11 +423,14 @@ def _squeeze_dim(
     if tensor._value is None:
         return tensor._shallow_copy()
 
-    dim = [d for d in dim if d != 0 and d != -tensor.dim()]
+    value = tensor.as_tensor()
+    for d in dim[::-1]:
+        if d != 0 and d != -tensor.dim():
+            value = value.squeeze(d)
 
     return tensor._from_data(
         tensor._map,
-        aten.squeeze.dims(tensor._value, dim),
+        value,
         tensor._min_key,
         tensor._max_key,
         num_keys=tensor.size(0),
@@ -374,10 +445,18 @@ def _slice(
     start: Optional[int] = None,
     end: Optional[int] = None,
     step: int = 1,
-) -> Union[HashTensor, Tensor]:
+) -> HashTensor:
 
     if dim == 0 or dim == -tensor.dim():
-        return aten.slice.Tensor(tensor.as_tensor(), dim, start, end, step)
+        copy = start is None or (start == 0 or start <= -tensor.size(0))
+        copy &= end is None or end > tensor.size(0)
+        copy &= step == 1
+        if copy:
+            return tensor._shallow_copy()
+
+        key = aten.slice.Tensor(tensor._key, 0, start, end, step)
+        value = aten.slice.Tensor(tensor.as_tensor(), 0, start, end, step)
+        return tensor.__class__(key, value)
 
     return tensor._from_data(
         tensor._map,
@@ -387,3 +466,102 @@ def _slice(
         num_keys=tensor.size(0),
         dtype=tensor.dtype,
     )
+
+
+# Since PyTorch does only allow PyTorch tensors as indices in `index_select`,
+# we need to create a wrapper function and monkey patch `index_select` :(
+_old_index_select = torch.index_select
+
+
+def _new_index_select(
+    input: Tensor,
+    dim: int,
+    index: Tensor,
+    *,
+    out: Optional[Tensor] = None,
+) -> Tensor:
+
+    if dim < -input.dim() or dim >= input.dim():
+        raise IndexError(f"Dimension out of range (expected to be in range of "
+                         f"[{-input.dim()}, {input.dim()-1}], but got {dim})")
+
+    # We convert any index tensor in the first dimension into a tensor. This
+    # means that downstream handling (i.e. in `aten.index_select.default`)
+    # needs to take this pre-conversion into account.
+    if isinstance(input, HashTensor) and (dim == 0 or dim == -input.dim()):
+        index = as_key_tensor(index, device=input.device)
+    return _old_index_select(input, dim, index, out=out)
+
+
+torch.index_select = _new_index_select  # type: ignore
+
+
+@implements(aten.index_select.default)
+def _index_select(
+    tensor: HashTensor,
+    dim: int,
+    index: Tensor,
+) -> Union[HashTensor, Tensor]:
+
+    if dim == 0 or dim == -tensor.dim():
+        return tensor._get(index)
+
+    return tensor._from_data(
+        tensor._map,
+        aten.index_select.default(tensor.as_tensor(), dim, index),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )
+
+
+# Since PyTorch does only allow PyTorch tensors as indices in `select`, we need
+# to create a wrapper function and monkey patch `select` :(
+_old_select = torch.select
+
+
+def _new_select(
+    input: Tensor,
+    dim: int,
+    index: int,
+) -> Tensor:
+
+    if dim < -input.dim() or dim >= input.dim():
+        raise IndexError(f"Dimension out of range (expected to be in range of "
+                         f"[{-input.dim()}, {input.dim()-1}], but got {dim})")
+
+    # We convert any index in the first dimension into an integer. This means
+    # that downstream handling (i.e. in `aten.select.int`) needs to take this
+    # pre-conversion into account.
+    if isinstance(input, HashTensor) and (dim == 0 or dim == -input.dim()):
+        index = int(as_key_tensor([index]))
+    return _old_select(input, dim, index)
+
+
+torch.select = _new_select  # type: ignore
+
+
+@implements(aten.select.int)
+def _select(
+    tensor: HashTensor,
+    dim: int,
+    index: int,
+) -> Union[HashTensor, Tensor]:
+
+    if dim == 0 or dim == -tensor.dim():
+        key = torch.tensor(
+            [index],
+            dtype=tensor._min_key.dtype,
+            device=tensor._min_key.device,
+        )
+        return tensor._get(key).squeeze(0)
+
+    return tensor._from_data(
+        tensor._map,
+        aten.select.int(tensor.as_tensor(), dim, index),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )

torch_geometric/metrics/link_pred.py

@@ -715,22 +715,32 @@ class LinkPredPersonalization(_LinkPredMetric):
 
     Args:
         k (int): The number of top-:math:`k` predictions to evaluate against.
+        max_src_nodes (int, optional): The maximum source nodes to consider to
+            compute pair-wise dissimilarity. If specified,
+            Personalization @ :math:`k` is approximated to avoid computation
+            blowup due to quadratic complexity. (default: :obj:`2**12`)
         batch_size (int, optional): The batch size to determine how many pairs
             of user recommendations should be processed at once.
             (default: :obj:`2**16`)
     """
     higher_is_better: bool = True
 
-    def __init__(self, k: int, batch_size: int = 2**16) -> None:
+    def __init__(
+        self,
+        k: int,
+        max_src_nodes: Optional[int] = 2**12,
+        batch_size: int = 2**16,
+    ) -> None:
         super().__init__(k)
+        self.max_src_nodes = max_src_nodes
         self.batch_size = batch_size
 
         if WITH_TORCHMETRICS:
             self.add_state('preds', default=[], dist_reduce_fx='cat')
-            self.add_state('dev_tensor', torch.empty(0), dist_reduce_fx='sum')
+            self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
         else:
             self.preds: List[Tensor] = []
-            self.register_buffer('dev_tensor', torch.empty(0))
+            self.register_buffer('total', torch.tensor(0))
 
     def update(
         self,
@@ -738,11 +748,21 @@ class LinkPredPersonalization(_LinkPredMetric):
         edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
         edge_label_weight: Optional[Tensor] = None,
     ) -> None:
+
         # NOTE Move to CPU to avoid memory blowup.
-        self.preds.append(pred_index_mat[:, :self.k].cpu())
+        pred_index_mat = pred_index_mat[:, :self.k].cpu()
+
+        if self.max_src_nodes is None:
+            self.preds.append(pred_index_mat)
+            self.total += pred_index_mat.size(0)
+        elif self.total < self.max_src_nodes:
+            remaining = int(self.max_src_nodes - self.total)
+            pred_index_mat = pred_index_mat[:remaining]
+            self.preds.append(pred_index_mat)
+            self.total += pred_index_mat.size(0)
 
     def compute(self) -> Tensor:
-        device = self.dev_tensor.device
+        device = self.total.device
         score = torch.tensor(0.0, device=device)
         total = torch.tensor(0, device=device)
 
@@ -786,6 +806,7 @@ class LinkPredPersonalization(_LinkPredMetric):
 
     def _reset(self) -> None:
         self.preds = []
+        self.total.zero_()
 
 
 class LinkPredAveragePopularity(_LinkPredMetric):

torch_geometric/testing/__init__.py

@@ -22,6 +22,7 @@ from .decorators import (
     withDevice,
     withCUDA,
     withMETIS,
+    withHashTensor,
     disableExtensions,
     withoutExtensions,
 )
@@ -53,6 +54,7 @@ __all__ = [
     'withDevice',
     'withCUDA',
     'withMETIS',
+    'withHashTensor',
     'disableExtensions',
     'withoutExtensions',
     'assert_module',

torch_geometric/testing/decorators.py

@@ -10,6 +10,7 @@ from packaging.requirements import Requirement
 from packaging.version import Version
 
 import torch_geometric
+import torch_geometric.typing
 from torch_geometric.typing import WITH_METIS, WITH_PYG_LIB, WITH_TORCH_SPARSE
 from torch_geometric.visualization.graph import has_graphviz
 
@@ -265,6 +266,17 @@ def withMETIS(func: Callable) -> Callable:
     )(func)
 
 
+def withHashTensor(func: Callable) -> Callable:
+    r"""A decorator to only test in case :class:`HashTensor` is available."""
+    import pytest
+
+    return pytest.mark.skipif(
+        not torch_geometric.typing.WITH_CPU_HASH_MAP
+        and not has_package('pandas'),
+        reason="HashTensor dependencies not available",
+    )(func)
+
+
 def disableExtensions(func: Callable) -> Callable:
     r"""A decorator to temporarily disable the usage of the
     :obj:`torch_scatter`, :obj:`torch_sparse` and :obj:`pyg_lib` extension