jax/jax.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Intro to Jax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(9.0, DeviceArray(6., dtype=float32))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from jax import grad\n",
    "\n",
    "def f(x):\n",
    "    return x**2\n",
    "\n",
    "grad_f = grad(f)\n",
    "\n",
    "f(3.), grad_f(3.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray(4., dtype=float32), DeviceArray(1., dtype=float32))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def f(x, y):\n",
    "    return y * x**2\n",
    "\n",
    "grad_f = grad(f, argnums=(0, 1))\n",
    "\n",
    "dx, dy = grad_f(1., 2.)\n",
    "dx, dy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import hessian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[-5.905694 ,  0.       ],\n",
       "             [ 0.       , -0.7950852]], dtype=float32)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def f(x):\n",
    "    return np.sum(x**2)\n",
    "\n",
    "grad_f = grad(f)\n",
    "hess_f = hessian(f)\n",
    "\n",
    "hessian(np.array([0.1, 0.2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(2., dtype=float32)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def g(x):\n",
    "    return np.sum(x**2)\n",
    "\n",
    "grad_g = grad(g)\n",
    "\n",
    "grad_g(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.65 s, sys: 449 ms, total: 3.1 s\n",
      "Wall time: 2.28 s\n"
     ]
    }
   ],
   "source": [
    "%time _ = np.stack([grad_g(i) for i in np.linspace(0, 10, 1000)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import jit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(2., dtype=float32)"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grad_g = jit(grad(g))\n",
    "grad_g(1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 280 ms, sys: 61.8 ms, total: 342 ms\n",
      "Wall time: 238 ms\n"
     ]
    }
   ],
   "source": [
    "%time _ = np.stack([grad_g(i) for i in np.linspace(0, 10, 1000)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import vmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 22.7 ms, sys: 3.52 ms, total: 26.3 ms\n",
      "Wall time: 23.9 ms\n"
     ]
    }
   ],
   "source": [
    "%time _ = vmap(grad_g)(np.linspace(0, 10, 1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Full Train Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(params, inputs):\n",
    "    return inputs @ params \n",
    "\n",
    "def mse(params, inputs, targets):\n",
    "    preds = predict(params, inputs)\n",
    "    return np.mean((preds - targets)**2)\n",
    "\n",
    "grad_fun = jit(grad(mse))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 660.49it/s]\n"
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "from numpy import zeros\n",
    "from numpy.random import normal\n",
    "\n",
    "n, k = 10_000, 5\n",
    "X = np.concatenate([np.ones((n, 1)), normal(0, 1, (n, k))], axis=1)\n",
    "true_w = normal(0, 5, (k + 1,))\n",
    "y = X @ true_w\n",
    "W = normal(0, 1, (k + 1,))\n",
    "stepsize = 0.02\n",
    "n_step = 100\n",
    "hist_gd = zeros((n_step,))\n",
    "for i in tqdm.tqdm(range(n_step)):\n",
    "    dW = grad_fun(W, inputs=X, targets=y)\n",
    "    W -= dW*stepsize\n",
    "    hist_gd[i] = mse(W, inputs=X, targets=y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 4.44226847, -6.76762119,  6.00692163, -1.13450748, -5.7014697 ,\n",
       "        -4.67501268]),\n",
       " DeviceArray([ 4.323477 , -6.633852 ,  5.9252224, -1.1087155, -5.583535 ,\n",
       "              -4.5851116], dtype=float32))"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_w, W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABIQAAAD4CAYAAACdfT2fAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjAsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy8GearUAAAgAElEQVR4nO3dd3gc1b3/8c93i3qxbMlykdxNccFgC9MJoZqEBC5wEwgtJMQ3gYSQHnLvTW4Jv5BKQi6BEHpCDRBwAoEQQkw3lo0rtsG9W7KtZvXdPb8/diSvhNxUPCvt+/U8+8zMmTMzX9nPMtaHOWfMOScAAAAAAACkjoDfBQAAAAAAAODwIhACAAAAAABIMQRCAAAAAAAAKYZACAAAAAAAIMUQCAEAAAAAAKSYkN8FSFJhYaEbM2aM32UAAAAAAAAMGAsWLNjpnCvqal9SBEJjxoxReXm532UAAAAAAAAMGGa2YV/7GDIGAAAAAACQYgiEAAAAAAAAUgyBEAAAAAAAQIohEAIAAAAAAEgxBEIAAAAAAAAphkAIAAAAAAAgxRAIAQAAAAAApBgCoV7U0BLR7voWv8sAAAAAAADYLwKhXhKJxnTh/72hm59e4ncpAAAAAAAA+0Ug1EtCwYAumVGiF5fv0F+XbvO7HAAAAAAAgH0iEOpF1506VpNH5On7c5arpqHV73IAAAAAAAC6RCDUi0LBgH58yTHaXd+iH/11hd/lAAAAAAAAdOmAgZCZ3WdmFWa2rFP7V8xspZktN7OfJLTfbGarzWyVmZ3XF0Unsykj83XdaWP12PxNenPNTr/LAQAAAAAA+JCDeULoAUmzEhvM7KOSLpQ0zTk3WdLPvPZJki6TNNk75jdmFuzNgvuDr519hEYPydLNTy9VU2vU73IAAAAAAAA6OGAg5Jx7VdLuTs1fknSrc67Z61PhtV8o6THnXLNzbp2k1ZJm9mK9/UJGOKgfXTxVG3Y16La/v+93OQAAAAAAAB10dw6hIySdZmbzzGyumR3vtY+UtCmh32av7UPMbLaZlZtZeWVlZTfLSF4njy/Up8tKdc9r67RsS43f5QAAAAAAALTrbiAUkjRY0omSviXpCTOzQzmBc+5u51yZc66sqKiom2Ukt+997GgNzk7Td55aokg05nc5AAAAAAAAkrofCG2W9LSLe0dSTFKhpC2SShP6lXhtKSk/K6z/+eRkLd9aq3teX+d3OQAAAAAAAJK6Hwg9I+mjkmRmR0hKk7RT0hxJl5lZupmNlTRR0ju9UWh/NWvKMJ07qVi3vfS+1u+s97scAAAAAACAg3rt/KOS3pJ0pJltNrPPS7pP0jjvVfSPSbrGe1pouaQnJL0n6QVJNzjnUvo1W2am/71oitKCAd389FI55/wuCQAAAAAApDhLhoCirKzMlZeX+11Gn3pk3kZ9709L9eNLpurTx4/yuxwAAAAAADDAmdkC51xZV/u6O2QMh+iy40s1c+xg3fLcClXUNvldDgAAAAAASGEEQodJIGC69eKpaorE9IM5y/0uBwAAAAAApDACocNoXFGOvnrWRP112Xa9sGy73+UAAAAAAIAURSB0mM0+fZyOHp6n/3x2marqW/wuBwAAAAAApCACocMsHAzop5ceo+qGFv37M7x1DAAAAAAAHH4EQj6YMjJfXzvnCD2/dLueWbTF73IAAAAAAECKIRDyyb+dPl5lowv0/WeWa0t1o9/lAAAAAACAFEIg5JNgwPSLTx2rmHP6xhOLFIsxdAwAAAAAABweBEI+GjUkSz/4xGS9vXa37ntjnd/lAAAAAACAFEEg5LN/LSvROZOK9ZMXVmnV9jq/ywEAAAAAACmAQMhnZqYfXTxVeZkh3fT4IjVHon6XBAAAAAAABjgCoSRQmJOuWy8+Riu21eq2lz7wuxwAAAAAADDAEQglibMnFevymaX67atr9M663X6XAwAAAAAABjACoSTyHx+fpNKCLH39iUWqa2r1uxwAAAAAADBAEQglkez0kG779DRtrW7U//z5Pb/LAQAAAAAAAxSBUJKZMXqwrj9jgv64YLNeWLbd73IAAAAAAMAAdMBAyMzuM7MKM1vWxb5vmJkzs0Jv28zsdjNbbWZLzGx6XxQ90N141kRNGZmn7/1pqSrqmvwuBwAAAAAADDAH84TQA5JmdW40s1JJ50ramNB8vqSJ3me2pDt7XmLqSQsF9MtPH6v65oi++9RSOef8LgkAAAAAAAwgBwyEnHOvSurqtVe3Sfq2pMS04kJJD7m4tyUNMrPhvVJpipkwNFc3n3+U/rGyQve/sd7vcgAAAAAAwADSrTmEzOxCSVucc4s77RopaVPC9mavratzzDazcjMrr6ys7E4ZA941J4/R2UcX60d/XaFFm6r9LgcAAAAAAAwQhxwImVmWpO9J+n5PLuycu9s5V+acKysqKurJqQYsM9PP/3WahuZm6IaHF6q6ocXvkgAAAAAAwADQnSeExksaK2mxma2XVCJpoZkNk7RFUmlC3xKvDd2UnxXWHVdMV0Vdk775xyXMJwQAAAAAAHrskAMh59xS59xQ59wY59wYxYeFTXfObZc0R9LV3tvGTpRU45zb1rslp55jSwfp5vOP1t9X7NA9r63zuxwAAAAAANDPHcxr5x+V9JakI81ss5l9fj/dn5e0VtJqSb+TdH2vVAlde8oYzZo8TD9+YaUWbKjyuxwAAAAAANCPWTIMQSorK3Pl5eV+l5H0ahpbdcGvX1M06vTcjaepIDvN75IAAAAAAECSMrMFzrmyrvZ16y1j8Ed+Zli/+cwM7dzTom/8cbFiMf/DPAAAAAAA0P8QCPUzU0vy9R8XHK1/rKzQ3a+t9bscAAAAAADQDxEI9UNXnThaH586XD99cZXmr9/tdzkAAAAAAKCfIRDqh8xMt14yVaUFmfrKI+9q155mv0sCAAAAAAD9CIFQP5WbEdYdV0zX7oYWfe0J5hMCAAAAAAAHj0CoH5s8Il8/+MQkvfp+pe6cu8bvcgAAAAAAQD9BINTPfWbmKH1y2gj9/G+r9OaanX6XAwAAAAAA+gECoX7OzPT/Lp6qcUU5uuHhhdq0u8HvkgAAAAAAQJIjEBoActJD+t3VZYrGnL7wULnqmyN+lwQAAAAAAJIYgdAAMbYwW//3mel6f0edvsEk0wAAAAAAYD8IhAaQ048o0vc+drReWL5dt//jA7/LAQAAAAAASSrkdwHoXZ8/daxWbKvTL//+gY4alqtZU4b7XRIAAAAAAEgyPCE0wJiZbvmXKTq2dJC+/sRirdxe63dJAAAAAAAgyRAIDUAZ4aDuvmqGcjNCuu7Bcu2ub/G7JAAAAAAAkEQIhAaooXkZuvuqMlXUNev6hxeoNRrzuyQAAAAAAJAkCIQGsGmlg/TjS6bq7bW79b9/ec/vcgAAAAAAQJI4YCBkZveZWYWZLUto+6mZrTSzJWb2JzMblLDvZjNbbWarzOy8viocB+dfjivR7NPH6aG3NujRdzb6XQ4AAAAAAEgCB/OE0AOSZnVqe0nSFOfcMZLel3SzJJnZJEmXSZrsHfMbMwv2WrXolu/MOkofOaJI3392md5Zt9vvcgAAAAAAgM8OGAg5516VtLtT29+ccxFv821JJd76hZIec841O+fWSVotaWYv1otuCAZMt19+nEoLsvSlPyzQ5qoGv0sCAAAAAAA+6o05hD4n6a/e+khJmxL2bfbaPsTMZptZuZmVV1ZW9kIZ2J/8zLB+d02ZWiIxXfdguWqbWv0uCQAAAAAA+KRHgZCZ/bukiKSHD/VY59zdzrky51xZUVFRT8rAQRpflKM7r5yh1RV79MXfL1BLhDePAQAAAACQirodCJnZZyVdIOkK55zzmrdIKk3oVuK1IUmcOrFQP7n0GL25Zpe+/eRixWLuwAcBAAAAAIABpVuBkJnNkvRtSZ90ziVOSDNH0mVmlm5mYyVNlPROz8tEb7p4eom+dd6RembRVv3kxVV+lwMAAAAAAA6z0IE6mNmjks6QVGhmmyX9QPG3iqVLesnMJOlt59wXnXPLzewJSe8pPpTsBudctK+KR/ddf8Z4ba1u1F1z12jEoAxdfdIYv0sCAAAAAACHyQEDIefc5V0037uf/rdIuqUnRaHvmZn+58Ip2lHbrB/MWa7ivAydN3mY32UBAAAAAIDDoDfeMoZ+Khgw/fry4zStZJBufPRdLdhQ5XdJAAAAAADgMCAQSnGZaUHde02Zhudn6LoH52tt5R6/SwIAAAAAAH2MQAgakpOuBz83UwEzXXP/O6qoa/K7JAAAAAAA0IcIhCBJGj0kW/d99njtrGvR5x8oV31zxO+SAAAAAABAHyEQQrtppYN0xxXH6b1ttbrhkYVqjcb8LgkAAAAAAPQBAiF0cOZRxfrhRVP0z1WV+vc/LZVzzu+SAAAAAABALzvga+eRei6fOUrba5r0q5c/UHZ6SN+/YJLMzO+yAAAAAABALyEQQpduOnui6poiuu+NdcpKC+pb5x3ld0kAAAAAAKCXEAihS2am/7zgaDW2RnXHK2uUGQ7qy2dO9LssAAAAAADQCwiEsE9mplsumqLm1qh+9rf3lREO6rrTxvldFgAAAAAA6CECIexXIGD6yaXHqCkS1Q+fW6HMtKCuOGG032UBAAAAAIAeIBDCAYWCAf3y08epqXWB/uOZZcoIBXXJjBK/ywIAAAAAAN3Ea+dxUNJCAf3miuk6efwQfevJxXpuyTa/SwIAAAAAAN1EIISDlhEO6ndXl2n6qAJ99bF39fKKHX6XBAAAAAAAuoFACIckKy2k+649XpNG5OlLf1io1z/Y6XdJAAAAAADgEBEI4ZDlZYT14LUzNa4oW194qFzz1+/2uyQAAAAAAHAIDhgImdl9ZlZhZssS2gab2Utm9oG3LPDazcxuN7PVZrbEzKb3ZfHwT0F2mn7/+RM0fFCGrr1/vsoJhQAAAAAA6DcO5gmhByTN6tT2XUkvO+cmSnrZ25ak8yVN9D6zJd3ZO2UiGRXlpuvh605QUW66rrr3Hb25muFjAAAAAAD0BwcMhJxzr0rq/PjHhZIe9NYflHRRQvtDLu5tSYPMbHhvFYvkMzw/U4//24kqHZypax+Yr1dWVfhdEgAAAAAAOIDuziFU7Jxre+/4dknF3vpISZsS+m322j7EzGabWbmZlVdWVnazDCSDobkZemz2SRpflKPZD5XrxeXb/S4JAAAAAADsR48nlXbOOUmuG8fd7Zwrc86VFRUV9bQM+Gxwdpoe/cKJmjwiX9c/vFB/XrzV75IAAAAAAMA+dDcQ2tE2FMxbto0T2iKpNKFfideGFJCfFdYfrjtBM0YX6KuPvas/lm868EEAAAAAAOCw624gNEfSNd76NZKeTWi/2nvb2ImSahKGliEF5KSH9OC1M3XKhEJ968kl+v3bG/wuCQAAAAAAdHIwr51/VNJbko40s81m9nlJt0o6x8w+kHS2ty1Jz0taK2m1pN9Jur5PqkZSy0wL6ndXl+mso4bqP59ZpnteW+t3SQAAAAAAIEHoQB2cc5fvY9dZXfR1km7oaVHo/zLCQd155Qzd9Pi7+uFzK9TUGtWXz5zod1kAAAAAAEAHEQgB3ZUWCuj2y45TemiJfva399XYGtU3zz1SZuZ3aQAAAAAApDQCIfSpUDCgn//rNGWEA7rjlTWqbYzovz45WcEAoRAAAAAAAH4hEEKfCwRM/+9fpiovI6zfvrpWO2qbdPvlxykjHPS7NAAAAAAAUlJ33zIGHBIz080fO1o/+MQkvbRih664Z56q6lv8LgsAAAAAgJREIITD6tpTxuo3n5mupVtqdMldb2rT7ga/SwIAAAAAIOUQCOGwO3/qcP3h8ydoZ12zLr7zTS3bUuN3SQAAAAAApBQCIfhi5tjBeupLJyscMH36t2/p1fcr/S4JAAAAAICUQSAE30wsztXT15+i0sFZ+twD8/XUgs1+lwQAAAAAQEogEIKvhuVn6IkvnqQTxg3WN/64WHe8slrOOb/LAgAAAABgQCMQgu/yMsK6/7MzddGxI/TTF1fpP59dpmiMUAgAAAAAgL4S8rsAQJLSQgH94lPHqjg/Q7+du1Zbq5v0q8uOVW5G2O/SAAAAAAAYcHhCCEkjEDDdfP7R+t+Lpmju+5X6l9+8qfU76/0uCwAAAACAAYdACEnnqhNH6/efm6mde5p14R1v6I3VO/0uCQAAAACAAYVACEnp5AmFmnPDqRqWl6Gr73tHD7yxjsmmAQAAAADoJQRCSFqjhmTpqetP1kePHKr/+vN7uvnppWqJxPwuCwAAAACAfo9ACEktJz2ku6+aoa+cOUGPzd+kz/zube3c0+x3WQAAAAAA9Gs9CoTM7GtmttzMlpnZo2aWYWZjzWyema02s8fNLK23ikVqCgRM3zj3SP368uO0bGuNPvnr17V8a43fZQEAAAAA0G91OxAys5GSbpRU5pybIiko6TJJP5Z0m3NugqQqSZ/vjUKBT0wboSe/eLIk6dI739JzS7b5XBEAAAAAAP1TT4eMhSRlmllIUpakbZLOlPSkt/9BSRf18BpAuykj8/Xsl0/VpBF5uuGRhfr531YpGmOyaQAAAAAADkW3AyHn3BZJP5O0UfEgqEbSAknVzrmI122zpJFdHW9ms82s3MzKKysru1sGUlBRbroe+cIJ+nRZqX79j9W6+r55qqxjXiEAAAAAAA5WT4aMFUi6UNJYSSMkZUuadbDHO+fuds6VOefKioqKulsGUlR6KKgfX3qMfnLpMSpfX6WP3f6a3l67y++yAAAAAADoF3oyZOxsSeucc5XOuVZJT0s6RdIgbwiZJJVI2tLDGoF9+lRZqZ798inKzQjpM797W3e8sloxhpABAAAAALBfPQmENko60cyyzMwknSXpPUmvSLrU63ONpGd7ViKwf0cNy9OcL5+qjx8zQj99cZU+9+B87a5v8bssAAAAAACSVk/mEJqn+OTRCyUt9c51t6TvSPq6ma2WNETSvb1QJ7BfOekh3X7ZsfrhRVP05upd+vjtr2nBht1+lwUAAAAAQFIy5/wfXlNWVubKy8v9LgMDxLItNbr+4YXaWt2o78w6StedNlbxh9gAAAAAAEgdZrbAOVfW1b6evnYeSDpTRubrz185VWcdPVS3PL9Cs3+/QDUNrX6XBQAAAABA0iAQwoCUnxnWXVfO0PcvmKRXVlbo479+TeXrGUIGAAAAAIBEIIQBzMz0uVPH6o9fPElm0qd++5Z++uJKtURifpcGAAAAAICvCIQw4B03qkDP33iaLp1RojteWaOL73xDqyvq/C4LAAAAAADfEAghJeRmhPWTS6fpritnaGt1kz5+++t64I11isX8n1QdAAAAAIDDjUAIKWXWlGF64abTdPL4IfqvP7+na+5/R9trmvwuCwAAAACAw4pACClnaG6G7vvs8frhRVNUvr5K5/3yVT23ZJvfZQEAAAAAcNgQCCElmZmuPHG0nrvxVI0pzNYNjyzU1x5fpNomXk8PAAAAABj4CISQ0sYV5ejJL56km86eqDmLt+r8X76m1z/Y6XdZAAAAAAD0KQIhpLxwMKCbzj5CT37xJKWHArry3nn61h8Xq6aBp4UAAAAAAAMTgRDgOW5UgZ7/6mm6/ozxevrdLTrrF3P1/NJtco43kQEAAAAABhYCISBBRjiob886SnO+fIqG5afr+ocX6t9+v0A7ankTGQAAAABg4CAQAroweUS+nrn+FN18/lGa+36lzv7FXD36zkbFYjwtBAAAAADo/wiEgH0IBQP6t4+M14s3na4pI/J189NL9Zl73tb6nfV+lwYAAAAAQI8QCAEHMKYwW4984QTdevFULd9aq/N++arumrtGkWjM79IAAAAAAOgWAiHgIJiZLps5Sn//+kd0xpFFuvWvK/WJ/3tD89fv9rs0AAAAAAAOWY8CITMbZGZPmtlKM1thZieZ2WAze8nMPvCWBb1VLOC34rwM/faqMt115QzVNrbqX+96S197fJEqmHQaAAAAANCP9PQJoV9JesE5d5SkaZJWSPqupJedcxMlvextAwPKrCnD9Pevf0RfOXOCnluyTWf+fK7ueW2tWhlGBgAAAADoB8y57r01yczyJS2SNM4lnMTMVkk6wzm3zcyGS/qnc+7I/Z2rrKzMlZeXd6sOwG/rd9brv/+8XK+sqtSEoTn6709O1ikTCv0uCwAAAACQ4sxsgXOurKt9PXlCaKykSkn3m9m7ZnaPmWVLKnbObfP6bJdU3INrAElvTGG27r92pu69pkwtkZiuuGeebnh4obZWN/pdGgAAAAAAXepJIBSSNF3Snc654yTVq9PwMO/JoS4fQTKz2WZWbmbllZWVPSgDSA5nHV2sv33tdH39nCP09xU7dNbP5+qOV1arORL1uzQAAAAAADroSSC0WdJm59w8b/tJxQOiHd5QMXnLiq4Ods7d7Zwrc86VFRUV9aAMIHlkhIO68ayJ+vvXP6LTjyjUT19cpXNve1XPL92m7g7PBAAAAACgt3U7EHLObZe0ycza5gc6S9J7kuZIusZru0bSsz2qEOiHSgdn6bdXlemhz81URiio6x9eqIvvfJPX1AMAAAAAkkK3J5WWJDM7VtI9ktIkrZV0reIh0xOSRknaIOlTzrn9/hbMpNIYyKIxp6cWbNbPX1qlHbXNOndSsb5z/lEaX5Tjd2kAAAAAgAFsf5NK9ygQ6i0EQkgFDS0R3ff6Ot01d60aW6O6fGapvnrWESrKTfe7NAAAAADAAEQgBCSRnXuadfvLH+iReRuVHgpo9unj9YXTxyorLeR3aQAAAACAAYRACEhCayv36CcvrNILy7draG66vnbOEbp0RonCwZ7M9Q4AAAAAQNz+AiF+8wR8Mq4oR3ddNUNPfvEklRRk6uanl+rsX8zVUws2KxKN+V0eAAAAAGAAIxACfFY2ZrCe+tLJ+t3VZcpOC+kbf1ysc257Vc+8u0XRmP9P8AEAAAAABh4CISAJmJnOmVSsv3zlVN115QylhwK66fFFOve2uZqzeKtiBEMAAAAAgF5EIAQkkUDANGvKMD1/42n6zRXTFTDTjY++q1m/elXPLdlGMAQAAAAA6BUEQkASCgRMH5s6XC/cdLpuv/w4RWJONzyyUB+7/TW9sGw7wRAAAAAAoEd4yxjQD0RjTn9evFW/evkDrdtZryOKc/TFj4zXJ6aN4K1kAAAAAIAu8dp5YICIRGOas3ir7pq7Ru/v2KORgzJ13Wlj9enjS5WVFvK7PAAAAABAEiEQAgaYWMzplVUVuvOfa1S+oUoFWWF99uSxuvqk0SrITvO7PAAAAABAEiAQAgaw+et3665/rtHLKyuUGQ7q8pmjdN1pYzViUKbfpQEAAAAAfEQgBKSAldtr9du5azVn8VaZpAuPHanZp4/TkcNy/S4NAAAAAOADAiEghWza3aB7X1+nx+ZvVFNrTKdMGKLPnjxWZx41VMGA+V0eAAAAAOAwIRACUtDu+hY9+s5G/f6tDdpe26RRg7N09Umj9anjS5WXEfa7PAAAAABAHyMQAlJYazSmF5dv1wNvrFf5hiplpQV16YwSXXPyGI0vyvG7PAAAAABAHyEQAiBJWrq5Rve/uU5/WbxNLdGYPnJEka49ZYxOn1ikAMPJAAAAAGBA6dNAyMyCksolbXHOXWBmYyU9JmmIpAWSrnLOtezvHARCwOFVWdesR+Zt1B/mbVBlXbPGFmbr8pmlumR6iYbkpPtdHgAAAACgF/R1IPR1SWWS8rxA6AlJTzvnHjOzuyQtds7dub9zEAgB/miJxPT80m16eN4GzV9fpXDQdN7kYfrMzFE6cdwQnhoCAAAAgH6szwIhMyuR9KCkWyR9XdInJFVKGuaci5jZSZL+yzl33v7OQyAE+O+DHXV69J1NemrhZtU0tmrMkCxdNnOULp1RokKeGgIAAACAfqcvA6EnJf1IUq6kb0r6rKS3nXMTvP2lkv7qnJuyv/MQCAHJo6k1qheWbdcj72zUO+t2Kxw0nTOpWJfPHKVTxhfy1BAAAAAA9BP7C4RCPTjpBZIqnHMLzOyMbhw/W9JsSRo1alR3ywDQyzLCQV103EhddNxIra7Yo8fe2ainFm7W80u3q3Rwpi6ZXqKLjyvRqCFZfpcKAAAAAOimbj8hZGY/knSVpIikDEl5kv4k6TwxZAwYUJojUb24fIcen79Rb67ZJeek48cU6OLpJfrY1OHKzwz7XSIAAAAAoJM+f+2894TQN71Jpf8o6amESaWXOOd+s7/jCYSA/mNrdaOeWbRFTy/cotUVe5QWCuicScW6ZPpInTaxSOFgwO8SAQAAAAA6/IHQOMVfOz9Y0ruSrnTONe/veAIhoP9xzmnplho9vXCLnl20RVUNrSrMSdMnp43UxdNHavKIPJkx3xAAAAAA+KXPA6GeIhAC+reWSExz36/U0ws36+UVFWqJxjS+KFsfP2aEPnHMcE0szvW7RAAAAABIOQRCAA6b6oYW/WXJNv1lyVbNW7dbzklHFufqgmOG6+PHDNe4ohy/SwQAAACAlEAgBMAXFbVNen7pNv1lyTaVb6iSJE0anqcLpg3XBVNH8KYyAAAAAOhDBEIAfLetplHPLYmHQ4s2VUuSppXk6/ypw3XupGKeHAIAAACAXkYgBCCpbNrdoOeWxoeVLdtSK0maODRH504u1rmThmnqyHwFAkxIDQAAAAA9QSAEIGltrmrQS+/t0N+W79A763crGnMalpehcyYV69zJxTph7BClhXiVPQAAAAAcKgIhAP1CVX2L/rGyQn97b7vmvl+pptaYcjNCOvOooTpnUrFOm1ik/Myw32UCAAAAQL9AIASg32lsier11Tv1t+Xb9fcVO1TV0KpgwDRjdIE+euRQffSoIh1ZnCszhpYBAAAAQFcIhAD0a5FoTIs2VeuVVRV6ZWWl3tsWn3doeH6GzjhyqD56ZJFOmVCo7PSQz5UCAAAAQPIgEAIwoGyvadLc9+Ph0Ourd2pPc0ThoGnm2MH66JFDderEQp4eAgAAAJDyCIQADFgtkZjKN+zW3FWVemVVhd7fsUeSVJiTrlMnDNEpEwp16sRCDc/P9LlSAAAAADi8CIQApIwt1Y16Y/XO9s/OPS2SpPFF2Tp1QqFOmR2/xE0AABK7SURBVFCoE8cPUV4Gk1MDAAAAGNgIhACkpFjMadWOOr2xeqdeX71T89buVmNrVMGAaVpJvk4cN0QnjBuistEFzD8EAAAAYMAhEAIASc2RqN7dWN0eEC3ZXKNozCkYME0Zma8Txw7WCeMGq2zMYJ4gAgAAANDvEQgBQBfqmyNasKFK89bt0jvrdmvRpmq1Rp0CJk0akacTxg7RCWMH6/gxg1WQneZ3uQAAAABwSAiEAOAgNLVGtXBjleat3a1563Zp4cZqtURikuJzEM0YXdD+GVeYo0CAt5gBAAAASF77C4SYNAMAPBnhoE4eX6iTxxdKig8xW7ypRuUbdmvhhiq99N4OPVG+WZI0KCus6aPi4dD0UQWaVpqvrDT+kwoAAACgf+C3FwDYh/RQUDPHDtbMsYMlSc45rd1ZrwUbqrRwQ5UWbKjSP1ZWSJKCAdPRw3N1TMkgHVsySNNKB2nC0BwFeYoIAAAAQBLq9pAxMyuV9JCkYklO0t3OuV+Z2WBJj0saI2m9pE8556r2dy6GjAHor2oaWrVwU5UWrK/Sok3VWry5WnVNEUlSVlpQU0bm69jSQTqmJF/TSgappCBTZoREAAAAAPpen8whZGbDJQ13zi00s1xJCyRdJOmzknY75241s+9KKnDOfWd/5yIQAjBQxGJO63bVa/Gmai3ZXKNFm6r13tZatUTjcxENzk7TMSX5mjoyX5NH5GnyiHxCIgAAAAB9ok/mEHLObZO0zVuvM7MVkkZKulDSGV63ByX9U9J+AyEAGCgCAdP4ohyNL8rRxdNLJEktkZhWba/Tos3VXlBUrVffr1TMy+PzM8OaPCJPUxJCorGF2Qw3AwAAANBneuUtY2Y2RtKrkqZI2uicG+S1m6Sqtu1Ox8yWNFuSRo0aNWPDhg09rgMA+oum1qhWbq/Tsi01Wr61Vsu31mjl9rr2t5plpQV19PA8HT08V0cNiy+PKM5VbkbY58oBAAAA9Bd9+tp5M8uRNFfSLc65p82sOjEAMrMq51zB/s7BkDEAkFqjMa2u2KPlW2u9oCgeErXNSSRJJQWZ7QHRkcPiYdGYIVkKBQM+Vg4AAAAgGfXZa+fNLCzpKUkPO+ee9pp3mNlw59w2b56hip5cAwBSRTgY8J4KytOlM+LDzZxz2lrTpJXbarVye138s61Wr6yqUNQbc5YeCmjC0BxNHJqjicW57eujhzDsDAAAAEDXuh0IecPB7pW0wjn3i4RdcyRdI+lWb/lsjyoEgBRmZho5KFMjB2XqrKOL29ubWqNaU7lHK7fVaeX2Wr2/Y4/mr6/SM4u2tvdJCwU0rjBbE4tz42HR0BxN8IKitBBPFAEAAACprCdvGTtV0muSlkqKec3fkzRP0hOSRknaoPhr53fv71wMGQOA3rGnOaI1FXv0/o46ra7Yow8q9uiDijpt2t3Y3idgUungLI0rzNa4ohyNK8rWuMIcjS/KVlFuOm88AwAAAAaIvnrL2OuS9vVbw1ndPS8AoPty0kOaVjpI00o7zuXf0BLRmop6ra6s09rKeq2trNeayj16a+0uNbXG2vvlpoc0tihb4wqzNXpItsYUZsWXQ7JVkBUmLAIAAAAGiB7NIQQA6B+y0kKaWpKvqSX5HdpjMaetNY1eSLRH63bWa+3Oes1fX6VnF29V4kOkuRkhjRmSrdFDstqXo4dka9TgLA3NTVeA+YoAAACAfoNACABSWCBgKinIUklBlk4/oqjDvuZIVJt2N2rDrnqt39XQvly6pUZ/Xba9fVJrKT5fUUlBpkoLslQ6uG2ZpVGDs1RakKX8rPDh/tEAAAAA7AeBEACgS+mhoCZ4E1F31hqNaWt1o9btrNemqkZt3t2gTVUN2ri7QYs2VaumsbVD/9yMkEoKsrwJsjM0siBTI7zJskcOylRhDk8YAQAAAIcTgRAA4JCFgwGNHhKfZ6grtU2t2rS7wfs0alNVg7ZUNWpzVYPmrd2luuZIh/5pwYCGD8rQyEHxoGh4foaG58eXw/IzNCI/U3mZIeYwAgAAAHoJgRAAoNflZYQ1eUS+Jo/I73J/bVOrtlQ1amt1o7ZUN2pLlbesbtRrH1Sqsq5ZsU4vwcwMB+NB0aAMDcuLh0XFeekampeh4rz4emFOusLBwGH4CQEAAID+jUAIAHDY5WWElTc8rKOH53W5PxKNqaKuWdtqmrS9pknbaho7rL+5Zqd21DZ9KDQykwpz0lWcl67i3AwvLEpXUW66inLiy6F5GSrMSVN6KHgYflIAAAAgOREIAQCSTigY0Ahv+Ni+RGNOu+qbVVHbrO01TdpR16Qdtc2qqG3Sjtombatp0uLN1dq5p6XL4/Mzwx2CoqLc+BNGQ3LSVOQth+Ska0h2mjLChEcAAAAYWAiEAAD9UjBgGpqboaG5GZoysuuhaVJ8Auxde1pUWdesyj1N8WXbZ098uXhztSrrmtXQEu3yHDnpIQ3JSYsHRtlpGpKTpsHZaSrIiq8XZKVpSHa6CrLDGpKdrsw0AiQAAAAkNwIhAMCAFg4GNMybnFrad3AkSQ0tEe3a06Kde5q1a0+LdtU3a2en7Q27GrRwY5WqGloV7TxmzZMRDrQHRAVZaRqUlaaCrHDCsm19b3teBpNmAwAA4PAhEAIAwJOVFlLW4JBKB2cdsG8s5lTXFNGu+mZVNbRo156W+LK+RVX1e5dVDfE3rlU1tKq2qVWu6wxJAYsPY8vPDCs/K02DvPVBWeG97d4nLzOsvIyw8rPCyssIKTstpECAMAkAAAAHj0AIAIBuCARM+VnxUOZgRWNONY2tqmpoUXVDi6rq29ZbVdPYqurGFtU0RlTt7V+/q141jfF9+wqSpHiYlJsRVl5mKD5hd0ZYuRkh5XrLvIT1nC7ac9JDykoL8oQSAABACiEQAgDgMAkGTIOz4/MPHYpYzKmuOaIa7ymj2sa2ZSRhO6Kaxr37NuxqUF1Tq+qaItrTEtlvoCTF39CWkxYPjHLSE5beJ7vDMqistL1tWenB9n3ZaUFlp4cUDgZ68CcFAACAvkYgBABAkgsErH24WHfEYk71LRHVNbV94kFRrbesb45oT3Okw3rb9vaapvj2QQZLbcJBi4dGaUFleU8gZaUFlZ0WUmbCsq09My3eJzMc7Nge9trTgsoIx/eHg8bTTAAAAD1EIAQAwAAXCJg3TKx7gVIb55waW6Oqb462B0cNLYnrEe1pjqqhOaL6lqgaW+LLhpZ4v4bmqLbXNqmxJar6logamqNqaI3uc3LufQkGTJlhLyBKCygjtDcwiodGgfh6KKiMcEAZaW3rCfvCQaWHAu3L9E7bGeGg0sPxcxNAAQCAgYhACAAAHBSz+FM/WWkhFeWm98o5nXNqjTo1tkTV0BoPjhpbovFlazxUamiJqqk1psbWqJpa4/sbW+OfpoT1hpaoahpatKM1pqZIvG9Ta0xNrVE1R2I9qjM9FOgQHKWFAkoPBdvb09r2h4JKCwWUFtzblpb4CXZsCwf39k08rvO+cDCgcNDa25hEHAAA9BSBEAAA8I2ZKS1kSgsFlK+ePcG0P7GYU3MkHg7Fw6KYmiMdA6POy2YvgGqJxOLb7Z+2/fH1lkhMtU0RtURiaolE1RKNeevx/i2RmCKH+BTUgYQC1h4StQVGoYTA6MPbplDbMhDfl+b1CQUS9wcUDsTXQwGL709oSzw+FNi7L9j5mE59Ql6fcCCgoNceDMSXPH0FAIA/+iwQMrNZkn4lKSjpHufcrX11LQAAgP0JBEyZ3lxEfojGnFqj8RCpJRprD41aox2Do7btluje9dZo/NjWaLxf23pr1LX3jUQ79ktcb2qNqq4pppaoUyQaD6fiIVVMkajzjneKxOLHHW7BhHBo73JviJS4b1/bQe+YoCm+DEihQPxJqsS2YMAUsIRjvPWAtx5ob1PC+t5P27EBU4fzmCUcZwn9vL5t5w54ffb2N5ntrSv+2Xtu62rd62ft14ovre063jpBGwDgQPokEDKzoKQ7JJ0jabOk+WY2xzn3Xl9cDwAAIJnFg4P43EXJzDmnSMy1B0SRqFOrt0xcb/WCpWjbvljbcbEPHR+NxY+Leu1Rr2/UC6A6b8e8GqLeedu3E87X3h51ijqn1taYorH4fFTtH+c6bnttsS72xVzb0u+/gd7TFiAFvHAoYHvDqsR98e3E/fH+H+6z9zymvfsT+ymxvxKP63jtvefZuy61Xbvj9RL7tK23L722xOskniexzrY/E0vc751Lndu199xqv97enysxcOvcP3FbHepVh7q9y3Y4f4frJrTpQz+7Es7R8Vpt5/1QHdr757P33B1rVBfnVofjrdO5PnyNthVLuM6Hju/Uv3NdHfdZF30/fH7tY3+HtoS+XR3f+RoHqimxrfO19lVDl7V22mld9DmYc9o+NqxTlV39zHv3df2z7Stf7nzuLk+6j+O7OmVXQXbX/Q6yln30Pdh+4cDAH6LdV08IzZS02jm3VpLM7DFJF0oiEAIAAEhSZubNVSTFH/JOLc7FQ6FohyBqb4jUvnTxYYht/eKBkjpte/1cQr+YvOPj52q7lnN7z+s6HRtLXO+07drbO/aNOicltru2n63jsdGYk1PC/lhCf+09f7xfvI/rcG2vTR2v4zrXF5Oiin2obtfhOnuPTawpvr3376atv0u4dvv5vHUlHp+wX965Ol/bO6TjddvPBSBVPfWlkzVjdIHfZfSpvgqERkralLC9WdIJiR3MbLak2d7mHjNb1Ue1HG6Fknb6XQTQj/CdAQ4N3xng0PCdAQ4N3xlAUtmPD7prsn9nRu9rh2+TSjvn7pZ0t1/X7ytmVu6cK/O7DqC/4DsDHBq+M8Ch4TsDHBq+M8Ch6c/fmUAfnXeLpNKE7RKvDQAAAAAAAD7rq0BovqSJZjbWzNIkXSZpTh9dCwAAAAAAAIegT4aMOeciZvZlSS8qPiPhfc655X1xrSQ04IbBAX2M7wxwaPjOAIeG7wxwaPjOAIem335nzDF9PgAAAAAAQErpqyFjAAAAAAAASFIEQgAAAAAAACmGQKiXmNksM1tlZqvN7Lt+1wMkGzMrNbNXzOw9M1tuZl/12geb2Utm9oG3LPC7ViCZmFnQzN41s79422PNbJ53v3nce3kDAElmNsjMnjSzlWa2wsxO4j4D7JuZfc37d9kyM3vUzDK4zwB7mdl9ZlZhZssS2rq8r1jc7d53Z4mZTfev8oNDINQLzCwo6Q5J50uaJOlyM5vkb1VA0olI+oZzbpKkEyXd4H1PvivpZefcREkve9sA9vqqpBUJ2z+WdJtzboKkKkmf96UqIDn9StILzrmjJE1T/LvDfQbogpmNlHSjpDLn3BTFXwZ0mbjPAIkekDSrU9u+7ivnS5rofWZLuvMw1dhtBEK9Y6ak1c65tc65FkmPSbrQ55qApOKc2+acW+it1yn+j/SRin9XHvS6PSjpIn8qBJKPmZVI+rike7xtk3SmpCe9LnxnAI+Z5Us6XdK9kuSca3HOVYv7DLA/IUmZZhaSlCVpm7jPAO2cc69K2t2peV/3lQslPeTi3pY0yMyGH55Ku4dAqHeMlLQpYXuz1wagC2Y2RtJxkuZJKnbObfN2bZdU7FNZQDL6paRvS4p520MkVTvnIt429xtgr7GSKiXd7w2zvMfMssV9BuiSc26LpJ9J2qh4EFQjaYG4zwAHsq/7Sr/LBQiEABxWZpYj6SlJNznnahP3OeecJOdLYUCSMbMLJFU45xb4XQvQT4QkTZd0p3PuOEn16jQ8jPsMsJc378mFioepIyRl68NDYwDsR3+/rxAI9Y4tkkoTtku8NgAJzCyseBj0sHPuaa95R9ujlN6ywq/6gCRziqRPmtl6xYcin6n4/CiDvEf7Je43QKLNkjY75+Z5208qHhBxnwG6drakdc65Sudcq6SnFb/3cJ8B9m9f95V+lwsQCPWO+ZImejPypyk+Gdscn2sCkoo398m9klY4536RsGuOpGu89WskPXu4awOSkXPuZudciXNujOL3lX84566Q9IqkS71ufGcAj3Nuu6RNZnak13SWpPfEfQbYl42STjSzLO/faW3fGe4zwP7t674yR9LV3tvGTpRUkzC0LClZ/Akn9JSZfUzxuR6Cku5zzt3ic0lAUjGzUyW9Jmmp9s6H8j3F5xF6QtIoSRskfco513niNiClmdkZkr7pnLvAzMYp/sTQYEnvSrrSOdfsZ31AsjCzYxWfhD1N0lpJ1yr+P0C5zwBdMLP/lvRpxd8G+66k6xSf84T7DCDJzB6VdIakQkk7JP1A0jPq4r7iBav/p/jQywZJ1zrnyv2o+2ARCAEAAAAAAKQYhowBAAAAAACkGAIhAAAAAACAFEMgBAAAAAAAkGIIhAAAAAAAAFIMgRAAAAAAAECKIRACAAAAAABIMQRCAAAAAAAAKeb/A77x7SAD470RAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 1440x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pylab as plt \n",
    "\n",
    "plt.figure(figsize=(20, 4))\n",
    "plt.plot(hist_gd)\n",
    "plt.ylim(0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Closing Notes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(-1., dtype=float32)"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def abs_val(x):\n",
    "    if x > 0:\n",
    "        return x\n",
    "    return -x\n",
    "\n",
    "grad(abs_val)(-2.0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}