pyg-nightly 2.7.0.dev20250228__tar.gz → 2.7.0.dev20250302__tar.gz

{pyg_nightly-2.7.0.dev20250228 → pyg_nightly-2.7.0.dev20250302}/PKG-INFO

@@ -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 → pyg_nightly-2.7.0.dev20250302}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250228"
+version="2.7.0.dev20250302"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]

{pyg_nightly-2.7.0.dev20250228 → pyg_nightly-2.7.0.dev20250302}/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',

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

@@ -0,0 +1,100 @@
+from typing import Optional
+
+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, 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))
+
+        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("blhm,blhd->bhmd", ks, vs)
+        attention_num = torch.einsum("bnhm,bhmd->bnhd", qs, kvs)
+        attention_num += N * vs
+
+        # denominator
+        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))
+        attention_normalizer += torch.ones_like(attention_normalizer) * N
+        attn_output = attention_num / attention_normalizer
+
+        return attn_output.mean(dim=2)
+
+    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.dev20250228 → pyg_nightly-2.7.0.dev20250302}/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

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

@@ -1,78 +0,0 @@
-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})')