pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__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/__init__.py

@@ -25,6 +25,7 @@ from .deep_sets import DeepSetsAggregation
 from .set_transformer import SetTransformerAggregation
 from .lcm import LCMAggregation
 from .variance_preserving import VariancePreservingAggregation
+from .patch_transformer import PatchTransformerAggregation
 
 __all__ = classes = [
     'Aggregation',
@@ -53,4 +54,5 @@ __all__ = classes = [
     'SetTransformerAggregation',
     'LCMAggregation',
     'VariancePreservingAggregation',
+    'PatchTransformerAggregation',
 ]

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

@@ -0,0 +1,149 @@
+import math
+from typing import List, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.experimental import disable_dynamic_shapes
+from torch_geometric.nn.aggr import Aggregation
+from torch_geometric.nn.aggr.utils import MultiheadAttentionBlock
+from torch_geometric.nn.encoding import PositionalEncoding
+from torch_geometric.utils import scatter
+
+
+class PatchTransformerAggregation(Aggregation):
+    r"""Performs patch transformer aggregation in which the elements to
+    aggregate are processed by multi-head attention blocks across patches, as
+    described in the `"Simplifying Temporal Heterogeneous Network for
+    Continuous-Time Link Prediction"
+    <https://dl.acm.org/doi/pdf/10.1145/3583780.3615059>`_ paper.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        patch_size (int): Number of elements in a patch.
+        hidden_channels (int): Intermediate size of each sample.
+        num_transformer_blocks (int, optional): Number of transformer blocks
+            (default: :obj:`1`).
+        heads (int, optional): Number of multi-head-attentions.
+            (default: :obj:`1`)
+        dropout (float, optional): Dropout probability of attention weights.
+            (default: :obj:`0.0`)
+        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,
+        in_channels: int,
+        out_channels: int,
+        patch_size: int,
+        hidden_channels: int,
+        num_transformer_blocks: int = 1,
+        heads: int = 1,
+        dropout: float = 0.0,
+        aggr: Union[str, List[str]] = 'mean',
+        device: Optional[torch.device] = None,
+    ) -> None:
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.patch_size = patch_size
+        self.aggrs = [aggr] if isinstance(aggr, str) else aggr
+
+        assert len(self.aggrs) > 0
+        for aggr in self.aggrs:
+            assert aggr in ['sum', 'mean', 'min', 'max', 'var', 'std']
+
+        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, device=device)
+
+        self.blocks = torch.nn.ModuleList([
+            MultiheadAttentionBlock(
+                channels=hidden_channels,
+                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:
+        self.lin.reset_parameters()
+        self.pad_projector.reset_parameters()
+        self.pe.reset_parameters()
+        for block in self.blocks:
+            block.reset_parameters()
+        self.fc.reset_parameters()
+
+    @disable_dynamic_shapes(required_args=['dim_size', 'max_num_elements'])
+    def forward(
+        self,
+        x: Tensor,
+        index: Tensor,
+        ptr: Optional[Tensor] = None,
+        dim_size: Optional[int] = None,
+        dim: int = -2,
+        max_num_elements: Optional[int] = None,
+    ) -> Tensor:
+
+        if max_num_elements is None:
+            if ptr is not None:
+                count = ptr.diff()
+            else:
+                count = scatter(torch.ones_like(index), index, dim=0,
+                                dim_size=dim_size, reduce='sum')
+            max_num_elements = int(count.max()) + 1
+
+        # Set `max_num_elements` to a multiple of `patch_size`:
+        max_num_elements = (math.floor(max_num_elements / self.patch_size) *
+                            self.patch_size)
+
+        x = self.lin(x)
+
+        # TODO If groups are heavily unbalanced, this will create a lot of
+        # "empty" patches. Try to figure out a way to fix this.
+        # [batch_size, num_patches * patch_size, hidden_channels]
+        x, _ = self.to_dense_batch(x, index, ptr, dim_size, dim,
+                                   max_num_elements=max_num_elements)
+
+        # [batch_size, num_patches, patch_size * hidden_channels]
+        x = x.view(x.size(0), max_num_elements // self.patch_size,
+                   self.patch_size * x.size(-1))
+
+        # [batch_size, num_patches, hidden_channels]
+        x = self.pad_projector(x)
+
+        x = x + self.pe(torch.arange(x.size(1), device=x.device))
+
+        # [batch_size, num_patches, hidden_channels]
+        for block in self.blocks:
+            x = block(x, x)
+
+        # [batch_size, hidden_channels]
+        outs: List[Tensor] = []
+        for aggr in self.aggrs:
+            out = getattr(torch, aggr)(x, dim=1)
+            outs.append(out[0] if isinstance(out, tuple) else out)
+        out = torch.cat(outs, dim=1) if len(outs) > 1 else outs[0]
+
+        # [batch_size, out_channels]
+        return self.fc(out)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, patch_size={self.patch_size})')

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/collect.jinja

@@ -98,13 +98,16 @@ def {{collect_name}}(
 
 {%- if 'edge_weight' in collect_param_dict and
     collect_param_dict['edge_weight'].type_repr.endswith('Tensor') %}
-    assert edge_weight is not None
+    if torch.jit.is_scripting():
+        assert edge_weight is not None
 {%- elif 'edge_attr' in collect_param_dict and
     collect_param_dict['edge_attr'].type_repr.endswith('Tensor') %}
-    assert edge_attr is not None
+    if torch.jit.is_scripting():
+        assert edge_attr is not None
 {%- elif 'edge_type' in collect_param_dict and
     collect_param_dict['edge_type'].type_repr.endswith('Tensor') %}
-    assert edge_type is not None
+    if torch.jit.is_scripting():
+        assert edge_type is not None
 {%- endif %}
 
     # Collect user-defined arguments:

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