Pre-LN vs post-LN block

Background

GPT-2 switched the LayerNorm placement from the original Transformer:

• post-LN (original): y = LN(x + f(x)) — LN sits on the residual path.
• pre-LN (GPT-2): y = x + f(LN(x)) — LN sits inside the residual; the skip stays an exact identity, so gradients flow back unscaled.

Problem statement

Implement pre_and_post_ln(x, f, gamma, beta, eps=1e-5) returning both block outputs.

Input

• x — (T, C) input.
• f — sublayer callable mapping (T, C) → (T, C).
• gamma, beta — LayerNorm params, shape (C,).
• eps — LayerNorm epsilon.

Output

Returns (pre, post):

• pre = x + f(LN(x)),
• post = LN(x + f(x)),

where LN(h) normalises each row over features then applies γ/β.

Examples

Example 1 — zero sublayer isolates the ordering

Input:  x = [[1, 2, 3, 4]], f = (h -> zeros), gamma = [1,1,1,1], beta = [0,0,0,0]
Output: pre  = [[1, 2, 3, 4]]                      # x + 0
        post = [[-1.3416, -0.4472, 0.4472, 1.3416]] # LN(x)

Constraints

• pre = x + f(LN(x)); post = LN(x + f(x)).
• LN(h) = γ·(h-μ)/√(var+eps) + β, per-row over the feature axis (population variance).
• Return the two (T, C) arrays in that order.

Notes

• With f≡0, pre returns x untouched (identity skip) while post normalises the whole stream — a clean illustration of where the LN sits.