pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251228__py3-none-any.whl

torch_geometric/nn/__init__.py

@@ -8,6 +8,7 @@ from .encoding import PositionalEncoding, TemporalEncoding
 from .summary import summary
 
 from .aggr import *  # noqa
+from .attention import *  # noqa
 from .conv import *  # noqa
 from .pool import *  # noqa
 from .glob import *  # noqa

torch_geometric/nn/aggr/base.py

@@ -135,7 +135,7 @@ class Aggregation(torch.nn.Module):
                 if index.numel() > 0 and dim_size <= int(index.max()):
                     raise ValueError(f"Encountered invalid 'dim_size' (got "
                                      f"'{dim_size}' but expected "
-                                     f">= '{int(index.max()) + 1}')")
+                                     f">= '{int(index.max()) + 1}')") from e
             raise e
 
     def __repr__(self) -> str:

torch_geometric/nn/aggr/equilibrium.py

@@ -52,7 +52,7 @@ class MomentumOptimizer(torch.nn.Module):
     layer. It is based on an unrolled Nesterov momentum algorithm.
 
     Args:
-        learning_rate (flaot): learning rate for optimizer.
+        learning_rate (float): learning rate for optimizer.
         momentum (float): momentum for optimizer.
         learnable (bool): If :obj:`True` then the :obj:`learning_rate` and
             :obj:`momentum` will be learnable parameters. If False they

torch_geometric/nn/aggr/fused.py

@@ -216,7 +216,7 @@ class FusedAggregation(Aggregation):
         outs: List[Optional[Tensor]] = []
 
         # Iterate over all reduction ops to compute first results:
-        for i, reduce in enumerate(self.reduce_ops):
+        for reduce in self.reduce_ops:
             if reduce is None:
                 outs.append(None)
                 continue

torch_geometric/nn/aggr/patch_transformer.py

@@ -32,6 +32,8 @@ class PatchTransformerAggregation(Aggregation):
         aggr (str or list[str], optional): The aggregation module, *e.g.*,
             :obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
             :obj:`"var"`, :obj:`"std"`. (default: :obj:`"mean"`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -43,6 +45,7 @@ class PatchTransformerAggregation(Aggregation):
         heads: int = 1,
         dropout: float = 0.0,
         aggr: Union[str, List[str]] = 'mean',
+        device: Optional[torch.device] = None,
     ) -> None:
         super().__init__()
 
@@ -55,12 +58,13 @@ class PatchTransformerAggregation(Aggregation):
         for aggr in self.aggrs:
             assert aggr in ['sum', 'mean', 'min', 'max', 'var', 'std']
 
-        self.lin = torch.nn.Linear(in_channels, hidden_channels)
+        self.lin = torch.nn.Linear(in_channels, hidden_channels, device=device)
         self.pad_projector = torch.nn.Linear(
             patch_size * hidden_channels,
             hidden_channels,
+            device=device,
         )
-        self.pe = PositionalEncoding(hidden_channels)
+        self.pe = PositionalEncoding(hidden_channels, device=device)
 
         self.blocks = torch.nn.ModuleList([
             MultiheadAttentionBlock(
@@ -68,12 +72,14 @@ class PatchTransformerAggregation(Aggregation):
                 heads=heads,
                 layer_norm=True,
                 dropout=dropout,
+                device=device,
             ) for _ in range(num_transformer_blocks)
         ])
 
         self.fc = torch.nn.Linear(
             hidden_channels * len(self.aggrs),
             out_channels,
+            device=device,
         )
 
     def reset_parameters(self) -> None:

torch_geometric/nn/aggr/set_transformer.py

@@ -38,7 +38,7 @@ class SetTransformerAggregation(Aggregation):
             (default: :obj:`1`)
         concat (bool, optional): If set to :obj:`False`, the seed embeddings
             are averaged instead of concatenated. (default: :obj:`True`)
-        norm (str, optional): If set to :obj:`True`, will apply layer
+        layer_norm (str, optional): If set to :obj:`True`, will apply layer
             normalization. (default: :obj:`False`)
         dropout (float, optional): Dropout probability of attention weights.
             (default: :obj:`0`)

torch_geometric/nn/aggr/utils.py

@@ -26,9 +26,11 @@ class MultiheadAttentionBlock(torch.nn.Module):
             normalization. (default: :obj:`True`)
         dropout (float, optional): Dropout probability of attention weights.
             (default: :obj:`0`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(self, channels: int, heads: int = 1, layer_norm: bool = True,
-                 dropout: float = 0.0):
+                 dropout: float = 0.0, device: Optional[torch.device] = None):
         super().__init__()
 
         self.channels = channels
@@ -40,10 +42,13 @@ class MultiheadAttentionBlock(torch.nn.Module):
             heads,
             batch_first=True,
             dropout=dropout,
+            device=device,
         )
