pyg-nightly 2.7.0.dev20250703__py3-none-any.whl → 2.7.0.dev20250704__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250703.dist-info → pyg_nightly-2.7.0.dev20250704.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=8_AcpgPpDfVr7gwK3sYPD-pu9JYUmziOQzs9SPRvqcE,2250
+torch_geometric/__init__.py,sha256=GOuL0XBOcsFqK-Q-c_STDpzZAG-vsctiDiU_Tg9W3t8,2250
 torch_geometric/_compile.py,sha256=9yqMTBKatZPr40WavJz9FjNi7pQj8YZAZOyZmmRGXgc,1351
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -119,7 +119,7 @@ torch_geometric/datasets/medshapenet.py,sha256=eCBCXKpueweCwDSf_Q4_MwVA3IbJd04FS
 torch_geometric/datasets/mixhop_synthetic_dataset.py,sha256=4NNvTHUvvV6pcqQCyVDS5XhppXUeF2H9GTfFoc49eyU,3951
 torch_geometric/datasets/mnist_superpixels.py,sha256=o2ArbZ0_OE0u8VCaHmWwvngESlOFr9oM9dSEP_tjAS4,3340
 torch_geometric/datasets/modelnet.py,sha256=-qmLjlQiKVWmtHefAIIE97dQxEcaBfetMJnvgYZuwkg,5347
-torch_geometric/datasets/molecule_gpt_dataset.py,sha256=gVZv14PuZCanE4oxxHlqRNrvzGv6_KN318q5yFA3lS0,18797
+torch_geometric/datasets/molecule_gpt_dataset.py,sha256=TFBduE3_3xxTFSHL3tirV-OAlBjSi6iHPOHJGQ_-tug,18785
 torch_geometric/datasets/molecule_net.py,sha256=pMzaJzd-LbBncZ0VoC87HfA8d1F4NwCWTb5YKvLM890,7404
 torch_geometric/datasets/movie_lens.py,sha256=M4Bu0Xus8IkW8GYzjxPxSdPXNbcCCx9cu6cncxBvLx8,4033
 torch_geometric/datasets/movie_lens_100k.py,sha256=eTpBAteM3jqTEtiwLxmhVj4r8JvftvPx8Hvs-3ZIHlU,6057
@@ -335,8 +335,9 @@ 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=SQvdc0g6p_E2j0prA14MW2ekjEDvV-g545N0Q85uc-o,8625
 torch_geometric/nn/aggr/variance_preserving.py,sha256=fu-U_aGYpVLpgSFvVg0ONMe6nqoyv8tZ6Y35qMYTf9w,1126
-torch_geometric/nn/attention/__init__.py,sha256=wLKTmlfP7qL9sZHy4cmDFHEtdwa-MEKE1dT51L1_w10,192
+torch_geometric/nn/attention/__init__.py,sha256=w-jDQFpVqARJKjttTgKkD9kkAqRJl4MpASCfiNYIfr0,263
 torch_geometric/nn/attention/performer.py,sha256=2PCDn4_-oNTao2-DkXIaoi18anP01OxRELF2pvp-jk8,7357
+torch_geometric/nn/attention/polynormer.py,sha256=uBxGs0nldp6oGlByqbxgEk23VeXLEd6B3myS5BOKDRs,3998
 torch_geometric/nn/attention/qformer.py,sha256=7J-pWm_vpumK38IC-iCBz4oqL-BEIofEIxJ0wfjWq9A,2338
 torch_geometric/nn/attention/sgformer.py,sha256=OBC5HQxbY289bPDtwN8UbPH46To2GRTeVN-najogD-o,3747
 torch_geometric/nn/conv/__init__.py,sha256=8CK-DFG2PEo2ZaFyg-IUlQH8ecQoDDi556uv3ugeQyc,3572
@@ -431,7 +432,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=fbHQauZw9Snvl2PuN5cjZoAW8SwUl6E-p2IOmwUKB3A,2395
+torch_geometric/nn/models/__init__.py,sha256=71Hqc-ZMfCKn9lelFYDjpHXapbEa0wqVAd2OXCb1y5o,2448
 torch_geometric/nn/models/attentive_fp.py,sha256=1z3iTV2O5W9tqHFAdno8FeBFeXmuG-TDZk4lwwVh3Ac,6634
 torch_geometric/nn/models/attract_repel.py,sha256=h9OyogT0NY0xiT0DkpJHMxH6ZUmo8R-CmwZdKEwq8Ek,5277
 torch_geometric/nn/models/autoencoder.py,sha256=nGje-zty78Y3hxOJ9o0_6QziJjOvBlknk6z0_fDQwQU,10770
