{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "92420776",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Embeddings - Converting Tokens to Dense Vector Representations\n",
    "\n",
    "Welcome to the Embeddings module! You'll implement the systems that convert discrete tokens into rich vector representations that capture semantic meaning for language models.\n",
    "\n",
    "## Learning Goals\n",
    "- Systems understanding: How embedding tables scale with vocabulary size and affect model memory\n",
    "- Core implementation skill: Build embedding layers with efficient lookup operations\n",
    "- Pattern recognition: Understand how positional encoding enables sequence understanding\n",
    "- Framework connection: See how your implementations match PyTorch's embedding systems\n",
    "- Performance insight: Learn how embedding lookup patterns affect cache efficiency and memory bandwidth\n",
    "\n",
    "## Build → Use → Reflect\n",
    "1. **Build**: Embedding layer with lookup table and positional encoding systems\n",
    "2. **Use**: Transform token sequences into rich vector representations for language processing\n",
    "3. **Reflect**: How do embedding choices determine model capacity and computational efficiency?\n",
    "\n",
    "## What You'll Achieve\n",
    "By the end of this module, you'll understand:\n",
    "- Deep technical understanding of how discrete tokens become continuous vector representations\n",
    "- Practical capability to implement embedding systems that handle large vocabularies efficiently\n",
    "- Systems insight into how embedding dimensions affect model capacity and memory usage\n",
    "- Performance consideration of how embedding lookup patterns affect training and inference speed\n",
    "- Connection to production systems like transformer embedding layers and their optimization techniques\n",
    "\n",
    "## Systems Reality Check\n",
    "💡 **Production Context**: Modern language models have embedding tables with billions of parameters (GPT-3: 50k vocab × 12k dim = 600M embedding params)\n",
    "⚡ **Performance Note**: Embedding lookups are memory-bandwidth bound - efficient access patterns are critical for high-throughput training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da397b8",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "embeddings-imports",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| default_exp core.embeddings\n",
    "\n",
    "#| export\n",
    "import math\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "from typing import Union, List, Optional, Tuple\n",
    "\n",
    "# Import our Tensor class - try from package first, then from local module\n",
    "try:\n",
    "    from tinytorch.core.tensor import Tensor\n",
    "except ImportError:\n",
    "    # For development, import from local tensor module\n",
    "    sys.path.append(os.path.join(os.path.dirname(__file__), '..', '02_tensor'))\n",
    "    from tensor_dev import Tensor\n",
    "\n",
    "# Try to import tokenization classes\n",
    "try:\n",
    "    from tinytorch.core.tokenization import CharTokenizer, BPETokenizer\n",
    "except ImportError:\n",
    "    # For development, import from local module\n",
    "    sys.path.append(os.path.join(os.path.dirname(__file__), '..', '11_tokenization'))\n",
    "    try:\n",
    "        from tokenization_dev import CharTokenizer, BPETokenizer\n",
    "    except ImportError:\n",
    "        # Create minimal mock classes if not available\n",
    "        class CharTokenizer:\n",
    "            def __init__(self): \n",
    "                self.vocab_size = 256\n",
    "        class BPETokenizer:\n",
    "            def __init__(self, vocab_size=1000):\n",
    "                self.vocab_size = vocab_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83e2b76d",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "embeddings-welcome",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "print(\"🎯 TinyTorch Embeddings Module\")\n",
    "print(f\"NumPy version: {np.__version__}\")\n",
    "print(\"Ready to build embedding systems!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53f64bfc",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 📦 Where This Code Lives in the Final Package\n",
    "\n",
    "**Learning Side:** You work in `modules/source/12_embeddings/embeddings_dev.py`  \n",
    "**Building Side:** Code exports to `tinytorch.core.embeddings`\n",
    "\n",
    "```python\n",
    "# Final package structure:\n",
    "from tinytorch.core.embeddings import Embedding, PositionalEncoding\n",
    "from tinytorch.core.tokenization import CharTokenizer, BPETokenizer  # Previous module\n",
    "from tinytorch.core.attention import MultiHeadAttention  # Next module\n",
    "```\n",
    "\n",
    "**Why this matters:**\n",
    "- **Learning:** Focused modules for deep understanding\n",
    "- **Production:** Proper organization like PyTorch's `torch.nn.Embedding`\n",
    "- **Consistency:** All embedding tools live together in `core.embeddings`\n",
    "- **Integration:** Works seamlessly with tokenization and attention systems"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43aa5503",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## What are Embeddings?\n",
    "\n",
    "### The Problem: Discrete to Continuous\n",
    "Tokens are discrete symbols, but neural networks work best with continuous vectors:\n",
    "```\n",
    "Token: 42 → Dense Vector: [0.1, -0.3, 0.8, 0.2, ...]\n",
    "```\n",
    "\n",
    "### Embedding Table (Lookup Table)\n",
    "An embedding layer is essentially a learnable lookup table:\n",
    "```\n",
    "Vocabulary size: 50,000 tokens\n",
    "Embedding dimension: 512\n",
    "Total parameters: 50,000 × 512 = 25.6M parameters\n",
    "```\n",
    "\n",
    "### Why Embeddings Work\n",
    "- **Similarity**: Similar words get similar vectors through training\n",
    "- **Composition**: Vector operations capture semantic relationships\n",
    "- **Learning**: Gradients update embeddings to improve task performance\n",
    "\n",
    "### Positional Encoding\n",
    "Since transformers lack inherent position awareness, we add positional information:\n",
    "```\n",
    "Token embedding + Positional encoding = Position-aware representation\n",
    "```\n",
    "\n",
    "### Systems Trade-offs\n",
    "- **Embedding dimension**: Higher = more capacity, more memory\n",
    "- **Vocabulary size**: Larger = more parameters, better coverage\n",
    "- **Lookup efficiency**: Memory access patterns affect performance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c001050e",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Embedding Layer Implementation\n",
    "\n",
    "Let's start with the core embedding layer - a learnable lookup table that converts token indices to dense vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8a0101a",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "embedding-layer",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class Embedding:\n",
    "    \"\"\"\n",
    "    Embedding layer that converts token indices to dense vector representations.\n",
    "    \n",
    "    This is the foundation of modern language models - a learnable lookup table\n",
    "    that maps discrete tokens to continuous vectors that capture semantic meaning.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, vocab_size: int, embedding_dim: int, \n",
    "                 padding_idx: Optional[int] = None, \n",
    "                 init_type: str = 'uniform'):\n",
    "        \"\"\"\n",
    "        Initialize embedding layer with learnable parameters.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Store configuration parameters\n",
    "        2. Initialize embedding table with chosen initialization\n",
    "        3. Handle special padding token if specified\n",
    "        4. Set up for gradient tracking (will connect to autograd later)\n",
    "        \n",
    "        DESIGN DECISIONS:\n",
    "        - Embedding table shape: (vocab_size, embedding_dim)\n",
    "        - Initialization affects training dynamics\n",
    "        - Padding idx gets zero gradient to stay constant\n",
    "        \n",
    "        Args:\n",
    "            vocab_size: Number of tokens in vocabulary\n",
    "            embedding_dim: Size of dense vector for each token\n",
    "            padding_idx: Optional token index that should remain zero\n",
    "            init_type: Initialization strategy ('uniform', 'normal', 'xavier')\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.padding_idx = padding_idx\n",
    "        self.init_type = init_type\n",
    "        \n",
    "        # Initialize embedding table based on strategy\n",
    "        if init_type == 'uniform':\n",
    "            # Uniform initialization in [-1/sqrt(dim), 1/sqrt(dim)]\n",
    "            bound = 1.0 / math.sqrt(embedding_dim)\n",
    "            self.weight = Tensor(np.random.uniform(-bound, bound, (vocab_size, embedding_dim)))\n",
    "        elif init_type == 'normal':\n",
    "            # Normal initialization with std=1/sqrt(dim)\n",
    "            std = 1.0 / math.sqrt(embedding_dim)\n",
    "            self.weight = Tensor(np.random.normal(0, std, (vocab_size, embedding_dim)))\n",
    "        elif init_type == 'xavier':\n",
    "            # Xavier/Glorot initialization\n",
    "            bound = math.sqrt(6.0 / (vocab_size + embedding_dim))\n",
    "            self.weight = Tensor(np.random.uniform(-bound, bound, (vocab_size, embedding_dim)))\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown init_type: {init_type}\")\n",
    "        \n",
    "        # Set padding token to zero if specified\n",
    "        if padding_idx is not None:\n",
    "            self.weight.data[padding_idx] = 0.0\n",
    "        \n",
    "        # Track parameters for optimization\n",
    "        self.parameters = [self.weight]\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def forward(self, input_ids: Union[Tensor, List[int], np.ndarray]) -> Tensor:\n",
    "        \"\"\"\n",
    "        Look up embeddings for input token indices.\n",
    "        \n",
    "        TODO: Implement embedding lookup.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Convert input to numpy array if needed\n",
    "        2. Validate token indices are within vocabulary\n",
    "        3. Use advanced indexing to look up embeddings\n",
    "        4. Return tensor with shape (batch_size, seq_len, embedding_dim)\n",
    "        \n",
    "        EXAMPLE:\n",
    "        embed = Embedding(vocab_size=100, embedding_dim=64)\n",
    "        tokens = Tensor([[1, 2, 3], [4, 5, 6]])  # Shape: (2, 3)\n",
    "        embeddings = embed.forward(tokens)  # Shape: (2, 3, 64)\n",
    "        \n",
    "        IMPLEMENTATION HINTS:\n",
    "        - Handle both Tensor and list inputs\n",
    "        - Use numpy advanced indexing: weight[indices]\n",
    "        - Preserve batch and sequence dimensions\n",
    "        \n",
    "        Args:\n",
    "            input_ids: Token indices with shape (batch_size, seq_len) or (seq_len,)\n",
    "            \n",
    "        Returns:\n",
    "            Embeddings with shape (*input_shape, embedding_dim)\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Convert input to numpy array\n",
    "        if isinstance(input_ids, Tensor):\n",
    "            indices = input_ids.data\n",
    "        elif isinstance(input_ids, list):\n",
    "            indices = np.array(input_ids)\n",
    "        else:\n",
    "            indices = input_ids\n",
    "        \n",
    "        # Validate indices\n",
    "        indices = indices.astype(int)\n",
    "        if np.any(indices < 0) or np.any(indices >= self.vocab_size):\n",
    "            raise ValueError(f\"Token indices must be in range [0, {self.vocab_size})\")\n",
    "        \n",
    "        # Look up embeddings using advanced indexing\n",
    "        # self.weight.data has shape (vocab_size, embedding_dim)\n",
    "        # indices has shape (...), result has shape (..., embedding_dim)\n",
    "        embeddings = self.weight.data[indices]\n",
    "        \n",
    "        return Tensor(embeddings)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, input_ids: Union[Tensor, List[int], np.ndarray]) -> Tensor:\n",
    "        \"\"\"Make the layer callable.\"\"\"\n",
    "        return self.forward(input_ids)\n",
    "    \n",
    "    def get_memory_usage(self):\n",
    "        \"\"\"\n",
    "        Calculate memory usage of embedding table.\n",
    "        \n",
    "        This function is PROVIDED to show memory analysis.\n",
    "        \"\"\"\n",
    "        # Embedding table memory\n",
    "        weight_memory_mb = self.weight.data.nbytes / (1024 * 1024)\n",
    "        \n",
    "        # Memory per token\n",
    "        memory_per_token_kb = (self.embedding_dim * 4) / 1024  # 4 bytes per float32\n",
    "        \n",
    "        return {\n",
    "            'total_memory_mb': weight_memory_mb,\n",
    "            'memory_per_token_kb': memory_per_token_kb,\n",
    "            'total_parameters': self.vocab_size * self.embedding_dim,\n",
    "            'vocab_size': self.vocab_size,\n",
    "            'embedding_dim': self.embedding_dim\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b2e152",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Test Your Embedding Layer Implementation\n",
    "\n",
    "Once you implement the Embedding forward method above, run this cell to test it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aee0534c",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": true,
     "grade_id": "test-embedding-immediate",
     "locked": true,
     "points": 15,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_embedding_layer():\n",
    "    \"\"\"Unit test for the embedding layer.\"\"\"\n",
    "    print(\"🔬 Unit Test: Embedding Layer...\")\n",
    "    \n",
    "    # Create embedding layer\n",
    "    vocab_size = 100\n",
    "    embedding_dim = 64\n",
    "    embed = Embedding(vocab_size=vocab_size, embedding_dim=embedding_dim)\n",
    "    \n",
    "    # Test single token\n",
    "    single_token = [5]\n",
    "    single_embedding = embed.forward(single_token)\n",
    "    assert single_embedding.shape == (1, embedding_dim), f\"Expected shape (1, {embedding_dim}), got {single_embedding.shape}\"\n",
    "    \n",
    "    # Test sequence of tokens\n",
    "    token_sequence = [1, 2, 3, 5, 10]\n",
    "    sequence_embeddings = embed.forward(token_sequence)\n",
    "    expected_shape = (len(token_sequence), embedding_dim)\n",
    "    assert sequence_embeddings.shape == expected_shape, f\"Expected shape {expected_shape}, got {sequence_embeddings.shape}\"\n",
    "    \n",
    "    # Test batch of sequences\n",
    "    batch_tokens = [[1, 2, 3], [4, 5, 6]]\n",
    "    batch_embeddings = embed.forward(batch_tokens)\n",
    "    assert batch_embeddings.shape == (2, 3, embedding_dim), f\"Expected shape (2, 3, {embedding_dim}), got {batch_embeddings.shape}\"\n",
    "    \n",
    "    # Test with Tensor input\n",
    "    tensor_input = Tensor(np.array([[7, 8, 9], [10, 11, 12]]))\n",
    "    tensor_embeddings = embed.forward(tensor_input)\n",
    "    assert tensor_embeddings.shape == (2, 3, embedding_dim), \"Should handle Tensor input\"\n",
    "    \n",
    "    # Test embedding lookup consistency\n",
    "    token_5_embed_1 = embed.forward([5])\n",
    "    token_5_embed_2 = embed.forward([5])\n",
    "    assert np.allclose(token_5_embed_1.data, token_5_embed_2.data), \"Same token should give same embedding\"\n",
    "    \n",
    "    # Test different tokens give different embeddings (with high probability)\n",
    "    token_1_embed = embed.forward([1])\n",
    "    token_2_embed = embed.forward([2])\n",
    "    assert not np.allclose(token_1_embed.data, token_2_embed.data, atol=1e-3), \"Different tokens should give different embeddings\"\n",
    "    \n",
    "    # Test initialization bounds\n",
    "    assert np.all(np.abs(embed.weight.data) <= 1.0), \"Uniform initialization should be bounded\"\n",
    "    \n",
    "    # Test padding token (if specified)\n",
    "    embed_with_padding = Embedding(vocab_size=50, embedding_dim=32, padding_idx=0)\n",
    "    assert np.allclose(embed_with_padding.weight.data[0], 0.0), \"Padding token should be zero\"\n",
    "    \n",
    "    # Test parameter tracking\n",
    "    assert len(embed.parameters) == 1, \"Should track embedding weight parameter\"\n",
    "    assert embed.parameters[0] is embed.weight, \"Should track weight tensor\"\n",
    "    \n",
    "    # Test memory usage calculation\n",
    "    memory_stats = embed.get_memory_usage()\n",
    "    assert 'total_memory_mb' in memory_stats, \"Should provide memory statistics\"\n",
    "    assert memory_stats['total_parameters'] == vocab_size * embedding_dim, \"Should calculate parameters correctly\"\n",
    "    \n",
    "    print(\"✅ Embedding layer tests passed!\")\n",
    "    print(f\"✅ Handles various input shapes correctly\")\n",
    "    print(f\"✅ Consistent lookup and parameter tracking\")\n",
    "    print(f\"✅ Memory usage: {memory_stats['total_memory_mb']:.2f}MB\")\n",
    "\n",
    "# Test function defined (called in main block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e75ef30e",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Positional Encoding Implementation\n",
    "\n",
    "Transformers need explicit position information since attention is position-agnostic. Let's implement sinusoidal positional encoding used in the original transformer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23bcb128",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "positional-encoding",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class PositionalEncoding:\n",
    "    \"\"\"\n",
    "    Sinusoidal positional encoding that adds position information to embeddings.\n",
    "    \n",
    "    Uses sine and cosine functions of different frequencies to create\n",
    "    unique position representations that the model can learn to use.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, embedding_dim: int, max_seq_length: int = 5000, \n",
    "                 dropout: float = 0.0):\n",
    "        \"\"\"\n",
    "        Initialize positional encoding with sinusoidal patterns.\n",
    "        \n",
    "        TODO: Implement positional encoding initialization.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Create position matrix (max_seq_length, embedding_dim)\n",
    "        2. For each position and dimension:\n",
    "           - Calculate frequency based on dimension\n",
    "           - Apply sine to even dimensions, cosine to odd dimensions\n",
    "        3. Store the precomputed positional encodings\n",
    "        \n",
    "        MATHEMATICAL FOUNDATION:\n",
    "        PE(pos, 2i) = sin(pos / 10000^(2i/d_model))\n",
    "        PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))\n",
    "        \n",
    "        Where:\n",
    "        - pos = position in sequence\n",
    "        - i = dimension index\n",
    "        - d_model = embedding_dim\n",
    "        \n",
    "        Args:\n",
    "            embedding_dim: Dimension of embeddings (must be even)\n",
    "            max_seq_length: Maximum sequence length to precompute\n",
    "            dropout: Dropout rate (for future use)\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.max_seq_length = max_seq_length\n",
    "        self.dropout = dropout\n",
    "        \n",
    "        # Create positional encoding matrix\n",
    "        pe = np.zeros((max_seq_length, embedding_dim))\n",
    "        \n",
    "        # Create position vector (0, 1, 2, ..., max_seq_length-1)\n",
    "        position = np.arange(0, max_seq_length).reshape(-1, 1)  # Shape: (max_seq_length, 1)\n",
    "        \n",
    "        # Create dimension indices for frequency calculation\n",
    "        # div_term calculates 10000^(2i/d_model) for i = 0, 1, 2, ...\n",
    "        div_term = np.exp(np.arange(0, embedding_dim, 2) * \n",
    "                         -(math.log(10000.0) / embedding_dim))\n",
    "        \n",
    "        # Apply sine to even dimensions (0, 2, 4, ...)\n",
    "        pe[:, 0::2] = np.sin(position * div_term)\n",
    "        \n",
    "        # Apply cosine to odd dimensions (1, 3, 5, ...)\n",
    "        if embedding_dim % 2 == 1:\n",
    "            # Handle odd embedding_dim - cosine gets one less dimension\n",
    "            pe[:, 1::2] = np.cos(position * div_term[:-1])\n",
    "        else:\n",
    "            pe[:, 1::2] = np.cos(position * div_term)\n",
    "        \n",
    "        # Store as tensor\n",
    "        self.pe = Tensor(pe)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def forward(self, embeddings: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Add positional encoding to embeddings.\n",
    "        \n",
    "        TODO: Implement positional encoding addition.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Get sequence length from embeddings shape\n",
    "        2. Extract relevant positional encodings\n",
    "        3. Add positional encodings to embeddings\n",
    "        4. Return position-aware embeddings\n",
    "        \n",
    "        EXAMPLE:\n",
    "        pos_enc = PositionalEncoding(embedding_dim=64)\n",
    "        embeddings = Tensor(np.random.randn(2, 10, 64))  # (batch, seq, dim)\n",
    "        pos_embeddings = pos_enc.forward(embeddings)\n",
    "        \n",
    "        Args:\n",
    "            embeddings: Input embeddings with shape (batch_size, seq_len, embedding_dim)\n",
    "            \n",
    "        Returns:\n",
    "            Position-aware embeddings with same shape as input\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Get sequence length from embeddings\n",
    "        if len(embeddings.shape) == 3:\n",
    "            batch_size, seq_length, embed_dim = embeddings.shape\n",
    "        elif len(embeddings.shape) == 2:\n",
    "            seq_length, embed_dim = embeddings.shape\n",
    "            batch_size = None\n",
    "        else:\n",
    "            raise ValueError(f\"Expected 2D or 3D embeddings, got shape {embeddings.shape}\")\n",
    "        \n",
    "        if embed_dim != self.embedding_dim:\n",
    "            raise ValueError(f\"Embedding dim mismatch: expected {self.embedding_dim}, got {embed_dim}\")\n",
    "        \n",
    "        if seq_length > self.max_seq_length:\n",
    "            raise ValueError(f\"Sequence length {seq_length} exceeds max {self.max_seq_length}\")\n",
    "        \n",
    "        # Extract positional encodings for this sequence length\n",
    "        position_encodings = self.pe.data[:seq_length, :]\n",
    "        \n",
    "        # Add positional encodings to embeddings\n",
    "        if batch_size is not None:\n",
    "            # Broadcast positional encodings across batch dimension\n",
    "            # embeddings: (batch, seq, dim) + position_encodings: (seq, dim)\n",
    "            result = embeddings.data + position_encodings[np.newaxis, :, :]\n",
    "        else:\n",
    "            # embeddings: (seq, dim) + position_encodings: (seq, dim)\n",
    "            result = embeddings.data + position_encodings\n",
    "        \n",
    "        return Tensor(result)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, embeddings: Tensor) -> Tensor:\n",
    "        \"\"\"Make the class callable.\"\"\"\n",
    "        return self.forward(embeddings)\n",
    "    \n",
    "    def visualize_encoding(self, seq_length: int = 100, dims_to_show: int = 10) -> None:\n",
    "        \"\"\"\n",
    "        Visualize positional encoding patterns.\n",
    "        \n",
    "        This function is PROVIDED to show encoding patterns.\n",
    "        \"\"\"\n",
    "        print(f\"📊 POSITIONAL ENCODING VISUALIZATION\")\n",
    "        print(f\"Sequence length: {seq_length}, Dimensions shown: {dims_to_show}\")\n",
    "        print(\"=\" * 60)\n",
    "        \n",
    "        # Get subset of positional encodings\n",
    "        pe_subset = self.pe.data[:seq_length, :dims_to_show]\n",
    "        \n",
    "        # Show patterns for first few positions\n",
    "        print(\"First 10 positions, first 10 dimensions:\")\n",
    "        print(\"Pos\", end=\"\")\n",
    "        for d in range(min(dims_to_show, 10)):\n",
    "            print(f\"    Dim{d:2d}\", end=\"\")\n",
    "        print()\n",
    "        \n",
    "        for pos in range(min(seq_length, 10)):\n",
    "            print(f\"{pos:3d}\", end=\"\")\n",
    "            for d in range(min(dims_to_show, 10)):\n",
    "                print(f\"{pe_subset[pos, d]:8.3f}\", end=\"\")\n",
    "            print()\n",
    "        \n",
    "        # Show frequency analysis\n",
    "        print(f\"\\n📈 FREQUENCY ANALYSIS:\")\n",
    "        print(\"Even dimensions (sine): Lower frequencies for early dimensions\")\n",
    "        print(\"Odd dimensions (cosine): Same frequencies, phase-shifted\")\n",
    "        \n",
    "        # Calculate frequency range\n",
    "        min_freq = 1.0 / 10000\n",
    "        max_freq = 1.0\n",
    "        print(f\"Frequency range: {min_freq:.6f} to {max_freq:.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae0b3f88",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Test Your Positional Encoding Implementation\n",
    "\n",
    "Once you implement the PositionalEncoding methods above, run this cell to test it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fe34d59",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": true,
     "grade_id": "test-positional-encoding-immediate",
     "locked": true,
     "points": 15,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_positional_encoding():\n",
    "    \"\"\"Unit test for positional encoding.\"\"\"\n",
    "    print(\"🔬 Unit Test: Positional Encoding...\")\n",
    "    \n",
    "    # Create positional encoding\n",
    "    embedding_dim = 64\n",
    "    max_seq_length = 100\n",
    "    pos_enc = PositionalEncoding(embedding_dim=embedding_dim, max_seq_length=max_seq_length)\n",
    "    \n",
    "    # Test initialization\n",
    "    assert pos_enc.pe.shape == (max_seq_length, embedding_dim), f\"Expected shape ({max_seq_length}, {embedding_dim})\"\n",
    "    \n",
    "    # Test that different positions have different encodings\n",
    "    pos_0 = pos_enc.pe.data[0]\n",
    "    pos_1 = pos_enc.pe.data[1]\n",
    "    assert not np.allclose(pos_0, pos_1), \"Different positions should have different encodings\"\n",
    "    \n",
    "    # Test sine/cosine pattern\n",
    "    # Even dimensions should use sine, odd should use cosine\n",
    "    # This is hard to test directly, but we can check the encoding is reasonable\n",
    "    assert not np.any(np.isnan(pos_enc.pe.data)), \"Positional encodings should not contain NaN\"\n",
    "    assert not np.any(np.isinf(pos_enc.pe.data)), \"Positional encodings should not contain inf\"\n",
    "    \n",
    "    # Test forward pass with 3D input (batch, seq, dim)\n",
    "    batch_size = 2\n",
    "    seq_length = 10\n",
    "    embeddings = Tensor(np.random.randn(batch_size, seq_length, embedding_dim))\n",
    "    \n",
    "    pos_embeddings = pos_enc.forward(embeddings)\n",
    "    assert pos_embeddings.shape == embeddings.shape, \"Output shape should match input shape\"\n",
    "    \n",
    "    # Test forward pass with 2D input (seq, dim)\n",
    "    embeddings_2d = Tensor(np.random.randn(seq_length, embedding_dim))\n",
    "    pos_embeddings_2d = pos_enc.forward(embeddings_2d)\n",
    "    assert pos_embeddings_2d.shape == embeddings_2d.shape, \"2D output shape should match input\"\n",
    "    \n",
    "    # Test that positional encoding is actually added\n",
    "    original_mean = np.mean(embeddings.data)\n",
    "    pos_mean = np.mean(pos_embeddings.data)\n",
    "    assert abs(pos_mean - original_mean) > 1e-6, \"Positional encoding should change the embeddings\"\n",
    "    \n",
    "    # Test sequence length validation\n",
    "    try:\n",
    "        long_embeddings = Tensor(np.random.randn(max_seq_length + 10, embedding_dim))\n",
    "        pos_enc.forward(long_embeddings)\n",
    "        assert False, \"Should raise error for sequence longer than max_seq_length\"\n",
    "    except ValueError:\n",
    "        pass  # Expected behavior\n",
    "    \n",
    "    # Test embedding dimension validation\n",
    "    try:\n",
    "        wrong_dim_embeddings = Tensor(np.random.randn(seq_length, embedding_dim + 10))\n",
    "        pos_enc.forward(wrong_dim_embeddings)\n",
    "        assert False, \"Should raise error for wrong embedding dimension\"\n",
    "    except ValueError:\n",
    "        pass  # Expected behavior\n",
    "    \n",
    "    # Test deterministic behavior\n",
    "    pos_embeddings_1 = pos_enc.forward(embeddings)\n",
    "    pos_embeddings_2 = pos_enc.forward(embeddings)\n",
    "    assert np.allclose(pos_embeddings_1.data, pos_embeddings_2.data), \"Should be deterministic\"\n",
    "    \n",
    "    # Test callable interface\n",
    "    pos_embeddings_callable = pos_enc(embeddings)\n",
    "    assert np.allclose(pos_embeddings_callable.data, pos_embeddings.data), \"Callable interface should work\"\n",
    "    \n",
    "    print(\"✅ Positional encoding tests passed!\")\n",
    "    print(f\"✅ Handles 2D and 3D inputs correctly\")\n",
    "    print(f\"✅ Proper validation and deterministic behavior\")\n",
    "    print(f\"✅ Encoding dimension: {embedding_dim}, Max length: {max_seq_length}\")\n",
    "\n",
    "# Test function defined (called in main block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bc6f623",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Learned Positional Embeddings\n",
    "\n",
    "Some models use learned positional embeddings instead of fixed sinusoidal ones. Let's implement this alternative approach:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c50a89a",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "learned-positional",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class LearnedPositionalEmbedding:\n",
    "    \"\"\"\n",
    "    Learned positional embeddings - another embedding table for positions.\n",
    "    \n",
    "    Unlike sinusoidal encoding, these are learned parameters that\n",
    "    the model optimizes during training. Used in models like BERT.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, max_seq_length: int, embedding_dim: int):\n",
    "        \"\"\"\n",
    "        Initialize learned positional embeddings.\n",
    "        \n",
    "        TODO: Implement learned positional embedding initialization.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Create embedding layer for positions (0, 1, 2, ..., max_seq_length-1)\n",
    "        2. Initialize with small random values\n",
    "        3. Set up parameter tracking for optimization\n",
    "        \n",
    "        This is essentially an Embedding layer where the \"vocabulary\"\n",
    "        is the set of possible positions in a sequence.\n",
    "        \n",
    "        Args:\n",
    "            max_seq_length: Maximum sequence length supported\n",
    "            embedding_dim: Dimension of position embeddings\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.max_seq_length = max_seq_length\n",
    "        self.embedding_dim = embedding_dim\n",
    "        \n",
    "        # Create learned positional embedding table\n",
    "        # This is like an embedding layer for positions\n",
    "        self.position_embedding = Embedding(\n",
    "            vocab_size=max_seq_length,\n",
    "            embedding_dim=embedding_dim,\n",
    "            init_type='normal'\n",
    "        )\n",
    "        \n",
    "        # Track parameters for optimization\n",
    "        self.parameters = self.position_embedding.parameters\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def forward(self, embeddings: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Add learned positional embeddings to input embeddings.\n",
    "        \n",
    "        TODO: Implement learned positional embedding addition.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Get sequence length from input shape\n",
    "        2. Create position indices [0, 1, 2, ..., seq_length-1]\n",
    "        3. Look up position embeddings using position indices\n",
    "        4. Add position embeddings to input embeddings\n",
    "        \n",
    "        EXAMPLE:\n",
    "        learned_pos = LearnedPositionalEmbedding(max_seq_length=100, embedding_dim=64)\n",
    "        embeddings = Tensor(np.random.randn(2, 10, 64))  # (batch, seq, dim)\n",
    "        pos_embeddings = learned_pos.forward(embeddings)\n",
    "        \n",
    "        Args:\n",
    "            embeddings: Input embeddings with shape (batch_size, seq_len, embedding_dim)\n",
    "            \n",
    "        Returns:\n",
    "            Position-aware embeddings with same shape as input\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Get sequence length from embeddings\n",
    "        if len(embeddings.shape) == 3:\n",
    "            batch_size, seq_length, embed_dim = embeddings.shape\n",
    "        elif len(embeddings.shape) == 2:\n",
    "            seq_length, embed_dim = embeddings.shape\n",
    "            batch_size = None\n",
    "        else:\n",
    "            raise ValueError(f\"Expected 2D or 3D embeddings, got shape {embeddings.shape}\")\n",
    "        \n",
    "        if embed_dim != self.embedding_dim:\n",
    "            raise ValueError(f\"Embedding dim mismatch: expected {self.embedding_dim}, got {embed_dim}\")\n",
    "        \n",
    "        if seq_length > self.max_seq_length:\n",
    "            raise ValueError(f\"Sequence length {seq_length} exceeds max {self.max_seq_length}\")\n",
    "        \n",
    "        # Create position indices [0, 1, 2, ..., seq_length-1]\n",
    "        position_ids = list(range(seq_length))\n",
    "        \n",
    "        # Look up position embeddings\n",
    "        position_embeddings = self.position_embedding.forward(position_ids)\n",
    "        \n",
    "        # Add position embeddings to input embeddings\n",
    "        if batch_size is not None:\n",
    "            # Broadcast across batch dimension\n",
    "            result = embeddings.data + position_embeddings.data[np.newaxis, :, :]\n",
    "        else:\n",
    "            result = embeddings.data + position_embeddings.data\n",
    "        \n",
    "        return Tensor(result)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, embeddings: Tensor) -> Tensor:\n",
    "        \"\"\"Make the class callable.\"\"\"\n",
    "        return self.forward(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4811c8f5",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Test Your Learned Positional Embedding Implementation\n",
    "\n",
    "Once you implement the LearnedPositionalEmbedding methods above, run this cell to test it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8509e91",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": true,
     "grade_id": "test-learned-positional-immediate",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_learned_positional_embedding():\n",
    "    \"\"\"Unit test for learned positional embeddings.\"\"\"\n",
    "    print(\"🔬 Unit Test: Learned Positional Embeddings...\")\n",
    "    \n",
    "    # Create learned positional embedding\n",
    "    max_seq_length = 50\n",
    "    embedding_dim = 32\n",
    "    learned_pos = LearnedPositionalEmbedding(max_seq_length=max_seq_length, embedding_dim=embedding_dim)\n",
    "    \n",
    "    # Test initialization\n",
    "    assert learned_pos.position_embedding.vocab_size == max_seq_length, \"Should have position for each sequence position\"\n",
    "    assert learned_pos.position_embedding.embedding_dim == embedding_dim, \"Should match embedding dimension\"\n",
    "    \n",
    "    # Test parameter tracking\n",
    "    assert len(learned_pos.parameters) == 1, \"Should track position embedding parameters\"\n",
    "    assert learned_pos.parameters[0] is learned_pos.position_embedding.weight, \"Should track weight tensor\"\n",
    "    \n",
    "    # Test forward pass with 3D input\n",
    "    batch_size = 3\n",
    "    seq_length = 10\n",
    "    embeddings = Tensor(np.random.randn(batch_size, seq_length, embedding_dim))\n",
    "    \n",
    "    pos_embeddings = learned_pos.forward(embeddings)\n",
    "    assert pos_embeddings.shape == embeddings.shape, \"Output shape should match input shape\"\n",
    "    \n",
    "    # Test forward pass with 2D input\n",
    "    embeddings_2d = Tensor(np.random.randn(seq_length, embedding_dim))\n",
    "    pos_embeddings_2d = learned_pos.forward(embeddings_2d)\n",
    "    assert pos_embeddings_2d.shape == embeddings_2d.shape, \"2D output shape should match input\"\n",
    "    \n",
    "    # Test that position embeddings are actually added\n",
    "    original_mean = np.mean(embeddings.data)\n",
    "    pos_mean = np.mean(pos_embeddings.data)\n",
    "    assert abs(pos_mean - original_mean) > 1e-6, \"Position embeddings should change the input\"\n",
    "    \n",
    "    # Test that different sequence lengths give different results\n",
    "    short_embeddings = Tensor(np.random.randn(batch_size, 5, embedding_dim))\n",
    "    long_embeddings = Tensor(np.random.randn(batch_size, 15, embedding_dim))\n",
    "    \n",
    "    short_pos = learned_pos.forward(short_embeddings)\n",
    "    long_pos = learned_pos.forward(long_embeddings)\n",
    "    \n",
    "    # The first 5 positions should be the same\n",
    "    assert np.allclose(short_pos.data, long_pos.data[:, :5, :]), \"Same positions should have same embeddings\"\n",
    "    \n",
    "    # Test sequence length validation\n",
    "    try:\n",
    "        too_long_embeddings = Tensor(np.random.randn(batch_size, max_seq_length + 5, embedding_dim))\n",
    "        learned_pos.forward(too_long_embeddings)\n",
    "        assert False, \"Should raise error for sequence longer than max_seq_length\"\n",
    "    except ValueError:\n",
    "        pass  # Expected behavior\n",
    "    \n",
    "    # Test embedding dimension validation\n",
    "    try:\n",
    "        wrong_dim_embeddings = Tensor(np.random.randn(batch_size, seq_length, embedding_dim + 5))\n",
    "        learned_pos.forward(wrong_dim_embeddings)\n",
    "        assert False, \"Should raise error for wrong embedding dimension\"\n",
    "    except ValueError:\n",
    "        pass  # Expected behavior\n",
    "    \n",
    "    # Test callable interface\n",
    "    pos_embeddings_callable = learned_pos(embeddings)\n",
    "    assert np.allclose(pos_embeddings_callable.data, pos_embeddings.data), \"Callable interface should work\"\n",
    "    \n",
    "    print(\"✅ Learned positional embedding tests passed!\")\n",
    "    print(f\"✅ Parameter tracking and optimization ready\")\n",
    "    print(f\"✅ Handles various input shapes correctly\")\n",
    "    print(f\"✅ Max sequence length: {max_seq_length}, Embedding dim: {embedding_dim}\")\n",
    "\n",
    "# Test function defined (called in main block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edbe4f83",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 ML Systems: Performance Analysis & Embedding Scaling\n",
    "\n",
    "Now let's develop systems engineering skills by analyzing embedding performance and understanding how embedding choices affect downstream ML system efficiency.\n",
    "\n",
    "### **Learning Outcome**: *\"I understand how embedding table size affects model memory, training speed, and language understanding capacity\"*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ad087b6",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "embedding-profiler",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "import time\n",
    "\n",
    "class EmbeddingProfiler:\n",
    "    \"\"\"\n",
    "    Performance profiling toolkit for embedding systems.\n",
    "    \n",
    "    Helps ML engineers understand memory usage, lookup performance,\n",
    "    and scaling characteristics of embedding layers.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.results = {}\n",
    "    \n",
    "    def measure_lookup_performance(self, embedding_layer: Embedding, \n",
    "                                  batch_sizes: List[int], seq_lengths: List[int]) -> Dict:\n",
    "        \"\"\"\n",
    "        Measure embedding lookup performance across different batch sizes and sequence lengths.\n",
    "        \n",
    "        TODO: Implement embedding lookup performance measurement.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Create test token indices for each (batch_size, seq_length) combination\n",
    "        2. Measure time to perform embedding lookup\n",
    "        3. Calculate throughput metrics (tokens/second, memory bandwidth)\n",
    "        4. Return comprehensive performance analysis\n",
    "        \n",
    "        METRICS TO CALCULATE:\n",
    "        - Lookup time (milliseconds)\n",
    "        - Tokens per second throughput\n",
    "        - Memory bandwidth utilization\n",
    "        - Scaling patterns with batch size and sequence length\n",
    "        \n",
    "        Args:\n",
    "            embedding_layer: Embedding layer to test\n",
    "            batch_sizes: List of batch sizes to test\n",
    "            seq_lengths: List of sequence lengths to test\n",
    "            \n",
    "        Returns:\n",
    "            Dictionary with performance metrics for each configuration\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        results = {}\n",
    "        vocab_size = embedding_layer.vocab_size\n",
    "        \n",
    "        for batch_size in batch_sizes:\n",
    "            for seq_length in seq_lengths:\n",
    "                # Create random token indices\n",
    "                token_indices = np.random.randint(0, vocab_size, (batch_size, seq_length))\n",
    "                \n",
    "                # Measure lookup performance\n",
    "                start_time = time.time()\n",
    "                embeddings = embedding_layer.forward(token_indices)\n",
    "                end_time = time.time()\n",
    "                \n",
    "                # Calculate metrics\n",
    "                lookup_time_ms = (end_time - start_time) * 1000\n",
    "                total_tokens = batch_size * seq_length\n",
    "                tokens_per_second = total_tokens / (end_time - start_time) if end_time > start_time else 0\n",
    "                \n",
    "                # Memory calculations\n",
    "                input_memory_mb = token_indices.nbytes / (1024 * 1024)\n",
    "                output_memory_mb = embeddings.data.nbytes / (1024 * 1024)\n",
    "                memory_bandwidth_mb_s = (input_memory_mb + output_memory_mb) / (end_time - start_time) if end_time > start_time else 0\n",
    "                \n",
    "                config_key = f\"batch_{batch_size}_seq_{seq_length}\"\n",
    "                results[config_key] = {\n",
    "                    'batch_size': batch_size,\n",
    "                    'seq_length': seq_length,\n",
    "                    'total_tokens': total_tokens,\n",
    "                    'lookup_time_ms': lookup_time_ms,\n",
    "                    'tokens_per_second': tokens_per_second,\n",
    "                    'input_memory_mb': input_memory_mb,\n",
    "                    'output_memory_mb': output_memory_mb,\n",
    "                    'memory_bandwidth_mb_s': memory_bandwidth_mb_s,\n",
    "                    'time_per_token_us': lookup_time_ms * 1000 / total_tokens if total_tokens > 0 else 0\n",
    "                }\n",
    "        \n",
    "        return results\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def analyze_memory_scaling(self, vocab_sizes: List[int], embedding_dims: List[int]) -> Dict:\n",
    "        \"\"\"\n",
    "        Analyze how embedding memory usage scales with vocabulary size and embedding dimension.\n",
    "        \n",
    "        This function is PROVIDED to show memory scaling analysis.\n",
    "        \"\"\"\n",
    "        print(\"📊 EMBEDDING MEMORY SCALING ANALYSIS\")\n",
    "        print(\"=\" * 60)\n",
    "        \n",
    "        scaling_results = {}\n",
    "        \n",
    "        print(f\"{'Vocab Size':<12} {'Embed Dim':<10} {'Parameters':<12} {'Memory (MB)':<12} {'Lookup Time':<12}\")\n",
    "        print(\"-\" * 70)\n",
    "        \n",
    "        for vocab_size in vocab_sizes:\n",
    "            for embed_dim in embedding_dims:\n",
    "                # Create embedding layer\n",
    "                embed = Embedding(vocab_size=vocab_size, embedding_dim=embed_dim)\n",
    "                \n",
    "                # Calculate memory usage\n",
    "                memory_stats = embed.get_memory_usage()\n",
    "                total_memory_mb = memory_stats['total_memory_mb']\n",
    "                total_params = memory_stats['total_parameters']\n",
    "                \n",
    "                # Measure lookup time\n",
    "                test_tokens = np.random.randint(0, vocab_size, (32, 64))  # Standard batch\n",
    "                start_time = time.time()\n",
    "                _ = embed.forward(test_tokens)\n",
    "                lookup_time_ms = (time.time() - start_time) * 1000\n",
    "                \n",
    "                # Store results\n",
    "                config_key = f\"vocab_{vocab_size}_dim_{embed_dim}\"\n",
    "                scaling_results[config_key] = {\n",
    "                    'vocab_size': vocab_size,\n",
    "                    'embedding_dim': embed_dim,\n",
    "                    'total_parameters': total_params,\n",
    "                    'memory_mb': total_memory_mb,\n",
    "                    'lookup_time_ms': lookup_time_ms\n",
    "                }\n",
    "                \n",
    "                print(f\"{vocab_size:<12,} {embed_dim:<10} {total_params:<12,} {total_memory_mb:<12.2f} {lookup_time_ms:<12.2f}\")\n",
    "        \n",
    "        # Analyze scaling patterns\n",
    "        print(f\"\\n📈 SCALING INSIGHTS:\")\n",
    "        if len(vocab_sizes) > 1 and len(embedding_dims) > 1:\n",
    "            # Compare scaling with vocab size (fixed embedding dim)\n",
    "            fixed_dim = embedding_dims[0]\n",
    "            small_vocab = min(vocab_sizes)\n",
    "            large_vocab = max(vocab_sizes)\n",
    "            \n",
    "            small_key = f\"vocab_{small_vocab}_dim_{fixed_dim}\"\n",
    "            large_key = f\"vocab_{large_vocab}_dim_{fixed_dim}\"\n",
    "            \n",
    "            if small_key in scaling_results and large_key in scaling_results:\n",
    "                vocab_ratio = large_vocab / small_vocab\n",
    "                memory_ratio = scaling_results[large_key]['memory_mb'] / scaling_results[small_key]['memory_mb']\n",
    "                print(f\"   Vocabulary scaling: {vocab_ratio:.1f}x vocab → {memory_ratio:.1f}x memory (Linear)\")\n",
    "            \n",
    "            # Compare scaling with embedding dim (fixed vocab)\n",
    "            fixed_vocab = vocab_sizes[0]\n",
    "            small_dim = min(embedding_dims)\n",
    "            large_dim = max(embedding_dims)\n",
    "            \n",
    "            small_key = f\"vocab_{fixed_vocab}_dim_{small_dim}\"\n",
    "            large_key = f\"vocab_{fixed_vocab}_dim_{large_dim}\"\n",
    "            \n",
    "            if small_key in scaling_results and large_key in scaling_results:\n",
    "                dim_ratio = large_dim / small_dim\n",
    "                memory_ratio = scaling_results[large_key]['memory_mb'] / scaling_results[small_key]['memory_mb']\n",
    "                print(f\"   Dimension scaling: {dim_ratio:.1f}x dim → {memory_ratio:.1f}x memory (Linear)\")\n",
    "        \n",
    "        return scaling_results\n",
    "    \n",
    "    def compare_positional_encodings(self, seq_length: int = 100, embedding_dim: int = 256) -> Dict:\n",
    "        \"\"\"\n",
    "        Compare performance and characteristics of different positional encoding approaches.\n",
    "        \n",
    "        This function is PROVIDED to show positional encoding comparison.\n",
    "        \"\"\"\n",
    "        print(f\"\\n🔍 POSITIONAL ENCODING COMPARISON\")\n",
    "        print(\"=\" * 50)\n",
    "        \n",
    "        # Create test embeddings\n",
    "        batch_size = 16\n",
    "        embeddings = Tensor(np.random.randn(batch_size, seq_length, embedding_dim))\n",
    "        \n",
    "        # Test sinusoidal positional encoding\n",
    "        sinusoidal_pe = PositionalEncoding(embedding_dim=embedding_dim, max_seq_length=seq_length*2)\n",
    "        start_time = time.time()\n",
    "        sin_result = sinusoidal_pe.forward(embeddings)\n",
    "        sin_time = (time.time() - start_time) * 1000\n",
    "        \n",
    "        # Test learned positional embedding\n",
    "        learned_pe = LearnedPositionalEmbedding(max_seq_length=seq_length*2, embedding_dim=embedding_dim)\n",
    "        start_time = time.time()\n",
    "        learned_result = learned_pe.forward(embeddings)\n",
    "        learned_time = (time.time() - start_time) * 1000\n",
    "        \n",
    "        # Calculate memory usage\n",
    "        sin_memory = 0  # No learnable parameters\n",
    "        learned_memory = learned_pe.position_embedding.get_memory_usage()['total_memory_mb']\n",
    "        \n",
    "        results = {\n",
    "            'sinusoidal': {\n",
    "                'computation_time_ms': sin_time,\n",
    "                'memory_usage_mb': sin_memory,\n",
    "                'parameters': 0,\n",
    "                'deterministic': True,\n",
    "                'extrapolation': 'Good (can handle longer sequences)'\n",
    "            },\n",
    "            'learned': {\n",
    "                'computation_time_ms': learned_time,\n",
    "                'memory_usage_mb': learned_memory,\n",
    "                'parameters': seq_length * 2 * embedding_dim,\n",
    "                'deterministic': False,\n",
    "                'extrapolation': 'Limited (fixed max sequence length)'\n",
    "            }\n",
    "        }\n",
    "        \n",
    "        print(f\"📊 COMPARISON RESULTS:\")\n",
    "        print(f\"{'Method':<12} {'Time (ms)':<10} {'Memory (MB)':<12} {'Parameters':<12} {'Extrapolation'}\")\n",
    "        print(\"-\" * 70)\n",
    "        print(f\"{'Sinusoidal':<12} {sin_time:<10.2f} {sin_memory:<12.2f} {0:<12,} {'Good'}\")\n",
    "        print(f\"{'Learned':<12} {learned_time:<10.2f} {learned_memory:<12.2f} {results['learned']['parameters']:<12,} {'Limited'}\")\n",
    "        \n",
    "        print(f\"\\n💡 INSIGHTS:\")\n",
    "        print(f\"   - Sinusoidal: Zero parameters, deterministic, good extrapolation\")\n",
    "        print(f\"   - Learned: Requires parameters, model-specific, limited extrapolation\")\n",
    "        print(f\"   - Choice depends on: model capacity, sequence length requirements, extrapolation needs\")\n",
    "        \n",
    "        return results\n",
    "\n",
    "def analyze_embedding_system_design():\n",
    "    \"\"\"\n",
    "    Comprehensive analysis of embedding system design choices and their impact.\n",
    "    \n",
    "    This function is PROVIDED to show systems-level design thinking.\n",
    "    \"\"\"\n",
    "    print(\"🏗️ EMBEDDING SYSTEM DESIGN ANALYSIS\")\n",
    "    print(\"=\" * 60)\n",
    "    \n",
    "    # Example model configurations\n",
    "    model_configs = [\n",
    "        {'name': 'Small GPT', 'vocab_size': 10000, 'embed_dim': 256, 'seq_length': 512},\n",
    "        {'name': 'Medium GPT', 'vocab_size': 50000, 'embed_dim': 512, 'seq_length': 1024},\n",
    "        {'name': 'Large GPT', 'vocab_size': 50000, 'embed_dim': 1024, 'seq_length': 2048}\n",
    "    ]\n",
    "    \n",
    "    print(f\"📋 MODEL CONFIGURATION COMPARISON:\")\n",
    "    print(f\"{'Model':<12} {'Vocab Size':<10} {'Embed Dim':<10} {'Seq Len':<8} {'Embed Params':<12} {'Memory (MB)'}\")\n",
    "    print(\"-\" * 80)\n",
    "    \n",
    "    for config in model_configs:\n",
    "        # Calculate embedding parameters\n",
    "        embed_params = config['vocab_size'] * config['embed_dim']\n",
    "        \n",
    "        # Calculate memory usage\n",
    "        embed_memory_mb = embed_params * 4 / (1024 * 1024)  # 4 bytes per float32\n",
    "        \n",
    "        print(f\"{config['name']:<12} {config['vocab_size']:<10,} {config['embed_dim']:<10} \"\n",
    "              f\"{config['seq_length']:<8} {embed_params:<12,} {embed_memory_mb:<10.1f}\")\n",
    "    \n",
    "    print(f\"\\n🎯 DESIGN TRADE-OFFS:\")\n",
    "    print(f\"   1. Vocabulary Size:\")\n",
    "    print(f\"      - Larger vocab: Better text coverage, more parameters\")\n",
    "    print(f\"      - Smaller vocab: Longer sequences, more compute\")\n",
    "    print(f\"   2. Embedding Dimension:\")\n",
    "    print(f\"      - Higher dim: More model capacity, more memory\")\n",
    "    print(f\"      - Lower dim: Faster computation, potential bottleneck\")\n",
    "    print(f\"   3. Position Encoding:\")\n",
    "    print(f\"      - Sinusoidal: No parameters, good extrapolation\")\n",
    "    print(f\"      - Learned: Model-specific, limited to training length\")\n",
    "    print(f\"   4. Memory Scaling:\")\n",
    "    print(f\"      - Embedding table: O(vocab_size × embed_dim)\")\n",
    "    print(f\"      - Sequence processing: O(batch_size × seq_length × embed_dim)\")\n",
    "    print(f\"      - Total memory dominated by model size, not embedding table\")\n",
    "    \n",
    "    print(f\"\\n🏭 PRODUCTION CONSIDERATIONS:\")\n",
    "    print(f\"   - GPU memory limits affect maximum embedding table size\")\n",
    "    print(f\"   - Embedding lookup is memory-bandwidth bound\")\n",
    "    print(f\"   - Vocabulary size affects tokenization and model download size\")\n",
    "    print(f\"   - Position encoding choice affects sequence length flexibility\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a52101b6",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Test: Embedding Performance Analysis\n",
    "\n",
    "Let's test our embedding profiler with realistic performance scenarios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36fbea2a",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "test-embedding-profiler",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def test_embedding_profiler():\n",
    "    \"\"\"Test embedding profiler with various scenarios.\"\"\"\n",
    "    print(\"🔬 Unit Test: Embedding Performance Profiler...\")\n",
    "    \n",
    "    profiler = EmbeddingProfiler()\n",
    "    \n",
    "    # Create test embedding layer\n",
    "    vocab_size = 1000\n",
    "    embedding_dim = 128\n",
    "    embed = Embedding(vocab_size=vocab_size, embedding_dim=embedding_dim)\n",
    "    \n",
    "    # Test lookup performance measurement\n",
    "    batch_sizes = [8, 16]\n",
    "    seq_lengths = [32, 64]\n",
    "    \n",
    "    performance_results = profiler.measure_lookup_performance(embed, batch_sizes, seq_lengths)\n",
    "    \n",
    "    # Verify results structure\n",
    "    expected_configs = len(batch_sizes) * len(seq_lengths)\n",
    "    assert len(performance_results) == expected_configs, f\"Should test {expected_configs} configurations\"\n",
    "    \n",
    "    for config, metrics in performance_results.items():\n",
    "        # Verify all required metrics are present\n",
    "        required_keys = ['batch_size', 'seq_length', 'total_tokens', 'lookup_time_ms', \n",
    "                        'tokens_per_second', 'memory_bandwidth_mb_s']\n",
    "        for key in required_keys:\n",
    "            assert key in metrics, f\"Missing metric: {key} in {config}\"\n",
    "            assert isinstance(metrics[key], (int, float)), f\"Invalid metric type for {key}\"\n",
    "        \n",
    "        # Verify reasonable values\n",
    "        assert metrics['total_tokens'] > 0, \"Should count tokens\"\n",
    "        assert metrics['lookup_time_ms'] >= 0, \"Time should be non-negative\"\n",
    "        assert metrics['tokens_per_second'] >= 0, \"Throughput should be non-negative\"\n",
    "    \n",
    "    print(\"✅ Lookup performance measurement test passed\")\n",
    "    \n",
    "    # Test memory scaling analysis\n",
    "    vocab_sizes = [500, 1000]\n",
    "    embedding_dims = [64, 128]\n",
    "    \n",
    "    scaling_results = profiler.analyze_memory_scaling(vocab_sizes, embedding_dims)\n",
    "    \n",
    "    # Verify scaling results\n",
    "    expected_configs = len(vocab_sizes) * len(embedding_dims)\n",
    "    assert len(scaling_results) == expected_configs, f\"Should test {expected_configs} configurations\"\n",
    "    \n",
    "    for config, metrics in scaling_results.items():\n",
    "        assert 'total_parameters' in metrics, \"Should include parameter count\"\n",
    "        assert 'memory_mb' in metrics, \"Should include memory usage\"\n",
    "        assert metrics['total_parameters'] > 0, \"Should have parameters\"\n",
    "        assert metrics['memory_mb'] > 0, \"Should use memory\"\n",
    "    \n",
    "    print(\"✅ Memory scaling analysis test passed\")\n",
    "    \n",
    "    # Test positional encoding comparison\n",
    "    comparison_results = profiler.compare_positional_encodings(seq_length=50, embedding_dim=64)\n",
    "    \n",
    "    # Verify comparison results\n",
    "    assert 'sinusoidal' in comparison_results, \"Should test sinusoidal encoding\"\n",
    "    assert 'learned' in comparison_results, \"Should test learned encoding\"\n",
    "    \n",
    "    for method, metrics in comparison_results.items():\n",
    "        assert 'computation_time_ms' in metrics, \"Should measure computation time\"\n",
    "        assert 'memory_usage_mb' in metrics, \"Should measure memory usage\"\n",
    "        assert 'parameters' in metrics, \"Should count parameters\"\n",
    "    \n",
    "    print(\"✅ Positional encoding comparison test passed\")\n",
    "    print(\"🎯 Embedding Profiler: All tests passed!\")\n",
    "\n",
    "# Test function defined (called in main block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b2b90b6",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Integration Testing: Complete Embedding Pipeline\n",
    "\n",
    "Let's test how all our embedding components work together in a realistic language processing pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4798be5",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "test-embedding-integration",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def test_embedding_integration():\n",
    "    \"\"\"Test complete embedding pipeline with tokenization integration.\"\"\"\n",
    "    print(\"🧪 Integration Test: Complete Embedding Pipeline...\")\n",
    "    \n",
    "    # Create tokenizer\n",
    "    tokenizer = CharTokenizer()\n",
    "    \n",
    "    # Create embedding layer\n",
    "    embed = Embedding(vocab_size=tokenizer.vocab_size, embedding_dim=128, padding_idx=0)\n",
    "    \n",
    "    # Create positional encoding\n",
    "    pos_encoding = PositionalEncoding(embedding_dim=128, max_seq_length=100)\n",
    "    \n",
    "    # Test text processing pipeline\n",
    "    texts = [\n",
    "        \"Hello world!\",\n",
    "        \"This is a test.\",\n",
    "        \"Short text.\",\n",
    "        \"A longer piece of text to test the pipeline.\"\n",
    "    ]\n",
    "    \n",
    "    print(f\"  Processing {len(texts)} texts through complete pipeline...\")\n",
    "    \n",
    "    # Step 1: Tokenize texts\n",
    "    tokenized = []\n",
    "    for text in texts:\n",
    "        tokens = tokenizer.encode(text, add_special_tokens=True)\n",
    "        tokenized.append(tokens)\n",
    "    \n",
    "    # Step 2: Pad sequences for batch processing\n",
    "    padded_sequences = tokenizer.pad_sequences(tokenized, max_length=20)\n",
    "    batch_tokens = Tensor(np.array(padded_sequences))\n",
    "    \n",
    "    print(f\"    Batch shape: {batch_tokens.shape}\")\n",
    "    \n",
    "    # Step 3: Embedding lookup\n",
    "    embeddings = embed.forward(batch_tokens)\n",
    "    print(f\"    Embeddings shape: {embeddings.shape}\")\n",
    "    \n",
    "    # Step 4: Add positional encoding\n",
    "    pos_embeddings = pos_encoding.forward(embeddings)\n",
    "    print(f\"    Position-aware embeddings shape: {pos_embeddings.shape}\")\n",
    "    \n",
    "    # Verify pipeline correctness\n",
    "    expected_shape = (len(texts), 20, 128)  # (batch, seq_len, embed_dim)\n",
    "    assert pos_embeddings.shape == expected_shape, f\"Expected {expected_shape}, got {pos_embeddings.shape}\"\n",
    "    \n",
    "    # Test that padding tokens have correct embeddings (should be zero from embedding layer)\n",
    "    padding_token_id = tokenizer.char_to_idx['<PAD>']\n",
    "    \n",
    "    # Find positions with padding tokens\n",
    "    padding_positions = (batch_tokens.data == padding_token_id)\n",
    "    \n",
    "    if np.any(padding_positions):\n",
    "        # Get embeddings for padding positions\n",
    "        padding_embeddings = embeddings.data[padding_positions]\n",
    "        \n",
    "        # Padding embeddings should be close to zero (from embedding initialization)\n",
    "        # Note: they won't be exactly zero because we add positional encoding\n",
    "        print(f\"    Padding token embeddings found: {np.sum(padding_positions)} positions\")\n",
    "    \n",
    "    # Test different sequence lengths\n",
    "    short_text = \"Hi!\"\n",
    "    short_tokens = tokenizer.encode(short_text, add_special_tokens=True)\n",
    "    short_tensor = Tensor(np.array([short_tokens]))  # Add batch dimension\n",
    "    \n",
    "    short_embeddings = embed.forward(short_tensor)\n",
    "    short_pos_embeddings = pos_encoding.forward(short_embeddings)\n",
    "    \n",
    "    print(f\"    Short text processing: {short_pos_embeddings.shape}\")\n",
    "    \n",
    "    # Test memory efficiency\n",
    "    large_batch_size = 32\n",
    "    large_seq_length = 50\n",
    "    large_tokens = np.random.randint(0, tokenizer.vocab_size, (large_batch_size, large_seq_length))\n",
    "    large_tensor = Tensor(large_tokens)\n",
    "    \n",
    "    start_time = time.time()\n",
    "    large_embeddings = embed.forward(large_tensor)\n",
    "    large_pos_embeddings = pos_encoding.forward(large_embeddings)\n",
    "    processing_time = time.time() - start_time\n",
    "    \n",
    "    print(f\"    Large batch processing: {large_pos_embeddings.shape} in {processing_time*1000:.2f}ms\")\n",
    "    \n",
    "    # Calculate memory usage\n",
    "    embedding_memory = embed.get_memory_usage()\n",
    "    total_memory_mb = embedding_memory['total_memory_mb']\n",
    "    \n",
    "    print(f\"    Embedding table memory: {total_memory_mb:.2f}MB\")\n",
    "    print(f\"    Sequence memory: {large_pos_embeddings.data.nbytes / (1024*1024):.2f}MB\")\n",
    "    \n",
    "    print(\"✅ Complete embedding pipeline integration test passed!\")\n",
    "    print(f\"✅ Tokenization → Embedding → Positional Encoding pipeline works\")\n",
    "    print(f\"✅ Handles various batch sizes and sequence lengths\")\n",
    "    print(f\"✅ Memory usage is reasonable for production systems\")\n",
    "\n",
    "# Test function defined (called in main block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bba8138f",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## Main Execution Block\n",
    "\n",
    "All embedding tests and demonstrations are run from here when the module is executed directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19963dff",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "embeddings-main",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    # Run all unit tests\n",
    "    test_unit_embedding_layer()\n",
    "    test_unit_positional_encoding()\n",
    "    test_unit_learned_positional_embedding()\n",
    "    test_embedding_profiler()\n",
    "    test_embedding_integration()\n",
    "    \n",
    "    print(\"\\n\" + \"=\"*60)\n",
    "    print(\"🔍 EMBEDDING SYSTEMS ANALYSIS\")\n",
    "    print(\"=\"*60)\n",
    "    \n",
    "    # Performance analysis\n",
    "    profiler = EmbeddingProfiler()\n",
    "    \n",
    "    # Test different embedding configurations\n",
    "    print(\"\\n📊 EMBEDDING PERFORMANCE COMPARISON:\")\n",
    "    \n",
    "    # Compare embedding layers with different sizes\n",
    "    vocab_sizes = [1000, 5000, 10000]\n",
    "    embedding_dims = [128, 256, 512]\n",
    "    \n",
    "    scaling_results = profiler.analyze_memory_scaling(vocab_sizes, embedding_dims)\n",
    "    \n",
    "    # Compare positional encoding approaches\n",
    "    print(\"\\n\" + \"=\"*60)\n",
    "    pos_comparison = profiler.compare_positional_encodings(seq_length=128, embedding_dim=256)\n",
    "    \n",
    "    # Systems design analysis\n",
    "    print(\"\\n\" + \"=\"*60)\n",
    "    analyze_embedding_system_design()\n",
    "    \n",
    "    # Demonstrate realistic language model embedding setup\n",
    "    print(\"\\n\" + \"=\"*60)\n",
    "    print(\"🏗️ REALISTIC LANGUAGE MODEL EMBEDDING SETUP\")\n",
    "    print(\"=\"*60)\n",
    "    \n",
    "    # Create realistic configuration\n",
    "    vocab_size = 10000  # 10k vocabulary\n",
    "    embedding_dim = 256  # 256-dim embeddings\n",
    "    max_seq_length = 512  # 512 token sequences\n",
    "    \n",
    "    print(f\"Model configuration:\")\n",
    "    print(f\"  Vocabulary size: {vocab_size:,}\")\n",
    "    print(f\"  Embedding dimension: {embedding_dim}\")\n",
    "    print(f\"  Max sequence length: {max_seq_length}\")\n",
    "    \n",
    "    # Create components\n",
    "    embedding_layer = Embedding(vocab_size=vocab_size, embedding_dim=embedding_dim, padding_idx=0)\n",
    "    pos_encoding = PositionalEncoding(embedding_dim=embedding_dim, max_seq_length=max_seq_length)\n",
    "    \n",
    "    # Calculate memory requirements\n",
    "    embed_memory = embedding_layer.get_memory_usage()\n",
    "    \n",
    "    print(f\"\\nMemory analysis:\")\n",
    "    print(f\"  Embedding table: {embed_memory['total_memory_mb']:.1f}MB\")\n",
    "    print(f\"  Parameters: {embed_memory['total_parameters']:,}\")\n",
    "    \n",
    "    # Simulate batch processing\n",
    "    batch_size = 32\n",
    "    seq_length = 256\n",
    "    test_tokens = np.random.randint(0, vocab_size, (batch_size, seq_length))\n",
    "    \n",
    "    start_time = time.time()\n",
    "    embeddings = embedding_layer.forward(test_tokens)\n",
    "    pos_embeddings = pos_encoding.forward(embeddings)\n",
    "    total_time = time.time() - start_time\n",
    "    \n",
    "    sequence_memory_mb = pos_embeddings.data.nbytes / (1024 * 1024)\n",
    "    \n",
    "    print(f\"\\nBatch processing:\")\n",
    "    print(f\"  Batch size: {batch_size}, Sequence length: {seq_length}\")\n",
    "    print(f\"  Processing time: {total_time*1000:.2f}ms\")\n",
    "    print(f\"  Sequence memory: {sequence_memory_mb:.1f}MB\")\n",
    "    print(f\"  Throughput: {(batch_size * seq_length) / total_time:.0f} tokens/second\")\n",
    "    \n",
    "    print(\"\\n\" + \"=\"*60)\n",
    "    print(\"🎯 EMBEDDINGS MODULE COMPLETE!\")\n",
    "    print(\"=\"*60)\n",
    "    print(\"All embedding tests passed!\")\n",
    "    print(\"Ready for attention mechanism integration!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7dd5edd0",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🤔 ML Systems Thinking: Interactive Questions\n",
    "\n",
    "Now that you've built the embedding systems that convert tokens to rich vector representations, let's connect this work to broader ML systems challenges. These questions help you think critically about how embedding design scales to production language processing systems.\n",
    "\n",
    "Take time to reflect thoughtfully on each question - your insights will help you understand how embedding choices connect to real-world ML systems engineering."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae828478",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### Question 1: Embedding Memory Optimization and Model Scaling\n",
    "\n",
    "**Context**: Your embedding implementations demonstrate how vocabulary size and embedding dimension directly impact model parameters and memory usage. In production language models, embedding tables often contain billions of parameters (GPT-3's embedding table alone has ~600M parameters), making memory optimization critical for deployment and training efficiency.\n",
    "\n",
    "**Reflection Question**: Design a memory-optimized embedding system for a production language model that needs to handle a 100k vocabulary with 1024-dimensional embeddings while operating under GPU memory constraints. How would you implement embedding compression techniques, design efficient lookup patterns for high-throughput training, and handle dynamic vocabulary expansion for domain adaptation? Consider the challenges of maintaining embedding quality while reducing memory footprint and optimizing for both training and inference scenarios.\n",
    "\n",
    "Think about: embedding compression techniques, memory-efficient lookup patterns, dynamic vocabulary management, and quality-memory trade-offs.\n",
    "\n",
    "*Target length: 150-300 words*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd58f225",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "question-1-embedding-memory",
     "locked": false,
     "points": 10,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "YOUR REFLECTION ON EMBEDDING MEMORY OPTIMIZATION:\n",
    "\n",
    "TODO: Replace this text with your thoughtful response about memory-optimized embedding system design.\n",
    "\n",
    "Consider addressing:\n",
    "- How would you implement embedding compression for a 100k × 1024 vocabulary under GPU constraints?\n",
    "- What techniques would you use to optimize lookup patterns for high-throughput training?\n",
    "- How would you design dynamic vocabulary expansion while maintaining memory efficiency?\n",
    "- What trade-offs would you make between embedding quality and memory footprint?\n",
    "- How would you optimize differently for training vs inference scenarios?\n",
    "\n",
    "Write a technical analysis connecting your embedding implementations to real memory optimization challenges.\n",
    "\n",
    "GRADING RUBRIC (Instructor Use):\n",
    "- Demonstrates understanding of embedding memory scaling and optimization (3 points)\n",
    "- Designs practical approaches to compression and efficient lookup patterns (3 points)\n",
    "- Addresses dynamic vocabulary and quality-memory trade-offs (2 points)\n",
    "- Shows systems thinking about production memory constraints (2 points)\n",
    "- Clear technical reasoning with memory optimization insights (bonus points for innovative approaches)\n",
    "\"\"\"\n",
    "\n",
    "### BEGIN SOLUTION\n",
    "# Student response area - instructor will replace this section during grading setup\n",
    "# This is a manually graded question requiring technical analysis of embedding memory optimization\n",
    "# Students should demonstrate understanding of large-scale embedding systems and memory efficiency\n",
    "### END SOLUTION"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b9bdf8",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### Question 2: Positional Encoding and Sequence Length Scalability\n",
    "\n",
    "**Context**: Your positional encoding implementations show the trade-offs between fixed sinusoidal patterns and learned position embeddings. Production language models increasingly need to handle variable sequence lengths efficiently while maintaining consistent position representations across different tasks and deployment scenarios.\n",
    "\n",
    "**Reflection Question**: Architect a positional encoding system for a production transformer that needs to efficiently handle sequences ranging from 512 tokens (typical sentences) to 32k tokens (long documents) while maintaining training stability and inference efficiency. How would you design hybrid positional encoding that combines the benefits of sinusoidal and learned approaches, implement efficient position computation for variable-length sequences, and optimize for both memory usage and computational efficiency across different sequence length distributions?\n",
    "\n",
    "Think about: hybrid encoding strategies, variable-length optimization, memory-efficient position computation, and sequence length distribution handling.\n",
    "\n",
    "*Target length: 150-300 words*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d23a6a",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "question-2-positional-encoding",
     "locked": false,
     "points": 10,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "YOUR REFLECTION ON POSITIONAL ENCODING AND SEQUENCE SCALABILITY:\n",
    "\n",
    "TODO: Replace this text with your thoughtful response about scalable positional encoding system design.\n",
    "\n",
    "Consider addressing:\n",
    "- How would you design hybrid positional encoding for sequences from 512 to 32k tokens?\n",
    "- What strategies would you use to optimize position computation for variable-length sequences?\n",
    "- How would you balance memory efficiency with computational performance?\n",
    "- What approaches would you use to handle different sequence length distributions?\n",
    "- How would you maintain training stability across diverse sequence lengths?\n",
    "\n",
    "Write an architectural analysis connecting your positional encoding work to scalable sequence processing.\n",
    "\n",
    "GRADING RUBRIC (Instructor Use):\n",
    "- Shows understanding of positional encoding scalability challenges (3 points)\n",
    "- Designs practical approaches to hybrid encoding and variable-length optimization (3 points)\n",
    "- Addresses memory and computational efficiency considerations (2 points)\n",
    "- Demonstrates systems thinking about sequence length distribution handling (2 points)\n",
    "- Clear architectural reasoning with scalability insights (bonus points for comprehensive system design)\n",
    "\"\"\"\n",
    "\n",
    "### BEGIN SOLUTION\n",
    "# Student response area - instructor will replace this section during grading setup\n",
    "# This is a manually graded question requiring understanding of positional encoding scalability\n",
    "# Students should demonstrate knowledge of sequence length optimization and hybrid approaches\n",
    "### END SOLUTION"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b67ef1",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### Question 3: Embedding Pipeline Integration and Training Efficiency\n",
    "\n",
    "**Context**: Your embedding pipeline integration demonstrates how tokenization, embedding lookup, and positional encoding work together in language model preprocessing. In production training systems, the embedding pipeline often becomes a bottleneck due to memory bandwidth limitations and the need to process billions of tokens efficiently during training.\n",
    "\n",
    "**Reflection Question**: Design an embedding pipeline optimization strategy for large-scale language model training that processes 1 trillion tokens efficiently while maintaining high GPU utilization and minimizing memory bandwidth bottlenecks. How would you implement pipeline parallelism for embedding operations, optimize batch processing for mixed sequence lengths, and design efficient gradient updates for massive embedding tables? Consider the challenges of coordinating embedding updates across distributed training nodes while maintaining numerical stability and convergence.\n",
    "\n",
    "Think about: pipeline parallelism, batch optimization, gradient update efficiency, and distributed training coordination.\n",
    "\n",
    "*Target length: 150-300 words*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5394f77",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "question-3-pipeline-integration",
     "locked": false,
     "points": 10,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "YOUR REFLECTION ON EMBEDDING PIPELINE INTEGRATION:\n",
    "\n",
    "TODO: Replace this text with your thoughtful response about embedding pipeline optimization for large-scale training.\n",
    "\n",
    "Consider addressing:\n",
    "- How would you implement pipeline parallelism for processing 1 trillion tokens efficiently?\n",
    "- What strategies would you use to optimize batch processing for mixed sequence lengths?\n",
    "- How would you design efficient gradient updates for massive embedding tables?\n",
    "- What approaches would you use for coordinating embedding updates across distributed nodes?\n",
    "- How would you maintain GPU utilization while minimizing memory bandwidth bottlenecks?\n",
    "\n",
    "Write a design analysis connecting your embedding pipeline to large-scale training optimization.\n",
    "\n",
    "GRADING RUBRIC (Instructor Use):\n",
    "- Understands embedding pipeline bottlenecks and optimization challenges (3 points)\n",
    "- Designs practical approaches to pipeline parallelism and batch optimization (3 points)\n",
    "- Addresses distributed training and gradient update efficiency (2 points)\n",
    "- Shows systems thinking about large-scale training coordination (2 points)\n",
    "- Clear design reasoning with pipeline optimization insights (bonus points for innovative approaches)\n",
    "\"\"\"\n",
    "\n",
    "### BEGIN SOLUTION\n",
    "# Student response area - instructor will replace this section during grading setup\n",
    "# This is a manually graded question requiring understanding of large-scale embedding pipeline optimization\n",
    "# Students should demonstrate knowledge of distributed training and pipeline efficiency\n",
    "### END SOLUTION"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bd47eab",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 MODULE SUMMARY: Embeddings\n",
    "\n",
    "Congratulations! You have successfully implemented comprehensive embedding systems for language processing:\n",
    "\n",
    "### ✅ What You Have Built\n",
    "- **Embedding Layer**: Learnable lookup table converting tokens to dense vector representations\n",
    "- **Positional Encoding**: Sinusoidal position information for sequence understanding\n",
    "- **Learned Positional Embeddings**: Trainable position representations for model-specific optimization\n",
    "- **Memory-Efficient Lookups**: Optimized embedding access patterns for production systems\n",
    "- **Performance Analysis**: Comprehensive profiling and scaling analysis tools\n",
    "- **🆕 Integration Pipeline**: Complete tokenization → embedding → positional encoding workflow\n",
    "- **🆕 Systems Optimization**: Memory usage analysis and performance optimization techniques\n",
    "\n",
    "### ✅ Key Learning Outcomes\n",
    "- **Understanding**: How discrete tokens become continuous vector representations\n",
    "- **Implementation**: Built embedding systems from scratch with efficient lookup operations\n",
    "- **Systems Insight**: How embedding table size affects model memory and training efficiency\n",
    "- **Performance Engineering**: Measured and optimized embedding lookup patterns and memory usage\n",
    "- **Production Context**: Understanding real-world embedding challenges and optimization techniques\n",
    "\n",
    "### ✅ Technical Mastery\n",
    "- **Embedding Lookup**: Efficient table lookup with various initialization strategies\n",
    "- **Positional Encoding**: Mathematical sine/cosine patterns for position representation\n",
    "- **Memory Scaling**: Understanding O(vocab_size × embedding_dim) parameter scaling\n",
    "- **Performance Optimization**: Cache-friendly access patterns and memory bandwidth optimization\n",
    "- **🆕 Integration Design**: Seamless pipeline from text processing to vector representations\n",
    "\n",
    "### ✅ Professional Skills Developed\n",
    "- **Systems Architecture**: Designing embedding systems for production scale\n",
    "- **Memory Engineering**: Optimizing large parameter tables for efficient access\n",
    "- **Performance Analysis**: Measuring and improving embedding pipeline throughput\n",
    "- **Integration Thinking**: Connecting embedding systems with tokenization and attention\n",
    "\n",
    "### ✅ Ready for Next Steps\n",
    "Your embedding systems are now ready to power:\n",
    "- **Attention Mechanisms**: Processing sequence representations with attention\n",
    "- **Transformer Models**: Complete language model architectures\n",
    "- **Language Understanding**: Rich semantic representations for NLP tasks\n",
    "- **🧠 Sequence Processing**: Foundation for advanced sequence modeling\n",
    "\n",
    "### 🔗 Connection to Real ML Systems\n",
    "Your implementations mirror production systems:\n",
    "- **PyTorch Embeddings**: `torch.nn.Embedding` and `torch.nn.functional.embedding`\n",
    "- **Transformer Models**: All modern language models use similar embedding approaches\n",
    "- **Production Optimizations**: Memory mapping, gradient checkpointing, and distributed embeddings\n",
    "- **Industry Applications**: GPT, BERT, and other transformer models rely on these foundations\n",
    "\n",
    "### 🎯 The Power of Dense Representations\n",
    "You have unlocked the bridge between discrete tokens and continuous understanding:\n",
    "- **Before**: Tokens were sparse, discrete symbols\n",
    "- **After**: Tokens become rich, continuous vectors that capture semantic relationships\n",
    "\n",
    "**Next Module**: Attention - Processing sequences with the mechanism that revolutionized language understanding!\n",
    "\n",
    "Your embedding systems provide the rich vector representations that attention mechanisms need to understand language. Now let's build the attention that makes transformers work!"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
