torchcell.nn.self_attention_block.SelfAttentionBlock

class SelfAttentionBlock(hidden_dim: int, num_heads: int = 8, dropout: float = 0.1, activation: Module = GELU(approximate='none'))

Bases: Module

Self-Attention Block (SAB) with device-adaptive implementation.

Uses flex_attention on GPU and standard attention on CPU.

forward(x: Tensor) → Tensor

Forward pass for the SAB.

Parameters:

x – Input tensor of shape [batch_size, seq_len, hidden_dim]

Returns:

Output tensor of same shape as input