TinyTorch/modules/source/04_layers/layers_dev.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c7474e2",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Layers - Neural Network Building Blocks and Composition Patterns\n",
    "\n",
    "Welcome to the Layers module! You'll build the fundamental components that stack together to form any neural network architecture, from simple perceptrons to transformers.\n",
    "\n",
    "## Learning Goals\n",
    "- Systems understanding: How layer composition creates complex function approximators and why stacking enables deep learning\n",
    "- Core implementation skill: Build matrix multiplication and Dense layers with proper parameter management\n",
    "- Pattern recognition: Understand how different layer types solve different computational problems\n",
    "- Framework connection: See how your layer implementations mirror PyTorch's nn.Module design patterns\n",
    "- Performance insight: Learn why layer computation order and memory layout determine training speed\n",
    "\n",
    "## Build → Use → Reflect\n",
    "1. **Build**: Matrix multiplication primitives and Dense layers with parameter initialization strategies\n",
    "2. **Use**: Compose layers into multi-layer networks and observe how data transforms through the stack\n",
    "3. **Reflect**: Why does layer depth enable more complex functions, and when does it hurt performance?\n",
    "\n",
    "## What You'll Achieve\n",
    "By the end of this module, you'll understand:\n",
    "- Deep technical understanding of how matrix operations enable neural networks to learn arbitrary functions\n",
    "- Practical capability to build and compose layers into complex architectures\n",
    "- Systems insight into why layer composition is the fundamental pattern for scalable ML systems\n",
    "- Performance consideration of how layer size and depth affect memory usage and computational cost\n",
    "- Connection to production ML systems and how frameworks optimize layer execution for different hardware\n",
    "\n",
    "## Systems Reality Check\n",
    "💡 **Production Context**: PyTorch's nn.Linear uses optimized BLAS operations and can automatically select GPU vs CPU execution based on data size\n",
    "⚡ **Performance Note**: Large matrix multiplications can be memory-bound rather than compute-bound - understanding this shapes how production systems optimize layer execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7346937b",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "layers-imports",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| default_exp core.layers\n",
    "\n",
    "#| export\n",
    "import numpy as np\n",
    "import sys\n",
    "import os\n",
    "from typing import Union, Tuple, Optional, Any\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__), '..', '02_tensor'))\n",
    "    from tensor_dev import Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa051980",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "layers-setup",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "print(\"🔥 TinyTorch Layers 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 neural network layers!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91373183",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## Where This Code Lives in the Final Package\n",
    "\n",
    "**Learning Side:** You work in modules/source/04_layers/layers_dev.py  \n",
    "**Building Side:** Code exports to tinytorch.core.layers\n",
    "\n",
    "```python\n",
    "# Final package structure:\n",
    "from tinytorch.core.layers import Dense, matmul  # All layer types together!\n",
    "from tinytorch.core.tensor import Tensor  # The foundation\n",
    "from tinytorch.core.activations import ReLU, Sigmoid  # Nonlinearity\n",
    "```\n",
    "\n",
    "**Why this matters:**\n",
    "- **Learning:** Focused modules for deep understanding\n",
    "- **Production:** Proper organization like PyTorch's torch.nn.Linear\n",
    "- **Consistency:** All layer types live together in core.layers\n",
    "- **Integration:** Works seamlessly with tensors and activations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e003e27",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "# Matrix Multiplication - The Heart of Neural Networks\n",
    "\n",
    "Every neural network operation ultimately reduces to matrix multiplication. Let's build the foundation that powers everything from simple perceptrons to transformers.\n",
    "\n",
    "## Why Matrix Multiplication Matters\n",
    "\n",
    "🧠 **Neural Network Core**: Every layer applies: output = input @ weights + bias  \n",
    "⚡ **Parallel Processing**: Matrix ops utilize vectorized CPU instructions and GPU parallelism  \n",
    "🏗️ **Scalable Architecture**: Stacking matrix operations creates arbitrarily complex function approximators  \n",
    "📈 **Performance Critical**: 90%+ of neural network compute time is spent in matrix multiplication  \n",
    "\n",
    "## Learning Objectives\n",
    "By implementing matrix multiplication, you'll understand:\n",
    "- How neural networks transform data through linear algebra\n",
    "- Why matrix operations are the building blocks of all modern ML frameworks\n",
    "- How proper implementation affects performance by orders of magnitude\n",
    "- The connection between mathematical operations and computational efficiency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73267b42",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "matmul-implementation",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "def matmul(a: Tensor, b: Tensor) -> Tensor:\n",
    "    \"\"\"\n",
    "    Matrix multiplication for tensors.\n",
    "    \n",
    "    Args:\n",
    "        a: Left tensor (shape: ..., m, k)\n",
    "        b: Right tensor (shape: ..., k, n)\n",
    "    \n",
    "    Returns:\n",
    "        Result tensor (shape: ..., m, n)\n",
    "    \n",
    "    TODO: Implement matrix multiplication using numpy's @ operator.\n",
    "    \n",
    "    STEP-BY-STEP IMPLEMENTATION:\n",
    "    1. Extract numpy arrays from both tensors using .data\n",
    "    2. Perform matrix multiplication: result_data = a_data @ b_data\n",
    "    3. Wrap result in a new Tensor and return\n",
    "    \n",
    "    LEARNING CONNECTIONS:\n",
    "    - This is the core operation in Dense layers: output = input @ weights\n",
    "    - PyTorch uses optimized BLAS libraries for this operation\n",
    "    - GPU implementations parallelize this across thousands of cores\n",
    "    - Understanding this operation is key to neural network performance\n",
    "    \n",
    "    EXAMPLE:\n",
    "    ```python\n",
    "    a = Tensor([[1, 2], [3, 4]])  # shape (2, 2)\n",
    "    b = Tensor([[5, 6], [7, 8]])  # shape (2, 2)\n",
    "    result = matmul(a, b)\n",
    "    # result.data = [[19, 22], [43, 50]]\n",
    "    ```\n",
    "    \n",
    "    IMPLEMENTATION HINTS:\n",
    "    - Use the @ operator for clean matrix multiplication\n",
    "    - Ensure you return a Tensor, not a numpy array\n",
    "    - The operation should work for any compatible matrix shapes\n",
    "    \"\"\"\n",
    "    ### BEGIN SOLUTION\n",
    "    # Extract numpy data from tensors\n",
    "    a_data = a.data\n",
    "    b_data = b.data\n",
    "    \n",
    "    # Perform matrix multiplication\n",
    "    result_data = a_data @ b_data\n",
    "    \n",
    "    # Wrap result in a Tensor\n",
    "    return Tensor(result_data)\n",
    "    ### END SOLUTION"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ee7a2cc",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Testing Matrix Multiplication\n",
    "\n",
    "Let's verify our matrix multiplication works correctly with some test cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aed9ac48",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-matmul",
     "locked": true,
     "points": 2,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def test_matmul():\n",
    "    \"\"\"Test matrix multiplication implementation.\"\"\"\n",
    "    print(\"🧪 Testing Matrix Multiplication...\")\n",
    "    \n",
    "    # Test case 1: Simple 2x2 matrices\n",
    "    a = Tensor([[1, 2], [3, 4]])\n",
    "    b = Tensor([[5, 6], [7, 8]])\n",
    "    result = matmul(a, b)\n",
    "    expected = np.array([[19, 22], [43, 50]])\n",
    "    \n",
    "    assert np.allclose(result.data, expected), f\"Expected {expected}, got {result.data}\"\n",
    "    print(\"✅ 2x2 matrix multiplication\")\n",
    "    \n",
    "    # Test case 2: Non-square matrices\n",
    "    a = Tensor([[1, 2, 3], [4, 5, 6]])  # 2x3\n",
    "    b = Tensor([[7, 8], [9, 10], [11, 12]])  # 3x2\n",
    "    result = matmul(a, b)\n",
    "    expected = np.array([[58, 64], [139, 154]])\n",
    "    \n",
    "    assert np.allclose(result.data, expected), f\"Expected {expected}, got {result.data}\"\n",
    "    print(\"✅ Non-square matrix multiplication\")\n",
    "    \n",
    "    # Test case 3: Vector-matrix multiplication\n",
    "    a = Tensor([[1, 2, 3]])  # 1x3 (row vector)\n",
    "    b = Tensor([[4], [5], [6]])  # 3x1 (column vector)\n",
    "    result = matmul(a, b)\n",
    "    expected = np.array([[32]])  # 1*4 + 2*5 + 3*6 = 32\n",
    "    \n",
    "    assert np.allclose(result.data, expected), f\"Expected {expected}, got {result.data}\"\n",
    "    print(\"✅ Vector-matrix multiplication\")\n",
    "    \n",
    "    print(\"🎉 All matrix multiplication tests passed!\")\n",
    "\n",
    "test_matmul()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce6f0c4c",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "# Dense Layer - The Fundamental Neural Network Component\n",
    "\n",
    "Dense layers (also called Linear or Fully Connected layers) are the building blocks of neural networks. They apply the transformation: **output = input @ weights + bias**\n",
    "\n",
    "## Why Dense Layers Matter\n",
    "\n",
    "🧠 **Universal Function Approximators**: Dense layers can approximate any continuous function when stacked  \n",
    "🔧 **Parameter Learning**: Weights and biases are learned through backpropagation  \n",
    "🏗️ **Modular Design**: Dense layers compose into complex architectures (MLPs, transformers, etc.)  \n",
    "⚡ **Computational Efficiency**: Matrix operations leverage optimized linear algebra libraries  \n",
    "\n",
    "## Learning Objectives\n",
    "By implementing Dense layers, you'll understand:\n",
    "- How neural networks learn through adjustable parameters\n",
    "- The mathematical foundation underlying all neural network layers\n",
    "- Why proper parameter initialization is crucial for training success\n",
    "- How layer composition enables complex function approximation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6f7319a",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "dense-implementation",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class Dense:\n",
    "    \"\"\"\n",
    "    Dense (Fully Connected) Layer implementation.\n",
    "    \n",
    "    Applies the transformation: output = input @ weights + bias\n",
    "    \n",
    "    This is equivalent to PyTorch's nn.Linear layer.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, input_size: int, output_size: int, use_bias: bool = True):\n",
    "        \"\"\"\n",
    "        Initialize Dense layer with random weights and optional bias.\n",
    "        \n",
    "        Args:\n",
    "            input_size: Number of input features\n",
    "            output_size: Number of output features  \n",
    "            use_bias: Whether to include bias term\n",
    "        \n",
    "        TODO: Implement Dense layer initialization.\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Store input_size and output_size as instance variables\n",
    "        2. Initialize weights as Tensor with shape (input_size, output_size)\n",
    "        3. Use small random values: np.random.randn(...) * 0.1\n",
    "        4. Initialize bias as Tensor with shape (output_size,) if use_bias is True\n",
    "        5. Set bias to None if use_bias is False\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - Small random initialization prevents symmetry breaking\n",
    "        - Weight shape (input_size, output_size) enables matrix multiplication\n",
    "        - Bias allows shifting the output (like y-intercept in linear regression)\n",
    "        - PyTorch uses more sophisticated initialization (Xavier, Kaiming)\n",
    "        \n",
    "        IMPLEMENTATION HINTS:\n",
    "        - Use np.random.randn() for Gaussian random numbers\n",
    "        - Scale by 0.1 to keep initial values small\n",
    "        - Remember to wrap numpy arrays in Tensor()\n",
    "        - Store use_bias flag for forward pass logic\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        self.input_size = input_size\n",
    "        self.output_size = output_size\n",
    "        self.use_bias = use_bias\n",
    "        \n",
    "        # Initialize weights with small random values\n",
    "        # Shape: (input_size, output_size) for matrix multiplication\n",
    "        weight_data = np.random.randn(input_size, output_size) * 0.1\n",
    "        self.weights = Tensor(weight_data)\n",
    "        \n",
    "        # Initialize bias if requested\n",
    "        if use_bias:\n",
    "            bias_data = np.random.randn(output_size) * 0.1\n",
    "            self.bias = Tensor(bias_data)\n",
    "        else:\n",
    "            self.bias = None\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Forward pass through the Dense layer.\n",
    "        \n",
    "        Args:\n",
    "            x: Input tensor (shape: ..., input_size)\n",
    "        \n",
    "        Returns:\n",
    "            Output tensor (shape: ..., output_size)\n",
    "        \n",
    "        TODO: Implement forward pass: output = input @ weights + bias\n",
    "        \n",
    "        STEP-BY-STEP IMPLEMENTATION:\n",
    "        1. Perform matrix multiplication: output = matmul(x, self.weights)\n",
    "        2. If bias exists, add it: output = output + bias\n",
    "        3. Return the result\n",
    "        \n",
    "        LEARNING CONNECTIONS:\n",
    "        - This is the core neural network transformation\n",
    "        - Matrix multiplication scales input features to output features\n",
    "        - Bias provides offset (like y-intercept in linear equations)\n",
    "        - Broadcasting handles different batch sizes automatically\n",
    "        \n",
    "        IMPLEMENTATION HINTS:\n",
    "        - Use the matmul function you implemented above\n",
    "        - Check if self.bias is not None before adding\n",
    "        - Tensor addition should work automatically via broadcasting\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Matrix multiplication: input @ weights\n",
    "        output = matmul(x, self.weights)\n",
    "        \n",
    "        # Add bias if it exists\n",
    "        if self.bias is not None:\n",
    "            output = output + self.bias\n",
    "        \n",
    "        return output\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"Make the layer callable: layer(x) instead of layer.forward(x)\"\"\"\n",
    "        return self.forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18b34b1b",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Testing Dense Layer\n",
    "\n",
    "Let's verify our Dense layer works correctly with comprehensive tests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a737b71",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-dense",
     "locked": true,
     "points": 3,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def test_dense_layer():\n",
    "    \"\"\"Test Dense layer implementation.\"\"\"\n",
    "    print(\"🧪 Testing Dense Layer...\")\n",
    "    \n",
    "    # Test case 1: Basic functionality\n",
    "    layer = Dense(input_size=3, output_size=2)\n",
    "    input_tensor = Tensor([[1.0, 2.0, 3.0]])  # Shape: (1, 3)\n",
    "    output = layer.forward(input_tensor)\n",
    "    \n",
    "    # Check output shape\n",
    "    assert output.shape == (1, 2), f\"Expected shape (1, 2), got {output.shape}\"\n",
    "    print(\"✅ Output shape correct\")\n",
    "    \n",
    "    # Test case 2: No bias\n",
    "    layer_no_bias = Dense(input_size=2, output_size=3, use_bias=False)\n",
    "    assert layer_no_bias.bias is None, \"Bias should be None when use_bias=False\"\n",
    "    print(\"✅ No bias option works\")\n",
    "    \n",
    "    # Test case 3: Multiple samples (batch processing)\n",
    "    batch_input = Tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])  # Shape: (3, 2)\n",
    "    layer_batch = Dense(input_size=2, output_size=2)\n",
    "    batch_output = layer_batch.forward(batch_input)\n",
    "    \n",
    "    assert batch_output.shape == (3, 2), f\"Expected shape (3, 2), got {batch_output.shape}\"\n",
    "    print(\"✅ Batch processing works\")\n",
    "    \n",
    "    # Test case 4: Callable interface\n",
    "    callable_output = layer_batch(batch_input)\n",
    "    assert np.allclose(callable_output.data, batch_output.data), \"Callable interface should match forward()\"\n",
    "    print(\"✅ Callable interface works\")\n",
    "    \n",
    "    # Test case 5: Parameter initialization\n",
    "    layer_init = Dense(input_size=10, output_size=5)\n",
    "    assert layer_init.weights.shape == (10, 5), f\"Expected weights shape (10, 5), got {layer_init.weights.shape}\"\n",
    "    assert layer_init.bias.shape == (5,), f\"Expected bias shape (5,), got {layer_init.bias.shape}\"\n",
    "    \n",
    "    # Check that weights are reasonably small (good initialization)\n",
    "    assert np.abs(layer_init.weights.data).mean() < 1.0, \"Weights should be small for good initialization\"\n",
    "    print(\"✅ Parameter initialization correct\")\n",
    "    \n",
    "    print(\"🎉 All Dense layer tests passed!\")\n",
    "\n",
    "test_dense_layer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0c11bbc",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "# Systems Analysis: Memory and Performance Characteristics\n",
    "\n",
    "Let's analyze the memory usage and computational complexity of our layer implementations.\n",
    "\n",
    "## Memory Analysis\n",
    "- **Dense Layer Storage**: input_size × output_size weights + output_size bias terms\n",
    "- **Forward Pass Memory**: Input tensor + weight tensor + output tensor (temporary storage)\n",
    "- **Scaling Behavior**: Memory grows quadratically with layer size\n",
    "\n",
    "## Computational Complexity\n",
    "- **Matrix Multiplication**: O(batch_size × input_size × output_size)\n",
    "- **Bias Addition**: O(batch_size × output_size)\n",
    "- **Total**: Dominated by matrix multiplication for large layers\n",
    "\n",
    "## Production Insights\n",
    "In production ML systems:\n",
    "- **Memory Management**: PyTorch uses memory pools to avoid frequent allocation/deallocation\n",
    "- **Compute Optimization**: BLAS libraries (MKL, OpenBLAS) optimize matrix operations for specific hardware\n",
    "- **GPU Acceleration**: CUDA kernels parallelize matrix operations across thousands of cores\n",
    "- **Mixed Precision**: Using float16 instead of float32 can halve memory usage with minimal accuracy loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c23259fb",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "memory-analysis",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def analyze_layer_memory():\n",
    "    \"\"\"Analyze memory usage of different layer sizes.\"\"\"\n",
    "    print(\"📊 Layer Memory Analysis\")\n",
    "    print(\"=\" * 40)\n",
    "    \n",
    "    layer_sizes = [(10, 10), (100, 100), (1000, 1000), (784, 128), (128, 10)]\n",
    "    \n",
    "    for input_size, output_size in layer_sizes:\n",
    "        # Calculate parameter count\n",
    "        weight_params = input_size * output_size\n",
    "        bias_params = output_size\n",
    "        total_params = weight_params + bias_params\n",
    "        \n",
    "        # Calculate memory usage (assuming float32 = 4 bytes)\n",
    "        memory_mb = total_params * 4 / (1024 * 1024)\n",
    "        \n",
    "        print(f\"  {input_size:4d} → {output_size:4d}: {total_params:,} params, {memory_mb:.3f} MB\")\n",
    "    \n",
    "    print(\"\\n🔍 Key Insights:\")\n",
    "    print(\"  • Memory grows quadratically with layer width\")\n",
    "    print(\"  • Large layers (1000×1000) use significant memory\")\n",
    "    print(\"  • Modern networks balance width vs depth for efficiency\")\n",
    "\n",
    "analyze_layer_memory()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cff8208b",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "# ML Systems Thinking: Interactive Questions\n",
    "\n",
    "Let's explore the deeper implications of our layer implementations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "052d28f8",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "systems-thinking",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def explore_layer_scaling():\n",
    "    \"\"\"Explore how layer operations scale with size.\"\"\"\n",
    "    print(\"🤔 Scaling Analysis: Matrix Multiplication Performance\")\n",
    "    print(\"=\" * 55)\n",
    "    \n",
    "    sizes = [64, 128, 256, 512]\n",
    "    \n",
    "    for size in sizes:\n",
    "        # Estimate FLOPs for square matrix multiplication\n",
    "        flops = 2 * size * size * size  # 2 operations per multiply-add\n",
    "        \n",
    "        # Estimate memory bandwidth (reading A, B, writing C)\n",
    "        memory_ops = 3 * size * size  # Elements read/written\n",
    "        memory_mb = memory_ops * 4 / (1024 * 1024)  # float32 = 4 bytes\n",
    "        \n",
    "        print(f\"  Size {size:3d}×{size:3d}: {flops/1e6:.1f} MFLOPS, {memory_mb:.2f} MB transfers\")\n",
    "    \n",
    "    print(\"\\n💡 Performance Insights:\")\n",
    "    print(\"  • FLOPs grow cubically (O(n³)) with matrix size\")\n",
    "    print(\"  • Memory bandwidth grows quadratically (O(n²))\")\n",
    "    print(\"  • Large matrices become memory-bound, not compute-bound\")\n",
    "    print(\"  • This is why GPUs excel: high memory bandwidth + parallel compute\")\n",
    "\n",
    "explore_layer_scaling()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85912da1",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🤔 ML Systems Thinking: Interactive Questions\n",
    "\n",
    "Now that you've implemented the core components, let's think about their implications for ML systems:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "568bc5e5",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "question-1",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Question 1: Memory vs Computation Trade-offs\n",
    "\"\"\"\n",
    "🤔 **Question 1: Memory vs Computation Analysis**\n",
    "\n",
    "You're designing a neural network for deployment on a mobile device with limited memory (1GB RAM) but decent compute power.\n",
    "\n",
    "You have two architecture options:\n",
    "A) Wide network: 784 → 2048 → 2048 → 10 (3 layers, wide)\n",
    "B) Deep network: 784 → 256 → 256 → 256 → 256 → 10 (5 layers, narrow)\n",
    "\n",
    "Calculate the memory requirements for each option and explain which you'd choose for mobile deployment and why.\n",
    "\n",
    "Consider:\n",
    "- Parameter storage requirements\n",
    "- Intermediate activation storage during forward pass\n",
    "- Training vs inference memory requirements\n",
    "- How your choice affects model capacity and accuracy\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9310cc11",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "question-2",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Question 2: Performance Optimization\n",
    "\"\"\"\n",
    "🤔 **Question 2: Production Performance Optimization**\n",
    "\n",
    "Your Dense layer implementation works correctly, but you notice it's slower than PyTorch's nn.Linear on the same hardware.\n",
    "\n",
    "Investigate and explain:\n",
    "1. Why might our implementation be slower? (Hint: think about underlying linear algebra libraries)\n",
    "2. What optimization techniques do production frameworks use?\n",
    "3. How would you modify our implementation to approach production performance?\n",
    "4. When might our simple implementation actually be preferable?\n",
    "\n",
    "Research areas to consider:\n",
    "- BLAS (Basic Linear Algebra Subprograms) libraries\n",
    "- Memory layout and cache efficiency\n",
    "- Vectorization and SIMD instructions\n",
    "- GPU kernel optimization\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b20402dd",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "question-3",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Question 3: Scaling and Architecture Design\n",
    "\"\"\"\n",
    "🤔 **Question 3: Systems Architecture Scaling**\n",
    "\n",
    "Modern transformer models like GPT-3 have billions of parameters, primarily in Dense layers.\n",
    "\n",
    "Analyze the scaling challenges:\n",
    "1. How does memory requirement scale with model size? Calculate the memory needed for a 175B parameter model.\n",
    "2. What are the computational bottlenecks during training vs inference?\n",
    "3. How do systems like distributed training address these scaling challenges?\n",
    "4. Why do large models use techniques like gradient checkpointing and model parallelism?\n",
    "\n",
    "Systems considerations:\n",
    "- Memory hierarchy (L1/L2/L3 cache, RAM, storage)\n",
    "- Network bandwidth for distributed training\n",
    "- GPU memory constraints and model sharding\n",
    "- Inference optimization for production serving\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11a35821",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "# Comprehensive Testing and Integration\n",
    "\n",
    "Let's run a comprehensive test suite to verify all our implementations work correctly together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92844ec1",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "comprehensive-tests",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def run_comprehensive_tests():\n",
    "    \"\"\"Run comprehensive tests of all layer functionality.\"\"\"\n",
    "    print(\"🔬 Comprehensive Layer Testing Suite\")\n",
    "    print(\"=\" * 45)\n",
    "    \n",
    "    # Test 1: Matrix multiplication edge cases\n",
    "    print(\"\\n1. Matrix Multiplication Edge Cases:\")\n",
    "    \n",
    "    # Single element\n",
    "    a = Tensor([[5]])\n",
    "    b = Tensor([[3]])\n",
    "    result = matmul(a, b)\n",
    "    assert result.data[0, 0] == 15, \"Single element multiplication failed\"\n",
    "    print(\"   ✅ Single element multiplication\")\n",
    "    \n",
    "    # Identity matrix\n",
    "    identity = Tensor([[1, 0], [0, 1]])\n",
    "    test_matrix = Tensor([[2, 3], [4, 5]])\n",
    "    result = matmul(test_matrix, identity)\n",
    "    assert np.allclose(result.data, test_matrix.data), \"Identity multiplication failed\"\n",
    "    print(\"   ✅ Identity matrix multiplication\")\n",
    "    \n",
    "    # Test 2: Dense layer composition\n",
    "    print(\"\\n2. Dense Layer Composition:\")\n",
    "    \n",
    "    # Create a simple 2-layer network\n",
    "    layer1 = Dense(4, 3)\n",
    "    layer2 = Dense(3, 2)\n",
    "    \n",
    "    # Test data flow\n",
    "    input_data = Tensor([[1, 2, 3, 4]])\n",
    "    hidden = layer1(input_data)\n",
    "    output = layer2(hidden)\n",
    "    \n",
    "    assert output.shape == (1, 2), f\"Expected final output shape (1, 2), got {output.shape}\"\n",
    "    print(\"   ✅ Multi-layer composition\")\n",
    "    \n",
    "    # Test 3: Batch processing\n",
    "    print(\"\\n3. Batch Processing:\")\n",
    "    \n",
    "    batch_size = 10\n",
    "    batch_input = Tensor(np.random.randn(batch_size, 4))\n",
    "    batch_hidden = layer1(batch_input)\n",
    "    batch_output = layer2(batch_hidden)\n",
    "    \n",
    "    assert batch_output.shape == (batch_size, 2), f\"Expected batch output shape ({batch_size}, 2), got {batch_output.shape}\"\n",
    "    print(\"   ✅ Batch processing\")\n",
    "    \n",
    "    # Test 4: Parameter access and modification\n",
    "    print(\"\\n4. Parameter Management:\")\n",
    "    \n",
    "    layer = Dense(5, 3)\n",
    "    original_weights = layer.weights.data.copy()\n",
    "    \n",
    "    # Simulate parameter update\n",
    "    layer.weights = Tensor(original_weights + 0.1)\n",
    "    \n",
    "    assert not np.allclose(layer.weights.data, original_weights), \"Parameter update failed\"\n",
    "    print(\"   ✅ Parameter modification\")\n",
    "    \n",
    "    print(\"\\n🎉 All comprehensive tests passed!\")\n",
    "    print(\"   Your layer implementations are ready for neural network construction!\")\n",
    "\n",
    "run_comprehensive_tests()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01d5c09c",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Module Summary\n",
    "\n",
    "## 🎯 What You've Accomplished\n",
    "\n",
    "You've successfully implemented the fundamental building blocks of neural networks:\n",
    "\n",
    "### ✅ **Core Implementations**\n",
    "- **Matrix Multiplication**: The computational primitive underlying all neural network operations\n",
    "- **Dense Layer**: Complete implementation with proper parameter initialization and forward propagation\n",
    "- **Composition Patterns**: How layers stack together to form complex function approximators\n",
    "\n",
    "### ✅ **Systems Understanding**\n",
    "- **Memory Analysis**: How layer size affects memory usage and why this matters for deployment\n",
    "- **Performance Characteristics**: Understanding computational complexity and scaling behavior\n",
    "- **Production Context**: Connection to real-world ML systems and optimization techniques\n",
    "\n",
    "### ✅ **ML Engineering Skills**\n",
    "- **Parameter Management**: How neural networks store and update learnable parameters\n",
    "- **Batch Processing**: Efficient handling of multiple data samples simultaneously\n",
    "- **Architecture Design**: Trade-offs between network width, depth, and resource requirements\n",
    "\n",
    "## 🔗 **Connection to Production ML Systems**\n",
    "\n",
    "Your implementations mirror the core concepts used in:\n",
    "- **PyTorch's nn.Linear**: Same mathematical operations with production optimizations\n",
    "- **TensorFlow's Dense layers**: Identical parameter structure and forward pass logic\n",
    "- **Transformer architectures**: Dense layers form the foundation of modern language models\n",
    "- **Computer vision models**: ConvNets use similar principles with spatial structure\n",
    "\n",
    "## 🚀 **What's Next**\n",
    "\n",
    "With solid layer implementations, you're ready to:\n",
    "- **Compose** these layers into complete neural networks\n",
    "- **Add** nonlinear activations to enable complex function approximation\n",
    "- **Implement** training algorithms to learn from data\n",
    "- **Scale** to larger, more sophisticated architectures\n",
    "\n",
    "## 💡 **Key Systems Insights**\n",
    "\n",
    "1. **Matrix multiplication is the computational bottleneck** in neural networks\n",
    "2. **Memory layout and access patterns** often matter more than raw compute power\n",
    "3. **Layer composition** is the fundamental abstraction for building complex ML systems\n",
    "4. **Parameter initialization and management** directly affects training success\n",
    "\n",
    "You now understand the mathematical and computational foundations that enable neural networks to learn complex patterns from data!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb70ede",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "final-demo",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    print(\"🔥 TinyTorch Layers Module - Final Demo\")\n",
    "    print(\"=\" * 50)\n",
    "    \n",
    "    # Create a simple neural network architecture\n",
    "    print(\"\\n🏗️ Building a 3-layer neural network:\")\n",
    "    layer1 = Dense(784, 128)  # Input layer (like MNIST images)\n",
    "    layer2 = Dense(128, 64)   # Hidden layer\n",
    "    layer3 = Dense(64, 10)    # Output layer (10 classes)\n",
    "    \n",
    "    print(f\"  Layer 1: {layer1.input_size} → {layer1.output_size} ({layer1.weights.data.size:,} parameters)\")\n",
    "    print(f\"  Layer 2: {layer2.input_size} → {layer2.output_size} ({layer2.weights.data.size:,} parameters)\")\n",
    "    print(f\"  Layer 3: {layer3.input_size} → {layer3.output_size} ({layer3.weights.data.size:,} parameters)\")\n",
    "    \n",
    "    # Simulate forward pass\n",
    "    print(\"\\n🚀 Forward pass through network:\")\n",
    "    batch_size = 32\n",
    "    input_data = Tensor(np.random.randn(batch_size, 784))\n",
    "    \n",
    "    print(f\"  Input shape: {input_data.shape}\")\n",
    "    hidden1 = layer1(input_data)\n",
    "    print(f\"  After layer 1: {hidden1.shape}\")\n",
    "    hidden2 = layer2(hidden1)\n",
    "    print(f\"  After layer 2: {hidden2.shape}\")\n",
    "    output = layer3(hidden2)\n",
    "    print(f\"  Final output: {output.shape}\")\n",
    "    \n",
    "    # Calculate total parameters\n",
    "    total_params = (layer1.weights.data.size + layer1.bias.data.size + \n",
    "                   layer2.weights.data.size + layer2.bias.data.size +\n",
    "                   layer3.weights.data.size + layer3.bias.data.size)\n",
    "    \n",
    "    print(f\"\\n📊 Network Statistics:\")\n",
    "    print(f\"  Total parameters: {total_params:,}\")\n",
    "    print(f\"  Memory usage: ~{total_params * 4 / 1024 / 1024:.2f} MB (float32)\")\n",
    "    print(f\"  Forward pass: {batch_size} samples processed simultaneously\")\n",
    "    \n",
    "    print(\"\\n✅ Neural network construction complete!\")\n",
    "    print(\"Ready for activation functions and training algorithms!\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}