pyg-nightly 2.7.0.dev20250214__tar.gz → 2.7.0.dev20250216__tar.gz

{pyg_nightly-2.7.0.dev20250214 → pyg_nightly-2.7.0.dev20250216}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pyg-nightly
-Version: 2.7.0.dev20250214
+Version: 2.7.0.dev20250216
 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>
@@ -22,6 +22,7 @@ Requires-Dist: psutil>=5.8.0
 Requires-Dist: pyparsing
 Requires-Dist: requests
 Requires-Dist: tqdm
+Requires-Dist: xxhash
 Requires-Dist: matplotlib ; extra == "benchmark"
 Requires-Dist: networkx ; extra == "benchmark"
 Requires-Dist: pandas ; extra == "benchmark"

{pyg_nightly-2.7.0.dev20250214 → pyg_nightly-2.7.0.dev20250216}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250214"
+version="2.7.0.dev20250216"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]
@@ -37,6 +37,7 @@ dependencies=[
     "pyparsing",
     "requests",
     "tqdm",
+    "xxhash",
 ]
 
 [project.optional-dependencies]

{pyg_nightly-2.7.0.dev20250214 → pyg_nightly-2.7.0.dev20250216}/torch_geometric/__init__.py

@@ -7,6 +7,7 @@ from ._compile import compile, is_compiling
 from ._onnx import is_in_onnx_export
 from .index import Index
 from .edge_index import EdgeIndex
+from .hash_tensor import HashTensor
 from .seed import seed_everything
 from .home import get_home_dir, set_home_dir
 from .device import is_mps_available, is_xpu_available, device
@@ -30,11 +31,12 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250214'
+__version__ = '2.7.0.dev20250216'
 
 __all__ = [
     'Index',
     'EdgeIndex',
+    'HashTensor',
     'seed_everything',
     'get_home_dir',
     'set_home_dir',
@@ -67,4 +69,5 @@ if torch_geometric.typing.WITH_PT24:
         EdgeIndex,
         torch_geometric.edge_index.SortOrder,
         torch_geometric.edge_index.CatMetadata,
+        HashTensor,
     ])

pyg_nightly-2.7.0.dev20250216/torch_geometric/hash_tensor.py

@@ -0,0 +1,133 @@
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
+
+import torch
+import torch.utils._pytree as pytree
+import xxhash
+from torch import Tensor
+
+from torch_geometric.typing import CPUHashMap, CUDAHashMap
+
+HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
+
+
+def as_key_tensor(
+    key: Any,
+    *,
+    device: Optional[torch.device] = None,
+) -> 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)
+
+    if key.element_size() == 1:
+        key = key.view(torch.uint8)
+    elif key.element_size() == 2:
+        key = key.view(torch.int16)
+    elif key.element_size() == 4:
+        key = key.view(torch.int32)
+    elif key.element_size() == 8:
+        key = key.view(torch.int64)
+    else:
+        raise ValueError(f"Received invalid dtype '{key.dtype}' with "
+                         f"{key.element_size()} bytes")
+
+    return key
+
+
+class HashTensor(Tensor):
+    _map: Union[Tensor, CPUHashMap, CUDAHashMap]
+    _value: Optional[Tensor]
+
+    @staticmethod
+    def __new__(
+        cls: Type,
+        key: Any,
+        value: Optional[Any] = None,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ) -> 'HashTensor':
+
+        if value is not None:
+            value = torch.as_tensor(value, dtype=dtype, device=device)
+            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 "
+                             f"one-dimensional (got {key.dim()} dimensions)")
+
+        if not key.is_contiguous():
+            raise ValueError(f"'key' data in '{cls.__name__}' needs to be "
+                             f"contiguous")
+
+        if value is not None:
+            if key.device != value.device:
+                raise ValueError(f"'key' and 'value' data in '{cls.__name__}' "
+                                 f"are expected to be on the same device (got "
+                                 f"'{key.device}' and '{value.device}')")
+
+            if key.numel() != value.size(0):
+                raise ValueError(f"'key' and 'value' data in '{cls.__name__}' "
+                                 f"are expected to have the same size in the "
+                                 f"first dimension (got {key.size(0)} and "
+                                 f"{value.size(0)})")
+
+            dtype = value.dtype
+            size = value.size()
+            stride = value.stride()
+            layout = value.layout
+            requires_grad = value.requires_grad
+        else:
+            dtype = dtype or torch.int64
+            size = torch.Size([key.numel()])
+            stride = (1, )
+            layout = torch.strided
+            requires_grad = False
+
+        out = Tensor._make_wrapper_subclass(  # type: ignore
+            cls,
+            size=size,
+            strides=stride,
+            dtype=dtype,
+            device=device,
+            layout=layout,
+            requires_grad=requires_grad,
+        )
+
+        out._value = value
+
+        return out
+
+    def as_tensor(self) -> Tensor:
+        r"""Zero-copies the :class:`HashTensor` representation back to a
+        :class:`torch.Tensor` representation.
+        """
+        if self._value is not None:
+            return self._value
+        return torch.arange(self.size(0), dtype=self.dtype, device=self.device)
+
+    @classmethod
+    def __torch_dispatch__(
+        cls: Type,
+        func: Callable[..., Any],
+        types: Iterable[Type[Any]],
+        args: Iterable[Tuple[Any, ...]] = (),
+        kwargs: Optional[Dict[Any, Any]] = None,
+    ) -> Any:
+        # Hold a number of `HANDLED_FUNCTIONS` that implement specific
+        # functions for valid `HashTensor` routines.
+        if func in HANDLED_FUNCTIONS:
+            return HANDLED_FUNCTIONS[func](*args, **(kwargs or {}))
+
+        # For all other PyTorch functions, we treat them as vanilla tensors.
+        args = pytree.tree_map_only(HashTensor, lambda x: x.as_tensor(), args)
+        if kwargs is not None:
+            kwargs = pytree.tree_map_only(HashTensor, lambda x: x.as_tensor(),
+                                          kwargs)
+        return func(*args, **(kwargs or {}))

