Encoder Architecture › Stacking Encoder Layers

class TransformerEncoder:
    """
    Complete transformer encoder.

    Stacks multiple encoder layers with optional final normalization.
    """

    def __init__(self, n_layers, d_model, n_heads, d_ff):
        """
        Args:
            n_layers: Number of encoder layers
            d_model: Model dimension
            n_heads: Number of attention heads
            d_ff: Feed-forward hidden dimension
        """
        self.n_layers = n_layers
        self.d_model = d_model

        # Stack of encoder layers
        self.layers = [
            EncoderLayer(d_model, n_heads, d_ff) for _ in range(n_layers)
        ]

        # Final layer normalization (used in pre-norm configuration)
        self.final_norm = LayerNorm(d_model)

    def __call__(self, x, mask=None, return_all_layers=False):
        """
        Forward pass through all encoder layers.

        Args:
            x: Input of shape (seq_len, d_model)
            mask: Optional padding mask
            return_all_layers: Whether to return intermediate representations

        Returns:
            Final output (and optionally all layer outputs)
        """
        all_outputs = [x]
        all_attentions = []

        for layer in self.layers:
            x, attn = layer(x, mask)
            all_outputs.append(x)
            all_attentions.append(attn)

        # Apply final normalization
        x = self.final_norm(x)

        if return_all_layers:
            return x, all_outputs, all_attentions
        return x


# Build a 6-layer encoder
n_layers = 6
encoder = TransformerEncoder(n_layers, d_model, n_heads, d_ff)

# Process a sequence
seq_len = 8
x = np.random.randn(seq_len, d_model)
output, all_outputs, all_attentions = encoder(x, return_all_layers=True)