torchcell.nn.masked_attention_block.MaskedAttentionBlock

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

Bases: Module

Memory-efficient Masked Attention Block that uses FlexAttention on GPU.

This module only handles adjacency masks. For edge attributes, use NodeSelfAttention.

Parameters:

• hidden_dim – Dimensionality of the input features

• num_heads – Number of attention heads

• dropout – Dropout probability

• activation – Activation function to use in the feed-forward network

• 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) → Tensor

Forward pass applying masked attention using adjacency mask.

Parameters:

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

• adj_mask – Adjacency mask tensor with shape [batch_size, seq_len, seq_len] True values indicate allowed attention connections

Returns:

Output tensor with same shape as input