torchcell.nn.masked_attention_block.NodeSelfAttention

class NodeSelfAttention(hidden_dim: int, num_heads: int = 8, dropout: float = 0.1, activation: Module = GELU(approximate='none'), mode: str = 'node', compile_block_mask: bool = False)

Bases: Module

Masked attention block with learned edge attribute projections.

Uses FlexAttention on GPU for efficiency when available.

This module supports: 1. Adjacency masks to control attention flow 2. Edge attributes projected through per-head MLPs to modulate attention scores

Parameters:

• hidden_dim – Dimensionality of the input features

• num_heads – Number of attention heads

• dropout – Dropout probability

• activation – Activation function for the MLP

• mode – Mode of attention (‘node’ for node-level attention)

• compile_block_mask – Whether to compile the block mask function

forward(x: Tensor, adj_mask: Tensor, edge_attr: Tensor | dict[tuple[int, int], float] | None = None, edge_index: Tensor | None = None) → Tensor

Forward pass with masked attention and edge attribute projections.

Parameters:

• x – Input tensor with shape [batch_size, seq_len, hidden_dim]

• adj_mask – Adjacency mask with shape [batch_size, seq_len, seq_len]

• edge_attr – Optional edge attributes as tensor [num_edges] or dictionary mapping (src, dst) tuples to scalar values

• edge_index – Optional edge indices as tensor with shape [2, num_edges] Required if edge_attr is a tensor

Returns:

Output tensor with same shape as input