pyg-nightly 2.7.0.dev20250228__py3-none-any.whl → 2.7.0.dev20250302__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250228.dist-info → pyg_nightly-2.7.0.dev20250302.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=55I0x87kSkECyYdIRxH0kF2v6J9VGbG_99GlaLPMQXQ,1978
+torch_geometric/__init__.py,sha256=-ySUkCuMvYC1iL9it4tdObi1jhTkJrpNJ86BUKiAxE0,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
@@ -334,7 +334,7 @@ torch_geometric/nn/aggr/variance_preserving.py,sha256=fu-U_aGYpVLpgSFvVg0ONMe6nq
 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/attention/sgformer.py,sha256=U4R_tGF1IWyKDlV0t4jFNatW2-payc-hPGM3sFdiYBE,3812
 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
@@ -458,7 +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/sgformer.py,sha256=V-F3J-yEdjuhkSsCkQePG_ByNHAu0BcKpknVxjCI3KY,6761
 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
@@ -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.dev20250228.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250228.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
-pyg_nightly-2.7.0.dev20250228.dist-info/METADATA,sha256=WqOwZeA_woFXbV3pQXk5CAIZeKr1txi5OpmJnrtXtww,63021
-pyg_nightly-2.7.0.dev20250228.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250302.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250302.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
+pyg_nightly-2.7.0.dev20250302.dist-info/METADATA,sha256=sLyKGhdjP0OYneFdjiW5cyxTIb8wDONfzlPRcuwVrYw,63021
+pyg_nightly-2.7.0.dev20250302.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.dev20250228'
+__version__ = '2.7.0.dev20250302'
 
 __all__ = [
     'Index',

torch_geometric/nn/attention/sgformer.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import Tensor
 
@@ -38,34 +40,54 @@ class SGFormerAttention(torch.nn.Module):
         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)
+    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
+        qs, ks, vs = self.q(x), self.k(x), self.v(x)
+        # reshape and permute q, k and v to proper shape
+        # (b, n, num_heads * head_channels) to (b, n, num_heads, head_channels)
+        qs, ks, vs = map(
+            lambda t: t.reshape(B, N, self.heads, self.head_channels),
+            (qs, ks, vs))
 
-        # 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]
+        if mask is not None:
+            mask = mask[:, :, None, None]
+            vs.masked_fill_(~mask, 0.)
+        # replace 0's with epsilon
+        epsilon = 1e-6
+        qs[qs == 0] = epsilon
+        ks[ks == 0] = epsilon
+        # normalize input, shape not changed
+        qs, ks = map(
+            lambda t: t / torch.linalg.norm(t, ord=2, dim=-1, keepdim=True),
+            (qs, ks))
 
         # numerator
-        kvs = torch.einsum("lhm,lhd->hmd", ks, vs)
-        attention_num = torch.einsum("nhm,hmd->nhd", qs, kvs)  # [N, H, D]
+        kvs = torch.einsum("blhm,blhd->bhmd", ks, vs)
+        attention_num = torch.einsum("bnhm,bhmd->bnhd", qs, kvs)
         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]
-
+        all_ones = torch.ones([B, N]).to(ks.device)
+        ks_sum = torch.einsum("blhm,bl->bhm", ks, all_ones)
+        attention_normalizer = torch.einsum("bnhm,bhm->bnh", qs, ks_sum)
         # attentive aggregated results
-        attention_normalizer = torch.unsqueeze(
-            attention_normalizer, len(attention_normalizer.shape))  # [N, H, 1]
+        attention_normalizer = torch.unsqueeze(attention_normalizer,
+                                               len(attention_normalizer.shape))
         attention_normalizer += torch.ones_like(attention_normalizer) * N
-        attn_output = attention_num / attention_normalizer  # [N, H, D]
+        attn_output = attention_num / attention_normalizer
 
-        return attn_output.mean(dim=1)
+        return attn_output.mean(dim=2)
 
     def reset_parameters(self):
         self.q.reset_parameters()

torch_geometric/nn/models/sgformer.py

@@ -1,8 +1,12 @@
+from typing import Optional
+
 import torch
 import torch.nn.functional as F
+from torch import Tensor
 
 from torch_geometric.nn.attention import SGFormerAttention
 from torch_geometric.nn.conv import GCNConv
+from torch_geometric.utils import to_dense_batch
 
 
 class GraphModule(torch.nn.Module):
@@ -84,7 +88,11 @@ class SGModule(torch.nn.Module):
         for fc in self.fcs:
             fc.reset_parameters()
 
-    def forward(self, x):
+    def forward(self, x: Tensor, batch: Tensor):
+        # to dense batch expects sorted batch
+        batch, indices = batch.sort(stable=True)
+        x = x[indices]
+        x, mask = to_dense_batch(x, batch)
         layer_ = []
 
         # input MLP layer
@@ -97,14 +105,17 @@ class SGModule(torch.nn.Module):
         layer_.append(x)
 
         for i, attn in enumerate(self.attns):
-            x = attn(x)
+            x = attn(x, mask)
             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
+        x_mask = x[mask]
+        # reverse the sorting
+        unsorted_x_mask = x_mask[indices.argsort()]
+        return unsorted_x_mask
 
 
 class SGFormer(torch.nn.Module):
@@ -179,8 +190,22 @@ class SGFormer(torch.nn.Module):
         self.graph_conv.reset_parameters()
         self.fc.reset_parameters()
 
-    def forward(self, x, edge_index):
-        x1 = self.trans_conv(x)
+    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.
+        """
+        x1 = self.trans_conv(x, batch)
         x2 = self.graph_conv(x, edge_index)
         if self.aggregate == 'add':
             x = self.graph_weight * x2 + (1 - self.graph_weight) * x1