{pyg_nightly-2.7.0.dev20250214 → pyg_nightly-2.7.0.dev20250216}/torch_geometric/nn/attention/__init__.py

@@ -1,7 +1,9 @@
 from .performer import PerformerAttention
 from .qformer import QFormer
+from .sgformer import SGFormerAttention
 
 __all__ = [
     'PerformerAttention',
     'QFormer',
+    'SGFormerAttention',
 ]

pyg_nightly-2.7.0.dev20250216/torch_geometric/nn/attention/sgformer.py

@@ -0,0 +1,78 @@
+import torch
+from torch import Tensor
+
+
+class SGFormerAttention(torch.nn.Module):
+    r"""The simple global attention mechanism from the
+    `"SGFormer: Simplifying and Empowering Transformers for
+    Large-Graph Representations"
+    <https://arxiv.org/abs/2306.10759>`_ paper.
+
+    Args:
+        channels (int): Size of each input sample.
+        heads (int, optional): Number of parallel attention heads.
+            (default: :obj:`1.`)
+        head_channels (int, optional): Size of each attention head.
+            (default: :obj:`64.`)
+        qkv_bias (bool, optional): If specified, add bias to query, key
+            and value in the self attention. (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        channels: int,
+        heads: int = 1,
+        head_channels: int = 64,
+        qkv_bias: bool = False,
+    ) -> None:
+        super().__init__()
+        assert channels % heads == 0
+        assert heads == 1, 'The number of heads are fixed as 1.'
+        if head_channels is None:
+            head_channels = channels // heads
+
+        self.heads = heads
+        self.head_channels = head_channels
+
+        inner_channels = head_channels * heads
+        self.q = torch.nn.Linear(channels, inner_channels, bias=qkv_bias)
+        self.k = torch.nn.Linear(channels, inner_channels, bias=qkv_bias)
+        self.v = torch.nn.Linear(channels, inner_channels, bias=qkv_bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # feature transformation
+        qs = self.q(x).reshape(-1, self.heads, self.head_channels)
+        ks = self.k(x).reshape(-1, self.heads, self.head_channels)
+        vs = self.v(x).reshape(-1, self.heads, self.head_channels)
+
+        # normalize input
+        qs = qs / torch.norm(qs, p=2)  # [N, H, M]
+        ks = ks / torch.norm(ks, p=2)  # [L, H, M]
+        N = qs.shape[0]
+
+        # numerator
+        kvs = torch.einsum("lhm,lhd->hmd", ks, vs)
+        attention_num = torch.einsum("nhm,hmd->nhd", qs, kvs)  # [N, H, D]
+        attention_num += N * vs
+
+        # denominator
+        all_ones = torch.ones([ks.shape[0]]).to(ks.device)
+        ks_sum = torch.einsum("lhm,l->hm", ks, all_ones)
+        attention_normalizer = torch.einsum("nhm,hm->nh", qs, ks_sum)  # [N, H]
+
+        # attentive aggregated results
+        attention_normalizer = torch.unsqueeze(
+            attention_normalizer, len(attention_normalizer.shape))  # [N, H, 1]
+        attention_normalizer += torch.ones_like(attention_normalizer) * N
+        attn_output = attention_num / attention_normalizer  # [N, H, D]
+
+        return attn_output.mean(dim=1)
+
+    def reset_parameters(self):
+        self.q.reset_parameters()
+        self.k.reset_parameters()
+        self.v.reset_parameters()
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}('
+                f'heads={self.heads}, '
+                f'head_channels={self.head_channels})')

{pyg_nightly-2.7.0.dev20250214 → pyg_nightly-2.7.0.dev20250216}/torch_geometric/nn/models/__init__.py

@@ -32,6 +32,7 @@ from .g_retriever import GRetriever
 from .git_mol import GITMol
 from .molecule_gpt import MoleculeGPT
 from .glem import GLEM
+from .sgformer import SGFormer
 # Deprecated:
 from torch_geometric.explain.algorithm.captum import (to_captum_input,
                                                       captum_output_to_dicts)
@@ -82,4 +83,5 @@ __all__ = classes = [
     'GITMol',
     'MoleculeGPT',
     'GLEM',
+    'SGFormer',
 ]

pyg_nightly-2.7.0.dev20250216/torch_geometric/nn/models/sgformer.py

@@ -0,0 +1,190 @@
+import torch
+import torch.nn.functional as F
+
+from torch_geometric.nn.attention import SGFormerAttention
+from torch_geometric.nn.conv import GCNConv
+
+
+class GraphModule(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        num_layers=2,
+        dropout=0.5,
+    ):
+        super().__init__()
+
+        self.convs = torch.nn.ModuleList()
+        self.fcs = torch.nn.ModuleList()
+        self.fcs.append(torch.nn.Linear(in_channels, hidden_channels))
+
+        self.bns = torch.nn.ModuleList()
+        self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
+        for _ in range(num_layers):
+            self.convs.append(GCNConv(hidden_channels, hidden_channels))
+            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
+
+        self.dropout = dropout
+        self.activation = F.relu
+
+    def reset_parameters(self):
+        for conv in self.convs:
+            conv.reset_parameters()
+        for bn in self.bns:
+            bn.reset_parameters()
+        for fc in self.fcs:
+            fc.reset_parameters()
+
+    def forward(self, x, edge_index):
+        x = self.fcs[0](x)
+        x = self.bns[0](x)
+        x = self.activation(x)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        last_x = x
+
+        for i, conv in enumerate(self.convs):
+            x = conv(x, edge_index)
+            x = self.bns[i + 1](x)
+            x = self.activation(x)
+            x = F.dropout(x, p=self.dropout, training=self.training)
+            x = x + last_x
+        return x
+
+
+class SGModule(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        num_layers=2,
+        num_heads=1,
+        dropout=0.5,
+    ):
+        super().__init__()
+
+        self.attns = torch.nn.ModuleList()
+        self.fcs = torch.nn.ModuleList()
+        self.fcs.append(torch.nn.Linear(in_channels, hidden_channels))
+        self.bns = torch.nn.ModuleList()
+        self.bns.append(torch.nn.LayerNorm(hidden_channels))
+        for _ in range(num_layers):
+            self.attns.append(
+                SGFormerAttention(hidden_channels, num_heads, hidden_channels))
+            self.bns.append(torch.nn.LayerNorm(hidden_channels))
+
+        self.dropout = dropout
+        self.activation = F.relu
+
+    def reset_parameters(self):
+        for attn in self.attns:
+            attn.reset_parameters()
+        for bn in self.bns:
+            bn.reset_parameters()
+        for fc in self.fcs:
+            fc.reset_parameters()
+
+    def forward(self, x):
+        layer_ = []
+
+        # input MLP layer
+        x = self.fcs[0](x)
+        x = self.bns[0](x)
+        x = self.activation(x)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+
+        # store as residual link
+        layer_.append(x)
+
+        for i, attn in enumerate(self.attns):
+            x = attn(x)
+            x = (x + layer_[i]) / 2.
+            x = self.bns[i + 1](x)
+            x = self.activation(x)
+            x = F.dropout(x, p=self.dropout, training=self.training)
+            layer_.append(x)
+
+        return x
+
+
+class SGFormer(torch.nn.Module):
+    r"""The sgformer module from the
+    `"SGFormer: Simplifying and Empowering Transformers for
+    Large-Graph Representations"
+    <https://arxiv.org/abs/2306.10759>`_ paper.
+
+    Args:
+        in_channels (int): Input channels.
+        hidden_channels (int): Hidden channels.
+        out_channels (int): Output channels.
+        trans_num_layers (int): The number of layers for all-pair attention.
+            (default: :obj:`2`)
+        trans_num_heads (int): The number of heads for attention.
+            (default: :obj:`1`)
+        trans_dropout (float): Global dropout rate.
+            (default: :obj:`0.5`)
+        gnn_num_layers (int): The number of layers for GNN.
+            (default: :obj:`3`)
+        gnn_dropout (float): GNN dropout rate.
+            (default: :obj:`0.5`)
+        graph_weight (float): The weight balance global and gnn module.
+            (default: :obj:`0.5`)
+        aggregate (str): Aggregate type.
+            (default: :obj:`add`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        trans_num_layers: int = 2,
+        trans_num_heads: int = 1,
+        trans_dropout: float = 0.5,
+        gnn_num_layers: int = 3,
+        gnn_dropout: float = 0.5,
+        graph_weight: float = 0.5,
+        aggregate: str = 'add',
+    ):
+        super().__init__()
+        self.trans_conv = SGModule(
+            in_channels,
+            hidden_channels,
+            trans_num_layers,
+            trans_num_heads,
+            trans_dropout,
+        )
+        self.graph_conv = GraphModule(
+            in_channels,
+            hidden_channels,
+            gnn_num_layers,
+            gnn_dropout,
+        )
+        self.graph_weight = graph_weight
+
+        self.aggregate = aggregate
+
+        if aggregate == 'add':
+            self.fc = torch.nn.Linear(hidden_channels, out_channels)
+        elif aggregate == 'cat':
+            self.fc = torch.nn.Linear(2 * hidden_channels, out_channels)
+        else:
+            raise ValueError(f'Invalid aggregate type:{aggregate}')
+
+        self.params1 = list(self.trans_conv.parameters())
+        self.params2 = list(self.graph_conv.parameters())
+        self.params2.extend(list(self.fc.parameters()))
+
+    def reset_parameters(self) -> None:
+        self.trans_conv.reset_parameters()
+        self.graph_conv.reset_parameters()
+        self.fc.reset_parameters()
+
+    def forward(self, x, edge_index):
+        x1 = self.trans_conv(x)
+        x2 = self.graph_conv(x, edge_index)
+        if self.aggregate == 'add':
+            x = self.graph_weight * x2 + (1 - self.graph_weight) * x1
+        else:
+            x = torch.cat((x1, x2), dim=1)
+        x = self.fc(x)
+        return F.log_softmax(x, dim=-1)

