{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "763d8283",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Module 13: Transformers - Complete Transformer Architecture\n",
    "\n",
    "Welcome to Module 13! You're about to build the complete transformer architecture that powers modern language models like GPT, Claude, and ChatGPT.\n",
    "\n",
    "## 🔗 Prerequisites & Progress\n",
    "**You've Built**: Tokenization, embeddings, attention mechanisms, and all foundational components\n",
    "**You'll Build**: TransformerBlock, complete GPT architecture, and autoregressive generation\n",
    "**You'll Enable**: Full language model training and text generation capabilities\n",
    "\n",
    "**Connection Map**:\n",
    "```\n",
    "Tokenization + Embeddings + Attention → Transformers → Language Generation\n",
    "(text→numbers)  (learnable vectors) (sequence modeling)  (complete models)\n",
    "```\n",
    "\n",
    "## Learning Objectives\n",
    "By the end of this module, you will:\n",
    "1. Implement complete TransformerBlock with attention, MLP, and layer normalization\n",
    "2. Build full GPT architecture with multiple transformer blocks\n",
    "3. Add autoregressive text generation capability\n",
    "4. Understand parameter scaling in large language models\n",
    "5. Test transformer components and generation pipeline\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0857efbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp models.transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b58c4de",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import numpy as np\n",
    "from tinytorch.core.tensor import Tensor\n",
    "from tinytorch.core.layers import Linear\n",
    "from tinytorch.core.attention import MultiHeadAttention\n",
    "from tinytorch.core.activations import GELU"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b35ba8b8",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 📦 Where This Code Lives in the Final Package\n",
    "\n",
    "**Learning Side:** You work in `modules/13_transformers/transformers_dev.py`\n",
    "**Building Side:** Code exports to `tinytorch.models.transformer`\n",
    "\n",
    "```python\n",
    "# How to use this module:\n",
    "from tinytorch.models.transformer import TransformerBlock, GPT, LayerNorm, MLP\n",
    "```\n",
    "\n",
    "**Why this matters:**\n",
    "- **Learning:** Complete transformer system showcasing how all components work together\n",
    "- **Production:** Matches PyTorch's transformer implementation with proper model organization\n",
    "- **Consistency:** All transformer components and generation logic in models.transformer\n",
    "- **Integration:** Demonstrates the power of modular design by combining all previous modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e36e4f2c",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "from typing import Optional, List\n",
    "\n",
    "# Import from previous modules - following proper dependency chain\n",
    "# Note: Actual imports happen in try/except blocks below with fallback implementations\n",
    "from tinytorch.core.tensor import Tensor\n",
    "from tinytorch.core.layers import Linear\n",
    "# MultiHeadAttention import happens in try/except below\n",
    "\n",
    "# For development, we'll use minimal implementations if imports fail\n",
    "try:\n",
    "    from tinytorch.core.tensor import Tensor\n",
    "except ImportError:\n",
    "    print(\"Warning: Using minimal Tensor implementation for development\")\n",
    "    class Tensor:\n",
    "        \"\"\"Minimal Tensor class for transformer development.\"\"\"\n",
    "        def __init__(self, data, requires_grad=False):\n",
    "            self.data = np.array(data)\n",
    "            self.shape = self.data.shape\n",
    "            self.size = self.data.size\n",
    "            self.requires_grad = requires_grad\n",
    "            self.grad = None\n",
    "\n",
    "        def __add__(self, other):\n",
    "            if isinstance(other, Tensor):\n",
    "                return Tensor(self.data + other.data)\n",
    "            return Tensor(self.data + other)\n",
    "\n",
    "        def __mul__(self, other):\n",
    "            if isinstance(other, Tensor):\n",
    "                return Tensor(self.data * other.data)\n",
    "            return Tensor(self.data * other)\n",
    "\n",
    "        def matmul(self, other):\n",
    "            return Tensor(np.dot(self.data, other.data))\n",
    "\n",
    "        def sum(self, axis=None, keepdims=False):\n",
    "            return Tensor(self.data.sum(axis=axis, keepdims=keepdims))\n",
    "\n",
    "        def mean(self, axis=None, keepdims=False):\n",
    "            return Tensor(self.data.mean(axis=axis, keepdims=keepdims))\n",
    "\n",
    "        def reshape(self, *shape):\n",
    "            return Tensor(self.data.reshape(shape))\n",
    "\n",
    "        def __repr__(self):\n",
    "            return f\"Tensor(data={self.data}, shape={self.shape})\"\n",
    "\n",
    "try:\n",
    "    from tinytorch.core.layers import Linear\n",
    "except ImportError:\n",
    "    class Linear:\n",
    "        \"\"\"Minimal Linear layer for development.\"\"\"\n",
    "        def __init__(self, in_features, out_features, bias=True):\n",
    "            std = math.sqrt(2.0 / (in_features + out_features))\n",
    "            self.weight = Tensor(np.random.normal(0, std, (in_features, out_features)))\n",
    "            self.bias = Tensor(np.zeros(out_features)) if bias else None\n",
    "\n",
    "        def forward(self, x):\n",
    "            output = x.matmul(self.weight)\n",
    "            if self.bias is not None:\n",
    "                output = output + self.bias\n",
    "            return output\n",
    "\n",
    "        def parameters(self):\n",
    "            params = [self.weight]\n",
    "            if self.bias is not None:\n",
    "                params.append(self.bias)\n",
    "            return params\n",
    "\n",
    "try:\n",
    "    from tinytorch.core.attention import MultiHeadAttention\n",
    "except ImportError:\n",
    "    class MultiHeadAttention:\n",
    "        \"\"\"Minimal MultiHeadAttention for development.\"\"\"\n",
    "        def __init__(self, embed_dim, num_heads):\n",
    "            assert embed_dim % num_heads == 0\n",
    "            self.embed_dim = embed_dim\n",
    "            self.num_heads = num_heads\n",
    "            self.head_dim = embed_dim // num_heads\n",
    "\n",
    "            self.q_proj = Linear(embed_dim, embed_dim)\n",
    "            self.k_proj = Linear(embed_dim, embed_dim)\n",
    "            self.v_proj = Linear(embed_dim, embed_dim)\n",
    "            self.out_proj = Linear(embed_dim, embed_dim)\n",
    "\n",
    "        def forward(self, query, key, value, mask=None):\n",
    "            batch_size, seq_len, embed_dim = query.shape\n",
    "\n",
    "            # Linear projections\n",
    "            Q = self.q_proj.forward(query)\n",
    "            K = self.k_proj.forward(key)\n",
    "            V = self.v_proj.forward(value)\n",
    "\n",
    "            # Reshape for multi-head attention\n",
    "            Q = Q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)\n",
    "            K = K.reshape(batch_size, seq_len, self.num_heads, self.head_dim)\n",
    "            V = V.reshape(batch_size, seq_len, self.num_heads, self.head_dim)\n",
    "\n",
    "            # Transpose to (batch_size, num_heads, seq_len, head_dim)\n",
    "            Q = Tensor(np.transpose(Q.data, (0, 2, 1, 3)))\n",
    "            K = Tensor(np.transpose(K.data, (0, 2, 1, 3)))\n",
    "            V = Tensor(np.transpose(V.data, (0, 2, 1, 3)))\n",
    "\n",
    "            # Scaled dot-product attention\n",
    "            scores = Tensor(np.matmul(Q.data, np.transpose(K.data, (0, 1, 3, 2))))\n",
    "            scores = scores * (1.0 / math.sqrt(self.head_dim))\n",
    "\n",
    "            # Apply causal mask for autoregressive generation\n",
    "            if mask is not None:\n",
    "                scores = Tensor(scores.data + mask.data)\n",
    "\n",
    "            # Softmax\n",
    "            attention_weights = self._softmax(scores)\n",
    "\n",
    "            # Apply attention to values\n",
    "            out = Tensor(np.matmul(attention_weights.data, V.data))\n",
    "\n",
    "            # Transpose back and reshape\n",
    "            out = Tensor(np.transpose(out.data, (0, 2, 1, 3)))\n",
    "            out = out.reshape(batch_size, seq_len, embed_dim)\n",
    "\n",
    "            # Final linear projection\n",
    "            return self.out_proj.forward(out)\n",
    "\n",
    "        def _softmax(self, x):\n",
    "            \"\"\"Numerically stable softmax.\"\"\"\n",
    "            exp_x = Tensor(np.exp(x.data - np.max(x.data, axis=-1, keepdims=True)))\n",
    "            return Tensor(exp_x.data / np.sum(exp_x.data, axis=-1, keepdims=True))\n",
    "\n",
    "        def parameters(self):\n",
    "            params = []\n",
    "            params.extend(self.q_proj.parameters())\n",
    "            params.extend(self.k_proj.parameters())\n",
    "            params.extend(self.v_proj.parameters())\n",
    "            params.extend(self.out_proj.parameters())\n",
    "            return params\n",
    "\n",
    "try:\n",
    "    from tinytorch.core.embeddings import Embedding\n",
    "except ImportError:\n",
    "    class Embedding:\n",
    "        \"\"\"Minimal Embedding layer for development.\"\"\"\n",
    "        def __init__(self, vocab_size, embed_dim):\n",
    "            self.vocab_size = vocab_size\n",
    "            self.embed_dim = embed_dim\n",
    "            self.weight = Tensor(np.random.normal(0, 0.02, (vocab_size, embed_dim)))\n",
    "\n",
    "        def forward(self, indices):\n",
    "            return Tensor(self.weight.data[indices.data.astype(int)])\n",
    "\n",
    "        def parameters(self):\n",
    "            return [self.weight]\n",
    "\n",
    "def gelu(x):\n",
    "    \"\"\"GELU activation function.\"\"\"\n",
    "    return Tensor(0.5 * x.data * (1 + np.tanh(np.sqrt(2 / np.pi) * (x.data + 0.044715 * x.data**3))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77ba5604",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 1. Introduction: What are Transformers?\n",
    "\n",
    "Transformers are the revolutionary architecture that powers modern AI language models like GPT, ChatGPT, and Claude. The key breakthrough is **self-attention**, which allows every token in a sequence to directly interact with every other token, creating rich contextual understanding.\n",
    "\n",
    "### The Transformer Revolution\n",
    "\n",
    "Before transformers, language models used RNNs or CNNs that processed text sequentially or locally. Transformers changed everything by processing all positions in parallel while maintaining global context.\n",
    "\n",
    "### Complete GPT Architecture Overview\n",
    "\n",
    "```\n",
    "┌─────────────────────────────────────────────────────────────────┐\n",
    "│  COMPLETE GPT ARCHITECTURE: From Text to Generation            │\n",
    "├─────────────────────────────────────────────────────────────────┤\n",
    "│                                                                 │\n",
    "│  INPUT: \"Hello world\"  →  Token IDs: [15496, 1917]             │\n",
    "│                                ↓                               │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │                EMBEDDING LAYER                            │  │\n",
    "│  │                                                           │  │\n",
    "│  │  ┌─────────────┐       ┌─────────────────────────────┐   │  │\n",
    "│  │  │Token Embed  │   +   │ Positional Embedding        │   │  │\n",
    "│  │  │15496→[0.1,  │       │ pos_0→[0.05, -0.02, ...]   │   │  │\n",
    "│  │  │     0.3,..]│       │ pos_1→[0.12,  0.08, ...]   │   │  │\n",
    "│  │  │1917→[0.2,   │       │                             │   │  │\n",
    "│  │  │    -0.1,..]│       │                             │   │  │\n",
    "│  │  └─────────────┘       └─────────────────────────────┘   │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                ↓                               │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │              TRANSFORMER BLOCK 1                         │  │\n",
    "│  │                                                           │  │\n",
    "│  │  x → LayerNorm → MultiHeadAttention → + x → result       │  │\n",
    "│  │  │                                      ↑               │  │\n",
    "│  │  │              residual connection     │               │  │\n",
    "│  │  └──────────────────────────────────────┘               │  │\n",
    "│  │  │                                                       │  │\n",
    "│  │  result → LayerNorm → MLP (Feed Forward) → + result      │  │\n",
    "│  │  │                                           ↑           │  │\n",
    "│  │  │                residual connection       │           │  │\n",
    "│  │  └───────────────────────────────────────────┘           │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                ↓                               │\n",
    "│              TRANSFORMER BLOCK 2 (same pattern)                │\n",
    "│                                ↓                               │\n",
    "│                      ... (more blocks) ...                     │\n",
    "│                                ↓                               │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │                   OUTPUT HEAD                             │  │\n",
    "│  │                                                           │  │\n",
    "│  │  final_hidden → LayerNorm → Linear(embed_dim, vocab_size) │  │\n",
    "│  │                              ↓                           │  │\n",
    "│  │               Vocabulary Logits: [0.1, 0.05, 0.8, ...]   │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                ↓                               │\n",
    "│  OUTPUT: Next Token Probabilities                              │\n",
    "│  \"Hello\" → 10%,  \"world\" → 5%,  \"!\" → 80%,  ...               │\n",
    "│                                                                 │\n",
    "└─────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "### Why Transformers Dominate\n",
    "\n",
    "**Parallel Processing**: Unlike RNNs that process tokens one by one, transformers process all positions simultaneously. This makes training much faster.\n",
    "\n",
    "**Global Context**: Every token can directly attend to every other token in the sequence, capturing long-range dependencies that RNNs struggle with.\n",
    "\n",
    "**Scalability**: Performance predictably improves with more parameters and data. This enabled the scaling laws that led to GPT-3, GPT-4, and beyond.\n",
    "\n",
    "**Residual Connections**: Allow training very deep networks (100+ layers) by providing gradient highways.\n",
    "\n",
    "### The Building Blocks We'll Implement\n",
    "\n",
    "1. **LayerNorm**: Stabilizes training by normalizing activations\n",
    "2. **Multi-Layer Perceptron (MLP)**: Provides non-linear transformation\n",
    "3. **TransformerBlock**: Combines attention + MLP with residuals\n",
    "4. **GPT**: Complete model with embeddings and generation capability"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4f69559",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 2. Foundations: Essential Transformer Mathematics\n",
    "\n",
    "### Layer Normalization: The Stability Engine\n",
    "\n",
    "Layer Normalization is crucial for training deep transformer networks. Unlike batch normalization (which normalizes across the batch), layer norm normalizes across the feature dimension for each individual sample.\n",
    "\n",
    "```\n",
    "Mathematical Formula:\n",
    "output = (x - μ) / σ * γ + β\n",
    "\n",
    "where:\n",
    "  μ = mean(x, axis=features)     # Mean across feature dimension\n",
    "  σ = sqrt(var(x) + ε)          # Standard deviation + small epsilon\n",
    "  γ = learnable scale parameter  # Initialized to 1.0\n",
    "  β = learnable shift parameter  # Initialized to 0.0\n",
    "```\n",
    "\n",
    "**Why Layer Norm Works:**\n",
    "- **Independence**: Each sample normalized independently (good for variable batch sizes)\n",
    "- **Stability**: Prevents internal covariate shift that breaks training\n",
    "- **Gradient Flow**: Helps gradients flow better through deep networks\n",
    "\n",
    "### Residual Connections: The Gradient Highway\n",
    "\n",
    "Residual connections are the secret to training deep networks. They create \"gradient highways\" that allow information to flow directly through the network.\n",
    "\n",
    "```\n",
    "┌─────────────────────────────────────────────────────────────────┐\n",
    "│  RESIDUAL CONNECTIONS: The Gradient Highway System             │\n",
    "├─────────────────────────────────────────────────────────────────┤\n",
    "│                                                                 │\n",
    "│  PRE-NORM ARCHITECTURE (Modern Standard):                      │\n",
    "│                                                                 │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │                ATTENTION SUB-LAYER                       │  │\n",
    "│  │                                                           │  │\n",
    "│  │  Input (x) ────┬─→ LayerNorm ─→ MultiHeadAttention ─┐    │  │\n",
    "│  │                │                                     │    │  │\n",
    "│  │                │         ┌─────────────────────────────┘    │  │\n",
    "│  │                │         ▼                               │  │\n",
    "│  │                └────→ ADD ─→ Output to next sub-layer    │  │\n",
    "│  │                      (x + attention_output)             │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                ↓                               │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │                   MLP SUB-LAYER                          │  │\n",
    "│  │                                                           │  │\n",
    "│  │  Input (x) ────┬─→ LayerNorm ─→ MLP (Feed Forward) ─┐    │  │\n",
    "│  │                │                                     │    │  │\n",
    "│  │                │         ┌─────────────────────────────┘    │  │\n",
    "│  │                │         ▼                               │  │\n",
    "│  │                └────→ ADD ─→ Final Output                │  │\n",
    "│  │                      (x + mlp_output)                   │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                                                 │\n",
    "│  KEY INSIGHT: Each sub-layer ADDS to the residual stream       │\n",
    "│  rather than replacing it, preserving information flow!        │\n",
    "│                                                                 │\n",
    "└─────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "**Gradient Flow Visualization:**\n",
    "```\n",
    "Backward Pass Without Residuals:    With Residuals:\n",
    "Loss                                Loss\n",
    " │ gradients get smaller             │ gradients stay strong\n",
    " ↓ at each layer                    ↓ via residual paths\n",
    "Layer N  ← tiny gradients          Layer N  ← strong gradients\n",
    " │                                  │     ↗ (direct path)\n",
    " ↓                                  ↓   ↗\n",
    "Layer 2  ← vanishing                Layer 2  ← strong gradients\n",
    " │                                  │     ↗\n",
    " ↓                                  ↓   ↗\n",
    "Layer 1  ← gone!                   Layer 1  ← strong gradients\n",
    "```\n",
    "\n",
    "### Feed-Forward Network (MLP): The Thinking Layer\n",
    "\n",
    "The MLP provides the actual \"thinking\" in each transformer block. It's a simple two-layer network with a specific expansion pattern.\n",
    "\n",
    "```\n",
    "MLP Architecture:\n",
    "Input (embed_dim) → Linear → GELU → Linear → Output (embed_dim)\n",
    "       512           2048      2048    512\n",
    "                   (4x expansion)\n",
    "\n",
    "Mathematical Formula:\n",
    "FFN(x) = Linear₂(GELU(Linear₁(x)))\n",
    "       = W₂ · GELU(W₁ · x + b₁) + b₂\n",
    "\n",
    "where:\n",
    "  W₁: (embed_dim, 4*embed_dim)  # Expansion matrix\n",
    "  W₂: (4*embed_dim, embed_dim)  # Contraction matrix\n",
    "  GELU: smooth activation function (better than ReLU for language)\n",
    "```\n",
    "\n",
    "**Why 4x Expansion?**\n",
    "- **Capacity**: More parameters = more representation power\n",
    "- **Non-linearity**: GELU activation creates complex transformations\n",
    "- **Information Bottleneck**: Forces the model to compress useful information\n",
    "\n",
    "### The Complete Transformer Block Data Flow\n",
    "\n",
    "```\n",
    "Input Tensor (batch, seq_len, embed_dim)\n",
    "          ↓\n",
    "    ┌─────────────────────────────────────┐\n",
    "    │ ATTENTION SUB-LAYER                 │\n",
    "    │                                     │\n",
    "    │ x₁ = LayerNorm(x₀)                  │\n",
    "    │ attention_out = MultiHeadAttn(x₁)   │\n",
    "    │ x₂ = x₀ + attention_out  (residual) │\n",
    "    └─────────────────────────────────────┘\n",
    "          ↓\n",
    "    ┌─────────────────────────────────────┐\n",
    "    │ MLP SUB-LAYER                       │\n",
    "    │                                     │\n",
    "    │ x₃ = LayerNorm(x₂)                  │\n",
    "    │ mlp_out = MLP(x₃)                   │\n",
    "    │ x₄ = x₂ + mlp_out    (residual)     │\n",
    "    └─────────────────────────────────────┘\n",
    "          ↓\n",
    "Output Tensor (batch, seq_len, embed_dim)\n",
    "```\n",
    "\n",
    "**Key Insight**: Each sub-layer (attention and MLP) gets a \"clean\" normalized input but adds its contribution to the residual stream. This creates a stable training dynamic."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a837896",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 3. Implementation: Building Transformer Components\n",
    "\n",
    "Now we'll implement each transformer component with a clear understanding of their role in the overall architecture. We'll follow the pattern: **Explanation → Implementation → Test** for each component.\n",
    "\n",
    "Each component serves a specific purpose:\n",
    "- **LayerNorm**: Stabilizes training and normalizes activations\n",
    "- **MLP**: Provides non-linear transformation and \"thinking\" capacity\n",
    "- **TransformerBlock**: Combines attention with MLP using residual connections\n",
    "- **GPT**: Complete autoregressive language model for text generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76f36a18",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### Understanding Layer Normalization\n",
    "\n",
    "Layer Normalization is the foundation of stable transformer training. Unlike batch normalization, it normalizes each sample independently across its feature dimensions.\n",
    "\n",
    "#### Why Layer Norm is Essential\n",
    "\n",
    "Without normalization, deep networks suffer from \"internal covariate shift\" - the distribution of inputs to each layer changes during training, making learning unstable.\n",
    "\n",
    "#### Layer Norm Visualization\n",
    "\n",
    "```\n",
    "┌─────────────────────────────────────────────────────────────────┐\n",
    "│  LAYER NORMALIZATION: Stabilizing Deep Networks               │\n",
    "├─────────────────────────────────────────────────────────────────┤\n",
    "│                                                                 │\n",
    "│  INPUT TENSOR: (batch=2, seq=3, features=4)                    │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │ Sample 1: [[1.0,  2.0,  3.0,  4.0],     ← Position 0     │  │\n",
    "│  │            [5.0,  6.0,  7.0,  8.0],     ← Position 1     │  │\n",
    "│  │            [9.0, 10.0, 11.0, 12.0]]     ← Position 2     │  │\n",
    "│  │                                                           │  │\n",
    "│  │ Sample 2: [[13., 14., 15., 16.],         ← Position 0     │  │\n",
    "│  │            [17., 18., 19., 20.],         ← Position 1     │  │\n",
    "│  │            [21., 22., 23., 24.]]         ← Position 2     │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                ↓                               │\n",
    "│           NORMALIZE ACROSS FEATURES (per position)             │\n",
    "│                                ↓                               │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │ AFTER NORMALIZATION: Each position → mean=0, std=1        │  │\n",
    "│  │                                                           │  │\n",
    "│  │ Sample 1: [[-1.34, -0.45,  0.45,  1.34],                 │  │\n",
    "│  │            [-1.34, -0.45,  0.45,  1.34],                 │  │\n",
    "│  │            [-1.34, -0.45,  0.45,  1.34]]                 │  │\n",
    "│  │                                                           │  │\n",
    "│  │ Sample 2: [[-1.34, -0.45,  0.45,  1.34],                 │  │\n",
    "│  │            [-1.34, -0.45,  0.45,  1.34],                 │  │\n",
    "│  │            [-1.34, -0.45,  0.45,  1.34]]                 │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                ↓                               │\n",
    "│            APPLY LEARNABLE PARAMETERS: γ * norm + β            │\n",
    "│                                ↓                               │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │ FINAL OUTPUT: Model can learn any desired distribution    │  │\n",
    "│  │ γ (scale) and β (shift) are learned during training       │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                                                 │\n",
    "│  KEY INSIGHT: Unlike batch norm, each sample normalized        │\n",
    "│  independently - perfect for variable-length sequences!        │\n",
    "│                                                                 │\n",
    "└─────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "#### Key Properties\n",
    "- **Per-sample normalization**: Each sequence position normalized independently\n",
    "- **Learnable parameters**: γ (scale) and β (shift) allow the model to recover any desired distribution\n",
    "- **Gradient friendly**: Helps gradients flow smoothly through deep networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6878edf0",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "layer-norm",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class LayerNorm:\n",
    "    \"\"\"\n",
    "    Layer Normalization for transformer blocks.\n",
    "\n",
    "    Normalizes across the feature dimension (last axis) for each sample independently,\n",
    "    unlike batch normalization which normalizes across the batch dimension.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, normalized_shape, eps=1e-5):\n",
    "        \"\"\"\n",
    "        Initialize LayerNorm with learnable parameters.\n",
    "\n",
    "        TODO: Set up normalization parameters\n",
    "\n",
    "        APPROACH:\n",
    "        1. Store the shape to normalize over (usually embed_dim)\n",
    "        2. Initialize learnable scale (gamma) and shift (beta) parameters\n",
    "        3. Set small epsilon for numerical stability\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> ln = LayerNorm(512)  # For 512-dimensional embeddings\n",
    "        >>> x = Tensor(np.random.randn(2, 10, 512))  # (batch, seq, features)\n",
    "        >>> normalized = ln.forward(x)\n",
    "        >>> # Each (2, 10) sample normalized independently across 512 features\n",
    "\n",
    "        HINTS:\n",
    "        - gamma should start at 1.0 (identity scaling)\n",
    "        - beta should start at 0.0 (no shift)\n",
    "        - eps prevents division by zero in variance calculation\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.normalized_shape = normalized_shape\n",
    "        self.eps = eps\n",
    "\n",
    "        # Learnable parameters: scale and shift\n",
    "        # CRITICAL: requires_grad=True so optimizer can train these!\n",
    "        self.gamma = Tensor(np.ones(normalized_shape), requires_grad=True)  # Scale parameter\n",
    "        self.beta = Tensor(np.zeros(normalized_shape), requires_grad=True)  # Shift parameter\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Apply layer normalization.\n",
    "\n",
    "        TODO: Implement layer normalization formula\n",
    "\n",
    "        APPROACH:\n",
    "        1. Compute mean and variance across the last dimension\n",
    "        2. Normalize: (x - mean) / sqrt(variance + eps)\n",
    "        3. Apply learnable scale and shift: gamma * normalized + beta\n",
    "\n",
    "        MATHEMATICAL FORMULA:\n",
    "        y = (x - μ) / σ * γ + β\n",
    "        where μ = mean(x), σ = sqrt(var(x) + ε)\n",
    "\n",
    "        HINT: Use keepdims=True to maintain tensor dimensions for broadcasting\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # CRITICAL: Use Tensor operations (not .data) to maintain gradient flow!\n",
    "        # Compute statistics across last dimension (features)\n",
    "        mean = x.mean(axis=-1, keepdims=True)\n",
    "\n",
    "        # Compute variance: E[(x - μ)²]\n",
    "        diff = x - mean  # Tensor subtraction maintains gradient\n",
    "        variance = (diff * diff).mean(axis=-1, keepdims=True)  # Tensor ops maintain gradient\n",
    "\n",
    "        # Normalize: (x - mean) / sqrt(variance + eps)\n",
    "        # Note: sqrt and division need to preserve gradient flow\n",
    "        std_data = np.sqrt(variance.data + self.eps)\n",
    "        normalized = diff * Tensor(1.0 / std_data)  # Scale by reciprocal to maintain gradient\n",
    "\n",
    "        # Apply learnable transformation\n",
    "        output = normalized * self.gamma + self.beta\n",
    "        return output\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def parameters(self):\n",
    "        \"\"\"Return learnable parameters.\"\"\"\n",
    "        return [self.gamma, self.beta]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b57594b0",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🔬 Unit Test: Layer Normalization\n",
    "This test validates our LayerNorm implementation works correctly.\n",
    "**What we're testing**: Normalization statistics and parameter learning\n",
    "**Why it matters**: Essential for transformer stability and training\n",
    "**Expected**: Mean ≈ 0, std ≈ 1 after normalization, learnable parameters work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f187ea71",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-layer-norm",
     "locked": true,
     "points": 10
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_layer_norm():\n",
    "    \"\"\"🔬 Test LayerNorm implementation.\"\"\"\n",
    "    print(\"🔬 Unit Test: Layer Normalization...\")\n",
    "\n",
    "    # Test basic normalization\n",
    "    ln = LayerNorm(4)\n",
    "    x = Tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]])  # (2, 4)\n",
    "\n",
    "    normalized = ln.forward(x)\n",
    "\n",
    "    # Check output shape\n",
    "    assert normalized.shape == (2, 4)\n",
    "\n",
    "    # Check normalization properties (approximately)\n",
    "    # For each sample, mean should be close to 0, std close to 1\n",
    "    for i in range(2):\n",
    "        sample_mean = np.mean(normalized.data[i])\n",
    "        sample_std = np.std(normalized.data[i])\n",
    "        assert abs(sample_mean) < 1e-5, f\"Mean should be ~0, got {sample_mean}\"\n",
    "        assert abs(sample_std - 1.0) < 1e-4, f\"Std should be ~1, got {sample_std}\"\n",
    "\n",
    "    # Test parameter shapes\n",
    "    params = ln.parameters()\n",
    "    assert len(params) == 2\n",
    "    assert params[0].shape == (4,)  # gamma\n",
    "    assert params[1].shape == (4,)  # beta\n",
    "\n",
    "    print(\"✅ LayerNorm works correctly!\")\n",
    "\n",
    "# Run test immediately when developing this module\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_layer_norm()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20fa9a45",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### Understanding the Multi-Layer Perceptron (MLP)\n",
    "\n",
    "The MLP is where the \"thinking\" happens in each transformer block. It's a simple feed-forward network that provides non-linear transformation capacity.\n",
    "\n",
    "#### The Role of MLP in Transformers\n",
    "\n",
    "While attention handles relationships between tokens, the MLP processes each position independently, adding computational depth and non-linearity.\n",
    "\n",
    "#### MLP Architecture and Information Flow\n",
    "\n",
    "```\n",
    "Information Flow Through MLP:\n",
    "\n",
    "Input: (batch, seq_len, embed_dim=512)\n",
    "         ↓\n",
    "┌─────────────────────────────────────────────┐\n",
    "│ Linear Layer 1: Expansion                   │\n",
    "│ Weight: (512, 2048)  Bias: (2048,)         │\n",
    "│ Output: (batch, seq_len, 2048)              │\n",
    "└─────────────────────────────────────────────┘\n",
    "         ↓\n",
    "┌─────────────────────────────────────────────┐\n",
    "│ GELU Activation                             │\n",
    "│ Smooth, differentiable activation           │\n",
    "│ Better than ReLU for language modeling     │\n",
    "└─────────────────────────────────────────────┘\n",
    "         ↓\n",
    "┌─────────────────────────────────────────────┐\n",
    "│ Linear Layer 2: Contraction                 │\n",
    "│ Weight: (2048, 512)  Bias: (512,)          │\n",
    "│ Output: (batch, seq_len, 512)               │\n",
    "└─────────────────────────────────────────────┘\n",
    "         ↓\n",
    "Output: (batch, seq_len, embed_dim=512)\n",
    "```\n",
    "\n",
    "#### Why 4x Expansion?\n",
    "\n",
    "```\n",
    "Parameter Count Analysis:\n",
    "\n",
    "Embed Dim: 512\n",
    "MLP Hidden: 2048 (4x expansion)\n",
    "\n",
    "Parameters:\n",
    "- Linear1: 512 × 2048 + 2048 = 1,050,624\n",
    "- Linear2: 2048 × 512 + 512 = 1,049,088\n",
    "- Total MLP: ~2.1M parameters\n",
    "\n",
    "For comparison:\n",
    "- Attention (same embed_dim): ~1.5M parameters\n",
    "- MLP has MORE parameters → more computational capacity\n",
    "```\n",
    "\n",
    "#### GELU vs ReLU\n",
    "\n",
    "```\n",
    "Activation Function Comparison:\n",
    "\n",
    "ReLU(x) = max(0, x)        # Hard cutoff at 0\n",
    "         ┌────\n",
    "         │\n",
    "    ─────┘\n",
    "         0\n",
    "\n",
    "GELU(x) ≈ x * Φ(x)         # Smooth, probabilistic\n",
    "         ╭────\n",
    "        ╱\n",
    "    ───╱\n",
    "      ╱\n",
    "     0\n",
    "\n",
    "GELU is smoother and provides better gradients for language modeling.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36edc347",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "mlp",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class MLP:\n",
    "    \"\"\"\n",
    "    Multi-Layer Perceptron (Feed-Forward Network) for transformer blocks.\n",
    "\n",
    "    Standard pattern: Linear -> GELU -> Linear with expansion ratio of 4:1.\n",
    "    This provides the non-linear transformation in each transformer block.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, embed_dim, hidden_dim=None, dropout_prob=0.1):\n",
    "        \"\"\"\n",
    "        Initialize MLP with two linear layers.\n",
    "\n",
    "        TODO: Set up the feed-forward network layers\n",
    "\n",
    "        APPROACH:\n",
    "        1. First layer expands from embed_dim to hidden_dim (usually 4x larger)\n",
    "        2. Second layer projects back to embed_dim\n",
    "        3. Use GELU activation (smoother than ReLU, preferred in transformers)\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> mlp = MLP(512)  # Will create 512 -> 2048 -> 512 network\n",
    "        >>> x = Tensor(np.random.randn(2, 10, 512))\n",
    "        >>> output = mlp.forward(x)\n",
    "        >>> assert output.shape == (2, 10, 512)\n",
    "\n",
    "        HINT: Standard transformer MLP uses 4x expansion (hidden_dim = 4 * embed_dim)\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        if hidden_dim is None:\n",
    "            hidden_dim = 4 * embed_dim  # Standard 4x expansion\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        # Two-layer feed-forward network\n",
    "        self.linear1 = Linear(embed_dim, hidden_dim)\n",
    "        self.linear2 = Linear(hidden_dim, embed_dim)\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Forward pass through MLP.\n",
    "\n",
    "        TODO: Implement the feed-forward computation\n",
    "\n",
    "        APPROACH:\n",
    "        1. First linear transformation: embed_dim -> hidden_dim\n",
    "        2. Apply GELU activation (smooth, differentiable)\n",
    "        3. Second linear transformation: hidden_dim -> embed_dim\n",
    "\n",
    "        COMPUTATION FLOW:\n",
    "        x -> Linear -> GELU -> Linear -> output\n",
    "\n",
    "        HINT: GELU activation is implemented above as a function\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # First linear layer with expansion\n",
    "        hidden = self.linear1.forward(x)\n",
    "\n",
    "        # GELU activation\n",
    "        hidden = gelu(hidden)\n",
    "\n",
    "        # Second linear layer back to original size\n",
    "        output = self.linear2.forward(hidden)\n",
    "\n",
    "        return output\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def parameters(self):\n",
    "        \"\"\"Return all learnable parameters.\"\"\"\n",
    "        params = []\n",
    "        params.extend(self.linear1.parameters())\n",
    "        params.extend(self.linear2.parameters())\n",
    "        return params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51e920ba",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🔬 Unit Test: MLP (Feed-Forward Network)\n",
    "This test validates our MLP implementation works correctly.\n",
    "**What we're testing**: Shape preservation and parameter counting\n",
    "**Why it matters**: MLP provides the non-linear transformation in transformers\n",
    "**Expected**: Input/output shapes match, correct parameter count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daa33cf0",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-mlp",
     "locked": true,
     "points": 10
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_mlp():\n",
    "    \"\"\"🔬 Test MLP implementation.\"\"\"\n",
    "    print(\"🔬 Unit Test: MLP (Feed-Forward Network)...\")\n",
    "\n",
    "    # Test MLP with standard 4x expansion\n",
    "    embed_dim = 64\n",
    "    mlp = MLP(embed_dim)\n",
    "\n",
    "    # Test forward pass\n",
    "    batch_size, seq_len = 2, 10\n",
    "    x = Tensor(np.random.randn(batch_size, seq_len, embed_dim))\n",
    "    output = mlp.forward(x)\n",
    "\n",
    "    # Check shape preservation\n",
    "    assert output.shape == (batch_size, seq_len, embed_dim)\n",
    "\n",
    "    # Check hidden dimension is 4x\n",
    "    assert mlp.hidden_dim == 4 * embed_dim\n",
    "\n",
    "    # Test parameter counting\n",
    "    params = mlp.parameters()\n",
    "    expected_params = 4  # 2 weights + 2 biases\n",
    "    assert len(params) == expected_params\n",
    "\n",
    "    # Test custom hidden dimension\n",
    "    custom_mlp = MLP(embed_dim, hidden_dim=128)\n",
    "    assert custom_mlp.hidden_dim == 128\n",
    "\n",
    "    print(\"✅ MLP works correctly!\")\n",
    "\n",
    "# Run test immediately when developing this module\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_mlp()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f7a5449",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### Understanding the Complete Transformer Block\n",
    "\n",
    "The TransformerBlock is the core building unit of GPT and other transformer models. It combines self-attention with feed-forward processing using a carefully designed residual architecture.\n",
    "\n",
    "#### Pre-Norm vs Post-Norm Architecture\n",
    "\n",
    "Modern transformers use \"pre-norm\" architecture where LayerNorm comes BEFORE the sub-layers, not after. This provides better training stability.\n",
    "\n",
    "```\n",
    "Pre-Norm Architecture (What We Implement):\n",
    "┌─────────────────────────────────────────────────────────┐\n",
    "│                     INPUT (x)                          │\n",
    "│                       │                                │\n",
    "│       ┌───────────────┴───────────────┐                │\n",
    "│       │                               │                │\n",
    "│       ▼                               │                │\n",
    "│  LayerNorm                            │                │\n",
    "│       │                               │                │\n",
    "│       ▼                               │                │\n",
    "│ MultiHeadAttention                    │                │\n",
    "│       │                               │                │\n",
    "│       └───────────────┬───────────────┘                │\n",
    "│                       │           (residual connection) │\n",
    "│                       ▼                                │\n",
    "│                  x + attention                         │\n",
    "│                       │                                │\n",
    "│       ┌───────────────┴───────────────┐                │\n",
    "│       │                               │                │\n",
    "│       ▼                               │                │\n",
    "│  LayerNorm                            │                │\n",
    "│       │                               │                │\n",
    "│       ▼                               │                │\n",
    "│      MLP                              │                │\n",
    "│       │                               │                │\n",
    "│       └───────────────┬───────────────┘                │\n",
    "│                       │           (residual connection) │\n",
    "│                       ▼                                │\n",
    "│                   x + mlp                              │\n",
    "│                       │                                │\n",
    "│                       ▼                                │\n",
    "│                    OUTPUT                              │\n",
    "└─────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "#### Why Pre-Norm Works Better\n",
    "\n",
    "**Training Stability**: LayerNorm before operations provides clean, normalized inputs to attention and MLP layers.\n",
    "\n",
    "**Gradient Flow**: Residual connections carry gradients directly from output to input, bypassing the normalized operations.\n",
    "\n",
    "**Deeper Networks**: Pre-norm enables training much deeper networks (100+ layers) compared to post-norm.\n",
    "\n",
    "#### Information Processing in Transformer Block\n",
    "\n",
    "```\n",
    "Step-by-Step Data Transformation:\n",
    "\n",
    "1. Input Processing:\n",
    "   x₀: (batch, seq_len, embed_dim) # Original input\n",
    "\n",
    "2. Attention Sub-layer:\n",
    "   x₁ = LayerNorm(x₀)               # Normalize input\n",
    "   attn_out = MultiHeadAttn(x₁)     # Self-attention\n",
    "   x₂ = x₀ + attn_out               # Residual connection\n",
    "\n",
    "3. MLP Sub-layer:\n",
    "   x₃ = LayerNorm(x₂)               # Normalize again\n",
    "   mlp_out = MLP(x₃)                # Feed-forward\n",
    "   x₄ = x₂ + mlp_out                # Final residual\n",
    "\n",
    "4. Output:\n",
    "   return x₄                        # Ready for next block\n",
    "```\n",
    "\n",
    "#### Residual Stream Concept\n",
    "\n",
    "Think of the residual connections as a \"stream\" that carries information through the network:\n",
    "\n",
    "```\n",
    "Residual Stream Flow:\n",
    "\n",
    "Layer 1: [original embeddings] ─┐\n",
    "                                 ├─→ + attention info ─┐\n",
    "Attention adds information ──────┘                      │\n",
    "                                                        ├─→ + MLP info ─┐\n",
    "MLP adds information ───────────────────────────────────┘               │\n",
    "                                                                        │\n",
    "Layer 2: carries accumulated information ──────────────────────────────┘\n",
    "```\n",
    "\n",
    "Each layer adds information to this stream rather than replacing it, creating a rich representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b54f39c",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "transformer-block",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class TransformerBlock:\n",
    "    \"\"\"\n",
    "    Complete Transformer Block with self-attention, MLP, and residual connections.\n",
    "\n",
    "    This is the core building block of GPT and other transformer models.\n",
    "    Each block processes the input sequence and passes it to the next block.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, embed_dim, num_heads, mlp_ratio=4, dropout_prob=0.1):\n",
    "        \"\"\"\n",
    "        Initialize a complete transformer block.\n",
    "\n",
    "        TODO: Set up all components of the transformer block\n",
    "\n",
    "        APPROACH:\n",
    "        1. Multi-head self-attention for sequence modeling\n",
    "        2. First layer normalization (pre-norm architecture)\n",
    "        3. MLP with specified expansion ratio\n",
    "        4. Second layer normalization\n",
    "\n",
    "        TRANSFORMER BLOCK ARCHITECTURE:\n",
    "        x → LayerNorm → MultiHeadAttention → + (residual) →\n",
    "            LayerNorm → MLP → + (residual) → output\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> block = TransformerBlock(embed_dim=512, num_heads=8)\n",
    "        >>> x = Tensor(np.random.randn(2, 10, 512))  # (batch, seq, embed)\n",
    "        >>> output = block.forward(x)\n",
    "        >>> assert output.shape == (2, 10, 512)\n",
    "\n",
    "        HINT: We use pre-norm architecture (LayerNorm before attention/MLP)\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "\n",
    "        # Multi-head self-attention\n",
    "        self.attention = MultiHeadAttention(embed_dim, num_heads)\n",
    "\n",
    "        # Layer normalizations (pre-norm architecture)\n",
    "        self.ln1 = LayerNorm(embed_dim)  # Before attention\n",
    "        self.ln2 = LayerNorm(embed_dim)  # Before MLP\n",
    "\n",
    "        # Feed-forward network\n",
    "        hidden_dim = int(embed_dim * mlp_ratio)\n",
    "        self.mlp = MLP(embed_dim, hidden_dim)\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        \"\"\"\n",
    "        Forward pass through transformer block.\n",
    "\n",
    "        TODO: Implement the complete transformer block computation\n",
    "\n",
    "        APPROACH:\n",
    "        1. Apply layer norm, then self-attention, then add residual\n",
    "        2. Apply layer norm, then MLP, then add residual\n",
    "        3. Return the transformed sequence\n",
    "\n",
    "        COMPUTATION FLOW:\n",
    "        x → ln1 → attention → + x → ln2 → mlp → + → output\n",
    "\n",
    "        RESIDUAL CONNECTIONS:\n",
    "        These are crucial for training deep networks - they allow gradients\n",
    "        to flow directly through the network during backpropagation.\n",
    "\n",
    "        HINT: Store intermediate results to add residual connections properly\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # First sub-layer: Multi-head self-attention with residual connection\n",
    "        # Pre-norm: LayerNorm before attention\n",
    "        normed1 = self.ln1.forward(x)\n",
    "        # Self-attention: query, key, value are all the same (normed1)\n",
    "        attention_out = self.attention.forward(normed1, normed1, normed1, mask)\n",
    "\n",
    "        # Residual connection\n",
    "        x = x + attention_out\n",
    "\n",
    "        # Second sub-layer: MLP with residual connection\n",
    "        # Pre-norm: LayerNorm before MLP\n",
    "        normed2 = self.ln2.forward(x)\n",
    "        mlp_out = self.mlp.forward(normed2)\n",
    "\n",
    "        # Residual connection\n",
    "        output = x + mlp_out\n",
    "\n",
    "        return output\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def parameters(self):\n",
    "        \"\"\"Return all learnable parameters.\"\"\"\n",
    "        params = []\n",
    "        params.extend(self.attention.parameters())\n",
    "        params.extend(self.ln1.parameters())\n",
    "        params.extend(self.ln2.parameters())\n",
    "        params.extend(self.mlp.parameters())\n",
    "        return params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78bc4bf0",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🔬 Unit Test: Transformer Block\n",
    "This test validates our complete TransformerBlock implementation.\n",
    "**What we're testing**: Shape preservation, residual connections, parameter counting\n",
    "**Why it matters**: This is the core component that will be stacked to create GPT\n",
    "**Expected**: Input/output shapes match, all components work together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f8fa7e8",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-transformer-block",
     "locked": true,
     "points": 15
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_transformer_block():\n",
    "    \"\"\"🔬 Test TransformerBlock implementation.\"\"\"\n",
    "    print(\"🔬 Unit Test: Transformer Block...\")\n",
    "\n",
    "    # Test transformer block\n",
    "    embed_dim = 64\n",
    "    num_heads = 4\n",
    "    block = TransformerBlock(embed_dim, num_heads)\n",
    "\n",
    "    # Test forward pass\n",
    "    batch_size, seq_len = 2, 8\n",
    "    x = Tensor(np.random.randn(batch_size, seq_len, embed_dim))\n",
    "    output = block.forward(x)\n",
    "\n",
    "    # Check shape preservation\n",
    "    assert output.shape == (batch_size, seq_len, embed_dim)\n",
    "\n",
    "    # Test with causal mask (for autoregressive generation)\n",
    "    mask = Tensor(np.triu(np.ones((seq_len, seq_len)) * -np.inf, k=1))\n",
    "    masked_output = block.forward(x, mask)\n",
    "    assert masked_output.shape == (batch_size, seq_len, embed_dim)\n",
    "\n",
    "    # Test parameter counting\n",
    "    params = block.parameters()\n",
    "    expected_components = 4  # attention, ln1, ln2, mlp parameters\n",
    "    assert len(params) > expected_components  # Should have parameters from all components\n",
    "\n",
    "    # Test different configurations\n",
    "    large_block = TransformerBlock(embed_dim=128, num_heads=8, mlp_ratio=2)\n",
    "    assert large_block.mlp.hidden_dim == 256  # 128 * 2\n",
    "\n",
    "    print(\"✅ TransformerBlock works correctly!\")\n",
    "\n",
    "# Run test immediately when developing this module\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_transformer_block()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d30f17d2",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### Understanding the Complete GPT Architecture\n",
    "\n",
    "GPT (Generative Pre-trained Transformer) is the complete language model that combines all our components into a text generation system. It's designed for **autoregressive** generation - predicting the next token based on all previous tokens.\n",
    "\n",
    "#### GPT's Autoregressive Nature\n",
    "\n",
    "GPT generates text one token at a time, using all previously generated tokens as context:\n",
    "\n",
    "```\n",
    "Autoregressive Generation Process:\n",
    "\n",
    "Step 1: \"The cat\" → model predicts → \"sat\"\n",
    "Step 2: \"The cat sat\" → model predicts → \"on\"\n",
    "Step 3: \"The cat sat on\" → model predicts → \"the\"\n",
    "Step 4: \"The cat sat on the\" → model predicts → \"mat\"\n",
    "\n",
    "Result: \"The cat sat on the mat\"\n",
    "```\n",
    "\n",
    "#### Complete GPT Architecture\n",
    "\n",
    "```\n",
    "┌─────────────────────────────────────────────────────────────┐\n",
    "│                     GPT ARCHITECTURE                       │\n",
    "│                                                             │\n",
    "│ Input: Token IDs [15496, 1917, ...]                        │\n",
    "│                    │                                        │\n",
    "│ ┌──────────────────┴──────────────────┐                    │\n",
    "│ │          EMBEDDING LAYER            │                    │\n",
    "│ │  ┌─────────────┐  ┌─────────────────┐│                    │\n",
    "│ │  │Token Embed  │+│Position Embed   ││                    │\n",
    "│ │  │vocab→vector ││ │sequence→vector  ││                    │\n",
    "│ │  └─────────────┘  └─────────────────┘│                    │\n",
    "│ └──────────────────┬──────────────────┘                    │\n",
    "│                    │                                        │\n",
    "│ ┌──────────────────┴──────────────────┐                    │\n",
    "│ │        TRANSFORMER BLOCK 1          │                    │\n",
    "│ │ ┌─────────┐  ┌─────────┐  ┌───────┐ │                    │\n",
    "│ │ │LayerNorm│→│Attention│→│  +x   │ │                    │\n",
    "│ │ └─────────┘  └─────────┘  └───┬───┘ │                    │\n",
    "│ │                               │     │                    │\n",
    "│ │ ┌─────────┐  ┌─────────┐  ┌───▼───┐ │                    │\n",
    "│ │ │LayerNorm│→│   MLP   │→│  +x   │ │                    │\n",
    "│ │ └─────────┘  └─────────┘  └───────┘ │                    │\n",
    "│ └──────────────────┬──────────────────┘                    │\n",
    "│                    │                                        │\n",
    "│         ... (more transformer blocks) ...                  │\n",
    "│                    │                                        │\n",
    "│ ┌──────────────────┴──────────────────┐                    │\n",
    "│ │         OUTPUT HEAD                 │                    │\n",
    "│ │ ┌─────────┐  ┌─────────────────────┐ │                    │\n",
    "│ │ │LayerNorm│→│Linear(embed→vocab)  │ │                    │\n",
    "│ │ └─────────┘  └─────────────────────┘ │                    │\n",
    "│ └──────────────────┬──────────────────┘                    │\n",
    "│                    │                                        │\n",
    "│ Output: Vocabulary Logits [0.1, 0.05, 0.8, ...]           │\n",
    "└─────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "#### Causal Masking for Autoregressive Training\n",
    "\n",
    "During training, GPT sees the entire sequence but must not \"cheat\" by looking at future tokens:\n",
    "\n",
    "```\n",
    "┌─────────────────────────────────────────────────────────────────┐\n",
    "│  CAUSAL MASKING: Preventing Future Information Leakage         │\n",
    "├─────────────────────────────────────────────────────────────────┤\n",
    "│                                                                 │\n",
    "│  SEQUENCE: [\"The\", \"cat\", \"sat\", \"on\"]                         │\n",
    "│  POSITIONS:   0      1      2     3                            │\n",
    "│                                                                 │\n",
    "│  ATTENTION MATRIX (what each position can see):                │\n",
    "│  ┌──────────────────────────────────────────────────────────┐   │\n",
    "│  │      Pos:  0   1   2   3                               │   │\n",
    "│  │  Pos 0:  [ ✓   ✗   ✗   ✗ ]  ← \"The\" only sees itself   │   │\n",
    "│  │  Pos 1:  [ ✓   ✓   ✗   ✗ ]  ← \"cat\" sees \"The\" + self  │   │\n",
    "│  │  Pos 2:  [ ✓   ✓   ✓   ✗ ]  ← \"sat\" sees all previous  │   │\n",
    "│  │  Pos 3:  [ ✓   ✓   ✓   ✓ ]  ← \"on\" sees everything     │   │\n",
    "│  └──────────────────────────────────────────────────────────┘   │\n",
    "│                                                                 │\n",
    "│  IMPLEMENTATION: Upper triangular matrix with -∞               │\n",
    "│  ┌──────────────────────────────────────────────────────────┐   │\n",
    "│  │ [[  0, -∞, -∞, -∞],                                     │   │\n",
    "│  │  [  0,   0, -∞, -∞],                                     │   │\n",
    "│  │  [  0,   0,   0, -∞],                                     │   │\n",
    "│  │  [  0,   0,   0,   0]]                                    │   │\n",
    "│  │                                                          │   │\n",
    "│  │ After softmax: -∞ becomes 0 probability                 │   │\n",
    "│  └──────────────────────────────────────────────────────────┘   │\n",
    "│                                                                 │\n",
    "│  WHY THIS WORKS: During training, model sees entire sequence   │\n",
    "│  but mask ensures position i only attends to positions ≤ i     │\n",
    "│                                                                 │\n",
    "└─────────────────────────────────────────────────────────────────┘\n",
    "```\n",
    "\n",
    "#### Generation Temperature Control\n",
    "\n",
    "Temperature controls the randomness of generation:\n",
    "\n",
    "```\n",
    "Temperature Effects:\n",
    "\n",
    "Original logits: [1.0, 2.0, 3.0]\n",
    "\n",
    "Temperature = 0.1 (Conservative):\n",
    "Scaled: [10.0, 20.0, 30.0] → Sharp distribution\n",
    "Probs: [0.00, 0.00, 1.00] → Always picks highest\n",
    "\n",
    "Temperature = 1.0 (Balanced):\n",
    "Scaled: [1.0, 2.0, 3.0] → Moderate distribution\n",
    "Probs: [0.09, 0.24, 0.67] → Weighted sampling\n",
    "\n",
    "Temperature = 2.0 (Creative):\n",
    "Scaled: [0.5, 1.0, 1.5] → Flatter distribution\n",
    "Probs: [0.18, 0.33, 0.49] → More random\n",
    "```\n",
    "\n",
    "#### Model Scaling and Parameters\n",
    "\n",
    "```\n",
    "GPT Model Size Scaling:\n",
    "\n",
    "Tiny GPT (our implementation):\n",
    "- embed_dim: 64, layers: 2, heads: 4\n",
    "- Parameters: ~50K\n",
    "- Use case: Learning and experimentation\n",
    "\n",
    "GPT-2 Small:\n",
    "- embed_dim: 768, layers: 12, heads: 12\n",
    "- Parameters: 117M\n",
    "- Use case: Basic text generation\n",
    "\n",
    "GPT-3:\n",
    "- embed_dim: 12,288, layers: 96, heads: 96\n",
    "- Parameters: 175B\n",
    "- Use case: Advanced language understanding\n",
    "\n",
    "GPT-4 (estimated):\n",
    "- embed_dim: ~16,384, layers: ~120, heads: ~128\n",
    "- Parameters: ~1.7T\n",
    "- Use case: Reasoning and multimodal tasks\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d86de25",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "gpt",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class GPT:\n",
    "    \"\"\"\n",
    "    Complete GPT (Generative Pre-trained Transformer) model.\n",
    "\n",
    "    This combines embeddings, positional encoding, multiple transformer blocks,\n",
    "    and a language modeling head for text generation.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, embed_dim, num_layers, num_heads, max_seq_len=1024):\n",
    "        \"\"\"\n",
    "        Initialize complete GPT model.\n",
    "\n",
    "        TODO: Set up all components of the GPT architecture\n",
    "\n",
    "        APPROACH:\n",
    "        1. Token embedding layer to convert tokens to vectors\n",
    "        2. Positional embedding to add position information\n",
    "        3. Stack of transformer blocks (the main computation)\n",
    "        4. Final layer norm and language modeling head\n",
    "\n",
    "        GPT ARCHITECTURE:\n",
    "        tokens → embedding → + pos_embedding →\n",
    "                transformer_blocks → layer_norm → lm_head → logits\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> model = GPT(vocab_size=1000, embed_dim=256, num_layers=6, num_heads=8)\n",
    "        >>> tokens = Tensor(np.random.randint(0, 1000, (2, 10)))  # (batch, seq)\n",
    "        >>> logits = model.forward(tokens)\n",
    "        >>> assert logits.shape == (2, 10, 1000)  # (batch, seq, vocab)\n",
    "\n",
    "        HINTS:\n",
    "        - Positional embeddings are learned, not fixed sinusoidal\n",
    "        - Final layer norm stabilizes training\n",
    "        - Language modeling head shares weights with token embedding (tie_weights)\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_layers = num_layers\n",
    "        self.num_heads = num_heads\n",
    "        self.max_seq_len = max_seq_len\n",
    "\n",
    "        # Token and positional embeddings\n",
    "        self.token_embedding = Embedding(vocab_size, embed_dim)\n",
    "        self.position_embedding = Embedding(max_seq_len, embed_dim)\n",
    "\n",
    "        # Stack of transformer blocks\n",
    "        self.blocks = []\n",
    "        for _ in range(num_layers):\n",
    "            block = TransformerBlock(embed_dim, num_heads)\n",
    "            self.blocks.append(block)\n",
    "\n",
    "        # Final layer normalization\n",
    "        self.ln_f = LayerNorm(embed_dim)\n",
    "\n",
    "        # Language modeling head (projects to vocabulary)\n",
    "        self.lm_head = Linear(embed_dim, vocab_size, bias=False)\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def forward(self, tokens):\n",
    "        \"\"\"\n",
    "        Forward pass through GPT model.\n",
    "\n",
    "        TODO: Implement the complete GPT forward pass\n",
    "\n",
    "        APPROACH:\n",
    "        1. Get token embeddings and positional embeddings\n",
    "        2. Add them together (broadcasting handles different shapes)\n",
    "        3. Pass through all transformer blocks sequentially\n",
    "        4. Apply final layer norm and language modeling head\n",
    "\n",
    "        COMPUTATION FLOW:\n",
    "        tokens → embed + pos_embed → blocks → ln_f → lm_head → logits\n",
    "\n",
    "        CAUSAL MASKING:\n",
    "        For autoregressive generation, we need to prevent tokens from\n",
    "        seeing future tokens. This is handled by the attention mask.\n",
    "\n",
    "        HINT: Create position indices as range(seq_len) for positional embedding\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        batch_size, seq_len = tokens.shape\n",
    "\n",
    "        # Token embeddings\n",
    "        token_emb = self.token_embedding.forward(tokens)\n",
    "\n",
    "        # Positional embeddings\n",
    "        positions = Tensor(np.arange(seq_len).reshape(1, seq_len))\n",
    "        pos_emb = self.position_embedding.forward(positions)\n",
    "\n",
    "        # Combine embeddings\n",
    "        x = token_emb + pos_emb\n",
    "\n",
    "        # Create causal mask for autoregressive generation\n",
    "        mask = self._create_causal_mask(seq_len)\n",
    "\n",
    "        # Pass through transformer blocks\n",
    "        for block in self.blocks:\n",
    "            x = block.forward(x, mask)\n",
    "\n",
    "        # Final layer normalization\n",
    "        x = self.ln_f.forward(x)\n",
    "\n",
    "        # Language modeling head\n",
    "        logits = self.lm_head.forward(x)\n",
    "\n",
    "        return logits\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def _create_causal_mask(self, seq_len):\n",
    "        \"\"\"Create causal mask to prevent attending to future positions.\"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Upper triangular matrix filled with -inf\n",
    "        mask = np.triu(np.ones((seq_len, seq_len)) * -np.inf, k=1)\n",
    "        return Tensor(mask)\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def generate(self, prompt_tokens, max_new_tokens=50, temperature=1.0):\n",
    "        \"\"\"\n",
    "        Generate text autoregressively.\n",
    "\n",
    "        TODO: Implement autoregressive text generation\n",
    "\n",
    "        APPROACH:\n",
    "        1. Start with prompt tokens\n",
    "        2. For each new position:\n",
    "           - Run forward pass to get logits\n",
    "           - Sample next token from logits\n",
    "           - Append to sequence\n",
    "        3. Return generated sequence\n",
    "\n",
    "        AUTOREGRESSIVE GENERATION:\n",
    "        At each step, the model predicts the next token based on all\n",
    "        previous tokens. This is how GPT generates coherent text.\n",
    "\n",
    "        EXAMPLE:\n",
    "        >>> model = GPT(vocab_size=100, embed_dim=64, num_layers=2, num_heads=4)\n",
    "        >>> prompt = Tensor([[1, 2, 3]])  # Some token sequence\n",
    "        >>> generated = model.generate(prompt, max_new_tokens=5)\n",
    "        >>> assert generated.shape[1] == 3 + 5  # original + new tokens\n",
    "\n",
    "        HINT: Use np.random.choice with temperature for sampling\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        current_tokens = Tensor(prompt_tokens.data.copy())\n",
    "\n",
    "        for _ in range(max_new_tokens):\n",
    "            # Get logits for current sequence\n",
    "            logits = self.forward(current_tokens)\n",
    "\n",
    "            # Get logits for last position (next token prediction)\n",
    "            last_logits = logits.data[:, -1, :]  # (batch_size, vocab_size)\n",
    "\n",
    "            # Apply temperature scaling\n",
    "            scaled_logits = last_logits / temperature\n",
    "\n",
    "            # Convert to probabilities (softmax)\n",
    "            exp_logits = np.exp(scaled_logits - np.max(scaled_logits, axis=-1, keepdims=True))\n",
    "            probs = exp_logits / np.sum(exp_logits, axis=-1, keepdims=True)\n",
    "\n",
    "            # Sample next token\n",
    "            next_token = np.array([[np.random.choice(self.vocab_size, p=probs[0])]])\n",
    "\n",
    "            # Append to sequence\n",
    "            current_tokens = Tensor(np.concatenate([current_tokens.data, next_token], axis=1))\n",
    "\n",
    "        return current_tokens\n",
    "        ### END SOLUTION\n",
    "\n",
    "    def parameters(self):\n",
    "        \"\"\"Return all learnable parameters.\"\"\"\n",
    "        params = []\n",
    "        params.extend(self.token_embedding.parameters())\n",
    "        params.extend(self.position_embedding.parameters())\n",
    "\n",
    "        for block in self.blocks:\n",
    "            params.extend(block.parameters())\n",
    "\n",
    "        params.extend(self.ln_f.parameters())\n",
    "        params.extend(self.lm_head.parameters())\n",
    "\n",
    "        return params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6994ec05",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🔬 Unit Test: GPT Model\n",
    "This test validates our complete GPT implementation.\n",
    "**What we're testing**: Model forward pass, shape consistency, generation capability\n",
    "**Why it matters**: This is the complete language model that ties everything together\n",
    "**Expected**: Correct output shapes, generation works, parameter counting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "377dc692",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-gpt",
     "locked": true,
     "points": 20
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_gpt():\n",
    "    \"\"\"🔬 Test GPT model implementation.\"\"\"\n",
    "    print(\"🔬 Unit Test: GPT Model...\")\n",
    "\n",
    "    # Test small GPT model\n",
    "    vocab_size = 100\n",
    "    embed_dim = 64\n",
    "    num_layers = 2\n",
    "    num_heads = 4\n",
    "\n",
    "    model = GPT(vocab_size, embed_dim, num_layers, num_heads)\n",
    "\n",
    "    # Test forward pass\n",
    "    batch_size, seq_len = 2, 8\n",
    "    tokens = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))\n",
    "    logits = model.forward(tokens)\n",
    "\n",
    "    # Check output shape\n",
    "    expected_shape = (batch_size, seq_len, vocab_size)\n",
    "    assert logits.shape == expected_shape\n",
    "\n",
    "    # Test generation\n",
    "    prompt = Tensor(np.random.randint(0, vocab_size, (1, 5)))\n",
    "    generated = model.generate(prompt, max_new_tokens=3)\n",
    "\n",
    "    # Check generation shape\n",
    "    assert generated.shape == (1, 8)  # 5 prompt + 3 new tokens\n",
    "\n",
    "    # Test parameter counting\n",
    "    params = model.parameters()\n",
    "    assert len(params) > 10  # Should have many parameters from all components\n",
    "\n",
    "    # Test different model sizes\n",
    "    larger_model = GPT(vocab_size=200, embed_dim=128, num_layers=4, num_heads=8)\n",
    "    test_tokens = Tensor(np.random.randint(0, 200, (1, 10)))\n",
    "    larger_logits = larger_model.forward(test_tokens)\n",
    "    assert larger_logits.shape == (1, 10, 200)\n",
    "\n",
    "    print(\"✅ GPT model works correctly!\")\n",
    "\n",
    "# Run test immediately when developing this module\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_gpt()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66fa0b98",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 4. Integration: Complete Transformer Workflow\n",
    "\n",
    "Now that we've built all the components, let's see how they work together in a complete language modeling pipeline. This demonstrates the full power of the transformer architecture.\n",
    "\n",
    "### The Language Modeling Pipeline\n",
    "\n",
    "```\n",
    "Complete Workflow Visualization:\n",
    "\n",
    "1. Text Input:\n",
    "   \"hello world\" → Tokenization → [15496, 1917]\n",
    "\n",
    "2. Model Processing:\n",
    "   [15496, 1917]\n",
    "        ↓ Token Embedding\n",
    "   [[0.1, 0.5, ...], [0.3, -0.2, ...]]  # Vector representations\n",
    "        ↓ + Position Embedding\n",
    "   [[0.2, 0.7, ...], [0.1, -0.4, ...]]  # With position info\n",
    "        ↓ Transformer Block 1\n",
    "   [[0.3, 0.2, ...], [0.5, -0.1, ...]]  # After attention + MLP\n",
    "        ↓ Transformer Block 2\n",
    "   [[0.1, 0.9, ...], [0.7, 0.3, ...]]   # Further processed\n",
    "        ↓ Final LayerNorm + LM Head\n",
    "   [[0.1, 0.05, 0.8, ...], [...]]       # Probability over vocab\n",
    "\n",
    "3. Generation:\n",
    "   Model predicts next token: \"!\" (token 33)\n",
    "   New sequence: \"hello world!\"\n",
    "```\n",
    "\n",
    "This integration demo will show:\n",
    "- **Character-level tokenization** for simplicity\n",
    "- **Forward pass** through all components\n",
    "- **Autoregressive generation** in action\n",
    "- **Temperature effects** on creativity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6381a082",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "integration-demo",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def demonstrate_transformer_integration():\n",
    "    \"\"\"\n",
    "    Demonstrate complete transformer pipeline.\n",
    "\n",
    "    This simulates training a small language model on a simple vocabulary.\n",
    "    \"\"\"\n",
    "    print(\"🔗 Integration Demo: Complete Language Model Pipeline\")\n",
    "    print(\"Building a mini-GPT for character-level text generation\")\n",
    "\n",
    "    # Create a small vocabulary (character-level)\n",
    "    vocab = list(\"abcdefghijklmnopqrstuvwxyz .\")\n",
    "    vocab_size = len(vocab)\n",
    "    char_to_idx = {char: i for i, char in enumerate(vocab)}\n",
    "    idx_to_char = {i: char for i, char in enumerate(vocab)}\n",
    "\n",
    "    print(f\"Vocabulary size: {vocab_size}\")\n",
    "    print(f\"Characters: {''.join(vocab)}\")\n",
    "\n",
    "    # Create model\n",
    "    model = GPT(\n",
    "        vocab_size=vocab_size,\n",
    "        embed_dim=64,\n",
    "        num_layers=2,\n",
    "        num_heads=4,\n",
    "        max_seq_len=32\n",
    "    )\n",
    "\n",
    "    # Sample text encoding\n",
    "    text = \"hello world.\"\n",
    "    tokens = [char_to_idx[char] for char in text]\n",
    "    input_tokens = Tensor(np.array([tokens]))\n",
    "\n",
    "    print(f\"\\nOriginal text: '{text}'\")\n",
    "    print(f\"Tokenized: {tokens}\")\n",
    "    print(f\"Input shape: {input_tokens.shape}\")\n",
    "\n",
    "    # Forward pass\n",
    "    logits = model.forward(input_tokens)\n",
    "    print(f\"Output logits shape: {logits.shape}\")\n",
    "    print(f\"Each position predicts next token from {vocab_size} possibilities\")\n",
    "\n",
    "    # Generation demo\n",
    "    prompt_text = \"hello\"\n",
    "    prompt_tokens = [char_to_idx[char] for char in prompt_text]\n",
    "    prompt = Tensor(np.array([prompt_tokens]))\n",
    "\n",
    "    print(f\"\\nGeneration demo:\")\n",
    "    print(f\"Prompt: '{prompt_text}'\")\n",
    "\n",
    "    generated = model.generate(prompt, max_new_tokens=8, temperature=1.0)\n",
    "    generated_text = ''.join([idx_to_char[idx] for idx in generated.data[0]])\n",
    "\n",
    "    print(f\"Generated: '{generated_text}'\")\n",
    "    print(\"(Note: Untrained model produces random text)\")\n",
    "\n",
    "    return model\n",
    "\n",
    "demonstrate_transformer_integration()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "540a7b4d",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 5. Systems Analysis: Parameter Scaling and Memory\n",
    "\n",
    "Transformer models scale dramatically with size, leading to both opportunities and challenges. Let's analyze the computational and memory requirements to understand why training large language models requires massive infrastructure.\n",
    "\n",
    "### The Scaling Laws Revolution\n",
    "\n",
    "One of the key discoveries in modern AI is that transformer performance follows predictable scaling laws:\n",
    "\n",
    "```\n",
    "Scaling Laws Pattern:\n",
    "Performance ∝ Parameters^α × Data^β × Compute^γ\n",
    "\n",
    "where α ≈ 0.7, β ≈ 0.8, γ ≈ 0.5\n",
    "\n",
    "This means:\n",
    "- 10× more parameters → ~5× better performance\n",
    "- 10× more data → ~6× better performance\n",
    "- 10× more compute → ~3× better performance\n",
    "```\n",
    "\n",
    "### Memory Scaling Analysis\n",
    "\n",
    "Memory requirements grow in different ways for different components:\n",
    "\n",
    "```\n",
    "Memory Scaling by Component:\n",
    "\n",
    "1. Parameter Memory (Linear with model size):\n",
    "   - Embeddings: vocab_size × embed_dim\n",
    "   - Transformer blocks: ~4 × embed_dim²\n",
    "   - Total: O(embed_dim²)\n",
    "\n",
    "2. Attention Memory (Quadratic with sequence length):\n",
    "   - Attention matrices: batch × heads × seq_len²\n",
    "   - This is why long context is expensive!\n",
    "   - Total: O(seq_len²)\n",
    "\n",
    "3. Activation Memory (Linear with batch size):\n",
    "   - Forward pass activations for backprop\n",
    "   - Scales with: batch × seq_len × embed_dim\n",
    "   - Total: O(batch_size)\n",
    "```\n",
    "\n",
    "### The Attention Memory Wall\n",
    "\n",
    "```\n",
    "┌─────────────────────────────────────────────────────────────────┐\n",
    "│  ATTENTION MEMORY WALL: Why Long Context is Expensive          │\n",
    "├─────────────────────────────────────────────────────────────────┤\n",
    "│                                                                 │\n",
    "│  MEMORY USAGE BY SEQUENCE LENGTH (Quadratic Growth):           │\n",
    "│                                                                 │\n",
    "│  1K tokens:   [▓] 16 MB        ← Manageable                    │\n",
    "│  2K tokens:   [▓▓▓▓] 64 MB     ← 4× memory (quadratic!)        │\n",
    "│  4K tokens:   [▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓▓] 256 MB   ← 16× memory         │\n",
    "│  8K tokens:   [████████████████████████████████] 1 GB          │\n",
    "│  16K tokens:  ████████████████████████████████████████████████████████████████ 4 GB │\n",
    "│  32K tokens:  ████████████████████████████████████████████████████████████████████████████████████████████████████████████████ 16 GB │\n",
    "│                                                                 │\n",
    "│  REAL-WORLD CONTEXT LIMITS:                                    │\n",
    "│  ┌───────────────────────────────────────────────────────────┐  │\n",
    "│  │ GPT-3:     2K tokens  (limited by memory)                │  │\n",
    "│  │ GPT-4:     8K tokens  (32K with optimizations)           │  │\n",
    "│  │ Claude-3:  200K tokens (special techniques required!)    │  │\n",
    "│  │ GPT-4o:    128K tokens (efficient attention)             │  │\n",
    "│  └───────────────────────────────────────────────────────────┘  │\n",
    "│                                                                 │\n",
    "│  MATHEMATICAL SCALING:                                         │\n",
    "│  Memory = batch_size × num_heads × seq_len² × 4 bytes          │\n",
    "│                                   ↑                           │\n",
    "│                          This is the killer!                  │\n",
    "│                                                                 │\n",
    "└─────────────────────────────────────────────────────────────────┘\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0849dfd0",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "analyze-scaling",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def analyze_parameter_scaling():\n",
    "    \"\"\"📊 Analyze how parameter count scales with model dimensions.\"\"\"\n",
    "    print(\"📊 Analyzing Parameter Scaling in Transformers...\")\n",
    "    print(\"Understanding why model size affects performance and cost\\n\")\n",
    "\n",
    "    # Test different model sizes\n",
    "    configs = [\n",
    "        {\"name\": \"Tiny\", \"embed_dim\": 64, \"num_layers\": 2, \"num_heads\": 4},\n",
    "        {\"name\": \"Small\", \"embed_dim\": 128, \"num_layers\": 4, \"num_heads\": 8},\n",
    "        {\"name\": \"Medium\", \"embed_dim\": 256, \"num_layers\": 8, \"num_heads\": 16},\n",
    "        {\"name\": \"Large\", \"embed_dim\": 512, \"num_layers\": 12, \"num_heads\": 16},\n",
    "    ]\n",
    "\n",
    "    vocab_size = 50000  # Typical vocabulary size\n",
    "\n",
    "    for config in configs:\n",
    "        model = GPT(\n",
    "            vocab_size=vocab_size,\n",
    "            embed_dim=config[\"embed_dim\"],\n",
    "            num_layers=config[\"num_layers\"],\n",
    "            num_heads=config[\"num_heads\"]\n",
    "        )\n",
    "\n",
    "        # Count parameters\n",
    "        total_params = 0\n",
    "        for param in model.parameters():\n",
    "            total_params += param.size\n",
    "\n",
    "        # Calculate memory requirements (4 bytes per float32 parameter)\n",
    "        memory_mb = (total_params * 4) / (1024 * 1024)\n",
    "\n",
    "        print(f\"{config['name']} Model:\")\n",
    "        print(f\"  Parameters: {total_params:,}\")\n",
    "        print(f\"  Memory: {memory_mb:.1f} MB\")\n",
    "        print(f\"  Embed dim: {config['embed_dim']}, Layers: {config['num_layers']}\")\n",
    "        print()\n",
    "\n",
    "    print(\"💡 Parameter scaling is roughly quadratic with embedding dimension\")\n",
    "    print(\"🚀 Real GPT-3 has 175B parameters, requiring ~350GB memory!\")\n",
    "\n",
    "analyze_parameter_scaling()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d83a8fb",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "analyze-attention-memory",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def analyze_attention_memory():\n",
    "    \"\"\"📊 Analyze attention memory complexity with sequence length.\"\"\"\n",
    "    print(\"📊 Analyzing Attention Memory Complexity...\")\n",
    "    print(\"Why long context is expensive and how it scales\\n\")\n",
    "\n",
    "    embed_dim = 512\n",
    "    num_heads = 8\n",
    "    batch_size = 4\n",
    "\n",
    "    # Test different sequence lengths\n",
    "    sequence_lengths = [128, 256, 512, 1024, 2048]\n",
    "\n",
    "    print(\"Attention Matrix Memory Usage:\")\n",
    "    print(\"Seq Len | Attention Matrix Size | Memory (MB)\")\n",
    "    print(\"-\" * 45)\n",
    "\n",
    "    for seq_len in sequence_lengths:\n",
    "        # Attention matrix is (batch_size, num_heads, seq_len, seq_len)\n",
    "        attention_elements = batch_size * num_heads * seq_len * seq_len\n",
    "\n",
    "        # 4 bytes per float32\n",
    "        memory_bytes = attention_elements * 4\n",
    "        memory_mb = memory_bytes / (1024 * 1024)\n",
    "\n",
    "        print(f\"{seq_len:6d} | {seq_len}×{seq_len} × {batch_size}×{num_heads} | {memory_mb:8.1f}\")\n",
    "\n",
    "    print()\n",
    "    print(\"💡 Attention memory grows quadratically with sequence length\")\n",
    "    print(\"🚀 This is why techniques like FlashAttention are crucial for long sequences\")\n",
    "\n",
    "analyze_attention_memory()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61c047e3",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🧪 Module Integration Test\n",
    "\n",
    "Final validation that everything works together correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f23223b",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-module",
     "locked": true,
     "points": 25
    }
   },
   "outputs": [],
   "source": [
    "def test_module():\n",
    "    \"\"\"\n",
    "    Comprehensive test of entire module functionality.\n",
    "\n",
    "    This final test runs before module summary to ensure:\n",
    "    - All unit tests pass\n",
    "    - Functions work together correctly\n",
    "    - Module is ready for integration with TinyTorch\n",
    "    \"\"\"\n",
    "    print(\"🧪 RUNNING MODULE INTEGRATION TEST\")\n",
    "    print(\"=\" * 50)\n",
    "\n",
    "    # Run all unit tests\n",
    "    print(\"Running unit tests...\")\n",
    "    test_unit_layer_norm()\n",
    "    test_unit_mlp()\n",
    "    test_unit_transformer_block()\n",
    "    test_unit_gpt()\n",
    "\n",
    "    print(\"\\nRunning integration scenarios...\")\n",
    "\n",
    "    # Test complete transformer training scenario\n",
    "    print(\"🔬 Integration Test: Full Training Pipeline...\")\n",
    "\n",
    "    # Create model and data\n",
    "    vocab_size = 50\n",
    "    embed_dim = 64\n",
    "    num_layers = 2\n",
    "    num_heads = 4\n",
    "\n",
    "    model = GPT(vocab_size, embed_dim, num_layers, num_heads)\n",
    "\n",
    "    # Test batch processing\n",
    "    batch_size = 3\n",
    "    seq_len = 16\n",
    "    tokens = Tensor(np.random.randint(0, vocab_size, (batch_size, seq_len)))\n",
    "\n",
    "    # Forward pass\n",
    "    logits = model.forward(tokens)\n",
    "    assert logits.shape == (batch_size, seq_len, vocab_size)\n",
    "\n",
    "    # Test generation with different temperatures\n",
    "    prompt = Tensor(np.random.randint(0, vocab_size, (1, 8)))\n",
    "\n",
    "    # Conservative generation\n",
    "    conservative = model.generate(prompt, max_new_tokens=5, temperature=0.1)\n",
    "    assert conservative.shape == (1, 13)\n",
    "\n",
    "    # Creative generation\n",
    "    creative = model.generate(prompt, max_new_tokens=5, temperature=2.0)\n",
    "    assert creative.shape == (1, 13)\n",
    "\n",
    "    # Test parameter counting consistency\n",
    "    total_params = sum(param.size for param in model.parameters())\n",
    "    assert total_params > 1000  # Should have substantial parameters\n",
    "\n",
    "    print(\"✅ Full transformer pipeline works!\")\n",
    "\n",
    "    print(\"\\n\" + \"=\" * 50)\n",
    "    print(\"🎉 ALL TESTS PASSED! Module ready for export.\")\n",
    "    print(\"Run: tito module complete 13\")\n",
    "\n",
    "# Call the comprehensive test\n",
    "test_module()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9c5a7f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    print(\"🚀 Running Transformers module...\")\n",
    "    test_module()\n",
    "    print(\"✅ Module validation complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "203f8df1",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🤔 ML Systems Thinking: Transformer Architecture Foundations\n",
    "\n",
    "### Question 1: Attention Memory Complexity\n",
    "You implemented multi-head attention that computes attention matrices of size (batch, heads, seq_len, seq_len).\n",
    "\n",
    "For a model with seq_len=1024, batch_size=4, num_heads=8:\n",
    "- How many elements in the attention matrix? _____\n",
    "- If each element is 4 bytes (float32), how much memory per layer? _____ MB\n",
    "- Why does doubling sequence length quadruple attention memory? _____\n",
    "\n",
    "### Question 2: Residual Connection Benefits\n",
    "Your TransformerBlock uses residual connections (x + attention_output, x + mlp_output).\n",
    "\n",
    "- What happens to gradients during backpropagation without residual connections? _____\n",
    "- How do residual connections help train deeper networks? _____\n",
    "- Why is pre-norm (LayerNorm before operations) preferred over post-norm? _____\n",
    "\n",
    "### Question 3: Parameter Scaling Analysis\n",
    "Your GPT model combines embeddings, transformer blocks, and output projection.\n",
    "\n",
    "For embed_dim=512, vocab_size=10000, num_layers=6:\n",
    "- Token embedding parameters: _____ (vocab_size × embed_dim)\n",
    "- Approximate parameters per transformer block: _____ (hint: ~4 × embed_dim²)\n",
    "- Total model parameters: approximately _____ million\n",
    "\n",
    "### Question 4: Autoregressive Generation Efficiency\n",
    "Your generate() method processes the full sequence for each new token.\n",
    "\n",
    "- Why is this inefficient for long sequences? _____\n",
    "- What optimization caches key-value pairs to avoid recomputation? _____\n",
    "- How would this change the computational complexity from O(n²) to O(n)? _____"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13761f1f",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 MODULE SUMMARY: Transformers\n",
    "\n",
    "Congratulations! You've built the complete transformer architecture that powers modern language models like GPT, Claude, and ChatGPT!\n",
    "\n",
    "### Key Accomplishments\n",
    "- Built LayerNorm for stable training across deep transformer networks\n",
    "- Implemented MLP (feed-forward) networks with GELU activation and 4x expansion\n",
    "- Created complete TransformerBlock with self-attention, residual connections, and pre-norm architecture\n",
    "- Built full GPT model with embeddings, positional encoding, and autoregressive generation\n",
    "- Discovered attention memory scaling and parameter distribution patterns\n",
    "- All tests pass ✅ (validated by `test_module()`)\n",
    "\n",
    "### Ready for Next Steps\n",
    "Your transformer implementation is the capstone of the language modeling pipeline.\n",
    "Export with: `tito module complete 13`\n",
    "\n",
    "**Next**: Module 14 will add profiling and optimization techniques to make your transformers production-ready!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