-        self.lin = Linear(channels, channels)
-        self.layer_norm1 = LayerNorm(channels) if layer_norm else None
-        self.layer_norm2 = LayerNorm(channels) if layer_norm else None
+        self.lin = Linear(channels, channels, device=device)
+        self.layer_norm1 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
+        self.layer_norm2 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
 
     def reset_parameters(self):
         self.attn._reset_parameters()

torch_geometric/nn/attention/__init__.py

@@ -1,3 +1,11 @@
 from .performer import PerformerAttention
+from .qformer import QFormer
+from .sgformer import SGFormerAttention
+from .polynormer import PolynormerAttention
 
-__all__ = ['PerformerAttention']
+__all__ = classes = [
+    '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/attention/qformer.py

@@ -0,0 +1,71 @@
+from typing import Callable
+
+import torch
+
+
+class QFormer(torch.nn.Module):
+    r"""The Querying Transformer (Q-Former) from
+    `"BLIP-2: Bootstrapping Language-Image Pre-training
+    with Frozen Image Encoders and Large Language Models"
+    <https://arxiv.org/pdf/2301.12597>`_ paper.
+
+    Args:
+        input_dim (int): The number of features in the input.
+        hidden_dim (int): The dimension of the fnn in the encoder layer.
+        output_dim (int): The final output dimension.
+        num_heads (int): The number of multi-attention-heads.
+        num_layers (int): The number of sub-encoder-layers in the encoder.
+        dropout (int): The dropout value in each encoder layer.
+
+
+    .. note::
+        This is a simplified version of the original Q-Former implementation.
+    """
+    def __init__(
+            self,
+            input_dim: int,
+            hidden_dim: int,
+            output_dim: int,
+            num_heads: int,
+            num_layers: int,
+            dropout: float = 0.0,
+            activation: Callable = torch.nn.ReLU(),
+    ) -> None:
+
+        super().__init__()
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+
+        self.layer_norm = torch.nn.LayerNorm(input_dim)
+        self.encoder_layer = torch.nn.TransformerEncoderLayer(
+            d_model=input_dim,
+            nhead=num_heads,
+            dim_feedforward=hidden_dim,
+            dropout=dropout,
+            activation=activation,
+            batch_first=True,
+        )
+        self.encoder = torch.nn.TransformerEncoder(
+            self.encoder_layer,
+            num_layers=num_layers,
+        )
+        self.project = torch.nn.Linear(input_dim, output_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Forward pass.
+
+        Args:
+            x (torch.Tensor): Input sequence to the encoder layer.
+                :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with
+                batch-size :math:`B`, sequence length :math:`N`,
+                and feature dimension :math:`F`.
+        """
+        x = self.layer_norm(x)
+        x = self.encoder(x)
+        out = self.project(x)
+        return out
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}('
+                f'num_heads={self.num_heads}, '
+                f'num_layers={self.num_layers})')

torch_geometric/nn/attention/sgformer.py

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

torch_geometric/nn/conv/__init__.py

@@ -61,6 +61,7 @@ from .gps_conv import GPSConv
 from .antisymmetric_conv import AntiSymmetricConv
 from .dir_gnn_conv import DirGNNConv
 from .mixhop_conv import MixHopConv
+from .meshcnn_conv import MeshCNNConv
 
 import torch_geometric.nn.conv.utils  # noqa
 
@@ -131,6 +132,7 @@ __all__ = [
     'AntiSymmetricConv',
     'DirGNNConv',
     'MixHopConv',
+    'MeshCNNConv',
 ]
 
 classes = __all__

torch_geometric/nn/conv/appnp.py

@@ -109,7 +109,7 @@ class APPNP(MessagePassing):
                     edge_index = cache
 
         h = x
-        for k in range(self.K):
+        for _ in range(self.K):
             if self.dropout > 0 and self.training:
                 if isinstance(edge_index, Tensor):
                     if is_torch_sparse_tensor(edge_index):

torch_geometric/nn/conv/cugraph/gat_conv.py

@@ -26,6 +26,9 @@ class CuGraphGATConv(CuGraphModule):  # pragma: no cover
     :class:`~torch_geometric.nn.conv.GATConv` based on the :obj:`cugraph-ops`
     package that fuses message passing computation for accelerated execution
     and lower memory footprint.
+    The current method to enable :obj:`cugraph-ops`
+    is to use `The NVIDIA PyG Container
+    <https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pyg>`_.
     """
     def __init__(
         self,
@@ -67,6 +70,7 @@ class CuGraphGATConv(CuGraphModule):  # pragma: no cover
         self,
         x: Tensor,
         edge_index: EdgeIndex,
+        edge_attr: Tensor,
         max_num_neighbors: Optional[int] = None,
     ) -> Tensor:
         graph = self.get_cugraph(edge_index, max_num_neighbors)
@@ -75,10 +79,12 @@ class CuGraphGATConv(CuGraphModule):  # pragma: no cover
 
         if LEGACY_MODE:
             out = GATConvAgg(x, self.att, graph, self.heads, 'LeakyReLU',
-                             self.negative_slope, False, self.concat)
+                             self.negative_slope, False, self.concat,
+                             edge_feat=edge_attr)
         else:
             out = GATConvAgg(x, self.att, graph, self.heads, 'LeakyReLU',
-                             self.negative_slope, self.concat)
+                             self.negative_slope, self.concat,
+                             edge_feat=edge_attr)
 
         if self.bias is not None:
             out = out + self.bias

