Composable transformer block with configurable attention/mixing function.

Implements the standard pre-norm transformer block pattern:

norm -> attention_fn -> residual -> norm -> FFN -> residual

The caller provides the attention/mixing function as a callback, making this block reusable across GQA, FNet, Performer, Nystromformer, LinearTransformer, and any future attention variant.

Architecture

Input
  |
  +---> LayerNorm/RMSNorm -> attention_fn(x) ---+
  |                                              |
  +<----- Residual + Dropout <-------------------+
  |
  +---> LayerNorm/RMSNorm -> FFN(x) ------------+
  |                                              |
  +<----- Residual + Dropout <-------------------+
  |
Output

Usage

# Single block with custom attention
block = TransformerBlock.layer(input,
  attention_fn: fn x, name -> build_my_attention(x, name) end,
  hidden_size: 256,
  name: "block_1"
)

# Stack N blocks
output = TransformerBlock.stack(input, 4,
  attention_fn: fn x, name -> build_my_attention(x, name) end,
  hidden_size: 256,
  name: "transformer"
)

Design

Follows the callback pattern established by Edifice.SSM.Common.build_block/3, where the caller provides the core mixing/attention function and this module handles the surrounding structure (normalization, residuals, FFN, dropout).