attention_blocks

Attention and transformer blocks used in graphnet models.

class graphnet.models.components.attention_blocks.DropPath(drop_prob)[source]

Bases: LightningModule

Drop paths (Stochastic Depth) per sample.

Construct DropPath.

Parameters:

drop_prob (float, default: 0.0) – Probability of dropping a path during training. If 0.0, no paths are dropped. Defaults to None.

forward(x, event_id, num_events)[source]

Forward pass.

Parameters:

  • x (Tensor) – Input tensor; one drop decision per row of the batch dimension.

  • event_id (Optional[Tensor], default: None) – Optional [N] event index of each token of a packed [1, N, D] input; decisions are then drawn per event.

  • num_events (Optional[int], default: None) – Number of events in event_id.

Return type:

Tensor

extra_repr()[source]

Return extra representation of the module.

Return type:

str

class graphnet.models.components.attention_blocks.Mlp(in_features, hidden_features, out_features, activation=<class 'torch.nn.modules.activation.GELU'>, dropout_prob)[source]

Bases: LightningModule

Multi-Layer Perceptron (MLP) module.

Construct Mlp.

Parameters:

  • in_features (int) – Number of input features.

  • hidden_features (Optional[int], default: None) – Number of hidden features. Defaults to None. If None, it is set to the value of in_features.

  • out_features (Optional[int], default: None) – Number of output features. Defaults to None. If None, it is set to the value of in_features.

  • activation (Module, default: <class 'torch.nn.modules.activation.GELU'>) – Activation layer. Defaults to nn.GELU.

  • dropout_prob (float, default: 0.0) – Dropout probability. Defaults to 0.0.

forward(x)[source]

Forward pass.

Return type:

Tensor

Parameters:

x (Tensor)

class graphnet.models.components.attention_blocks.Attention_rel(input_dim, num_heads, qkv_bias, qk_scale, attn_drop, proj_drop, attn_head_dim)[source]

Bases: LightningModule

Attention mechanism with relative position bias.

Construct ‘Attention_rel’.

Parameters:

  • input_dim (int) – Dimension of the input tensor.

  • num_heads (int, default: 8) – the number of attention heads to use (default: 8)

  • qkv_bias (bool, default: False) – whether to add bias to the query, key, and value projections. Defaults to False.

  • qk_scale (Optional[float], default: None) – a scaling factor that multiplies the dot product of query and key vectors. Defaults to None. If None, computed as :math: head_dim^(-1/2).

  • attn_drop (float, default: 0.0) – the dropout probability for the attention weights. Defaults to 0.0.

  • proj_drop (float, default: 0.0) – the dropout probability for the output of the attention module. Defaults to 0.0.

  • attn_head_dim (Optional[int], default: None) – the feature dimensionality of each attention head. Defaults to None. If None, computed as dim // num_heads.

forward(q, k, v, rel_pos_bias, key_padding_mask)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • q (Tensor)

  • k (Tensor)

  • v (Tensor)

  • rel_pos_bias (Tensor | None)

  • key_padding_mask (Tensor | None)

class graphnet.models.components.attention_blocks.Block_rel(input_dim, num_heads, mlp_ratio, qkv_bias, qk_scale, dropout, attn_drop, drop_path, init_values, activation=<class 'torch.nn.modules.activation.GELU'>, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>, attn_head_dim)[source]

Bases: LightningModule

Implementation of BEiTv2 Block.

Construct ‘Block_rel’.

Parameters:

  • input_dim (int) – Dimension of the input tensor.

  • num_heads (int) – Number of attention heads to use in the Attention_rel

  • layer.

  • mlp_ratio (float, default: 4.0) – Ratio of the hidden size of the feedforward network to the input size in the Mlp layer.

  • qkv_bias (bool, default: False) – Whether or not to include bias terms in the query, key, and value matrices in the Attention_rel layer.

  • qk_scale (Optional[float], default: None) – Scaling factor for the dot product of the query and key matrices in the Attention_rel layer.

  • dropout (float, default: 0.0) – Dropout probability to use in the Mlp layer.

  • attn_drop (float, default: 0.0) – Dropout probability to use in the Attention_rel layer.

  • drop_path (float, default: 0.0) – Probability of applying drop path regularization to the output of the layer.

  • init_values (Optional[float], default: None) – Initial value to use for the gamma_1 and gamma_2 parameters if not None.

  • activation (Module, default: <class 'torch.nn.modules.activation.GELU'>) – Activation function to use in the Mlp layer.

  • norm_layer (Module, default: <class 'torch.nn.modules.normalization.LayerNorm'>) – Normalization layer to use.

  • attn_head_dim (Optional[int], default: None) – Dimension of the attention head outputs in the Attention_rel layer.

forward(x, key_padding_mask, rel_pos_bias, kv)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • key_padding_mask (Tensor | None)

  • rel_pos_bias (Tensor | None)

  • kv (Tensor | None)