torch_geometric/nn/conv/cugraph/rgcn_conv.py

@@ -29,6 +29,9 @@ class CuGraphRGCNConv(CuGraphModule):  # pragma: no cover
     :class:`~torch_geometric.nn.conv.RGCNConv` based on the :obj:`cugraph-ops`
     package that fuses message passing computation for accelerated execution
     and lower memory footprint.
+    The current method to enable :obj:`cugraph-ops`
+    is to use `The NVIDIA PyG Container
+    <https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pyg>`_.
     """
     def __init__(self, in_channels: int, out_channels: int, num_relations: int,
                  num_bases: Optional[int] = None, aggr: str = 'mean',

torch_geometric/nn/conv/cugraph/sage_conv.py

@@ -27,6 +27,9 @@ class CuGraphSAGEConv(CuGraphModule):  # pragma: no cover
     :class:`~torch_geometric.nn.conv.SAGEConv` based on the :obj:`cugraph-ops`
     package that fuses message passing computation for accelerated execution
     and lower memory footprint.
+    The current method to enable :obj:`cugraph-ops`
+    is to use `The NVIDIA PyG Container
+    <https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pyg>`_.
     """
     def __init__(
         self,

torch_geometric/nn/conv/dna_conv.py

@@ -163,7 +163,7 @@ class MultiHead(Attention):
     def __repr__(self) -> str:  # pragma: no cover
         return (f'{self.__class__.__name__}({self.in_channels}, '
                 f'{self.out_channels}, heads={self.heads}, '
-                f'groups={self.groups}, dropout={self.droput}, '
+                f'groups={self.groups}, dropout={self.dropout}, '
                 f'bias={self.bias})')
 
 

torch_geometric/nn/conv/eg_conv.py

@@ -81,7 +81,7 @@ class EGConv(MessagePassing):
         self,
         in_channels: int,
         out_channels: int,
-        aggregators: List[str] = ['symnorm'],
+        aggregators: Optional[List[str]] = None,
         num_heads: int = 8,
         num_bases: int = 4,
         cached: bool = False,
@@ -96,23 +96,23 @@ class EGConv(MessagePassing):
                              f"divisible by the number of heads "
                              f"(got {num_heads})")
 
-        for a in aggregators:
-            if a not in ['sum', 'mean', 'symnorm', 'min', 'max', 'var', 'std']:
-                raise ValueError(f"Unsupported aggregator: '{a}'")
-
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.num_heads = num_heads
         self.num_bases = num_bases
         self.cached = cached
         self.add_self_loops = add_self_loops
-        self.aggregators = aggregators
+        self.aggregators = aggregators or ['symnorm']
+
+        for a in self.aggregators:
+            if a not in ['sum', 'mean', 'symnorm', 'min', 'max', 'var', 'std']:
+                raise ValueError(f"Unsupported aggregator: '{a}'")
 
         self.bases_lin = Linear(in_channels,
                                 (out_channels // num_heads) * num_bases,
                                 bias=False, weight_initializer='glorot')
         self.comb_lin = Linear(in_channels,
-                               num_heads * num_bases * len(aggregators))
+                               num_heads * num_bases * len(self.aggregators))
 
         if bias:
             self.bias = Parameter(torch.empty(out_channels))

torch_geometric/nn/conv/gen_conv.py

@@ -178,7 +178,7 @@ class GENConv(MessagePassing):
             self.lin_dst = Linear(in_channels[1], out_channels, bias=bias)
 
         channels = [out_channels]
-        for i in range(num_layers - 1):
+        for _ in range(num_layers - 1):
             channels.append(out_channels * expansion)
         channels.append(out_channels)
         self.mlp = MLP(channels, norm=norm, bias=bias)

torch_geometric/nn/conv/gravnet_conv.py

@@ -63,7 +63,8 @@ class GravNetConv(MessagePassing):
         if num_workers is not None:
             warnings.warn(
                 "'num_workers' attribute in '{self.__class__.__name__}' is "
-                "deprecated and will be removed in a future release")
+                "deprecated and will be removed in a future release",
+                stacklevel=2)
 
         self.in_channels = in_channels
         self.out_channels = out_channels

torch_geometric/nn/conv/hetero_conv.py

@@ -77,7 +77,8 @@ class HeteroConv(torch.nn.Module):
                 f"There exist node types ({src_node_types - dst_node_types}) "
                 f"whose representations do not get updated during message "
                 f"passing as they do not occur as destination type in any "
-                f"edge type. This may lead to unexpected behavior.")
+                f"edge type. This may lead to unexpected behavior.",
+                stacklevel=2)
 
         self.convs = ModuleDict(convs)
         self.aggr = aggr