Complete Encoder Implementation

import numpy as np

class EncoderLayer:
    """Single encoder layer"""
    
    def __init__(self, d_model, num_heads, d_ff):
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff
        
        # Multi-head attention (simplified - would need full implementation)
        # self.attention = MultiHeadAttention(d_model, num_heads)
        
        # Feed-forward network
        # self.ffn = FeedForwardNetwork(d_model, d_ff)
        
        # Layer normalizations
        self.layer_norm1 = LayerNormalization(d_model)
        self.layer_norm2 = LayerNormalization(d_model)
    
    def forward(self, x, attention_fn, ffn_fn):
        """
        Forward pass through encoder layer
        
        Parameters:
        x: Input (batch, seq_len, d_model)
        attention_fn: Function for multi-head attention
        ffn_fn: Function for feed-forward network
        """
        # Sublayer 1: Multi-head self-attention with residual
        x_norm1 = self.layer_norm1.forward(x)
        attention_output = attention_fn(x_norm1)
        x = x + attention_output  # Residual connection
        
        # Sublayer 2: FFN with residual
        x_norm2 = self.layer_norm2.forward(x)
        ffn_output = ffn_fn(x_norm2)
        x = x + ffn_output  # Residual connection
        
        return x

class Encoder:
    """Complete encoder stack"""
    
    def __init__(self, num_layers, d_model, num_heads, d_ff):
        """
        Initialize encoder
        
        Parameters:
        num_layers: Number of encoder layers
        d_model: Model dimension
        num_heads: Number of attention heads
        d_ff: Feed-forward dimension
        """
        self.num_layers = num_layers
        self.layers = [
            EncoderLayer(d_model, num_heads, d_ff)
            for _ in range(num_layers)
        ]
    
    def forward(self, x, attention_fn, ffn_fn):
        """
        Forward pass through encoder stack
        
        Parameters:
        x: Input embeddings (batch, seq_len, d_model)
        attention_fn: Function for attention
        ffn_fn: Function for FFN
        """
        # Process through each layer sequentially
        for layer in self.layers:
            x = layer.forward(x, attention_fn, ffn_fn)
        
        return x

# Example usage
num_layers, d_model, num_heads, d_ff = 12, 512, 8, 2048
encoder = Encoder(num_layers, d_model, num_heads, d_ff)

# Input: (batch=2, seq_len=10, d_model=512)
x = np.random.randn(2, 10, 512)

# Forward pass (would need actual attention_fn and ffn_fn)
# output = encoder.forward(x, attention_fn, ffn_fn)
print(f"Encoder with {num_layers} layers created")
print(f"Input shape: {x.shape}")