{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c9bc6eb",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# DataLoader - Efficient Data Pipeline and Batch Processing Systems\n",
    "\n",
    "Welcome to the DataLoader module! You'll build the data infrastructure that feeds neural networks, understanding how I/O optimization and memory management determine training speed.\n",
    "\n",
    "## Learning Goals\n",
    "- Systems understanding: How data I/O becomes the bottleneck in ML training and why efficient data pipelines are critical for system performance\n",
    "- Core implementation skill: Build Dataset and DataLoader classes with batching, shuffling, and memory-efficient iteration patterns\n",
    "- Pattern recognition: Understand the universal Dataset/DataLoader abstraction used across all ML frameworks\n",
    "- Framework connection: See how your implementation mirrors PyTorch's data loading infrastructure and optimization strategies\n",
    "- Performance insight: Learn why data loading parallelization and prefetching are essential for GPU utilization in production systems\n",
    "\n",
    "## Build → Use → Reflect\n",
    "1. **Build**: Complete Dataset and DataLoader classes with efficient batching, shuffling, and real dataset support (CIFAR-10)\n",
    "2. **Use**: Load large-scale image datasets and feed them to neural networks with proper batch processing\n",
    "3. **Reflect**: Why does data loading speed often determine training speed more than model computation?\n",
    "\n",
    "## What You'll Achieve\n",
    "By the end of this module, you'll understand:\n",
    "- Deep technical understanding of how efficient data pipelines enable scalable ML training\n",
    "- Practical capability to build data loading systems that handle datasets larger than memory\n",
    "- Systems insight into why data engineering is often the limiting factor in ML system performance\n",
    "- Performance consideration of how batch size, shuffling, and prefetching affect training throughput and convergence\n",
    "- Connection to production ML systems and how frameworks optimize data loading for different storage systems\n",
    "\n",
    "## Systems Reality Check\n",
    "💡 **Production Context**: PyTorch's DataLoader uses multiprocessing and memory pinning to overlap data loading with GPU computation, achieving near-zero data loading overhead\n",
    "⚡ **Performance Note**: Modern GPUs can process data faster than storage systems can provide it - data loading optimization is critical for hardware utilization in production training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92c9d8b6",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "dataloader-imports",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| default_exp core.dataloader\n",
    "\n",
    "#| export\n",
    "import numpy as np\n",
    "import sys\n",
    "import os\n",
    "from typing import Tuple, Optional, Iterator\n",
    "import urllib.request\n",
    "import tarfile\n",
    "import pickle\n",
    "import time\n",
    "\n",
    "# Import our building blocks - try package first, then local modules\n",
    "try:\n",
    "    from tinytorch.core.tensor import Tensor\n",
    "except ImportError:\n",
    "    # For development, import from local modules\n",
    "    sys.path.append(os.path.join(os.path.dirname(__file__), '..', '01_tensor'))\n",
    "    from tensor_dev import Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2959209b",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "dataloader-welcome",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "print(\"🔥 TinyTorch DataLoader Module\")\n",
    "print(f\"NumPy version: {np.__version__}\")\n",
    "print(f\"Python version: {sys.version_info.major}.{sys.version_info.minor}\")\n",
    "print(\"Ready to build data pipelines!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f2d9467",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 📦 Where This Code Lives in the Final Package\n",
    "\n",
    "**Learning Side:** You work in `modules/source/06_dataloader/dataloader_dev.py`  \n",
    "**Building Side:** Code exports to `tinytorch.core.dataloader`\n",
    "\n",
    "```python\n",
    "# Final package structure:\n",
    "from tinytorch.core.dataloader import Dataset, DataLoader  # Data loading utilities!\n",
    "from tinytorch.core.tensor import Tensor  # Foundation\n",
    "from tinytorch.core.networks import Sequential  # Models to train\n",
    "```\n",
    "\n",
    "**Why this matters:**\n",
    "- **Learning:** Focused modules for deep understanding of data pipelines\n",
    "- **Production:** Proper organization like PyTorch's `torch.utils.data`\n",
    "- **Consistency:** All data loading utilities live together in `core.dataloader`\n",
    "- **Integration:** Works seamlessly with tensors and networks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b07e46b",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🔧 DEVELOPMENT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52c9b734",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## Step 1: Understanding Data Pipelines\n",
    "\n",
    "### What are Data Pipelines?\n",
    "**Data pipelines** are the systems that efficiently move data from storage to your model. They're the foundation of all machine learning systems.\n",
    "\n",
    "### The Data Pipeline Equation\n",
    "```\n",
    "Raw Data → Load → Transform → Batch → Model → Predictions\n",
    "```\n",
    "\n",
    "### Why Data Pipelines Matter\n",
    "- **Performance**: Efficient loading prevents GPU starvation\n",
    "- **Scalability**: Handle datasets larger than memory\n",
    "- **Consistency**: Reproducible data processing\n",
    "- **Flexibility**: Easy to switch between datasets\n",
    "\n",
    "### Real-World Challenges\n",
    "- **Memory constraints**: Datasets often exceed available RAM\n",
    "- **I/O bottlenecks**: Disk access is much slower than computation\n",
    "- **Batch processing**: Neural networks need batched data for efficiency\n",
    "- **Shuffling**: Random order prevents overfitting\n",
    "\n",
    "### Systems Thinking\n",
    "- **Memory efficiency**: Handle datasets larger than RAM\n",
    "- **I/O optimization**: Read from disk efficiently\n",
    "- **Batching strategies**: Trade-offs between memory and speed\n",
    "- **Caching**: When to cache vs recompute\n",
    "\n",
    "### Visual Intuition\n",
    "```\n",
    "Raw Files: [image1.jpg, image2.jpg, image3.jpg, ...]\n",
    "Load: [Tensor(32x32x3), Tensor(32x32x3), Tensor(32x32x3), ...]\n",
    "Batch: [Tensor(32, 32, 32, 3)]  # 32 images at once\n",
    "Model: Process batch efficiently\n",
    "```\n",
    "\n",
    "Let's start by building the most fundamental component: **Dataset**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d07094e6",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 2: Building the Dataset Interface\n",
    "\n",
    "### What is a Dataset?\n",
    "A **Dataset** is an abstract interface that provides consistent access to data. It's the foundation of all data loading systems.\n",
    "\n",
    "### Why Abstract Interfaces Matter\n",
    "- **Consistency**: Same interface for all data types\n",
    "- **Flexibility**: Easy to switch between datasets\n",
    "- **Testability**: Easy to create test datasets\n",
    "- **Extensibility**: Easy to add new data sources\n",
    "\n",
    "### The Dataset Pattern\n",
    "```python\n",
    "class Dataset:\n",
    "    def __getitem__(self, index):  # Get single sample\n",
    "        return data, label\n",
    "    \n",
    "    def __len__(self):  # Get dataset size\n",
    "        return total_samples\n",
    "```\n",
    "\n",
    "### Real-World Usage\n",
    "- **Computer vision**: ImageNet, CIFAR-10, custom image datasets\n",
    "- **NLP**: Text datasets, tokenized sequences\n",
    "- **Audio**: Audio files, spectrograms\n",
    "- **Time series**: Sequential data with proper windowing\n",
    "\n",
    "Let's implement the Dataset interface!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "275c4926",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "dataset-class",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class Dataset:\n",
    "    \"\"\"\n",
    "    Base Dataset class: Abstract interface for all datasets.\n",
    "    \n",
    "    The fundamental abstraction for data loading in TinyTorch.\n",
    "    Students implement concrete datasets by inheriting from this class.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __getitem__(self, index: int) -> Tuple[Tensor, Tensor]:\n",
    "        \"\"\"\n",
    "        Get a single sample and label by index.\n",
    "        \n",
    "        Args:\n",
    "            index: Index of the sample to retrieve\n",
    "            \n",
    "        Returns:\n",
    "            Tuple of (data, label) tensors\n",
    "            \n",
    "        TODO: Implement abstract method for getting samples.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. This is an abstract method - subclasses will implement it\n",
    "        2. Return a tuple of (data, label) tensors\n",
    "        3. Data should be the input features, label should be the target\n",
    "        \n",
    "        EXAMPLE:\n",
    "        dataset[0] should return (Tensor(image_data), Tensor(label))\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - **PyTorch Integration**: This follows the exact same pattern as torch.utils.data.Dataset\n",
    "        - **Production Data**: Real datasets like ImageNet, CIFAR-10 use this interface\n",
    "        - **Memory Efficiency**: On-demand loading prevents loading entire dataset into memory\n",
    "        - **Batching Foundation**: DataLoader uses __getitem__ to create batches efficiently\n",
    "        \n",
    "        HINTS:\n",
    "        - This is an abstract method that subclasses must override\n",
    "        - Always return a tuple of (data, label) tensors\n",
    "        - Data contains the input features, label contains the target\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # This is an abstract method - subclasses must implement it\n",
    "        raise NotImplementedError(\"Subclasses must implement __getitem__\")\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __len__(self) -> int:\n",
    "        \"\"\"\n",
    "        Get the total number of samples in the dataset.\n",
    "        \n",
    "        TODO: Implement abstract method for getting dataset size.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. This is an abstract method - subclasses will implement it\n",
    "        2. Return the total number of samples in the dataset\n",
    "        \n",
    "        EXAMPLE:\n",
    "        len(dataset) should return 50000 for CIFAR-10 training set\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - **Memory Planning**: DataLoader uses len() to calculate number of batches\n",
    "        - **Progress Tracking**: Training loops use len() for progress bars and epoch calculations\n",
    "        - **Distributed Training**: Multi-GPU systems need dataset size for work distribution\n",
    "        - **Statistical Sampling**: Some training strategies require knowing total dataset size\n",
    "        \n",
    "        HINTS:\n",
    "        - This is an abstract method that subclasses must override\n",
    "        - Return an integer representing the total number of samples\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # This is an abstract method - subclasses must implement it\n",
    "        raise NotImplementedError(\"Subclasses must implement __len__\")\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def get_sample_shape(self) -> Tuple[int, ...]:\n",
    "        \"\"\"\n",
    "        Get the shape of a single data sample.\n",
    "        \n",
    "        TODO: Implement method to get sample shape.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Get the first sample using self[0]\n",
    "        2. Extract the data part (first element of tuple)\n",
    "        3. Return the shape of the data tensor\n",
    "        \n",
    "        EXAMPLE:\n",
    "        For CIFAR-10: returns (3, 32, 32) for RGB images\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - **Model Architecture**: Neural networks need to know input shape for first layer\n",
    "        - **Batch Planning**: Systems use sample shape to calculate memory requirements\n",
    "        - **Preprocessing Validation**: Ensures all samples have consistent shape\n",
    "        - **Framework Integration**: Similar to PyTorch's dataset shape inspection\n",
    "        \n",
    "        HINTS:\n",
    "        - Use self[0] to get the first sample\n",
    "        - Extract data from the (data, label) tuple\n",
    "        - Return data.shape\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Get the first sample to determine shape\n",
    "        data, _ = self[0]\n",
    "        return data.shape\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def get_num_classes(self) -> int:\n",
    "        \"\"\"\n",
    "        Get the number of classes in the dataset.\n",
    "        \n",
    "        TODO: Implement abstract method for getting number of classes.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. This is an abstract method - subclasses will implement it\n",
    "        2. Return the number of unique classes in the dataset\n",
    "        \n",
    "        EXAMPLE:\n",
    "        For CIFAR-10: returns 10 (classes 0-9)\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - **Output Layer Design**: Neural networks need num_classes for final layer size\n",
    "        - **Loss Function Setup**: CrossEntropyLoss uses num_classes for proper computation\n",
    "        - **Evaluation Metrics**: Accuracy calculation depends on number of classes\n",
    "        - **Model Validation**: Ensures model predictions match expected class range\n",
    "        \n",
    "        HINTS:\n",
    "        - This is an abstract method that subclasses must override\n",
    "        - Return the number of unique classes/categories\n",
    "        \"\"\"\n",
    "        # This is an abstract method - subclasses must implement it\n",
    "        raise NotImplementedError(\"Subclasses must implement get_num_classes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06c34e75",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### 🧪 Unit Test: Dataset Interface\n",
    "\n",
    "Let's understand the Dataset interface! While we can't test the abstract class directly, we'll create a simple test dataset.\n",
    "\n",
    "**This is a unit test** - it tests the Dataset interface pattern in isolation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e349589",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": true,
     "grade_id": "test-dataset-interface-immediate",
     "locked": true,
     "points": 5,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test Dataset interface with a simple implementation\n",
    "print(\"🔬 Unit Test: Dataset Interface...\")\n",
    "\n",
    "# Create a minimal test dataset\n",
    "class TestDataset(Dataset):\n",
    "    def __init__(self, size=5):\n",
    "        self.size = size\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        # Simple test data: features are [index, index*2], label is index % 2\n",
    "        data = Tensor([index, index * 2])\n",
    "        label = Tensor([index % 2])\n",
    "        return data, label\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "    \n",
    "    def get_num_classes(self):\n",
    "        return 2\n",
    "\n",
    "# Test the interface (moved to main block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "261ad6cc",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 3: Building the DataLoader\n",
    "\n",
    "### What is a DataLoader?\n",
    "A **DataLoader** efficiently batches and iterates through datasets. It's the bridge between individual samples and the batched data that neural networks expect.\n",
    "\n",
    "### Why DataLoaders Matter\n",
    "- **Batching**: Groups samples for efficient GPU computation\n",
    "- **Shuffling**: Randomizes data order to prevent overfitting\n",
    "- **Memory efficiency**: Loads data on-demand rather than all at once\n",
    "- **Iteration**: Provides clean interface for training loops\n",
    "\n",
    "### The DataLoader Pattern\n",
    "```python\n",
    "DataLoader(dataset, batch_size=32, shuffle=True)\n",
    "for batch_data, batch_labels in dataloader:\n",
    "    # batch_data.shape: (32, ...)\n",
    "    # batch_labels.shape: (32,)\n",
    "    # Train on batch\n",
    "```\n",
    "\n",
    "### Real-World Applications\n",
    "- **Training loops**: Feed batches to neural networks\n",
    "- **Validation**: Evaluate models on held-out data\n",
    "- **Inference**: Process large datasets efficiently\n",
    "- **Data analysis**: Explore datasets systematically\n",
    "\n",
    "### Systems Thinking\n",
    "- **Batch size**: Trade-off between memory and speed\n",
    "- **Shuffling**: Prevents overfitting to data order\n",
    "- **Iteration**: Efficient looping through data\n",
    "- **Memory**: Manage large datasets that don't fit in RAM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7607154",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "dataloader-class",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class DataLoader:\n",
    "    \"\"\"\n",
    "    DataLoader: Efficiently batch and iterate through datasets.\n",
    "    \n",
    "    Provides batching, shuffling, and efficient iteration over datasets.\n",
    "    Essential for training neural networks efficiently.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, dataset: Dataset, batch_size: int = 32, shuffle: bool = True):\n",
    "        \"\"\"\n",
    "        Initialize DataLoader.\n",
    "        \n",
    "        Args:\n",
    "            dataset: Dataset to load from\n",
    "            batch_size: Number of samples per batch\n",
    "            shuffle: Whether to shuffle data each epoch\n",
    "            \n",
    "        TODO: Store configuration and dataset.\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Store dataset as self.dataset\n",
    "        2. Store batch_size as self.batch_size\n",
    "        3. Store shuffle as self.shuffle\n",
    "        \n",
    "        EXAMPLE:\n",
    "        DataLoader(dataset, batch_size=32, shuffle=True)\n",
    "        \n",
    "        HINTS:\n",
    "        - Store all parameters as instance variables\n",
    "        - These will be used in __iter__ for batching\n",
    "        \"\"\"\n",
    "        # Input validation\n",
    "        if dataset is None:\n",
    "            raise TypeError(\"Dataset cannot be None\")\n",
    "        if not isinstance(batch_size, int) or batch_size <= 0:\n",
    "            raise ValueError(f\"Batch size must be a positive integer, got {batch_size}\")\n",
    "        \n",
    "        self.dataset = dataset\n",
    "        self.batch_size = batch_size\n",
    "        self.shuffle = shuffle\n",
    "    \n",
    "    def __iter__(self) -> Iterator[Tuple[Tensor, Tensor]]:\n",
    "        \"\"\"\n",
    "        Iterate through dataset in batches.\n",
    "        \n",
    "        Returns:\n",
    "            Iterator yielding (batch_data, batch_labels) tuples\n",
    "            \n",
    "        TODO: Implement batching and shuffling logic.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Create indices list: list(range(len(dataset)))\n",
    "        2. Shuffle indices if self.shuffle is True\n",
    "        3. Loop through indices in batch_size chunks\n",
    "        4. For each batch: collect samples, stack them, yield batch\n",
    "        \n",
    "        EXAMPLE:\n",
    "        for batch_data, batch_labels in dataloader:\n",
    "            # batch_data.shape: (batch_size, ...)\n",
    "            # batch_labels.shape: (batch_size,)\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - **GPU Efficiency**: Batching maximizes GPU utilization by processing multiple samples together\n",
    "        - **Training Stability**: Shuffling prevents overfitting to data order and improves generalization\n",
    "        - **Memory Management**: Batches fit in GPU memory while full dataset may not\n",
    "        - **Gradient Estimation**: Batch gradients provide better estimates than single-sample gradients\n",
    "        \n",
    "        HINTS:\n",
    "        - Use list(range(len(self.dataset))) for indices\n",
    "        - Use np.random.shuffle() if self.shuffle is True\n",
    "        - Loop in chunks of self.batch_size\n",
    "        - Collect samples and stack with np.stack()\n",
    "        \"\"\"\n",
    "        # Create indices for all samples\n",
    "        indices = list(range(len(self.dataset)))\n",
    "        \n",
    "        # Shuffle if requested\n",
    "        if self.shuffle:\n",
    "            np.random.shuffle(indices)\n",
    "        \n",
    "        # Iterate through indices in batches\n",
    "        for i in range(0, len(indices), self.batch_size):\n",
    "            batch_indices = indices[i:i + self.batch_size]\n",
    "            \n",
    "            # Collect samples for this batch\n",
    "            batch_data = []\n",
    "            batch_labels = []\n",
    "            \n",
    "            for idx in batch_indices:\n",
    "                data, label = self.dataset[idx]\n",
    "                batch_data.append(data.data)\n",
    "                batch_labels.append(label.data)\n",
    "            \n",
    "            # Stack into batch tensors\n",
    "            batch_data_array = np.stack(batch_data, axis=0)\n",
    "            batch_labels_array = np.stack(batch_labels, axis=0)\n",
    "            \n",
    "            yield Tensor(batch_data_array), Tensor(batch_labels_array)\n",
    "    \n",
    "    def __len__(self) -> int:\n",
    "        \"\"\"\n",
    "        Get the number of batches per epoch.\n",
    "        \n",
    "        TODO: Calculate number of batches.\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Get dataset size: len(self.dataset)\n",
    "        2. Divide by batch_size and round up\n",
    "        3. Use ceiling division: (n + batch_size - 1) // batch_size\n",
    "        \n",
    "        EXAMPLE:\n",
    "        Dataset size 100, batch size 32 → 4 batches\n",
    "        \n",
    "        HINTS:\n",
    "        - Use len(self.dataset) for dataset size\n",
    "        - Use ceiling division for exact batch count\n",
    "        - Formula: (dataset_size + batch_size - 1) // batch_size\n",
    "        \"\"\"\n",
    "        # Calculate number of batches using ceiling division\n",
    "        dataset_size = len(self.dataset)\n",
    "        return (dataset_size + self.batch_size - 1) // self.batch_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec802471",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### 🧪 Unit Test: DataLoader\n",
    "\n",
    "Let's test your DataLoader implementation! This is the heart of efficient data loading for neural networks.\n",
    "\n",
    "**This is a unit test** - it tests the DataLoader class in isolation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb2f9065",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-dataloader-immediate",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test DataLoader immediately after implementation\n",
    "print(\"🔬 Unit Test: DataLoader...\")\n",
    "\n",
    "# Use the test dataset from before\n",
    "class TestDataset(Dataset):\n",
    "    def __init__(self, size=10):\n",
    "        self.size = size\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        data = Tensor([index, index * 2])\n",
    "        label = Tensor([index % 3])  # 3 classes\n",
    "        return data, label\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.size\n",
    "    \n",
    "    def get_num_classes(self):\n",
    "        return 3\n",
    "\n",
    "# Test basic DataLoader functionality\n",
    "try:\n",
    "    dataset = TestDataset(size=10)\n",
    "    dataloader = DataLoader(dataset, batch_size=3, shuffle=False)\n",
    "    \n",
    "    print(f\"DataLoader created: batch_size={dataloader.batch_size}, shuffle={dataloader.shuffle}\")\n",
    "    print(f\"Number of batches: {len(dataloader)}\")\n",
    "    \n",
    "    # Test __len__\n",
    "    expected_batches = (10 + 3 - 1) // 3  # Ceiling division: 4 batches\n",
    "    assert len(dataloader) == expected_batches, f\"Should have {expected_batches} batches, got {len(dataloader)}\"\n",
    "    print(\"✅ DataLoader __len__ works correctly\")\n",
    "    \n",
    "    # Test iteration\n",
    "    batch_count = 0\n",
    "    total_samples = 0\n",
    "    \n",
    "    for batch_data, batch_labels in dataloader:\n",
    "        batch_count += 1\n",
    "        batch_size = batch_data.shape[0]\n",
    "        total_samples += batch_size\n",
    "        \n",
    "        print(f\"Batch {batch_count}: data shape {batch_data.shape}, labels shape {batch_labels.shape}\")\n",
    "        \n",
    "        # Verify batch dimensions\n",
    "        assert len(batch_data.shape) == 2, f\"Batch data should be 2D, got {batch_data.shape}\"\n",
    "        assert len(batch_labels.shape) == 2, f\"Batch labels should be 2D, got {batch_labels.shape}\"\n",
    "        assert batch_data.shape[1] == 2, f\"Each sample should have 2 features, got {batch_data.shape[1]}\"\n",
    "        assert batch_labels.shape[1] == 1, f\"Each label should have 1 element, got {batch_labels.shape[1]}\"\n",
    "        \n",
    "    assert batch_count == expected_batches, f\"Should iterate {expected_batches} times, got {batch_count}\"\n",
    "    assert total_samples == 10, f\"Should process 10 total samples, got {total_samples}\"\n",
    "    print(\"✅ DataLoader iteration works correctly\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"❌ DataLoader test failed: {e}\")\n",
    "    raise\n",
    "\n",
    "# Test shuffling\n",
    "try:\n",
    "    dataloader_shuffle = DataLoader(dataset, batch_size=5, shuffle=True)\n",
    "    dataloader_no_shuffle = DataLoader(dataset, batch_size=5, shuffle=False)\n",
    "    \n",
    "    # Get first batch from each\n",
    "    batch1_shuffle = next(iter(dataloader_shuffle))\n",
    "    batch1_no_shuffle = next(iter(dataloader_no_shuffle))\n",
    "    \n",
    "    print(\"✅ DataLoader shuffling parameter works\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"❌ DataLoader shuffling test failed: {e}\")\n",
    "    raise\n",
    "\n",
    "# Test different batch sizes\n",
    "try:\n",
    "    small_loader = DataLoader(dataset, batch_size=2, shuffle=False)\n",
    "    large_loader = DataLoader(dataset, batch_size=8, shuffle=False)\n",
    "    \n",
    "    assert len(small_loader) == 5, f\"Small loader should have 5 batches, got {len(small_loader)}\"\n",
    "    assert len(large_loader) == 2, f\"Large loader should have 2 batches, got {len(large_loader)}\"\n",
    "    print(\"✅ DataLoader handles different batch sizes correctly\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"❌ DataLoader batch size test failed: {e}\")\n",
    "    raise\n",
    "\n",
    "# Show the DataLoader behavior\n",
    "print(\"🎯 DataLoader behavior:\")\n",
    "print(\"   Batches data for efficient processing\")\n",
    "print(\"   Handles shuffling and iteration\")\n",
    "print(\"   Provides clean interface for training loops\")\n",
    "print(\"📈 Progress: Dataset interface ✓, DataLoader ✓\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a834dfd9",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 4: Creating a Simple Dataset Example\n",
    "\n",
    "### Why We Need Concrete Examples\n",
    "Abstract classes are great for interfaces, but we need concrete implementations to understand how they work. Let's create a simple dataset for testing.\n",
    "\n",
    "### Design Principles\n",
    "- **Simple**: Easy to understand and debug\n",
    "- **Configurable**: Adjustable size and properties\n",
    "- **Predictable**: Deterministic data for testing\n",
    "- **Educational**: Shows the Dataset pattern clearly\n",
    "\n",
    "### Real-World Connection\n",
    "This pattern is used for:\n",
    "- **CIFAR-10**: 32x32 RGB images with 10 classes\n",
    "- **ImageNet**: High-resolution images with 1000 classes\n",
    "- **MNIST**: 28x28 grayscale digits with 10 classes\n",
    "- **Custom datasets**: Your own data following this pattern"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e77a02",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "simple-dataset",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class SimpleDataset(Dataset):\n",
    "    \"\"\"\n",
    "    Simple dataset for testing and demonstration.\n",
    "    \n",
    "    Generates synthetic data with configurable size and properties.\n",
    "    Perfect for understanding the Dataset pattern.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, size: int = 100, num_features: int = 4, num_classes: int = 3):\n",
    "        \"\"\"\n",
    "        Initialize SimpleDataset.\n",
    "        \n",
    "        Args:\n",
    "            size: Number of samples in the dataset\n",
    "            num_features: Number of features per sample\n",
    "            num_classes: Number of classes\n",
    "            \n",
    "        TODO: Initialize the dataset with synthetic data.\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Store the configuration parameters\n",
    "        2. Generate synthetic data and labels\n",
    "        3. Make data deterministic for testing\n",
    "        \n",
    "        EXAMPLE:\n",
    "        SimpleDataset(size=100, num_features=4, num_classes=3)\n",
    "        creates 100 samples with 4 features each, 3 classes\n",
    "        \n",
    "        HINTS:\n",
    "        - Store size, num_features, num_classes as instance variables\n",
    "        - Use np.random.seed() for reproducible data\n",
    "        - Generate random data with np.random.randn()\n",
    "        - Generate random labels with np.random.randint()\n",
    "        \"\"\"\n",
    "        self.size = size\n",
    "        self.num_features = num_features\n",
    "        self.num_classes = num_classes\n",
    "        \n",
    "        # Generate synthetic data (deterministic for testing)\n",
    "        np.random.seed(42)  # For reproducible data\n",
    "        self.data = np.random.randn(size, num_features).astype(np.float32)\n",
    "        self.labels = np.random.randint(0, num_classes, size=size)\n",
    "    \n",
    "    def __getitem__(self, index: int) -> Tuple[Tensor, Tensor]:\n",
    "        \"\"\"\n",
    "        Get a sample by index.\n",
    "        \n",
    "        Args:\n",
    "            index: Index of the sample\n",
    "            \n",
    "        Returns:\n",
    "            Tuple of (data, label) tensors\n",
    "            \n",
    "        TODO: Return the sample at the given index.\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Get data sample from self.data[index]\n",
    "        2. Get label from self.labels[index]\n",
    "        3. Convert both to Tensors and return as tuple\n",
    "        \n",
    "        EXAMPLE:\n",
    "        dataset[0] returns (Tensor(features), Tensor(label))\n",
    "        \n",
    "        HINTS:\n",
    "        - Use self.data[index] for the data\n",
    "        - Use self.labels[index] for the label\n",
    "        - Convert to Tensors: Tensor(data), Tensor(label)\n",
    "        \"\"\"\n",
    "        data = self.data[index]\n",
    "        label = self.labels[index]\n",
    "        return Tensor(data), Tensor(label)\n",
    "    \n",
    "    def __len__(self) -> int:\n",
    "        \"\"\"\n",
    "        Get the dataset size.\n",
    "        \n",
    "        TODO: Return the dataset size.\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Return self.size\n",
    "        \n",
    "        EXAMPLE:\n",
    "        len(dataset) returns 100 for dataset with 100 samples\n",
    "        \n",
    "        HINTS:\n",
    "        - Simply return self.size\n",
    "        \"\"\"\n",
    "        return self.size\n",
    "    \n",
    "    def get_num_classes(self) -> int:\n",
    "        \"\"\"\n",
    "        Get the number of classes.\n",
    "        \n",
    "        TODO: Return the number of classes.\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Return self.num_classes\n",
    "        \n",
    "        EXAMPLE:\n",
    "        dataset.get_num_classes() returns 3 for 3-class dataset\n",
    "        \n",
    "        HINTS:\n",
    "        - Simply return self.num_classes\n",
    "        \"\"\"\n",
    "        return self.num_classes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b88878e6",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 4b: CIFAR-10 Dataset - Real Data for CNNs\n",
    "\n",
    "### Download and Load Real Computer Vision Data\n",
    "Let's implement loading CIFAR-10, the dataset we'll use to achieve our north star goal of 75% accuracy!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "417df9df",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "cifar10",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "def download_cifar10(root: str = \"./data\") -> str:\n",
    "    \"\"\"\n",
    "    Download CIFAR-10 dataset.\n",
    "    \n",
    "    TODO: Download and extract CIFAR-10.\n",
    "    \n",
    "    HINTS:\n",
    "    - URL: https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz\n",
    "    - Use urllib.request.urlretrieve()\n",
    "    - Extract with tarfile\n",
    "    \"\"\"\n",
    "    ### BEGIN SOLUTION\n",
    "    os.makedirs(root, exist_ok=True)\n",
    "    dataset_dir = os.path.join(root, \"cifar-10-batches-py\")\n",
    "    \n",
    "    if os.path.exists(dataset_dir):\n",
    "        print(f\"✅ CIFAR-10 found at {dataset_dir}\")\n",
    "        return dataset_dir\n",
    "    \n",
    "    url = \"https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz\"\n",
    "    tar_path = os.path.join(root, \"cifar-10.tar.gz\")\n",
    "    \n",
    "    print(f\"📥 Downloading CIFAR-10 (~170MB)...\")\n",
    "    urllib.request.urlretrieve(url, tar_path)\n",
    "    print(\"✅ Downloaded!\")\n",
    "    \n",
    "    print(\"📦 Extracting...\")\n",
    "    with tarfile.open(tar_path, 'r:gz') as tar:\n",
    "        tar.extractall(root)\n",
    "    print(\"✅ Ready!\")\n",
    "    \n",
    "    return dataset_dir\n",
    "    ### END SOLUTION\n",
    "\n",
    "class CIFAR10Dataset(Dataset):\n",
    "    \"\"\"CIFAR-10 dataset for CNN training.\"\"\"\n",
    "    \n",
    "    def __init__(self, root=\"./data\", train=True, download=False):\n",
    "        \"\"\"Load CIFAR-10 data.\"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        if download:\n",
    "            dataset_dir = download_cifar10(root)\n",
    "        else:\n",
    "            dataset_dir = os.path.join(root, \"cifar-10-batches-py\")\n",
    "        \n",
    "        if train:\n",
    "            data_list = []\n",
    "            label_list = []\n",
    "            for i in range(1, 6):\n",
    "                with open(os.path.join(dataset_dir, f\"data_batch_{i}\"), 'rb') as f:\n",
    "                    batch = pickle.load(f, encoding='bytes')\n",
    "                    data_list.append(batch[b'data'])\n",
    "                    label_list.extend(batch[b'labels'])\n",
    "            self.data = np.concatenate(data_list)\n",
    "            self.labels = np.array(label_list)\n",
    "        else:\n",
    "            with open(os.path.join(dataset_dir, \"test_batch\"), 'rb') as f:\n",
    "                batch = pickle.load(f, encoding='bytes')\n",
    "                self.data = batch[b'data']\n",
    "                self.labels = np.array(batch[b'labels'])\n",
    "        \n",
    "        # Reshape to (N, 3, 32, 32) and normalize\n",
    "        self.data = self.data.reshape(-1, 3, 32, 32).astype(np.float32) / 255.0\n",
    "        print(f\"✅ Loaded {len(self.data):,} images\")\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return Tensor(self.data[idx]), Tensor(self.labels[idx])\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def get_num_classes(self):\n",
    "        return 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "480db551",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### 🧪 Unit Test: SimpleDataset\n",
    "\n",
    "Let's test your SimpleDataset implementation! This concrete example shows how the Dataset pattern works.\n",
    "\n",
    "**This is a unit test** - it tests the SimpleDataset class in isolation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e73cdb0",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-simple-dataset-immediate",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test SimpleDataset immediately after implementation\n",
    "print(\"🔬 Unit Test: SimpleDataset...\")\n",
    "\n",
    "try:\n",
    "    # Create dataset\n",
    "    dataset = SimpleDataset(size=20, num_features=5, num_classes=4)\n",
    "    \n",
    "    print(f\"Dataset created: size={len(dataset)}, features={dataset.num_features}, classes={dataset.get_num_classes()}\")\n",
    "        \n",
    "        # Test basic properties\n",
    "    assert len(dataset) == 20, f\"Dataset length should be 20, got {len(dataset)}\"\n",
    "    assert dataset.get_num_classes() == 4, f\"Should have 4 classes, got {dataset.get_num_classes()}\"\n",
    "    print(\"✅ SimpleDataset basic properties work correctly\")\n",
    "        \n",
    "    # Test sample access\n",
    "    data, label = dataset[0]\n",
    "    assert isinstance(data, Tensor), \"Data should be a Tensor\"\n",
    "    assert isinstance(label, Tensor), \"Label should be a Tensor\"\n",
    "    assert data.shape == (5,), f\"Data shape should be (5,), got {data.shape}\"\n",
    "    assert label.shape == (), f\"Label shape should be (), got {label.shape}\"\n",
    "    print(\"✅ SimpleDataset sample access works correctly\")\n",
    "    \n",
    "    # Test sample shape\n",
    "    sample_shape = dataset.get_sample_shape()\n",
    "    assert sample_shape == (5,), f\"Sample shape should be (5,), got {sample_shape}\"\n",
    "    print(\"✅ SimpleDataset get_sample_shape works correctly\")\n",
    "    \n",
    "    # Test multiple samples\n",
    "    for i in range(5):\n",
    "            data, label = dataset[i]\n",
    "            assert data.shape == (5,), f\"Data shape should be (5,) for sample {i}, got {data.shape}\"\n",
    "            assert 0 <= label.data < 4, f\"Label should be in [0, 3] for sample {i}, got {label.data}\"\n",
    "    print(\"✅ SimpleDataset multiple samples work correctly\")\n",
    "    \n",
    "    # Test deterministic data (same seed should give same data)\n",
    "    dataset2 = SimpleDataset(size=20, num_features=5, num_classes=4)\n",
    "    data1, label1 = dataset[0]\n",
    "    data2, label2 = dataset2[0]\n",
    "    assert np.array_equal(data1.data, data2.data), \"Data should be deterministic\"\n",
    "    assert np.array_equal(label1.data, label2.data), \"Labels should be deterministic\"\n",
    "    print(\"✅ SimpleDataset data is deterministic\")\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"❌ SimpleDataset test failed: {e}\")\n",
    "\n",
    "# Show the SimpleDataset behavior\n",
    "print(\"🎯 SimpleDataset behavior:\")\n",
    "print(\"   Generates synthetic data for testing\")\n",
    "print(\"   Implements complete Dataset interface\")\n",
    "print(\"   Provides deterministic data for reproducibility\")\n",
    "print(\"📈 Progress: Dataset interface ✓, DataLoader ✓, SimpleDataset ✓\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "243297c6",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## Step 5: Comprehensive Test - Complete Data Pipeline\n",
    "\n",
    "### Real-World Data Pipeline Applications\n",
    "Let's test our data loading components in realistic scenarios:\n",
    "\n",
    "#### **Training Pipeline**\n",
    "```python\n",
    "# The standard ML training pattern\n",
    "dataset = SimpleDataset(size=1000, num_features=10, num_classes=5)\n",
    "dataloader = DataLoader(dataset, batch_size=32, shuffle=True)\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    for batch_data, batch_labels in dataloader:\n",
    "        # Train model on batch\n",
    "        pass\n",
    "```\n",
    "\n",
    "#### **Validation Pipeline**\n",
    "```python\n",
    "# Validation without shuffling\n",
    "val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False)\n",
    "\n",
    "for batch_data, batch_labels in val_loader:\n",
    "    # Evaluate model on batch\n",
    "    pass\n",
    "```\n",
    "\n",
    "#### **Data Analysis Pipeline**\n",
    "```python\n",
    "# Systematic data exploration\n",
    "for batch_data, batch_labels in dataloader:\n",
    "    # Analyze batch statistics\n",
    "    pass\n",
    "```\n",
    "\n",
    "This comprehensive test ensures our data loading components work together for real ML applications!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c994c580",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-comprehensive",
     "locked": true,
     "points": 15,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Comprehensive test - complete data pipeline applications\n",
    "print(\"🔬 Comprehensive Test: Complete Data Pipeline...\")\n",
    "\n",
    "try:\n",
    "    # Test 1: Training Data Pipeline\n",
    "    print(\"\\n1. Training Data Pipeline Test:\")\n",
    "    \n",
    "    # Create training dataset\n",
    "    train_dataset = SimpleDataset(size=100, num_features=8, num_classes=5)\n",
    "    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)\n",
    "    \n",
    "    # Simulate training epoch\n",
    "    epoch_samples = 0\n",
    "    epoch_batches = 0\n",
    "    \n",
    "    for batch_data, batch_labels in train_loader:\n",
    "        epoch_batches += 1\n",
    "        epoch_samples += batch_data.shape[0]\n",
    "        \n",
    "        # Verify batch properties\n",
    "        assert batch_data.shape[1] == 8, f\"Features should be 8, got {batch_data.shape[1]}\"\n",
    "        assert len(batch_labels.shape) == 1, f\"Labels should be 1D, got shape {batch_labels.shape}\"\n",
    "        assert isinstance(batch_data, Tensor), \"Batch data should be Tensor\"\n",
    "        assert isinstance(batch_labels, Tensor), \"Batch labels should be Tensor\"\n",
    "    \n",
    "    assert epoch_samples == 100, f\"Should process 100 samples, got {epoch_samples}\"\n",
    "    expected_batches = (100 + 16 - 1) // 16\n",
    "    assert epoch_batches == expected_batches, f\"Should have {expected_batches} batches, got {epoch_batches}\"\n",
    "    print(\"✅ Training pipeline works correctly\")\n",
    "    \n",
    "    # Test 2: Validation Data Pipeline\n",
    "    print(\"\\n2. Validation Data Pipeline Test:\")\n",
    "    \n",
    "    # Create validation dataset (no shuffling)\n",
    "    val_dataset = SimpleDataset(size=50, num_features=8, num_classes=5)\n",
    "    val_loader = DataLoader(val_dataset, batch_size=10, shuffle=False)\n",
    "    \n",
    "    # Simulate validation\n",
    "    val_samples = 0\n",
    "    val_batches = 0\n",
    "    \n",
    "    for batch_data, batch_labels in val_loader:\n",
    "        val_batches += 1\n",
    "        val_samples += batch_data.shape[0]\n",
    "        \n",
    "        # Verify consistent batch processing\n",
    "        assert batch_data.shape[1] == 8, \"Validation features should match training\"\n",
    "        assert len(batch_labels.shape) == 1, \"Validation labels should be 1D\"\n",
    "        \n",
    "    assert val_samples == 50, f\"Should process 50 validation samples, got {val_samples}\"\n",
    "    assert val_batches == 5, f\"Should have 5 validation batches, got {val_batches}\"\n",
    "    print(\"✅ Validation pipeline works correctly\")\n",
    "    \n",
    "    # Test 3: Different Dataset Configurations\n",
    "    print(\"\\n3. Dataset Configuration Test:\")\n",
    "    \n",
    "    # Test different configurations\n",
    "    configs = [\n",
    "        (200, 4, 3),   # Medium dataset\n",
    "        (50, 12, 10),  # High-dimensional features\n",
    "        (1000, 2, 2),  # Large dataset, simple features\n",
    "    ]\n",
    "    \n",
    "    for size, features, classes in configs:\n",
    "        dataset = SimpleDataset(size=size, num_features=features, num_classes=classes)\n",
    "        loader = DataLoader(dataset, batch_size=32, shuffle=True)\n",
    "        \n",
    "        # Test one batch\n",
    "        batch_data, batch_labels = next(iter(loader))\n",
    "        \n",
    "        assert batch_data.shape[1] == features, f\"Features mismatch for config {configs}\"\n",
    "        assert len(dataset) == size, f\"Size mismatch for config {configs}\"\n",
    "        assert dataset.get_num_classes() == classes, f\"Classes mismatch for config {configs}\"\n",
    "    \n",
    "    print(\"✅ Different dataset configurations work correctly\")\n",
    "    \n",
    "    # Test 4: Memory Efficiency Simulation\n",
    "    print(\"\\n4. Memory Efficiency Test:\")\n",
    "    \n",
    "    # Create larger dataset to test memory efficiency\n",
    "    large_dataset = SimpleDataset(size=500, num_features=20, num_classes=10)\n",
    "    large_loader = DataLoader(large_dataset, batch_size=50, shuffle=True)\n",
    "    \n",
    "    # Process all batches to ensure memory efficiency\n",
    "    processed_samples = 0\n",
    "    max_batch_size = 0\n",
    "    \n",
    "    for batch_data, batch_labels in large_loader:\n",
    "        processed_samples += batch_data.shape[0]\n",
    "        max_batch_size = max(max_batch_size, batch_data.shape[0])\n",
    "        \n",
    "        # Verify memory usage stays reasonable\n",
    "        assert batch_data.shape[0] <= 50, f\"Batch size should not exceed 50, got {batch_data.shape[0]}\"\n",
    "    \n",
    "    assert processed_samples == 500, f\"Should process all 500 samples, got {processed_samples}\"\n",
    "    print(\"✅ Memory efficiency works correctly\")\n",
    "    \n",
    "    # Test 5: Multi-Epoch Training Simulation\n",
    "    print(\"\\n5. Multi-Epoch Training Test:\")\n",
    "    \n",
    "    # Simulate multiple epochs\n",
    "    dataset = SimpleDataset(size=60, num_features=6, num_classes=3)\n",
    "    loader = DataLoader(dataset, batch_size=20, shuffle=True)\n",
    "    \n",
    "    for epoch in range(3):\n",
    "        epoch_samples = 0\n",
    "        for batch_data, batch_labels in loader:\n",
    "            epoch_samples += batch_data.shape[0]\n",
    "            \n",
    "            # Verify shapes remain consistent across epochs\n",
    "            assert batch_data.shape[1] == 6, f\"Features should be 6 in epoch {epoch}\"\n",
    "            assert len(batch_labels.shape) == 1, f\"Labels should be 1D in epoch {epoch}\"\n",
    "        \n",
    "        assert epoch_samples == 60, f\"Should process 60 samples in epoch {epoch}, got {epoch_samples}\"\n",
    "    \n",
    "    print(\"✅ Multi-epoch training works correctly\")\n",
    "    \n",
    "    print(\"\\n🎉 Comprehensive test passed! Your data pipeline works correctly for:\")\n",
    "    print(\"  • Large-scale dataset handling\")\n",
    "    print(\"  • Batch processing with multiple workers\")\n",
    "    print(\"  • Shuffling and sampling strategies\")\n",
    "    print(\"  • Memory-efficient data loading\")\n",
    "    print(\"  • Complete training pipeline integration\")\n",
    "    print(\"📈 Progress: Production-ready data pipeline ✓\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"❌ Comprehensive test failed: {e}\")\n",
    "    raise\n",
    "\n",
    "print(\"📈 Final Progress: Complete data pipeline ready for production ML!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d090c1",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Unit Test: Dataset Interface Implementation\n",
    "\n",
    "This test validates the abstract Dataset interface, ensuring proper inheritance, method implementation, and interface compliance for creating custom datasets in the TinyTorch data loading pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62c32031",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def test_unit_dataset_interface():\n",
    "    \"\"\"Unit test for the Dataset abstract interface implementation.\"\"\"\n",
    "    print(\"🔬 Unit Test: Dataset Interface...\")\n",
    "    \n",
    "    # Test TestDataset implementation\n",
    "    dataset = TestDataset(size=5)\n",
    "    \n",
    "    # Test basic interface\n",
    "    assert len(dataset) == 5, \"Dataset should have correct length\"\n",
    "    \n",
    "    # Test data access\n",
    "    sample, label = dataset[0]\n",
    "    assert isinstance(sample, Tensor), \"Sample should be Tensor\"\n",
    "    assert isinstance(label, Tensor), \"Label should be Tensor\"\n",
    "    \n",
    "    print(\"✅ Dataset interface works correctly\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbbce516",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Unit Test: DataLoader Implementation\n",
    "\n",
    "This test validates the DataLoader class functionality, ensuring proper batch creation, iteration capability, and integration with datasets for efficient data loading in machine learning training pipelines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0025080",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def test_unit_dataloader():\n",
    "    \"\"\"Unit test for the DataLoader implementation.\"\"\"\n",
    "    print(\"🔬 Unit Test: DataLoader...\")\n",
    "    \n",
    "    # Test DataLoader with TestDataset\n",
    "    dataset = TestDataset(size=10)\n",
    "    loader = DataLoader(dataset, batch_size=3, shuffle=False)\n",
    "    \n",
    "    # Test iteration\n",
    "    batches = list(loader)\n",
    "    assert len(batches) >= 3, \"Should have at least 3 batches\"\n",
    "    \n",
    "    # Test batch shapes\n",
    "    batch_data, batch_labels = batches[0]\n",
    "    assert batch_data.shape[0] <= 3, \"Batch size should be <= 3\"\n",
    "    assert batch_labels.shape[0] <= 3, \"Batch labels should match data\"\n",
    "    \n",
    "    print(\"✅ DataLoader works correctly\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfc685e4",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Unit Test: Simple Dataset Implementation\n",
    "\n",
    "This test validates the SimpleDataset class, ensuring it can handle real-world data scenarios including proper data storage, indexing, and compatibility with the DataLoader for practical machine learning workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cc885b1",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def test_unit_simple_dataset():\n",
    "    \"\"\"Unit test for the SimpleDataset implementation.\"\"\"\n",
    "    print(\"🔬 Unit Test: SimpleDataset...\")\n",
    "    \n",
    "    # Test SimpleDataset\n",
    "    dataset = SimpleDataset(size=100, num_features=4, num_classes=3)\n",
    "    \n",
    "    # Test properties\n",
    "    assert len(dataset) == 100, \"Dataset should have correct size\"\n",
    "    assert dataset.get_num_classes() == 3, \"Should have correct number of classes\"\n",
    "    \n",
    "    # Test data access\n",
    "    sample, label = dataset[0]\n",
    "    assert sample.shape == (4,), \"Sample should have correct features\"\n",
    "    assert 0 <= label.data < 3, \"Label should be valid class\"\n",
    "    \n",
    "    print(\"✅ SimpleDataset works correctly\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bd59540",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🧪 Unit Test: Complete Data Pipeline Integration\n",
    "\n",
    "This comprehensive test validates the entire data pipeline from dataset creation through DataLoader batching, ensuring all components work together seamlessly for end-to-end machine learning data processing workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c63e6cd",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def test_unit_dataloader_pipeline():\n",
    "    \"\"\"Comprehensive unit test for the complete data pipeline.\"\"\"\n",
    "    print(\"🔬 Comprehensive Test: Data Pipeline...\")\n",
    "    \n",
    "    # Test complete pipeline\n",
    "    dataset = SimpleDataset(size=50, num_features=10, num_classes=5)\n",
    "    loader = DataLoader(dataset, batch_size=8, shuffle=True)\n",
    "    \n",
    "    total_samples = 0\n",
    "    for batch_data, batch_labels in loader:\n",
    "        assert isinstance(batch_data, Tensor), \"Batch data should be Tensor\"\n",
    "        assert isinstance(batch_labels, Tensor), \"Batch labels should be Tensor\"\n",
    "        assert batch_data.shape[1] == 10, \"Features should be correct\"\n",
    "        total_samples += batch_data.shape[0]\n",
    "    \n",
    "    assert total_samples == 50, \"Should process all samples\"\n",
    "    \n",
    "    print(\"✅ Data pipeline integration works correctly\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63acc83f",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "307992df",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🧪 Module Testing\n",
    "\n",
    "Time to test your implementation! This section uses TinyTorch's standardized testing framework to ensure your implementation works correctly.\n",
    "\n",
    "**This testing section is locked** - it provides consistent feedback across all modules and cannot be modified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd73bc81",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "standardized-testing",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# =============================================================================\n",
    "# STANDARDIZED MODULE TESTING - DO NOT MODIFY\n",
    "# This cell is locked to ensure consistent testing across all TinyTorch modules\n",
    "# ============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3171e7ee",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🔬 Integration Test: DataLoader with Tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924540fd",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def test_module_dataloader_tensor_yield():\n",
    "    \"\"\"\n",
    "    Integration test for the DataLoader and Tensor classes.\n",
    "    \n",
    "    Tests that the DataLoader correctly yields batches of Tensors.\n",
    "    \"\"\"\n",
    "    print(\"🔬 Running Integration Test: DataLoader with Tensors...\")\n",
    "\n",
    "    # 1. Create a simple dataset\n",
    "    dataset = SimpleDataset(size=50, num_features=8, num_classes=4)\n",
    "\n",
    "    # 2. Create a DataLoader\n",
    "    dataloader = DataLoader(dataset, batch_size=10, shuffle=False)\n",
    "\n",
    "    # 3. Get one batch from the dataloader\n",
    "    data_batch, labels_batch = next(iter(dataloader))\n",
    "\n",
    "    # 4. Assert the batch contents are correct\n",
    "    assert isinstance(data_batch, Tensor), \"Data batch should be a Tensor\"\n",
    "    assert data_batch.shape == (10, 8), f\"Expected data shape (10, 8), but got {data_batch.shape}\"\n",
    "    \n",
    "    assert isinstance(labels_batch, Tensor), \"Labels batch should be a Tensor\"\n",
    "    assert labels_batch.shape == (10,), f\"Expected labels shape (10,), but got {labels_batch.shape}\"\n",
    "\n",
    "    print(\"✅ Integration Test Passed: DataLoader correctly yields batches of Tensors.\")\n",
    "\n",
    "# Test function defined (called in main block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8b23ef0",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 📊 ML Systems: I/O Pipeline Optimization & Bottleneck Analysis\n",
    "\n",
    "Now that you have data loading systems, let's develop **I/O optimization skills**. This section teaches you to identify and fix data loading bottlenecks that can dramatically slow down training in production systems.\n",
    "\n",
    "### **Learning Outcome**: *\"I can identify and fix I/O bottlenecks that limit training speed\"*\n",
    "\n",
    "---\n",
    "\n",
    "## Data Pipeline Profiler (Medium Guided Implementation)\n",
    "\n",
    "As an ML systems engineer, you need to ensure data loading doesn't become the bottleneck. Training GPUs can process data much faster than traditional storage can provide it. Let's build tools to measure and optimize data pipeline performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac8f7b9",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import os\n",
    "import threading\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "\n",
    "class DataPipelineProfiler:\n",
    "    \"\"\"\n",
    "    I/O pipeline profiling toolkit for data loading systems.\n",
    "    \n",
    "    Helps ML engineers identify bottlenecks in data loading pipelines\n",
    "    and optimize throughput for high-performance training systems.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self):\n",
    "        self.profiling_history = []\n",
    "        self.bottleneck_threshold = 0.1  # seconds per batch\n",
    "        \n",
    "    def time_dataloader_iteration(self, dataloader, num_batches=10):\n",
    "        \"\"\"\n",
    "        Time how long it takes to iterate through DataLoader batches.\n",
    "        \n",
    "        TODO: Implement DataLoader timing analysis.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Record start time\n",
    "        2. Iterate through specified number of batches\n",
    "        3. Time each batch loading\n",
    "        4. Calculate statistics (total, average, min, max times)\n",
    "        5. Identify if data loading is a bottleneck\n",
    "        6. Return comprehensive timing analysis\n",
    "        \n",
    "        EXAMPLE:\n",
    "        profiler = DataPipelineProfiler()\n",
    "        timing = profiler.time_dataloader_iteration(my_dataloader, 20)\n",
    "        print(f\"Avg batch time: {timing['avg_batch_time']:.3f}s\")\n",
    "        print(f\"Bottleneck: {timing['is_bottleneck']}\")\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - **Production Optimization**: Fast GPUs often wait for slow data loading\n",
    "        - **System Bottlenecks**: Data loading can limit training speed more than model complexity\n",
    "        - **Resource Planning**: Understanding I/O vs compute trade-offs for hardware selection\n",
    "        - **Pipeline Tuning**: Multi-worker data loading and prefetching strategies\n",
    "        \n",
    "        HINTS:\n",
    "        - Use enumerate(dataloader) to get batches\n",
    "        - Time each batch: start = time.time(), batch = next(iter), end = time.time()\n",
    "        - Break after num_batches to avoid processing entire dataset\n",
    "        - Calculate: total_time, avg_time, min_time, max_time\n",
    "        - Bottleneck if avg_time > self.bottleneck_threshold\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        batch_times = []\n",
    "        total_start = time.time()\n",
    "        \n",
    "        try:\n",
    "            dataloader_iter = iter(dataloader)\n",
    "            for i in range(num_batches):\n",
    "                batch_start = time.time()\n",
    "                try:\n",
    "                    batch = next(dataloader_iter)\n",
    "                    batch_end = time.time()\n",
    "                    batch_time = batch_end - batch_start\n",
    "                    batch_times.append(batch_time)\n",
    "                except StopIteration:\n",
    "                    print(f\"   DataLoader exhausted after {i} batches\")\n",
    "                    break\n",
    "        except Exception as e:\n",
    "            print(f\"   Error during iteration: {e}\")\n",
    "            return {'error': str(e)}\n",
    "        \n",
    "        total_end = time.time()\n",
    "        total_time = total_end - total_start\n",
    "        \n",
    "        if batch_times:\n",
    "            avg_batch_time = sum(batch_times) / len(batch_times)\n",
    "            min_batch_time = min(batch_times)\n",
    "            max_batch_time = max(batch_times)\n",
    "            \n",
    "            # Check if data loading is a bottleneck\n",
    "            is_bottleneck = avg_batch_time > self.bottleneck_threshold\n",
    "            \n",
    "            # Calculate throughput\n",
    "            batches_per_second = len(batch_times) / total_time if total_time > 0 else 0\n",
    "            \n",
    "            return {\n",
    "                'total_time': total_time,\n",
    "                'num_batches': len(batch_times),\n",
    "                'avg_batch_time': avg_batch_time,\n",
    "                'min_batch_time': min_batch_time,\n",
    "                'max_batch_time': max_batch_time,\n",
    "                'batches_per_second': batches_per_second,\n",
    "                'is_bottleneck': is_bottleneck,\n",
    "                'bottleneck_threshold': self.bottleneck_threshold\n",
    "            }\n",
    "        else:\n",
    "            return {'error': 'No batches processed'}\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def analyze_batch_size_scaling(self, dataset, batch_sizes=[16, 32, 64, 128]):\n",
    "        \"\"\"\n",
    "        Analyze how batch size affects data loading performance.\n",
    "        \n",
    "        TODO: Implement batch size scaling analysis.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. For each batch size, create a DataLoader\n",
    "        2. Time the data loading for each configuration\n",
    "        3. Calculate throughput (samples/second) for each\n",
    "        4. Identify optimal batch size for I/O performance\n",
    "        5. Return scaling analysis with recommendations\n",
    "        \n",
    "        EXAMPLE:\n",
    "        profiler = DataPipelineProfiler()\n",
    "        analysis = profiler.analyze_batch_size_scaling(my_dataset, [16, 32, 64])\n",
    "        print(f\"Optimal batch size: {analysis['optimal_batch_size']}\")\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - **Memory vs Throughput**: Larger batches improve throughput but consume more memory\n",
    "        - **Hardware Optimization**: Optimal batch size depends on GPU memory and compute units\n",
    "        - **Training Dynamics**: Batch size affects gradient noise and convergence behavior\n",
    "        - **Production Scaling**: Understanding batch size impact on serving latency and cost\n",
    "        \n",
    "        HINTS:\n",
    "        - Create DataLoader: DataLoader(dataset, batch_size=bs, shuffle=False)\n",
    "        - Time with self.time_dataloader_iteration()\n",
    "        - Calculate: samples_per_second = batch_size * batches_per_second\n",
    "        - Find batch size with highest samples/second\n",
    "        - Consider memory constraints vs throughput\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        scaling_results = []\n",
    "        \n",
    "        for batch_size in batch_sizes:\n",
    "            print(f\"   Testing batch size {batch_size}...\")\n",
    "            \n",
    "            # Create DataLoader with current batch size\n",
    "            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "            \n",
    "            # Time the data loading\n",
    "            timing_result = self.time_dataloader_iteration(dataloader, num_batches=min(10, len(dataset)//batch_size))\n",
    "            \n",
    "            if 'error' not in timing_result:\n",
    "                # Calculate throughput metrics\n",
    "                samples_per_second = batch_size * timing_result['batches_per_second']\n",
    "                \n",
    "                result = {\n",
    "                    'batch_size': batch_size,\n",
    "                    'avg_batch_time': timing_result['avg_batch_time'],\n",
    "                    'batches_per_second': timing_result['batches_per_second'],\n",
    "                    'samples_per_second': samples_per_second,\n",
    "                    'is_bottleneck': timing_result['is_bottleneck']\n",
    "                }\n",
    "                scaling_results.append(result)\n",
    "        \n",
    "        # Find optimal batch size (highest throughput)\n",
    "        if scaling_results:\n",
    "            optimal = max(scaling_results, key=lambda x: x['samples_per_second'])\n",
    "            optimal_batch_size = optimal['batch_size']\n",
    "            \n",
    "            return {\n",
    "                'scaling_results': scaling_results,\n",
    "                'optimal_batch_size': optimal_batch_size,\n",
    "                'max_throughput': optimal['samples_per_second']\n",
    "            }\n",
    "        else:\n",
    "            return {'error': 'No valid results obtained'}\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def compare_io_strategies(self, dataset, strategies=['sequential', 'shuffled']):\n",
    "        \"\"\"\n",
    "        Compare different I/O strategies for data loading performance.\n",
    "        \n",
    "        This function is PROVIDED to demonstrate I/O optimization analysis.\n",
    "        Students use it to understand different data loading patterns.\n",
    "        \"\"\"\n",
    "        print(\"📊 I/O STRATEGY COMPARISON\")\n",
    "        print(\"=\" * 40)\n",
    "        \n",
    "        results = {}\n",
    "        batch_size = 32  # Standard batch size for comparison\n",
    "        \n",
    "        for strategy in strategies:\n",
    "            print(f\"\\n🔍 Testing {strategy.upper()} strategy...\")\n",
    "            \n",
    "            if strategy == 'sequential':\n",
    "                dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "            elif strategy == 'shuffled':\n",
    "                dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "            else:\n",
    "                print(f\"   Unknown strategy: {strategy}\")\n",
    "                continue\n",
    "            \n",
    "            # Time the strategy\n",
    "            timing_result = self.time_dataloader_iteration(dataloader, num_batches=20)\n",
    "            \n",
    "            if 'error' not in timing_result:\n",
    "                results[strategy] = timing_result\n",
    "                print(f\"   Avg batch time: {timing_result['avg_batch_time']:.3f}s\")\n",
    "                print(f\"   Throughput: {timing_result['batches_per_second']:.1f} batches/sec\")\n",
    "                print(f\"   Bottleneck: {'Yes' if timing_result['is_bottleneck'] else 'No'}\")\n",
    "        \n",
    "        # Compare strategies\n",
    "        if len(results) >= 2:\n",
    "            fastest = min(results.items(), key=lambda x: x[1]['avg_batch_time'])\n",
    "            slowest = max(results.items(), key=lambda x: x[1]['avg_batch_time'])\n",
    "            \n",
    "            speedup = slowest[1]['avg_batch_time'] / fastest[1]['avg_batch_time']\n",
    "            \n",
    "            print(f\"\\n🎯 STRATEGY ANALYSIS:\")\n",
    "            print(f\"   Fastest: {fastest[0]} ({fastest[1]['avg_batch_time']:.3f}s)\")\n",
    "            print(f\"   Slowest: {slowest[0]} ({slowest[1]['avg_batch_time']:.3f}s)\")\n",
    "            print(f\"   Speedup: {speedup:.1f}x\")\n",
    "        \n",
    "        return results\n",
    "    \n",
    "    def simulate_compute_vs_io_balance(self, dataloader, simulated_compute_time=0.05):\n",
    "        \"\"\"\n",
    "        Simulate the balance between data loading and compute time.\n",
    "        \n",
    "        This function is PROVIDED to show I/O vs compute analysis.\n",
    "        Students use it to understand when I/O becomes a bottleneck.\n",
    "        \"\"\"\n",
    "        print(\"⚖️  COMPUTE vs I/O BALANCE ANALYSIS\")\n",
    "        print(\"=\" * 45)\n",
    "        \n",
    "        print(f\"Simulated compute time per batch: {simulated_compute_time:.3f}s\")\n",
    "        print(f\"(This represents GPU processing time)\")\n",
    "        \n",
    "        # Time data loading\n",
    "        io_timing = self.time_dataloader_iteration(dataloader, num_batches=15)\n",
    "        \n",
    "        if 'error' in io_timing:\n",
    "            print(f\"Error in timing: {io_timing['error']}\")\n",
    "            return\n",
    "        \n",
    "        avg_io_time = io_timing['avg_batch_time']\n",
    "        \n",
    "        print(f\"\\n📊 TIMING ANALYSIS:\")\n",
    "        print(f\"   Data loading time: {avg_io_time:.3f}s per batch\")\n",
    "        print(f\"   Simulated compute: {simulated_compute_time:.3f}s per batch\")\n",
    "        \n",
    "        # Determine bottleneck\n",
    "        if avg_io_time > simulated_compute_time:\n",
    "            bottleneck = \"I/O\"\n",
    "            utilization = simulated_compute_time / avg_io_time * 100\n",
    "            print(f\"\\n🚨 BOTTLENECK: {bottleneck}\")\n",
    "            print(f\"   GPU utilization: {utilization:.1f}%\")\n",
    "            print(f\"   GPU waiting for data: {avg_io_time - simulated_compute_time:.3f}s per batch\")\n",
    "        else:\n",
    "            bottleneck = \"Compute\"\n",
    "            utilization = avg_io_time / simulated_compute_time * 100\n",
    "            print(f\"\\n✅ BOTTLENECK: {bottleneck}\")\n",
    "            print(f\"   I/O utilization: {utilization:.1f}%\")\n",
    "            print(f\"   I/O waiting for GPU: {simulated_compute_time - avg_io_time:.3f}s per batch\")\n",
    "        \n",
    "        # Calculate training impact\n",
    "        total_cycle_time = max(avg_io_time, simulated_compute_time)\n",
    "        efficiency = min(avg_io_time, simulated_compute_time) / total_cycle_time * 100\n",
    "        \n",
    "        print(f\"\\n🎯 TRAINING IMPACT:\")\n",
    "        print(f\"   Pipeline efficiency: {efficiency:.1f}%\")\n",
    "        print(f\"   Total cycle time: {total_cycle_time:.3f}s\")\n",
    "        \n",
    "        if bottleneck == \"I/O\":\n",
    "            print(f\"   💡 Recommendation: Optimize data loading\")\n",
    "            print(f\"      - Increase batch size\")\n",
    "            print(f\"      - Use data prefetching\")\n",
    "            print(f\"      - Faster storage (SSD vs HDD)\")\n",
    "        else:\n",
    "            print(f\"   💡 Recommendation: I/O is well optimized\")\n",
    "            print(f\"      - Consider larger models or batch sizes\")\n",
    "            print(f\"      - Focus on compute optimization\")\n",
    "        \n",
    "        return {\n",
    "            'io_time': avg_io_time,\n",
    "            'compute_time': simulated_compute_time,\n",
    "            'bottleneck': bottleneck,\n",
    "            'efficiency': efficiency,\n",
    "            'total_cycle_time': total_cycle_time\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad2c8bd8",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### 🎯 Learning Activity 1: DataLoader Performance Profiling (Medium Guided Implementation)\n",
    "\n",
    "**Goal**: Learn to measure data loading performance and identify I/O bottlenecks that can slow down training.\n",
    "\n",
    "Complete the missing implementations in the `DataPipelineProfiler` class above, then use your profiler to analyze data loading performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b50e007",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the data pipeline profiler\n",
    "profiler = DataPipelineProfiler()\n",
    "\n",
    "# Only run tests when module is executed directly\n",
    "if __name__ == '__main__':\n",
    "    print(\"📊 DATA PIPELINE PERFORMANCE ANALYSIS\")\n",
    "    print(\"=\" * 50)\n",
    "\n",
    "    # Create test dataset and dataloader\n",
    "    test_dataset = TensorDataset([\n",
    "        Tensor(np.random.randn(100)) for _ in range(1000)  # 1000 samples\n",
    "    ], [\n",
    "        Tensor([i % 10]) for i in range(1000)  # Labels\n",
    "    ])\n",
    "\n",
    "    # Test 1: Basic DataLoader timing\n",
    "    print(\"⏱️  Basic DataLoader Timing:\")\n",
    "    basic_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=False)\n",
    "\n",
    "# Students use their implemented timing function\n",
    "timing_result = profiler.time_dataloader_iteration(basic_dataloader, num_batches=25)\n",
    "\n",
    "if 'error' not in timing_result:\n",
    "    print(f\"   Average batch time: {timing_result['avg_batch_time']:.3f}s\")\n",
    "    print(f\"   Throughput: {timing_result['batches_per_second']:.1f} batches/sec\")\n",
    "    print(f\"   Bottleneck detected: {'Yes' if timing_result['is_bottleneck'] else 'No'}\")\n",
    "    \n",
    "    # Calculate samples per second\n",
    "    samples_per_sec = 32 * timing_result['batches_per_second']\n",
    "    print(f\"   Samples/second: {samples_per_sec:.1f}\")\n",
    "else:\n",
    "    print(f\"   Error: {timing_result['error']}\")\n",
    "\n",
    "# Test 2: Batch size scaling analysis\n",
    "print(f\"\\n📈 Batch Size Scaling Analysis:\")\n",
    "\n",
    "# Students use their implemented scaling analysis\n",
    "scaling_analysis = profiler.analyze_batch_size_scaling(test_dataset, [16, 32, 64, 128])\n",
    "\n",
    "if 'error' not in scaling_analysis:\n",
    "    print(f\"   Optimal batch size: {scaling_analysis['optimal_batch_size']}\")\n",
    "    print(f\"   Max throughput: {scaling_analysis['max_throughput']:.1f} samples/sec\")\n",
    "    \n",
    "    print(f\"\\n   📊 Detailed Results:\")\n",
    "    for result in scaling_analysis['scaling_results']:\n",
    "        print(f\"      Batch {result['batch_size']:3d}: {result['samples_per_second']:6.1f} samples/sec\")\n",
    "else:\n",
    "    print(f\"   Error: {scaling_analysis['error']}\")\n",
    "\n",
    "print(f\"\\n💡 I/O PERFORMANCE INSIGHTS:\")\n",
    "print(f\"   - Larger batches often improve throughput (better amortization)\")\n",
    "print(f\"   - But memory constraints limit maximum batch size\")\n",
    "print(f\"   - Sweet spot balances throughput vs memory usage\")\n",
    "print(f\"   - Real systems: GPU memory determines practical limits\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92ef4498",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### 🎯 Learning Activity 2: Production I/O Optimization Analysis (Review & Understand)\n",
    "\n",
    "**Goal**: Understand how I/O performance affects real training systems and learn optimization strategies used in production."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74695654",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare different I/O strategies\n",
    "io_comparison = profiler.compare_io_strategies(test_dataset, ['sequential', 'shuffled'])\n",
    "\n",
    "# Simulate compute vs I/O balance with different scenarios\n",
    "print(f\"\\n⚖️  COMPUTE vs I/O SCENARIOS:\")\n",
    "print(f\"=\" * 40)\n",
    "\n",
    "# Test different compute scenarios\n",
    "compute_scenarios = [\n",
    "    (0.01, \"Fast GPU (V100/A100)\"),\n",
    "    (0.05, \"Medium GPU (RTX 3080)\"),\n",
    "    (0.1, \"CPU-only training\"),\n",
    "    (0.2, \"Complex model/large batch\")\n",
    "]\n",
    "\n",
    "sample_dataloader = DataLoader(test_dataset, batch_size=64, shuffle=False)\n",
    "\n",
    "for compute_time, scenario_name in compute_scenarios:\n",
    "    print(f\"\\n🖥️  {scenario_name}:\")\n",
    "    balance_analysis = profiler.simulate_compute_vs_io_balance(sample_dataloader, compute_time)\n",
    "\n",
    "print(f\"\\n🎯 PRODUCTION I/O OPTIMIZATION LESSONS:\")\n",
    "print(f\"=\" * 50)\n",
    "\n",
    "print(f\"\\n1. 📊 I/O BOTTLENECK IDENTIFICATION:\")\n",
    "print(f\"   - Fast GPUs often bottlenecked by data loading\")\n",
    "print(f\"   - CPU training rarely I/O bottlenecked\")\n",
    "print(f\"   - Modern GPUs process data faster than storage provides it\")\n",
    "\n",
    "print(f\"\\n2. 🚀 OPTIMIZATION STRATEGIES:\")\n",
    "print(f\"   - Data prefetching: Load next batch while GPU computes\")\n",
    "print(f\"   - Parallel workers: Multiple threads/processes for loading\")\n",
    "print(f\"   - Faster storage: NVMe SSD vs SATA vs network storage\")\n",
    "print(f\"   - Data caching: Keep frequently used data in memory\")\n",
    "\n",
    "print(f\"\\n3. 🏗️ ARCHITECTURE DECISIONS:\")\n",
    "print(f\"   - Batch size: Larger batches amortize I/O overhead\")\n",
    "print(f\"   - Data format: Preprocessed vs on-the-fly transformation\")\n",
    "print(f\"   - Storage location: Local vs network vs cloud storage\")\n",
    "\n",
    "print(f\"\\n4. 💰 COST IMPLICATIONS:\")\n",
    "print(f\"   - I/O bottlenecks waste expensive GPU time\")\n",
    "print(f\"   - GPU utilization directly affects training costs\")\n",
    "print(f\"   - Faster storage investment pays off in GPU efficiency\")\n",
    "\n",
    "print(f\"\\n💡 SYSTEMS ENGINEERING INSIGHT:\")\n",
    "print(f\"I/O optimization is often the highest-impact performance improvement:\")\n",
    "print(f\"- GPUs are expensive → maximize their utilization\")\n",
    "print(f\"- Data loading is often the limiting factor\")\n",
    "print(f\"- 10% I/O improvement = 10% faster training = 10% cost reduction\")\n",
    "print(f\"- Modern ML systems spend significant effort on data pipeline optimization\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # Test the dataset interface demonstration\n",
    "    try:\n",
    "        test_dataset = TestDataset(size=5)\n",
    "        print(f\"Dataset created with size: {len(test_dataset)}\")\n",
    "        \n",
    "        # Test __getitem__\n",
    "        data, label = test_dataset[0]\n",
    "        print(f\"Sample 0: data={data}, label={label}\")\n",
    "        assert isinstance(data, Tensor), \"Data should be a Tensor\"\n",
    "        assert isinstance(label, Tensor), \"Label should be a Tensor\"\n",
    "        print(\"✅ Dataset __getitem__ works correctly\")\n",
    "        \n",
    "        # Test __len__\n",
    "        assert len(test_dataset) == 5, f\"Dataset length should be 5, got {len(test_dataset)}\"\n",
    "        print(\"✅ Dataset __len__ works correctly\")\n",
    "        \n",
    "        # Test get_num_classes\n",
    "        num_classes = test_dataset.get_num_classes()\n",
    "        assert num_classes == 2, f\"Number of classes should be 2, got {num_classes}\"\n",
    "        print(\"✅ Dataset get_num_classes works correctly\")\n",
    "        \n",
    "        # Test get_sample_shape\n",
    "        sample_shape = test_dataset.get_sample_shape()\n",
    "        assert sample_shape == (3,), f\"Sample shape should be (3,), got {sample_shape}\"\n",
    "        print(\"✅ Dataset get_sample_shape works correctly\")\n",
    "        \n",
    "        print(\"🎯 Dataset interface pattern:\")\n",
    "        print(\"   __getitem__: Returns (data, label) tuple\")\n",
    "        print(\"   __len__: Returns dataset size\")\n",
    "        print(\"   get_num_classes: Returns number of classes\")\n",
    "        print(\"   get_sample_shape: Returns shape of data samples\")\n",
    "        print(\"📈 Progress: Dataset interface ✓\")\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"❌ Dataset interface test failed: {e}\")\n",
    "        raise\n",
    "    \n",
    "    # Run all tests\n",
    "    test_unit_dataset_interface()\n",
    "    test_unit_dataloader()\n",
    "    test_unit_simple_dataset()\n",
    "    test_unit_dataloader_pipeline()\n",
    "    test_module_dataloader_tensor_yield()\n",
    "    \n",
    "    print(\"All tests passed!\")\n",
    "    print(\"dataloader_dev module complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27bce6e8",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🤔 ML Systems Thinking Questions\n",
    "\n",
    "### System Design\n",
    "1. How does TinyTorch's DataLoader design compare to PyTorch's DataLoader and TensorFlow's tf.data API in terms of flexibility and performance?\n",
    "2. What are the trade-offs between memory-mapped files, streaming data loading, and in-memory caching for large-scale ML datasets?\n",
    "3. How would you design a data loading system that efficiently handles both structured (tabular) and unstructured (images, text) data?\n",
    "\n",
    "### Production ML\n",
    "1. How would you implement fault-tolerant data loading that can handle network failures and corrupted files in production environments?\n",
    "2. What strategies would you use to ensure data consistency and prevent data leakage when loading from constantly updating production databases?\n",
    "3. How would you design a data pipeline that supports both batch inference and real-time prediction serving?\n",
    "\n",
    "### Framework Design\n",
    "1. What design patterns enable efficient data preprocessing that can be distributed across multiple worker processes without blocking training?\n",
    "2. How would you implement dynamic batching that adapts batch sizes based on available memory and model complexity?\n",
    "3. What abstractions would you create to support different data formats (images, audio, text) while maintaining a unified loading interface?\n",
    "\n",
    "### Performance & Scale\n",
    "1. How do different data loading strategies (synchronous vs asynchronous, single vs multi-threaded) impact training throughput on different hardware?\n",
    "2. What are the bottlenecks when loading data for distributed training across multiple machines, and how would you optimize data transfer?\n",
    "3. How would you implement data loading that scales efficiently from small datasets (MB) to massive datasets (TB) without code changes?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0abe9e82",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 MODULE SUMMARY: Data Loading and Processing\n",
    "\n",
    "Congratulations! You've successfully implemented professional data loading systems:\n",
    "\n",
    "### What You've Accomplished\n",
    "✅ **DataLoader Class**: Efficient batch processing with memory management\n",
    "✅ **Dataset Integration**: Seamless compatibility with Tensor operations\n",
    "✅ **Batch Processing**: Optimized data loading for training\n",
    "✅ **Memory Management**: Efficient handling of large datasets\n",
    "✅ **Real Applications**: Image classification, regression, and more\n",
    "\n",
    "### Key Concepts You've Learned\n",
    "- **Batch processing**: How to efficiently process data in chunks\n",
    "- **Memory management**: Handling large datasets without memory overflow\n",
    "- **Data iteration**: Creating efficient data loading pipelines\n",
    "- **Integration patterns**: How data loaders work with neural networks\n",
    "- **Performance optimization**: Balancing speed and memory usage\n",
    "\n",
    "### Professional Skills Developed\n",
    "- **Data engineering**: Building robust data processing pipelines\n",
    "- **Memory optimization**: Efficient handling of large datasets\n",
    "- **API design**: Clean interfaces for data loading operations\n",
    "- **Integration testing**: Ensuring data loaders work with neural networks\n",
    "\n",
    "### Ready for Advanced Applications\n",
    "Your data loading implementations now enable:\n",
    "- **Large-scale training**: Processing datasets too big for memory\n",
    "- **Real-time learning**: Streaming data for online learning\n",
    "- **Multi-modal data**: Handling images, text, and structured data\n",
    "- **Production systems**: Robust data pipelines for deployment\n",
    "\n",
    "### Connection to Real ML Systems\n",
    "Your implementations mirror production systems:\n",
    "- **PyTorch**: `torch.utils.data.DataLoader` provides identical functionality\n",
    "- **TensorFlow**: `tf.data.Dataset` implements similar concepts\n",
    "- **Industry Standard**: Every major ML framework uses these exact patterns\n",
    "\n",
    "### Next Steps\n",
    "1. **Export your code**: `tito export 09_dataloader`\n",
    "2. **Test your implementation**: `tito test 09_dataloader`\n",
    "3. **Build training pipelines**: Combine with neural networks for complete ML systems\n",
    "4. **Move to Module 9**: Add automatic differentiation for training!\n",
    "\n",
    "**Ready for autograd?** Your data loading systems are now ready for real training!"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
