torch_geometric.nn.attention.SGFormerAttention

class SGFormerAttention(channels: int, heads: int = 1, head_channels: int = 64, qkv_bias: bool = False)[source]

Bases: Module

The simple global attention mechanism from the “SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations” paper.

Parameters:

  • channels (int) – Size of each input sample.

  • heads (int, optional) – Number of parallel attention heads. (default: 1.)

  • head_channels (int, optional) – Size of each attention head. (default: 64.)

  • qkv_bias (bool, optional) – If specified, add bias to query, key and value in the self attention. (default: False)

forward(x: Tensor, mask: Optional[Tensor] = None) → Tensor[source]

Forward pass.

Parameters:

  • x (torch.Tensor) – Node feature tensor \(\mathbf{X} \in \mathbb{R}^{B \times N \times F}\), with batch-size \(B\), (maximum) number of nodes \(N\) for each graph, and feature dimension \(F\).

  • mask (torch.Tensor, optional) – Mask matrix \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) indicating the valid nodes for each graph. (default: None)

Return type:

Tensor

reset_parameters()[source]