pyg-nightly 2.7.0.dev20250214__py3-none-any.whl → 2.7.0.dev20250216__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250214.dist-info → pyg_nightly-2.7.0.dev20250216.dist-info}/METADATA

@@ -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.dist-info → pyg_nightly-2.7.0.dev20250216.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=wEIwjfiE7YtSKe1hUelXwhBp2yBtwx2u2gRooK6Gb_s,1904
+torch_geometric/__init__.py,sha256=QlYg3d7vRWC0k_OzPCmAanyTUBiTvikPuzM9feT0REk,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,6 +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=VBuz9n16ouSk2u4DifpCIW5MXiuoH5UGO4rPj-Celjw,4418
 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
@@ -18,7 +19,7 @@ torch_geometric/logging.py,sha256=HmHHLiCcM64k-6UYNOSfXPIeSGNAyiGGcn8cD8tlyuQ,85
 torch_geometric/resolver.py,sha256=fn-_6mCpI2xv7eDZnIFcYrHOn0IrwbkWFLDb9laQrWI,1270
 torch_geometric/seed.py,sha256=MJLbVwpb9i8mK3oi32sS__Cq-dRq_afTeoOL_HoA9ko,372
 torch_geometric/template.py,sha256=rqjDWgcSAgTCiV4bkOjWRPaO4PpUdC_RXigzxxBqAu8,1060
-torch_geometric/typing.py,sha256=SzuZPdeYLt7_lFUHHcAlaggxqLA0VZ_kx8s0iy_tMIw,14429
+torch_geometric/typing.py,sha256=mtSM6QhCsrohstnyvqMuxEIajCkhcvkQKOU4uVu-nDs,15596
 torch_geometric/warnings.py,sha256=t114CbkrmiqkXaavx5g7OO52dLdktf-U__B5QqYIQvI,413
 torch_geometric/contrib/__init__.py,sha256=0pWkmXfZtbdr-AKwlii5LTFggTEH-MCrSKpZxrtPlVs,352
 torch_geometric/contrib/datasets/__init__.py,sha256=lrGnWsEiJf5zsBRmshGZZFN_uYR2ezDjbj9n9nCpvtk,23
@@ -330,9 +331,10 @@ torch_geometric/nn/aggr/set_transformer.py,sha256=FG7_JizpFX14M6VSCwLSjYXYdJ1ZiQ
 torch_geometric/nn/aggr/sort.py,sha256=bvOOWnFkNOBOZih4rqVZQsjfeDX3vmXo1bpPSFD846w,2507
 torch_geometric/nn/aggr/utils.py,sha256=CLJ-ZrVWYIOBpdhQBLAz94dj3cMKKKc3qwGr4DFbiCU,8338
 torch_geometric/nn/aggr/variance_preserving.py,sha256=fu-U_aGYpVLpgSFvVg0ONMe6nqoyv8tZ6Y35qMYTf9w,1126
-torch_geometric/nn/attention/__init__.py,sha256=1lCB7zh7uM6FkpW81S9U4CvxTwpCkz59KatPTIE9UmA,127
+torch_geometric/nn/attention/__init__.py,sha256=wLKTmlfP7qL9sZHy4cmDFHEtdwa-MEKE1dT51L1_w10,192
 torch_geometric/nn/attention/performer.py,sha256=2PCDn4_-oNTao2-DkXIaoi18anP01OxRELF2pvp-jk8,7357
 torch_geometric/nn/attention/qformer.py,sha256=7J-pWm_vpumK38IC-iCBz4oqL-BEIofEIxJ0wfjWq9A,2338
+torch_geometric/nn/attention/sgformer.py,sha256=QT_3bVOPsqPY0fWVdwyr8E0LL9u9TwC_sv0_2EggdpA,2893
 torch_geometric/nn/conv/__init__.py,sha256=37zTdt0gfSAUPMtwXjZg5mWx_itojJVFNODYR1h1ch0,3515
 torch_geometric/nn/conv/agnn_conv.py,sha256=5nEPLx_BBHcDaO6HWzLuHfXc0Yd_reKynAOH0Iq09lU,3077
 torch_geometric/nn/conv/antisymmetric_conv.py,sha256=dhA6sCETy1jlXReYJZBSyToOcL_mZ1wL10fMIb8Ppuw,4387
