Complete Mixer Block with Residuals

A complete MLP-Mixer block combines token mixing and channel mixing with residual connections:

$X_1 = X + \text{TokenMix}(X)$ $X_2 = X_1 + \text{ChannelMix}(X_1)$

Where:

• TokenMix: $W_{\text{token}} \cdot X$ (mix across N tokens)
• ChannelMix: MLP with GELU applied per-token

The residual connections ensure stable training and allow the network to learn identity mappings when beneficial.

Task: Implement a complete Mixer block with both mixing stages and residual connections.