{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "720f94f1",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "# Module 2: Activations - Nonlinearity in Neural Networks\n",
    "\n",
    "Welcome to the Activations module! This is where neural networks get their power through nonlinearity.\n",
    "\n",
    "## Learning Goals\n",
    "- Understand why activation functions are essential for neural networks\n",
    "- Implement the four most important activation functions: ReLU, Sigmoid, Tanh, and Softmax\n",
    "- Visualize how activations transform data and enable complex learning\n",
    "- See how activations work with layers to build powerful networks\n",
    "- Master the NBGrader workflow with comprehensive testing\n",
    "\n",
    "## Build → Use → Understand\n",
    "1. **Build**: Activation functions that add nonlinearity\n",
    "2. **Use**: Transform tensors and see immediate results\n",
    "3. **Understand**: How nonlinearity enables complex pattern learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c0ecb71",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "activations-imports",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| default_exp core.activations\n",
    "\n",
    "#| export\n",
    "import math\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import sys\n",
    "from typing import Union, List\n",
    "\n",
    "# Import our Tensor class - try from package first, then from local module\n",
    "try:\n",
    "    from tinytorch.core.tensor import Tensor\n",
    "except ImportError:\n",
    "    # For development, import from local tensor module\n",
    "    sys.path.append(os.path.join(os.path.dirname(__file__), '..', '01_tensor'))\n",
    "    from tensor_dev import Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd3c4277",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "activations-setup",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| export\n",
    "def _should_show_plots():\n",
    "    \"\"\"Check if we should show plots (disable during testing)\"\"\"\n",
    "    # Check multiple conditions that indicate we're in test mode\n",
    "    is_pytest = (\n",
    "        'pytest' in sys.modules or\n",
    "        'test' in sys.argv or\n",
    "        os.environ.get('PYTEST_CURRENT_TEST') is not None or\n",
    "        any('test' in arg for arg in sys.argv) or\n",
    "        any('pytest' in arg for arg in sys.argv)\n",
    "    )\n",
    "    \n",
    "    # Show plots in development mode (when not in test mode)\n",
    "    return not is_pytest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d08aa85",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "activations-visualization",
     "locked": false,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| export\n",
    "def visualize_activation_function(activation_fn, name: str, x_range: tuple = (-5, 5), num_points: int = 100):\n",
    "    \"\"\"Visualize an activation function's behavior\"\"\"\n",
    "    if not _should_show_plots():\n",
    "        return\n",
    "        \n",
    "    try:\n",
    "        \n",
    "        # Generate input values\n",
    "        x_vals = np.linspace(x_range[0], x_range[1], num_points)\n",
    "        \n",
    "        # Apply activation function\n",
    "        y_vals = []\n",
    "        for x in x_vals:\n",
    "            input_tensor = Tensor([[x]])\n",
    "            output = activation_fn(input_tensor)\n",
    "            y_vals.append(output.data.item())\n",
    "        \n",
    "        # Create plot\n",
    "        plt.figure(figsize=(10, 6))\n",
    "        plt.plot(x_vals, y_vals, 'b-', linewidth=2, label=f'{name} Activation')\n",
    "        plt.grid(True, alpha=0.3)\n",
    "        plt.xlabel('Input (x)')\n",
    "        plt.ylabel(f'{name}(x)')\n",
    "        plt.title(f'{name} Activation Function')\n",
    "        plt.legend()\n",
    "        plt.show()\n",
    "        \n",
    "    except ImportError:\n",
    "        print(\"   📊 Matplotlib not available - skipping visualization\")\n",
    "    except Exception as e:\n",
    "        print(f\"   ⚠️  Visualization error: {e}\")\n",
    "\n",
    "def visualize_activation_on_data(activation_fn, name: str, data: Tensor):\n",
    "    \"\"\"Show activation function applied to sample data\"\"\"\n",
    "    if not _should_show_plots():\n",
    "        return\n",
    "        \n",
    "    try:\n",
    "        output = activation_fn(data)\n",
    "        print(f\"   📊 {name} Example:\")\n",
    "        print(f\"      Input:  {data.data.flatten()}\")\n",
    "        print(f\"      Output: {output.data.flatten()}\")\n",
    "        print(f\"      Range:  [{output.data.min():.3f}, {output.data.max():.3f}]\")\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"   ⚠️  Data visualization error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a29b0c94",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## Step 1: What is an Activation Function?\n",
    "\n",
    "### Definition\n",
    "An **activation function** is a mathematical function that adds nonlinearity to neural networks. It transforms the output of a layer before passing it to the next layer.\n",
    "\n",
    "### Why Activation Functions Matter\n",
    "**Without activation functions, neural networks are just linear transformations!**\n",
    "\n",
    "```\n",
    "Linear → Linear → Linear = Still Linear\n",
    "```\n",
    "\n",
    "No matter how many layers you stack, without activation functions, you can only learn linear relationships. Activation functions introduce the nonlinearity that allows neural networks to:\n",
    "- Learn complex patterns\n",
    "- Approximate any continuous function\n",
    "- Solve non-linear problems\n",
    "\n",
    "### Visual Analogy\n",
    "Think of activation functions as **decision makers** at each neuron:\n",
    "- **ReLU**: \"If positive, pass it through; if negative, block it\"\n",
    "- **Sigmoid**: \"Squash everything between 0 and 1\"\n",
    "- **Tanh**: \"Squash everything between -1 and 1\"\n",
    "- **Softmax**: \"Convert to probabilities that sum to 1\"\n",
    "\n",
    "### Connection to Previous Modules\n",
    "In Module 1 (Tensor), we learned how to store and manipulate data. Now we add the nonlinear functions that make neural networks powerful."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b3cce52",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 2: ReLU - The Workhorse of Deep Learning\n",
    "\n",
    "### What is ReLU?\n",
    "**ReLU (Rectified Linear Unit)** is the most popular activation function in deep learning.\n",
    "\n",
    "**Mathematical Definition:**\n",
    "```\n",
    "f(x) = max(0, x)\n",
    "```\n",
    "\n",
    "**In Plain English:**\n",
    "- If input is positive → pass it through unchanged\n",
    "- If input is negative → output zero\n",
    "\n",
    "### Why ReLU is Popular\n",
    "1. **Simple**: Easy to compute and understand\n",
    "2. **Fast**: No expensive operations (no exponentials)\n",
    "3. **Sparse**: Outputs many zeros, creating sparse representations\n",
    "4. **Gradient-friendly**: Gradient is either 0 or 1 (no vanishing gradient for positive inputs)\n",
    "\n",
    "### Real-World Analogy\n",
    "ReLU is like a **one-way valve** - it only lets positive \"pressure\" through, blocking negative values completely.\n",
    "\n",
    "### When to Use ReLU\n",
    "- **Hidden layers** in most neural networks\n",
    "- **Convolutional layers** in image processing\n",
    "- **When you want sparse activations**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4300f9b3",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "relu-class",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class ReLU:\n",
    "    \"\"\"\n",
    "    ReLU Activation Function: f(x) = max(0, x)\n",
    "    \n",
    "    The most popular activation function in deep learning.\n",
    "    Simple, fast, and effective for most applications.\n",
    "    \"\"\"\n",
    "    \n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Apply ReLU activation: f(x) = max(0, x)\n",
    "        \n",
    "        TODO: Implement ReLU activation\n",
    "        \n",
    "        APPROACH:\n",
    "        1. For each element in the input tensor, apply max(0, element)\n",
    "        2. Return a new Tensor with the results\n",
    "        \n",
    "        EXAMPLE:\n",
    "        Input: Tensor([[-1, 0, 1, 2, -3]])\n",
    "        Expected: Tensor([[0, 0, 1, 2, 0]])\n",
    "        \n",
    "        HINTS:\n",
    "        - Use np.maximum(0, x.data) for element-wise max\n",
    "        - Remember to return a new Tensor object\n",
    "        - The shape should remain the same as input\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        result = np.maximum(0, x.data)\n",
    "        return Tensor(result)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"Make the class callable: relu(x) instead of relu.forward(x)\"\"\"\n",
    "        return self.forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533c471b",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 3: Sigmoid - The Smooth Squasher\n",
    "\n",
    "### What is Sigmoid?\n",
    "**Sigmoid** is a smooth S-shaped function that squashes inputs to the range (0, 1).\n",
    "\n",
    "**Mathematical Definition:**\n",
    "```\n",
    "f(x) = 1 / (1 + e^(-x))\n",
    "```\n",
    "\n",
    "**Properties:**\n",
    "- **Range**: (0, 1) - never exactly 0 or 1\n",
    "- **Smooth**: Differentiable everywhere\n",
    "- **Monotonic**: Always increasing\n",
    "- **Centered**: Around 0.5\n",
    "\n",
    "### Why Sigmoid is Useful\n",
    "1. **Probabilistic**: Output can be interpreted as probabilities\n",
    "2. **Bounded**: Output is always between 0 and 1\n",
    "3. **Smooth**: Good for gradient-based optimization\n",
    "4. **Historical**: Was the standard before ReLU\n",
    "\n",
    "### Real-World Analogy\n",
    "Sigmoid is like a **soft switch** - it gradually turns on as input increases, unlike ReLU's hard cutoff.\n",
    "\n",
    "### When to Use Sigmoid\n",
    "- **Binary classification** (output layer)\n",
    "- **Gates** in LSTM/GRU networks\n",
    "- **When you need probabilistic outputs**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbe9f91c",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "sigmoid-class",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class Sigmoid:\n",
    "    \"\"\"\n",
    "    Sigmoid Activation Function: f(x) = 1 / (1 + e^(-x))\n",
    "    \n",
    "    Smooth S-shaped function that squashes inputs to (0, 1).\n",
    "    Useful for binary classification and probabilistic outputs.\n",
    "    \"\"\"\n",
    "    \n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Apply Sigmoid activation: f(x) = 1 / (1 + e^(-x))\n",
    "        \n",
    "        TODO: Implement Sigmoid activation with numerical stability\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Clip input values to prevent overflow (e.g., between -500 and 500)\n",
    "        2. Apply the sigmoid formula: 1 / (1 + exp(-x))\n",
    "        3. Return a new Tensor with the results\n",
    "        \n",
    "        EXAMPLE:\n",
    "        Input: Tensor([[-2, 0, 2]])\n",
    "        Expected: Tensor([[0.119, 0.5, 0.881]]) (approximately)\n",
    "        \n",
    "        HINTS:\n",
    "        - Use np.clip(x.data, -500, 500) for numerical stability\n",
    "        - Use np.exp() for the exponential function\n",
    "        - Be careful with very large/small inputs to avoid overflow\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Clip for numerical stability\n",
    "        clipped = np.clip(x.data, -500, 500)\n",
    "        result = 1 / (1 + np.exp(-clipped))\n",
    "        return Tensor(result)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"Make the class callable: sigmoid(x) instead of sigmoid.forward(x)\"\"\"\n",
    "        return self.forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67dc777f",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 4: Tanh - The Zero-Centered Squasher\n",
    "\n",
    "### What is Tanh?\n",
    "**Tanh (Hyperbolic Tangent)** is similar to Sigmoid but centered around zero.\n",
    "\n",
    "**Mathematical Definition:**\n",
    "```\n",
    "f(x) = tanh(x) = (e^x - e^(-x)) / (e^x + e^(-x))\n",
    "```\n",
    "\n",
    "**Properties:**\n",
    "- **Range**: (-1, 1) - symmetric around zero\n",
    "- **Zero-centered**: Output averages to zero\n",
    "- **Smooth**: Differentiable everywhere\n",
    "- **Stronger gradients**: Than sigmoid in some regions\n",
    "\n",
    "### Why Tanh is Useful\n",
    "1. **Zero-centered**: Better for training (gradients don't all have same sign)\n",
    "2. **Symmetric**: Treats positive and negative inputs equally\n",
    "3. **Stronger gradients**: Can help with training dynamics\n",
    "4. **Bounded**: Output is always between -1 and 1\n",
    "\n",
    "### Real-World Analogy\n",
    "Tanh is like a **balanced scale** - it can tip positive or negative, with zero as the neutral point.\n",
    "\n",
    "### When to Use Tanh\n",
    "- **Hidden layers** (alternative to ReLU)\n",
    "- **RNNs** (traditional choice)\n",
    "- **When you need zero-centered outputs**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e982bfbd",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "tanh-class",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class Tanh:\n",
    "    \"\"\"\n",
    "    Tanh Activation Function: f(x) = tanh(x)\n",
    "    \n",
    "    Zero-centered S-shaped function that squashes inputs to (-1, 1).\n",
    "    Better than sigmoid for hidden layers due to zero-centered outputs.\n",
    "    \"\"\"\n",
    "    \n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Apply Tanh activation: f(x) = tanh(x)\n",
    "        \n",
    "        TODO: Implement Tanh activation\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Use NumPy's tanh function for numerical stability\n",
    "        2. Apply to the tensor data\n",
    "        3. Return a new Tensor with the results\n",
    "        \n",
    "        EXAMPLE:\n",
    "        Input: Tensor([[-2, 0, 2]])\n",
    "        Expected: Tensor([[-0.964, 0.0, 0.964]]) (approximately)\n",
    "        \n",
    "        HINTS:\n",
    "        - Use np.tanh(x.data) - NumPy handles the math\n",
    "        - Much simpler than implementing the formula manually\n",
    "        - NumPy's tanh is numerically stable\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        result = np.tanh(x.data)\n",
    "        return Tensor(result)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"Make the class callable: tanh(x) instead of tanh.forward(x)\"\"\"\n",
    "        return self.forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "726ae88b",
   "metadata": {
    "cell_marker": "\"\"\"",
    "lines_to_next_cell": 1
   },
   "source": [
    "## Step 5: Softmax - The Probability Converter\n",
    "\n",
    "### What is Softmax?\n",
    "**Softmax** converts a vector of numbers into a probability distribution.\n",
    "\n",
    "**Mathematical Definition:**\n",
    "```\n",
    "f(x_i) = e^(x_i) / Σ(e^(x_j)) for all j\n",
    "```\n",
    "\n",
    "**Properties:**\n",
    "- **Probabilities**: All outputs sum to 1\n",
    "- **Non-negative**: All outputs are ≥ 0\n",
    "- **Differentiable**: Smooth everywhere\n",
    "- **Competitive**: Amplifies differences between inputs\n",
    "\n",
    "### Why Softmax is Essential\n",
    "1. **Multi-class classification**: Converts logits to probabilities\n",
    "2. **Attention mechanisms**: Focuses on important elements\n",
    "3. **Interpretable**: Output can be understood as confidence\n",
    "4. **Competitive**: Emphasizes the largest input\n",
    "\n",
    "### Real-World Analogy\n",
    "Softmax is like **dividing a pie** - it takes any set of numbers and converts them into slices that sum to 100%.\n",
    "\n",
    "### When to Use Softmax\n",
    "- **Multi-class classification** (output layer)\n",
    "- **Attention mechanisms** in transformers\n",
    "- **When you need probability distributions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a99d93cc",
   "metadata": {
    "lines_to_next_cell": 1,
    "nbgrader": {
     "grade": false,
     "grade_id": "softmax-class",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "#| export\n",
    "class Softmax:\n",
    "    \"\"\"\n",
    "    Softmax Activation Function: f(x_i) = e^(x_i) / Σ(e^(x_j))\n",
    "    \n",
    "    Converts a vector of numbers into a probability distribution.\n",
    "    Essential for multi-class classification and attention mechanisms.\n",
    "    \"\"\"\n",
    "    \n",
    "    def forward(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"\n",
    "        Apply Softmax activation: f(x_i) = e^(x_i) / Σ(e^(x_j))\n",
    "        \n",
    "        TODO: Implement Softmax activation with numerical stability\n",
    "        \n",
    "        APPROACH:\n",
    "        1. Subtract max value from inputs for numerical stability\n",
    "        2. Compute exponentials: e^(x_i - max)\n",
    "        3. Divide by sum of exponentials\n",
    "        4. Return a new Tensor with the results\n",
    "        \n",
    "        EXAMPLE:\n",
    "        Input: Tensor([[1, 2, 3]])\n",
    "        Expected: Tensor([[0.09, 0.24, 0.67]]) (approximately, sums to 1)\n",
    "        \n",
    "        HINTS:\n",
    "        - Use np.max(x.data, axis=-1, keepdims=True) for stability\n",
    "        - Use np.exp() for exponentials\n",
    "        - Use np.sum() for the denominator\n",
    "        - Make sure the result sums to 1 along the last axis\n",
    "        \"\"\"\n",
    "        ### BEGIN SOLUTION\n",
    "        # Subtract max for numerical stability\n",
    "        x_max = np.max(x.data, axis=-1, keepdims=True)\n",
    "        x_shifted = x.data - x_max\n",
    "        \n",
    "        # Compute softmax\n",
    "        exp_x = np.exp(x_shifted)\n",
    "        sum_exp = np.sum(exp_x, axis=-1, keepdims=True)\n",
    "        result = exp_x / sum_exp\n",
    "        \n",
    "        return Tensor(result)\n",
    "        ### END SOLUTION\n",
    "    \n",
    "    def __call__(self, x: Tensor) -> Tensor:\n",
    "        \"\"\"Make the class callable: softmax(x) instead of softmax.forward(x)\"\"\"\n",
    "        return self.forward(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d37cb352",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "### 🧪 Test Your Activation Functions\n",
    "\n",
    "Once you implement the activation functions above, run these cells to test them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "067e766c",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-relu",
     "locked": true,
     "points": 20,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test ReLU activation\n",
    "print(\"Testing ReLU activation...\")\n",
    "\n",
    "relu = ReLU()\n",
    "\n",
    "# Test basic functionality\n",
    "input_tensor = Tensor([[-2, -1, 0, 1, 2]])\n",
    "output = relu(input_tensor)\n",
    "expected = np.array([[0, 0, 0, 1, 2]])\n",
    "assert np.array_equal(output.data, expected), f\"ReLU failed: expected {expected}, got {output.data}\"\n",
    "\n",
    "# Test with matrix\n",
    "matrix_input = Tensor([[-1, 2], [3, -4]])\n",
    "matrix_output = relu(matrix_input)\n",
    "expected_matrix = np.array([[0, 2], [3, 0]])\n",
    "assert np.array_equal(matrix_output.data, expected_matrix), f\"ReLU matrix failed: expected {expected_matrix}, got {matrix_output.data}\"\n",
    "\n",
    "# Test shape preservation\n",
    "assert output.shape == input_tensor.shape, f\"ReLU should preserve shape: input {input_tensor.shape}, output {output.shape}\"\n",
    "\n",
    "print(\"✅ ReLU tests passed!\")\n",
    "print(f\"✅ ReLU({input_tensor.data.flatten()}) = {output.data.flatten()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e01b7261",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-sigmoid",
     "locked": true,
     "points": 20,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test Sigmoid activation\n",
    "print(\"Testing Sigmoid activation...\")\n",
    "\n",
    "sigmoid = Sigmoid()\n",
    "\n",
    "# Test basic functionality\n",
    "input_tensor = Tensor([[0]])\n",
    "output = sigmoid(input_tensor)\n",
    "expected_value = 0.5\n",
    "assert abs(output.data.item() - expected_value) < 1e-6, f\"Sigmoid(0) should be 0.5, got {output.data.item()}\"\n",
    "\n",
    "# Test range bounds (allowing for floating-point precision at extremes)\n",
    "large_input = Tensor([[100]])\n",
    "large_output = sigmoid(large_input)\n",
    "assert 0 < large_output.data.item() <= 1, f\"Sigmoid output should be in (0,1], got {large_output.data.item()}\"\n",
    "\n",
    "small_input = Tensor([[-100]])\n",
    "small_output = sigmoid(small_input)\n",
    "assert 0 <= small_output.data.item() < 1, f\"Sigmoid output should be in [0,1), got {small_output.data.item()}\"\n",
    "\n",
    "# Test with multiple values\n",
    "multi_input = Tensor([[-2, 0, 2]])\n",
    "multi_output = sigmoid(multi_input)\n",
    "assert multi_output.shape == multi_input.shape, \"Sigmoid should preserve shape\"\n",
    "assert np.all((multi_output.data > 0) & (multi_output.data < 1)), \"All sigmoid outputs should be in (0,1)\"\n",
    "\n",
    "print(\"✅ Sigmoid tests passed!\")\n",
    "print(f\"✅ Sigmoid({multi_input.data.flatten()}) = {multi_output.data.flatten()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ca2fa6f",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-tanh",
     "locked": true,
     "points": 20,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test Tanh activation\n",
    "print(\"Testing Tanh activation...\")\n",
    "\n",
    "tanh = Tanh()\n",
    "\n",
    "# Test basic functionality\n",
    "input_tensor = Tensor([[0]])\n",
    "output = tanh(input_tensor)\n",
    "expected_value = 0.0\n",
    "assert abs(output.data.item() - expected_value) < 1e-6, f\"Tanh(0) should be 0.0, got {output.data.item()}\"\n",
    "\n",
    "# Test range bounds (allowing for floating-point precision at extremes)\n",
    "large_input = Tensor([[100]])\n",
    "large_output = tanh(large_input)\n",
    "assert -1 <= large_output.data.item() <= 1, f\"Tanh output should be in [-1,1], got {large_output.data.item()}\"\n",
    "\n",
    "small_input = Tensor([[-100]])\n",
    "small_output = tanh(small_input)\n",
    "assert -1 <= small_output.data.item() <= 1, f\"Tanh output should be in [-1,1], got {small_output.data.item()}\"\n",
    "\n",
    "# Test symmetry: tanh(-x) = -tanh(x)\n",
    "test_input = Tensor([[2]])\n",
    "pos_output = tanh(test_input)\n",
    "neg_input = Tensor([[-2]])\n",
    "neg_output = tanh(neg_input)\n",
    "assert abs(pos_output.data.item() + neg_output.data.item()) < 1e-6, \"Tanh should be symmetric: tanh(-x) = -tanh(x)\"\n",
    "\n",
    "print(\"✅ Tanh tests passed!\")\n",
    "print(f\"✅ Tanh(±2) = ±{abs(pos_output.data.item()):.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50795506",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-softmax",
     "locked": true,
     "points": 20,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test Softmax activation\n",
    "print(\"Testing Softmax activation...\")\n",
    "\n",
    "softmax = Softmax()\n",
    "\n",
    "# Test basic functionality\n",
    "input_tensor = Tensor([[1, 2, 3]])\n",
    "output = softmax(input_tensor)\n",
    "\n",
    "# Check that outputs sum to 1\n",
    "sum_output = np.sum(output.data)\n",
    "assert abs(sum_output - 1.0) < 1e-6, f\"Softmax outputs should sum to 1, got {sum_output}\"\n",
    "\n",
    "# Check that all outputs are positive\n",
    "assert np.all(output.data > 0), \"All softmax outputs should be positive\"\n",
    "\n",
    "# Check that larger inputs give larger outputs\n",
    "assert output.data[0, 2] > output.data[0, 1] > output.data[0, 0], \"Softmax should preserve order\"\n",
    "\n",
    "# Test with matrix (multiple rows)\n",
    "matrix_input = Tensor([[1, 2], [3, 4]])\n",
    "matrix_output = softmax(matrix_input)\n",
    "row_sums = np.sum(matrix_output.data, axis=1)\n",
    "assert np.allclose(row_sums, 1.0), f\"Each row should sum to 1, got {row_sums}\"\n",
    "\n",
    "print(\"✅ Softmax tests passed!\")\n",
    "print(f\"✅ Softmax({input_tensor.data.flatten()}) = {output.data.flatten()}\")\n",
    "print(f\"✅ Sum = {np.sum(output.data):.6f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8dfc085",
   "metadata": {
    "nbgrader": {
     "grade": true,
     "grade_id": "test-activation-integration",
     "locked": true,
     "points": 20,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# Test activation function integration\n",
    "print(\"Testing activation function integration...\")\n",
    "\n",
    "# Create test data\n",
    "test_data = Tensor([[-2, -1, 0, 1, 2]])\n",
    "\n",
    "# Test all activations\n",
    "relu = ReLU()\n",
    "sigmoid = Sigmoid()\n",
    "tanh = Tanh()\n",
    "softmax = Softmax()\n",
    "\n",
    "# Apply all activations\n",
    "relu_out = relu(test_data)\n",
    "sigmoid_out = sigmoid(test_data)\n",
    "tanh_out = tanh(test_data)\n",
    "softmax_out = softmax(test_data)\n",
    "\n",
    "# Check shapes are preserved\n",
    "assert relu_out.shape == test_data.shape, \"ReLU should preserve shape\"\n",
    "assert sigmoid_out.shape == test_data.shape, \"Sigmoid should preserve shape\"\n",
    "assert tanh_out.shape == test_data.shape, \"Tanh should preserve shape\"\n",
    "assert softmax_out.shape == test_data.shape, \"Softmax should preserve shape\"\n",
    "\n",
    "# Check ranges (allowing for floating-point precision at extremes)\n",
    "assert np.all(relu_out.data >= 0), \"ReLU outputs should be non-negative\"\n",
    "assert np.all((sigmoid_out.data >= 0) & (sigmoid_out.data <= 1)), \"Sigmoid outputs should be in [0,1]\"\n",
    "assert np.all((tanh_out.data >= -1) & (tanh_out.data <= 1)), \"Tanh outputs should be in [-1,1]\"\n",
    "assert np.all(softmax_out.data > 0), \"Softmax outputs should be positive\"\n",
    "\n",
    "# Test chaining (composition)\n",
    "chained = relu(sigmoid(test_data))\n",
    "assert chained.shape == test_data.shape, \"Chained activations should preserve shape\"\n",
    "\n",
    "print(\"✅ Activation integration tests passed!\")\n",
    "print(f\"✅ All activation functions work correctly\")\n",
    "print(f\"✅ Input:   {test_data.data.flatten()}\")\n",
    "print(f\"✅ ReLU:    {relu_out.data.flatten()}\")\n",
    "print(f\"✅ Sigmoid: {sigmoid_out.data.flatten()}\")\n",
    "print(f\"✅ Tanh:    {tanh_out.data.flatten()}\")\n",
    "print(f\"✅ Softmax: {softmax_out.data.flatten()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa5f40bb",
   "metadata": {
    "cell_marker": "\"\"\""
   },
   "source": [
    "## 🎯 Module Summary\n",
    "\n",
    "Congratulations! You've successfully implemented the core activation functions for TinyTorch:\n",
    "\n",
    "### What You've Accomplished\n",
    "✅ **ReLU**: The workhorse activation for hidden layers  \n",
    "✅ **Sigmoid**: Smooth probabilistic outputs for binary classification  \n",
    "✅ **Tanh**: Zero-centered activation for better training dynamics  \n",
    "✅ **Softmax**: Probability distributions for multi-class classification  \n",
    "✅ **Integration**: All functions work together and preserve tensor shapes  \n",
    "\n",
    "### Key Concepts You've Learned\n",
    "- **Nonlinearity** is essential for neural networks to learn complex patterns\n",
    "- **ReLU** is simple, fast, and effective for most hidden layers\n",
    "- **Sigmoid** squashes outputs to (0,1) for probabilistic interpretation\n",
    "- **Tanh** is zero-centered and often better than sigmoid for hidden layers\n",
    "- **Softmax** converts logits to probability distributions\n",
    "- **Numerical stability** is crucial for functions with exponentials\n",
    "\n",
    "### Next Steps\n",
    "1. **Export your code**: `tito package nbdev --export 02_activations`\n",
    "2. **Test your implementation**: `tito module test 02_activations`\n",
    "3. **Use your activations**: \n",
    "   ```python\n",
    "   from tinytorch.core.activations import ReLU, Sigmoid, Tanh, Softmax\n",
    "   from tinytorch.core.tensor import Tensor\n",
    "   \n",
    "   relu = ReLU()\n",
    "   x = Tensor([[-1, 0, 1, 2]])\n",
    "   y = relu(x)  # Your activation in action!\n",
    "   ```\n",
    "4. **Move to Module 3**: Start building neural network layers!\n",
    "\n",
    "**Ready for the next challenge?** Let's combine tensors and activations to build the fundamental building blocks of neural networks!"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
