torch_geometric.nn.models.Polynormer

class Polynormer(in_channels: int, hidden_channels: int, out_channels: int, local_layers: int = 7, global_layers: int = 2, in_dropout: float = 0.15, dropout: float = 0.5, global_dropout: float = 0.5, heads: int = 1, beta: float = 0.9, qk_shared: bool = False, pre_ln: bool = False, post_bn: bool = True, local_attn: bool = False)

Bases: Module

The polynormer module from the “Polynormer: polynomial-expressive graph transformer in linear time” paper.

Parameters:

• in_channels (int) – Input channels.

• hidden_channels (int) – Hidden channels.

• out_channels (int) – Output channels.

• local_layers (int) – The number of local attention layers. (default: 7)

• global_layers (int) – The number of global attention layers. (default: 2)

• in_dropout (float) – Input dropout rate. (default: 0.15)

• dropout (float) – Dropout rate. (default: 0.5)

• global_dropout (float) – Global dropout rate. (default: 0.5)

• heads (int) – The number of heads. (default: 1)

• beta (float) – Aggregate type. (default: 0.9)

• qk_shared (bool optional) – Whether weight of query and key are shared. (default: True)

• pre_ln (bool) – Pre layer normalization. (default: False)

• post_bn (bool) – Post batch normalization. (default: True)

• local_attn (bool) – Whether use local attention. (default: False)

forward(x: Tensor, edge_index: Tensor, batch: Optional[Tensor]) → Tensor

Forward pass.

Parameters:

• x (torch.Tensor) – The input node features.

• edge_index (torch.Tensor or SparseTensor) – The edge indices.

• batch (torch.Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each element to a specific example.

Return type:

Tensor

reset_parameters() → None

Return type:

None