core/nn/transformer_decoder_block library

Pre-LayerNorm Transformer decoder block (seq2seq style).

Layout (three pre-norm sub-layers, standard for encoder-decoder Transformers):

h1 = x     + dropout(selfAttn(ln1(x), causalMask))
h2 = h1    + dropout(crossAttn(ln2(h1), memory))
out = h2   + dropout(ffn(ln3(h2)))

The self-attention is masked (causal) — the decoder can only look at earlier positions. The cross-attention is unmasked — the whole encoder output (memory) is visible at every decoder step.

Accepts either 2D ([Sd, embedDim] decoder, [Se, kvEmbedDim] memory) or 3D ([B, Sd, embedDim] and [B, Se, kvEmbedDim]).

Classes

TransformerDecoderBlock