RNN / LSTM

import torch
import torch.nn as nn

# ===== Tensor shapes =====
# input: (seq_len, batch, input_size) by default
# Or set batch_first=True for (batch, seq_len, input_size)

# ===== nn.LSTM =====
lstm = nn.LSTM(input_size=10, hidden_size=64, num_layers=2, batch_first=True, dropout=0.2)

batch, seq_len, input_size = 16, 100, 10
x = torch.randn(batch, seq_len, input_size)

# Hidden + cell initialised to zero by default
output, (h_n, c_n) = lstm(x)
# output: (batch, seq_len, hidden_size) — last layer hidden state at each timestep
# h_n: (num_layers, batch, hidden_size) — final hidden for each layer
# c_n: (num_layers, batch, hidden_size) — final cell

# ===== Classification on top =====
class Classifier(nn.Module):
    def __init__(self, vocab_size, embed_dim=64, hidden=128, num_classes=2):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden, batch_first=True, num_layers=1, bidirectional=True)
        self.fc = nn.Linear(hidden * 2, num_classes)   # *2 for bidirectional

    def forward(self, ids):
        x = self.emb(ids)                              # (B, T, E)
        out, (h, c) = self.lstm(x)
        # Use the last timestep:
        last = out[:, -1, :]
        return self.fc(last)

# ===== Variable-length sequences (padding + packing) =====
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence

# Pad batch of variable-length tensors:
batch = [torch.tensor([1,2,3]), torch.tensor([4,5,6,7,8]), torch.tensor([9])]
lengths = torch.tensor([len(s) for s in batch])
padded = pad_sequence(batch, batch_first=True, padding_value=0)

# Pack so the LSTM skips padding:
packed = pack_padded_sequence(padded, lengths, batch_first=True, enforce_sorted=False)
output_packed, (h_n, c_n) = lstm(packed)
# Unpack to a padded tensor:
output, lengths_back = pad_packed_sequence(output_packed, batch_first=True)

# ===== GRU (lighter than LSTM, often comparable) =====
gru = nn.GRU(input_size=10, hidden_size=64, batch_first=True)
out, h_n = gru(x)   # GRU has no cell state

# ===== Vanilla RNN (rarely used directly) =====
rnn = nn.RNN(input_size=10, hidden_size=64, batch_first=True, nonlinearity='tanh')

# ===== Training tips =====
# - Gradient clipping to avoid explosion:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# - Teacher forcing for sequence-to-sequence training (feed ground truth as next input)

# - Use the appropriate loss:
#   classification:  CrossEntropyLoss on logits
#   language model:  CrossEntropyLoss with reshape over (B*T, V)
#   regression:      MSELoss

# ===== When RNNs vs Transformers =====
# Transformers dominate for most NLP tasks in 2026.
# RNNs still useful for:
#   - Small models / edge devices
#   - Streaming inputs (one timestep at a time)
#   - Time series with strong recurrence assumptions

# ===== Patterns to internalise =====
# - batch_first=True everywhere for consistency
# - pack_padded_sequence for variable-length batches
# - Gradient clipping (max_norm 1.0 or 5.0)
# - Bidirectional for classification; unidirectional for generation

# ===== Pitfalls =====
# - Wrong shape order (seq_len vs batch) — set batch_first=True
# - Treating padding as real input -> packs help
# - Vanishing gradients in vanilla RNN over long sequences -> use LSTM/GRU
# - Not zeroing or detaching hidden state between batches when treating them as independent

import torch.nn as nn
lstm = nn.LSTM(input_size=64, hidden_size=128, num_layers=2, batch_first=True, dropout=0.2)