Exercises/2.1.First-NN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "54a1d03d-4c0c-4f30-ba7e-1beab91a7bb1",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/makagan/TRISEP_Tutorial/blob/main/Exercises/2.1.First-NN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15d1ce10-57f4-4253-8470-8c5e065600ae",
   "metadata": {},
   "source": [
    "# Our First NN\n",
    "\n",
    "Adapted from [notebooks](https://github.com/odsl-team/block-course-apr24-ML/tree/main/Exercises) developed by Nicole Hartman\n",
    "\n",
    "**Goal:** Code up a NN from scratch (numpy)\n",
    "\n",
    "**Table of Contents**\n",
    "1. Forward pass\n",
    "\n",
    "-------- lunch break -------------\n",
    "\n",
    "2. Computational graphs interlude\n",
    "3. NN gradients\n",
    "4. Checking gradients in jax\n",
    "5. **Bonus:** Set up the training loop for the model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9ab3ec40-2f2c-4e71-b96c-c5b6380e076f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae043544-48e9-4ea6-85b1-80805ef58f14",
   "metadata": {},
   "source": [
    "**Step 0:** Data generation\n",
    "\n",
    "**Plan:** Let's build on the same example as Lukas's regression problem from yesterday, the $3^{rd}$ order polynomial:\n",
    "\n",
    "$$y = \\theta_0 x^3 + \\theta_1 x^2 + \\theta_2 x + \\theta_3$$,\n",
    "\n",
    "with $\\theta = [5, 4, -2, -0.7]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "273fed51-5d47-4bc0-9427-c56f73530df1",
   "metadata": {},
   "outputs": [],
   "source": [
    "coeffs_true = [5, 4, -2, -0.7]\n",
    "\n",
    "def generate_data(N):\n",
    "    '''\n",
    "    Same function as yesterday\n",
    "    '''\n",
    "    x = np.random.uniform(low=-1, high=1, size=N)\n",
    "    \n",
    "    # y ~ N(mu=f(x), std=0.2)\n",
    "    mu = np.polyval(coeffs_true,x)\n",
    "    std = 0.2 * np.random.randn(N)\n",
    "    y = mu + std\n",
    "    \n",
    "    return x,y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "943029dc-52bb-4486-901f-43f7263fe587",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And, as before, let's try fitting a 4^{th} order polynomial to it,\n",
    "# so stack the features so X is a matrix (N, 5):\n",
    "# N = # of training examples\n",
    "\n",
    "def make_features(N, degree=4):\n",
    "    x,y = generate_data(N)\n",
    "    X = np.column_stack([x**i for i in reversed(range(degree+1))])\n",
    "    return X,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "142473be-d545-4a3b-82e8-d432770e5979",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'y')"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXsAAAEmCAYAAACDLjAiAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8g+/7EAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAvBElEQVR4nO3de3hTVbo/8G9a2pSWJqVQSKuFcilgKVAugrUOKJbDbRBmzuNoBUV0dGBgFHA8wnOOAuOl6MyDnDky4DgoHpHLjD8QRC2Hi4hCuVioUGAQahGEBoYWklJoWpP9+wN37CU72Un3TrKzv5/nyR9N995Z3YQ3K+9a610GQRAEEBFRRIsKdQOIiEh9DPZERDrAYE9EpAMM9kREOsBgT0SkAwz2REQ6wGBPRKQDDPZERDrQJtQNCCaXy4ULFy4gMTERBoMh1M0hImo1QRBQU1ODtLQ0REVJ9991FewvXLiA9PT0UDeDiEhx586dw6233ir5e10F+8TERAA3b4rJZApxa4iIWs9utyM9Pd0d36ToKtiLqRuTycRgT0QRxVdqmgO0REQ6wGBPRKQDmgr258+fx5QpU9ChQwe0bdsW/fr1w1dffRXqZhERtYrTJaC4vAqbSs+juLwKTpfylec1k7O/cuUK8vLycM899+DTTz9FSkoKTp06hfbt24e6aUREASsqq8Sij46j0lbnfi7VHIcFE7IwJjtVsdcxaGXzknnz5mHPnj344osvAr6G3W6H2WyGzWbjAC0RhVxRWSVmrD6E5kFYHGpdPmWQz4AvN65pJo2zefNmDBkyBPfffz86deqEgQMH4q233vJ6jsPhgN1ub/IgIgoHTpeARR8dbxHoAbifW/TRccVSOpoJ9t9++y2WL1+OzMxMbN26FTNmzMBTTz2Fd999V/KcwsJCmM1m94MLqogoXByoqG6SumlOAFBpq8OBimpFXk8zaZzY2FgMGTIEe/fudT/31FNP4eDBgyguLvZ4jsPhgMPhcP8sLj5gGoeIQm1T6Xk8va7U53H//WAOJubcIvn7iEvjpKamIisrq8lzt912G86ePSt5jtFodC+g4kIqIgonnRLjFD3OF80E+7y8PJw8ebLJc9988w26du0aohYREQVuaLdkpJqlA7kBN2flDO2WrMjraSbYz5kzB/v27cMrr7yC06dPY82aNfjrX/+KmTNnhrppRER+i44y4L4B3mfaLJiQhegoZSr0aibY33777di4cSPWrl2L7OxsvPjii1i6dCkmT54c6qYREfmtqKwSf91dIfn7J4d30+c8eyVwnj0RhQOnS8Bdr+6UnI1jAGAxx+HL50b67NlH3AAtEVGkCPa0S4DBnogoqJwuAXtOX5Z17KUa6Q8Ef2mmNg4RkdZ5qoPjjVLTLgEGeyKioJCqg+OJmLNXatolwDQOEZHqvNXBkaLktEuAwZ6ISHW+BmSbm53fS9FplwCDPRGR6vwdaM3oGK94GxjsiYhU5u9Aq5IDsyIGeyIilYl1cORk4JWsh9MYgz0RkcqiowxYMCFL1gCt0gOzIgZ7IqIgGJVlQVJ8jNdj2sfHYFSWRZXXZ7AnIgqCAxXVuHq9wesxV643KFoioTEGeyIilYWqREJjXEFLRKSiUJZIaIzBnohIJf6USABu5uzVmIkDMI1DRKSKQEokqLm5CIM9EZEK/C2RAABXOUBLRKQtgQ60qjVAy2BPRKSCQAda1Rqg1UywX7hwIQwGQ5NHnz59Qt0sIiKPhnZL9rmIqjED1CuVAGhsNk7fvn2xfft2989t2miq+UREHonFEdQqlQBoLNi3adMGFos6S4mJiJQkZ8WsyGKOw4IJWYrXsG9MU8H+1KlTSEtLQ1xcHHJzc1FYWIguXbpIHu9wOOBwONw/2+32YDSTiEj2QOvYbAveeGiQaj16kWZy9sOGDcOqVatQVFSE5cuXo6KiAj/72c9QU1MjeU5hYSHMZrP7kZ6eHsQWE5GeyR1o/bTMim3HrSq3BjAIgqDmPH7VXL16FV27dsWSJUvw+OOPezzGU88+PT0dNpsNJpMpWE0lIh1yugTc9epOWG11PhdLpZrj8OVzIwPq3dvtdpjNZp9xTTM9++aSkpLQq1cvnD59WvIYo9EIk8nU5EFEFAz+1LCvtNWptphKpNlgf+3aNZSXlyM1Vb0BDSKi1hiTnYrH8jJkHavWYiqRZoL973//e3z++ec4c+YM9u7di1/84heIjo5GQUFBqJtGRCRJ7mYkai2mEmlmNs7333+PgoICVFVVISUlBXfddRf27duHlJSUUDeNiEiSuP+sVO7egJtTL9VaTCXSTLBft25dqJtAROQ3MXc/Y/UhGNC0smUwFlOJNJPGISLSqjHZqVg+ZRAs5qapGos5DsunDFJ1MZVIMz17IiItG5OdilFZFhyoqMalmjp0SryZulG7Ry9isCciUonTJbQI7rk9OoSkLQz2REQq8LT3bGoQauBIYc6eiEhh4t6zzXeqstrqMGP1IRSVVQa9TQz2REQK8rb3rPjcoo+Ow+kKbqUaBnsiIgX52ntWQHDKIzTHYE9EpCC5ZQ/ULo/QHIM9EZGC5JY9ULs8QnMM9kRECrpSW+/zGDX3mpXCYE9EpBCnS8CLHx/3edzz428L2mIqEYM9EZFCfA3OitonGIPQmqYY7ImIFBKug7MAgz0RkWLCdXAWYLAnIlKMWLteKhtvQGgGZwEGeyIixYi16wG0CPjBrF3vCYM9EZGCwqF2vSeseklEpLBQ1673RLM9+8WLF8NgMGD27NmhbgoRUQvRUQbk9uiAiTm3ILdHh5AGekCjwf7gwYN488030b9//1A3hYhIEzQX7K9du4bJkyfjrbfeQvv27UPdHCIiN6dLQHF5FTaVnkdxeVXQyxh7o7mc/cyZMzF+/Hjk5+fjpZde8nqsw+GAw+Fw/2y329VuHhHpVLjtTNWcpnr269atw6FDh1BYWCjr+MLCQpjNZvcjPT1d5RYSkR6F485UzWkm2J87dw5PP/003n//fcTFyVt9Nn/+fNhsNvfj3LlzKreSiPQmXHemak4zaZySkhJcunQJgwYNcj/ndDqxe/duvPHGG3A4HIiOjm5yjtFohNEY/IJDRKQf/uxMldujQ/Aa1oxmgv29996Lo0ePNnlu2rRp6NOnD5577rkWgZ6IKBjCufhZY5oJ9omJicjOzm7yXEJCAjp06NDieSKiYAnn4meNaSZnT0QUjsK5+FljmunZe7Jr165QN4GIdMjpEpqUQnh+/G2YueYwDECTgdpQFz9rTNPBnogo2DzNp7eY4jC+fyq+PHUZV280/PR8GM2zZ7AnIpJJnE/ffBKl1V6HLUd+mkuf1DYG0/IyMGtkZsh79CLm7ImIZPA2n745240GLN1+CtuOW1Vvl1wM9kREMsjdTBwIr8VUIgZ7IiIZ/J0n33gxVThgsCcikiHQefKhXkwlYrAnIpLB13x6KaFeTCVisCciksHbZuKehMtiKhGDPRGRTFKbiTcXToupRJxnT0Tkh+abiZ+5XIu1B87Cav9po6RwWkwlYrAnIvKieWmEod2S3ZuJi2aNzPR4TDhhsCcikiB3q8HmwT8cMWdPROSB1FaDlbY6TF99CJ8cuRCilgWGwZ6IqBk5pRFmrT2MT46Efm9ZuRjsiYiakVMawSUAv10THpuJy8FgT0TUjD+rXsOp/o03DPZERM34s+o1nOrfeMNgT0TUzNBuybCYjLKPD5f6N95oJtgvX74c/fv3h8lkgslkQm5uLj799NNQN4uIIlB0lAEFQ7vIPj5c6t94o5lgf+utt2Lx4sUoKSnBV199hZEjR2LixIk4duxYqJtGRBEoo2OCrOOS2saETf0bbzSzqGrChAlNfn755ZexfPly7Nu3D3379g1Rq4goUsntrU/Lywi71bKeaCbYN+Z0OvGPf/wDtbW1yM3NlTzO4XDA4fipXoXdbg9G84goAoglja22Osn59u3jYzBrZGZQ2xUozaRxAODo0aNo164djEYjpk+fjo0bNyIrK0vy+MLCQpjNZvcjPT09iK0lIi3zVdLYAKDwl/000asHAIMgCOE/QfRH9fX1OHv2LGw2Gz744AP87W9/w+effy4Z8D317NPT02Gz2WAymYLVbCLSMLn1cULFbrfDbDb7jGuaCvbN5efno0ePHnjzzTdlHS/3phARNSZV+TIcyI1rmszZi1wuV5OeOxGRGrRQ1dIXzQT7+fPnY+zYsejSpQtqamqwZs0a7Nq1C1u3bg1104iIwp5mgv2lS5fwyCOPoLKyEmazGf3798fWrVsxatSoUDeNiCjsaSbYr1y5MtRNICLSLE1NvSQiosAw2BMR6QCDPRGRDmgmZ09EpJZwnkevFAZ7ItK1cF8hqxSmcYhIt4rKKjFj9aEW+81abXWYsVo7+8vKwWBPRLrkdAlY9NFxjxUtxee0sr+sHAz2RKRLByqqW/ToGxOgnf1l5WCwJyJdkrtv7PbjVpVbEhx+B/upU6di9+7darSFiCho5O5EtbH0fESkcvwO9jabDfn5+cjMzMQrr7yC8+fPq9EuIiJVDe2WjOSEWJ/HVdc2REQqx+9g/+GHH+L8+fOYMWMG1q9fj4yMDIwdOxYffPABGhoa1GgjEZHioqMMuG+AvKmVVtsNlVujvoBy9ikpKZg7dy6+/vpr7N+/Hz179sTDDz+MtLQ0zJkzB6dOnVK6nUREreZ0CSgur8Km0vMoLq/CLUnxss6rrq1XuWXqa9WiqsrKSmzbtg3btm1DdHQ0xo0bh6NHjyIrKwuvvfYa5syZo1Q7iYhaxdPiKWO0vFWyye2MajUraPwO9g0NDdi8eTPeeecd/N///R/69++P2bNn46GHHnJvibVx40Y89thjDPZEFBbExVPNh1kdTnkDrxaTvMHccOZ3sE9NTYXL5UJBQQEOHDiAnJycFsfcc889SEpKUqB5RESt423xlByp5pu1crTO72D/+uuv4/7770dcnPQnXVJSEioqKlrVMCIiJfhaPOXLgglZEVEUze9g//DDD6vRDiIiVchdPOXJY3kZEVMMTTMraAsLC3H77bcjMTERnTp1wqRJk3Dy5MlQN4uIwpzcxVOejMqyKNiS0NJMsP/8888xc+ZM7Nu3D9u2bUNDQwP+7d/+DbW1taFuGhGFsaHdkpFq9j/gt4+PiYhcvUgz9eyLioqa/Lxq1Sp06tQJJSUlGD58eIhaRUThLjrKgAUTsjB99SG/ztN+gYSmNNOzb85mswEAkpMj55OXiNQxJjsVf3loEPwZZ716PTLKJIg0GexdLhdmz56NvLw8ZGdnSx7ncDhgt9ubPIhIn8b1T8UbBQP9Oqc1g7vhRpPBfubMmSgrK8O6deu8HldYWAiz2ex+pKenB6mFRBSOxvVPw4opg5CcECPr+NYM7oYbzQX7WbNmYcuWLfjss89w6623ej12/vz5sNls7se5c+eC1EoiCkdOlwBz21hM6J/m89hIWUwl0swArSAI+N3vfoeNGzdi165d6Natm89zjEYjjEbt17QgotYrKqvEws3HYbXLS808P/62iFhMJdJMsJ85cybWrFmDTZs2ITExEVbrzd1jzGYz2rZtG+LWEVE4Kyqr9Hs2TvuEyOooaiaNs3z5cthsNtx9991ITU11P9avXx/qphFRGHO6BMzbcNTv8yJpcBbQUM9eECJt1isRqc3pErBy97e4et3/jZUiaXAW0FCwJyLyR1FZJeZtOOp3oDcAsETY4CzAYE9EESiQHD1wM9ADkVPpsjEGeyKKKE6XgIWbjwV0rsUchwUTsiKm0mVjDPZEFFEOVFTDanf4dU58bDTeemQI7ujeIeJ69CIGex1wugQcqKjGpZo6dEq8mYuM1Dc0USCzaJb8agDyenZUoTXhg8FeIwIN2J42WU6N4K+qRP7MokmOj8Erv+yni/8LDPYa4ClgJ7WNwbS8DMwamSkZ9KU2Wbba6jBj9SEsnzJIF29y0peh3ZJhMRl9pnKSE2Kwb34+YttoZrlRq+jjr9QwMWA330Pz6o0GvL79FAa/tA1FZZUtzvO2ybL43KKPjsPp4voFiizRUQYsvK+vz+Ne+UU/3QR6gME+rHkL2KKr1xswY/WhFgHf1ybLAoBKW11E1esmEo3JTsWKKYOQFN+yumX7+Bis0OG3WqZxwkjzvLxLELwGbJEAYOHmYxiVZXGndOQOUkXaknAi0ZjsVIzKsmBfeRWKv70MwIDcHh0iesaNNwz2YUIqLy+X1e7AGztP4+n8TADyB6kibUk4UWPRUQbkZXZEXmZkz7SRg2mcMOAtL++P17d/407niJssS/VfDIi8et1EJI3BPsTk5OX9IQ66ipssA2gR8CN5STgRecZgH2K+BlL91XjQdUx2KpZPGYTOpqZ1uTubjJx2SaQzzNkHkaeFUWoMkLa8plTfnkibuCrcfwz2QSK1kvX+wd730W1seGZH7D512edx4qCr1KKqi3YuqiLt4qrwwDCNEwRSA7CVtjr8eedp2df5+txVj/OGG4syAFdqHbIXVdX/4EJxeRU2lZ5HcXlVk0VWTpcg+TuiUJD6vySuCve0wJBuYs9eZUoOwNrqfvB5jEsAfrvmMObkX5O1qGrYK9txpdHmDhaTEQVDu8B2owEfll5AdW29+3fsPVEo+erAGHCzA5MQ0wb7z1RB7/PqmzMIOtrvz263w2w2w2azwWQyBeU1i8urUPDWPsWuZwBgbtsGtrof4O1fLiE2GrX1TsVeV3xtAEz/UNA5XQJW7anAix+f8PvcpPgYLI7gYmdy45qm0ji7d+/GhAkTkJaWBoPBgA8//DDUTfJJ6QFYAcDVG94DPQDFA7342gBr6lBwFZVV4q5XdwYU6IGbJUWmM8WjrWBfW1uLAQMGYNmyZaFuimyhXKHazhit+Lwb1tShYJLK0QdC750UTeXsx44di7Fjx4a6GX4RV7JabXWKLZyS62eZHVFUdlGVa7OmDqlN6QWHYiclt0cHha6oLZrq2fvL4XDAbrc3eQSbt5WsgTAASGobLevYKcMysHzKIL9q7MjFmjqkNqUXHAKA1XZD0etpiaZ69v4qLCzEokWLWnUNJRZviCtZm88N9pcBP+bNDVEAvOfkk+JjcEePm7MQEo0xmLxyf8Cv27wNFnMcBndtj+Lyqhb3pfH96phgBAzA5WsOLnwhv6nx7bHx7DK9iehgP3/+fMydO9f9s91uR3p6uuzz5SzeEIOb1XYD1bX1SG5nhMXUMrCJ5VbFQHjm8nW8vv0bv/4ec3wMrl5vwNXrvgukLf5lP/fr39Gjg2KpJAHA4K7tMfy1nU12Ako1x+G+AanY/HWl5Acap26SP9T49pjczuj7oAil2amXBoMBGzduxKRJk2Sf48/US6nVp42nHwKQ7K17C2xOl4C7Xt3pVy//P8f1wcovz8Bq935O58RYLJqY3eSDpVNiHK7UOjBzzeGgjxt4Mie/F2aN7MlePnkVyP8TX9Y+cUfE5ezlxrWI7tkHSs7ijXkbjsJ2vUEyeFba6jB99SE8npeB/CxLkzTHqj0Vfr+BbTcafAZ6AFjywEDU1DW0+E+Sao7Dk8O7Yc2Bc6iRsThLTa9v/wZrD3yHhff1ZS+fJEVHGXDfgFS8ubtCkevpvaS3poL9tWvXcPr0T+UFKioqUFpaiuTkZHTp0kWx15GzpZ+cVAoArNxzBiv3nJGV5vBOXi9454mLeHvPGY+bjP91dwUeye2Kd4u/C+D1lWW1O1ifh7wqKqvEXxUI9CzpfZOmZuN89dVXGDhwIAYOHAgAmDt3LgYOHIgXXnhB0ddRY2Co0laHN3f736MHgOSEGNlfPTeWnvdaD+ejIxf8fn016X3uM3mm5LRLizmOnQporGd/9913IxhDDOE2rXDigFtwR3fvg6wGAMkJsajyMttAAFBd24DkhFifsxKS42NQLfPbS6AaL9DK7dGBZWvJTYlpl7Pu6Ym8nh35PvqRpnr2weJrS79gS0uKk7Xz1MScNFnXm5STBoOH64hm35uJ53+eFUhTA3Kpps69JL7grX14el0pCt7ah7te3an7Je56te24tdXXyOzcDrk9WARNxGDvgZzA6qvUsJLE8QFxvr7F3PSbh/g1dVSWRdb1RmVZPF5HtP6rc9jxz0uta7QfzlyulSwBPX31Ifz39lNM9ehIUVkl3t5zptXXCbdv6KGmqTROMEkthLL8OKXy8Nkris0S8MVgkJ6v33xBk69Uj8X80/Eul4Dfrjnc4rhKWx22HAlOjzo5IRZr9p/1mpvl7B39cLoELNx8rFXXaPw+p58w2HshFVi3llmDFugBtBicjY4yeBywFb+RzFh96KfVtj9qPiPB6RICriKoJLkrGjl7Rx/e2Hm6yWK9QOl95o0nTOP4EB1lwNBuyeiUGIdLNXX4845TmLX2kM/zlHqbtY+PwR3d5S8C8ZXqEQOlGnVHgoGzdyJXUVml36vKm0tqG8MOgQT27H3wVDJBjgRjG1xz+F68lGCMRq1Dus5NYaOyB3J5S/WItFi1svnsHYoc4lTL1lo2eRDyenZUoEWRh8HeC6mSCXLICfQAkGhsgz/++wD8YcvxJitkW1tHRirVI9Ly4JUWP6jIOyW+aaaa4/z6Fqw3DPYSlK6lLcVqd6B9Qiz2zBsZ1Dnmoayz31pa/qAiz5T4AGee3jsGewnBzGlfqqnz2RNXmrfB3HDFWRaRq7Uf4HPyezFP7wMHaCUEM1UQqp6q1GBuqjkOvxneDckJsSFplxQBwIO3K1cDicLHldrAZ+BYTEbMGtlTwdZEJvbsJQQjAIdDT9XbYG4fiwlz/v51yNrmyevbv8G6g2dZFz+CtGYasAHAwvv6Mn0jA4O9hGDltMMhzyiVQrKY24agNb5ZbXXuOfe+Zh1R+GtNyvSpezP5oS8Tg70EtXPaWti1KZAPvCgDIAjqjgE03lNg4WZlZzFR8LWmDs67e8/gttRE/nvLwJy9F1I57UB1SIjFY3kZWPvEHfjyuZFh/wb1Z7N0sbDaEz/rJuv41hL3FGi+oYvY62cBNW1obR2cqzca+O8tk2a3JQyEP9sSNtZ8E+0Nh7/H/zt0Xvb5j+VlYFSj3aq0xtPCsigD0Hgha+MedaAL0ZQijoV8+dxITd5vvVBq20G9/3tzW0IFiTntorJK/P6Dr2W/OZPiY7D4l/3Cvgfvi6dB3MFd26Pkuysec+UtN1evxdt7KmC7EZztELnSVhuUmt7Mf295GOxlCmQ17bKCQcjLjIyl254Gcb39x2p8vNMlYO2Bc0EL9iKr7UZQX4/8o/T0Zq6s9o45exn8XU1rwI9Lt9nLAHCzBydns3SlvfjxCeZyw9iZy7WKXo8rq73TXLBftmwZMjIyEBcXh2HDhuHAgQOqv6Y/Xze5uXFLoepxXamt5+BdmLpZ4fKUItcSO1dcWe2dptI469evx9y5c7FixQoMGzYMS5cuxejRo3Hy5El06tRJtdf1J1hZOPWvhVD1uMQpmos+Oo5RWRaPH77i4LvVdgPVtfVIbmeExcT5+mpw32t7HV7c0roNSkTsXMmnqWC/ZMkSPPHEE5g2bRoAYMWKFfj444/x9ttvY968eaq9rtxg9fz42/BoXje+6ZqRM1/fYjLi8bu64dyVG6i50YCNpRd8XvfxvAxsLD2P6lrpjdHFwbt931a1KH3rbdYQ5+srx+kS8MbOU3hnzxlcvSF/E/uf909FyXdXvM4CY+dKPs0E+/r6epSUlGD+/Pnu56KiopCfn4/i4mKP5zgcDjgcP9XcsNvtAb22r2AlTv1ioPdMzg5ajbccLC6vkhXs87MsyL41CXPWl/o8dub7h7D433+aGeVrwL3yx/n6s/N7IaNjPFfnBqiorBLzNhx176Psj1FZnfHfDw6UPQuMvNNMzv7y5ctwOp3o3Llzk+c7d+4Mq9XzCrzCwkKYzWb3Iz09PaDXlrMBOb9Geid3By3gpw9XqbvZOEdrMcn71tV48Y24z6mvAXcBN2vxPL2uFAVv7cNdr+5k/t8P4gdqIIEeuPmNWpzVNTHnFuT26IDYNlFNfub/Ofk007MPxPz58zF37lz3z3a7PeCA72sDcn6N9E3ODlqAf3vp+lvSYdFHx/HPypqA9jmtbFSTh//e3rVmP4hwKBAYiTQT7Dt27Ijo6GhcvHixyfMXL16ExWLxeI7RaITRaFSsDXKDFUmTW7df7odr4w8GX8T8/dIdgc8CEeB9wFdvGq8ub/z/IdAFU/ymrB7NBPvY2FgMHjwYO3bswKRJkwAALpcLO3bswKxZs4LWjmBvMqJncj9cx2Sn4tc/64a3vqgISru4WvMmTwPc4sC24wdXQNfkN2X1aCbYA8DcuXMxdepUDBkyBEOHDsXSpUtRW1vrnp1DkUfOh+snRyqx8svgBHqR3ldrSg1wW90D25myr9UhIRb/Nf42WMxt+U1ZRZoK9g888AD+9a9/4YUXXoDVakVOTg6KiopaDNqSfhSVVeK3a3yncJSm59Wa3vLx4nPv7j0DiykOF+2+x1Ieyc3AfTm3MMirjFUvSbOUqpror+SEGBz8z1G6DU7F5VUoeGufz+PaGaNxzeGUtR8E1zUETm5c08zUS6LmgrkpfGMvTczWbaAH5KewrjmcAABzfIzPY7kPgfoY7EmzQpE3/83wbhjXPy3orxtO/E1hGSDgvceG4vUHcpCc4Dnwiz3/RR8dh9Olm2RDUDHYk2a1Jm+eGOffcJW4C9f8cVkBv6ZWOV0CisursKn0PIrLq1BV498ahSvXf8DT60txtuq6rNIWByqqW9li8kRTA7REjYkLqgJJ5Sz8eRbmbTyKBqe8XqQA4G9fVGBw1/a6yit7ml5pCCCDVV1bj9e3fyPrWL3PdFILe/akWeKCqkCy5/+65pAd6BtTKs3QvLccjqkLcXpl8w9Ttad06Hmmk5rYsydNE1fa+ltsa3HRSb9fS+72d81XlTYv3nWlth4vfux5MZI/3xqkVq8qoTXlDgLFMgnqYrAnzRNX2gZSRjcQ3tIMcjZn98TqZ90db6tXlUgzqT3TyVfNI1Ie0zgUEaKjDHg6vxdKnh+FtU/cgcfyMpCcEKvKa0mlGaTSHnIyNP7MRvnkyAVM9/A6Sk5fVDNvPic/U1b1U1IWe/YUUcTyCrk9OuA/x2dh1Z4KvPjxCUVf40pty9koSqQ95KSJPjlSiVlrD0ue72tnLrnUyJuLaZpZIzMxa2QmCwoGGXv2FLGiowzomKhc1VPRix+faNL7droErNpToVjaQ6pXLZaG8NbxV2r6oq89BfzVPE3TvE49A7362LOniKZGD7XSVodVeyrQMdGIM5drsfbA2YDq40tp3ObG+7Y+/+FR2ddQIg3z4O1dZE+XbI7bB4YfBnuKaP7Oxc9OM6Hsgu/tK5VODYmS4mPcs1G87ZHri/iBEciMnda8rujPD+SgQ2Ic0zRhhMGeIpo4F3+6j81NxHzy/LG3YfLK/cFpnAfT7ry5j7GvPXK9Ebds9GfGjtMlYF95FVbvP4NPyy42v6RbXEwU6hq816o3AHj503/iy+dGMsCHEebsKeKNyU7FiimDkCRRkKtxPvmOHh0UzVX7I6ltG8wa2bPVg70LJmRh23Grx5lBnmbsFJVVYvBL2zB55X6vgR6Az0APsOxBuGKwJ10Yk52Kkv8ahTn5vZDUtmnQbzztz9vm8qozGLDtuLVVW/r95aGBGJVl8VlvXpziWVRWiemt2BTcG5Y9CC9M45Bu3JyLn4lZI3t6zWNL7X+rNtv1BsxYfQjT8jICOv/pezMxrn8aisurvLZb7HnvK6/Cws3HAmusDCx7EF4Y7El35Gx12Hj/2z2nL+ONz06r3i5xnvym0gt+n5vUtg1+d+/NrQC3HbfKOqf428uKziISsexBeGIah0iC+KGQ2bmdIteLa+M7MSQAqKqt97sE87S8mwO7TpeAD2V/WCifqGLZg/ClmWD/8ssv484770R8fDySkpJC3RzSEf836/Cs7gf5Q641dT/49ZoZHRMA3KxpU11b7/P4hNhomNv63kHKXyx7EL40k8apr6/H/fffj9zcXKxcuTLUzSEd8WeufoeEWCwYn4Wn/l6qfsMaET+Q5A6K1tY78fInJ2QVafMlOSEGz/+8LywmzqcPZ5oJ9osWLQIArFq1KrQNId0RZ+j4mvduAPDixGz8YYt6g56eXrNxfrxjO//KQ7Q20BsAvPKLfuzJa4Bm0jiBcDgcsNvtTR5EgRBn6KSaPad0Un9MX7RPiFVl0NMbMT9eVFaJZ4L4jcJiMjJloyGa6dkHorCw0P2NgKi1Gs/QsdpuoLq2HsntjE3SF5tKzwetPUnxMVj8y5u96tasuA3EnPxemDWyJ1M2GhLSYD9v3jy8+uqrXo85ceIE+vTpE9D158+fj7lz57p/ttvtSE9PD+haRIDvaZvBnFu+rGAQ8jI7BnVXKSU3SKHgCmmwf+aZZ/Doo496PaZ79+4BX99oNMJoVL7ELZGUod2SYTEZVU/lpJrjcMePHzpK7yqVnBCD6tqGJj//IucW5GdZOACrYSEN9ikpKUhJSQllE4gUFR1lwML7+vosvNZajeexW+3KrvIVZ9awYmVk0UzO/uzZs6iursbZs2fhdDpRWloKAOjZsyfatVNm0QuREsTCa/5ugi7XnPxeTdIo1deU/RZhMcX5XGFM2qOZYP/CCy/g3Xffdf88cOBAAMBnn32Gu+++O0StIvJMHMyVUzbYH/GxUZg1smeT55Tca7dDQizLHEQozUy9XLVqFQRBaPFgoKdwFR1lQF5mRyyfMgQrPEzbTDXH4S8PSU/n9OTuXiktUioWc1tF2gsAE3PSmLKJUJrp2RNpWeNpm81z4VFRkJ3jnzwso8Vzclb4Nh90lTIqyyKrHaQ9munZE2md1CbbY7JT8ZeHBvosS5YUH+OegdP8ugsmZEmebwDw0sRsn5uypLJSZURjsCcKA+P6p2HZQ4O8HrP4l/0kUyxSK3zFlb3j+qdJbspi+PHBSpWRzSAIQrAW3YWc3W6H2WyGzWaDyWQKdXOIWigqq8TCzceazNO3mIxYeF9fWQuZfG0w7s++tKQNcuMagz1RmPEVsMP9+hRccuMaB2iJwoycnbTC+foUnpizJyLSAQZ7IiIdYLAnItIBXeXsxbFobmJCRJFCjGe+5troKtjX1NQAAGvaE1HEqampgdlslvy9rqZeulwuXLhwAYmJiTAY5E81Ezc9OXfuHKdsNsN7I433RhrvjTR/740gCKipqUFaWhqioqQz87rq2UdFReHWW28N+HyTycQ3pgTeG2m8N9J4b6T5c2+89ehFHKAlItIBBnsiIh1gsJfBaDRiwYIF3M/WA94babw30nhvpKl1b3Q1QEtEpFfs2RMR6QCDPRGRDjDYExHpAIM9EZEOMNh78PLLL+POO+9EfHw8kpKSZJ0jCAJeeOEFpKamom3btsjPz8epU6fUbWiIVFdXY/LkyTCZTEhKSsLjjz+Oa9eueT3n7rvvhsFgaPKYPn16kFqsnmXLliEjIwNxcXEYNmwYDhw44PX4f/zjH+jTpw/i4uLQr18/fPLJJ0FqafD5c29WrVrV4v0RFxcnebyW7d69GxMmTEBaWhoMBgM+/PBDn+fs2rULgwYNgtFoRM+ePbFq1Sq/X5fB3oP6+nrcf//9mDFjhuxzXnvtNfz5z3/GihUrsH//fiQkJGD06NGoq6vzfbLGTJ48GceOHcO2bduwZcsW7N69G08++aTP85544glUVla6H6+99loQWque9evXY+7cuViwYAEOHTqEAQMGYPTo0bh06ZLH4/fu3YuCggI8/vjjOHz4MCZNmoRJkyahrKwsyC1Xn7/3Bri5YrTx++O7774LYouDp7a2FgMGDMCyZctkHV9RUYHx48fjnnvuQWlpKWbPno1f//rX2Lp1q38vLJCkd955RzCbzT6Pc7lcgsViEf74xz+6n7t69apgNBqFtWvXqtjC4Dt+/LgAQDh48KD7uU8//VQwGAzC+fPnJc8bMWKE8PTTTwehhcEzdOhQYebMme6fnU6nkJaWJhQWFno8/le/+pUwfvz4Js8NGzZM+M1vfqNqO0PB33sj9/9apAEgbNy40esx//Ef/yH07du3yXMPPPCAMHr0aL9eiz17BVRUVMBqtSI/P9/9nNlsxrBhw1BcXBzClimvuLgYSUlJGDJkiPu5/Px8REVFYf/+/V7Pff/999GxY0dkZ2dj/vz5uH79utrNVU19fT1KSkqa/JtHRUUhPz9f8t+8uLi4yfEAMHr06Ih7jwRybwDg2rVr6Nq1K9LT0zFx4kQcO3YsGM0Ne0q9b3RVCE0tVqsVANC5c+cmz3fu3Nn9u0hhtVrRqVOnJs+1adMGycnJXv/Whx56CF27dkVaWhqOHDmC5557DidPnsSGDRvUbrIqLl++DKfT6fHf/J///KfHc6xWqy7eI4Hcm969e+Ptt99G//79YbPZ8Kc//Ql33nknjh071qrihZFA6n1jt9tx48YNtG3bVtZ1dNOznzdvXosBoOYPqTeiHqh9f5588kmMHj0a/fr1w+TJk/G///u/2LhxI8rLyxX8K0ircnNz8cgjjyAnJwcjRozAhg0bkJKSgjfffDPUTYsYuunZP/PMM3j00Ue9HtO9e/eArm2xWAAAFy9eRGpqqvv5ixcvIicnJ6BrBpvc+2OxWFoMsv3www+orq523wc5hg0bBgA4ffo0evTo4Xd7Q61jx46Ijo7GxYsXmzx/8eJFyftgsVj8Ol6rArk3zcXExGDgwIE4ffq0Gk3UFKn3jclkkt2rB3QU7FNSUpCSkqLKtbt16waLxYIdO3a4g7vdbsf+/fv9mtETSnLvT25uLq5evYqSkhIMHjwYALBz5064XC53AJejtLQUAJp8OGpJbGwsBg8ejB07dmDSpEkAbm6Os2PHDsyaNcvjObm5udixYwdmz57tfm7btm3Izc0NQouDJ5B705zT6cTRo0cxbtw4FVuqDbm5uS2m6Ab0vvF39FgPvvvuO+Hw4cPCokWLhHbt2gmHDx8WDh8+LNTU1LiP6d27t7Bhwwb3z4sXLxaSkpKETZs2CUeOHBEmTpwodOvWTbhx40Yo/gRVjRkzRhg4cKCwf/9+4csvvxQyMzOFgoIC9++///57oXfv3sL+/fsFQRCE06dPC3/4wx+Er776SqioqBA2bdokdO/eXRg+fHio/gRFrFu3TjAajcKqVauE48ePC08++aSQlJQkWK1WQRAE4eGHHxbmzZvnPn7Pnj1CmzZthD/96U/CiRMnhAULFggxMTHC0aNHQ/UnqMbfe7No0SJh69atQnl5uVBSUiI8+OCDQlxcnHDs2LFQ/QmqqampcccUAMKSJUuEw4cPC999950gCIIwb9484eGHH3Yf/+233wrx8fHCs88+K5w4cUJYtmyZEB0dLRQVFfn1ugz2HkydOlUA0OLx2WefuY8BILzzzjvun10ul/D8888LnTt3FoxGo3DvvfcKJ0+eDH7jg6CqqkooKCgQ2rVrJ5hMJmHatGlNPggrKiqa3K+zZ88Kw4cPF5KTkwWj0Sj07NlTePbZZwWbzRaiv0A5//M//yN06dJFiI2NFYYOHSrs27fP/bsRI0YIU6dObXL83//+d6FXr15CbGys0LdvX+Hjjz8OcouDx597M3v2bPexnTt3FsaNGyccOnQoBK1W32effeYxvoj3Y+rUqcKIESNanJOTkyPExsYK3bt3bxJ75GKJYyIiHdDNbBwiIj1jsCci0gEGeyIiHWCwJyLSAQZ7IiIdYLAnItIBBnsiIh1gsCci0gEGeyIiHWCwJyLSAQZ7IoX861//gsViwSuvvOJ+bu/evYiNjcWOHTtC2DIigLVxiBT0ySefYNKkSdi7dy969+6NnJwcTJw4EUuWLAl100jnGOyJFDZz5kxs374dQ4YMwdGjR3Hw4EEYjcZQN4t0jsGeSGE3btxAdnY2zp07h5KSEvTr1y/UTSJizp5IaeXl5bhw4QJcLhfOnDkT6uYQAWDPnkhR9fX1GDp0KHJyctC7d28sXboUR48eRadOnULdNNI5BnsiBT377LP44IMP8PXXX6Ndu3YYMWIEzGYztmzZEuqmkc4xjUOkkF27dmHp0qV47733YDKZEBUVhffeew9ffPEFli9fHurmkc6xZ09EpAPs2RMR6QCDPRGRDjDYExHpAIM9EZEOMNgTEekAgz0RkQ4w2BMR6QCDPRGRDjDYExHpAIM9EZEOMNgTEekAgz0RkQ78fxUSBonWAYDlAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 400x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# sanity check plot\n",
    "plt.figure(figsize=(4,3))\n",
    "N=200\n",
    "X,y = make_features(N)\n",
    "\n",
    "plt.scatter(X[:,-2],y)\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32d4cf94-1812-4b1e-bdce-8c96b0f9c021",
   "metadata": {},
   "source": [
    "## Step 1: Multi-layer perceptron (forward pass)\n",
    "\n",
    "Define a simple feed-forward NN that accepts a input matrix\n",
    "\n",
    "$$X \\in \\mathbb{R}^{N \\times d}$$\n",
    "\n",
    "(where for the example just above $N=200$, $d=5$)\n",
    "\n",
    "and constructs an MLP with a single hidden layer with $H = 16$ units, and a ReLU non-linearity.\n",
    "\n",
    "**Hint:** this is just three equations:\n",
    "\n",
    "$$z = X W^T + b_1$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "648c20e1-2788-41e3-bb41-555cc7931d22",
   "metadata": {},
   "source": [
    "$$h = \\mathrm{ReLU}(z_1)$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e6cdfbf-0eea-4034-95cf-2cc45e6518a1",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "\n",
    "$$y_\\textrm{pred} =h W_2^T  + b_2$$\n",
    "\n",
    "where $W_1 \\in \\mathbb{R}^{H \\times d}, b_1 \\in \\mathbb{R}^{H}, W_2 \\in \\mathbb{R}^{1 \\times H}, b_2 \\in \\mathbb{R}$,\n",
    "$z \\in \\mathbb{R}^{N\\times H}$, $h \\in \\mathbb{R}^{N \\times H}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b4fdd7f-9696-4240-bdcc-5d3286eadf3e",
   "metadata": {},
   "source": [
    "**Potentially useful functions:** \n",
    "- `np.einsum`\n",
    "- `np.where`\n",
    "- `np.random.randn`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "id": "1dd4b75b-c2a8-48e2-ab9f-c1b460ccf765",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(16, 5)\n",
      "(16,)\n",
      "(1, 16)\n",
      "(1,)\n"
     ]
    }
   ],
   "source": [
    "# Initialize the parameters (randomly) to have the desired shape\n",
    "H = 16\n",
    "d = X.shape[1]\n",
    "\n",
    "'''\n",
    "TO DO: Initialize the parameters\n",
    "'''\n",
    "\n",
    "# 1st layer of the NN\n",
    "W1 = np.random.randn(H, d)\n",
    "b1 = np.random.randn(H)\n",
    "\n",
    "# Now the 2nd layer of the NN\n",
    "W2 = np.random.randn(1,H)\n",
    "b2 = np.random.randn(1)\n",
    "\n",
    "for v in [W1,b1,W2,b2]:\n",
    "    print(v.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ab05ccd-f9af-4df1-b66a-cabcd937c219",
   "metadata": {},
   "source": [
    "$$ReLU(x) = \\max(x,0)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "id": "0b5e09c8-11f3-4898-bf71-efa74a168335",
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu(x):\n",
    "    return np.where(x>0, x, 0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "id": "9f91caa6-1558-4cda-bf1b-b6660d98f13b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 3, 0])"
      ]
     },
     "execution_count": 203,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = np.array([1,3,-1])\n",
    "relu(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "id": "0141bf7a-7df2-4c01-aa04-2c07baf64be9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(200, 16)\n",
      "(200, 16)\n",
      "(200, 1)\n"
     ]
    }
   ],
   "source": [
    "# Set up the feedforward model\n",
    "\n",
    "# z = \"W1x + b1\"\n",
    "z = np.einsum('hj,ij->ih',W1,X) + b1\n",
    "print(z.shape)\n",
    "\n",
    "h = relu(z)\n",
    "print(h.shape)\n",
    "y_pred = np.einsum('oh,ih->io',W2,h) + b2\n",
    "print(y_pred.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce40a3e3-39c1-4b92-bd02-ad7c3bb5b1c5",
   "metadata": {},
   "source": [
    "We're thinking about a regression task right now, so no need to have a non-linearity on the last layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 207,
   "id": "be087011-6d57-496b-8fa5-9f68ba3c7f45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sanity check: Does the output have the dimensionality we expect?\n",
    "assert y_pred.shape[0] == N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3549f9ce-4878-43ca-be2f-00cfe25597df",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "28a5a566-c2bf-4d72-a60c-98b4a33d6883",
   "metadata": {},
   "source": [
    "## Step 2: Computational graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64e1e4d5-6731-4583-b71a-6cdd3ae3d227",
   "metadata": {},
   "source": [
    "**Motivating example:**  Let's code up the model we had worked out in lecture.\n",
    "\n",
    "$$f(x,y,z) = (x+y) \\cdot z$$\n",
    "\n",
    "<span style=\"color:red\">$$f_1(x,y) = q = x+y$$</span>\n",
    "\n",
    "<span style=\"color:blue\">$$f_2(q,z) = q \\cdot z$$</span>\n",
    "\n",
    "<img src=\"toy-ex.pdf\" width=650/>\n",
    "\n",
    "**To Do:** Implement `f1`, `f2` and `f` functions in numpy with both a forward and a backward mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "bf1022aa-71a4-4ab1-a4fd-9b5f542549d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f1(x,y, grad=None):\n",
    "    '''\n",
    "    f1(x,y) = x+y\n",
    "    \n",
    "    Inputs:\n",
    "    - x,y: The inputs to the functoin\n",
    "    - grad: if not None, this is the \"upstream gradient\"\n",
    "\n",
    "    Outpus:\n",
    "    - If grad is passed, return f(x,y), grad\n",
    "    - Else, just the \"normal\" fct, return f1(x,y)\n",
    "    '''\n",
    "    \n",
    "    out = \n",
    "    if grad:\n",
    "        raise NotImplementedError\n",
    "    else:\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "c0e4b3e8-116e-4727-81d8-69e2843d35a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f2(q,z, grad=None):\n",
    "    '''\n",
    "    f2(q,z) = q*z\n",
    "    '''\n",
    "\n",
    "    out =\n",
    "    if grad:\n",
    "       raise NotImplementedError\n",
    "    else: \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "5941b7b7-de97-49f8-992e-672758648f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x,y,z, grad=None):\n",
    "    '''\n",
    "    f(x,y,z) = f2(q,x)\n",
    "    '''\n",
    "\n",
    "    # Need to run the program 1x to get the values to eval the grad\n",
    "    q = \n",
    "    out = \n",
    "\n",
    "    if grad:\n",
    "\n",
    "        '''\n",
    "        Fill in\n",
    "        '''\n",
    "        \n",
    "        return out, [dfdx, dfdy, dfdz]\n",
    "    \n",
    "    else:\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2b1e2a2-a2fe-4a67-8997-1052be88420e",
   "metadata": {},
   "source": [
    "Check the forward mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "id": "fce6e1be-3907-4442-b162-33a1e56f5e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "x,y,z = 5,-2,4\n",
    "print(f(x,y,z))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaf31f60-adf1-4e51-bea5-087feaafc3dc",
   "metadata": {},
   "source": [
    "\n",
    "**What about the derivative?**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "f12d235d-33a6-40a6-9936-f28a361bb8b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdf = 1\n",
    "out, grad = f(x,y,z, grad=dfdf)\n",
    "\n",
    "for xi, dfi in zip(['x','y','z'],grad):\n",
    "    print(f'df/d{xi}= {dfi}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73adc084-eb39-4831-a7d5-2d36b60002f2",
   "metadata": {},
   "source": [
    "## Step 3: Gradients of a NN\n",
    "\n",
    "OK, let's re-code our \"neural network\" up, but this time with the gradient computation so we can do _back prop_. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "d6ea1e0c-4b15-4715-b0d4-fdc9be7d8995",
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear(W,X,b, grad=None):\n",
    "    '''\n",
    "    Linear layer mapping H = W x + b\n",
    "    - W: array (m,n)\n",
    "    - x: array (bs, n)\n",
    "    - b: array (m)\n",
    "\n",
    "    Outputs:\n",
    "    - H: array (bs,m)\n",
    "\n",
    "    If grad is true, return tuple: out, (grad_W, grad_b, grad_H)\n",
    "    - grad_W: array (bs, m, n)\n",
    "    - grad_b: array (bs, m)\n",
    "    - grad_X: array (bs,    n)\n",
    "    '''\n",
    "\n",
    "    assert X.shape[1] == W.shape[1]\n",
    "    assert W.shape[0] == b.shape[0]\n",
    "\n",
    "    '''\n",
    "    TO DO: Fill in\n",
    "    '''\n",
    "    raise NotImplementedError\n",
    "    \n",
    "    H = \n",
    "    \n",
    "    if grad is not None:\n",
    "        \n",
    "        grad_W = \n",
    "        grad_b = \n",
    "        grad_X = \n",
    "        \n",
    "        return H, (grad_W, grad_b, grad_X)\n",
    "    else:\n",
    "        return H"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "849c6d4b-25b8-45f0-9d63-4ea1ac07f1fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu(z,grad=None):\n",
    "    '''\n",
    "    Non-linearity: better gradient flow\n",
    "    '''\n",
    "    out = \n",
    "\n",
    "    if grad is not None:\n",
    "        '''\n",
    "        To do: Calculate the local grad and return it too\n",
    "        ''' \n",
    "        raise NotImplementedError \n",
    "    else:\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "2ce982f0-1f53-433e-bbd2-9c73a88786c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def myNN(X, param_dict, grad=None):\n",
    "    '''\n",
    "    Build an MLP with a single hidden layer.\n",
    "    '''\n",
    "\n",
    "    # Unpack the parameter dict\n",
    "    W1, b1 = param_dict[\"W1\"], param_dict[\"b1\"]\n",
    "    W2, b2 = param_dict[\"W2\"], param_dict[\"b2\"]\n",
    "\n",
    "    # forward pass\n",
    "    '''\n",
    "    TO DO: Fill in\n",
    "    '''\n",
    "\n",
    "    if grad is not None:\n",
    "        \n",
    "        # reverse pass \n",
    "        '''\n",
    "        TO DO: Fill in\n",
    "        '''\n",
    "        \n",
    "\n",
    "        grad_dict ={'W1':grad_W1.mean(axis=0),\n",
    "                    'b1':grad_b1.mean(axis=0),\n",
    "                    'W2':grad_W2.mean(axis=0),\n",
    "                    'b2':grad_b2.mean(axis=0)}       \n",
    "\n",
    "        return out.mean(), grad_dict\n",
    "    else:\n",
    "        return out.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "0f3b8cd7-a2b1-4f27-bbce-ea738e1d9dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "param_dict ={'W1':W1,'b1':b1,'W2':W2,'b2':b2}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "id": "ebb18297-1d3b-46db-9951-fee8f090b29d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# forward mode\n",
    "out = myNN(X, param_dict)\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "id": "94dbc5e0-d713-4f6a-bf67-4dafce0b882c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reverse mode\n",
    "dfdf = np.ones((1,1))\n",
    "out, grad_dict = myNN(X, param_dict,np.ones((1,1)))\n",
    "out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "id": "9001d687-79e4-41ec-9f69-7b609eb0a261",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k in param_dict.keys():\n",
    "    print(k,grad_dict[k].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8690fe9-2597-4df2-ab74-81197e9d643e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "15a631ce-7123-41c8-be8f-00d201acb35b",
   "metadata": {},
   "source": [
    "## Step 4: Cross check the gradients in jax\n",
    "\n",
    "`Jax` is an **automatic differentiation** library where it can automatically keep track of the gradient propagation for us instead of us needing to keep track of everything manually with these \"forward\" and reverse modes.\n",
    "\n",
    "We'll cover another auto diff library, `pytorch` in more detail tomorrow, but we'll use jax here for x-checking the gradients as it's a very transparent interface for these types of checks.\n",
    "\n",
    "Also, it's super cute b/c the code stays _almost identical_, just replace `np` with `jnp`, so we can just reuse the forward pass NN code we wrote back in step 1.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "ce259cf1-fcc5-4a84-85e4-914a3da0e7cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import grad\n",
    "from jax import numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4adbdf9-a6c5-4a25-a896-3c47b30c66c2",
   "metadata": {},
   "source": [
    "Illustrative example: Calculate the gradient of :\n",
    "\n",
    "$$g(x) = x^2, \\qquad g'(x) = 2x$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "id": "4d4aa2ee-f2f7-4cf6-a03e-63a2ee910a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "def g(xi):\n",
    "    return xi**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "id": "b1fc2082-c8a0-42ea-9c4a-8e3fc196c2ca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(6., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set up the gradient function\n",
    "grad_g = grad(g) # 2x\n",
    "\n",
    "# To evaluate the gradient, need to eval at a specific point, $x$\n",
    "grad_g(3.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "id": "a46fa584-082b-4adb-8470-35b158205b2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's evaluate the same parameters... just need to type cast to numpy\n",
    "param_dict_jax = {k:jnp.array(v) for k,v in param_dict.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "4eb3cc40-afca-4b55-abaa-6aa116afff68",
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu_jax(z):\n",
    "    return jnp.where(z>0,z,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "id": "fb565620-2c7f-49ab-ac19-5eb2ae5e20d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def linear_jax(W,X,b):\n",
    "    '''\n",
    "    Linear layer mapping H = W x + b\n",
    "    - W: array (m,n)\n",
    "    - x: array (bs, n)\n",
    "    - b: array (m)\n",
    "\n",
    "    Outputs:\n",
    "    - H: array (bs,m)\n",
    "    '''\n",
    "    \n",
    "    return jnp.einsum('mn,bn->bm',W,X) + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "id": "e93e469d-4e10-4a94-b28f-e45ceeaad333",
   "metadata": {},
   "outputs": [],
   "source": [
    "# (X, param_dict)\n",
    "def myNN_jax(param_dict):\n",
    "\n",
    "    # Unpack the parameters\n",
    "    W1 = param_dict['W1']\n",
    "    b1 = param_dict['b1']\n",
    "    \n",
    "    W2 = param_dict['W2']\n",
    "    b2 = param_dict['b2']\n",
    "\n",
    "    # forward pass\n",
    "    z1 = linear_jax(W1, X, b1)\n",
    "    h1 = relu_jax(z1)\n",
    "\n",
    "    z2 = linear_jax(W2, h1, b2)\n",
    "    \n",
    "    return z2.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "id": "e2a3138f-7645-4dcd-a8b9-1682e79e82dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_nn = grad(myNN_jax)\n",
    "\n",
    "# And... evaluate it!\n",
    "grad_dict_jax = grad_nn(param_dict_jax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "id": "6dba56d4-1c27-4015-995d-686ed5d1f977",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "W1 (16, 5) (16, 5)\n",
      "W2 (1, 16) (1, 16)\n",
      "b1 (16,) (16,)\n",
      "b2 (1,) (1,)\n"
     ]
    }
   ],
   "source": [
    "for k, v in grad_dict_jax.items():\n",
    "    print(k, v.shape, grad_dict[k].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41d47fb8-5695-489c-ac1c-e8386eb7a3b4",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "id": "bef05af7-35c9-45ff-a5fa-55be0b29a117",
   "metadata": {},
   "outputs": [],
   "source": [
    "for k, v in grad_dict_jax.items():\n",
    "    print(k, np.all(np.isclose(v, grad_dict[k])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ce891dd-bd9e-4280-af1a-49cacb077cf3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "dd303782-2064-4251-9b4d-7017543bc0ab",
   "metadata": {},
   "source": [
    "Awesome, we've played around with auto diff (AD), coded up both some intro examples and a simple NN, and cross checked these gradients with one of the autodiff packages on the market.\n",
    "\n",
    "Tomorrow, we'll stick with these AD libraries  that do these \"gutsy\" bits for us, and move on to training, optimizers and regularization tricks. But, having this detail oriented understanding is also great for building our inution for why a lot of these tricks are motivated and \"work\", and we'll highligh that tomorrow as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1c0b09c-0039-4e48-bd96-f14c28c05207",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}