RNN, Transformers & Attention

import torch.nn as nn

# Built-in LSTM
lstm = nn.LSTM(input_size=100, hidden_size=256, num_layers=2, 
               batch_first=True, bidirectional=True)
output, (h_n, c_n) = lstm(x)  # x shape: (batch, seq, feature)

# Manual LSTM cell (for understanding)
class LSTMCell(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.fc = nn.Linear(input_size + hidden_size, hidden_size * 4)
        
    def forward(self, x, h, c):
        gates = self.fc(torch.cat([x, h], dim=1))
        f, i, o, c_tilde = gates.chunk(4, dim=1)
        f, i, o = torch.sigmoid(f), torch.sigmoid(i), torch.sigmoid(o)
        c = f * c + i * torch.tanh(c_tilde)
        h = o * torch.tanh(c)
        return h, c

# PyTorch GRU
gru = nn.GRU(input_size, hidden_size, num_layers, batch_first=True)
output, h_n = gru(x)

# TensorFlow/Keras
tf.keras.layers.GRU(units=128, return_sequences=True)

import numpy as np

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V: (..., seq_len, d_k)
    mask: (..., seq_len, seq_len) optional
    """
    d_k = Q.shape[-1]
    scores = np.matmul(Q, K.transpose(0,1,3,2)) / np.sqrt(d_k)  # (..., seq_len, seq_len)
    
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)  # causal mask
    
    attention_weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    output = np.matmul(attention_weights, V)
    return output, attention_weights

import torch
import math

def positional_encoding(seq_len, d_model):
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                        -(math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe.unsqueeze(0)  # (1, seq_len, d_model)

import torch.nn as nn
import torch.nn.functional as F

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.activation = F.relu
        
    def forward(self, src, src_mask=None, src_key_padding_mask=None):
        # Self-attention block with residual + norm
        x = src
        attn_out, _ = self.self_attn(x, x, x, attn_mask=src_mask,
                                     key_padding_mask=src_key_padding_mask)
        x = x + self.dropout1(attn_out)
        x = self.norm1(x)
        
        # Feedforward block with residual + norm
        ff_out = self.linear2(self.dropout2(self.activation(self.linear1(x))))
        x = x + self.dropout2(ff_out)
        x = self.norm2(x)
        return x

class BahdanauAttention(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.W_a = nn.Linear(hidden_dim * 2, hidden_dim)  # [s; h]
        self.v_a = nn.Linear(hidden_dim, 1, bias=False)
    
    def forward(self, query, encoder_outputs):
        # query: decoder hidden (batch, hidden)
        # encoder_outputs: (batch, seq_len, hidden)
        seq_len = encoder_outputs.size(1)
        query = query.unsqueeze(1).repeat(1, seq_len, 1)  # (batch, seq_len, hidden)
        
        # Combine query and encoder outputs
        energy = torch.tanh(self.W_a(torch.cat((query, encoder_outputs), dim=2)))  # (batch, seq_len, hidden)
        scores = self.v_a(energy).squeeze(2)  # (batch, seq_len)
        attn_weights = torch.softmax(scores, dim=1)
        
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs).squeeze(1)
        return context, attn_weights

def luong_dot_attention(query, encoder_outputs):
    # query: (batch, 1, hidden)
    # encoder_outputs: (batch, seq_len, hidden)
    scores = torch.bmm(query, encoder_outputs.transpose(1, 2))  # (batch, 1, seq_len)
    attn_weights = torch.softmax(scores, dim=2)
    context = torch.bmm(attn_weights, encoder_outputs)
    return context, attn_weights

def scaled_dot_product_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    attention_weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, V)
    return output, attention_weights

def causal_mask(size):
    """Upper triangular matrix with zeros on diagonal and below, -inf above"""
    mask = torch.triu(torch.ones(size, size), diagonal=1).bool()
    return mask  # True where future tokens (to be masked)