ml_workshop-01-06-2020.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# ML workshop\n",
    "01/06/2020"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "## Bootstrap\n",
    "Before we get started, ensure you have your machine setup. Here I am using Python 3 with `pip` (not conda):\n",
    "\n",
    "```sh\n",
    "python3 -m venv ml-env\n",
    ". ml-env/bin/activate  # Windows: > ml-env\\Scripts\\activate\n",
    "pip install -U pip\n",
    "pip install -U setuptools wheel\n",
    "pip install jupyter\n",
    "jupyter notebook  # Opens an interactive web interface; then you should File⇒New Notebook.\n",
    "# Feel free to run `ipython`—or `python`—instead, and work just on the CLI.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "Alternatively, follow one of these links to run it in a cloud environment:\n",
    "\n",
    "| Google | Microsoft | Open-source |\n",
    "|--------|-----------|-------------|\n",
    "| [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/offscale/offscale-presentations/blob/master/ml_workshop-01-06-2020.ipynb) | [![Azure Notebooks](https://notebooks.azure.com/launch.svg)](https://docs.microsoft.com/en-us/azure/notebooks/quickstart-clone-jupyter-notebook) | [![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/offscale/offscale-presentations/master?filepath=ml_workshop-01-06-2020.ipynb) |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install jupyter matplotlib tensorflow tensorflow_datasets jax jaxlib torch torchvision scikit-learn\n",
    "# Some of the hosted Jupyter Notebook don't auto find and install dependencies via imports, so we do it explicitly^\n",
    "# Ensure you uncomment that top line in those environments, otherwise you'll get `ImportError`s."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## What is machine-learning and artificial-intelligence?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What's a neural network (NN), and what's a convolutional neural network (CNN)?\n",
    "\n",
    "### Convolution\n",
    "> A convolution is an integral that expresses the amount of overlap of one function $g$ as it is shifted over another function $f$. \n",
    "![Convolution visualisation](https://mathworld.wolfram.com/images/gifs/convgaus.gif)\n",
    "https://mathworld.wolfram.com/Convolution.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"475px\"\n",
       "            src=\"https://docs.google.com/viewerng/viewer?url=https://github.com/brohrer/public-hosting/raw/master/How_CNNs_work.pdf&embedded=true\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x10d943c90>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tempfile import gettempdir\n",
    "from os import path\n",
    "\n",
    "from IPython.display import display, IFrame\n",
    "\n",
    "display(\n",
    "    IFrame(\n",
    "        \"https://docs.google.com/viewerng/viewer?url={url}&embedded=true\".format(\n",
    "            url=\"https://github.com/brohrer/public-hosting/raw/master/How_CNNs_work.pdf\"\n",
    "        ),\n",
    "        height=\"475px\",\n",
    "        width=\"100%\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## Why should I care?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "## How do I ML?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "There are some popular offerrings, of which I'll go through 4 today:\n",
    "- TensorFlow (Google)\n",
    "- JAX (Google)\n",
    "- PyTorch (Facebook)\n",
    "- scikit.learn (David Cournapeau then INRIA)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "|                                          | License      | CPU | GPU | TPU | Web | Mobile |\n",
    "|------------------------------------------|--------------|-----|-----|-----|-----|--------|\n",
    "| [JAX](https://jax.readthedocs.io)        | Apache-2.0   | ✔   | ✔   | ✔   | ✔   | ✔      |\n",
    "| [TensorFlow](https://www.tensorflow.org) | Apache-2.0   | ✔   | ✔   | ✔   | ✔   | ✔      |\n",
    "| [PyTorch](https://pytorch.org)           | BSD 3-Clause | ✔   | ✔   | ✔   | ✘*  | ✔      |\n",
    "| [scikit.learn](https://scikit-learn.org) | New BSD      | ✔   | ✘   | ✘   | ✘†  | ✘‡     |\n",
    "\n",
    "*Efforts to port to WASM, through another language: https://github.com/sinkingsugar/nimtorch\n",
    "\n",
    "†Efforts to port to WASM: https://github.com/iodide-project/pyodide\n",
    "\n",
    "‡Nothing built into the framework, but using third-party libraries you can."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# MNIST\n",
    "Very popular dataset in the machine-learning sphere, with the goal of figuring out from an image which arabic numeral it is.\n",
    "\n",
    "![https://upload.wikimedia.org/wikipedia/commons/2/27/MnistExamples.png](https://upload.wikimedia.org/wikipedia/commons/2/27/MnistExamples.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualisation (3D, drag and right-clik): https://www.cs.ryerson.ca/~aharley/vis/\n",
    "\n",
    "Node/math based: https://jalammar.github.io/visual-interactive-guide-basics-neural-networks https://jalammar.github.io/feedforward-neural-networks-visual-interactive\n",
    "\n",
    "Input nodes, hidden nodes, output nodes: http://nn-mnist.sennabaum.com/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAgQAAAIFCAYAAACtXuUzAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAgAElEQVR4nO3de7yVc97/8feH2iUih1ChDBUR6cBNSeSUwz2MDBlyT/wY4RZ3Y7jnNsjD3GgcxpRkQmZ+TaPRwaFIjdNw49ZOSI1D6aCioqNqKn3vP1rMuq7vd+117bXW3muvtV/Px2Mej76f/bmu/Wke31Yf1/70vcw5JwAAUL/tUOwCAABA8dEQAAAAGgIAAEBDAAAAREMAAABEQwAAACQ1qE6ymfFvFOFxzlmxa8gH+xoZrHTONS92EflgbyMk02c2TwgAIGxhsQsAahMNAQAAoCEAAAA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEDVfLkRAADlol27dl5s+PDhkXXv3r29nNGjR3uxgQMHerFNmzblXlwR8IQAAADQEAAAABoCAAAgGgIAACCGCgEA9dRxxx3nxU466aTI2jnn5Vx66aVe7Ntvv/ViV199dWS9efPm6pZYq3hCAAAAaAgAAAANAQAAEA0BAAAQQ4W1rm/fvl5s3LhxXuzKK6/0Yr///e9rpCagunbaaafI+qGHHvJymjRp4sX69evnxbZt21a4woAMTj/9dC/2wAMPFOz+AwYM8GJz5syJrO+///6Cfb+awBMCAABAQwAAAGgIAACAaAgAAIAYKqx1F110kRcLnYS1xx571EY5QFZm5sVGjhwZWV988cWJ7vXf//3fXmzWrFm5FQZkEBpoHTJkiBdr2rRpjdZxyy23RNYMFQIAgDqPhgAAANAQAAAAZghqXOvWrSPrPn36eDmVlZVe7E9/+lON1QRUR4cOHbxYkpmBtWvXerGvvvqqIDUBVRk/frwX69q1qxcLzW/FhWZcOnXqlKiOBg1K669YnhAAAAAaAgAAQEMAAABEQwAAAFTGQ4Whw1RCkgyV5OPf//3fI+uKigovZ/78+V5s8eLFNVYTUB3nn39+TtctWrTIi7GvUWiXX365F+vVq1fO94t/Hp9wwgleTmho8eSTT/Zi8aHCgw46yMuZN29edUusMTwhAAAANAQAAICGAAAAiIYAAACojIcKQ0MloTdN/exnP4us33rrrYLW0bFjx6w5vO0Nddl1112XNWfr1q1eLPRmQyBf/fv3j6yHDRvm5TRs2DDRvT799FMvdtppp0XW69ev93KSnrjZqFGjyDr09xJDhQAAoE6hIQAAADQEAACAhgAAAKiMhwo3btzoxUIDfvFTqPIZKtxvv/2y3n/dunVezhNPPJHz9wQKqVmzZl5st912y3rdihUrvNjYsWMLUhPqr1atWnmxm2++ObJOOkC4bNkyL3bllVd6sQULFiQrLge9e/f2Yo8++miNfb/q4gkBAACgIQAAADQEAABANAQAAEBlPFS4fPnyWv+e5557rheLD7zMmDHDywkNuwDFMGTIkJyu++CDDwpcCeqb0FD2lClTvFi7du1yuv8999zjxV555ZWc7pWrww47rFa/X3XxhAAAANAQAAAAGgIAAKAyniHYY489av17tmzZMmtObf/MCqiOyy+/PKfrfvvb3xa4EtQ3oQN6cv2Ze+gNsqNHj87pXoVUF2qoCk8IAAAADQEAAKAhAAAAoiEAAAAq46HC0CFBZlaw+4fewnXVVVdl/Z6PPfZYwWoAimX16tWR9bRp04pUCUrRaaed5sVOOeWUnO71zTffeLFzzjnHi61Zsyan+4eE/i5J8vdL6G23dQlPCAAAAA0BAACgIQAAAKIhAAAAKpOhwkaNGnmxK664wos557xYv379Ius2bdp4OaFTD4844ggv1rRpUy/27rvvRtafffaZlwMUQ6dOnbxY/O2cmQwfPjyy3rp1a0FqQvlp1qyZFxs1apQXC30+h8SHCC+99FIvZ/HixQmry66iosKL7b333l4sVP+3334bWS9ZsqRgddUEnhAAAAAaAgAAQEMAAABEQwAAAFQmQ4UXXXSRF0v6+uOOHTtG1qFhwaTDLiF33XVXZL1t27ac7wUU0j333OPFGjTwPxK2bNnixeJDhUAmoaHvJK+Kz+TZZ5+NrCdOnJjzvZK49tprvVivXr0SXbtp06bI+vnnny9ESTWGJwQAAICGAAAA0BAAAADREAAAAJXJUGG3bt282IYNG7xY6NXDS5cujay//vprL2flypVe7KmnnkpU2wsvvJAoD6hJrVu39mLHHnusFwsN0H766ade7IsvvihMYSg7PXv2jKyfeeaZnO8V2o9TpkzJ+X65OOuss3K+Nn7KYdeuXb2cGTNm5Hz/QuMJAQAAoCEAAAA0BAAAQGUyQzBw4MBEsVz17dvXi5mZF5swYYIXW7t2bcHqAHI1ePBgL7bzzjsnujZ0gBGQybBhwyLr0Ftgk5o/f74XGzNmTM73S+LEE0+MrLt3757zveIH0a1atSrne9UGnhAAAAAaAgAAQEMAAABEQwAAAFQmQ4U1LfQ2xdCBGe+8805tlANUW9K3s4WMHj26YHWg/I0bNy6yvv3223O+15NPPplvOVW6+OKLvdhtt90WWe+444453//WW2+NrOfNm5fzvWoDTwgAAAANAQAAoCEAAACiIQAAAGKoMJETTjjBi4WGCl999dXaKAfI6sgjj4ys27Vrl+i6SZMm1UQ5qEcK+SbM+NsCJemyyy6LrLt06eLlLF682IuFBmvjb2bM9D3j4icQSv4wpSTde++9We9Vl/CEAAAA0BAAAAAaAgAAIBoCAAAghgo9nTt39mINGvj/N7344ote7K233qqRmoDqir+CtmHDhomuGzJkSE2UA+Qk9NruXO2wg//fv6HhwLgvv/zSi913331e7De/+U1uhdUhPCEAAAA0BAAAgIYAAABIstABOxmTzZInl6hp06Z5sd69e3uxLVu2eLFBgwZ5sREjRhSmsDrMOWfFriEfpb6vd9llFy/20UcfRdYtWrTwclatWuXFQnmbN2/Oo7qSVumc61rsIvJRjL3dsmXLyHrq1KleTocOHWqrnO+Z+R9TK1as8GKPPPJIZP3oo496OQsWLChYXcWQ6TObJwQAAICGAAAA0BAAAADREAAAAHEwkSc0ZBmKffjhh17sqaeeqpGagKqE3mQYGg6M+5//+R8vVo8HCFEgS5cujaxDbxS88MILvdgtt9zixfbZZ5+cahg9erQXe+6557zYm2++6cUK+bbGUsMTAgAAQEMAAABoCAAAgGgIAACAGCr0HHrooV7sm2++8WI/+tGPvFjo1Cugpp199tk5XTdq1KgCVwL4Qidihk5wrQ+nutZ1PCEAAAA0BAAAgIYAAACIhgAAAIjXH3tWrlzpxUJDMW3btq2NckoCrz8urr322suLxU/SDP05P+igg7xYaIC2HuP1xyhLvP4YAABkREMAAABoCAAAAA0BAAAQQ4UoAIYKUaYYKkRZYqgQAABkREMAAABoCAAAAA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAACQ1KCa+SslLayJQlCyWhe7gAJgXyOEvY1ylHFfV+v1xwAAoDzxIwMAAEBDAAAAaAgAAIBoCDxmtqOZvWtmz1WR84CZ9YzFHjSz9Wnra8xsQE3WCiRhZo+Z2XIzm50lb5CZ9U/9+nwz+9DMtplZ17ScjmY2uoZLBhIxs9PN7CMz+9TMbqoi7/vPbDM70MzeTl3zpJlVpOL1/jObhsB3naS5mb5oZntK+hfn3Gtpsa6Sdo+lPibp2hqpEKie0ZJOryrBzBpIGiDpT6nQbEk/kvRaep5z7gNJ+5nZAYUvE0jOzHaUNFxSH0kdJPUzsw6BvPhn9t2S7nfOHSxplaTLUvF6/5lNQ5DGzPaTdKakUVWknSfphbRrdpQ0VNKN6UnOuQ2SFpjZ0TVQKpBY6oPw6yxpJ0ma6ZzbmrpmrnPuowy5z0q6sIAlArk4WtKnzrn5zrnNkv4s6YeBvO8/s83MtH2vP5X62hOSzpH4zJZoCOIe0Pa/2LdVkdNdUmXa+hpJzzjnlgVyZ0g6vnDlATUmvq+rwr5GXdBK0uK09eepWFz63t5T0urvGt/ANfV6b9MQpJjZWZKWO+eyfSi2kLQidU1LSedL+l2G3OWSWhasSKDmfL+vE2Bfo5SwtxOiIfin7pL+1cwWaPujp5PM7P8H8jZKapz69VGSDpb0aeq6Jmb2aVpu41Q+UNel7+ts2NeoC5ZI2j9tvV8qFpe+t7+S1Cw1MxO6pl7vbRqCFOfczc65/ZxzbbT956MvOecuDqTO1fYmQM65yc65fZ1zbVLXbUgNqnynnbYPZwF13ff7OgH2NeqCdyS1Tf2rgQpt/9x+JpCX/pntJL0sqW/qa5dKejott17vbRqC6pssqVfC3O6SptVcKUB2ZjZW0puS2pvZ52Z2WSDteUk9064518w+l3SspMlmNjUt90Rt/3MAFE1qDuAaSVO1/S/9cc65DwOp8c/sX0i6IfU0d09Jj6Z9rV5/ZvMugxyY2euSznLOra4i5yhJNzjnLqm9yoDcmdlESTc65z6pIqeRpFcl9UgbzALqND6zk6EhyIGZHSNpo3Pu/SpyTpH0iXNuQa0VBuTBzNpL2if9jI1ATltJrZxzr9RaYUCe+MxOhoYAAAAwQwAAAGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgKQG1Uk2M1dThaB0Oees2DXkg32NDFY655oXu4h8sLcRkukzmycEABC2sNgFALWJhgAAANAQAAAAGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgqUGxCwBQ+po2berFrr76ai/261//2ostW7Yssu7QoYOXs2bNmjyqA8IaNWrkxd54443I+gc/+IGXc/LJJ3uxmTNnFq6wIuEJAQAAoCEAAAA0BAAAQDQEAABAkjnnkiebJU9GveGcs2LXkI9y3NfxQajQ0N95553nxRo3bpz1XqHYe++95+X0798/a52SZBbdPi1atPByvvzyy0T3KrBK51zXYnzjQinHvV1I++67rxdbunRp1utmz57txbp16+bF/vGPf+RWWA3L9JnNEwIAAEBDAAAAaAgAAIBoCAAAgMr4pMKXX37Zi/Xq1cuL3X333ZH1TTfdVFMlAdUSOkXtwAMP9GIjRozwYkcddVRkveuuu3o51RkojosPAh555JE53wsolttuuy2n60J/npo3b+7FPv/885zuXyw8IQAAADQEAACAhgAAAKgEZwjiP7uUpPbt23ux+M9QJWnbtm1e7Lrrrousv/32Wy9nwoQJXiz089ePPvrIi8WddNJJXix08MuCBQu82JQpUyLrLVu2ZP1+KA2hPTBu3DgvFtrXScTf4CZJ8+bN82KTJ0/2YqtXr/ZiU6dOzamOkCVLlkTWmzZtKti9ge+ce+65XuzKK6/0Yklma+bMmePFSm1eIIQnBAAAgIYAAADQEAAAANEQAAAAleBQYceOHb3Yu+++m/P9KioqIuvQwUR15bCiv/3tb5F1aEhm1apVtVUO8tCnT5/IOjTMF7Ju3TovFjqEa+jQoZF1aKgwqUsuuSRrzvr16xPdK/TWxb/+9a+R9Zo1a5IVBlTDIYccktN18aFXSRowYEC+5dRJPCEAAAA0BAAAgIYAAACIhgAAAEiy6rzxzMxyfz1ajlq3bh1Zhwao4jmZrF271ovFTy/cfffdvZyk/x+FTlFMcm1oiGq33XbLev+HH37Yyxk4cGDW71dozjn/N15CanpfH3bYYV5s5syZkXWDBv587//+7/96sb59+3qx0NBTIXXo0MGLXXXVVZF16JS266+/3ouF3gi3yy67RNYbN26sbok1pdI517XYReSjGJ/ZddXcuXO9WGjQMP6Zfeutt3o5d9xxR+EKK4JMn9k8IQAAADQEAACAhgAAAIiGAAAAqAROKrziiisi66QDhHfffbcXe+CBB7xYfIAp9HrimjZ79mwv9vHHH2e9LnTqG+qeI444wouFhgjjzjjjDC9WjJMoQ696vfbaayPrfv36eTmhAcINGzZ4sTo0RIgyEdqPbdu2zeleixcvzrecksETAgAAQEMAAABoCAAAgGgIAACA6thQYY8ePbzYoEGDcrrXgw8+6MWWL1+e9bqnn346p++Xj4MPPjhRXvwErdNOO83Lady4sRfbtGlTboWhII466qicruvSpYsXmz59er7l1Iif//znifLuvffeGq4EkG655RYvtsMOyf77d8WKFZH1hAkTClJTKeAJAQAAoCEAAAA0BAAAQHVshiD0M/74z8Q3b97s5QwbNsyLFeMAl1xddNFFifLibzucOnWql8O8QN0zZswYLzZ48OCs17344ouJ7v/cc895sfj+X7ZsmZczadIkL/bWW28l+p6XXnppZN2pUycv54svvvBit912W6L7A/kIvbU2qfvvvz+yDr0lt1zxhAAAANAQAAAAGgIAACAaAgAAoDo2VPjJJ594scMOOyyyXrdunZezZMmSGqupNuy6666J8uIHE6E0hN4WeOaZZ0bWd955p5cT2hcHHnhg1nuFxAdSJen666/3Yl999VXWe0nSbrvtFlmH9uaiRYu82JFHHunF3nvvvUTfEwi55JJLvNjee++d6Nr169d7sfp8eBZPCAAAAA0BAACgIQAAAKIhAAAAkqw6g2pmxlRbnoYMGeLFQm+Kq6io8GLxgcpzzjnHy3nllVdyLy5Hzjl/Yq2E1NV93bRpUy+WdKiwWbNmkXVoqDD0Zz9+AqEkNW/e3IvF75fPwOsHH3wQWYf+PEybNi3n++eh0jnXtRjfuFDq6t4upD/84Q9eLDRoGLJ69Wovls8ph6Ui02c2TwgAAAANAQAAoCEAAACiIQAAAKpjJxWWozvuuCOyvvnmm72c0MBXyKhRoyLrYgwQovaETuV8//33E8WSOPnkk73YlVdemejaysrKyHro0KFezhlnnOHFevfu7cWOOOKIyPovf/mLl9O5c2cvNn/+/Kx1ovzEX7V99tlnezlJh1zvueeegtRULnhCAAAAaAgAAAANAQAAEA0BAAAQQ4U5Cw0C/uQnP/Fi//Ef/5H1upCXXnrJi910000JqwOibrvtNi8WOhFwp5128mJvvPGGF4ufaBga8Bs3bpwX69Gjhxd77bXXIuvQa5932WUXL4b6qW3btpF1/FXc1TF58uR8yykrPCEAAAA0BAAAgIYAAACIGYJE2rRp48Vuv/12LxZ6w1aSAzI++ugjL/bTn/7Ui23dujXrvVD/NGzY0ItNmjQpsu7Tp4+XE9qbY8aM8WLXXHONF1uzZk11Svxe6IChuNmzZ3uxOXPm5PT9gKp0797di+V60Fc54AkBAACgIQAAADQEAABANAQAAECSJX0rlCSZWfLkEnX44Yd7sbvvvtuLnX766Tndf+LEiV5s8ODBXmzBggU53b8YnHPJTluqo+rqvt533329WN++fb3YBRdckPXa/fbbz8sJ7etQbOPGjVXWmcnOO+/sxWbMmOHF2rdvH1mHDvgaO3ZsTjXkqdI517UY37hQ6urezsf48eMj63PPPTfne33zzTderGnTpjnfr1Rk+szmCQEAAKAhAAAANAQAAEA0BAAAQJxUqFatWkXWjz76qJfTtWvuc0XxU95GjBiR871QPuJvFXzooYe8nPgbBaVkJ19K0vTp0yPrm2++2ct56qmnEt0rVx07dvRi7dq182JLliyJrKdMmVJjNaH0/eAHPyh2CWWLJwQAAICGAAAA0BAAAADREAAAADFUqOuuuy6y7tatm5cTGuRav369F7vpppu82KhRo/KoDuXgmGOO8WLDhg2LrLt06eLlmPmHid13331e7M477/Riq1atqk6JeTvggAO82OTJk71Y6Pd0xx13RNa5vloZqK7QybH1GU8IAAAADQEAAKAhAAAAqmczBPGfVUr+DEFoXiD0M83QQS8jR47MozqUq/POO8+Lde7cObJOeuDQ3LlzvVjo7Wyhn+kX0nHHHRdZh/48NGvWzIvNmzfPiz3yyCOFKwxl5YQTTvBihx56aE73ev/9971Y//79c7pXueIJAQAAoCEAAAA0BAAAQDQEAABAZTxUGBpouuiii7xYgwbR/wtCB6f8+c9/9mIMECKp0aNHe7Gzzz47sg69BTAkNIAXOoRo9913j6xD+zrpIGNI/H6bN2/2ckJvLQz9GQQyadKkiRerqKjI6V6hg7IQxRMCAABAQwAAAGgIAACAaAgAAIDKeKiwX79+XqxNmzZZr5s/f74X+/Wvf12IklBPzZkzx4t16tQpsu7Zs6eX0717dy8W2sM77bSTF+vbt281KvynUK2VlZVe7IsvvoisJ02a5OW89dZbOdUAfGfatGlebNCgQZH1Kaec4uWETsR89dVXC1dYmeIJAQAAoCEAAAA0BAAAQDQEAABAklXntDIzy/1os1rWp08fLxY6qSr++7/qqqu8HF7PWjXnnH8MXgkppX2NWlXpnOta7CLywd5GSKbPbJ4QAAAAGgIAAEBDAAAAREMAAABUxicVvvTSS17s7bff9mLt27fPeh0AAOWOJwQAAICGAAAA0BAAAACV8cFEqD0cTIQyxcFEKEscTAQAADKiIQAAADQEAACAhgAAAIiGAAAAiIYAAACIhgAAAIiGAAAAiIYAAACo+m87XClpYU0UgpLVutgFFAD7GiHsbZSjjPu6WkcXAwCA8sSPDAAAAA0BAACgIQAAAKIh+J6ZtTezWWn/W2tmgzLkDjKz/qlfP5l2zQIzm5WKdzSz0bX4WwCCzOx6M/vQzGab2Vgza5wh7wEz65n6dW8zm5na16+b2cGp+DVmNqA26wcyMbPHzGy5mc3Okpf+mX1+6s/DNjPrmpZT7z+zGSoMMLMdJS2RdIxzbmHsaw0kzZTU2Tm3Nfa1eyWtcc4NSa2nSxrgnFtUO5UDUWbWStLrkjo45zaa2ThJU5xzo2N5e0qa7Jz7l9T6Y0k/dM7NNbOBko52zv2bmTWR9IZz7qja/Z0AvlQDu17SH5xzh2fIiXxmm9mhkrZJGilpsHNuRlpuvf7M5glBWG9J8+LNQMpJkmYGmgGT9GNJY9PCz0q6sMaqBJJpIGmn1AdjE0lLAznnSXohbe0k7Zr69W7fXeOc2yBpgZkdXXPlAsk4516T9HWWtMhntnNurnPuowy59fozm4Yg7EJF/2JP111SZSB+vKQvnXOfpMVmpOJAUTjnlkj6jaRFkpZp+xOsFwOp8X19uaQpZva5pEsk3ZX2NfY1Skmmz+yQer23aQhizKxC0r9K+kuGlBaSVgTi/eQ3EcsltSxcdUD1mNnukn4o6UBt34s7m9nFgdT4vr5e0hnOuf0kPS7pvrSvsa9RSjJ9ZofU671NQ+Dro+2Pl77M8PWNkiJDWalHsT+S9GQst3EqHyiWkyV95pxb4ZzbImmCpOMCed/vazNrLulI59zbqa89GbuGfY1S4n1mV6Fe720aAl/ov/TTzZV0cCx2sqS/O+c+j8XbSapy+hWoYYsk/YuZNUnNufTW9j0cl76vV0nazczapdanxK5hX6OUhD6zM6nXe5uGII2Z7aztH34Tqkh7XlLPWCzTzMGJkiYXpjqg+lL/lf+Utk9Zf6Dtf+YfCaROltQrdc1WSf9P0ngze0/bZwh+npbbXdK0mqsaSMbMxkp6U1J7M/vczC4LpEU+s83s3NRszLGSJpvZ1LTcev2ZzT87zIGZTZR0Y2yAMJ7TSNKrknrE/0UCUBeZ2euSznLOra4i5yhJNzjnLqm9yoD88JmdDA1BDsysvaR9Uv/kJVNOW0mtnHOv1FphQB7M7BhJG51z71eRc4qkT5xzC2qtMCBPfGYnQ0MAAACYIQAAADQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQFKD6iSbmaupQlC6nHNW7Brywb5GBiudc82LXUQ+2NsIyfSZzRMCAAhbWOwCgNpEQwAAAGgIAAAADQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAAAQDQEAABANAQAAEA0BAACQ1KDYBRTC7373Oy/WpUuXRNe+8MILkfXChQu9nC+++MKLTZ06NWF1AIBycsghh3ixWbNmebF33nknsj7++ONrrKZC4AkBAACgIQAAADQEAABANAQAAEAlMFTYqFGjyHr48OFezoABA3K+/7HHHhtZO+e8nG3btnmxGTNmeLFf/epXXuzFF1/MuTYAQN3To0cPL7bjjjt6scMPPzyyPuigg7ycefPmFa6wPPGEAAAA0BAAAAAaAgAAIBoCAACgEhgqvPHGGyPrfAYIQ0JDhHE77OD3TUcffbQXCw089uvXL7IODSMCdUXPnj292IMPPujF2rdvH1nfcMMNXs6IESMKVxhQJH369PFioQHyBg38v043bNgQWW/atKlwhdUAnhAAAAAaAgAAQEMAAABUAjMELVu2zJozYcIEL/bee+95sfXr13uxP/7xj5F1/CAkSRozZowXO+6447xY6NCJRx55JLLu1q2bl/Ptt996MSBkl1128WJbt271YvG9GD8gRQrv4dAMQceOHbPWFT/gS2KGAKUpfsDQwIEDvZz999/fi4U+x//6179G1kuWLMmzuprFEz9Nef4AAAtLSURBVAIAAEBDAAAAaAgAAIBoCAAAgEpgqDA+mLRo0SIv55577vFihRzU69Wrlxd74YUXvNipp57qxTp16hRZ/+xnP/NyQgcaobw1adIksp4yZUqi6zZv3uzFDj74YC+2zz77RNaNGzf2cszMiyU5qCtk3bp1OV0H1DVDhgyJrM8666xE173zzjterH///gWpqbbwhAAAANAQAAAAGgIAACAaAgAAIMmqM0RkZrlNHJWhHj16eLHp06d7sYqKish6+fLlXk7ozYmh4cm6yjnnT6eVkGLs6z333DOyDu2LfIb+4m9VC51m+Pjjj2etS5IuuOACLxY/zS30RsTrr78+a511XKVzrmuxi8gHn9lVO+SQQ7xYZWVlZB0fAJbCQ+tnn322F3v++efzqK76unb1t2voDbuZPrN5QgAAAGgIAAAADQEAABANAQAAUAmcVFhXvf76615s6NChXuyXv/xlZL333nt7OW3atPFipTRUiOqLn+x35plnFvT+CxYsiKzXrl3r5SxdujTRvUJDr/HTEUP3B+qS0HDgrbfemigvbuzYsV6stgcIQzZs2JDX9TwhAAAANAQAAICGAAAAiIYAAACIocKCevrpp71YfKgwpGPHjl7stddeK0hNqJvirzEOvU67GJo1a+bFQkNW8VMU40OMQF0TOknwwgsvzHrd119/7cVGjhxZkJoKbc6cOXldzxMCAABAQwAAAGgIAACAmCGoE0I/23r44Ye9WOgNW0AhtW/f3ou1bNnSi8XfunjiiSd6OaG3KQK1oVevXl7siSeeSHRtfG/fcMMNXk7oYLpywBMCAABAQwAAAGgIAACAaAgAAIAYKiyoFStWeLGVK1dG1nvttZeXE39znCRVVFR4sY0bN+ZRHZBd6JCsJD744IMCVwLk7le/+pUXa9SoUaJrhw0bFlknHUYsBzwhAAAANAQAAICGAAAAiIYAAACoBIcKQ29jC52kFrJ161Yv9vHHH+dd03eaN2/uxUJDhHH333+/F2OAEMWQ61BhIf8cAdVx1VVXebEePXokunbhwoVe7L/+67/yrqlU8YQAAADQEAAAABoCAAAgGgIAAKASGCrs06dPZB0awGvXrl2ie23evNmL3X777ZH1lClTvJz33nsv0f1/+MMfJsqL45Q35CO07+LDgZ999pmX85Of/MSLHXLIITnVED/dTZK6dOnixUInyAHVsc8++0TWv/jFL7ychg0berHQUPnQoUO92Nq1a/OorrTxhAAAANAQAAAAGgIAACAaAgAAoBIYKnz66acj6wYNci859ErhO++8M7K+9dZbvZxnn33Wi02ePNmL3XjjjVlr2LJlixf7xz/+kfU6QJJGjRrlxS644AIvtvPOO2e9l5l5MedcojriA7qhP1tAvkKf9/HXEbdu3TrRvULD28OHD8+tsDLFEwIAAEBDAAAAaAgAAIAkS/ozQ0kys+TJBRI/UCXpz4uWLVvmxUI/Qzr11FNzKyxHc+fO9WKHHXZYrdYgSZ07d/Zi+++/f2Qdn9/IxDnn/zC6hBRjX+eqTZs2XmzEiBFe7KCDDoqsV65c6eWEZggOOOAAL7bvvvt6salTp0bWoTmGdevWebESU+mc61rsIvJRSns7pFOnTl7s3XffzXpd6BCiH//4x15s4sSJuRVW4jJ9ZvOEAAAA0BAAAAAaAgAAIBoCAACgEjiYaMiQIZH1yJEjvZzQ4RWVlZVe7IorrvBijRs3jqz/9re/eTmtWrXKWmdSbdu29WJLlizxYnPmzPFiHTp0KFgdzZo182LxIbMmTZoU7PuhMBYsWODF4m8ElaSmTZtG1kkH/F566SUvFhoqjL8VsQwGCFEH3XLLLTld99vf/taL1dcBwurgCQEAAKAhAAAANAQAAEA0BAAAQCUwVPj4449H1qGhqt///vde7KyzzvJiS5cu9WJvvvlmZL3HHntUs8LqCQ1AtmjRIlEsV4sWLfJiEyZM8GL33ntvwb4niivJkF/o1MNu3bolun/Dhg2rWxJQpa5d/UMhQwOzSUyaNCnfcuolnhAAAAAaAgAAQEMAAABEQwAAAFQCQ4VxL7/8she74YYbvNjQoUO9WGiI6thjj836PTdv3uzFQq/gvPPOO73Y3//+96z3DxkwYIAXq6ioiKxDpzG+8847Xmz16tVeLPQ6XNQvhx56qBdLejrl+PHjC10O6rnBgwd7sZ122inrddOnT/dib7/9dkFqqm94QgAAAGgIAAAADQEAAFAJzhCEPPPMM4linTp18mJHHHFE1vu/9tprXix0QFIh/ed//meN3h8IzdTE33iZybJlywpcDeqTvffe24slmecKueuuu7zYli1bcrpXfccTAgAAQEMAAABoCAAAgGgIAACAymSoMKlZs2YligH1wV577eXFnHOJrg0dEAYktfvuu3uxAw44IKd7bdu2Ld9ykMITAgAAQEMAAABoCAAAgGgIAACA6tlQIYB/ateuXaK80Kmc77//foGrQX3y2WefebGHHnrIiw0cONCLff3115H14sWLC1dYPccTAgAAQEMAAABoCAAAgGgIAACAGCoEkMU333zjxTZt2lSESlAuNm/e7MWuvvrqRDHUHJ4QAAAAGgIAAEBDAAAAREMAAADEUCGALMaPH1/sEgDUAp4QAAAAGgIAAEBDAAAAJJlzLnmyWfJk1BvOOSt2DflgXyODSudc12IXkQ/2NkIyfWbzhAAAANAQAAAAGgIAACAaAgAAIBoCAAAgGgIAACAaAgAAIBoCAAAgGgIAAKDqv+1wpaSFNVEISlbrYhdQAOxrhLC3UY4y7utqHV0MAADKEz8yAAAANAQAAICGAAAAiIbge2a2v5m9bGZzzOxDM7uuitxBZtY/9euhZvZ3M3vfzCaaWbNUvKOZja6l8oEgM2tvZrPS/rfWzAZlyE3f10+mXbPAzGal4uxr1Blm9piZLTez2Vny0vf2+anP+G1m1jUtp97vbYYKU8yshaQWzrmZZtZUUqWkc5xzc2J5DSTNlNTZObfVzE6V9FLq13dLknPuF6nc6ZIGOOcW1epvBggwsx0lLZF0jHNuYexrkX0d+9q9ktY454ak1uxr1Alm1lPSekl/cM4dniEn/pl9qKRtkkZKGuycm5GWW6/3Nk8IUpxzy5xzM1O/XidprqRWgdSTJM387kPTOfdi2gfoW5L2S8t9VtKFNVc1UC29Jc2LNwMpkX39HTMzST+WNDYtzL5GneCce03S11nS4p/Zc51zH2XIrdd7m4YgwMzaSDpK0tuBL3fX9qcHIQMkPZ+2niHp+ELWBuThQkX/Yk+XaV8fL+lL59wnaTH2NUpJVZ/ZcfV6b9MQxJjZLpLGSxrknFsbSGkhaUXgul9K2ippTFp4uaSWNVEnUB1mViHpXyX9JUNKcF9L6ie/iWBfo5Rk2tsh9XpvV/ekwrJmZg21vRkY45ybkCFto6TGsev+TdJZknq76FBG41Q+UGx9tP2x6ZcZvh7a1w0k/UhSl1gu+xqlxNvbVajXe5uGICX1s9JHJc11zt1XRepcSQenXXe6pBslneCc2xDLbSepyulXoJaE/ks/XWRfp5ws6e/Ouc9jcfY1Sklob2dSr/c2PzL4p+6SLpF0Uto/tzojkPe8pJ5p62GSmkqalrrm4bSvnShpco1VDCRgZjtLOkVSpqdekr+vpcwzB+xr1AlmNlbSm5Lam9nnZnZZIC2yt83sXDP7XNKxkiab2dS03Hq9t/lnhzkws4mSbowNWsVzGkl6VVKP+OQ2UBexr1Gu2NvJ0BDkwMzaS9on9U9eMuW0ldTKOfdKrRUG5IF9jXLF3k6GhgAAADBDAAAAaAgAAIBoCAAAgGgIAACAaAgAAICk/wNZrye0fV8Z5gAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 648x648 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Let's load in the dataset. This will be used for JAX and TensorFlow examples.\n",
    "# (Although `as_numpy` would work fine with PyTorch, we want to use their recommended setup.)\n",
    "\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "data_dir = path.join(gettempdir(), \"tfds\")\n",
    "\n",
    "# Fetch full datasets for evaluation\n",
    "# tfds.load returns tf.Tensors (or tf.data.Datasets if batch_size != -1)\n",
    "# You can convert them to NumPy arrays (or iterables of NumPy arrays) with tfds.dataset_as_numpy\n",
    "mnist_data, mnist_info = tfds.load(\n",
    "    name=\"mnist\", batch_size=-1, data_dir=data_dir, with_info=True\n",
    ")\n",
    "_ = tfds.show_examples(\n",
    "    *tfds.load(\"mnist\", split=\"train\", with_info=True, data_dir=mnist_info.data_dir)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tfds.core.DatasetInfo(\n",
       "    name='mnist',\n",
       "    version=3.0.1,\n",
       "    description='The MNIST database of handwritten digits.',\n",
       "    homepage='http://yann.lecun.com/exdb/mnist/',\n",
       "    features=FeaturesDict({\n",
       "        'image': Image(shape=(28, 28, 1), dtype=tf.uint8),\n",
       "        'label': ClassLabel(shape=(), dtype=tf.int64, num_classes=10),\n",
       "    }),\n",
       "    total_num_examples=70000,\n",
       "    splits={\n",
       "        'test': 10000,\n",
       "        'train': 60000,\n",
       "    },\n",
       "    supervised_keys=('image', 'label'),\n",
       "    citation=\"\"\"@article{lecun2010mnist,\n",
       "      title={MNIST handwritten digit database},\n",
       "      author={LeCun, Yann and Cortes, Corinna and Burges, CJ},\n",
       "      journal={ATT Labs [Online]. Available: http://yann. lecun. com/exdb/mnist},\n",
       "      volume={2},\n",
       "      year={2010}\n",
       "    }\"\"\",\n",
       "    redistribution_info=,\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mnist_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def globals_helper(k):\n",
    "    globals_helper.i += 1\n",
    "\n",
    "    return k if isinstance(k, (type(None), dict, float, int, list, set, str)) else None\n",
    "\n",
    "\n",
    "globals_helper.i = 0\n",
    "globals_before_jax = frozenset(filter(None, map(globals_helper, globals())))\n",
    "globals_i_before_jax = globals_helper.i\n",
    "globals_helper.i = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "# JAX\n",
    "![JAX](https://raw.githubusercontent.com/google/jax/master/images/jax_logo_250px.png)\n",
    "Think of it as numpy but for GPUs and TPUs. It is very focussed on composability.\n",
    "> It can differentiate through loops, branches, recursion, and closures, and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation) via grad as well as forward-mode differentiation, and the two can be composed arbitrarily to any order.\n",
    "\n",
    "Ref: The next are based off the [official Apache-2.0 licensed notebook by JAX](https://github.com/google/jax/blob/b6777c0/docs/notebooks/neural_network_with_tfds_data.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OksHydJDtbbI",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import jax.numpy as np\n",
    "from jax import grad, jit, vmap\n",
    "from jax import random\n",
    "from jax.scipy.special import logsumexp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MTVcKi-ZYB3R",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Hyperparameters\n",
    "Let's get a few bookkeeping items out of the way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-fmWA06xYE7d",
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# A helper function to randomly initialize weights and biases\n",
    "# for a dense neural network layer\n",
    "def random_layer_params(m, n, key, scale=1e-2):\n",
    "    w_key, b_key = random.split(key)\n",
    "    return scale * random.normal(w_key, (n, m)), scale * random.normal(b_key, (n,))\n",
    "\n",
    "\n",
    "# Initialize all layers for a fully-connected neural network with sizes \"sizes\"\n",
    "def init_network_params(sizes, key):\n",
    "    keys = random.split(key, len(sizes))\n",
    "    return [\n",
    "        random_layer_params(m, n, k) for m, n, k in zip(sizes[:-1], sizes[1:], keys)\n",
    "    ]\n",
    "\n",
    "\n",
    "layer_sizes = [784, 512, 512, 10]\n",
    "param_scale = 0.1\n",
    "step_size = 0.01\n",
    "batch_size = 128\n",
    "n_targets = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "BtoNk_yxWtIw",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Auto-batching predictions\n",
    "\n",
    "Let us first define our prediction function. Note that we're defining this for a _single_ image example. We then use JAX's `vmap` function to automatically handle mini-batches, with no performance penalty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7APc6tD7TiuZ",
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def relu(x):\n",
    "    return np.maximum(0, x)\n",
    "\n",
    "\n",
    "def predict(params, image):\n",
    "    # per-example predictions\n",
    "    activations = image\n",
    "    for w, b in params[:-1]:\n",
    "        outputs = np.dot(w, activations) + b\n",
    "        activations = relu(outputs)\n",
    "\n",
    "    final_w, final_b = params[-1]\n",
    "    logits = np.dot(final_w, activations) + final_b\n",
    "    return logits - logsumexp(logits)\n",
    "\n",
    "\n",
    "batched_predict = vmap(predict, in_axes=(None, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NwDuFqc9X7ER",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Utility and loss functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "local_metadata": {},
    "remote_metadata": {
     "colab": {},
     "colab_type": "code",
     "id": "6lTI6I4lWdh5"
    },
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def one_hot(x, k, dtype=np.float32):\n",
    "    \"\"\"Create a one-hot encoding of x of size k.\"\"\"\n",
    "    return np.array(x[:, None] == np.arange(k), dtype)\n",
    "\n",
    "\n",
    "def accuracy(params, images, targets):\n",
    "    target_class = np.argmax(targets, axis=1)\n",
    "    predicted_class = np.argmax(batched_predict(params, images), axis=1)\n",
    "    return np.mean(predicted_class == target_class)\n",
    "\n",
    "\n",
    "def loss(params, images, targets):\n",
    "    preds = batched_predict(params, images)\n",
    "    return -np.mean(preds * targets)\n",
    "\n",
    "\n",
    "@jit\n",
    "def update(params, x, y):\n",
    "    grads = grad(loss)(params, x, y)\n",
    "    return [\n",
    "        (w - step_size * dw, b - step_size * db)\n",
    "        for (w, b), (dw, db) in zip(params, grads)\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "umJJGZCC2oKl",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Data Loading with `tensorflow/datasets`\n",
    "\n",
    "JAX is laser-focused on program transformations and accelerator-backed NumPy, so we don't include data loading or munging in the JAX library. There are already a lot of great data loaders out there, so let's just use them instead of reinventing anything. We'll use the `tensorflow/datasets` data loader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "local_metadata": {},
    "remote_metadata": {
     "colab": {},
     "colab_type": "code",
     "id": "uWvo1EgZCvnK"
    },
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/samuel/repos/.venvs_py3/ml-env/lib/python3.7/site-packages/jax/lib/xla_bridge.py:116: UserWarning: No GPU/TPU found, falling back to CPU.\n",
      "  warnings.warn('No GPU/TPU found, falling back to CPU.')\n"
     ]
    }
   ],
   "source": [
    "mnist_data = tfds.as_numpy(mnist_data)\n",
    "train_data, test_data = mnist_data[\"train\"], mnist_data[\"test\"]\n",
    "num_labels = mnist_info.features[\"label\"].num_classes\n",
    "h, w, c = mnist_info.features[\"image\"].shape\n",
    "num_pixels = h * w * c\n",
    "\n",
    "# Full train set\n",
    "train_images, train_labels = train_data[\"image\"], train_data[\"label\"]\n",
    "train_images = np.reshape(train_images, (len(train_images), num_pixels))\n",
    "train_labels = one_hot(train_labels, num_labels)\n",
    "\n",
    "# Full test set\n",
    "test_images, test_labels = test_data[\"image\"], test_data[\"label\"]\n",
    "test_images = np.reshape(test_images, (len(test_images), num_pixels))\n",
    "test_labels = one_hot(test_labels, num_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7VMSC03gCvnO",
    "outputId": "e565586e-d598-4fa1-dd6f-10ba39617f6a",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train: (60000, 784) (60000, 10) \n",
      "Test:  (10000, 784) (10000, 10)\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"Train:\".ljust(6),\n",
    "    train_images.shape,\n",
    "    train_labels.shape,\n",
    "    \"\\nTest:\".ljust(7),\n",
    "    test_images.shape,\n",
    "    test_labels.shape,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xxPd6Qw3Z98v",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "X2DnZo3iYj18",
    "outputId": "bad334e0-127a-40fe-ec21-b0db77c73088",
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "def get_train_batches():\n",
    "    # as_supervised=True gives us the (image, label) as a tuple instead of a dict\n",
    "    ds = tfds.load(name=\"mnist\", split=\"train\", as_supervised=True, data_dir=data_dir)\n",
    "    # You can build up an arbitrary tf.data input pipeline\n",
    "    ds = ds.batch(batch_size).prefetch(1)\n",
    "    # tfds.dataset_as_numpy converts the tf.data.Dataset into an iterable of NumPy arrays\n",
    "    return tfds.as_numpy(ds)\n",
    "\n",
    "\n",
    "def train(num_epochs, callbacks):\n",
    "    params = init_network_params(layer_sizes, random.PRNGKey(0))\n",
    "    for epoch in range(num_epochs):\n",
    "        start_time = time.time()\n",
    "        for x, y in get_train_batches():\n",
    "            x = np.reshape(x, (len(x), num_pixels))\n",
    "            y = one_hot(y, num_labels)\n",
    "            params = update(params, x, y)\n",
    "        epoch_time = time.time() - start_time\n",
    "\n",
    "        train_acc = accuracy(params, train_images, train_labels)\n",
    "        test_acc = accuracy(params, test_images, test_labels)\n",
    "        for callback in callbacks:\n",
    "            callback(**locals())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "X2DnZo3iYj18",
    "outputId": "bad334e0-127a-40fe-ec21-b0db77c73088",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "print_callback = (\n",
    "    lambda epoch, epoch_time, num_epochs, train_acc, test_acc, **_kwargs: print(\n",
    "        \"Epoch {epoch} in {epoch_time:0.2f} sec\\n\"\n",
    "        \"Training set accuracy {train_acc:>10}\\n\"\n",
    "        \"Test set accuracy {test_acc:>22}{maybe_nl}\".format(\n",
    "            epoch=epoch,\n",
    "            epoch_time=epoch_time,\n",
    "            train_acc=format(float(train_acc), \"01.16f\"),\n",
    "            test_acc=format(float(test_acc), \"01.16f\"),\n",
    "            maybe_nl=\"\\n\" if num_epochs > epoch + 1 else \"\",\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "X2DnZo3iYj18",
    "outputId": "bad334e0-127a-40fe-ec21-b0db77c73088",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 in 5.98 sec\n",
      "Training set accuracy 0.9253333210945129\n",
      "Test set accuracy     0.9272000193595886\n",
      "\n",
      "Epoch 1 in 4.58 sec\n",
      "Training set accuracy 0.9428166747093201\n",
      "Test set accuracy     0.9413999915122986\n",
      "\n",
      "Epoch 2 in 4.54 sec\n",
      "Training set accuracy 0.9532666802406311\n",
      "Test set accuracy     0.9516999721527100\n",
      "\n",
      "Epoch 3 in 5.39 sec\n",
      "Training set accuracy 0.9599499702453613\n",
      "Test set accuracy     0.9552999734878540\n",
      "\n",
      "Epoch 4 in 6.37 sec\n",
      "Training set accuracy 0.9652333259582520\n",
      "Test set accuracy     0.9603999853134155\n",
      "\n",
      "Epoch 5 in 6.08 sec\n",
      "Training set accuracy 0.9691500067710876\n",
      "Test set accuracy     0.9631000161170959\n",
      "\n",
      "Epoch 6 in 4.92 sec\n",
      "Training set accuracy 0.9726166725158691\n",
      "Test set accuracy     0.9653000235557556\n",
      "\n",
      "Epoch 7 in 5.24 sec\n",
      "Training set accuracy 0.9754499793052673\n",
      "Test set accuracy     0.9666000008583069\n",
      "\n",
      "Epoch 8 in 5.28 sec\n",
      "Training set accuracy 0.9779833555221558\n",
      "Test set accuracy     0.9682000279426575\n",
      "\n",
      "Epoch 9 in 4.89 sec\n",
      "Training set accuracy 0.9802833199501038\n",
      "Test set accuracy     0.9689999818801880\n"
     ]
    }
   ],
   "source": [
    "train(num_epochs=10, callbacks=(print_callback,))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xC1CMcVNYwxm",
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "We've now used the whole of the JAX API: `grad` for derivatives, `jit` for speedups and `vmap` for auto-vectorization.\n",
    "We used NumPy to specify all of our computation, and borrowed the great data loaders from `tensorflow/datasets`, and ran the whole thing on either the CPU, GPU, or TPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JAX introduced 47 new symbols\n"
     ]
    }
   ],
   "source": [
    "globals_before_tensorflow = frozenset(filter(None, map(globals_helper, globals())))\n",
    "globals_i_before_tensorflow = globals_helper.i\n",
    "globals_helper.i = 0\n",
    "print(\n",
    "    \"JAX introduced\", globals_i_before_tensorflow - globals_i_before_jax, \"new symbols\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "![TensorFlow](https://www.gstatic.com/devrel-devsite/prod/vf4ca28c48392b1412e7b030290622a0dd55b62dec1202c59f119b1e23227c988/tensorflow/images/lockup.svg)\n",
    "> TensorFlow is an end-to-end open source platform for machine learning. It has a comprehensive, flexible ecosystem of tools, libraries and community resources that lets researchers push the state-of-the-art in ML and developers easily build and deploy ML powered applications.\n",
    "\n",
    "Ref: The next are based off the [official Apache-2.0 licensed notebook by TensorFlow](https://github.com/tensorflow/datasets/blob/8723d84/docs/keras_example.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "(ds_train, ds_test), ds_info = tfds.load(\n",
    "    \"mnist\",\n",
    "    split=[\"train\", \"test\"],\n",
    "    shuffle_files=True,\n",
    "    as_supervised=True,\n",
    "    with_info=True,\n",
    "    data_dir=data_dir,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "def normalize_img(image, label):\n",
    "    \"\"\"Normalizes images: `uint8` -> `float32`.\"\"\"\n",
    "    return tf.cast(image, tf.float32) / 255.0, label\n",
    "\n",
    "\n",
    "ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "ds_train = ds_train.cache()\n",
    "ds_train = ds_train.shuffle(ds_info.splits[\"train\"].num_examples)\n",
    "ds_train = ds_train.batch(128)\n",
    "ds_train = ds_train.prefetch(tf.data.experimental.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "ds_test = ds_test.map(normalize_img, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "ds_test = ds_test.batch(128)\n",
    "ds_test = ds_test.cache()\n",
    "ds_test = ds_test.prefetch(tf.data.experimental.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "model = tf.keras.models.Sequential(\n",
    "    [\n",
    "        tf.keras.layers.Flatten(input_shape=(28, 28, 1)),\n",
    "        tf.keras.layers.Dense(128, activation=\"relu\"),\n",
    "        tf.keras.layers.Dense(10, activation=\"softmax\"),\n",
    "    ]\n",
    ")\n",
    "model.compile(\n",
    "    loss=\"sparse_categorical_crossentropy\",\n",
    "    optimizer=tf.keras.optimizers.Adam(0.001),\n",
    "    metrics=[\"accuracy\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "469/469 [==============================] - 2s 4ms/step - loss: 0.3556 - accuracy: 0.9028 - val_loss: 0.2012 - val_accuracy: 0.9407\n",
      "Epoch 2/10\n",
      "469/469 [==============================] - 1s 2ms/step - loss: 0.1640 - accuracy: 0.9527 - val_loss: 0.1356 - val_accuracy: 0.9600\n",
      "Epoch 3/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.1179 - accuracy: 0.9661 - val_loss: 0.1119 - val_accuracy: 0.9664\n",
      "Epoch 4/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.0915 - accuracy: 0.9734 - val_loss: 0.1097 - val_accuracy: 0.9674\n",
      "Epoch 5/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.0740 - accuracy: 0.9788 - val_loss: 0.0909 - val_accuracy: 0.9722\n",
      "Epoch 6/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.0613 - accuracy: 0.9823 - val_loss: 0.0863 - val_accuracy: 0.9736\n",
      "Epoch 7/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.0513 - accuracy: 0.9854 - val_loss: 0.0794 - val_accuracy: 0.9763\n",
      "Epoch 8/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.0429 - accuracy: 0.9877 - val_loss: 0.0782 - val_accuracy: 0.9745\n",
      "Epoch 9/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.0365 - accuracy: 0.9898 - val_loss: 0.0780 - val_accuracy: 0.9758\n",
      "Epoch 10/10\n",
      "469/469 [==============================] - 1s 3ms/step - loss: 0.0312 - accuracy: 0.9918 - val_loss: 0.0758 - val_accuracy: 0.9778\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x150c19d50>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(ds_train, epochs=10, validation_data=ds_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow introduced 15 new symbols\n"
     ]
    }
   ],
   "source": [
    "globals_before_pytorch = frozenset(filter(None, map(globals_helper, globals())))\n",
    "globals_i_before_pytorch = globals_helper.i\n",
    "globals_helper.i = 0\n",
    "print(\n",
    "    \"TensorFlow introduced\",\n",
    "    globals_i_before_pytorch - globals_i_before_tensorflow,\n",
    "    \"new symbols\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "![PyTorch](https://raw.githubusercontent.com/pytorch/pytorch/master/docs/source/_static/img/pytorch-logo-dark.png)\n",
    "> An open source machine learning framework that accelerates the path from research prototyping to production deployment.\n",
    "\n",
    "Ref: The next are based off the [official BSD 3-Clause example by PyTorch](https://github.com/google/jax/blob/b6777c0/docs/notebooks/neural_network_with_tfds_data.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.optim.lr_scheduler import StepLR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.dropout1 = nn.Dropout2d(0.25)\n",
    "        self.dropout2 = nn.Dropout2d(0.5)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dropout1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def train(model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if batch_idx % 25 == 0:\n",
    "            print(\n",
    "                \"Train Epoch: {epoch} [{batch}/{train_loader_n:} ({pc:.0f}%)]\".format(\n",
    "                    epoch=epoch,\n",
    "                    batch=batch_idx * len(data),\n",
    "                    train_loader_n=len(train_loader.dataset),\n",
    "                    pc=100.0 * batch_idx / len(train_loader),\n",
    "                ),\n",
    "                \"Loss: {loss:.6f}\".ljust(35).format(loss=loss.item()),\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(\n",
    "                output, target, reduction=\"sum\"\n",
    "            ).item()  # sum up batch loss\n",
    "            pred = output.argmax(\n",
    "                dim=1, keepdim=True\n",
    "            )  # get the index of the max log-probability\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "\n",
    "    print(\n",
    "        \"\\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{test_loader_n} ({pc:.0f}%)\\n\".format(\n",
    "            test_loss=test_loss,\n",
    "            correct=correct,\n",
    "            test_loader_n=len(test_loader.dataset),\n",
    "            pc=100.0 * correct / len(test_loader.dataset),\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "\n",
    "kwargs = {\"num_workers\": 1, \"pin_memory\": True} if use_cuda else {}\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST(\n",
    "        data_dir,\n",
    "        train=True,\n",
    "        download=True,\n",
    "        transform=transforms.Compose(\n",
    "            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
    "        ),\n",
    "    ),\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    **kwargs\n",
    ")\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST(\n",
    "        data_dir,\n",
    "        train=False,\n",
    "        transform=transforms.Compose(\n",
    "            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]\n",
    "        ),\n",
    "    ),\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    **kwargs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 1 [0/60000 (0%)] Loss: 2.293000                   \n",
      "Train Epoch: 1 [3200/60000 (5%)] Loss: 0.629171                   \n",
      "Train Epoch: 1 [6400/60000 (11%)] Loss: 0.426245                   \n",
      "Train Epoch: 1 [9600/60000 (16%)] Loss: 0.245386                   \n",
      "Train Epoch: 1 [12800/60000 (21%)] Loss: 0.139739                   \n",
      "Train Epoch: 1 [16000/60000 (27%)] Loss: 0.225497                   \n",
      "Train Epoch: 1 [19200/60000 (32%)] Loss: 0.162955                   \n",
      "Train Epoch: 1 [22400/60000 (37%)] Loss: 0.190514                   \n",
      "Train Epoch: 1 [25600/60000 (43%)] Loss: 0.187234                   \n",
      "Train Epoch: 1 [28800/60000 (48%)] Loss: 0.202223                   \n",
      "Train Epoch: 1 [32000/60000 (53%)] Loss: 0.078650                   \n",
      "Train Epoch: 1 [35200/60000 (59%)] Loss: 0.115532                   \n",
      "Train Epoch: 1 [38400/60000 (64%)] Loss: 0.187219                   \n",
      "Train Epoch: 1 [41600/60000 (69%)] Loss: 0.092489                   \n",
      "Train Epoch: 1 [44800/60000 (75%)] Loss: 0.156758                   \n",
      "Train Epoch: 1 [48000/60000 (80%)] Loss: 0.191907                   \n",
      "Train Epoch: 1 [51200/60000 (85%)] Loss: 0.090627                   \n",
      "Train Epoch: 1 [54400/60000 (91%)] Loss: 0.045642                   \n",
      "Train Epoch: 1 [57600/60000 (96%)] Loss: 0.111965                   \n",
      "\n",
      "Test set: Average loss: 0.0489, Accuracy: 9841/10000 (98%)\n",
      "\n",
      "Train Epoch: 2 [0/60000 (0%)] Loss: 0.057380                   \n",
      "Train Epoch: 2 [3200/60000 (5%)] Loss: 0.054453                   \n",
      "Train Epoch: 2 [6400/60000 (11%)] Loss: 0.053452                   \n",
      "Train Epoch: 2 [9600/60000 (16%)] Loss: 0.035754                   \n",
      "Train Epoch: 2 [12800/60000 (21%)] Loss: 0.080681                   \n",
      "Train Epoch: 2 [16000/60000 (27%)] Loss: 0.232738                   \n",
      "Train Epoch: 2 [19200/60000 (32%)] Loss: 0.032548                   \n",
      "Train Epoch: 2 [22400/60000 (37%)] Loss: 0.033910                   \n",
      "Train Epoch: 2 [25600/60000 (43%)] Loss: 0.042153                   \n",
      "Train Epoch: 2 [28800/60000 (48%)] Loss: 0.120734                   \n",
      "Train Epoch: 2 [32000/60000 (53%)] Loss: 0.082842                   \n",
      "Train Epoch: 2 [35200/60000 (59%)] Loss: 0.116158                   \n",
      "Train Epoch: 2 [38400/60000 (64%)] Loss: 0.118132                   \n",
      "Train Epoch: 2 [41600/60000 (69%)] Loss: 0.139319                   \n",
      "Train Epoch: 2 [44800/60000 (75%)] Loss: 0.099373                   \n",
      "Train Epoch: 2 [48000/60000 (80%)] Loss: 0.059946                   \n",
      "Train Epoch: 2 [51200/60000 (85%)] Loss: 0.024202                   \n",
      "Train Epoch: 2 [54400/60000 (91%)] Loss: 0.114242                   \n",
      "Train Epoch: 2 [57600/60000 (96%)] Loss: 0.021467                   \n",
      "\n",
      "Test set: Average loss: 0.0367, Accuracy: 9873/10000 (99%)\n",
      "\n",
      "Train Epoch: 3 [0/60000 (0%)] Loss: 0.090901                   \n",
      "Train Epoch: 3 [3200/60000 (5%)] Loss: 0.099264                   \n",
      "Train Epoch: 3 [6400/60000 (11%)] Loss: 0.136969                   \n",
      "Train Epoch: 3 [9600/60000 (16%)] Loss: 0.085514                   \n",
      "Train Epoch: 3 [12800/60000 (21%)] Loss: 0.030009                   \n",
      "Train Epoch: 3 [16000/60000 (27%)] Loss: 0.154468                   \n",
      "Train Epoch: 3 [19200/60000 (32%)] Loss: 0.065474                   \n",
      "Train Epoch: 3 [22400/60000 (37%)] Loss: 0.057200                   \n",
      "Train Epoch: 3 [25600/60000 (43%)] Loss: 0.136812                   \n",
      "Train Epoch: 3 [28800/60000 (48%)] Loss: 0.012489                   \n",
      "Train Epoch: 3 [32000/60000 (53%)] Loss: 0.035273                   \n",
      "Train Epoch: 3 [35200/60000 (59%)] Loss: 0.057625                   \n",
      "Train Epoch: 3 [38400/60000 (64%)] Loss: 0.014466                   \n",
      "Train Epoch: 3 [41600/60000 (69%)] Loss: 0.014027                   \n",
      "Train Epoch: 3 [44800/60000 (75%)] Loss: 0.075264                   \n",
      "Train Epoch: 3 [48000/60000 (80%)] Loss: 0.023049                   \n",
      "Train Epoch: 3 [51200/60000 (85%)] Loss: 0.046660                   \n",
      "Train Epoch: 3 [54400/60000 (91%)] Loss: 0.041569                   \n",
      "Train Epoch: 3 [57600/60000 (96%)] Loss: 0.026383                   \n",
      "\n",
      "Test set: Average loss: 0.0337, Accuracy: 9886/10000 (99%)\n",
      "\n",
      "Train Epoch: 4 [0/60000 (0%)] Loss: 0.044014                   \n",
      "Train Epoch: 4 [3200/60000 (5%)] Loss: 0.047033                   \n",
      "Train Epoch: 4 [6400/60000 (11%)] Loss: 0.082856                   \n",
      "Train Epoch: 4 [9600/60000 (16%)] Loss: 0.053350                   \n",
      "Train Epoch: 4 [12800/60000 (21%)] Loss: 0.048778                   \n",
      "Train Epoch: 4 [16000/60000 (27%)] Loss: 0.036856                   \n",
      "Train Epoch: 4 [19200/60000 (32%)] Loss: 0.038749                   \n",
      "Train Epoch: 4 [22400/60000 (37%)] Loss: 0.076408                   \n",
      "Train Epoch: 4 [25600/60000 (43%)] Loss: 0.124678                   \n",
      "Train Epoch: 4 [28800/60000 (48%)] Loss: 0.024585                   \n",
      "Train Epoch: 4 [32000/60000 (53%)] Loss: 0.019719                   \n",
      "Train Epoch: 4 [35200/60000 (59%)] Loss: 0.046091                   \n",
      "Train Epoch: 4 [38400/60000 (64%)] Loss: 0.011804                   \n",
      "Train Epoch: 4 [41600/60000 (69%)] Loss: 0.009481                   \n",
      "Train Epoch: 4 [44800/60000 (75%)] Loss: 0.027054                   \n",
      "Train Epoch: 4 [48000/60000 (80%)] Loss: 0.056267                   \n",
      "Train Epoch: 4 [51200/60000 (85%)] Loss: 0.052366                   \n",
      "Train Epoch: 4 [54400/60000 (91%)] Loss: 0.045904                   \n",
      "Train Epoch: 4 [57600/60000 (96%)] Loss: 0.020108                   \n",
      "\n",
      "Test set: Average loss: 0.0287, Accuracy: 9898/10000 (99%)\n",
      "\n",
      "Train Epoch: 5 [0/60000 (0%)] Loss: 0.020414                   \n",
      "Train Epoch: 5 [3200/60000 (5%)] Loss: 0.038635                   \n",
      "Train Epoch: 5 [6400/60000 (11%)] Loss: 0.042979                   \n",
      "Train Epoch: 5 [9600/60000 (16%)] Loss: 0.042862                   \n",
      "Train Epoch: 5 [12800/60000 (21%)] Loss: 0.009988                   \n",
      "Train Epoch: 5 [16000/60000 (27%)] Loss: 0.032572                   \n",
      "Train Epoch: 5 [19200/60000 (32%)] Loss: 0.015156                   \n",
      "Train Epoch: 5 [22400/60000 (37%)] Loss: 0.033873                   \n",
      "Train Epoch: 5 [25600/60000 (43%)] Loss: 0.018086                   \n",
      "Train Epoch: 5 [28800/60000 (48%)] Loss: 0.033106                   \n",
      "Train Epoch: 5 [32000/60000 (53%)] Loss: 0.103487                   \n",
      "Train Epoch: 5 [35200/60000 (59%)] Loss: 0.011282                   \n",
      "Train Epoch: 5 [38400/60000 (64%)] Loss: 0.039687                   \n",
      "Train Epoch: 5 [41600/60000 (69%)] Loss: 0.046402                   \n",
      "Train Epoch: 5 [44800/60000 (75%)] Loss: 0.011683                   \n",
      "Train Epoch: 5 [48000/60000 (80%)] Loss: 0.021494                   \n",
      "Train Epoch: 5 [51200/60000 (85%)] Loss: 0.028144                   \n",
      "Train Epoch: 5 [54400/60000 (91%)] Loss: 0.060041                   \n",
      "Train Epoch: 5 [57600/60000 (96%)] Loss: 0.095917                   \n",
      "\n",
      "Test set: Average loss: 0.0303, Accuracy: 9897/10000 (99%)\n",
      "\n",
      "Train Epoch: 6 [0/60000 (0%)] Loss: 0.023260                   \n",
      "Train Epoch: 6 [3200/60000 (5%)] Loss: 0.018637                   \n",
      "Train Epoch: 6 [6400/60000 (11%)] Loss: 0.018472                   \n",
      "Train Epoch: 6 [9600/60000 (16%)] Loss: 0.013239                   \n",
      "Train Epoch: 6 [12800/60000 (21%)] Loss: 0.044616                   \n",
      "Train Epoch: 6 [16000/60000 (27%)] Loss: 0.035581                   \n",
      "Train Epoch: 6 [19200/60000 (32%)] Loss: 0.030014                   \n",
      "Train Epoch: 6 [22400/60000 (37%)] Loss: 0.070240                   \n",
      "Train Epoch: 6 [25600/60000 (43%)] Loss: 0.027967                   \n",
      "Train Epoch: 6 [28800/60000 (48%)] Loss: 0.022757                   \n",
      "Train Epoch: 6 [32000/60000 (53%)] Loss: 0.033938                   \n",
      "Train Epoch: 6 [35200/60000 (59%)] Loss: 0.015179                   \n",
      "Train Epoch: 6 [38400/60000 (64%)] Loss: 0.026299                   \n",
      "Train Epoch: 6 [41600/60000 (69%)] Loss: 0.015088                   \n",
      "Train Epoch: 6 [44800/60000 (75%)] Loss: 0.041853                   \n",
      "Train Epoch: 6 [48000/60000 (80%)] Loss: 0.119645                   \n",
      "Train Epoch: 6 [51200/60000 (85%)] Loss: 0.026324                   \n",
      "Train Epoch: 6 [54400/60000 (91%)] Loss: 0.007893                   \n",
      "Train Epoch: 6 [57600/60000 (96%)] Loss: 0.029690                   \n",
      "\n",
      "Test set: Average loss: 0.0266, Accuracy: 9907/10000 (99%)\n",
      "\n",
      "Train Epoch: 7 [0/60000 (0%)] Loss: 0.013343                   \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 7 [3200/60000 (5%)] Loss: 0.003926                   \n",
      "Train Epoch: 7 [6400/60000 (11%)] Loss: 0.021495                   \n",
      "Train Epoch: 7 [9600/60000 (16%)] Loss: 0.067986                   \n",
      "Train Epoch: 7 [12800/60000 (21%)] Loss: 0.040834                   \n",
      "Train Epoch: 7 [16000/60000 (27%)] Loss: 0.009835                   \n",
      "Train Epoch: 7 [19200/60000 (32%)] Loss: 0.032620                   \n",
      "Train Epoch: 7 [22400/60000 (37%)] Loss: 0.029635                   \n",
      "Train Epoch: 7 [25600/60000 (43%)] Loss: 0.018652                   \n",
      "Train Epoch: 7 [28800/60000 (48%)] Loss: 0.012598                   \n",
      "Train Epoch: 7 [32000/60000 (53%)] Loss: 0.042428                   \n",
      "Train Epoch: 7 [35200/60000 (59%)] Loss: 0.023380                   \n",
      "Train Epoch: 7 [38400/60000 (64%)] Loss: 0.024325                   \n",
      "Train Epoch: 7 [41600/60000 (69%)] Loss: 0.025021                   \n",
      "Train Epoch: 7 [44800/60000 (75%)] Loss: 0.020271                   \n",
      "Train Epoch: 7 [48000/60000 (80%)] Loss: 0.151270                   \n",
      "Train Epoch: 7 [51200/60000 (85%)] Loss: 0.008643                   \n",
      "Train Epoch: 7 [54400/60000 (91%)] Loss: 0.024911                   \n",
      "Train Epoch: 7 [57600/60000 (96%)] Loss: 0.050122                   \n",
      "\n",
      "Test set: Average loss: 0.0274, Accuracy: 9912/10000 (99%)\n",
      "\n",
      "Train Epoch: 8 [0/60000 (0%)] Loss: 0.035282                   \n",
      "Train Epoch: 8 [3200/60000 (5%)] Loss: 0.006967                   \n",
      "Train Epoch: 8 [6400/60000 (11%)] Loss: 0.060184                   \n",
      "Train Epoch: 8 [9600/60000 (16%)] Loss: 0.027839                   \n",
      "Train Epoch: 8 [12800/60000 (21%)] Loss: 0.018268                   \n",
      "Train Epoch: 8 [16000/60000 (27%)] Loss: 0.019240                   \n",
      "Train Epoch: 8 [19200/60000 (32%)] Loss: 0.031086                   \n",
      "Train Epoch: 8 [22400/60000 (37%)] Loss: 0.019107                   \n",
      "Train Epoch: 8 [25600/60000 (43%)] Loss: 0.014636                   \n",
      "Train Epoch: 8 [28800/60000 (48%)] Loss: 0.068243                   \n",
      "Train Epoch: 8 [32000/60000 (53%)] Loss: 0.011646                   \n",
      "Train Epoch: 8 [35200/60000 (59%)] Loss: 0.067783                   \n",
      "Train Epoch: 8 [38400/60000 (64%)] Loss: 0.006821                   \n",
      "Train Epoch: 8 [41600/60000 (69%)] Loss: 0.051031                   \n",
      "Train Epoch: 8 [44800/60000 (75%)] Loss: 0.010537                   \n",
      "Train Epoch: 8 [48000/60000 (80%)] Loss: 0.036322                   \n",
      "Train Epoch: 8 [51200/60000 (85%)] Loss: 0.027204                   \n",
      "Train Epoch: 8 [54400/60000 (91%)] Loss: 0.126911                   \n",
      "Train Epoch: 8 [57600/60000 (96%)] Loss: 0.005657                   \n",
      "\n",
      "Test set: Average loss: 0.0274, Accuracy: 9908/10000 (99%)\n",
      "\n",
      "Train Epoch: 9 [0/60000 (0%)] Loss: 0.028593                   \n",
      "Train Epoch: 9 [3200/60000 (5%)] Loss: 0.006153                   \n",
      "Train Epoch: 9 [6400/60000 (11%)] Loss: 0.066536                   \n",
      "Train Epoch: 9 [9600/60000 (16%)] Loss: 0.017906                   \n",
      "Train Epoch: 9 [12800/60000 (21%)] Loss: 0.026001                   \n",
      "Train Epoch: 9 [16000/60000 (27%)] Loss: 0.026464                   \n",
      "Train Epoch: 9 [19200/60000 (32%)] Loss: 0.062509                   \n",
      "Train Epoch: 9 [22400/60000 (37%)] Loss: 0.009449                   \n",
      "Train Epoch: 9 [25600/60000 (43%)] Loss: 0.006778                   \n",
      "Train Epoch: 9 [28800/60000 (48%)] Loss: 0.006699                   \n",
      "Train Epoch: 9 [32000/60000 (53%)] Loss: 0.089658                   \n",
      "Train Epoch: 9 [35200/60000 (59%)] Loss: 0.036081                   \n",
      "Train Epoch: 9 [38400/60000 (64%)] Loss: 0.111251                   \n",
      "Train Epoch: 9 [41600/60000 (69%)] Loss: 0.013245                   \n",
      "Train Epoch: 9 [44800/60000 (75%)] Loss: 0.009567                   \n",
      "Train Epoch: 9 [48000/60000 (80%)] Loss: 0.006871                   \n",
      "Train Epoch: 9 [51200/60000 (85%)] Loss: 0.028005                   \n",
      "Train Epoch: 9 [54400/60000 (91%)] Loss: 0.032505                   \n",
      "Train Epoch: 9 [57600/60000 (96%)] Loss: 0.035857                   \n",
      "\n",
      "Test set: Average loss: 0.0272, Accuracy: 9911/10000 (99%)\n",
      "\n",
      "Train Epoch: 10 [0/60000 (0%)] Loss: 0.036498                   \n",
      "Train Epoch: 10 [3200/60000 (5%)] Loss: 0.012918                   \n",
      "Train Epoch: 10 [6400/60000 (11%)] Loss: 0.090970                   \n",
      "Train Epoch: 10 [9600/60000 (16%)] Loss: 0.006573                   \n",
      "Train Epoch: 10 [12800/60000 (21%)] Loss: 0.038290                   \n",
      "Train Epoch: 10 [16000/60000 (27%)] Loss: 0.003858                   \n",
      "Train Epoch: 10 [19200/60000 (32%)] Loss: 0.014475                   \n",
      "Train Epoch: 10 [22400/60000 (37%)] Loss: 0.004153                   \n",
      "Train Epoch: 10 [25600/60000 (43%)] Loss: 0.029727                   \n",
      "Train Epoch: 10 [28800/60000 (48%)] Loss: 0.007603                   \n",
      "Train Epoch: 10 [32000/60000 (53%)] Loss: 0.056450                   \n",
      "Train Epoch: 10 [35200/60000 (59%)] Loss: 0.041987                   \n",
      "Train Epoch: 10 [38400/60000 (64%)] Loss: 0.022617                   \n",
      "Train Epoch: 10 [41600/60000 (69%)] Loss: 0.006550                   \n",
      "Train Epoch: 10 [44800/60000 (75%)] Loss: 0.047044                   \n",
      "Train Epoch: 10 [48000/60000 (80%)] Loss: 0.013737                   \n",
      "Train Epoch: 10 [51200/60000 (85%)] Loss: 0.010326                   \n",
      "Train Epoch: 10 [54400/60000 (91%)] Loss: 0.005790                   \n",
      "Train Epoch: 10 [57600/60000 (96%)] Loss: 0.007009                   \n",
      "\n",
      "Test set: Average loss: 0.0273, Accuracy: 9910/10000 (99%)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = Net().to(device)\n",
    "optimizer = optim.Adadelta(model.parameters(), lr=1.0)\n",
    "\n",
    "scheduler = StepLR(optimizer, step_size=1, gamma=0.7)\n",
    "for epoch in range(1, 11):\n",
    "    train(model, device, train_loader, optimizer, epoch)\n",
    "    test(model, device, test_loader)\n",
    "    scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow introduced 41 new symbols\n"
     ]
    }
   ],
   "source": [
    "globals_before_scikit_learn = frozenset(filter(None, map(globals_helper, globals())))\n",
    "globals_i_before_scikit_learn = globals_helper.i\n",
    "globals_helper.i = 0\n",
    "print(\n",
    "    \"TensorFlow introduced\",\n",
    "    globals_i_before_scikit_learn - globals_i_before_tensorflow,\n",
    "    \"new symbols\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "![scikit.learn](https://scikit-learn.org/stable/_static/scikit-learn-logo-small.png)\n",
    "> Built on NumPy, SciPy, and matplotlib.\n",
    "\n",
    "Ref: The next are based off the [official BSD licensed Python code by scikit.learn](https://scikit-learn.org/stable/_downloads/9b3be64651591413a73d3848e0317ffd/plot_mnist_filters.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from sklearn.datasets import fetch_openml\n",
    "from sklearn.exceptions import ConvergenceWarning\n",
    "from sklearn.neural_network import MLPClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Load data from https://www.openml.org/d/554\n",
    "X, y = fetch_openml(\"mnist_784\", version=1, return_X_y=True, data_home=data_dir)\n",
    "X = X / 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# rescale the data, use the traditional train/test split\n",
    "X_train, X_test = X[:60000], X[60000:]\n",
    "y_train, y_test = y[:60000], y[60000:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration 1, loss = 0.32009978\n",
      "Iteration 2, loss = 0.15347534\n",
      "Iteration 3, loss = 0.11544755\n",
      "Iteration 4, loss = 0.09279764\n",
      "Iteration 5, loss = 0.07889367\n",
      "Iteration 6, loss = 0.07170497\n",
      "Iteration 7, loss = 0.06282111\n",
      "Iteration 8, loss = 0.05530788\n",
      "Iteration 9, loss = 0.04960484\n",
      "Iteration 10, loss = 0.04645355\n",
      "Iteration 11, loss = 0.04082169\n",
      "Iteration 12, loss = 0.03828222\n",
      "Iteration 13, loss = 0.03557957\n",
      "Iteration 14, loss = 0.03054891\n",
      "Iteration 15, loss = 0.02924761\n",
      "Iteration 16, loss = 0.02610471\n",
      "Iteration 17, loss = 0.02363894\n",
      "Iteration 18, loss = 0.02208186\n",
      "Iteration 19, loss = 0.01932900\n",
      "Iteration 20, loss = 0.01830387\n",
      "Iteration 21, loss = 0.01639227\n",
      "Iteration 22, loss = 0.01392950\n",
      "Iteration 23, loss = 0.01270193\n",
      "Iteration 24, loss = 0.01234102\n",
      "Iteration 25, loss = 0.01081313\n",
      "Iteration 26, loss = 0.01028644\n",
      "Iteration 27, loss = 0.00896707\n",
      "Iteration 28, loss = 0.00744908\n",
      "Iteration 29, loss = 0.00707946\n",
      "Iteration 30, loss = 0.00573869\n",
      "Iteration 31, loss = 0.00499554\n",
      "Iteration 32, loss = 0.00477064\n",
      "Iteration 33, loss = 0.00395458\n",
      "Iteration 34, loss = 0.00355619\n",
      "Iteration 35, loss = 0.00375497\n",
      "Iteration 36, loss = 0.00304228\n",
      "Iteration 37, loss = 0.00264245\n",
      "Iteration 38, loss = 0.00241425\n",
      "Iteration 39, loss = 0.00234957\n",
      "Iteration 40, loss = 0.00233803\n",
      "Iteration 41, loss = 0.00204653\n",
      "Iteration 42, loss = 0.00199057\n",
      "Iteration 43, loss = 0.00190567\n",
      "Iteration 44, loss = 0.00180530\n",
      "Iteration 45, loss = 0.00175054\n",
      "Iteration 46, loss = 0.00168160\n",
      "Iteration 47, loss = 0.00162517\n",
      "Iteration 48, loss = 0.00159676\n",
      "Iteration 49, loss = 0.00154993\n",
      "Iteration 50, loss = 0.00152799\n",
      "Iteration 51, loss = 0.00146697\n",
      "Iteration 52, loss = 0.00145257\n",
      "Iteration 53, loss = 0.00143422\n",
      "Iteration 54, loss = 0.00135888\n",
      "Iteration 55, loss = 0.00134281\n",
      "Training loss did not improve more than tol=0.000100 for 10 consecutive epochs. Stopping.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "MLPClassifier(hidden_layer_sizes=(50,), learning_rate_init=0.1, max_iter=55,\n",
       "              random_state=1, solver='sgd', verbose=10)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mlp = MLPClassifier(\n",
    "    hidden_layer_sizes=(50,),\n",
    "    alpha=1e-4,\n",
    "    max_iter=55,\n",
    "    solver=\"sgd\",\n",
    "    verbose=10,\n",
    "    random_state=1,\n",
    "    learning_rate_init=0.1,\n",
    ")\n",
    "\n",
    "mlp.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training set score: 1.0\n",
      "Test set score:     0.9731\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    \"Training set score: {training_set_score:>3}\\n\"\n",
    "    \"Test set score: {test_set_score:>10}\".format(\n",
    "        training_set_score=mlp.score(X_train, y_train),\n",
    "        test_set_score=mlp.score(X_test, y_test),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scikit.learn introduced 20 new symbols\n"
     ]
    }
   ],
   "source": [
    "globals_now = frozenset(filter(None, map(globals_helper, globals())))\n",
    "globals_i_now = globals_helper.i\n",
    "globals_helper.i = 0\n",
    "print(\n",
    "    \"scikit.learn introduced\",\n",
    "    globals_i_now - globals_i_before_scikit_learn,\n",
    "    \"new symbols\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "celltoolbar": "Slideshow",
  "colab": {
   "collapsed_sections": [],
   "name": "neural-network-and-data-loading.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "ml-env (py3)",
   "language": "python",
   "name": "ml-env"
  },
  "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.7"
  },
  "nbdime-conflicts": {
   "local_diff": [
    {
     "key": "kernelspec",
     "op": "add",
     "value": {
      "display_name": "ml-env (py3)",
      "language": "python",
      "name": "ml-env"
     }
    },
    {
     "key": "language_info",
     "op": "add",
     "value": {
      "codemirror_mode": {
       "name": "ipython",
       "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.7"
     }
    }
   ],
   "remote_diff": [
    {
     "key": "kernelspec",
     "op": "add",
     "value": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
     }
    },
    {
     "key": "language_info",
     "op": "add",
     "value": {
      "codemirror_mode": {
       "name": "ipython",
       "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.4"
     }
    }
   ]
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}