{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3405f85e",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Module 05: Autograd ⚡ - The Gradient Engine\n",
    "\n",
    "Welcome to Module 05! Today you'll awaken the gradient engine and unlock automatic differentiation.\n",
    "\n",
    "## 🔗 Prerequisites & Progress\n",
    "**You've Built**: Tensor operations, activations, layers, and loss functions  \n",
    "**You'll Build**: The autograd system that computes gradients automatically  \n",
    "**You'll Enable**: Learning! Training! The ability to optimize neural networks!\n",
    "\n",
    "**Connection Map**:\n",
    "```\n",
    "Modules 01-04 → Autograd → Training (Module 06-07)\n",
    "(forward pass) (backward pass) (learning loops)\n",
    "```\n",
    "\n",
    "## Learning Objectives ⭐⭐\n",
    "By the end of this module, you will:\n",
    "1. **Enhance Tensor** with automatic differentiation capabilities\n",
    "2. **Build computation graphs** that track operations for gradient flow\n",
    "3. **Implement backward()** method for reverse-mode differentiation\n",
    "4. **Create Function classes** for operation-specific gradient rules\n",
    "5. **Test gradient correctness** with mathematical validation\n",
    "\n",
    "**CRITICAL**: This module enhances the existing Tensor class - no new wrapper classes needed!\n",
    "\n",
    "## 📦 Where This Code Lives in the Final Package\n",
    "\n",
    "**Learning Side:** You work in `modules/05_autograd/autograd_dev.py`  \n",
    "**Building Side:** Code exports to `tinytorch.core.autograd`\n",
    "\n",
    "```python\n",
    "# How to use this module:\n",
    "from tinytorch.core.autograd import Function, enable_autograd\n",
    "```\n",
    "\n",
    "**Why this matters:**\n",
    "- **Learning:** Complete autograd system enabling automatic differentiation\n",
    "- **Production:** PyTorch-style computational graph and backward pass\n",
    "- **Consistency:** All gradient operations in core.autograd\n",
    "- **Integration:** Enhances existing Tensor without breaking anything\n",
    "\n",
    "Let's build the gradient engine that makes neural networks learn! 🚀"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "261c3177",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "imports",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| default_exp core.autograd\n",
    "#| export\n",
    "\n",
    "import numpy as np\n",
    "from typing import Optional, List, Tuple\n",
    "import sys\n",
    "import os\n",
    "\n",
    "from tinytorch.core.tensor import Tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "984dc0f4",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 1. Introduction: What is Automatic Differentiation?\n",
    "\n",
    "Automatic differentiation (autograd) is the magic that makes neural networks learn. Instead of manually computing gradients for every parameter, autograd tracks operations and automatically computes gradients via the chain rule.\n",
    "\n",
    "### The Challenge\n",
    "In previous modules, you implemented layers and loss functions. To train a model, you need:\n",
    "```\n",
    "Loss = f(W₃, f(W₂, f(W₁, x)))\n",
    "∂Loss/∂W₁ = ?  ∂Loss/∂W₂ = ?  ∂Loss/∂W₃ = ?\n",
    "```\n",
    "\n",
    "Manual gradient computation becomes impossible for complex models with millions of parameters.\n",
    "\n",
    "### The Solution: Computational Graphs\n",
    "```\n",
    "Forward Pass:  x → Linear₁ → ReLU → Linear₂ → Loss\n",
    "Backward Pass: ∇x ← ∇Linear₁ ← ∇ReLU ← ∇Linear₂ ← ∇Loss\n",
    "```\n",
    "\n",
    "**Complete Autograd Process Visualization:**\n",
    "```\n",
    "┌─ FORWARD PASS ──────────────────────────────────────────────┐\n",
    "│                                                             │\n",
    "│ x ──┬── W₁ ──┐                                              │\n",
    "│     │        ├──[Linear₁]──→ z₁ ──[ReLU]──→ a₁ ──┬── W₂ ──┐ │\n",
    "│     └── b₁ ──┘                               │        ├─→ Loss\n",
    "│                                              └── b₂ ──┘ │\n",
    "│                                                             │\n",
    "└─ COMPUTATION GRAPH BUILT ──────────────────────────────────┘\n",
    "                             │\n",
    "                             ▼\n",
    "┌─ BACKWARD PASS ─────────────────────────────────────────────┐\n",
    "│                                                             │\n",
    "│∇x ←┬← ∇W₁ ←┐                                               │\n",
    "│    │       ├←[Linear₁]←─ ∇z₁ ←[ReLU]← ∇a₁ ←┬← ∇W₂ ←┐      │\n",
    "│    └← ∇b₁ ←┘                             │       ├← ∇Loss  │\n",
    "│                                          └← ∇b₂ ←┘      │\n",
    "│                                                             │\n",
    "└─ GRADIENTS COMPUTED ───────────────────────────────────────┘\n",
    "\n",
    "Key Insight: Each [operation] stores how to compute its backward pass.\n",
    "The chain rule automatically flows gradients through the entire graph.\n",
    "```\n",
    "\n",
    "Each operation records how to compute its backward pass. The chain rule connects them all."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4859deb3",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 2. Foundations: The Chain Rule in Action\n",
    "\n",
    "### Mathematical Foundation\n",
    "For composite functions: f(g(x)), the derivative is:\n",
    "```\n",
    "df/dx = (df/dg) × (dg/dx)\n",
    "```\n",
    "\n",
    "### Computational Graph Example\n",
    "```\n",
    "Simple computation: L = (x * y + 5)²\n",
    "\n",
    "Forward Pass:\n",
    "  x=2 ──┐\n",
    "        ├──[×]──→ z=6 ──[+5]──→ w=11 ──[²]──→ L=121\n",
    "  y=3 ──┘\n",
    "\n",
    "Backward Pass (Chain Rule in Action):\n",
    "  ∂L/∂x = ∂L/∂w × ∂w/∂z × ∂z/∂x\n",
    "        = 2w  ×  1  ×  y\n",
    "        = 2(11) × 1 × 3 = 66\n",
    "\n",
    "  ∂L/∂y = ∂L/∂w × ∂w/∂z × ∂z/∂y\n",
    "        = 2w  ×  1  ×  x\n",
    "        = 2(11) × 1 × 2 = 44\n",
    "\n",
    "Gradient Flow Visualization:\n",
    "  ∇x=66 ←──┐\n",
    "           ├──[×]←── ∇z=22 ←──[+]←── ∇w=22 ←──[²]←── ∇L=1\n",
    "  ∇y=44 ←──┘\n",
    "```\n",
    "\n",
    "### Memory Layout During Backpropagation\n",
    "```\n",
    "Computation Graph Memory Structure:\n",
    "┌─────────────────────────────────────────────────────────┐\n",
    "│ Forward Pass (stored for backward)                      │\n",
    "├─────────────────────────────────────────────────────────┤\n",
    "│ Node 1: x=2 (leaf, requires_grad=True) │ grad: None→66  │\n",
    "│ Node 2: y=3 (leaf, requires_grad=True) │ grad: None→44  │\n",
    "│ Node 3: z=x*y (MulFunction)            │ grad: None→22  │\n",
    "│         saved: (x=2, y=3)              │ inputs: [x,y]  │\n",
    "│ Node 4: w=z+5 (AddFunction)            │ grad: None→22  │\n",
    "│         saved: (z=6, 5)                │ inputs: [z]    │\n",
    "│ Node 5: L=w² (PowFunction)             │ grad: 1        │\n",
    "│         saved: (w=11)                  │ inputs: [w]    │\n",
    "└─────────────────────────────────────────────────────────┘\n",
    "\n",
    "Memory Cost: 2× parameters (data + gradients) + graph overhead\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfc1da56",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 3. Implementation: Building the Autograd Engine\n",
    "\n",
    "Let's implement the autograd system step by step. We'll enhance the existing Tensor class and create supporting infrastructure.\n",
    "\n",
    "### The Function Architecture\n",
    "\n",
    "Every differentiable operation needs two things:\n",
    "1. **Forward pass**: Compute the result\n",
    "2. **Backward pass**: Compute gradients for inputs\n",
    "\n",
    "```\n",
    "Function Class Design:\n",
    "┌─────────────────────────────────────┐\n",
    "│ Function (Base Class)               │\n",
    "├─────────────────────────────────────┤\n",
    "│ • saved_tensors    ← Store data     │\n",
    "│ • apply()          ← Compute grads  │\n",
    "└─────────────────────────────────────┘\n",
    "          ↑\n",
    "    ┌─────┴─────┬─────────┬──────────┐\n",
    "    │           │         │          │\n",
    "┌───▼────┐ ┌────▼───┐ ┌───▼────┐ ┌───▼────┐\n",
    "│  Add   │ │  Mul   │ │ Matmul │ │  Sum   │\n",
    "│Backward│ │Backward│ │Backward│ │Backward│\n",
    "└────────┘ └────────┘ └────────┘ └────────┘\n",
    "```\n",
    "\n",
    "Each operation inherits from Function and implements specific gradient rules."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a252129",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### Function Base Class - The Foundation of Autograd\n",
    "\n",
    "The Function class is the foundation that makes autograd possible. Every differentiable operation (addition, multiplication, etc.) inherits from this class.\n",
    "\n",
    "**Why Functions Matter:**\n",
    "- They remember inputs needed for backward pass\n",
    "- They implement gradient computation via apply()\n",
    "- They connect to form computation graphs\n",
    "- They enable the chain rule to flow gradients\n",
    "\n",
    "**The Pattern:**\n",
    "```\n",
    "Forward:  inputs → Function.forward() → output\n",
    "Backward: grad_output → Function.apply() → grad_inputs\n",
    "```\n",
    "\n",
    "This pattern enables the chain rule to flow gradients through complex computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7311a2dd",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "function-base",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class Function:\n",
    "    \"\"\"\n",
    "    Base class for differentiable operations.\n",
    "\n",
    "    Every operation that needs gradients (add, multiply, matmul, etc.)\n",
    "    will inherit from this class and implement the apply() method.\n",
    "    \n",
    "    **Key Concepts:**\n",
    "    - **saved_tensors**: Store inputs needed for backward pass\n",
    "    - **apply()**: Compute gradients using chain rule\n",
    "    - **next_functions**: Track computation graph connections\n",
    "    \n",
    "    **Example Usage:**\n",
    "    ```python\n",
    "    class AddBackward(Function):\n",
    "        def apply(self, grad_output):\n",
    "            # Addition distributes gradients equally\n",
    "            return grad_output, grad_output\n",
    "    ```\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, *tensors):\n",
    "        \"\"\"\n",
    "        Initialize function with input tensors.\n",
    "        \n",
    "        Args:\n",
    "            *tensors: Input tensors that will be saved for backward pass\n",
    "        \"\"\"\n",
    "        self.saved_tensors = tensors\n",
    "        self.next_functions = []\n",
    "\n",
    "        # Build computation graph connections\n",
    "        for t in tensors:\n",
    "            if isinstance(t, Tensor) and t.requires_grad:\n",
    "                if hasattr(t, '_grad_fn'):\n",
    "                    self.next_functions.append(t._grad_fn)\n",
    "\n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"\n",
    "        Compute gradients for inputs.\n",
    "        \n",
    "        Args:\n",
    "            grad_output: Gradient flowing backward from the output\n",
    "            \n",
    "        Returns:\n",
    "            Tuple of gradients for each input tensor\n",
    "            \n",
    "        **Must be implemented by subclasses**\n",
    "        \"\"\"\n",
    "        raise NotImplementedError(\"Each Function must implement apply() method\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c03db390",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### Operation Functions - Implementing Gradient Rules\n",
    "\n",
    "Now we'll implement specific operations that compute gradients correctly. Each operation has mathematical rules for how gradients flow backward.\n",
    "\n",
    "**Gradient Flow Visualization:**\n",
    "```\n",
    "Addition (z = a + b):\n",
    "    ∂z/∂a = 1    ∂z/∂b = 1\n",
    "\n",
    "    a ──┐           grad_a ←──┐\n",
    "        ├─[+]─→ z          ├─[+]←── grad_z\n",
    "    b ──┘           grad_b ←──┘\n",
    "\n",
    "Multiplication (z = a * b):\n",
    "    ∂z/∂a = b    ∂z/∂b = a\n",
    "\n",
    "    a ──┐           grad_a = grad_z * b\n",
    "        ├─[×]─→ z\n",
    "    b ──┘           grad_b = grad_z * a\n",
    "\n",
    "Matrix Multiplication (Z = A @ B):\n",
    "    ∂Z/∂A = grad_Z @ B.T\n",
    "    ∂Z/∂B = A.T @ grad_Z\n",
    "\n",
    "    A ──┐           grad_A = grad_Z @ B.T\n",
    "        ├─[@]─→ Z\n",
    "    B ──┘           grad_B = A.T @ grad_Z\n",
    "```\n",
    "\n",
    "Each operation stores the inputs it needs for computing gradients."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c58b717a",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### AddBackward - Gradient Rules for Addition\n",
    "\n",
    "Addition is the simplest gradient operation: gradients flow unchanged to both inputs.\n",
    "\n",
    "**Mathematical Principle:**\n",
    "```\n",
    "If z = a + b, then:\n",
    "∂z/∂a = 1  (gradient of z w.r.t. a)\n",
    "∂z/∂b = 1  (gradient of z w.r.t. b)\n",
    "\n",
    "By chain rule:\n",
    "∂Loss/∂a = ∂Loss/∂z × ∂z/∂a = grad_output × 1 = grad_output\n",
    "∂Loss/∂b = ∂Loss/∂z × ∂z/∂b = grad_output × 1 = grad_output\n",
    "```\n",
    "\n",
    "**Broadcasting Challenge:**\n",
    "When tensors have different shapes, NumPy broadcasts automatically in forward pass,\n",
    "but we must \"unbroadcast\" gradients in backward pass to match original shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a96c73",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "add-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class AddBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for tensor addition.\n",
    "    \n",
    "    **Mathematical Rule:** If z = a + b, then ∂z/∂a = 1 and ∂z/∂b = 1\n",
    "    \n",
    "    **Key Insight:** Addition distributes gradients equally to both inputs.\n",
    "    The gradient flowing backward is passed unchanged to each input.\n",
    "    \n",
    "    **Broadcasting Handling:** When input shapes differ due to broadcasting,\n",
    "    we sum gradients appropriately to match original tensor shapes.\n",
    "    \"\"\"\n",
    "\n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"\n",
    "        Compute gradients for addition.\n",
    "        \n",
    "        Args:\n",
    "            grad_output: Gradient flowing backward from output\n",
    "            \n",
    "        Returns:\n",
    "            Tuple of (grad_a, grad_b) for the two inputs\n",
    "            \n",
    "        **Mathematical Foundation:**\n",
    "        - ∂(a+b)/∂a = 1 → grad_a = grad_output\n",
    "        - ∂(a+b)/∂b = 1 → grad_b = grad_output\n",
    "        \"\"\"\n",
    "        a, b = self.saved_tensors\n",
    "        grad_a = grad_b = None\n",
    "\n",
    "        # Gradient for first input\n",
    "        if isinstance(a, Tensor) and a.requires_grad:\n",
    "            grad_a = grad_output\n",
    "\n",
    "        # Gradient for second input  \n",
    "        if isinstance(b, Tensor) and b.requires_grad:\n",
    "            grad_b = grad_output\n",
    "\n",
    "        return grad_a, grad_b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ddb8b58",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### MulBackward - Gradient Rules for Element-wise Multiplication\n",
    "\n",
    "Element-wise multiplication follows the product rule of calculus.\n",
    "\n",
    "**Mathematical Principle:**\n",
    "```\n",
    "If z = a * b (element-wise), then:\n",
    "∂z/∂a = b  (gradient w.r.t. a equals the other input)\n",
    "∂z/∂b = a  (gradient w.r.t. b equals the other input)\n",
    "\n",
    "By chain rule:\n",
    "∂Loss/∂a = grad_output * b\n",
    "∂Loss/∂b = grad_output * a\n",
    "```\n",
    "\n",
    "**Visual Example:**\n",
    "```\n",
    "Forward:  a=[2,3] * b=[4,5] = z=[8,15]\n",
    "Backward: grad_z=[1,1]\n",
    "          grad_a = grad_z * b = [1,1] * [4,5] = [4,5]\n",
    "          grad_b = grad_z * a = [1,1] * [2,3] = [2,3]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167d60c6",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "mul-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class MulBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for tensor multiplication.\n",
    "    \n",
    "    **Mathematical Rule:** If z = a * b, then ∂z/∂a = b and ∂z/∂b = a\n",
    "    \n",
    "    **Key Insight:** Each input's gradient equals the gradient output \n",
    "    multiplied by the OTHER input's value (product rule).\n",
    "    \n",
    "    **Applications:** Used in weight scaling, attention mechanisms,\n",
    "    and anywhere element-wise multiplication occurs.\n",
    "    \"\"\"\n",
    "\n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"\n",
    "        Compute gradients for multiplication.\n",
    "        \n",
    "        Args:\n",
    "            grad_output: Gradient flowing backward from output\n",
    "            \n",
    "        Returns:\n",
    "            Tuple of (grad_a, grad_b) for the two inputs\n",
    "            \n",
    "        **Mathematical Foundation:**\n",
    "        - ∂(a*b)/∂a = b → grad_a = grad_output * b\n",
    "        - ∂(a*b)/∂b = a → grad_b = grad_output * a\n",
    "        \"\"\"\n",
    "        a, b = self.saved_tensors\n",
    "        grad_a = grad_b = None\n",
    "\n",
    "        # Gradient for first input: grad_output * b\n",
    "        if isinstance(a, Tensor) and a.requires_grad:\n",
    "            if isinstance(b, Tensor):\n",
    "                grad_a = grad_output * b.data\n",
    "            else:\n",
    "                grad_a = grad_output * b\n",
    "\n",
    "        # Gradient for second input: grad_output * a\n",
    "        if isinstance(b, Tensor) and b.requires_grad:\n",
    "            grad_b = grad_output * a.data\n",
    "\n",
    "        return grad_a, grad_b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "526a5ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90e9e19c",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### MatmulBackward - Gradient Rules for Matrix Multiplication\n",
    "\n",
    "Matrix multiplication has more complex gradient rules based on matrix calculus.\n",
    "\n",
    "**Mathematical Principle:**\n",
    "```\n",
    "If Z = A @ B (matrix multiplication), then:\n",
    "∂Z/∂A = grad_Z @ B.T\n",
    "∂Z/∂B = A.T @ grad_Z\n",
    "```\n",
    "\n",
    "**Why These Rules Work:**\n",
    "```\n",
    "For element Z[i,j] = Σ_k A[i,k] * B[k,j]\n",
    "∂Z[i,j]/∂A[i,k] = B[k,j]  ← This gives us grad_Z @ B.T\n",
    "∂Z[i,j]/∂B[k,j] = A[i,k]  ← This gives us A.T @ grad_Z\n",
    "```\n",
    "\n",
    "**Dimension Analysis:**\n",
    "```\n",
    "Forward:  A(m×k) @ B(k×n) = Z(m×n)\n",
    "Backward: grad_Z(m×n) @ B.T(n×k) = grad_A(m×k) ✓\n",
    "          A.T(k×m) @ grad_Z(m×n) = grad_B(k×n) ✓\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c3ff8c4",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "matmul-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class MatmulBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for matrix multiplication.\n",
    "    \n",
    "    **Mathematical Rule:** If Z = A @ B, then:\n",
    "    - ∂Z/∂A = grad_Z @ B.T\n",
    "    - ∂Z/∂B = A.T @ grad_Z\n",
    "    \n",
    "    **Key Insight:** Matrix multiplication gradients involve transposing\n",
    "    one input and multiplying with the gradient output.\n",
    "    \n",
    "    **Applications:** Core operation in neural networks for weight updates\n",
    "    in linear layers, attention mechanisms, and transformers.\n",
    "    \"\"\"\n",
    "\n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"\n",
    "        Compute gradients for matrix multiplication.\n",
    "        \n",
    "        Args:\n",
    "            grad_output: Gradient flowing backward from output\n",
    "            \n",
    "        Returns:\n",
    "            Tuple of (grad_a, grad_b) for the two matrix inputs\n",
    "            \n",
    "        **Mathematical Foundation:**\n",
    "        - ∂(A@B)/∂A = grad_output @ B.T\n",
    "        - ∂(A@B)/∂B = A.T @ grad_output\n",
    "        \"\"\"\n",
    "        a, b = self.saved_tensors\n",
    "        grad_a = grad_b = None\n",
    "\n",
    "        # Gradient for first input: grad_output @ b.T\n",
    "        if isinstance(a, Tensor) and a.requires_grad:\n",
    "            grad_a = np.dot(grad_output, b.data.T)\n",
    "\n",
    "        # Gradient for second input: a.T @ grad_output\n",
    "        if isinstance(b, Tensor) and b.requires_grad:\n",
    "            grad_b = np.dot(a.data.T, grad_output)\n",
    "\n",
    "        return grad_a, grad_b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53f8163c",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### SumBackward - Gradient Rules for Reduction Operations\n",
    "\n",
    "Sum operations reduce tensor dimensions, so gradients must be broadcast back.\n",
    "\n",
    "**Mathematical Principle:**\n",
    "```\n",
    "If z = sum(a), then ∂z/∂a[i] = 1 for all i\n",
    "Gradient is broadcasted from scalar result back to input shape.\n",
    "```\n",
    "\n",
    "**Gradient Broadcasting Examples:**\n",
    "```\n",
    "Case 1: Full sum\n",
    "  Forward:  a=[1,2,3] → sum() → z=6 (scalar)\n",
    "  Backward: grad_z=1 → broadcast → grad_a=[1,1,1]\n",
    "\n",
    "Case 2: Axis sum\n",
    "  Forward:  a=[[1,2],[3,4]] → sum(axis=0) → z=[4,6]\n",
    "  Backward: grad_z=[1,1] → broadcast → grad_a=[[1,1],[1,1]]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6b4ae48",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "sum-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class SumBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for tensor sum.\n",
    "    \n",
    "    **Mathematical Rule:** If z = sum(a), then ∂z/∂a[i] = 1 for all i\n",
    "    \n",
    "    **Key Insight:** Sum distributes the gradient equally to all input elements.\n",
    "    The gradient is broadcast from the reduced output back to input shape.\n",
    "    \n",
    "    **Applications:** Used in loss functions, mean operations, and\n",
    "    anywhere tensor reduction occurs.\n",
    "    \"\"\"\n",
    "\n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"\n",
    "        Compute gradients for sum operation.\n",
    "        \n",
    "        Args:\n",
    "            grad_output: Gradient flowing backward from output\n",
    "            \n",
    "        Returns:\n",
    "            Tuple containing gradient for the input tensor\n",
    "            \n",
    "        **Mathematical Foundation:**\n",
    "        - ∂sum(a)/∂a[i] = 1 → grad_a = ones_like(a) * grad_output\n",
    "        \"\"\"\n",
    "        tensor, = self.saved_tensors\n",
    "\n",
    "        if isinstance(tensor, Tensor) and tensor.requires_grad:\n",
    "            # Gradient is 1 for all elements, scaled by grad_output\n",
    "            return np.ones_like(tensor.data) * grad_output,\n",
    "        return None,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a559da",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b7d62de",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7be03d75",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🔬 Unit Test: Function Classes\n",
    "This test validates our Function classes compute gradients correctly.\n",
    "**What we're testing**: Forward and backward passes for each operation\n",
    "**Why it matters**: These are the building blocks of autograd\n",
    "**Expected**: Correct gradients that satisfy mathematical definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2da6c55b",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-function-classes",
     "locked": true,
     "points": 15
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_function_classes():\n",
    "    \"\"\"🔬 Test Function classes.\"\"\"\n",
    "    print(\"🔬 Unit Test: Function Classes...\")\n",
    "\n",
    "    # Test AddBackward\n",
    "    a = Tensor([1, 2, 3], requires_grad=True)\n",
    "    b = Tensor([4, 5, 6], requires_grad=True)\n",
    "    add_func = AddBackward(a, b)\n",
    "    grad_output = np.array([1, 1, 1])\n",
    "    grad_a, grad_b = add_func.apply(grad_output)\n",
    "    assert np.allclose(grad_a, grad_output), f\"AddBackward grad_a failed: {grad_a}\"\n",
    "    assert np.allclose(grad_b, grad_output), f\"AddBackward grad_b failed: {grad_b}\"\n",
    "\n",
    "    # Test MulBackward\n",
    "    mul_func = MulBackward(a, b)\n",
    "    grad_a, grad_b = mul_func.apply(grad_output)\n",
    "    assert np.allclose(grad_a, b.data), f\"MulBackward grad_a failed: {grad_a}\"\n",
    "    assert np.allclose(grad_b, a.data), f\"MulBackward grad_b failed: {grad_b}\"\n",
    "\n",
    "    # Test MatmulBackward\n",
    "    a_mat = Tensor([[1, 2], [3, 4]], requires_grad=True)\n",
    "    b_mat = Tensor([[5, 6], [7, 8]], requires_grad=True)\n",
    "    matmul_func = MatmulBackward(a_mat, b_mat)\n",
    "    grad_output = np.ones((2, 2))\n",
    "    grad_a, grad_b = matmul_func.apply(grad_output)\n",
    "    assert grad_a.shape == a_mat.shape, f\"MatmulBackward grad_a shape: {grad_a.shape}\"\n",
    "    assert grad_b.shape == b_mat.shape, f\"MatmulBackward grad_b shape: {grad_b.shape}\"\n",
    "\n",
    "    print(\"✅ Function classes work correctly!\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_function_classes()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "503cbbfd",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 4. Enhancing Tensor with Autograd Capabilities\n",
    "\n",
    "Now we'll enhance the existing Tensor class to use these gradient functions and build computation graphs automatically.\n",
    "\n",
    "**Computation Graph Formation:**\n",
    "```\n",
    "Before Autograd:             After Autograd:\n",
    "  x → operation → y           x → [Function] → y\n",
    "                                     ↓\n",
    "                               Stores operation\n",
    "                               for backward pass\n",
    "```\n",
    "\n",
    "**The Enhancement Strategy:**\n",
    "1. **Add backward() method** - Triggers gradient computation\n",
    "2. **Enhance operations** - Replace simple ops with gradient-tracking versions\n",
    "3. **Track computation graphs** - Each tensor remembers how it was created\n",
    "4. **Maintain compatibility** - All existing code continues to work\n",
    "\n",
    "**Critical Design Decision:**\n",
    "We enhance the EXISTING Tensor class rather than creating a new one.\n",
    "This means:\n",
    "- ✅ All previous modules continue working unchanged\n",
    "- ✅ No import changes needed\n",
    "- ✅ Gradients are \"opt-in\" via requires_grad=True\n",
    "- ✅ No confusion between Tensor types"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23ee7914",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### The enable_autograd() Function\n",
    "\n",
    "This function is the magic that brings gradients to life! It enhances the existing Tensor class with autograd capabilities by:\n",
    "\n",
    "1. **Monkey-patching operations** - Replaces `__add__`, `__mul__`, etc. with gradient-aware versions\n",
    "2. **Adding backward() method** - Implements reverse-mode automatic differentiation\n",
    "3. **Maintaining compatibility** - All existing code continues to work unchanged\n",
    "\n",
    "**The Pattern:**\n",
    "```\n",
    "Original: x + y → simple addition\n",
    "Enhanced: x + y → addition + gradient tracking (if requires_grad=True)\n",
    "```\n",
    "\n",
    "This approach follows PyTorch 2.0 style - clean, modern, and educational."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebf8d15",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "relu-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class ReLUBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for ReLU activation.\n",
    "    \n",
    "    ReLU: f(x) = max(0, x)\n",
    "    Derivative: f'(x) = 1 if x > 0, else 0\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, input_tensor):\n",
    "        \"\"\"Initialize with input tensor.\"\"\"\n",
    "        super().__init__(input_tensor)\n",
    "    \n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"Compute gradient for ReLU.\"\"\"\n",
    "        tensor, = self.saved_tensors\n",
    "        \n",
    "        if isinstance(tensor, Tensor) and tensor.requires_grad:\n",
    "            # ReLU gradient: 1 if x > 0, else 0\n",
    "            relu_grad = (tensor.data > 0).astype(np.float32)\n",
    "            return grad_output * relu_grad,\n",
    "        return None,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9270d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb9b24ed",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "sigmoid-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class SigmoidBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for sigmoid activation.\n",
    "    \n",
    "    Sigmoid: σ(x) = 1/(1 + exp(-x))\n",
    "    Derivative: σ'(x) = σ(x) * (1 - σ(x))\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, input_tensor, output_tensor):\n",
    "        \"\"\"\n",
    "        Initialize with both input and output.\n",
    "        \n",
    "        Args:\n",
    "            input_tensor: Original input to sigmoid\n",
    "            output_tensor: Output of sigmoid (saves recomputation)\n",
    "        \"\"\"\n",
    "        super().__init__(input_tensor)\n",
    "        self.output_data = output_tensor.data\n",
    "    \n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"Compute gradient for sigmoid.\"\"\"\n",
    "        tensor, = self.saved_tensors\n",
    "        \n",
    "        if isinstance(tensor, Tensor) and tensor.requires_grad:\n",
    "            # σ'(x) = σ(x) * (1 - σ(x))\n",
    "            sigmoid_grad = self.output_data * (1 - self.output_data)\n",
    "            return grad_output * sigmoid_grad,\n",
    "        return None,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34e47d63",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "mse-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class MSEBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for Mean Squared Error Loss.\n",
    "    \n",
    "    MSE: L = mean((predictions - targets)²)\n",
    "    Derivative: ∂L/∂predictions = 2 * (predictions - targets) / N\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, predictions, targets):\n",
    "        \"\"\"Initialize with predictions and targets.\"\"\"\n",
    "        super().__init__(predictions)\n",
    "        self.targets_data = targets.data\n",
    "        self.num_samples = np.size(targets.data)\n",
    "    \n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"Compute gradient for MSE loss.\"\"\"\n",
    "        predictions, = self.saved_tensors\n",
    "        \n",
    "        if isinstance(predictions, Tensor) and predictions.requires_grad:\n",
    "            # Gradient: 2 * (predictions - targets) / N\n",
    "            grad = 2.0 * (predictions.data - self.targets_data) / self.num_samples\n",
    "            \n",
    "            return grad * grad_output,\n",
    "        return None,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7d1bfe9",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "bce-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class BCEBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for Binary Cross-Entropy Loss.\n",
    "    \n",
    "    BCE: L = -[y*log(p) + (1-y)*log(1-p)]\n",
    "    Derivative: ∂L/∂p = (p - y) / (p*(1-p)*N)\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, predictions, targets):\n",
    "        \"\"\"Initialize with predictions and targets.\"\"\"\n",
    "        super().__init__(predictions)\n",
    "        self.targets_data = targets.data\n",
    "        self.num_samples = np.size(targets.data)\n",
    "    \n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"Compute gradient for BCE loss.\"\"\"\n",
    "        predictions, = self.saved_tensors\n",
    "        \n",
    "        if isinstance(predictions, Tensor) and predictions.requires_grad:\n",
    "            eps = 1e-7\n",
    "            p = np.clip(predictions.data, eps, 1 - eps)\n",
    "            y = self.targets_data\n",
    "            \n",
    "            # Gradient: (p - y) / (p * (1-p) * N)\n",
    "            grad = (p - y) / (p * (1 - p) * self.num_samples)\n",
    "            \n",
    "            return grad * grad_output,\n",
    "        return None,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62bdddaa",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "ce-backward",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class CrossEntropyBackward(Function):\n",
    "    \"\"\"\n",
    "    Gradient computation for Cross-Entropy Loss.\n",
    "    \n",
    "    CrossEntropy: L = -mean(log_softmax(logits)[targets])\n",
    "    \n",
    "    The gradient with respect to logits is remarkably elegant:\n",
    "    ∂L/∂logits = (softmax(logits) - one_hot(targets)) / N\n",
    "    \n",
    "    This is one of the most beautiful results in machine learning:\n",
    "    - The gradient is simply the difference between predictions and targets\n",
    "    - It naturally scales with how wrong we are\n",
    "    - It's numerically stable when computed via softmax\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, logits, targets):\n",
    "        \"\"\"Initialize with logits and target class indices.\"\"\"\n",
    "        super().__init__(logits)\n",
    "        self.targets_data = targets.data.astype(int)\n",
    "        self.batch_size = logits.data.shape[0]\n",
    "        self.num_classes = logits.data.shape[1]\n",
    "    \n",
    "    def apply(self, grad_output):\n",
    "        \"\"\"Compute gradient for cross-entropy loss.\"\"\"\n",
    "        logits, = self.saved_tensors\n",
    "        \n",
    "        if isinstance(logits, Tensor) and logits.requires_grad:\n",
    "            # Compute softmax probabilities\n",
    "            # Using stable softmax: subtract max for numerical stability\n",
    "            logits_data = logits.data\n",
    "            max_logits = np.max(logits_data, axis=1, keepdims=True)\n",
    "            exp_logits = np.exp(logits_data - max_logits)\n",
    "            softmax = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)\n",
    "            \n",
    "            # Create one-hot encoding of targets\n",
    "            one_hot = np.zeros((self.batch_size, self.num_classes), dtype=np.float32)\n",
    "            one_hot[np.arange(self.batch_size), self.targets_data] = 1.0\n",
    "            \n",
    "            # Gradient: (softmax - one_hot) / batch_size\n",
    "            grad = (softmax - one_hot) / self.batch_size\n",
    "            \n",
    "            return grad * grad_output,\n",
    "        return None,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56acda3f",
   "metadata": {
    "nbgrader": {
     "grade": false,
     "grade_id": "enable-autograd",
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "def enable_autograd():\n",
    "    \"\"\"\n",
    "    Enable gradient tracking for all Tensor operations.\n",
    "\n",
    "    This function enhances the existing Tensor class with autograd capabilities.\n",
    "    Call this once to activate gradients globally.\n",
    "\n",
    "    **What it does:**\n",
    "    - Replaces Tensor operations with gradient-tracking versions\n",
    "    - Adds backward() method for reverse-mode differentiation\n",
    "    - Enables computation graph building\n",
    "    - Maintains full backward compatibility\n",
    "\n",
    "    **After calling this:**\n",
    "    - Tensor operations will track computation graphs\n",
    "    - backward() method becomes available\n",
    "    - Gradients will flow through operations\n",
    "    - requires_grad=True enables tracking per tensor\n",
    "\n",
    "    **Example:**\n",
    "    ```python\n",
    "    enable_autograd()  # Call once\n",
    "    x = Tensor([2.0], requires_grad=True)\n",
    "    y = x * 3\n",
    "    y.backward()\n",
    "    print(x.grad)  # [3.0]\n",
    "    ```\n",
    "    \"\"\"\n",
    "\n",
    "    # Check if already enabled\n",
    "    if hasattr(Tensor, '_autograd_enabled'):\n",
    "        print(\"⚠️ Autograd already enabled\")\n",
    "        return\n",
    "\n",
    "    # Store original operations\n",
    "    _original_add = Tensor.__add__\n",
    "    _original_mul = Tensor.__mul__\n",
    "    _original_matmul = Tensor.matmul if hasattr(Tensor, 'matmul') else None\n",
    "\n",
    "    # Enhanced operations that track gradients\n",
    "    def tracked_add(self, other):\n",
    "        \"\"\"\n",
    "        Addition with gradient tracking.\n",
    "        \n",
    "        Enhances the original __add__ method to build computation graphs\n",
    "        when requires_grad=True for any input.\n",
    "        \"\"\"\n",
    "        # Convert scalar to Tensor if needed\n",
    "        if not isinstance(other, Tensor):\n",
    "            other = Tensor(other)\n",
    "\n",
    "        # Call original operation\n",
    "        result = _original_add(self, other)\n",
    "\n",
    "        # Track gradient if needed\n",
    "        if self.requires_grad or other.requires_grad:\n",
    "            result.requires_grad = True\n",
    "            result._grad_fn = AddBackward(self, other)\n",
    "\n",
    "        return result\n",
    "\n",
    "    def tracked_mul(self, other):\n",
    "        \"\"\"\n",
    "        Multiplication with gradient tracking.\n",
    "        \n",
    "        Enhances the original __mul__ method to build computation graphs\n",
    "        when requires_grad=True for any input.\n",
    "        \"\"\"\n",
    "        # Convert scalar to Tensor if needed for consistency\n",
    "        if not isinstance(other, Tensor):\n",
    "            other_tensor = Tensor(other)\n",
    "        else:\n",
    "            other_tensor = other\n",
    "\n",
    "        # Call original operation\n",
    "        result = _original_mul(self, other)\n",
    "\n",
    "        # Track gradient if needed\n",
    "        if self.requires_grad or (isinstance(other, Tensor) and other.requires_grad):\n",
    "            result.requires_grad = True\n",
    "            result._grad_fn = MulBackward(self, other)\n",
    "\n",
    "        return result\n",
    "\n",
    "    def tracked_matmul(self, other):\n",
    "        \"\"\"\n",
    "        Matrix multiplication with gradient tracking.\n",
    "        \n",
    "        Enhances the original matmul method to build computation graphs\n",
    "        when requires_grad=True for any input.\n",
    "        \"\"\"\n",
    "        if _original_matmul:\n",
    "            result = _original_matmul(self, other)\n",
    "        else:\n",
    "            # Fallback if matmul doesn't exist\n",
    "            result = Tensor(np.dot(self.data, other.data))\n",
    "\n",
    "        # Track gradient if needed\n",
    "        if self.requires_grad or other.requires_grad:\n",
    "            result.requires_grad = True\n",
    "            result._grad_fn = MatmulBackward(self, other)\n",
    "\n",
    "        return result\n",
    "\n",
    "    def sum_op(self, axis=None, keepdims=False):\n",
    "        \"\"\"\n",
    "        Sum operation with gradient tracking.\n",
    "        \n",
    "        Creates a new sum method that builds computation graphs\n",
    "        when requires_grad=True.\n",
    "        \"\"\"\n",
    "        result_data = np.sum(self.data, axis=axis, keepdims=keepdims)\n",
    "        result = Tensor(result_data)\n",
    "\n",
    "        if self.requires_grad:\n",
    "            result.requires_grad = True\n",
    "            result._grad_fn = SumBackward(self)\n",
    "\n",
    "        return result\n",
    "\n",
    "    def backward(self, gradient=None):\n",
    "        \"\"\"\n",
    "        Compute gradients via backpropagation.\n",
    "\n",
    "        This is the key method that makes training possible!\n",
    "        It implements reverse-mode automatic differentiation.\n",
    "        \n",
    "        **Algorithm:**\n",
    "        1. Initialize gradient if not provided (for scalar outputs)\n",
    "        2. Accumulate gradient in self.grad\n",
    "        3. If this tensor has a _grad_fn, call it to propagate gradients\n",
    "        4. Recursively call backward() on parent tensors\n",
    "        \n",
    "        **Example:**\n",
    "        ```python\n",
    "        x = Tensor([2.0], requires_grad=True)\n",
    "        y = x * 3\n",
    "        y.backward()  # Computes gradients for x\n",
    "        print(x.grad)  # [3.0]\n",
    "        ```\n",
    "        \"\"\"\n",
    "        # Only compute gradients if required\n",
    "        if not self.requires_grad:\n",
    "            return\n",
    "\n",
    "        # Initialize gradient if not provided (for scalar outputs)\n",
    "        if gradient is None:\n",
    "            if self.data.size == 1:\n",
    "                gradient = np.ones_like(self.data)\n",
    "            else:\n",
    "                raise ValueError(\"backward() requires gradient for non-scalar outputs\")\n",
    "\n",
    "        # Initialize or accumulate gradient\n",
    "        if self.grad is None:\n",
    "            self.grad = np.zeros_like(self.data)\n",
    "        \n",
    "        # Handle broadcasting: sum gradient to match self.data shape\n",
    "        # This happens when operations broadcast tensors (e.g., adding bias to batch)\n",
    "        if gradient.shape != self.grad.shape:\n",
    "            # Step 1: Remove extra leading dimensions added during forward pass\n",
    "            # Example: gradient (batch_size, features) → self.grad (features,)\n",
    "            while gradient.ndim > self.grad.ndim:\n",
    "                gradient = gradient.sum(axis=0)\n",
    "            \n",
    "            # Step 2: Sum over dimensions that were size-1 in original tensor\n",
    "            # Example: bias with shape (1,) broadcast to (batch_size,) during forward\n",
    "            for i in range(gradient.ndim):\n",
    "                if self.grad.shape[i] == 1 and gradient.shape[i] != 1:\n",
    "                    gradient = gradient.sum(axis=i, keepdims=True)\n",
    "        \n",
    "        self.grad += gradient\n",
    "\n",
    "        # Propagate gradients through computation graph\n",
    "        if hasattr(self, '_grad_fn') and self._grad_fn:\n",
    "            grads = self._grad_fn.apply(gradient)\n",
    "\n",
    "            # Recursively call backward on parent tensors\n",
    "            for tensor, grad in zip(self._grad_fn.saved_tensors, grads):\n",
    "                if isinstance(tensor, Tensor) and tensor.requires_grad and grad is not None:\n",
    "                    tensor.backward(grad)\n",
    "\n",
    "    def zero_grad(self):\n",
    "        \"\"\"\n",
    "        Reset gradients to zero.\n",
    "        \n",
    "        Call this before each backward pass to prevent gradient accumulation\n",
    "        from previous iterations.\n",
    "        \"\"\"\n",
    "        self.grad = None\n",
    "\n",
    "    # Install enhanced operations\n",
    "    Tensor.__add__ = tracked_add\n",
    "    Tensor.__mul__ = tracked_mul\n",
    "    Tensor.matmul = tracked_matmul\n",
    "    Tensor.sum = sum_op\n",
    "    Tensor.backward = backward\n",
    "    Tensor.zero_grad = zero_grad\n",
    "\n",
    "    # Patch activations and losses to track gradients\n",
    "    try:\n",
    "        from tinytorch.core.activations import Sigmoid, ReLU\n",
    "        from tinytorch.core.losses import BinaryCrossEntropyLoss, MSELoss, CrossEntropyLoss\n",
    "        \n",
    "        # Store original methods\n",
    "        _original_sigmoid_forward = Sigmoid.forward\n",
    "        _original_relu_forward = ReLU.forward\n",
    "        _original_bce_forward = BinaryCrossEntropyLoss.forward\n",
    "        _original_mse_forward = MSELoss.forward\n",
    "        _original_ce_forward = CrossEntropyLoss.forward\n",
    "        \n",
    "        def tracked_sigmoid_forward(self, x):\n",
    "            \"\"\"Sigmoid with gradient tracking.\"\"\"\n",
    "            result_data = 1.0 / (1.0 + np.exp(-x.data))\n",
    "            result = Tensor(result_data)\n",
    "            \n",
    "            if x.requires_grad:\n",
    "                result.requires_grad = True\n",
    "                result._grad_fn = SigmoidBackward(x, result)\n",
    "            \n",
    "            return result\n",
    "        \n",
    "        def tracked_relu_forward(self, x):\n",
    "            \"\"\"ReLU with gradient tracking.\"\"\"\n",
    "            result_data = np.maximum(0, x.data)\n",
    "            result = Tensor(result_data)\n",
    "            \n",
    "            if x.requires_grad:\n",
    "                result.requires_grad = True\n",
    "                result._grad_fn = ReLUBackward(x)\n",
    "            \n",
    "            return result\n",
    "        \n",
    "        def tracked_bce_forward(self, predictions, targets):\n",
    "            \"\"\"Binary cross-entropy with gradient tracking.\"\"\"\n",
    "            # Compute BCE loss\n",
    "            eps = 1e-7\n",
    "            clamped_preds = np.clip(predictions.data, eps, 1 - eps)\n",
    "            log_preds = np.log(clamped_preds)\n",
    "            log_one_minus_preds = np.log(1 - clamped_preds)\n",
    "            bce_per_sample = -(targets.data * log_preds + (1 - targets.data) * log_one_minus_preds)\n",
    "            bce_loss = np.mean(bce_per_sample)\n",
    "            \n",
    "            result = Tensor(bce_loss)\n",
    "            \n",
    "            if predictions.requires_grad:\n",
    "                result.requires_grad = True\n",
    "                result._grad_fn = BCEBackward(predictions, targets)\n",
    "            \n",
    "            return result\n",
    "        \n",
    "        def tracked_mse_forward(self, predictions, targets):\n",
    "            \"\"\"MSE loss with gradient tracking.\"\"\"\n",
    "            # Compute MSE loss\n",
    "            diff = predictions.data - targets.data\n",
    "            squared_diff = diff ** 2\n",
    "            mse = np.mean(squared_diff)\n",
    "            \n",
    "            result = Tensor(mse)\n",
    "            \n",
    "            if predictions.requires_grad:\n",
    "                result.requires_grad = True\n",
    "                result._grad_fn = MSEBackward(predictions, targets)\n",
    "            \n",
    "            return result\n",
    "        \n",
    "        def tracked_ce_forward(self, logits, targets):\n",
    "            \"\"\"Cross-entropy loss with gradient tracking.\"\"\"\n",
    "            from tinytorch.core.losses import log_softmax\n",
    "            \n",
    "            # Compute log-softmax for numerical stability\n",
    "            log_probs = log_softmax(logits, dim=-1)\n",
    "            \n",
    "            # Select log-probabilities for correct classes\n",
    "            batch_size = logits.shape[0]\n",
    "            target_indices = targets.data.astype(int)\n",
    "            selected_log_probs = log_probs.data[np.arange(batch_size), target_indices]\n",
    "            \n",
    "            # Return negative mean\n",
    "            ce_loss = -np.mean(selected_log_probs)\n",
    "            \n",
    "            result = Tensor(ce_loss)\n",
    "            \n",
    "            if logits.requires_grad:\n",
    "                result.requires_grad = True\n",
    "                result._grad_fn = CrossEntropyBackward(logits, targets)\n",
    "            \n",
    "            return result\n",
    "        \n",
    "        # Install patched methods\n",
    "        Sigmoid.forward = tracked_sigmoid_forward\n",
    "        ReLU.forward = tracked_relu_forward\n",
    "        BinaryCrossEntropyLoss.forward = tracked_bce_forward\n",
    "        MSELoss.forward = tracked_mse_forward\n",
    "        CrossEntropyLoss.forward = tracked_ce_forward\n",
    "        \n",
    "    except ImportError:\n",
    "        # Activations/losses not yet available (happens during module development)\n",
    "        pass\n",
    "\n",
    "    # Mark as enabled\n",
    "    Tensor._autograd_enabled = True\n",
    "\n",
    "    print(\"✅ Autograd enabled! Tensors now track gradients.\")\n",
    "    print(\"   - Operations build computation graphs\")\n",
    "    print(\"   - backward() computes gradients\")\n",
    "    print(\"   - requires_grad=True enables tracking\")\n",
    "\n",
    "# Auto-enable when module is imported\n",
    "enable_autograd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9ff4aea",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "### 🔬 Unit Test: Tensor Autograd Enhancement\n",
    "This test validates our enhanced Tensor class computes gradients correctly.\n",
    "**What we're testing**: Gradient computation and chain rule implementation\n",
    "**Why it matters**: This is the core of automatic differentiation\n",
    "**Expected**: Correct gradients for various operations and computation graphs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4222797",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-tensor-autograd",
     "locked": true,
     "points": 20
    }
   },
   "outputs": [],
   "source": [
    "def test_unit_tensor_autograd():\n",
    "    \"\"\"🔬 Test Tensor autograd enhancement.\"\"\"\n",
    "    print(\"🔬 Unit Test: Tensor Autograd Enhancement...\")\n",
    "\n",
    "    # Test simple gradient computation\n",
    "    x = Tensor([2.0], requires_grad=True)\n",
    "    y = x * 3\n",
    "    z = y + 1  # z = 3x + 1, so dz/dx = 3\n",
    "\n",
    "    z.backward()\n",
    "    assert np.allclose(x.grad, [3.0]), f\"Expected [3.0], got {x.grad}\"\n",
    "\n",
    "    # Test matrix multiplication gradients\n",
    "    a = Tensor([[1.0, 2.0]], requires_grad=True)  # 1x2\n",
    "    b = Tensor([[3.0], [4.0]], requires_grad=True)  # 2x1\n",
    "    c = a.matmul(b)  # 1x1, result = [[11.0]]\n",
    "\n",
    "    c.backward()\n",
    "    assert np.allclose(a.grad, [[3.0, 4.0]]), f\"Expected [[3.0, 4.0]], got {a.grad}\"\n",
    "    assert np.allclose(b.grad, [[1.0], [2.0]]), f\"Expected [[1.0], [2.0]], got {b.grad}\"\n",
    "\n",
    "    # Test computation graph with multiple operations\n",
    "    x = Tensor([1.0, 2.0], requires_grad=True)\n",
    "    y = x * 2      # y = [2, 4]\n",
    "    z = y.sum()    # z = 6\n",
    "\n",
    "    z.backward()\n",
    "    assert np.allclose(x.grad, [2.0, 2.0]), f\"Expected [2.0, 2.0], got {x.grad}\"\n",
    "\n",
    "    print(\"✅ Tensor autograd enhancement works correctly!\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    test_unit_tensor_autograd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96acf9fa",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## 🧪 Module Integration Test\n",
    "\n",
    "Final validation that everything works together correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec61fc12",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": true,
     "grade_id": "module-integration",
     "locked": true,
     "points": 25
    }
   },
   "outputs": [],
   "source": [
    "def test_module():\n",
    "    \"\"\"\n",
    "    Comprehensive test of entire module functionality.\n",
    "\n",
    "    This final test runs before module summary to ensure:\n",
    "    - All unit tests pass\n",
    "    - Autograd works for complex computation graphs\n",
    "    - Module is ready for integration with TinyTorch\n",
    "    \"\"\"\n",
    "    print(\"🧪 RUNNING MODULE INTEGRATION TEST\")\n",
    "    print(\"=\" * 50)\n",
    "\n",
    "    # Run all unit tests\n",
    "    print(\"Running unit tests...\")\n",
    "    test_unit_function_classes()\n",
    "    test_unit_tensor_autograd()\n",
    "\n",
    "    print(\"\\nRunning integration scenarios...\")\n",
    "\n",
    "    # Test 1: Multi-layer computation graph\n",
    "    print(\"🔬 Integration Test: Multi-layer Neural Network...\")\n",
    "\n",
    "    # Create a 3-layer computation: x -> Linear -> Linear -> Linear -> loss\n",
    "    x = Tensor([[1.0, 2.0]], requires_grad=True)\n",
    "    W1 = Tensor([[0.5, 0.3, 0.1], [0.2, 0.4, 0.6]], requires_grad=True)\n",
    "    b1 = Tensor([[0.1, 0.2, 0.3]], requires_grad=True)\n",
    "\n",
    "    # First layer\n",
    "    h1 = x.matmul(W1) + b1\n",
    "    assert h1.shape == (1, 3)\n",
    "    assert h1.requires_grad == True\n",
    "\n",
    "    # Second layer\n",
    "    W2 = Tensor([[0.1], [0.2], [0.3]], requires_grad=True)\n",
    "    h2 = h1.matmul(W2)\n",
    "    assert h2.shape == (1, 1)\n",
    "\n",
    "    # Compute simple loss (just square the output for testing)\n",
    "    loss = h2 * h2\n",
    "\n",
    "    # Backward pass\n",
    "    loss.backward()\n",
    "\n",
    "    # Verify all parameters have gradients\n",
    "    assert x.grad is not None\n",
    "    assert W1.grad is not None\n",
    "    assert b1.grad is not None\n",
    "    assert W2.grad is not None\n",
    "    assert x.grad.shape == x.shape\n",
    "    assert W1.grad.shape == W1.shape\n",
    "\n",
    "    print(\"✅ Multi-layer neural network gradients work!\")\n",
    "\n",
    "    # Test 2: Gradient accumulation\n",
    "    print(\"🔬 Integration Test: Gradient Accumulation...\")\n",
    "\n",
    "    x = Tensor([2.0], requires_grad=True)\n",
    "\n",
    "    # First computation\n",
    "    y1 = x * 3\n",
    "    y1.backward()\n",
    "    first_grad = x.grad.copy()\n",
    "\n",
    "    # Second computation (should accumulate)\n",
    "    y2 = x * 5\n",
    "    y2.backward()\n",
    "\n",
    "    assert np.allclose(x.grad, first_grad + 5.0), \"Gradients should accumulate\"\n",
    "    print(\"✅ Gradient accumulation works!\")\n",
    "\n",
    "    # Test 3: Complex mathematical operations\n",
    "    print(\"🔬 Integration Test: Complex Operations...\")\n",
    "\n",
    "    a = Tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)\n",
    "    b = Tensor([[2.0, 1.0], [1.0, 2.0]], requires_grad=True)\n",
    "\n",
    "    # Complex computation: ((a @ b) + a) * b\n",
    "    temp1 = a.matmul(b)  # Matrix multiplication\n",
    "    temp2 = temp1 + a    # Addition\n",
    "    result = temp2 * b   # Element-wise multiplication\n",
    "    final = result.sum() # Sum reduction\n",
    "\n",
    "    final.backward()\n",
    "\n",
    "    assert a.grad is not None\n",
    "    assert b.grad is not None\n",
    "    assert a.grad.shape == a.shape\n",
    "    assert b.grad.shape == b.shape\n",
    "\n",
    "    print(\"✅ Complex mathematical operations work!\")\n",
    "\n",
    "    print(\"\\n\" + \"=\" * 50)\n",
    "    print(\"🎉 ALL TESTS PASSED! Module ready for export.\")\n",
    "    print(\"Run: tito module complete 05_autograd\")\n",
    "\n",
    "# Test function defined above, will be called in main block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aff36fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run comprehensive module test\n",
    "if __name__ == \"__main__\":\n",
    "    test_module()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5db854b",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 MODULE SUMMARY: Autograd Engine\n",
    "\n",
    "Congratulations! You've built the gradient engine that makes neural networks learn!\n",
    "\n",
    "### Key Accomplishments ⭐⭐\n",
    "- **Enhanced Tensor class** with backward() method (no new wrapper classes!)\n",
    "- **Built computation graph tracking** for automatic differentiation\n",
    "- **Implemented Function classes** (Add, Mul, Matmul, Sum) with correct gradients\n",
    "- **Created enable_autograd()** function that activates gradients globally\n",
    "- **Tested complex multi-layer** computation graphs with gradient propagation\n",
    "- **All tests pass** ✅ (validated by `test_module()`)\n",
    "\n",
    "### Ready for Next Steps 🚀\n",
    "Your autograd implementation enables optimization! The dormant gradient features from Module 01 are now fully active. Every tensor can track gradients, every operation builds computation graphs, and backward() computes gradients automatically.\n",
    "\n",
    "**What you can do now:**\n",
    "```python\n",
    "# Create tensors with gradient tracking\n",
    "x = Tensor([2.0], requires_grad=True)\n",
    "W = Tensor([[0.5, 0.3]], requires_grad=True)\n",
    "\n",
    "# Build computation graphs automatically\n",
    "y = x.matmul(W.T)  # Forward pass\n",
    "loss = (y - 1.0) ** 2  # Simple loss\n",
    "\n",
    "# Compute gradients automatically\n",
    "loss.backward()  # Magic happens here!\n",
    "\n",
    "# Access gradients\n",
    "print(f\"x.grad: {x.grad}\")  # Gradient w.r.t. x\n",
    "print(f\"W.grad: {W.grad}\")  # Gradient w.r.t. W\n",
    "```\n",
    "\n",
    "Export with: `tito module complete 05_autograd`\n",
    "\n",
    "**Next**: Module 06 will add optimizers (SGD, Adam) that use these gradients to actually train neural networks! 🎯\n",
    "\n",
    "### 📈 Progress: Autograd ✓\n",
    "```\n",
    "✅ Module 01: Tensor (Foundation)\n",
    "✅ Module 02: Activations (Non-linearities) \n",
    "✅ Module 03: Layers (Building blocks)\n",
    "✅ Module 04: Losses (Training objectives)\n",
    "✅ Module 05: Autograd (Gradient engine) ← YOU ARE HERE\n",
    "🔄 Module 06: Optimizers (Learning algorithms)\n",
    "🔄 Module 07: Training (Complete training loops)\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
