Post-Norm Transformer Block

Apply a post-norm Transformer block (the original design). With callables norm, attention and ff: first $h=\text{norm}(x+\text{attention}(x))h=\backslash text\{norm\}\backslash big(x+\backslash text\{attention\}(x)\backslash big)$, then return $\text{norm}(h+\text{ff}(h))\backslash text\{norm\}\backslash big(h+\backslash text\{ff\}(h)\backslash big)$.

Implement solve(x, norm, attention, ff).

Returns. Return an array with the shape of x.

Examples

Example 1

Input:

solve(np.array([1.0, 2.0]), lambda v: v - v.mean(), lambda v: v * 2, lambda v: v + 1)

Output:

[-3.0, 3.0]

Example 2

Input:

solve(np.array([1.0, 1.0]), lambda v: v, lambda v: v, lambda v: v * 0)

Output:

[2.0, 2.0]