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

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

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=Ht6O12SaZRLjd7aAphjjLY_T59xQTnRQ9nN1oSA1Xmk,1978
+torch_geometric/__init__.py,sha256=f29gN2VOOKFStDrxccL16W-vkm6S_5-CHkutK9E5NX4,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=T4fRmS6TD-j7PKa9LMralSbI6naqA4ctW8AnjWsREzw,9615
+torch_geometric/hash_tensor.py,sha256=aM5chfNpYWOU01nEkE2hBm_ZNlOwCjrZcYXjNtf6PMc,11836
 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
@@ -463,7 +463,7 @@ torch_geometric/nn/models/signed_gcn.py,sha256=J40CnedFIqtKI1LhW1ITSEFRbA_XiJZL6
 torch_geometric/nn/models/tgn.py,sha256=kEGdfLJybkbMT4UMoAh2nCzfX3_nDjfm1cicuPHEwAM,11878
 torch_geometric/nn/models/visnet.py,sha256=97OFMCsPDEI5BCSi7RhoRcU2CNRp7zck2tEzrltFZj4,43192
 torch_geometric/nn/nlp/__init__.py,sha256=q6CPUiJHcc9bXw90lyj-ID4F3kfW8uPM-SOxW9uCMHs,213
-torch_geometric/nn/nlp/llm.py,sha256=j03tyCO1ADgrzGhLqYOUcsy0haGbV4dmT9bdwEnESPE,12181
+torch_geometric/nn/nlp/llm.py,sha256=uVPoIB2_nU0s3N_Z3oVW2Q7XGsAisW1QGaaLe2Y__hw,12227
 torch_geometric/nn/nlp/sentence_transformer.py,sha256=q5M7SGtrUzoSiNhKCGFb7JatWiukdhNF6zdq2yiqxwE,4475
 torch_geometric/nn/nlp/vision_transformer.py,sha256=diVBefjIynzYs8WBlcpTeSVnw1PUecHY--B9Yd-W2hA,863
 torch_geometric/nn/norm/__init__.py,sha256=u2qIDrkbeuObGVXSAIftAlvSd6ouGTtxznCfD-59UiA,669
@@ -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.dev20250220.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250220.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
-pyg_nightly-2.7.0.dev20250220.dist-info/METADATA,sha256=SsXACQ7At6aou2g3LkqAlRX2zcjw3aTPw-mMliTKU1E,63021
-pyg_nightly-2.7.0.dev20250220.dist-info/RECORD,,
+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,,

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.dev20250220'
+__version__ = '2.7.0.dev20250221'
 
 __all__ = [
     'Index',

torch_geometric/hash_tensor.py

@@ -1,6 +1,16 @@
 import functools
 import warnings
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 import torch
 import torch.utils._pytree as pytree
@@ -197,6 +207,16 @@ class HashTensor(Tensor):
 
         return out
 
+    def _shallow_copy(self) -> 'HashTensor':
+        return self._from_data(
+            self._map,
+            self._value,
+            self._min_key,
+            self._max_key,
+            num_keys=self.size(0),
+            dtype=self.dtype,
+        )
+
     # Methods #################################################################
 
     def as_tensor(self) -> Tensor:
@@ -233,6 +253,11 @@ class HashTensor(Tensor):
         return func(*args, **(kwargs or {}))
 
 
+@implements(aten.alias.default)
+def _alias(tensor: HashTensor) -> HashTensor:
+    return tensor._shallow_copy()
+
+
 @implements(aten._to_copy.default)
 def _to_copy(
     tensor: HashTensor,
@@ -295,3 +320,70 @@ def _unsqueeze(tensor: HashTensor, dim: int) -> HashTensor:
         num_keys=tensor.size(0),
         dtype=tensor.dtype,
     )
+
+
+@implements(aten.squeeze.default)
+def _squeeze_default(tensor: HashTensor) -> HashTensor:
+    if tensor._value is None:
+        return tensor._shallow_copy()
+
+    return tensor._from_data(
+        tensor._map,
+        aten.squeeze.dims(tensor._value, list(range(1, tensor.dim()))),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )
+
+
+@implements(aten.squeeze.dim)
+@implements(getattr(aten.squeeze, 'dims', aten.squeeze.dim))
+def _squeeze_dim(
+    tensor: HashTensor,
+    dim: Union[int, List[int]],
+) -> HashTensor:
+    if isinstance(dim, int):
+        dim = [dim]
+
+    for d in dim:
+        if d < -tensor.dim() or d >= tensor.dim():
+            raise IndexError(f"Dimension out of range (expected to be in "
+                             f"range of [{-tensor.dim()}, {tensor.dim()-1}], "
+                             f"but got {d})")
+
+    if tensor._value is None:
+        return tensor._shallow_copy()
+
+    dim = [d for d in dim if d != 0 and d != -tensor.dim()]
+
+    return tensor._from_data(
+        tensor._map,
+        aten.squeeze.dims(tensor._value, dim),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )
+
+
+@implements(aten.slice.Tensor)
+def _slice(
+    tensor: HashTensor,
+    dim: int,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+    step: int = 1,
+) -> Union[HashTensor, Tensor]:
+
+    if dim == 0 or dim == -tensor.dim():
+        return aten.slice.Tensor(tensor.as_tensor(), dim, start, end, step)
+
+    return tensor._from_data(
+        tensor._map,
+        aten.slice.Tensor(tensor.as_tensor(), dim, start, end, step),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )

torch_geometric/nn/nlp/llm.py

@@ -49,15 +49,16 @@ def get_llm_kwargs(required_memory: int, dtype=torch.dtype) -> Dict[str, Any]:
 class LLM(torch.nn.Module):
     r"""A wrapper around a Large Language Model (LLM) from HuggingFace.
 
-    model_name (str): The HuggingFace model name, *e.g.*, :obj:`"llama2"` or
-        :obj:`"gemma"`.
-    num_params (int, optional): An integer representing how many parameters the
-        HuggingFace model has, in billions. This is used to automatically
-        allocate the correct number of GPUs needed, given the available GPU
-        memory of your GPUs. If not specified, the number of parameters
-        is determined using the `huggingface_hub` module.
-    dtype (torch.dtype, optional): The data type to use for the LLM.
-        (default :obj: `torch.bfloat16`)
+    Args:
+        model_name (str): The HuggingFace model name, *e.g.*, :obj:`"llama2"`
+            or :obj:`"gemma"`.
+        num_params (int, optional): An integer representing how many parameters
+            the HuggingFace model has, in billions. This is used to
+            automatically allocate the correct number of GPUs needed, given the
+            available GPU memory of your GPUs. If not specified, the number of
+            parameters is determined using the `huggingface_hub` module.
+        dtype (torch.dtype, optional): The data type to use for the LLM.
+            (default :obj: `torch.bfloat16`)
     """
     def __init__(
         self,