{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autodifferentiation and Gradient Descent in Burn\n",
    "\n",
    "This notebook demonstrates how to use automatic differentiation in Burn to compute gradients and implement gradient descent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [],
   "source": [
    "// Dependency declarations\n",
    ":dep burn = {path = \"../../crates/burn\"}\n",
    ":dep burn-flex = {path = \"../../crates/burn-flex\"}\n",
    ":dep burn-autodiff = {path = \"../../crates/burn-autodiff\"}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [],
   "source": [
    "// Import packages\n",
    "use burn::prelude::*;\n",
    "use burn_autodiff::Autodiff;\n",
    "use burn_flex::Flex;\n",
    "\n",
    "// Type alias: Autodiff<Flex> enables automatic differentiation\n",
    "type B = Autodiff<Flex>;\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Understanding require_grad()\n",
    "\n",
    "In Burn, tensors can be marked for gradient tracking using `.require_grad()`. This tells the framework to track operations on this tensor so gradients can be computed later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Regular tensor x: Tensor {\n",
      "  data:\n",
      "[1.0, 2.0, 3.0, 4.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Tensor y with require_grad: Tensor {\n",
      "  data:\n",
      "[1.0, 2.0, 3.0, 4.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "result = sum(y * 2) = Tensor {\n",
      "  data:\n",
      "[20.0],\n",
      "  shape:  [1],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "let device = Default::default();\n",
    "\n",
    "// Create a regular tensor - no gradient tracking\n",
    "let x: Tensor<B, 1> = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device);\n",
    "println!(\"Regular tensor x: {}\", x);\n",
    "\n",
    "// Create a tensor that requires gradient computation\n",
    "let y: Tensor<B, 1> = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device).require_grad();\n",
    "println!(\"Tensor y with require_grad: {}\", y);\n",
    "\n",
    "// Now let's do some operations on y\n",
    "let z = y.clone() * 2.0;\n",
    "let result = z.sum();\n",
    "println!(\"result = sum(y * 2) = {}\", result);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Computing Gradients with backward()\n",
    "\n",
    "The `.backward()` method computes the gradients of all tensors that have `require_grad()` set. It returns a gradients object that holds the computed gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y = Tensor {\n",
      "  data:\n",
      "[1.0, 2.0, 3.0, 4.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "dy/dx = Tensor {\n",
      "  data:\n",
      "[2.0, 2.0, 2.0, 2.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// Example: y = [1, 2, 3, 4]\n",
    "// z = y * 2 = [2, 4, 6, 8]\n",
    "// result = sum(z) = 20\n",
    "//\n",
    "// d(result)/d(y) = d(result)/dz * dz/dy = 1 * 2 = [2, 2, 2, 2]\n",
    "\n",
    "let device = Default::default();\n",
    "let y: Tensor<B, 1> = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device).require_grad();\n",
    "let z = y.clone() * 2.0;\n",
    "let result = z.sum();\n",
    "\n",
    "// Compute gradients\n",
    "let grads = result.backward();\n",
    "\n",
    "// Get gradient for y\n",
    "let y_grad = y.grad(&grads).unwrap();\n",
    "println!(\"y = {}\", y);\n",
    "println!(\"d(result)/dy = {}\", y_grad);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. More Complex Example: Quadratic Function\n",
    "Let's compute the gradient of a more complex function: f(x) = x²\n",
    "\n",
    "The derivative is: f'(x) = 2x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x = Tensor {\n",
      "  data:\n",
      "[1.0, 2.0, 3.0, 4.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "x^2 = Tensor {\n",
      "  data:\n",
      "[1.0, 4.0, 9.0, 16.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "d(x^2)/dx = Tensor {\n",
      "  data:\n",
      "[2.0, 4.0, 6.0, 8.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Expected: [2, 4, 6, 8]\n"
     ]
    }
   ],
   "source": [
    "// f(x) = x^2\n",
    "// f'(x) = 2x\n",
    "\n",
    "let device = Default::default();\n",
    "let x: Tensor<B, 1> = Tensor::from_floats([1.0, 2.0, 3.0, 4.0], &device).require_grad();\n",
    "let y = x.clone().powf_scalar(2.0);\n",
    "let result = y.clone().sum();\n",
    "\n",
    "let grads = result.backward();\n",
    "let x_grad = x.grad(&grads).unwrap();\n",
    "\n",
    "println!(\"x = {}\", x);\n",
    "println!(\"x^2 = {}\", y);\n",
    "println!(\"d(x^2)/dx = {}\", x_grad);\n",
    "\n",
    "// Verify: d(x^2)/dx should be [2, 4, 6, 8]\n",
    "println!(\"Expected: [2, 4, 6, 8]\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Chain Rule Example\n",
    "\n",
    "Let's verify the chain rule: f(g(x))' = f'(g(x)) * g'(x)\n",
    "\n",
    "Example: y = sin(x²), we want dy/dx\n",
    "\n",
    "Let u = x², y = sin(u)\n",
    "dy/du = cos(u), du/dx = 2x\n",
    "dy/dx = cos(x²) * 2x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x = Tensor {\n",
      "  data:\n",
      "[0.0, 1.0, 2.0, 3.0],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "y = sin(x^2) = Tensor {\n",
      "  data:\n",
      "[0.0, 0.84147096, -0.7568025, 0.4121185],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "dy/dx = Tensor {\n",
      "  data:\n",
      "[0.0, 1.0806046, -2.6145744, -5.4667816],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"flex\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n",
      "Expected (cos(x^2) * 2x): Tensor {\n",
      "  data:\n",
      "[0.0, 1.0806046, -2.6145744, -5.4667816],\n",
      "  shape:  [4],\n",
      "  device:  Cpu,\n",
      "  backend:  \"autodiff<flex>\",\n",
      "  kind:  \"Float\",\n",
      "  dtype:  \"f32\",\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "// y = sin(x^2)\n",
    "// dy/dx = cos(x^2) * 2x\n",
    "\n",
    "let device = Default::default();\n",
    "let x: Tensor<B, 1> = Tensor::from_floats([0.0, 1.0, 2.0, 3.0], &device).require_grad();\n",
    "\n",
    "// Forward pass\n",
    "let x_squared = x.clone().powf_scalar(2.0);\n",
    "let y = x_squared.sin();\n",
    "let result = y.clone().sum();\n",
    "\n",
    "// Backward pass\n",
    "let grads = result.backward();\n",
    "let x_grad = x.grad(&grads).unwrap();\n",
    "\n",
    "println!(\"x = {}\", x);\n",
    "println!(\"y = sin(x^2) = {}\", y);\n",
    "println!(\"dy/dx = {}\", x_grad);\n",
    "\n",
    "// Verify manually: cos(x^2) * 2x\n",
    "let expected_grad = x.clone().powf_scalar(2.0).cos() * (x.clone() * 2.0);\n",
    "println!(\"Expected (cos(x^2) * 2x): {}\", expected_grad);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Gradient Descent from Scratch\n",
    "\n",
    "Now let's implement the classic gradient descent algorithm to find the minimum of a function.\n",
    "\n",
    "We'll minimize: f(x) = (x - 3)²\n",
    "\n",
    "The minimum is at x = 3, where f(x) = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting gradient descent to minimize (x - 3)^2\n",
      "Expected minimum: x = 3\n",
      "---\n",
      "Iteration 0: x = 0.0000, loss = 9.0000\n",
      "Iteration 1: x = 0.6000, loss = 5.7600\n",
      "Iteration 2: x = 1.0800, loss = 3.6864\n",
      "Iteration 3: x = 1.4640, loss = 2.3593\n",
      "Iteration 4: x = 1.7712, loss = 1.5099\n",
      "Iteration 5: x = 2.0170, loss = 0.9664\n",
      "Iteration 6: x = 2.2136, loss = 0.6185\n",
      "Iteration 7: x = 2.3709, loss = 0.3958\n",
      "Iteration 8: x = 2.4967, loss = 0.2533\n",
      "Iteration 9: x = 2.5973, loss = 0.1621\n",
      "Iteration 10: x = 2.6779, loss = 0.1038\n",
      "Iteration 11: x = 2.7423, loss = 0.0664\n",
      "Iteration 12: x = 2.7938, loss = 0.0425\n",
      "Iteration 13: x = 2.8351, loss = 0.0272\n",
      "Iteration 14: x = 2.8681, loss = 0.0174\n",
      "Iteration 15: x = 2.8944, loss = 0.0111\n",
      "Iteration 16: x = 2.9156, loss = 0.0071\n",
      "Iteration 17: x = 2.9324, loss = 0.0046\n",
      "Iteration 18: x = 2.9460, loss = 0.0029\n",
      "Iteration 19: x = 2.9568, loss = 0.0019\n",
      "---\n",
      "Final x = 2.9654\n"
     ]
    }
   ],
   "source": [
    "// Target: minimize f(x) = (x - 3)^2\n",
    "// This has minimum at x = 3\n",
    "\n",
    "fn loss<B: Backend>(x: &Tensor<B, 1>) -> Tensor<B, 1> {\n",
    "    // f(x) = (x - 3)^2\n",
    "    (x.clone() - 3.0).powf_scalar(2.0)\n",
    "}\n",
    "\n",
    "let device = Default::default();\n",
    "// Start from x = 0\n",
    "let mut x_val: f32 = 0.0;\n",
    "\n",
    "let learning_rate: f32 = 0.1;\n",
    "\n",
    "println!(\"Starting gradient descent to minimize (x - 3)^2\");\n",
    "println!(\"Expected minimum: x = 3\");\n",
    "println!(\"---\");\n",
    "\n",
    "for i in 0..20 {\n",
    "    // Create a new tensor with current x value and require gradients\n",
    "    let x = Tensor::<B, 1>::from_floats([x_val], &device).require_grad();\n",
    "    \n",
    "    // Forward pass\n",
    "    let loss_value = loss(&x);\n",
    "    \n",
    "    // Get loss as f32 for printing\n",
    "    let loss_scalar: f32 = loss_value.clone().into_scalar::<f32>();\n",
    "    \n",
    "    println!(\"Iteration {}: x = {:.4}, loss = {:.4}\", i, x_val, loss_scalar);\n",
    "\n",
    "    // Backward pass\n",
    "    let grads = loss_value.backward();\n",
    "    let grad = x.grad(&grads).unwrap();\n",
    "    \n",
    "    // Update: x = x - learning_rate * gradient\n",
    "    let grad_val: f32 = grad.into_scalar::<f32>();\n",
    "    x_val = x_val - grad_val * learning_rate;\n",
    "}\n",
    "\n",
    "println!(\"---\");\n",
    "println!(\"Final x = {:.4}\", x_val);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Linear Regression with Gradient Descent\n",
    "\n",
    "Let's use gradient descent to fit a simple linear regression model: y = wx + b\n",
    "\n",
    "We'll generate synthetic data where the true relationship is y = 2x + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 100 data points\n",
      "True relationship: y = 2x + 1\n",
      "First 5 x values: [0.0, 0.1, 0.2, 0.3, 0.4]\n",
      "First 5 y values: [0.87993187, 0.98804677, 1.5366085, 1.7324162, 1.653858]\n"
     ]
    }
   ],
   "source": [
    "use burn::tensor::{Distribution, TensorData};\n",
    "\n",
    "let device = Default::default();\n",
    "// Generate synthetic data: y = 2x + 1 + noise\n",
    "let num_samples = 100;\n",
    "let x_data = TensorData::new((0..num_samples).map(|i| i as f32 / 10.0).collect(), [num_samples, 1]);\n",
    "// Generate noise using Burn's random tensor\n",
    "let noise = Tensor::<B, 2>::random([num_samples, 1], Distribution::Uniform(-0.25, 0.25), &device);\n",
    "\n",
    "let x = Tensor::<B, 2>::from(x_data);\n",
    "let y: Tensor<B, 2> = 2 * x.clone() + 1 + noise;\n",
    "\n",
    "println!(\"Generated {} data points\", num_samples);\n",
    "println!(\"True relationship: y = 2x + 1\");\n",
    "println!(\"First 5 x values: {}\", x.clone().slice([0..5, 0..1]));\n",
    "println!(\"First 5 y values: {}\", y.clone().slice([0..5, 0..1]));\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "rust"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training linear regression with gradient descent...\n",
      "Initial w = 0.5000, b = 0.5000\n",
      "Epoch   0: loss = 81.7705, w = 1.5358, b = 0.6586\n",
      "Epoch  20: loss = 0.0365, w = 2.0384, b = 0.7594\n",
      "Epoch  40: loss = 0.0341, w = 2.0351, b = 0.7810\n",
      "Epoch  60: loss = 0.0321, w = 2.0322, b = 0.8006\n",
      "Epoch  80: loss = 0.0305, w = 2.0295, b = 0.8184\n",
      "---\n",
      "Final: w = 2.0272, b = 0.8336\n",
      "True: w = 2.0, b = 1.0\n"
     ]
    }
   ],
   "source": [
    "// Initialize weights randomly\n",
    "let device = Default::default();\n",
    "let mut w_val: f32 = 0.5; // Start with reasonable initial values\n",
    "let mut b_val: f32 = 0.5;\n",
    "\n",
    "let learning_rate: f32 = 0.01;\n",
    "let num_epochs = 100;\n",
    "\n",
    "println!(\"Training linear regression with gradient descent...\");\n",
    "println!(\"Initial w = {:.4}, b = {:.4}\", w_val, b_val);\n",
    "\n",
    "for epoch in 0..num_epochs {\n",
    "    // Create tensors with current parameter values\n",
    "    let w = Tensor::<B, 2>::from_floats([[w_val]], &device).require_grad();\n",
    "    let b = Tensor::<B, 2>::from_floats([[b_val]], &device).require_grad();\n",
    "    \n",
    "    // Forward pass: y_pred = w * x + b\n",
    "    let y_pred = x.clone().matmul(w.clone()) + b.clone();\n",
    "    \n",
    "    // Compute loss: MSE = (1/n) * sum((y_pred - y)^2)\n",
    "    let loss = (y_pred.clone() - y.clone()).powf_scalar(2.0).mean();\n",
    "    \n",
    "    // Backward pass\n",
    "    let grads = loss.backward();\n",
    "    let w_grad = w.grad(&grads).unwrap();\n",
    "    let b_grad = b.grad(&grads).unwrap();\n",
    "    \n",
    "    // Update weights\n",
    "    let w_grad_val: f32 = w_grad.into_scalar::<f32>();\n",
    "    let b_grad_val: f32 = b_grad.into_scalar::<f32>();\n",
    "    w_val = w_val - w_grad_val * learning_rate;\n",
    "    b_val = b_val - b_grad_val * learning_rate;\n",
    "    \n",
    "    if epoch % 20 == 0 {\n",
    "        let loss_val: f32 = loss.clone().into_scalar::<f32>();\n",
    "        println!(\"Epoch {:3}: loss = {:.4}, w = {:.4}, b = {:.4}\", epoch, loss_val, w_val, b_val);\n",
    "    }\n",
    "}\n",
    "\n",
    "println!(\"---\");\n",
    "println!(\"Final: w = {:.4}, b = {:.4}\", w_val, b_val);\n",
    "println!(\"True: w = 2.0, b = 1.0\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we covered:\n",
    "\n",
    "- **require_grad()**: Mark tensors for gradient tracking\n",
    "- **backward()**: Compute gradients automatically using reverse-mode autodiff\n",
    "- **grad()**: Retrieve computed gradients\n",
    "- **Gradient Descent**: Implemented from scratch to minimize a quadratic function\n",
    "- **Linear Regression**: Used gradient descent to fit a linear model to data\n",
    "\n",
    "These concepts are the foundation of neural network training in Burn!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Rust",
   "language": "rust",
   "name": "rust"
  },
  "language_info": {
   "codemirror_mode": "rust",
   "file_extension": ".rs",
   "mimetype": "text/rust",
   "name": "Rust",
   "pygment_lexer": "rust",
   "version": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