@@ -461,6 +462,7 @@ torch_geometric/nn/models/molecule_gpt.py,sha256=k-XULH6jaurj-R2EE4sIWTkqlNqa3Cz
 torch_geometric/nn/models/neural_fingerprint.py,sha256=pTLJgU9Uh2Lnf9bggLj4cKI8YdEFcMF-9MALuubqbuQ,2378
 torch_geometric/nn/models/node2vec.py,sha256=81Ku4Rp4IwLEAy06KEgJ2fYtXXVL_uv_Hb8lBr6YXrE,7664
 torch_geometric/nn/models/pmlp.py,sha256=dcAASVSyQMMhItSfEJWPeAFh0R3tNCwAHwdrShwQ8o4,3538
+torch_geometric/nn/models/polynormer.py,sha256=mayWdzdolT5PCt_Oo7UGG-JUripMHWB2lUWF1bh6goU,7640
 torch_geometric/nn/models/protein_mpnn.py,sha256=QXHfltiJPmakpzgJKw_1vwCGBlszv9nfY4r4F38Sg9k,11031
 torch_geometric/nn/models/re_net.py,sha256=pz66q5b5BoGDNVQvpEGS2RGoeKvpjkYAv9r3WAuvITk,8986
 torch_geometric/nn/models/rect.py,sha256=2F3XyyvHTAEuqfJpiNB5M8pSGy738LhPiom5I-SDWqM,2808