{pyg_nightly-2.7.0.dev20250214 → pyg_nightly-2.7.0.dev20250216}/torch_geometric/typing.py

@@ -9,6 +9,11 @@ import numpy as np
 import torch
 from torch import Tensor
 
+try:
+    from typing import TypeAlias  # type: ignore
+except ImportError:
+    from typing_extensions import TypeAlias
+
 WITH_PT20 = int(torch.__version__.split('.')[0]) >= 2
 WITH_PT21 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 1
 WITH_PT22 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 2
@@ -64,6 +69,16 @@ try:
         pyg_lib.sampler.neighbor_sample).parameters)
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = ('edge_weight' in inspect.signature(
         pyg_lib.sampler.neighbor_sample).parameters)
+    try:
+        torch.classes.pyg.CPUHashMap
+        WITH_CPU_HASH_MAP = True
+    except Exception:
+        WITH_CPU_HASH_MAP = False
+    try:
+        torch.classes.pyg.CUDAHashMap
+        WITH_CUDA_HASH_MAP = True
+    except Exception:
+        WITH_CUDA_HASH_MAP = False
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
         warnings.warn(f"An issue occurred while importing 'pyg-lib'. "
@@ -78,6 +93,32 @@ except Exception as e:
     WITH_METIS = False
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = False
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = False
+    WITH_CPU_HASH_MAP = False
+    WITH_CUDA_HASH_MAP = False
+
+if WITH_CPU_HASH_MAP:
+    CPUHashMap: TypeAlias = torch.classes.pyg.CPUHashMap
+else:
+
+    class CPUHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+
+if WITH_CUDA_HASH_MAP:
+    CUDAHashMap: TypeAlias = torch.classes.pyg.CUDAHashMap
+else:
+
+    class CUDAHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
 
 try:
     import torch_scatter  # noqa