class graphnet.models.components.attention_blocks.Block(input_dim, num_heads, mlp_ratio, dropout, attn_drop, drop_path, init_values, activation=<class 'torch.nn.modules.activation.GELU'>, norm_layer=<class 'torch.nn.modules.normalization.LayerNorm'>)[source]

Bases: LightningModule

Transformer block.

Construct ‘Block’.

Parameters:

  • input_dim (int) – Dimension of the input tensor.

  • num_heads (int) – Number of attention heads to use in the MultiheadAttention layer.

  • mlp_ratio (float, default: 4.0) – Ratio of the hidden size of the feedforward network to the input size in the Mlp layer.

  • dropout (float, default: 0.0) – Dropout probability to use in the Mlp layer.

  • attn_drop (float, default: 0.0) – Dropout probability to use in the MultiheadAttention layer.

  • drop_path (float, default: 0.0) – Probability of applying drop path regularization to the output of the layer.

  • init_values (Optional[float], default: None) – Initial value to use for the gamma_1 and gamma_2 parameters if not None.

  • activation (Module, default: <class 'torch.nn.modules.activation.GELU'>) – Activation function to use in the Mlp layer.

  • norm_layer (Module, default: <class 'torch.nn.modules.normalization.LayerNorm'>) – Normalization layer to use.

forward(x, attn_mask, key_padding_mask)[source]

Forward pass.

Return type:

Tensor

Parameters:

  • x (Tensor)

  • attn_mask (Tensor | None)

  • key_padding_mask (Tensor | None)

class graphnet.models.components.attention_blocks.SwiGLU(dim, hidden_dim, dropout, bias)[source]

Bases: LightningModule

Feed-forward block with a SwiGLU gate and fused input projection.

Construct SwiGLU.

Parameters:

  • dim (int) – Input and output dimension.

  • hidden_dim (int) – Inner dimension of the gated projection.

  • dropout (float, default: 0.1) – Dropout applied to the block output.

  • bias (bool, default: True) – Whether the projections carry a bias.

forward(x)[source]

Apply the gated feed-forward transform.

Return type:

Tensor

Parameters:

x (Tensor)

class graphnet.models.components.attention_blocks.AttentionPool(dim)[source]

Bases: LightningModule

Pool a sequence to one vector, attending from a learned query.

Construct AttentionPool.

Parameters:

dim (int) – Feature dimension of the sequence and of the output.

forward(x, mask)[source]

Pool [B, S, D] to [B, D]; fully masked rows give zero.

Parameters:

  • x (Tensor) – [B, S, D] sequence.

  • mask (Optional[Tensor], default: None) – Optional bool [B, S]; True marks a valid element.

Return type:

Tensor

class graphnet.models.components.attention_blocks.RotaryTransformerBlock(d_model, nhead, dim_feedforward, dropout, layer_norm_eps, bias, ff_bias, qk_norm, drop_path_rate, layerscale_init)[source]

Bases: LightningModule

Pre-norm transformer block with rotary attention.

Uses RMSNorm, a fused QKV projection, per-head query/key norm, LayerScale and a SwiGLU feed-forward. There is no attention-matrix dropout, which flex_attention cannot express, so the packed and padded paths match.

Construct RotaryTransformerBlock.

Parameters:

  • d_model (int) – Token dimension.

  • nhead (int) – Number of heads; d_model must be divisible by it.

  • dim_feedforward (int) – Inner dimension of the SwiGLU network.

  • dropout (float, default: 0.1) – Dropout on the residual and feed-forward branches.

  • layer_norm_eps (float, default: 1e-05) – Epsilon of the RMSNorm layers.

  • bias (bool, default: True) – Whether the attention projections carry a bias.

  • ff_bias (bool, default: False) – Whether the feed-forward projections carry a bias.

  • qk_norm (bool, default: True) – Whether to RMS-normalize queries and keys per head.

  • drop_path_rate (float, default: 0.0) – Stochastic-depth rate of this block.

  • layerscale_init (float, default: 1e-05) – Initial value of the LayerScale parameters.

forward(x, rope_tables, attn_mask, block_mask, event_id, num_events)[source]

Apply the block.

Parameters:

  • x (Tensor) – [B, S, d_model] token features.

  • rope_tables (Tuple[Tensor, Tensor]) – (cos, sin) tables from RoPE4D.

  • attn_mask (Optional[Tensor], default: None) – Optional SDPA mask for the padded path.

  • block_mask (Any, default: None) – flex_attention BlockMask; selects the packed path.

  • event_id (Optional[Tensor], default: None) – [N] event index per token on the packed path.

  • num_events (Optional[int], default: None) – Number of events in event_id.

Return type:

Tensor

Returns:

[B, S, d_model] updated token features.