@@ -643,7 +645,7 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=b-HnVesXjyJ_L1N-DnjiRiRVf7lhwKaBQF_2i5YMVSU,208
 torch_geometric/visualization/graph.py,sha256=mfZHXYfiU-CWMtfawYc80IxVwVmtK9hbIkSKhM_j7oI,14311
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250703.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250703.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250703.dist-info/METADATA,sha256=lm-xVNswGkfSaczkhV_ANegM0oGwBEfhCv75gwj0X5Q,63005
-pyg_nightly-2.7.0.dev20250703.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250704.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250704.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250704.dist-info/METADATA,sha256=Nau44bIMI13OXEqYNOlI1hYfDz8FXpUoPydv-JxRW2Q,63005
+pyg_nightly-2.7.0.dev20250704.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.dev20250703'
+__version__ = '2.7.0.dev20250704'
 
 __all__ = [
     'Index',

torch_geometric/datasets/molecule_gpt_dataset.py

@@ -438,7 +438,7 @@ class MoleculeGPTDataset(InMemoryDataset):
         for mol in tqdm(suppl):
             if mol.HasProp('PUBCHEM_COMPOUND_CID'):
                 CID = mol.GetProp("PUBCHEM_COMPOUND_CID")
-                CAN_SMILES = mol.GetProp("PUBCHEM_OPENEYE_CAN_SMILES")
+                CAN_SMILES = mol.GetProp("PUBCHEM_SMILES")
 
                 m: Chem.Mol = Chem.MolFromSmiles(CAN_SMILES)
                 if m is None:

torch_geometric/nn/attention/__init__.py

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

torch_geometric/nn/attention/polynormer.py

@@ -0,0 +1,107 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class PolynormerAttention(torch.nn.Module):
+    r"""The polynomial-expressive attention mechanism from the
+    `"Polynormer: Polynomial-Expressive Graph Transformer in Linear Time"
+    <https://arxiv.org/abs/2403.01232>`_ paper.
+
+    Args:
+        channels (int): Size of each input sample.
+        heads (int, optional): Number of parallel attention heads.
+        head_channels (int, optional): Size of each attention head.
+            (default: :obj:`64.`)
+        beta (float, optional): Polynormer beta initialization.
+            (default: :obj:`0.9`)
+        qkv_bias (bool, optional): If specified, add bias to query, key
+            and value in the self attention. (default: :obj:`False`)
+        qk_shared (bool optional): Whether weight of query and key are shared.
+            (default: :obj:`True`)
+        dropout (float, optional): Dropout probability of the final
+            attention output. (default: :obj:`0.0`)
+    """
+    def __init__(
+        self,
+        channels: int,
+        heads: int,
+        head_channels: int = 64,
+        beta: float = 0.9,
+        qkv_bias: bool = False,
+        qk_shared: bool = True,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        self.head_channels = head_channels
+        self.heads = heads
+        self.beta = beta
+        self.qk_shared = qk_shared
+
+        inner_channels = heads * head_channels
+        self.h_lins = torch.nn.Linear(channels, inner_channels)
+        if not self.qk_shared:
+            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)
+        self.lns = torch.nn.LayerNorm(inner_channels)
+        self.lin_out = torch.nn.Linear(inner_channels, inner_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+
+    def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
+        r"""Forward pass.
+
+        Args:
+            x (torch.Tensor): Node feature tensor
+                :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with
+                batch-size :math:`B`, (maximum) number of nodes :math:`N` for
+                each graph, and feature dimension :math:`F`.
+            mask (torch.Tensor, optional): Mask matrix
+                :math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}` indicating
+                the valid nodes for each graph. (default: :obj:`None`)
+        """
+        B, N, *_ = x.shape
+        h = self.h_lins(x)
+        k = self.k(x).sigmoid().view(B, N, self.head_channels, self.heads)
+        if self.qk_shared:
+            q = k
+        else:
+            q = F.sigmoid(self.q(x)).view(B, N, self.head_channels, self.heads)
+        v = self.v(x).view(B, N, self.head_channels, self.heads)
+
+        if mask is not None:
+            mask = mask[:, :, None, None]
+            v.masked_fill_(~mask, 0.)
+
+        # numerator
+        kv = torch.einsum('bndh, bnmh -> bdmh', k, v)
+        num = torch.einsum('bndh, bdmh -> bnmh', q, kv)
+
+        # denominator
+        k_sum = torch.einsum('bndh -> bdh', k)
+        den = torch.einsum('bndh, bdh -> bnh', q, k_sum).unsqueeze(2)
+
+        # linear global attention based on kernel trick
+        x = (num / (den + 1e-6)).reshape(B, N, -1)
+        x = self.lns(x) * (h + self.beta)
+        x = F.relu(self.lin_out(x))
+        x = self.dropout(x)
+
+        return x
+
+    def reset_parameters(self) -> None:
+        self.h_lins.reset_parameters()
+        if not self.qk_shared:
+            self.q.reset_parameters()
+        self.k.reset_parameters()
+        self.v.reset_parameters()
+        self.lns.reset_parameters()
+        self.lin_out.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

@@ -35,6 +35,7 @@ from .molecule_gpt import MoleculeGPT
 from .protein_mpnn import ProteinMPNN
 from .glem import GLEM
 from .sgformer import SGFormer
+from .polynormer import Polynormer
 # Deprecated:
 from torch_geometric.explain.algorithm.captum import (to_captum_input,
                                                       captum_output_to_dicts)
@@ -90,5 +91,6 @@ __all__ = classes = [
     'ProteinMPNN',
     'GLEM',
     'SGFormer',
+    'Polynormer',
     'ARLinkPredictor',
 ]

torch_geometric/nn/models/polynormer.py

@@ -0,0 +1,206 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from torch_geometric.nn import GATConv, GCNConv
+from torch_geometric.nn.attention import PolynormerAttention
+from torch_geometric.utils import to_dense_batch
+
+
+class Polynormer(torch.nn.Module):
+    r"""The polynormer module from the
+    `"Polynormer: polynomial-expressive graph
+    transformer in linear time"
+    <https://arxiv.org/abs/2403.01232>`_ paper.
+
+    Args:
+        in_channels (int): Input channels.
+        hidden_channels (int): Hidden channels.
+        out_channels (int): Output channels.
+        local_layers (int): The number of local attention layers.
+            (default: :obj:`7`)
+        global_layers (int): The number of global attention layers.
+            (default: :obj:`2`)
+        in_dropout (float): Input dropout rate.
+            (default: :obj:`0.15`)
+        dropout (float): Dropout rate.
+            (default: :obj:`0.5`)
+        global_dropout (float): Global dropout rate.
+            (default: :obj:`0.5`)
+        heads (int): The number of heads.
+            (default: :obj:`1`)
+        beta (float): Aggregate type.
+            (default: :obj:`0.9`)
+        qk_shared (bool optional): Whether weight of query and key are shared.
+            (default: :obj:`True`)
+        pre_ln (bool): Pre layer normalization.
+            (default: :obj:`False`)
+        post_bn (bool): Post batch normlization.
+            (default: :obj:`True`)
+        local_attn (bool): Whether use local attention.
+            (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        local_layers: int = 7,
+        global_layers: int = 2,
+        in_dropout: float = 0.15,
+        dropout: float = 0.5,
+        global_dropout: float = 0.5,
+        heads: int = 1,
+        beta: float = 0.9,
+        qk_shared: bool = False,
+        pre_ln: bool = False,
+        post_bn: bool = True,
+        local_attn: bool = False,
+    ) -> None:
+        super().__init__()
+        self._global = False
+        self.in_drop = in_dropout
+        self.dropout = dropout
+        self.pre_ln = pre_ln
+        self.post_bn = post_bn
+
+        self.beta = beta
+
+        self.h_lins = torch.nn.ModuleList()
+        self.local_convs = torch.nn.ModuleList()
+        self.lins = torch.nn.ModuleList()
+        self.lns = torch.nn.ModuleList()
+        if self.pre_ln:
+            self.pre_lns = torch.nn.ModuleList()
+        if self.post_bn:
+            self.post_bns = torch.nn.ModuleList()
+
+        # first layer
+        inner_channels = heads * hidden_channels
+        self.h_lins.append(torch.nn.Linear(in_channels, inner_channels))
+        if local_attn:
+            self.local_convs.append(
+                GATConv(in_channels, hidden_channels, heads=heads, concat=True,
+                        add_self_loops=False, bias=False))
+        else:
+            self.local_convs.append(
+                GCNConv(in_channels, inner_channels, cached=False,
+                        normalize=True))
+
+        self.lins.append(torch.nn.Linear(in_channels, inner_channels))
+        self.lns.append(torch.nn.LayerNorm(inner_channels))
+        if self.pre_ln:
+            self.pre_lns.append(torch.nn.LayerNorm(in_channels))
+        if self.post_bn:
+            self.post_bns.append(torch.nn.BatchNorm1d(inner_channels))
+
+        # following layers
+        for _ in range(local_layers - 1):
+            self.h_lins.append(torch.nn.Linear(inner_channels, inner_channels))
+            if local_attn:
+                self.local_convs.append(
+                    GATConv(inner_channels, hidden_channels, heads=heads,
+                            concat=True, add_self_loops=False, bias=False))
+            else:
+                self.local_convs.append(
+                    GCNConv(inner_channels, inner_channels, cached=False,
+                            normalize=True))
+
+            self.lins.append(torch.nn.Linear(inner_channels, inner_channels))
+            self.lns.append(torch.nn.LayerNorm(inner_channels))
+            if self.pre_ln:
+                self.pre_lns.append(torch.nn.LayerNorm(heads *
+                                                       hidden_channels))
+            if self.post_bn:
+                self.post_bns.append(torch.nn.BatchNorm1d(inner_channels))
+
+        self.lin_in = torch.nn.Linear(in_channels, inner_channels)
+        self.ln = torch.nn.LayerNorm(inner_channels)
+
+        self.global_attn = torch.nn.ModuleList()
+        for _ in range(global_layers):
+            self.global_attn.append(
+                PolynormerAttention(
+                    channels=hidden_channels,
+                    heads=heads,
+                    head_channels=hidden_channels,
+                    beta=beta,
+                    dropout=global_dropout,
+                    qk_shared=qk_shared,
+                ))
+        self.pred_local = torch.nn.Linear(inner_channels, out_channels)
+        self.pred_global = torch.nn.Linear(inner_channels, out_channels)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        for local_conv in self.local_convs:
+            local_conv.reset_parameters()
+        for attn in self.global_attn:
+            attn.reset_parameters()
+        for lin in self.lins:
+            lin.reset_parameters()
+        for h_lin in self.h_lins:
+            h_lin.reset_parameters()
+        for ln in self.lns:
+            ln.reset_parameters()
+        if self.pre_ln:
+            for p_ln in self.pre_lns:
+                p_ln.reset_parameters()
+        if self.post_bn:
+            for p_bn in self.post_bns:
+                p_bn.reset_parameters()
+        self.lin_in.reset_parameters()
+        self.ln.reset_parameters()
+        self.pred_local.reset_parameters()
+        self.pred_global.reset_parameters()
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Optional[Tensor],
+    ) -> Tensor:
+        r"""Forward pass.
+
+        Args:
+            x (torch.Tensor): The input node features.
+            edge_index (torch.Tensor or SparseTensor): The edge indices.
+            batch (torch.Tensor, optional): The batch vector
+                :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
+                each element to a specific example.
+        """
+        x = F.dropout(x, p=self.in_drop, training=self.training)
+
+        # equivariant local attention
+        x_local = 0
+        for i, local_conv in enumerate(self.local_convs):
+            if self.pre_ln:
+                x = self.pre_lns[i](x)
+            h = self.h_lins[i](x)
+            h = F.relu(h)
+            x = local_conv(x, edge_index) + self.lins[i](x)
+            if self.post_bn:
+                x = self.post_bns[i](x)
+            x = F.relu(x)
+            x = F.dropout(x, p=self.dropout, training=self.training)
+            x = (1 - self.beta) * self.lns[i](h * x) + self.beta * x
+            x_local = x_local + x
+
+        # equivariant global attention
+        if self._global:
+            batch, indices = batch.sort()
+            rev_perm = torch.empty_like(indices)
+            rev_perm[indices] = torch.arange(len(indices),
+                                             device=indices.device)
+            x_local = self.ln(x_local[indices])
+            x_global, mask = to_dense_batch(x_local, batch)
+            for attn in self.global_attn:
+                x_global = attn(x_global, mask)
+            x = x_global[mask][rev_perm]
+            x = self.pred_global(x)
+        else:
+            x = self.pred_local(x_local)
+
+        return F.log_softmax(x, dim=-1)