@@ -424,7 +426,7 @@ torch_geometric/nn/kge/distmult.py,sha256=dGQ0bVzjreZgFN1lXE23_IIidsiOq7ehPrMb-N
 torch_geometric/nn/kge/loader.py,sha256=5Uc1j3OUMQnBYSHDqL7pLCty1siFLzoPkztigYO2zP8,771
 torch_geometric/nn/kge/rotate.py,sha256=XLuO1AbyTt5cJxr97ZzoyAyIEsHKesgW5TvDmnGJAao,3208
 torch_geometric/nn/kge/transe.py,sha256=jlejq5BLMm-sb1wWcLDp7pZqCdelWBgjDIC8ctbjSdU,3088
-torch_geometric/nn/models/__init__.py,sha256=vWMKzGBVxA1Fm0uGDLnH4jzYgfhK34CQTRJ-xi5pf5k,2150
+torch_geometric/nn/models/__init__.py,sha256=gmBRXrbjkxLv_g0hfI87bSXMgwFoIo3XgslpGGdMs3g,2197
 torch_geometric/nn/models/attentive_fp.py,sha256=tkgvw28wg9-JqHIfBllfCwTHrZIUiv85yZJcDqjz3z0,6634
 torch_geometric/nn/models/autoencoder.py,sha256=nGje-zty78Y3hxOJ9o0_6QziJjOvBlknk6z0_fDQwQU,10770
 torch_geometric/nn/models/basic_gnn.py,sha256=PGa0RUMyvrNy_5yRI2jX_zwPsmZXwOQWfsWvxOiHsSk,31225
@@ -456,6 +458,7 @@ torch_geometric/nn/models/re_net.py,sha256=pz66q5b5BoGDNVQvpEGS2RGoeKvpjkYAv9r3W
 torch_geometric/nn/models/rect.py,sha256=2F3XyyvHTAEuqfJpiNB5M8pSGy738LhPiom5I-SDWqM,2808
 torch_geometric/nn/models/rev_gnn.py,sha256=1b6wU-6YTuLsWn5p8c5LXQm2KugEAVcEYJKZbWTDvgQ,11796
 torch_geometric/nn/models/schnet.py,sha256=0aaHrVtxApdvn3RHCGLQJW1MbIb--CSYUrx9O3hDOZM,16656
+torch_geometric/nn/models/sgformer.py,sha256=QacmFjlTSUgk6T5Dk4NfG9zmxHZqr0ggx1ElNrPMJNc,5878
 torch_geometric/nn/models/signed_gcn.py,sha256=J40CnedFIqtKI1LhW1ITSEFRbA_XiJZL6lASrKwUEAI,9841
 torch_geometric/nn/models/tgn.py,sha256=kEGdfLJybkbMT4UMoAh2nCzfX3_nDjfm1cicuPHEwAM,11878
 torch_geometric/nn/models/visnet.py,sha256=97OFMCsPDEI5BCSi7RhoRcU2CNRp7zck2tEzrltFZj4,43192
@@ -630,6 +633,6 @@ 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.dev20250214.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
-pyg_nightly-2.7.0.dev20250214.dist-info/METADATA,sha256=7oAMP5u6qz3XuWwE8lms2zjkk_lGmt3Q7-3FswFCbUk,62977
-pyg_nightly-2.7.0.dev20250214.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250216.dist-info/WHEEL,sha256=CpUCUxeHQbRN5UGRQHYRJorO5Af-Qy_fHMctcQ8DSGI,82
+pyg_nightly-2.7.0.dev20250216.dist-info/METADATA,sha256=nrJ478w_yYVSpIyqlfR5lxdYTu6NOIzCv_atkMY5w5w,62999
+pyg_nightly-2.7.0.dev20250216.dist-info/RECORD,,

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,
     ])

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 {}))

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',
 ]

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})')

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',
 ]

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)

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