One decoder layer (post-norm, 3 sublayers)

Background

A decoder layer chains three Add & Norm sublayers — masked self-attention, then cross-attention to the encoder, then the FFN (post-norm):

$$\begin{aligned} x_1 &= \text{LN}(x + \text{SelfAttn}(x)) \\ x_2 &= \text{LN}(x_1 + \text{CrossAttn}(x_1)) \\ \text{out} &= \text{LN}(x_2 + \text{FFN}(x_2)) \end{aligned}$$

Here each sublayer is given to you as a function; LayerNorm standardises each row (zero mean, unit variance — affine γ/β covered separately).

Problem statement

Implement decoder_layer(x, self_attn, cross_attn, ffn, eps=1e-5) applying the three post-norm sublayers in order.

Input

• x — (n, d) array, the layer input.
• self_attn, cross_attn, ffn — callables mapping an (n, d) array to an (n, d) array.
• eps — LayerNorm epsilon.

Output

Returns an np.ndarray of shape (n, d).

Examples

Example 1 — all sublayers output zeros

Input:  x = [[1, 2], [3, 4]], self_attn = cross_attn = ffn = (lambda h: zeros_like(h))
Output: [[-1, 1], [-1, 1]]

Explanation: each Add & Norm reduces to LayerNorm of the (unchanged) residual; standardising each row of [[1,2],[3,4]] gives [[-1,1],[-1,1]] (and re-normalising is idempotent).

Constraints

• Apply, in order: x1 = LN(x + self_attn(x)), x2 = LN(x1 + cross_attn(x1)), out = LN(x2 + ffn(x2)).
• LayerNorm per row: (r - mean) / sqrt(var + eps) over the last axis (population variance).
• Return shape (n, d).

Notes

• The order is fixed: self-attention sees the target so far, cross-attention injects the source, then the FFN processes per position. Each is wrapped in a residual + norm.