21como-intro-to-deeplearning/2a-mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6387de74-7f19-446c-9c10-3d201a270a87",
   "metadata": {},
   "source": [
    "# Introduction to Deep Learning\n",
    "\n",
    "### Hands-on 2a: MNIST\n",
    "Filippo Vicentini and Giuseppe Carleo\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PhilipVinc/Lectures/blob/master/21como-intro-to-deeplearning/2a-mnist.ipynb) \n",
    "\n",
    "#### Link to the GitHub repository:\n",
    "https://github.com/PhilipVinc/IntroDeepLearning\n",
    "\n",
    "or \n",
    "\n",
    "https://bit.ly/3ki90ki"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a1b1bab-9c40-4b20-943f-149668966322",
   "metadata": {},
   "source": [
    "---------------\n",
    "\n",
    "The objective of this hands-on is to write and optimise an image-classifier that identifies handwritten digits.\n",
    "\n",
    "We will use for this the [MNIST dataset of handwritten digits](http://yann.lecun.com/exdb/mnist/)\n",
    "\n",
    "![title](images/mnist.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e558598f-2027-45fc-8b92-b8e8ae74215e",
   "metadata": {},
   "source": [
    "## 0 - Install Requirements\n",
    "\n",
    "For this example notebook we will need jax+flax+optax for the machine-learning part.\n",
    "\n",
    "For the dataset, instead, we will be using `tensorflow_datasets`, which is a submodule of `tensorflow` that makes it easy to download and load into memory large datasets (such as MNIST and many others).\n",
    "\n",
    "If you are running notebook locally, you need to also install `tensorflow` to make `tensorflow_datasets` work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d7895a56-6b30-49bc-b07e-1fe5a6a00b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Requirements\n",
    "#!pip install tensorflow_datasets flax jax optax tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f998ee7-4e9d-4538-b125-5222253ad15b",
   "metadata": {},
   "source": [
    "### 0.1 - Utility Functions\n",
    "\n",
    "Don't look here. There is nothing interesting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "561535f1-4d9a-4908-bbb7-e63c94d22d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility functions (don't worry. you don't need to understand this one)\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def show_img(img, ax=None, title=None):\n",
    "  \"\"\"Shows a single image.\n",
    "  \n",
    "  Must be stored as a 3d-tensor where the last dimension is 1 channel (greyscale)\n",
    "  \"\"\"\n",
    "  if ax is None:\n",
    "    ax = plt.gca()\n",
    "  ax.imshow(img[..., 0], cmap='gray')\n",
    "  ax.set_xticks([])\n",
    "  ax.set_yticks([])\n",
    "  if title:\n",
    "    ax.set_title(title)\n",
    "\n",
    "def show_img_grid(imgs, titles):\n",
    "  \"\"\"Shows a grid of images.\"\"\"\n",
    "  n = int(np.ceil(len(imgs)**.5))\n",
    "  _, axs = plt.subplots(n, n, figsize=(3 * n, 3 * n))\n",
    "  for i, (img, title) in enumerate(zip(imgs, titles)):\n",
    "    show_img(img, axs[i // n][i % n], title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "241cfd9e-6a5e-4257-b3b5-2e7bad914a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import \n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55c5988d-c3c3-43d0-abc8-6ceb1c09054d",
   "metadata": {},
   "source": [
    "## 1 - Setting up the dataset\n",
    "First of all, we need to download the dataset.\n",
    "\n",
    "The MNIST dataset is a standard dataset composed of several 28x28 black/white images representing numbers, and a label corresponding to the number that is represented there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "90e05be1-0fd8-4562-8251-f8caa0576d3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/filippovicentini/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/tensorflow_datasets/core/dataset_builder.py:622: get_single_element (from tensorflow.python.data.experimental.ops.get_single_element) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Dataset.get_single_element()`.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-09-17 14:35:43.177745: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "WARNING:tensorflow:From /Users/filippovicentini/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/tensorflow_datasets/core/dataset_builder.py:622: get_single_element (from tensorflow.python.data.experimental.ops.get_single_element) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use `tf.data.Dataset.get_single_element()`.\n",
      "2021-09-17 14:35:43.222810: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2)\n"
     ]
    }
   ],
   "source": [
    "# We use Tensorflow datasets to download and import data in a simple numpy-tensor format\n",
    "# It's just handy. You could use anything else.\n",
    "\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "# Specify the dataset we are interested in\n",
    "ds_builder = tfds.builder('mnist')\n",
    "\n",
    "# Download the data\n",
    "ds_builder.download_and_prepare()\n",
    "\n",
    "# Get the whole dataset's train set\n",
    "train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))\n",
    "test_ds = tfds.as_numpy(ds_builder.as_dataset(split='test', batch_size=-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a7b5869-071c-4431-b830-fd837948577d",
   "metadata": {},
   "source": [
    "The dataset is split into two sub-sets: the training dataset that we will use to 'train' our model, and the 'test' dataset, which the model *never sees* during training, but that we use to check that the model performs well.\n",
    "\n",
    "This is to verify that the model does not simply learn _by heart_ the images in the training dataset, but that it actually _learns_ to generalize and works correctly with images that he did not see before.\n",
    "\n",
    "We can inspect the shape of the training dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dc6521ad-6be7-46ab-9cf4-28d24da2d5d9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset keys: dict_keys(['image', 'label'])\n",
      "The training dataset has shape: (60000, 28, 28, 1) and dtype uint8\n",
      "The test     dataset has shape: (10000, 28, 28, 1) and dtype uint8\n",
      "\n",
      "The training labels have shape: (60000,) and dtype int64\n",
      "The test     labels have shape: (10000,) and dtype int64\n"
     ]
    }
   ],
   "source": [
    "print(\"dataset keys:\", train_ds.keys())\n",
    "print(f\"The training dataset has shape: {train_ds['image'].shape} and dtype {train_ds['image'].dtype}\")\n",
    "print(f\"The test     dataset has shape: {test_ds['image'].shape} and dtype {train_ds['image'].dtype}\")\n",
    "print(\"\")\n",
    "print(f\"The training labels have shape: {train_ds['label'].shape} and dtype {train_ds['label'].dtype}\")\n",
    "print(f\"The test     labels have shape: {test_ds['label'].shape} and dtype {test_ds['label'].dtype}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15665ab5-cb0c-4478-a48e-e0515cbe4fef",
   "metadata": {},
   "source": [
    "We can visualize it to understand it a bit more, using an utility function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e6146cb7-be9d-4206-8b60-ae27699f0444",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAgQAAAILCAYAAACXVIRDAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAAAqaUlEQVR4nO3deZRV1Zn+8eclUICCgBEHaAUHECUgYbAXShADRsthGRMSBRW6iUsCJnEIptFuo5IfGqCNMQFxgEg6y3aIAkYBcUy0oyaIghrURLAQGSLILEMx7N8fVXaXvPuaU/dc6k7fz1p3reLxDBvc4MNx1z4WQhAAAChvjfI9AAAAkH8UAgAAQCEAAAAUAgAAIAoBAAAQhQAAAKgMC4GZVZnZoATHBTM7Lst7ZH0ukATzGMWOOVx4yq4QFDozqzCzt83sw3yPBagvMzvdzJ43s01mVpXv8QD1ZTUmmNnHtZ8JZmb5HldDoBAUnmslrc33IIAsfSLpV6qZx0AxulzS1yWdJKm7pPMkjczngBpK2RYCMzvZzF42s41mttrMJptZxT6HnW1my8xsnZlNMrNGdc4fUfs3+Q1mNt/MOuRgTEdLukTSrWmvhfJQaPM4hPDnEMJvJC1Lcx2Uj0Kbw5KGS7othPBhCGGlpNsk/UvKaxaFsi0EkvZIulrSIZL6ShooafQ+x1wgqbeknpLOlzRCkszsfEnXS/qGpLaSXpT0QOwmZja2dqJHP/sc/sva627Pwc8P5aEQ5zFQH4U2h7tKWlznx4trs9IXQiirj6QqSYMi+VWSZtX5cZB0Vp0fj5b0bO3X8yR9p84/ayRpm6QOdc49rp7jukDSvNqvB0j6MN+/VnwK91Oo87jOtQZJqsr3rxOfwv0U6hxWTUHpUufHnWqvY/n+Ndvfn7J9QmBmnc3sCTNbY2abJd2imoZa14o6Xy+X1K726w6S7qjTLNdLMkntsxzLgZImSvpBNuejfBXSPAayUYBzeKukg+r8+CBJW0NtOyhlZVsIJE2V9I6kTiGEg1Tz2GnflaRH1vn6KEmrar9eIWlkCKF1nU/zEMJL+97EzK43s62ZPrWHdZLUUdKLZrZG0kxJR9T+BumYq58wSlIhzWMgG4U2h/+imgWFnzqpNit55VwIWkraLGmrmXWRNCpyzLVm1sbMjpR0paSHavO7JF1nZl0lycxamdm3YjcJIdwSQmiR6VN72FuqmfA9aj+XSfp77dcrIpcFPlVI81hm1sjMmklqUvNDaxZZIAbUVVBzWNJ/SbrGzNqbWTtJP5Q0Iyc/0wJXzoVgjKShkrZIulf/N8HqekzSQkmLJM2RNF2SQgizJE2Q9GDtI663JFVmO5AQwu4QwppPP6p57LW39sd7sr0uykLBzONa/VWzKHauav4mt13SUymvidJWaHP4bkmPS3qz9npzarOSZ2Xwv0UAAMA/UM5PCAAAQC0KAQAAoBAAAAAKAQAAEIUAAABIalyfg82Mb0lAGutCCG3zOQDmMFJiDqPYZZzDPCFAQ1qe7wEAKTGHUewyzmEKAQAAoBAAAAAKAQAAEIUAAACIQgAAAEQhAAAAohAAAABRCAAAgCgEAABAFAIAACAKAQAAEIUAAACIQgAAAFTP1x8DAFBsOnfu7LIpU6a4bODAgdHzZ8yY4bLRo0e7bMeOHfUfXAHhCQEAAKAQAAAACgEAABCFAAAAiEIAAADEdxkAAErcKaec4rKvfvWrLgshRM8fPny4y/bs2eOyK664wmXV1dVJhlgQeEIAAAAoBAAAgEIAAABEIQAAAGJRYcEYPHhwNH/44YddNnLkSJfde++9OR8TkEnz5s1dduedd0aPPeCAA1w2ZMgQl+3duzf9wFD2zjrrLJf9/Oc/z/l9RowY4bIlS5a47Pbbb8/5vfcXnhAAAAAKAQAAoBAAAABRCAAAgFhUWDCGDh0azWM7Zx188MH7ezjA/zIzl919990uu+SSSxJf89Zbb3XZokWL6jUuILZgddy4cS5r2bJlQwxHN9xwg8tYVAgAAIoKhQAAAFAIAAAAhQAAAIhFhXnRoUMHl1VWVkaPXbhwocv++7//O+djAjI58cQTXVafBYSbN2922ccff5xqTIAkPfrooy7r3bu3yzK91nhfmRa29ujRI9H5jRsX939SeUIAAAAoBAAAgEIAAABEIQAAACqTRYWxndYySbr4JI0f/OAHLquoqIgeu2zZMpetWLEi52MCMvnWt76V6vwPPvjAZcxh1Ndll13msgEDBmR9vdifraeddlr02NjixUGDBrkstqjw2GOPddnSpUuTDLHB8YQAAABQCAAAAIUAAACIQgAAAFQmiwpjC08yvZLyu9/9rsteeeWVnI6nW7duiY/llbDItyuvvDLRcbt3747msVcdA59n2LBhLps8ebLLmjRpkuh67733nsvOPPNMl23dujV6ftKdNZs2beqy2H9/WFQIAAAKFoUAAABQCAAAAIUAAACoTBYVbt++3WWZFvbFdqpKs6jwn/7pnxLdY8uWLdHzf/3rX2d9b6C+Wrdu7bJWrVolOnft2rXR/IEHHkgzJJSw9u3bR/PrrrvOZUkXEK5evdplI0eOdFlVVVWi66U1cOBAl02fPr1B7l1fPCEAAAAUAgAAQCEAAACiEAAAAFEIAACAyuS7DD766KO83fuCCy5wWWy17Kuvvho9P7ZiFthfxo0bl/W5b775Zg5HglIT+46ruXPnRo/t3Llz1veZOHGiy37/+99nfb20unbtmrd71xdPCAAAAIUAAABQCAAAgCgEAABAZbKo8OCDD87bvdu1a5fouHwuegE+ddlll2V97h133JHDkaDUxLbrTbvgbtGiRS6bMWNGqmvmWqGN5/PwhAAAAFAIAAAAhQAAAIhCAAAAVCaLCmO7BZpZzu8Te7f3qFGjEt37V7/6Vc7HA+wvGzdudNnTTz/d8ANBQTrzzDNddsYZZ6S65ieffOKyr3/96y7btGlTqvvExP7MTvrfkC1btuR6OPsNTwgAAACFAAAAUAgAAIAoBAAAQCW4qLBp06Yuu/zyy10WQoieP2TIEJd17NjRZbHdD7t37+6yli1buuz111932fvvvx8dD7C/9OjRw2WxV3PHTJkyxWW7d+9OOyQUodatW7ts2rRpLsv0Z25MbAHh8OHDXbZixYrE10yioqIimh966KEui/189uzZ47KVK1emH1gD4QkBAACgEAAAAAoBAAAQhQAAAKgEFxUOHTrUZfV5/XG3bt1cFlssWJ8FMvv66U9/6rK9e/dmfT0gGxMnTnRZ48b+j4Rdu3a5LLaoEOUptpA76WvfM3n88cddNmvWrFTXTOL73/9+NB8wYECi83fs2OGyefPmpRlSg+IJAQAAoBAAAAAKAQAAEIUAAACoBBcV9unTx2Xbtm1zWabXDa9atcpl69evd9m6detc9sgjjyQZop588slExwG50qFDB5f17dvXZbHFsu+9957L1qxZk5uBoaj079/fZb/73e+yvl6mxdlz587N+pppnHvuuanOj+102Lt3b5e9+uqrqe6zv/CEAAAAUAgAAACFAAAAiEIAAABUgosKR48enShLa/DgwS4zM5fNnDnTZZs3b875eIDPM2bMGJcdeOCBic6N7WiI8jR58mSXxV7xntSyZcui+f3335/1NZM6/fTTXXbqqaemumZsx9kNGzakumZD4gkBAACgEAAAAAoBAAAQhQAAAIhCAAAAVILfZdBQhg4d6rLYNpwLFixoiOEAnyvp+9xjZsyYkbNxoLg9/PDDLrv55puzvt5DDz2UZjiJXXLJJS676aabXPaFL3wh1X1uvPFGly1dujTVNRsSTwgAAACFAAAAUAgAAIAoBAAAQCwqzNppp53mstiiwj/84Q8NMRzgf5100kku69y5c6JzZ8+enePRoJSsWbMmp9erqKiI5t/5zndc1qtXL5etWLHCZbEFtP37909875jYlsSxBZa33XZb4msWIp4QAAAACgEAAKAQAAAAUQgAAIBYVJhIz549Xda4sf+le+qpp1z2yiuv7JcxAZnE3lnfpEmTROeOGzcu18MBMhozZkzOr9mokf97bmxRYMzf//73aP6zn/3MZf/5n/9Zv4EVAZ4QAAAACgEAAKAQAAAAUQgAAIBYVJjIhAkTXNayZUuXDRw40GWjRo1y2dSpU3MzMJS1Fi1aRPNjjjkm0fkbNmxw2ZIlS1KNCaVt7ty5LovNmRNPPLEhhhMV2zF23bp1LrvnnntcNn369Og1q6qqUo+rGPCEAAAAUAgAAACFAAAAiEIAAADEosJEYotUYtlf/vIXlz3yyCP7ZUxAplcaH3HEEYnOf+mll1xWXV2dakwobatWrXJZ7NXCF110kctuuOEGlx122GGpxjNjxgyXPfHEEy57+eWXXZbrVzmXAp4QAAAACgEAAKAQAAAAUQgAAIBYVJjICSec4LJPPvnEZd/4xjdctnbt2v0yJuC8885Ldf60adNyNBKUs9iOl7HdWNmhtfDxhAAAAFAIAAAAhQAAAIhCAAAAJFlsx72MB5slP7iExF6dGVtI06lTp4YYTjFbGELonc8BlNIcPuSQQ6J5bMfM2O/zY4891mWxxbL4DOYwil3GOcwTAgAAQCEAAAAUAgAAIAoBAAAQhQAAAIitixPJtJobyKfYd79I6d8xD6A88YQAAABQCAAAAIUAAACIQgAAAEQhAAAAohAAAABRCAAAgCgEAABAFAIAACAKAQAAEIUAAACIQgAAAEQhAAAAohAAAADV//XH6yQt3x8DQVnokO8BiDmMdJjDKHYZ57CFEBpyIAAAoADxvwwAAACFAAAAUAgAAIAoBAAAQBQCAAAgCgEAABCFAAAAiEIAAABEIQAAAKIQAAAAUQgAAIDKsBCYWZWZDUpwXDCz47K8R9bnAkkwj1HsmMOFp+wKQaEys3lmtrXOp9rM3sz3uID6MLOmZnaXmf3dzNab2eNm1j7f4wKSMrPTzex5M9tkZlX5Hk9DohAUiBBCZQihxacfSS9J+m2+xwXU05WS+krqLqmdpA2SfpnXEQH184mkX0m6Nt8DaWhlWwjM7GQze9nMNprZajObbGYV+xx2tpktM7N1ZjbJzBrVOX+Emb1tZhvMbL6Z5ew96WbWUdJXJP1Xrq6J0lSA8/hoSfNDCH8PIeyQ9JCkrimviRJWaHM4hPDnEMJvJC1Lc51iVLaFQNIeSVdLOkQ1f6MZKGn0PsdcIKm3pJ6Szpc0QpLM7HxJ10v6hqS2kl6U9EDsJmY2tnaiRz8ZxjZM0oshhKoUPz+Uh0Kbx9MlnWpm7czsAEkXS5qXm58qSlShzeHyFUIoq4+kKkmDIvlVkmbV+XGQdFadH4+W9Gzt1/MkfafOP2skaZukDnXOPS7FGN+T9C/5/rXiU7ifQp3HklpJerD23N2SXpd0cL5/vfgU3qdQ53Cdaw2SVJXvX6eG/JTtEwIz62xmT5jZGjPbLOkW1TTUulbU+Xq5av6fqCR1kHRHnWa5XpJJSr14ysz6STpc0iNpr4XSV4DzeIqkppK+KOlASTPFEwJ8jgKcw2WrbAuBpKmS3pHUKYRwkGoeO9k+xxxZ5+ujJK2q/XqFpJEhhNZ1Ps1DCC/texMzu36f7x74zCcyruGSZoYQYv8M2FehzeMekmaEENaHEHaqZkHhyWa27x/wwKcKbQ6XrXIuBC0lbZa01cy6SBoVOeZaM2tjZkeqZvX0Q7X5XZKuM7OukmRmrczsW7GbhBBuCXW+e2DfT91jzay5pG9LmpGTnyHKQaHN4wWShtVeq4lqHu+uCiGsy81PFyWooOawmTUys2aSmtT80JpFFjmWpHIuBGMkDZW0RdK9+r8JVtdjkhZKWiRpjmoWTCmEMEvSBEkP1j7iektSZQ7G9HVJGyU9n4NroTwU2jweI2mHpL9JWivpbNUsCAMyKbQ53F/SdklzVfM0Yrukp1JesyhY7eIJAABQxsr5CQEAAKhFIQAAABQCAABAIQAAAKIQAAAASY3rc7CZ8S0JSGNdCKFtPgfAHEZKzGEUu4xzmCcEaEjL8z0AICXmMIpdxjlMIQAAABQCAABAIQAAAKIQAAAAUQgAAIAoBAAAQBQCAAAgCgEAABCFAAAAiEIAAABEIQAAAKIQAAAAUQgAAIAoBAAAQBQCAAAgCgEAABCFAAAAiEIAAABEIQAAAJIa53sAAEpDy5Yto/kVV1zhsltuucVlq1evdtmJJ57osk2bNmUxOpSzpk2buuyPf/yjy4455pjo+YMGDXLZa6+9ln5gBYYnBAAAgEIAAAAoBAAAQBQCAAAgFhUCJSe2MCq24O+b3/ymy5o1a5boerFs8eLF0fEMGzbMZSEElx1xxBGJxsOiQtRXmzZtXNazZ8/E58+YMcNlffr0cdnOnTvrNa5CwxMCAABAIQAAABQCAAAgCgEAAFCZLCp8/vnnXTZgwIDosRMmTHDZ2LFjcz0kIKPYrmpHH320y6ZOnRo9/8tf/rLLDjroIJfFFvYlZWYuO+mkk7K+HrA/3XTTTanOj/3+adu2rcs+/PDDVPfJN54QAAAACgEAAKAQAAAAUQgAAICKfFFhbGHT8ccf77LYIqu9e/dGr3nllVe6bM+ePS6bOXOmy2KLtN59993offb11a9+NZrHdoSrqqpy2dy5c122a9euRPdG/sT+/T788MMui83h+oi96nXp0qUumzNnjss2btzosvnz56caT8zKlStdtmPHjpzfB6XtggsucNnIkSNdVp9FtUuWLHFZsS8gjOEJAQAAoBAAAAAKAQAAEIUAAACoyBcVduvWzWWvv/56qmtWVFS4LLZTYaHtXvjiiy+6LLa4ZsOGDQ0xHERUVla6LLaIL2bLli0ui+3AKUmTJk1yWWxRYVKXXnppouO2bt2a+Jqx1zE/++yzLuNVx6ivLl26ZH1ubGGrJI0YMSLraxYTnhAAAAAKAQAAoBAAAABRCAAAgIpoUWGHDh1cNnv27Kyvt3nz5mge28GwTZs2Lku6y1VsN8X67JAVW1TVqlUrl/Xv399l48ePd9no0aMT3xvZ69q1q8ti8zU2F/785z+7bPDgwS7LtAAq1xYuXOiyKVOmuCzTzm1XX321y1q0aOGyUaNGZTE64LOGDRuW9bn33HNPNF+zZk3W1ywmPCEAAAAUAgAAQCEAAACiEAAAABXRosLLL7/cZbGFhjETJkxw2c9//vPosdu3b3dZplcTN4S33nrLZX/9618TnRvbDQ4No3v37i5r3DjZb7ezzz7bZfncYTL26tfvf//7LhsyZEj0/LZt27ps27ZtLov93gM+T2zOderUKevrrVixIs1wih5PCAAAAIUAAABQCAAAgCgEAABAFAIAAKAC/S6Dfv36ueyqq67K+nq/+MUvXPbRRx8lPv+xxx7L+t5pHXfccYmOi22Be+aZZ7qsWbNmLtuxY0f9B4bP9eUvfznrc3v16uWyZ555Js1wGsS1116b+NjbbrttP44E5eKGG25wWaNGyf6eu3btWpfNnDkz9ZiKGU8IAAAAhQAAAFAIAACAKAQAAEAFuqgwtuAvthiuurraZZMnT3ZZPrd9TWvo0KGJjjMzl82fP99lLCBsGPfff7/LxowZk+jcp556KtFxTzzxRDSPzffVq1e7bPbs2S575ZVXEt17+PDhLuvRo0f02Ni75G+66aZE9wE+T5s2bbI+9/bbb3fZ5s2b0wyn6PGEAAAAUAgAAACFAAAAiEIAAABUoIsK//a3v7msa9euLtuyZYvLVq5cuV/GlC8HHXRQouNiOxUif5YsWeKyc845x2Xjx493Wezf+dFHH53oepnEFp1effXVLvv4448TXa9Vq1YuyzQHP/jgA5eddNJJLlu8eHGie6M8XXrppS479NBDE527detWl7FbpscTAgAAQCEAAAAUAgAAIAoBAABQgS4qjC1Oeuedd/IwkoYzbty4aH7FFVckOj+2wHL69OmpxoTs7dq1y2Xz5s1LlLVs2dJl9VlU2Lp1a5fFFhXGfp/FdiBs27Zt1teTpD59+rjstddec9mbb77pstgrlZ9++unofVDazjjjDJclfdXx7t27XRb7PVrueEIAAAAoBAAAgEIAAABEIQAAACrQRYWl7ic/+YnLrrvuuuixscVbMdOmTXPZ73//+3qNC4UhtkD0jTfeSJTVx6BBg1w2cuTIROcuXLjQZZMmTYoee/bZZ7ts4MCBLuvevbvLfvvb37qsZ8+eLlu2bFn03ihOsVdpn3feeS5LukPrxIkT0w6pLPCEAAAAUAgAAACFAAAAiEIAAADEosKcii0AvPjii132wx/+MNG5mTz33HMuGzt2bOLzUX5uuukml8V2AWzevLnL/vjHP7ostqNhpoV9Dz/8sMv69evnshdeeMFlsVdBt2jRInoflI5OnTq5LPbK7aTmzJmTZjhlgycEAACAQgAAACgEAABAFAIAACAWFWatY8eOLrv55ptddumll7os6e5akvTuu++67F//9V9dFnu9J0pbkyZNovns2bNdVllZ6bLYPLz//vtd9r3vfc9lmzZtSjDCzGK7Dca89dZbLluyZEmqe6P8nHrqqS5Lu9NnKeIJAQAAoBAAAAAKAQAAEIUAAACIQgAAAMR3GSTypS99yWUTJkxw2VlnneWypN9RMGvWrGg+ZswYl3344YeJronidPjhh7ts8ODBLrvwwgsTn79z506XxeZwLNu+fXv0PkkdeOCBLhs1alSic2+99VaX8R01pe+iiy7K6fUmTpzosqlTp+b0HqWAJwQAAIBCAAAAKAQAAEAUAgAAIBYVOu3bt3fZ9OnTXda7d++s7xHbCpYFLqWvefPmLrvzzjtdNnz4cJfVZ7vrZ555xmXXXXedyx555JHE10yjW7duLuvcubPLVq5c6bK5c+fulzGhsB1zzDH5HkJZ4gkBAACgEAAAAAoBAAAQhQAAAIhFhc6VV17psj59+rgstshr69atLhs7dqzLpk2bluXoUCz++Z//2WWTJ092Wa9evVxmZi772c9+5rLx48dH771hw4YkQ8y5o446KprPmTPHZbGf409+8hOXbdq0Kf3AUPYy7QSLz+IJAQAAoBAAAAAKAQAAEIUAAACojBcVxhYwSfFFhbEFhLHFTrHd4O6+++4sRodi981vftNlPXv2dFnSHQjffvttl7Vs2TJ6bKbFfbl0yimnuCw2/yWpdevWLlu6dKnL7rnnntTjQnE57bTTovkJJ5yQ9TXfeOMNlw0bNizr65UTnhAAAAAKAQAAoBAAAABRCAAAgMpkUWFsUdPQoUOjxzZu7H9JYruqPfjggy5jASE+NWPGDJedd955Lou9BjgmtuAu046Ebdq0cVlsDtfnlcpJrlddXR09NvYK40y//1BeDjjggGheUVGR9TVjO2MiGZ4QAAAACgEAAKAQAAAAUQgAAIDKZFHhkCFDXNaxY8fE5y9btsxlt9xyS5ohocQtWbLEZT169HBZ//79XXbqqae6LDZfmzdvHr334MGD//EAM4iNe+HChS5bs2aNy2bPnh295iuvvJL1eFDann766Wh+1VVXueyMM85wWWzHyz/84Q+px1WueEIAAAAoBAAAgEIAAABEIQAAAJKsPruVmVn2W5vlUWVlpcsy7WYV+/UYNWqUy3hVa1YWhhB653MAxTqHUTCYwyh2GecwTwgAAACFAAAAUAgAAIAoBAAAQBQCAACgMtm6+LnnnnPZn/70p+ixxx9/fKLzAQAoJTwhAAAAFAIAAEAhAAAAohAAAACVyaLCnTt3uqxv3755GAkAAIWJJwQAAIBCAAAAKAQAAEAUAgAAIAoBAAAQhQAAAIhCAAAARCEAAACiEAAAANV/p8J1kpbvj4GgLHTI9wDEHEY6zGEUu4xz2EIIDTkQAABQgPhfBgAAgEIAAAAoBAAAQBQCAAAgCgEAABCFAAAAiEIAAABEIQAAAKIQAAAAUQgAAIAoBAAAQGVYCMysyswGJTgumNlxWd4j63OBJJjHKHbM4cJTdoWgUJnZtWb2lpltMbP3zezafI8JqC8zm2dmW+t8qs3szXyPC0jKzE43s+fNbJOZVeV7PA2JQlA4TNIwSW0knSXpe2Z2UX6HBNRPCKEyhNDi04+klyT9Nt/jAurhE0m/klR2fykr20JgZieb2ctmttHMVpvZZDOr2Oews81smZmtM7NJZtaozvkjzOxtM9tgZvPNLNV70kMIE0MIr4UQdocQ3pX0mKRT01wTpa/Q5vE+Y+so6SuS/itX10TpKbQ5HEL4cwjhN5KWpblOMSrbQiBpj6SrJR0iqa+kgZJG73PMBZJ6S+op6XxJIyTJzM6XdL2kb0hqK+lFSQ/EbmJmY2snevST4RxTzR+kf0n3U0QZKNh5rJonXi+GEKpS/PxQ+gp5DpeXEEJZfSRVSRoUya+SNKvOj4Oks+r8eLSkZ2u/nifpO3X+WSNJ2yR1qHPucSnGeLOkxZKa5vvXi09hfopkHr8n6V/y/WvFpzA/hT6HJQ2SVJXvX6eG/JTtEwIz62xmT5jZGjPbLOkW1TTUulbU+Xq5pHa1X3eQdEedZrleNWsA2udgXN9Tzd+szgkh7Ex7PZS2Ap7H/SQdLumRtNdCaSvUOVyOyrYQSJoq6R1JnUIIB6nmsZPtc8yRdb4+StKq2q9XSBoZQmhd59M8hPDSvjcxs+v3WXX9mc8+x46QNFbSwBDChzn6eaK0Fdw8rjVc0swQQuyfAXUV6hwuO+VcCFpK2ixpq5l1kTQqcsy1ZtbGzI6UdKWkh2rzuyRdZ2ZdJcnMWpnZt2I3CSHcEuqsut738+lxZnaxaprxGSGEslvMgqwV1DyuvU5zSd+WNCMnP0OUuoKaw2bWyMyaSWpS80NrFlnkWJLKuRCMkTRU0hZJ9+r/Jlhdj0laKGmRpDmSpktSCGGWpAmSHqx9xPWWpMqU4/l/kr4oaUGd1npXymui9BXaPJakr0vaKOn5HFwLpa/Q5nB/SdslzVXN04jtkp5Kec2iYLWLJwAAQBkr5ycEAACgFoUAAABQCAAAAIUAAABIalyfg82MFYhIY10IoW0+B8AcRkrMYRS7jHOYJwRoSMvzPQAgJeYwil3GOUwhAAAAFAIAAEAhAAAAohAAAABRCAAAgCgEAABAFAIAACAKAQAAEIUAAACIQgAAAEQhAAAAohAAAABRCAAAgCgEAABAFAIAACAKAQAAEIUAAACIQgAAAEQhAAAAkhrnewC59stf/tJlvXr1Snz+k08+6bLly5e7bM2aNS6bP39+4vsAAIpDly5dXLZo0SKXLViwwGVf+cpX9seQ9gueEAAAAAoBAACgEAAAAFEIAACAimhRYdOmTV02ZcoUl40YMSLVffr27euyEILL9u7d67JXX33VZT/+8Y9d9tRTT2U5OgBAQ+vXr5/LvvCFL7jsS1/6ksuOPfZYly1dujQ3A8sxnhAAAAAKAQAAoBAAAABRCAAAgIpoUeGPfvQjl6VdQBgTW0AY06iR71Inn3yyy2ILH4cMGeKy2IJEoKH179/fZb/4xS9cdvzxx7vsmmuuiV5z6tSp6QcGNIDKyspoHlsc3rix/8/ntm3bXLZjx470A2sgPCEAAAAUAgAAQCEAAACiEAAAABXRosJ27dolOm7mzJkuW7x4scu2bt0aPf83v/mNy2K7JN5///0uO+WUU1wW26XqnnvucVmfPn2i49mzZ080R3lp0aKFy3bv3h09NjbnYjuoxeZrbFFht27dkgwxusunxKJCFKbYToOjR4+OHnvkkUe6LPZn87PPPuuylStXZjG6/OAJAQAAoBAAAAAKAQAAEIUAAACoiBYVxhYmffDBBy6bOHGiy/bHwrwBAwa47Mknn3TZ1772NZf16NHDZd/97nej94ntdIjidMABB7hs7ty5ic6trq522XHHHRc99rDDDnNZs2bNXGZmLku6U2fMli1bsj4XaGjjxo1z2bnnnpv4/AULFrhs2LBhqcaUbzwhAAAAFAIAAEAhAAAAohAAAABJVp9FRGaW/YqjMtCvXz+XPfPMMy6rqKhw2UcffRS9ZuyVyrHFlEViYQihdz4HkM85/MUvftFlsX/vaRf7xV63GtvV8L777nNZbIwXXnihy2K7vMVekyxJV199dTQvUmU9h4tVly5dXLZw4UKXxRb+SvGF6eedd57L5s2bl8XocqN3bz8tX3311dihGecwTwgAAACFAAAAUAgAAIAoBAAAQEW0U2Ex+J//+R+XTZo0yWX//u//7rJDDz00es2OHTu6rIgXFZa12E5+55xzTs7vU1VV5bLNmze7bNWqVYmuF1vYGtslMXYPoKHFFgbeeOONiY7L5IEHHnBZPhcQxmzbti31NXhCAAAAKAQAAIBCAAAARCEAAACiEAAAAPFdBvvdY4895rLYdxlk0q1bN5e98MILqcaE/KiurnbZk08+mYeRZNa6dWuXxVZjx7ZXjn13A9DQYlsKX3TRRYnOXb9+fTS/++67U42pISxZsiT1NXhCAAAAKAQAAIBCAAAARCEAAABiUWHBiy2Queuuu1wWe183UF/HH3+8y9q1a+eyEILLTj/99Og177vvvvQDAyIGDBjgsl//+teJzo3N4WuuuSZ6bGxb+lLEEwIAAEAhAAAAFAIAACAKAQAAEIsK97u1a9e6bN26dS475JBDoufH3jtfUVHhsu3bt2cxOuCzYjtjJvXmm2/mcCTAP/bjH//YZU2bNk107uTJk12WdEFiqeIJAQAAoBAAAAAKAQAAEIUAAACoyBcVxl7VGttVLWb37t3R/K9//WuaITlt27Z1WaYFhDG33367y1hAiP0lzaLCXP/eAeoaNWqUy/r165fo3OXLl7vsP/7jP1KPqdTwhAAAAFAIAAAAhQAAAIhCAAAAVESLCisrK10WW3DXuXPnRNerrq6O5jfffLPL5s6d67LFixcnus/555+f6LhM2P0NmWSaW7GFge+//77LLr74Ypd16dIl6/HEdn6TpF69erkstsMc8KnDDjvMZf/2b//msiZNmrgstmB80qRJLtu8eXOWoytdPCEAAAAUAgAAQCEAAACiEAAAABXRosLHHnvMZY0bZz/82CuEJWn8+PEuu/HGG132+OOPu2zOnDku+9GPfpRoPLt27YrmO3fuTHQ+Stu0adNcduGFF0aPPfDAAxNd08xcFkJIdG5sUW6m31NAJpn+DI+9hrhDhw6JrhlbiD1lypT6DaxM8YQAAABQCAAAAIUAAACIQgAAACRZ0kVEkmRmyQ/OsdhOa0kXmaxevdplmXYA/NrXvla/geXI22+/Hc27du3awCOp0bNnz2h+5JFHuiy24DODhSGE3tmPKr18zuE0Onbs6LKpU6dGjz322GNdtm7dOpfFFhUeddRRLjv88MNdNn/+fJdlWuS4ZcuWaF6kmMM51KNHj2j++uuvJzo/tivht7/9bZfNmjWrXuMqcRnnME8IAAAAhQAAAFAIAACAKAQAAEAUAgAAoCLaunjcuHEuu/vuu10W2wpz4cKFLrv88suj92nWrJnLXnzxRZe1b98+en62OnXqFM1XrlzpsiVLlrjsxBNPzOl4WrduHc1jK9MPOOCAnN4bXlVVlcsqKyujx7Zs2dJlSVf6P/fccy6LfZdBly5dsr4H8Kkbbrgh1fl33HGHy/iOguzxhAAAAFAIAAAAhQAAAIhCAAAAVESLCu+77z6XxRZa3XvvvS4799xzXbZq1arofV5++WWXHXzwwQlGmE6m94IfccQRibI0PvjgA5fNnDkzeuxtt92W03sj95Iu7otth9ynT59E5zZp0qQ+QwLUu7ffLTfTwtikZs+enep8fBZPCAAAAIUAAABQCAAAgCgEAABARbSoMOb555932TXXXOOySZMmuSy2oEqS+vbtm+je1dXVLou9w3v8+PEue+eddxLdI5MRI0a4rKKiwmWxHRoXLFjgso0bN7ps3bp12Q0OReOEE05wWdJdJx999NFcDwclbsyYMS5r3rx54vOfeeYZl/3pT39KNSZ8Fk8IAAAAhQAAAFAIAACAKAQAAEBFvqgw5ne/+12irEePHtHzu3fvnug+L7zwgstiOyfuD9dff32D3AelLbawNvZ665jVq1fneDQoJYceeqjLki7YzuSnP/2py3bt2pXqmvgsnhAAAAAKAQAAoBAAAABRCAAAgEpwUWFSixYtqlcOlJpDDjnEZSGEROfGdgkFPtWmTRuXHXXUUamuuXfv3lTn4x/jCQEAAKAQAAAACgEAABCFAAAAqIwXFQLlrnPnzomOi+3A+cYbb+R4NCgl77//vsvuvPNOl40ePTp6/vr16122YsWK9APD5+IJAQAAoBAAAAAKAQAAEIUAAACIRYUA/oFPPvnEZTt27MjDSFAsqqurXXbFFVckypA/PCEAAAAUAgAAQCEAAACiEAAAAFEIAACA+C4DAP/Ao48+mu8hAGgAPCEAAAAUAgAAQCEAAACiEAAAAEkWQkh+sFnygwFvYQihdz4HwBxGSsxhFLuMc5gnBAAAgEIAAAAoBAAAQBQCAAAgCgEAABCFAAAAiEIAAABEIQAAAKIQAAAA1f/1x+skLd8fA0FZ6JDvAYg5jHSYwyh2GedwvbYuBgAApYn/ZQAAACgEAACAQgAAAEQhAAAAohAAAABRCAAAgCgEAABAFAIAACAKAQAAkPT/AQB+5esB6/mhAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 648x648 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_img_grid(\n",
    "    [train_ds['image'][idx] for idx in range(9)],\n",
    "    [f'label={train_ds[\"label\"][idx]}' for idx in range(9)],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d357d877-89e9-41ea-8301-d5fe22693a90",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The mean     of the data stored in the images are:  33.318421449829934\n",
      "The variance of the data stored in the images are:  6172.850482291347\n"
     ]
    }
   ],
   "source": [
    "print(\"The mean     of the data stored in the images are: \", np.mean(train_ds['image']))\n",
    "print(\"The variance of the data stored in the images are: \", np.var(train_ds['image']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c9f1e8-b449-4589-a826-3dd5771db7ce",
   "metadata": {},
   "source": [
    "We have seen that the data is stored in uint8 (an *unsigned* 8-bit integer which can take values from 0 to 255 ).\n",
    "\n",
    "However it is often preferable when working with Neural Networks to work with floating-point values with values around 0 and variance approximately 1. The reasons are 2:\n",
    "\n",
    " - modern CPUs (and to an extent GPUs) are often faster at working with batches (blocks) of floating-point numbers rather than integers [caveats apply]\n",
    " - Many nonlinear functions used in machine-learning have the nonlinear crossover aroud ~0 or ~1/2, so we want our data to be spread around those values\n",
    " - Most research about how to initialize neural-network layers assumes that the input data has mean 0 and variance 1, so to exploit those results we have to rescale our data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "83684f21-5c8e-423d-8c61-9973988fc8cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Therefore... let's convert the data!\n",
    "train_ds['image'] = jnp.float32(train_ds['image']) / 255.\n",
    "test_ds['image'] = jnp.float32(test_ds['image']) / 255."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e5a79b56-21ac-49ec-9485-2590a18a2748",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The training dataset has shape: (60000, 28, 28, 1) and dtype float32\n",
      "The test     dataset has shape: (10000, 28, 28, 1) and dtype float32\n",
      "The mean     of the data stored in the images are:  0.13066326\n",
      "The variance of the data stored in the images are:  0.09471572\n"
     ]
    }
   ],
   "source": [
    "print(f\"The training dataset has shape: {train_ds['image'].shape} and dtype {train_ds['image'].dtype}\")\n",
    "print(f\"The test     dataset has shape: {test_ds['image'].shape} and dtype {train_ds['image'].dtype}\")\n",
    "\n",
    "print(\"The mean     of the data stored in the images are: \", np.mean(train_ds['image']))\n",
    "print(\"The variance of the data stored in the images are: \", np.var(train_ds['image']))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bf695fb-0fab-49f3-81fd-9a23fc7ea311",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "98e56e14-2a4e-4fad-9d47-809ba4b0a354",
   "metadata": {},
   "source": [
    "## 2 - The model (Neural Network)\n",
    "\n",
    "We want now to define the Model.\n",
    "We will use Flax to do that.\n",
    "\n",
    "We want our network to return a probability distribution for the input to correspond to one of several output labels.\n",
    "\n",
    "e.g: if $x$ is an image, then $f : \\mathbb{R}^{28\\times 28}\\rightarrow \\mathbb{R}^{10}$ and $f^{(i)}(x)$ is the probability that the image $x$ represents a $i\\in[0,9]$\n",
    "\n",
    "To make the output of the network a probability distribution, we can use a softmax function, defined as\n",
    "\n",
    "$$\n",
    "\\sigma_i(x) = \\frac{e^{x_i}}{\\sum_i^K e^{x_i} }  \\text{   for  } i\\in [1,K] \\text{ and } x\\in\\mathbb{R}^K\n",
    "$$\n",
    "\n",
    "We want to use a Feedforward network with 2 Dense Layers, relu-nonlinearity and output softmax using Flax.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8af5cf83-b3e2-4a1d-b204-73f383952134",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "# We import flax.linen as nn\n",
    "# The reason is that flax.nn is old and deprecated and will be removed one day\n",
    "import flax.linen as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "7bf8a1ee-844e-4342-8578-6a0df21cc422",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A Flax model must be a class sub-classing nn.Module\n",
    "class Model(nn.Module):\n",
    "    \n",
    "    # We can have some attributes of the Model. \n",
    "    # Those are considered compile-time constants and must be hashable\n",
    "    # They are useful to define some variables that might be changed often\n",
    "    hidden_width : int = 1024\n",
    "    \"\"\"\n",
    "    The width of the hidden dense layers in the neural network.\n",
    "    \"\"\"\n",
    "    \n",
    "    n_classes : int = 10\n",
    "    \"\"\"\n",
    "    Number of output classes for the classifier \n",
    "    \"\"\"\n",
    "    \n",
    "    # The body of the model must be defined using the `@nn.compact` decorator.\n",
    "    # Just think of it as boilerplate, and if you are curious, check out\n",
    "    # Flax documentation\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        \"\"\"\n",
    "        This function should evaluate the result of the model for an input image\n",
    "        x or a batch of images x.\n",
    "        \n",
    "        x has shape (28,28,1) or (N, 28, 28, 1)\n",
    "        \"\"\"\n",
    "        # we first ensure a single image is a 4-tensor\n",
    "        x = x.reshape(-1, 28, 28, 1)\n",
    "        # Then flatten the x/y/channels dimensions\n",
    "        x = x.reshape((x.shape[0], -1))\n",
    "        \n",
    "        # First dense layer with large output\n",
    "        x = nn.Dense(features=self.hidden_width)(x)\n",
    "        \n",
    "        # First nonlinear activation function\n",
    "        x = nn.relu(x)\n",
    "        \n",
    "        # Last dense layer with n_classes output\n",
    "        x = nn.Dense(features=self.n_classes)(x)\n",
    "        \n",
    "        # transform the output to be a log-probability distribution\n",
    "        return nn.log_softmax(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b1324af-5409-4a2c-9fc3-1ae9e09e009b",
   "metadata": {},
   "source": [
    "Let's initialize the model:\n",
    " \n",
    " - We need a seed for the RNG that generates the initial weights\n",
    " - We need a sample input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "add7d9bd-8028-43e0-b62f-d1c8f05b0de7",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 123\n",
    "\n",
    "model = Model(hidden_width = 1024, n_classes=10)\n",
    "\n",
    "key = jax.random.PRNGKey(seed)\n",
    "sample_input = jnp.ones([1, 28, 28, 1])\n",
    "\n",
    "pars = model.init(key, sample_input)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7604bf44-3b6e-4550-93bf-af5253a50364",
   "metadata": {},
   "source": [
    "we can inspect the parameters `pars`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "2ad04a49-1d07-4f16-b419-0cf3cf34ee78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1024, 10)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pars['params']['Dense_1']['kernel'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "b3e0b121-050c-443e-ac7f-9c639b2fb4e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[0.07704171, 0.08986449, 0.07054991, 0.04636909, 0.30741608,\n",
       "              0.064256  , 0.10086796, 0.06952595, 0.10017835, 0.07393045]],            dtype=float32)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# sample application:\n",
    "jnp.exp(model.apply(pars, jnp.ones([1, 28, 28, 1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "aef61d55-eb5d-46bd-b66e-8c4323e02a81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(1.0000001, dtype=float32)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.exp(model.apply(pars, jnp.ones([1, 28, 28, 1]))).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ebdbb9-8e06-4620-89f8-572f752ff73e",
   "metadata": {},
   "source": [
    "## 3 - Writing the loss function\n",
    "\n",
    "We now want to take as a loss function the distance between the _predicted_ probability given by the model $q_W^{(i)}(x)$ and the actualy probabilith $p^{(i)}(x)$.\n",
    "\n",
    "The actual probability is a delta function: it is zero for every label except for the correct one, for which it is 1.\n",
    "\n",
    "To perform this, we can use one-hot encoding, which takes an integer value in $i\\in[0..K]$ and returns a vector in $R^K$ where only the i-th component is 1 and the other are zero: $v_j = \\delta_{i,j}$.\n",
    "\n",
    "See the examples below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d168c699-dbc5-4f84-9f1f-40511e35050f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 becomes: [1. 0. 0. 0. 0.]\n",
      "1 becomes: [0. 1. 0. 0. 0.]\n",
      "2 becomes: [0. 0. 1. 0. 0.]\n",
      "3 becomes: [0. 0. 0. 1. 0.]\n",
      "4 becomes: [0. 0. 0. 0. 1.]\n"
     ]
    }
   ],
   "source": [
    "for i in range(5):\n",
    "    print(f\"{i} becomes: {jax.nn.one_hot(i, 5)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c72691d-9a4e-4826-a042-67b0aaa3238e",
   "metadata": {},
   "source": [
    "For the loss function, i'll draw from my vast knowledge of loss functions (aka: [here](https://optax.readthedocs.io/en/latest/api.html)) and choose `optax.softmax_cross_entropy`.\n",
    "\n",
    "However, for the sake of completeness, i'll write it down here by hand.\n",
    "\n",
    "`?optax.softmax_cross_entropy`\n",
    "> Computes the softmax cross entropy between sets of logits and labels.\n",
    ">\n",
    ">Measures the probability error in discrete classification tasks in which\n",
    ">the classes are mutually exclusive (each entry is in exactly one class).\n",
    ">For example, each CIFAR-10 image is labeled with one and only one label:\n",
    ">an image can be a dog or a truck, but not both.\n",
    ">\n",
    ">References:\n",
    "> [Goodfellow et al, 2016](http://www.deeplearningbook.org/contents/prob.html)\n",
    ">\n",
    ">Args:\n",
    ">\n",
    ">  logits: unnormalized log probabilities.\n",
    ">\n",
    ">  labels: a valid probability distribution (non-negative, sum to 1), e.g a\n",
    ">    one hot encoding of which class is the correct one for each input.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2b7a0b6-6157-4bf5-a6a0-39888bebe46b",
   "metadata": {},
   "source": [
    "The cross entropy between an approximate distribution $P_W(x)$ and a target distribution $Q(x)$ is defined as \n",
    "\n",
    "$$\n",
    "\\mathcal{L}(W) = - \\langle log P_W(x) \\rangle_{x \\approx Q(x)} = - \\sum_{x \\approx Q(x)} log P_W(x)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "e4112499-6b6f-445e-bd18-22fd5b0b2a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The loss function that we will use\n",
    "def cross_entropy(*, logits, labels):\n",
    "    one_hot_labels = jax.nn.one_hot(labels, num_classes=10)\n",
    "    return -jnp.mean(jnp.sum(one_hot_labels * logits, axis=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "e90668ed-0090-446a-81d3-4f6e1763b2e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(params, images, labels):\n",
    "    \"\"\"\n",
    "    Loss function minimised during training of the model.\n",
    "    \"\"\"\n",
    "    # compute the output of the model, which gives the \n",
    "    # log-probability distribution over the possible classes (0...9)\n",
    "    logits = model.apply({'params': params}, images)\n",
    "    # feed it to the cross_entropy\n",
    "    return cross_entropy(logits=logits, labels=labels)\n",
    "\n",
    "def compute_metrics(*, logits, labels):\n",
    "    \"\"\"\n",
    "    Compute metrics of the model during training.\n",
    "    \n",
    "    Returns the loss and the accuracy.\n",
    "    \"\"\"\n",
    "    loss = cross_entropy(logits=logits, labels=labels)\n",
    "    accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)\n",
    "    metrics = {\n",
    "      'loss': loss,\n",
    "      'accuracy': accuracy,\n",
    "    }\n",
    "    return metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff0ff54a-3d6a-42e7-bdf3-23a9b0d899fd",
   "metadata": {},
   "source": [
    "## 4 - Create the setup and training loop\n",
    "\n",
    "We need to define some functions to create the initial state and we need to define a function to execute one training step, and the whole training loop.\n",
    "\n",
    "For the optimiser, we use optimisers defined in "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "50082ccb-9022-496a-8898-f2b645116347",
   "metadata": {},
   "outputs": [],
   "source": [
    "import optax\n",
    "from flax.training import train_state  # Useful dataclass to keep train state\n",
    "\n",
    "def create_train_state(rng, optimiser):\n",
    "    \"\"\"Creates initial `TrainState`, holding the current parameters, state of the\n",
    "    optimiser and other metadata.\n",
    "    \"\"\"\n",
    "    # Construct the model parameters\n",
    "    params = model.init(rng, jnp.ones([1, 28, 28, 1]))['params']\n",
    "        \n",
    "    # Package all those informations in the model state\n",
    "    return train_state.TrainState.create(\n",
    "        apply_fn=model.apply, params=params, tx=optimiser)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "4d0ec7de-f082-4276-96e6-86c59712ba69",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def eval_metrics(params, batch):\n",
    "    \"\"\"\n",
    "    This function evaluates the metrics without training the model.\n",
    "    \n",
    "    Used to check the performance of the network on training and test datasets.\n",
    "    \"\"\"\n",
    "    logits = model.apply({'params': params}, batch['image'])\n",
    "    return compute_metrics(logits=logits, labels=batch['label'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "5f6b6254-da8a-4bd5-ad56-15b13a3c2bd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Partial is handy as it can be used to 'fix' some arguments to a function.\n",
    "# so partial(f, x)(y) == f(x,y)\n",
    "from functools import partial\n",
    "\n",
    "@jax.jit\n",
    "def train_step(state, batch):\n",
    "    \"\"\"\n",
    "    Train for a single step.\n",
    "    \n",
    "    The input images `batch` should not be too large, otherwise we will run\n",
    "    out of memory. Therefore the input should be 'batched', meaning should be\n",
    "    separated into small blocks of ~hundreds (instead of tens of thousands)\n",
    "    iamges.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Fix some arguments to the loss function (so that the only 'free' parameter is\n",
    "    # the parameters of the network.\n",
    "    _loss_fn = partial(loss_fn, images=batch['image'], labels=batch['label'])\n",
    "    \n",
    "    # construct the function returning the loss value and gradient.\n",
    "    val_grad_fn = jax.value_and_grad(_loss_fn)\n",
    "    \n",
    "    # compute loss and gradient\n",
    "    loss, grads = val_grad_fn(state.params)\n",
    "    \n",
    "    # update the state parameters with the new gradients\n",
    "    # objects are immutable so the output of this function is a different\n",
    "    # object than the starting one.\n",
    "    state = state.apply_gradients(grads=grads)\n",
    "    \n",
    "    # Evaluate the network again to get the log-probability distribution\n",
    "    # over the batch images\n",
    "    metrics = eval_metrics(state.params, batch)\n",
    "    \n",
    "    return state, metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "e2f17cf0-198d-4d57-88d9-de60e6e76492",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(state, train_ds, batch_size, epoch, rng, *, max_steps=None):\n",
    "    \"\"\"Train for a single `epoch`.\n",
    "    \n",
    "    And epoch is composed of several steps, where every step is taken by updating\n",
    "    the network parameters with a small mini-batch.\n",
    "    \"\"\"\n",
    "    \n",
    "    # total number of training images\n",
    "    train_ds_size = len(train_ds['image'])\n",
    "    \n",
    "    # Compute how many steps are present in this epoch.\n",
    "    # In one epoch we want to go through the whole dataset.\n",
    "    steps_per_epoch = train_ds_size // batch_size\n",
    "\n",
    "    # Truncate the number of steps (used to speed up training)\n",
    "    # Sometimes we might want not to go through the whole dataset\n",
    "    # in an epoch.\n",
    "    if max_steps is not None:\n",
    "        steps_per_epoch = min(steps_per_epoch, max_steps)\n",
    "\n",
    "    # generate a random permutation of the indices to shuffle the training\n",
    "    # dataset, and reshape it to a set of batches.\n",
    "    perms = jax.random.permutation(rng, train_ds_size)\n",
    "    perms = perms[:steps_per_epoch * batch_size]\n",
    "    perms = perms.reshape((steps_per_epoch, batch_size))\n",
    "    \n",
    "    # execute the training step for every mini-batch\n",
    "    batch_metrics = []\n",
    "    for perm in perms:\n",
    "        batch = {k: v[perm, ...] for k, v in train_ds.items()}\n",
    "        state, metrics = train_step(state, batch)\n",
    "        batch_metrics.append(metrics)\n",
    "\n",
    "    # compute mean of metrics across each batch in epoch.\n",
    "    batch_metrics_np = jax.device_get(batch_metrics)\n",
    "    epoch_metrics_np = {\n",
    "        k: np.mean([metrics[k] for metrics in batch_metrics_np])\n",
    "            for k in batch_metrics_np[0]}\n",
    "\n",
    "    return state, epoch_metrics_np\n",
    "\n",
    "\n",
    "def evaluate_model(params, test_ds):\n",
    "    \"\"\"\n",
    "    evaluate the performance of the model on the test dataset\n",
    "    \"\"\"\n",
    "    metrics = eval_metrics(params, test_ds)\n",
    "    metrics = jax.device_get(metrics)\n",
    "    summary = jax.tree_map(lambda x: x.item(), metrics)\n",
    "    return summary['loss'], summary['accuracy']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed1db647-2349-46d8-8f6f-8bebedf73872",
   "metadata": {},
   "source": [
    "# 5 - Running the optimisation\n",
    "\n",
    "We will now finally run the optimisation.  Below we will define all the required HyperParameters.\n",
    "\n",
    "For the optimiser, we pick one from the [optax](https://optax.readthedocs.io) package, which is very comprehensive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "d1123ddd-9083-4122-b47d-07513797aad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Definition of optimiser HyperParameters\n",
    "\n",
    "learning_rate = 0.1\n",
    "\"\"\"\n",
    "Standard SGD step size\n",
    "\"\"\"\n",
    "momentum = 0.9\n",
    "\"\"\"\n",
    "Amount of memntum. The maximum effective learning rate will be\n",
    "$ learning_rate * momentum/(1-momentum)$\n",
    "\"\"\"\n",
    "\n",
    "# Construct the optimiser\n",
    "# we use optimisers from the optax package which is a very comprehensive\n",
    "# optimiser library\n",
    "optimiser = optax.sgd(learning_rate, momentum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "afae251b-a760-4156-ae09-51e4476d769d",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 10\n",
    "\"\"\"\n",
    "\"\"\"\n",
    "\n",
    "batch_size = 32\n",
    "\"\"\"\n",
    "\"\"\"\n",
    "\n",
    "max_steps = 200\n",
    "\"\"\"\n",
    "Cutoff to the number of steps (minibatches) in an epoch\n",
    "\"\"\";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "6e7263bc-b1f0-4a71-8b2a-ff19daf8257b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the rng to get two keys, one to 'shuffle' the dataset at every iteration,\n",
    "# and one to initialise the network\n",
    "rng, init_rng = jax.random.split(jax.random.PRNGKey(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "6a077215-b821-45cf-8747-cd7154f4351a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "252f896f828d42a4a4eb1c84e1501488",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/21 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train epoch: 0, loss: 0.1496, accuracy: 95.72\n",
      " test epoch: 0, loss: 0.15, accuracy: 95.52\n",
      "train epoch: 1, loss: 0.0684, accuracy: 98.13\n",
      " test epoch: 1, loss: 0.11, accuracy: 96.63\n",
      "train epoch: 2, loss: 0.0497, accuracy: 98.63\n",
      " test epoch: 2, loss: 0.11, accuracy: 96.95\n",
      "train epoch: 3, loss: 0.0369, accuracy: 99.02\n",
      " test epoch: 3, loss: 0.12, accuracy: 96.91\n",
      "train epoch: 4, loss: 0.0276, accuracy: 99.29\n",
      " test epoch: 4, loss: 0.13, accuracy: 97.08\n",
      "train epoch: 5, loss: 0.0226, accuracy: 99.40\n",
      " test epoch: 5, loss: 0.12, accuracy: 97.31\n",
      "train epoch: 6, loss: 0.0181, accuracy: 99.55\n",
      " test epoch: 6, loss: 0.12, accuracy: 97.60\n",
      "train epoch: 7, loss: 0.0148, accuracy: 99.61\n",
      " test epoch: 7, loss: 0.11, accuracy: 97.83\n",
      "train epoch: 8, loss: 0.0139, accuracy: 99.64\n",
      " test epoch: 8, loss: 0.13, accuracy: 97.46\n",
      "train epoch: 9, loss: 0.0115, accuracy: 99.72\n",
      " test epoch: 9, loss: 0.14, accuracy: 97.35\n",
      "train epoch: 10, loss: 0.0120, accuracy: 99.71\n",
      " test epoch: 10, loss: 0.13, accuracy: 97.81\n",
      "train epoch: 11, loss: 0.0089, accuracy: 99.78\n",
      " test epoch: 11, loss: 0.17, accuracy: 97.39\n",
      "train epoch: 12, loss: 0.0091, accuracy: 99.77\n",
      " test epoch: 12, loss: 0.16, accuracy: 97.61\n",
      "train epoch: 13, loss: 0.0063, accuracy: 99.85\n",
      " test epoch: 13, loss: 0.15, accuracy: 97.70\n",
      "train epoch: 14, loss: 0.0058, accuracy: 99.87\n",
      " test epoch: 14, loss: 0.19, accuracy: 97.54\n",
      "train epoch: 15, loss: 0.0062, accuracy: 99.84\n",
      " test epoch: 15, loss: 0.20, accuracy: 97.36\n",
      "train epoch: 16, loss: 0.0062, accuracy: 99.84\n",
      " test epoch: 16, loss: 0.15, accuracy: 97.91\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m/var/folders/jy/14wyknxn18v9yncg8x189vn40000gn/T/ipykernel_75403/1238214855.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m         \u001b[0mrng\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_rng\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrng\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m         \u001b[0;31m# Run an optimization step over a training batch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m         \u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_metrics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_epoch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_ds\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_rng\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m         \u001b[0;31m# Evaluate on the test set after each training epoch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m         \u001b[0mtest_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_accuracy\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mevaluate_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_ds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/var/folders/jy/14wyknxn18v9yncg8x189vn40000gn/T/ipykernel_75403/3575831663.py\u001b[0m in \u001b[0;36mtrain_epoch\u001b[0;34m(state, train_ds, batch_size, epoch, rng, max_steps)\u001b[0m\n\u001b[1;32m     28\u001b[0m     \u001b[0mbatch_metrics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mperm\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mperms\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 30\u001b[0;31m         \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mperm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m...\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtrain_ds\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     31\u001b[0m         \u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmetrics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m         \u001b[0mbatch_metrics\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmetrics\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/var/folders/jy/14wyknxn18v9yncg8x189vn40000gn/T/ipykernel_75403/3575831663.py\u001b[0m in \u001b[0;36m<dictcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m     28\u001b[0m     \u001b[0mbatch_metrics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mperm\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mperms\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 30\u001b[0;31m         \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mperm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m...\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mk\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtrain_ds\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitems\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     31\u001b[0m         \u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmetrics\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m         \u001b[0mbatch_metrics\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmetrics\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_rewriting_take\u001b[0;34m(arr, idx, indices_are_sorted, unique_indices)\u001b[0m\n\u001b[1;32m   5150\u001b[0m   \u001b[0marr\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0masarray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5151\u001b[0m   \u001b[0mtreedef\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstatic_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdynamic_idx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_split_index_for_jit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 5152\u001b[0;31m   return _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,\n\u001b[0m\u001b[1;32m   5153\u001b[0m                  unique_indices)\n\u001b[1;32m   5154\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36m_gather\u001b[0;34m(arr, treedef, static_idx, dynamic_idx, indices_are_sorted, unique_indices)\u001b[0m\n\u001b[1;32m   5179\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5180\u001b[0m   \u001b[0;31m# This adds np.newaxis/None dimensions.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 5181\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mexpand_dims\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindexer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnewaxis_dims\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   5182\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5183\u001b[0m _Indexer = collections.namedtuple(\"_Indexer\", [\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py\u001b[0m in \u001b[0;36mexpand_dims\u001b[0;34m(a, axis)\u001b[0m\n\u001b[1;32m   1585\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mexpand_dims\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maxis\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mUnion\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mSequence\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1586\u001b[0m   \u001b[0m_check_arraylike\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"expand_dims\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1587\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpand_dims\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_ensure_index_tuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maxis\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1588\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1589\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/_src/lax/lax.py\u001b[0m in \u001b[0;36mexpand_dims\u001b[0;34m(array, dimensions)\u001b[0m\n\u001b[1;32m   3921\u001b[0m     \u001b[0mresult_shape\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minsert\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3922\u001b[0m   \u001b[0mbroadcast_dims\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mndim_out\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdims_set\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3923\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mbroadcast_in_dim\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult_shape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbroadcast_dims\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3924\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3925\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/_src/lax/lax.py\u001b[0m in \u001b[0;36mbroadcast_in_dim\u001b[0;34m(operand, shape, broadcast_dimensions)\u001b[0m\n\u001b[1;32m    746\u001b[0m       and isinstance(operand, (xla.DeviceArray, core.Tracer))):\n\u001b[1;32m    747\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0moperand\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 748\u001b[0;31m   return broadcast_in_dim_p.bind(\n\u001b[0m\u001b[1;32m    749\u001b[0m       \u001b[0moperand\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    750\u001b[0m       broadcast_dimensions=tuple(broadcast_dimensions))\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/core.py\u001b[0m in \u001b[0;36mbind\u001b[0;34m(self, *args, **params)\u001b[0m\n\u001b[1;32m    265\u001b[0m         args, used_axis_names(self, params) if self._dispatch_on_params else None)\n\u001b[1;32m    266\u001b[0m     \u001b[0mtracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtop_trace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull_raise\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 267\u001b[0;31m     \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtop_trace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess_primitive\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    268\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_lower\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmultiple_results\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mfull_lower\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    269\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess_primitive\u001b[0;34m(self, primitive, tracers, params)\u001b[0m\n\u001b[1;32m    610\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    611\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mprocess_primitive\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 612\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimpl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    613\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    614\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/interpreters/xla.py\u001b[0m in \u001b[0;36mapply_primitive\u001b[0;34m(prim, *args, **params)\u001b[0m\n\u001b[1;32m    273\u001b[0m   \u001b[0;34m\"\"\"Impl rule that compiles and runs a single primitive 'prim' using XLA.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    274\u001b[0m   \u001b[0mcompiled_fun\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mxla_primitive_callable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprim\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0munsafe_map\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg_spec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 275\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mcompiled_fun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    276\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    277\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/Documents/pythonenvs/ml_teach-env/lib/python3.9/site-packages/jax/interpreters/xla.py\u001b[0m in \u001b[0;36m_execute_compiled_primitive\u001b[0;34m(prim, compiled, result_handler, *args)\u001b[0m\n\u001b[1;32m    389\u001b[0m   \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompiled\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlocal_devices\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    390\u001b[0m   \u001b[0minput_bufs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_iterable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice_put\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mx\u001b[0m \u001b[0;32min\u001b[0m \u001b[0margs\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mx\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mtoken\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 391\u001b[0;31m   \u001b[0mout_bufs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompiled\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_bufs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    392\u001b[0m   \u001b[0mcheck_special\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprim\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_bufs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    393\u001b[0m   \u001b[0;32mreturn\u001b[0m \u001b[0mresult_handler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mout_bufs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# Import the TQDM progress bar module using automatic notebook detection\n",
    "# Otherwise it would not work..\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "state = create_train_state(init_rng, optimiser)\n",
    "\n",
    "metrics = {\"test_loss\" : [], \"test_accuracy\": [], \"train_loss\":[], \"train_accuracy\":[]}\n",
    "\n",
    "with tqdm(range(0, 20 + 1)) as pbar:\n",
    "    for epoch in pbar:\n",
    "        # Use a separate PRNG key to permute image data during shuffling\n",
    "        rng, input_rng = jax.random.split(rng)\n",
    "        # Run an optimization step over a training batch\n",
    "        state, train_metrics = train_epoch(state, train_ds, batch_size, epoch, input_rng)\n",
    "        # Evaluate on the test set after each training epoch\n",
    "        test_loss, test_accuracy = evaluate_model(state.params, test_ds)\n",
    "        pbar.write('train epoch: %d, loss: %.4f, accuracy: %.2f' % (epoch, train_metrics['loss'], train_metrics['accuracy'] * 100))\n",
    "        pbar.write(' test epoch: %d, loss: %.2f, accuracy: %.2f' % (epoch, test_loss, test_accuracy * 100))\n",
    "\n",
    "        # save data\n",
    "        metrics[\"train_loss\"].append(train_metrics[\"loss\"])\n",
    "        metrics[\"train_accuracy\"].append(train_metrics[\"accuracy\"])\n",
    "        metrics[\"test_loss\"].append(test_loss)\n",
    "        metrics[\"test_accuracy\"].append(test_accuracy)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3489436-f20c-4e1b-8d2c-05ac86769fc0",
   "metadata": {},
   "source": [
    "We now want to check the performance of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "aefca0cb-80ed-43df-84de-1b4073a84d54",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'Accuracy')"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAtcAAADRCAYAAAAHSJ5pAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/MnkTPAAAACXBIWXMAAAsTAAALEwEAmpwYAABbNUlEQVR4nO3dd3iUVfbA8e/JpPdOS+hFpHcBkWIB7BUVsa2KveyurrqWXd11RX+ubS2IiB0suCquoIhSVKRLFOkgkASUEEgo6TP398edQIAAAWbmnSTn8zzzzLz9TAhvzty591wxxqCUUkoppZQ6fiFOB6CUUkoppVRdocm1UkoppZRSPqLJtVJKKaWUUj6iybVSSimllFI+osm1UkoppZRSPqLJtVJKKaWUUj6iybVSSikARGSCiGwVkWWH2C4i8ryIrBWRn0Ske5VtV4vIGu/j6sBFrZRSwUXqSp3r1NRU07x5c6fDUEqpY7J48eJtxpg0J2MQkVOA3cBbxpiO1Ww/E7gdOBPoAzxnjOkjIsnAIqAnYIDFQA9jzI7DXU/v20qp2upw9+zQQAfjL82bN2fRokVOh6GUUsdERDY6HYMxZo6IND/MLudhE28DzBORRBFpBAwCvjLGbAcQka+AYcCkw11P79tKqdrqcPds7RailFKqppoA2VWWc7zrDrVeKaXqHU2ulVJKBYyIjBaRRSKyKC8vz+lwlFLK5zS5VkopVVO5QGaV5QzvukOtP4gxZpwxpqcxpmdamqNdzJVSyi/qTJ/r6pSXl5OTk0NJSYnTofhdZGQkGRkZhIWFOR2KUqrumgLcJiLvYQc0FhpjtojIl8C/RCTJu98ZwP1OBamUUk6q08l1Tk4OcXFxNG/eHBFxOhy/McaQn59PTk4OLVq0cDocpWqf7evhx3eg25WQXH//D4nIJOzgxFQRyQH+BoQBGGPGAlOxlULWAkXAtd5t20XkH8BC76kerRzcqJSqWzweQ2mFh5JyNyUVborL3JSUeyipcNt15d7lKs/F5W5Ky92UeI8rLfcQERZCTEQoMeEu73OofY6oumxfx0aEEhEaUmtyOb8m1yIyDHgOcAHjjTFjDtj+J+B6oALIA/5gjNno3XY18KB3138aY9482uuXlJTU+cQaQERISUlB+y8qdZS2r4c5/4asSWDcsGMDXDzB6agcY4y5/AjbDXDrIbZNAOrvD0+pOqTC7SFnRzG/bttz0GNzYTHHWsU5MiyEqDAX4aEhlFZ42FNaQbm7ZidzhQjR4a69SXdsRCjh3oTbJUJICISIECKCK0S8r9n3OsS7LGKP8e4fERrCI+cdVHn0uPgtuRYRF/AicDp25PhCEZlijFleZbcfgZ7GmCIRuRl4ErjUWzP1b1Spmeo99rA1Uw8Rx/G+lVqhvrxPpXyialLtCoM+N0JJIfz0PuzcAvGNnI5QKaX8yhjD1l2lrM+rTJx38+u2Pazftofs7UX7Jb1xkaG0TI2hV/MkMpKaEB3hIirMRWSYa2/CHBHmIjLUuxxe+drl3RZyyJbnMm+Svaesgj2lbnaXVlBUVmHXlbrZU1Zh13m37SmtoKjMvi53e3B7DG6Pocxt8BiDx2NwG4PHg102drsx2PXebW6PfR3mCuGR83z7s/Vny3VvYK0xZj2At4/eecDe5NoYM7PK/vOAUd7XQzmGmqnBqKCggIkTJ3LLLbcc1XFnnnkmEydOJDEx0T+BKVUf5a+Db/8NWe/tS6r73wlxDe22pRNh8RswWLsLK6WcUVhUznpvoluZ7G7KL6Lc7SE8NIQwVwhhLvE+73sdXrkcesCyd124K4TC4nLWb9vDBu+5i8rce68bERpCi9QY2qbHMbRDQ1qkxtAyNYYWqTEkx4T7rREvPDSE8NBwkmLC/XJ+J/gzua6u7mmfw+x/HTDtMMceVDNVREYDowGaNm16PLH6TUFBAS+99NJByXVFRQWhoYf+8U+dOtXfoSlVfxwuqa6U0granA6LX4cBf4bQunOjV0oFl+IyNxvybZK7/oBuF9v3lO3dL0QgIymaZinRRIa5KHd7qHAbytyevS235RWGcreHMu+2ytflbg/lbttqW8kVImQmRdE8NYbeLZK9yXMsLdJiaBQfSUiIfgvuC0ExoFFERmG7gAw8muOMMeOAcQA9e/YMynnc77vvPtatW0fXrl0JCwsjMjKSpKQkVq5cyerVqzn//PPJzs6mpKSEO++8k9GjRwP7Zi7bvXs3w4cP5+STT2bu3Lk0adKETz/9lKioKIffmVK1wEFJ9U3Q/479k+qqeo+Gdy+GFVOg08WBjVUp5ShjDNv3lO1tLf512x627y4j1NsyHBoihHpbikNDQrzr7eswl90WGuLdt8p6j4FN24v2drv4NW8Pmwv3r2LWID6C5ikxDO3QgBaVCW9qDJnJUUSEuo7rfbk9xptoe4gItf2dlX/5M7muUd1TETkNeAAYaIwprXLsoAOOnXU8wTzy2S8s37zzeE5xkBMbx/O3czocdp8xY8awbNkyli5dyqxZszjrrLNYtmzZ3qoeEyZMIDk5meLiYnr16sVFF11ESkrKfudYs2YNkyZN4tVXX2XEiBF89NFHjBo1qrrLKaXgEEn1nRDX4PDHtToVklvCgnGaXCtVR+0urdjXYpy3hw35la93s7OkYu9+YS4hOSbc9uet8FDhMbZl2OM5pgF98ZGhtEyLpU/LFG8CbR/NU2OIjfBfOuYKEVwhtv+zCgx/JtcLgTYi0gKbLF8GjKy6g4h0A14BhhljtlbZVGdrpvbu3Xu/cnnPP/88H3/8MQDZ2dmsWbPmoOS6RYsWdO3aFYAePXqwYcOGQIWrVO2Svw7mPGUHJh5NUl0pJAR63QBf3g+bl0Ljrv6MVinlB8YYSso95BYU7z9Qzztwb+uu0v32b5IYRYvUGM7r2sQmvGm2r3GTxChCXdW38la2BtuE23a/qPDs65ZRUbndOygwMzmapOgwLT5QT/gtuTbGVIjIbdhE2QVMMMb8IiKPAouMMVOA/wNigQ+9v3CbjDHn+qNm6pFamAMlJiZm7+tZs2YxY8YMfvjhB6Kjoxk0aFC1E95ERETsfe1yuSguLg5IrErVGvsl1eFw0s3Q746aJ9VVdR0J3/wDFrwK57/o+1iVUkdUUu5mZ3E5O0vKKSyuYGdJuXe5Yu/6nQes31VlfZnbs9/5UmLCaZEaw8C2aXuT5xapsXv7Mh+tytZgparj1z7Xxpip2EkHqq57uMrr0w5zbJ2omRoXF8euXbuq3VZYWEhSUhLR0dGsXLmSefPmBTg6peqANV/BxEuPP6muFJUIXS6DH9+FM/4B0ck+C1UpZe0qKSd7ezHZO4rI3l7Epu37nnMLiikp9xz2+HBXCPFRYcRHhRIfGUZCVBiZSVF2XWQYcZGhNE6MtH2XU2JIiNbZi1XgBMWAxrosJSWF/v3707FjR6KiomjQYN8f/WHDhjF27Fjat29Pu3btOOmkkxyMVKlayF0BXz5g+0lf8/nxJdVV9boBFk2AJW/ByXf55pxK1SMVbg9bCkvYdEDinL29iOwdxftVxABbR7lpcjRtG8Qx5IR0kmLCiY8M8ybLoXuT5spkWvsPq2CmyXUATJw4sdr1ERERTJs2rdptlf2qU1NTWbZs2d71d999t8/jU6rW+uk92LYKRrztu8QaoMGJ0HwALHwN+t0O+vWvUodU4faQlVPArFV5LNm0g03bi9hcULJfCbjQECEjKYrM5GiGNUmgaXI0mUnRNE22D21ZVnWJJtdKqdqpvARmjYHG3aH9Ob4/f+/R8MGVsPoLOOEs359fqVrs950lzF6dx+xVeXy7Jo+dJRWECHRonEC3zCTO62KT5ozkKJomR9MoIQqX1lBW9YQm10qp2mnRBCjMhvNeBH+MwG93JsQ3sWX5NLlW9Vy528OiDTtsQr06jxVbbGnb9LgIhnZoyKB26ZzcOlVboJVCk2ulVG1UshO+fQpaDoaWRzX3VM25QqHnH2zlkLxVkNbOP9dRKkjlFhQze1Ues1ZtZe66fHaXVhAaIvRsnsS9w05gYNs02jeK0/JySh1Ak2ulVO3zw4tQlA+nPnzkfY9Hj2tg9hO2LN9ZT/n3Wko5rLTCzYJftzN7lW2dXrN1N2DrQJ/btTED26bRr1UKcZHaOq3U4WhyrZSqXXbnwQ8vwInnQ5Pu/r1WTCp0vAiyJtlEPjLev9dTKoB2lpSzdFMBSzbtYMmmAhb+up3icjfhrhD6tEzm0l6ZDGqXRqu0WG2dVuooaHKtlKpdvv03lBfDkAcDc73eN9jkOmsS9LkxMNdUyseMMazL28OSTTv4cdMOlmwsYPXWXRhjhyy0axDHiJ4ZDGyXxkktU4gO1/RAqWOl/3v8rKCggIkTJ3LLLbcc9bHPPvsso0ePJjo62g+RKVUL7dgIi16DbqMgtU1grtmkBzTpaQc29rrBTpGuVJDbXVpBVnYBSzbu2NsyXVhcDkB8ZCjdmyVxVudGdG+aRJfMBO3qoZQPaXLtZwUFBbz00kvHnFyPGjVKk2ulKs0aAwgMvDew1+09Gj4eDb/OglZDAnttpY7AGMPG/CIWV0mkV/22k8oy023SYxnesSHdmybRvVkiLVNjCdGyeEr5jSbXfnbfffexbt06unbtyumnn056ejoffPABpaWlXHDBBTzyyCPs2bOHESNGkJOTg9vt5qGHHuL3339n8+bNDB48mNTUVGbOnOn0W1HKWb8vt10z+t0GCU0Ce+0O58P0B2D+OE2uVdAoq/Dw/sJNvDRrHVsKSwCIiwila9NEzhjShu7NkuiamUhClLZKKxVI9Se5nnYf/Pazb8/ZsBMMH3PYXcaMGcOyZctYunQp06dPZ/LkySxYsABjDOeeey5z5swhLy+Pxo0b8/nnnwNQWFhIQkICTz/9NDNnziQ1NdW3cStVG33zT4iIg5P/FPhrh0bYyiFznoIdGyCpeeBjUMrL7TF8/GMuz85YTc6OYno3T+aOU9vQvWkSrdNjdbIWpRymnQcDaPr06UyfPp1u3brRvXt3Vq5cyZo1a+jUqRNfffUV9957L99++y0JCQlOh6rqurxV4PE4HUXNZS+AVZ9D/zsgOtmZGHpcCxJip0RXygHGGKb9vIWhz87h7g+zSIwO480/9Ob9G0/i8t5NadcwThNrpYJA/Wm5PkILcyAYY7j//vu58caDKw4sWbKEqVOn8uCDD3Lqqafy8MN+rt+r6q9lH8HkP8AJZ8MFY21rcDAzBmb8HWLSoc/NzsWR0MROs77kLRh0P4TrWAgVGMYY5qzZxlNfruLn3EJap8fy8hXdGdaxoZbIUyoIacu1n8XFxbFr1y4Ahg4dyoQJE9i92xbmz83NZevWrWzevJno6GhGjRrFPffcw5IlSw46VimfcFfAzH/ZRHXVNHjtDNj+q9NRHd7ar2Hj9zDwLxAR62wsvUdDSQEsm+xsHKreWLhhO5eOm8fVExawo6iMf1/ShS/vOoXhnRppYq1UkKo/LdcOSUlJoX///nTs2JHhw4czcuRI+vbtC0BsbCzvvPMOa9eu5Z577iEkJISwsDBefvllAEaPHs2wYcNo3LixDmhUvvHTe5C/Fi6bCGHR8OE18OpguOQNaDnI4eCq4fHA13+HxGbQ/Wqno4Fm/SC9gy3L1+1KWyBYKT9YllvIU9NXMWtVHmlxEfzjvA5c2qsp4aHaJqZUsBNjjNMx+ETPnj3NokWL9lu3YsUK2rdv71BEgVff3q86ShVl8J8edtbBG76xieH29TBpJGxbDUMfgz43BVfCWNmF5cJXofMIp6OxFr8Bn90J134Bzfr67LQistgY09NnJ6wFqrtv13drt+7mma9W8/nPW0iMDuOmga24um9zosJdToemlKricPds/QisVH2x5E0o3GRnNqxMoJNbwvVfQbvh8MV98OmtUF7ibJyV3OW2QkiDjtDxYqej2afTJRCZYFuv6xgRGSYiq0RkrYjcV832ZiLytYj8JCKzRCSjyrYnRGSZ93FpYCOv/XJ2FHHPh1mc8cxsZq3ayh2ntmHOXwZz08BWmlgrVctotxCl6oPyYltGrmm/g+s0R8TBiLdh9hMwe4ytJHLZuxDX0JlYK/34tm1ZH/lBcM2KGB5ju4TMHws7t0B8I6cj8gkRcQEvAqcDOcBCEZlijFleZbengLeMMW+KyBDgceBKETkL6A50BSKAWSIyzRizM6BvohbauquEl2au4935GxER/tC/BTcPakVKbITToSmljlEQ/cVSSvnNwtdg92/7t1pXFRICg++3SfbWFTBuEOQsDniYe5UVwawnoGlfaHOGc3EcSq/rwOOGxa87HYkv9QbWGmPWG2PKgPeA8w7Y50TgG+/rmVW2nwjMMcZUGGP2AD8BwwIQc61ljGHi/E0MfHIWb8/byMU9Mpl9zyAePPtETayVquXqfHJdV/qUH0l9eZ/qGJTugu+eti3Wzfsfft8Tz7XdRFzh8PpwWDopMDEeaMEr9sPAqX8Lrj7glZJb2qR/0eu2L3vd0ATIrrKc411XVRZwoff1BUCciKR41w8TkWgRSQUGA5l+jrfWKil385fJP/HXj3+mZ/Mkvv7TQB6/sBONEqKcDk0p5QN1OrmOjIwkPz+/zieexhjy8/OJjIx0OhQVjOaPhaJ8GPxgzfZv0AFGz4KmfeCTm+CLv9oSfoFSvAO+ewbaDPXpgEGf6z0a9myFFVOcjiSQ7gYGisiPwEAgF3AbY6YDU4G5wCTgB8Bd3QlEZLSILBKRRXl5eQEKO3hkby/i4rFz+XBxDncMac0b1/ameWqM02EppXyoTve5zsjIICcnh/pwA4+MjCQjI+PIO6r6pXgHfP8faHcmZPSo+XHRyTDqvzD9QZj3Imz9BS5+PTCzI37/HJTshFODfCKlVkMguRXMfwU6BdGAy2OXy/6tzRnedXsZYzbjbbkWkVjgImNMgXfbY8Bj3m0TgdXVXcQYMw4YB7ZaiE/fQZCbtWord72/FLfH8NrVPTm1fQOnQ1JK+UGdTq7DwsJo0aKF02Eo5Zy5L0BpIQx+4OiPdYXB8CdstY7P/wSvDoHLJ0G6H8s97twC88baihwNO/rvOr4QEgK9b7BVVjb/CI27OR3R8VoItBGRFtik+jJgZNUdvF0+thtjPMD9wATveheQaIzJF5HOQGdgeiCDD2Yej+GFmWt5ZsZq2jWIY+yoHtparVQdVqe7hShVr+3ZBvNehg4XHl+i2v1KuOZzKC+C8afBys99F+OB5jwJnnIY/Ff/XcOXuo6EsBhYMN7pSI6bMaYCuA34ElgBfGCM+UVEHhWRc727DQJWichqoAHelmogDPhWRJZjW6VHec9X7xUWl3PDW4t4+qvVnN+1CR/f0l8Ta6XquDrdcq1UvfbdM1BRDIPuP/5zZfa2/bDfuwLeG2lbwgfc7dsSefnrYPGb0PMPkFxLvnGKTIAul8GP78Dpj0JMitMRHRdjzFRs3+mq6x6u8noycNDc78aYEmzFEFXFii07uemdxeTuKObR8zpw5UnNdMpypeoBbblWqi7auRkWjocul0NaW9+cM74xXDsNOl8GMx+DD6+Gou2+OTfYc4ZGwCn3+O6cgdD7BnCXwo9vOR2JCiIf/5jDBS99T0m5m/dvPImr+jbXxFqpekJbrpWqi779N3gqYOBffHvesEi4YCw07ARfPWS7iGT0hFanQutTbb/jkGOYTW5Llp3qfMDdEFfLBnmlt4fmA2wt8X53HNv7V3VGWYWHxz5fzps/bKR3i2ReGNmN9Dit5KRUfeLXlusaTKV7iogsEZEKEbn4gG1uEVnqffin1lXRdsh63y+nVsoxOzba7hXdr4Kk5r4/vwj0uw1Gz4YBf7LTlM96HMafCv/XCj68Bpa8bVvPa+rrRyEqCfrf4ft4A6HPjVCYDaumOR2JctBvhSVcNu4H3vxhI9ef3IJ3r++jibVS9ZDfWq5rOJXuJuAabO3UAxUbY7r6Kz4A5r0Ec/4PUlofXZkypYLZ7CdBQvzfvaJRZ/sY8iDsyYf1M2HdN7D2a/jlY7tPWnvbot1qCDTrb1u+D7ThO1g7A07/h+3DXBu1HQ7xGbBgHLQ/2+lolAPmrc/ntolLKCpz88LIbpzdubHTISmlHOLPbiF7p9IFEJHKqXT3JtfGmA3ebR4/xnFo/e/ELH4LmXo3XP+1bwdnKeWEbWshayL0udn2kQ6UmBRb67nTxWAMbF1uk+x1X9uE84cXIDTSJtitT4XWp0Gqty/4jEcgrrHtu1xbuULtlOhfPwJ5qyCtndMRqQAxxvDad7/y+LSVNEuJZtINJ9GmQZzTYSmlHOTPbLImU+keTqR3Fq95InK+TyPzGr8gjwf2XAKbl9iERKnabta/IDQKTv6jczGI2Fke+98BV30K926EKyZDj2tt14kv/wov9oZnOtrKIzkLYNB9EFbLp37ufhW4IuyHCVUv7Cmt4LZJP/LPz1dwevsGfHprf02slVJBPaCxmTEmV0RaAt+IyM/GmHVVdxCR0cBogKZNmx71BZqnxPDPkr7clzmX+Bl/h/bn1N6vpZX6bZl3UOCfITbN6Wj2CY+GNqfbB0DBpn2t2uvnQHoH6HqFszH6QkyqTbCr6/qi6pz1ebsZ/fZi1uft5r7hJ3DjKS21GohSCvBvcn3EqXQPxxiT631eLyKzgG7AugP2Oa5pdPu1TiHc5eLDtDu4bvm1MOsJGPavoz2NUsFh1uMQkQD9bnc6ksNLbAo9r7UPdwVgbLeKuuCsp5yOQAXAii07GTV+PgZ4+7o+9G+d6nRISqkg4s9uIXun0hWRcOxUujWq+iEiSSIS4X2dCvSnSl9tX4kOD6V3i2Tey0mGHlfDgldg60pfX0Ypa+0MyF3sn3PnLoaV/7NVPKKS/HMNf3CF2mnWlaolluUWcvmr8whzhTD5pr6aWCulDuK35LomU+mKSC8RyQEuAV4RkV+8h7cHFolIFjATGHNAlRGfGdQujTVbd7Ol5z0QHgvT/mIHZCnlKzs325kN37kIXj0Vpj8I5cW+vcY3j0FUMvS5ybfnVUrtlZVdwMhX5xETHsr7N55Ey7RYp0NSSgUhv5bHMMZMNca0Nca0MsY85l33sDFmivf1QmNMhjEmxhiTYozp4F0/1xjTyRjTxfv8mr9iHNTO9k39ZpPblhT7dTas+Mxfl1P1icdjZ0l8sY/tY3za321XiLn/gVdOgZxFvrnOxrm2//LJf4TIeN+cUym1n8UbdzBq/HwSosN4b/RJNEuJcTokpVSQqve151qlxdIkMYrZq/JsNYP0DvDlA1BW5HRoqjbbuhJeHw6f/xmadIdb5trk9+xn4MqP7e/Xa6fDjL9DRemxX8cY+OafENsAel3vs/CVUvss+HU7V702n5TYcN4f3ZfM5GinQ1JKBbF6n1yLCAPbpfH92m2UmRA480ko3ATfP+d0aKo2qiiFmY/D2JNh2yo4fyxc+Qkkt9y3T6shcMsP0G0UfPcMvDIQNv94bNdbPxM2fm+nDQ/XP/hK+drcddu4esICGiRE8v6NfWmcWMtLRiql/K7eJ9cAg9qmsafMzaKN26H5ydDxIvj+WTuNtFI1tWkejB0As8dAhwvgtkXQ9XJb9/lAkfFw7n9s/eeSQtsX+5vHoKKs5tczBr7+ByRk2gG5SimfmrM6j2tfX0hmchTvj+5Lg3gts6iUOjJNroF+rVMJc4ntGgJ2GmYJgekPOBuYqh1KCuF/f4QJQ+1AxSsmw0Wv2rrHR9LmdNuK3eUymPMkvDoYtvxUs+uummYnQBr4FwiNOL73oJTaz8yVW7n+rUW0TItl0g0nkRan/8eUUjWjyTUQGxFKr+bJzKpMrhOa2Ik4VnwG62Y6G5wKbis+swMWF78BJ91qE+XKyVJqKioRzn8JLn8P9uTZBHv2k+AuP/QxHg/MfMx2N+ly+fG8A6XUAab/8huj315EuwZxTLqhDymxmlgrpWpOk2uvQe3SWPX7LjYXeEuk9b0NkprDtHsPn+So+mnnFlte7/1REJ0K18+wExBFHEdprnbD4ZZ50OFCmziPPxV+P0QFyuUfw+/LYNBftU60Uj407ect3PLuEjo0TuCd6/uQGB3udEhKqVpGk2uvQe3SAZi92tt6HRYJw8bYQWkLxjkYmQoqHg8sfA1e7G0nhTnt7zB6JjTp4ZvzRyfbLiWXvmPrY79yCnz7b+9Mhl7uCjtoMq09dLzQN9dVSvHp0lxum/QjXTITefu63iRE6QdXpdTR0+Taq016LI0TIpm1auu+lW2HQevTYdYY2L310Aer+iFvNbxxJnz+J2jcFW72ltfzR8tx+3PglvnQ/mz4+lFbti9vld328weQvwYG/xVCXL6/tlL10EeLc/jj+0vp0SyJt/7Qm7hITayVUsdGk2uvfSX58il3eypX2tbr8mKY8YizASrneNy2D/TY/rB1BZz3Elw1BVJa+fe6MSlwyRtw8euwY4OtRPLds/bDXqMuNgFXqhoico6I6P29ht5fuIm7J2fRt1UKb1zbi5iIUKdDUkrVYnrzrWJg23R2l1aweOOOfStTW0PfW2DpO76bUS9Y7cm33Q3GDYZZT8Cu352OyHllRfDBVbYPdPtzbHm9bldUX17PXzpeCLfOtwMlZ/wNCjbCkIcCG4OqbS4F1ojIkyJygtPBBLN35m3k3o9+ZkCbNF67uhfR4ZpYK6WOjybXVfRvnUJoiOyrGlLplHsgtiFMvcf2ua1rCjbB1L/AMx1sjWZ3Gcz6l13+6HrIXmhrKtc3e7bBm+fAys9h2BNw8QSITXMmlth02w/74gkw6H5ofZozcahawRgzCugGrAPeEJEfRGS0iMQ5HFpQef37X3nwk2WcekI6467sQWSYdrNSSh0/Ta6riIsMo2fzpP37XQNExMHpj9qawkvfdSY4f/j9F/jvaHiuKyx6zbaQ3jIfbv4eblsMva6DVV/Aa6fBuEGwdCKUlzgddWDkr4Pxp9mKHJe+Ayfd5HREtqW640Uw6D5ttVZHZIzZCUwG3gMaARcAS0TkdkcDCxKvzlnPI58tZ2iHBrw8ShNrpZTvaHJ9gEHt0ln52y5+Kzwgiew8AjJPghl/h+ICJ0LzDWNg41x4dwS83A9W/A/63AR3Ztlay+neb5BTW8PwJ+DPK+DMp2y/809uhmdOtP3PC3OcfR/+lL3AJtalO+Hq/9lBhUrVIiJyroh8DMwCwoDexpjhQBfgz07GFgxemb2Ox6au4KzOjXhhZHfCQ/VPoVLKd/SOcoBB7ezX/rNXH9B6LQJnPglF+TD7CQciO04eD6ycCq+dAa8Ph9xFMPhB+OMyW585IaP64yLioPcNts/vVZ/aDxjfPwvPdob3r4Rfv/V9lxGPG7av985A+KNvz30ky6fYriBRiXDdV5DZK7DXV8o3LgKeMcZ0Msb8nzFmK4Axpgi4ztnQnFVa4ebf01dzWvsGPHdpV8Jc+mdQKeVbOnLjAO0axNEwPpJZq/K4tFfT/Tc26gI9roH5r0D3qyC9vSMxHpWKMlg2Gb5/DvJWQmJT2xLd9QoIj675eUSg5SD72LHRdiNZ8hasmALpHWwC3nkEhMfU/JzuClsFI2+l97HKPm9bAxXF+/breLHtlpPQpObnPhbzXoYv7oeMXna2xJgU/15PKf/5O7ClckFEooAGxpgNxpivD3WQiAwDngNcwHhjzJgDtjcDJgBpwHZglDEmx7vtSeAsbKPNV8CdxgTfYI0VW3ZR5vZwcY8mhGpirZTyA02uDyAiDGybxtSft1Du9hzcqjHkIfjlY5j2F1uOLVj7vpbuhiVvwg8vws5caNARLhwPHS4A13H+syc1s8nuwPts4j5/HPzvLlvJotuVtq92cst9+7vLbUt01QQ6b5VNot2l+/ZLyIS0dtDiFPuc2s5O1PL9c7BqKgz4E/S93U7w40seD0x/AOa9BCecDReNh7Ao315DqcD6EOhXZdntXXfIr2JExAW8CJwO5AALRWSKMabqNKFPAW8ZY94UkSHA48CVItIP6A909u73HTAQ2y0lqGRlFwDQJTPR0TiUUnWXJtfVGNQujfcXZfPjpgJ6t0jef2NMCgx5EKbebVttTzzPmSAPZc8227K+YByUFECzk+Gc52x1CV9/EAiPti343a6ETfNgwSu29feHF+31wmNsEp2/FjxVppBPbAZpJ0CrIfY5/QRIbWu7oByoaR9b+m76g/DNP2HJ2zD0X3DCWb55P+XF8N8bYMVn0OdmGPqYTsyi6oJQY0xZ5YIxpkxEjjSPd29grTFmPYCIvAecB1RNrk8E/uR9PRP4pPISQCQQDgi2n3dQ1vLMyi4gPS6ChvE+/pCulFJemlxXo3+bVG9Jvq0HJ9cAPa6FxW/Alw/YGRyPpnuFPxTmwLqZsO4b20+5oti2wPa/KzB9hkWgWV/72LkZFr1uq6qERtjkud0w+5zWzibRR9N1BCCpua3YsX4WTLsP3r/Cdk8Z9sS+AZjHYk8+TLoMchbC0MdtPXOl6oY8ETnXGDMFQETOA7Yd4ZgmQHaV5RygzwH7ZAEXYruOXADEiUiKMeYHEZmJ7YoiwAvGmBU+eB8+tzSngC6ZiUiwfuuolKr1NLmuRnxkGN2bJTFrVR5/GVZN8uYKheFP2qmwv38OBt8f2ABLd8GG720yvX4mbFtt18c2gC6Xwkm3QlrbwMZUKb4xDHnAPnyt5SC46Tvb33vmY7baSe/RtjRdVOLRnSt/Hbx7sf0wMOLN4PsGQqnjcxPwroi8gE12s4GrfHDeu4EXROQaYA6QC7hFpDXQHqgcGf2ViAwwxnx74AlEZDQwGqBp06YHbvarwuJy1uft4aLuhxjArZRSPqDJ9SEMapfGk1+sYuvOEtKr+/qweX9bc/j7Z6HrSNsP2V88bls1Y91Mm0xnzwdPBYRG2Th6XAMtB9sBlnW9NcYVCn1utD/7b/4J88fCzx/AqQ/b7ik16dKRswgmjrBVTq6aYrueKFWHGGPWASeJSKx3eXcNDssFMqssZ3jXVT3vZmzLNd5zX2SMKRCRG4B5ldcRkWlAX+Cg5NoYMw4YB9CzZ8+ADnhcllsIQJeMxEBeVikVrIyBst3Vd0s9DjUaKi0iMSIS4n3d1ltDNcynkQSZQW3TAZi1Ou/QO53+D5AQ+PKvvg9gxwbbveL9K+HJljD+VJj5T/tL0O92mxTeuwFGfQR9b4UGJ9b9xLqqmFQ451m4cbbtavLZnfDqYNv3+3BW/A/eOBsi4uH6GZpYqzpLRM4CbgH+JCIPi8jDRzhkIdBGRFp4+2dfBkw54JyplX8LgPuxlUMANgEDRSTU+7dhIBB03UKWegczdspIcDYQpVRwmP0EjBsMRdt9etqatlzPAQaISBIwHXsTvhS4wqfRBJH2jeJIj4tg9qo8RvTMrH6nhCZwyt3w9aPwVFsIj7X9iSufI6ouH2FbWBRsXbGvq8f29fYa8U3sJCYtB9tuETGpAfsZ1AqNusC102DZRzD9IZgwFDqNgNMfsV1Uqpr/Cky7F5r0sKX2nJrKXCk/E5GxQDQwGBgPXAwsONwxxpgKEbkN+BJbim+CMeYXEXkUWOTtvz0IeFxEDPbvwq3ewycDQ4CfsYMbvzDGfObzN3acsrILaJkaQ0JUnW4bUkrVxPxxMOtxW5o4Ksmnp65pci3GmCIRuQ54yRjzpIgs9WkkQaayJN+Xv/xGhdtz6HqofW+zXysUbLKtymV77KNom219rlwu2wXGc+QLh8VAiwHQ+0ZbTSO1Tf1qkT4WItDpYmg3HL57Br5/HlZ+Dqf82fY/d4XDVw/BDy9Au7NsqT2nB6Eq5V/9jDGdReQnY8wjIvJvYNqRDjLGTAWmHrDu4SqvJ2MT6QOPcwM3Hn/Y/pWVU0C/VtpAEfSKtttkR//2KX/56UOYdg+0OxPOed7nv2s1Tq5FpC+2pbpydq86X69sULt0Plycw9LsAno2r6ZqCNiKGKfcfeSTGQMVJd5E25uEl+7ePyFPzISM3hB6pIpZqlrhMbZMYldv6b6vH7UT3aS0gbVf2cGPw8ZoqT1VH5R4n4tEpDGQDzRyMB7H/VZYwu87S+miXUKCk8dt5zOY+x87ruiEs+Hc/0D0If721hZF2+GX/0Lny+w31sp5a76CT26ypYovnnD8c39Uo6ZnvAvbv+5j79eELbE1Tuu0k9uk4goRZq3KO3RyXVMitutHWJR27fC35BZw2bt2AOgX99nE+ox/2m8ZtCVE1Q+fiUgi8H/AEmxXjVcdjchhS3XymOBUVgRZE+38CNvX23kQel5nG0bGngwXvmoH7tdGhbnwzoV24rRFr8Plk+wsyco5m+bbsWzpJ8LlE/02YVyNkmtjzGxgNoB3MMs2Y8wdfokoiCREhdG9aSKzVm/l7qHtnA5HHa1Wg23pvp2b/VvNRakg4r1Hf22MKQA+EpH/AZHGmEJnI3NWVk4BYS6hfaN4p0NRYCc8W/AqLHwVivKhcXe45A044RzbkthtFHx0Hbx5Ngy8F065p3Z965i3Ct6+EEp3wmmPwLdP24Fzl70LTU9yOrr66fdfYOIldjzWqP9CpP++xapptZCJIhIvIjHAMmC5iNzjt6iCyKB26SzL3cnWXSVH3lkFH1eYJtaqXjHGeLDTmFcul9b3xBrsYMb2jeKJDKtFCVpdtG0tfHYXPNMBZo+xXSGvnQY3fAMdLtj3FX2T7nDjHOh0iR109uY5dsK02iB7oR1c7y6Daz6Hk++y1aki4221qh/fdTrC+mfHBvthJywarvzY7wUNapRcAycaY3YC52MHxbQArvRXUMFkYFv7DzBn9ZEmN1NKqaDxtYhcJDoNIQAej+GnnEKtb+0UY2yZ1PeugBd6wtKJ0PlSuHUhjHwPmvWrvsteRBxcOA7OHwubl9puIis/D3j4R2X1dPtBIDIRrpsOjTrb9Wlt4fqv7Xv99BY7LsjjdjTUGtuzDdbPhnkvw6e3wqSR9luHnZudjqxmdm+Ft863496u/DggDW417XMd5q1dej52WttybymmOq9D43jS4iKYtWorF/fQWb2UUrXCjcCfgAoRKcHO0miMMfWyT8T6bbvZXVqh/a0DzeOGlf+zgxRzFtoKIKfcA71vgNj0mp+n6+WQ0QsmXwvvjYReN9hxNGHVTPDmpKWTbPLZsCNcMfng9xidbOem+OJ++zPJWwUXvWZbtINBeQlsW2W7T1Q+ti6H3b/v2yc61Q7MXPU5TL3blrY94Wz7cGpm6MMpLrAt1rt/t/ODpLcPyGVrmly/AmwAsoA5ItIM2Hmkg0RkGPActrLIeGPMmAO2nwI8C3QGLvOWearcdjXwoHfxn8aYN2sYq0+JCKe0SWPGit8PX5JPKaWChDHGt9ON1XJLs22vmK6ZWikkIMqKYOm7dpDijl8hqTmc+ZSdzTg85tjOmdradq2Y8QjMexE2/WArPaQFyXio75+3JV9bDLT9qg81458rDM56CtJPgKl/gddOtwMdk1sGLlZjoGAj/L4ctlYm0sshfy0Yb2t6aKT92bY+zQ7+a9DBPio/MOSthpWf2YnZvn7EPlLb2iS7/dm2D73TX5yVF8Oky+2A0pHvQWavgF1ajDm2BmgRCTXGVBxmuwtYDZwO5GAnnrncGLO8yj7NgXjgbmBKZXItIsnAIqAndpT7YqCHMWbHoa7Xs2dPs2jRomN6L0fyWdZmbp/0Ix/d3JcezWp5WSClVFASkcXGmJ4+Otcp1a03xszxxfl9xZ/37aoe+mQZH/+Yy09/O4OQEO0p4zdF223XgYXjoXg7NOkJ/e+wCZcvByOunm5LqZUXw/AnoNuVziVyHs++eRQ6XAgXjLUlemti/Wz44Cob+4i37RwX/rJzCyyaAOtn2Qnrynbt25bUHNK9yXODE+3r5JY1L1FXmGu766z8DDZ8bxP0+CZwwln20ay//VARSO4KeH8UrP4CLn4NOl7k80sc7p5do5+ciCQAfwMqb9izgUeBww2S6Q2sNcas957jPeA8YG9ybYzZ4N124OwqQ4GvjDHbvdu/AoYBk2oSr68NaJNKiMDsVXmaXCulaoOqA84jsffjxdhZFOudn3IK6JyRoIm1v7jLYeFrduBhSaGdmKPf7bYqhj+S3rZnwE3fw8ejYcrttuzqOc/6tfpDtdzl8Olt8NN7duK3YWMg5Ci+3W450A7knHQ5vH2+bd3vea1vY8xZDPNfhl8+tt10mp5ku9mknwgNOtoW9EO1stdUQhPoM9o+irbD6i9td6Alb8OCcbb/ebvh9kNWqyH+n8TN47G/F6unwVn/9ktifSQ17RYyAVslZIR3+UrgdeDCwxzTBMiuspwD9Knh9ao7tsmBO4nIaGA0QNOm/qsdmRgdTremScxancefzgiSr6CUUuoQjDHnVF0WkUxsF7x6p7TCzfItO7nu5AB+7V6frJ4OX/4V8tdAy0Ew9F+2BdTf4hvBlZ/YWXln/gtyF8FFEwL31X/ZHvjgajuPwpCHYMCfj+2DREoruP4rmHwd/O8u26o89F/HN7GJuxyWfwrzx9q+7uFxdhK13jf4v/tJdLJN3rtebrsHrfvadh1ZNRWyJkFoFLQ+FdqfYxNuX38gMsYOFs2aCIMfgF7X+/b8NVTTf71Wxpiqqf8jwTD9uTFmHDAO7NeL/rzWoLZp/Pur1WzbXUpqbA2/8lFKqeCQAwRmJE+QWbFlF+Vuo/2tfW3rSpj+AKydAcmt4PL3oO2wwHbPCHHZGZKbD4CProfXh9mEqv9dR9eCfLSKtsO7l8DmJXbq7B5XH9/5IhNg5Pvw1cO2e8m2Vbbmd1TS0ce1+HVYMB52bbaJ9PAnbV/3422dPhbh0TaJbn+OTfg3fm8T7ZWf25ZtVwS0GwadRkCb02veneZwvnva9snvc5MdPOuQmibXxSJysjHmOwAR6Q8UH+GYXCCzynKGd11N5AKDDjh2Vg2P9YtB7dL591ermbM6jwu7a9UQpVTwEpH/YMergC252hU7U2O9k6UzM/pW0XaYNcb2qw6Pta2svW6A0HDnYmraB276Fj67ww6s+3U2XPAKxDX0/bUKsu2siwWb4NJ3bJ9iXwhxwdDHbDWLz+6C8afB5e/bgZxH8vty2/Xjpw9submWg+DsZ6DNGf79kHE0XGE2rpaDbMKfuxh+/tBODb/8U/sB48TzbKLdrP+xxb3odfj6UXuOoY87OqCypsn1TcBb3r7XADuAI31UWwi0EZEW2GT5MmBkDa/3JfAvEan82HYGdvp1x3RoHE9qbDizVmlyrZQKelVHCVYAk4wx3zsVjJOysgtIj4ugYXyQlW2rbdzldkDczH/ZWQd7XAuD/woxqU5HZkUlwiVvwpI3Ydp98HJ/O3lLZh9o2Nk3Zfu2rrBl3cr22HrJzfod/zkP1G2U/Sbg/Stg/BDbgt2qmqESHrft2zz/Zfh1ju1u0eUy22IboHJzxywkxHbfyexlP5ytnwU/fwA/f2SnvY9vYvtJdx5h+4XXJEn+5RP43x+hzVA4/yXHP1TUdPrzLKCLiMR7l3eKyF3AT4c5pkJEbsMmyi5ggjHmFxF5FFhkjJkiIr2Aj4Ek4BwRecQY08EYs11E/oFN0AEerRzc6JSQEFuS75tVW3F7DC4dGKOUCl6TgRJjbF0tEXGJSLQxpsjhuAJuaU4BXTIT0fl0jsOaGbZf9bZVttTcsMcD06/6aIlAj2sg8yQ72HG6t5pvSBg07AQZPW297IyekNTi6Fo2N82DiSNsEvuHaf59/836wg0z7UDHdy62P+/eo228JTttmcP5Y+2sg/FN4LS/Q/erbX/n2sYVCm1Os4+yPbBqmm2Bn/cSzH0e0tpD50vsTJ2Jhxhbt26m7RaU2cd+GAl0ZZJqHE8pvk3GGP+NIjxKgSjp9OnSXO58byn/vaUf3ZseZV8opZQ6DB+X4psHnGaM2e1djgWmG2P80NR27Px93y4sLqfLI9O5Z2g7bh1cg6/X1f7yVtukeu1XtjV16GOB71d9PHb9BjmL7KC+3MWQuwTK99ht0Sm2VGCG99Gkx6EH162aBh9eAwkZMOq/AZnhD4DSXfDf0XYwYLcrbY3wH9+1ZfQy+9hW6vbnBEUy6XN78mH5x/DTh5A9z65r2tcm2R0u2PdBImexnREzuYWdaj4qMWAhHncpvkOd9ziOrZVOaZO2tySfJtdKqSAWWZlYAxhjdouIn+tfBZ+fc2y1WJ32/CgVbYfZT9h+1WExcMZjtuXUyX7VxyKuoZ3QpP3ZdtldAXkrvAn3IlthZM2X3p3FToKS0QsyetjntPa2zN6UO6BRF7jiw8B2g4mIg0vfhW/+YQfqhYRBxwttUt2ke+DicEJMiq300et620L/82TbR/vzP8G0v9jJbdoOha//AbFpdubLACbWR3I8yXW9mP68qqSYcLpkJjJrdR5/PD0Ip/lUSilrj4h0N8YsARCRHhx5EHqdk5VTAECnDK0UUiPuCtuveta/bL3qHtfY6hvB0q/6eLlCbfeQhp321ZMuKbQt2pXJ9uppsPQduy0sGsqLbJ/nEW/bab8DLSQETvubHTiZkOGfQZrBLqm5rQoz4M/w28/7+mev/gJiG9iSjEH2czlsci0iu6g+iRYgyi8RBblBbdN59uvV5O8uJUVL8imlgtNdwIcishl7v24IXOpoRA5Yml1Ay7QYEqLq4NfmvuTx2BbcGX+3U0W3OMVWW2jY0enI/C8yAVoNtg+wdZJ3/Gq7G+QshMh4OOUvzrfaZ/ikx1jtJgKNOtvHaY9A9gKIbxy4bjpH4bDJtTHGgcKIwW1QuzSembGab9ds4/xuB81ro5RSjjPGLBSRE4DKWa9WGWPKnYwp0IwxLM0uYEDrOtLq6g8VpfDT+zDXW1s5qQVcNtHOsFhb+lX7moitD53c0g6kU8EpxGUHfgapICmAWHt0apJAckw4s1ZtdToUpZSqlojcCsQYY5YZY5YBsSJyi9NxBdJvO0vI21VKZ+0ScrCi7TDnKXimo50mOjQcLnwVbltoux/U18RaKR85nj7X9ZItyZfKnDXb8HgMIVqSTykVfG4wxrxYuWCM2SEiNwAvORhTQGVlewcz6uQx+2z/Fea9DD++bfsStz4N+t1uy+tpQq2Uz2hyfQwGtUvnk6Wb+Sm3kK5641ZKBR+XiIjx1loVERdQy0o9HJ+snALCXEL7RvFOh+K8nMW2ZvCKKSAuW86s323BWataqTpAk+tjcErbNERg1qqtmlwrpYLRF8D7IvKKd/lGYJqD8QRcVnYB7RvFExnmcjoUZ1QOUvz+edg0FyISoN8d0OdGOwhMKeU3mlwfg+SYcDpnJDJ7dR53naYl+ZRSQedeYDRwk3f5J2zFkHrB4zH8lFPIBfVx0Hl5ia3NPPcFyF8DCZl2iunuV9m6yUopv9Pk+hgNapvG89+sYceeMpJi6tW3rUqpIGeM8YjIfKAVMAJIBT5yNqrAWb9tN7tLK+pXf+ui7XbSlwXjYE8eNOwMF70GJ55XN2fwUyqIaXJ9jAa1S+O5r9cwZ00e53Wth60jSqmgIyJtgcu9j23A+wDGmMFOxhVoS72DGbtm1vFKIR4PZM+3k2osnQQVxdDmDDtIsfkAHaSolEM0uT5GnTMSSYoOY/YqTa6VUkFjJfAtcLYxZi2AiPyxpgeLyDDgOcAFjDfGjDlgezNgApAGbAdGGWNyRGQw8EyVXU8ALjPGfHIc7+WYZWUXEBsRSstUB2bU8zePx05u8svHsPwT2LUFXBHQ6WKbVKe3dzpCpeo9Ta6PkStEGNAmjdmr87Qkn1IqWFwIXAbMFJEvgPewMzQekbeiyIvA6UAOsFBEphhjllfZ7SngLWPMmyIyBHgcuNIYMxPo6j1PMrAWmO6bt3T0snIK6JyRUHfuyx6PnZr7l49h+aewM9cm1G1Ohw4XQNuh2p9aqSCiyfVxGNQujSlZm1m2uZDOGYlOh6OUque8LcWfiEgMcB52GvR0EXkZ+NgYc7iEtzew1hizHkBE3vOeo2pyfSLwJ+/rmcAn1ZznYmCaMabo2N/JsSspd7Niy06uH9DSicv7jjGQu9gm1L98AjtzwBVua1Of9ndoO8xOza2UCjqaXB+HU9qmATBrVZ4m10qpoGGM2QNMBCaKSBJwCbaCyOGS6yZAdpXlHKDPAftkYVvHnwMuAOJEJMUYk19ln8uApw91EREZja1kQtOmTWv0fo7Gii07KXcbuvh7ZsaKUtj1G0QnQ3isb/o3GwO5S2C5N6EuzIaQMGh9Kpz6ELQbDpF1vB+5UnWAJtfHITU2gs4ZCcxatZU7Tm3jdDhKKXUQY8wOYJz3cbzuBl4QkWuAOUAu4K7cKCKNgE7Al4eJZ28sPXv2ND6IaT9Z2QWAn2dmXDMDPv8jFGyyy64IiE7xPpL3vY5JPXhddKpdDo2wxxoDm3/c14e6YJNNqFsNgcEP2IQ6yo/vRSnlc5pcH6dBbdN4YeZaCorKSIzWknxKqVorF8isspzhXbeXMWYztuUaEYkFLjLGFFTZZQS2+0m5f0M9tKycQtLjImgYH+n7k+/eCl/cD8smQ2pbOOvfULYHivK9j+2wZxtsybLLJQWHPld4nE2yPRW2D3VIKLQcDAPvgxPOhKgk38evlAoITa6P05D2DXj+m7X89eOfeeqSLkSH649UKVUrLQTaiEgLbFJ9GTCy6g4ikgpsN8Z4gPuxlUOquty73jFZOQV0yUxEfFmGzhj48W2Y/hCUF8Gg++HkP+5rfT4UdwUU74CibVUS8AMScXcZtL4fTjjLJttKqVpPM8Hj1DUzkfuGn8ATX6zk121FjLuyB5nJ0U6HpZRSR8UYUyEit2G7dLiACcaYX0TkUWCRMWYKMAh4XEQMtlvIrZXHi0hzbMv37EDHXqmwuJz1eXu4qHuG7066bQ18dhds/A6a9Yezn4W0Gs7M6wqF2DT7UErVG5pc+8BNA1txQsM4bp/0I+e+8B0vjuxOv9apToellFJHxRgzFZh6wLqHq7yeDEw+xLEbsIMiHfNzjp08posvBphXlMJ3z8K3T0FYFJz7H+g6CkJCjv/cSqk6Te8SPjKoXTpTbjuZlNgIrpywgNe++xVjfD5WRyml1CFk5RQA0Ol4K4VsnAtjT4ZZ/4L258CtC6H7VZpYK6VqRO8UPtQiNYZPbu3PkBPS+cf/lnP3hz9RUu4+8oFKKaWO29LsAlqmxZAQFXZsJyjeAZ/dCa8Ph/ISuGIyXDwB4hr4NlClVJ2m3UJ8LDYilFdG9eD5b9bw7Iw1rN26i7FX9qBRQpTToSmlVJ1ljGFpdgEDjqVLnjG2FN60e+3gw36320GL4TG+D1QpVedpy7UfhIQId53Wlleu7MHarbs55z/fs2jDdqfDUkqpOuu3nSXk7So9+vrWBZtg4qUw+VqIbww3zIQz/qmJtVLqmGly7UdDOzTkk1v7Exvh4vJX5/Hu/I1Oh6SUUnXSUU8e466AH16EF/vAhu9g6ONw/dfQuKu/QlRK1ROaXPtZmwZxfHrryfRrlcoDHy/jrx//TFmFx+mwlFKqTlmaXUiYS2jfKO7IO2/5CcafCl/+FZoPgFvnQd9bbOk8pZQ6TppcB0BCdBgTrunFTQNbMXH+Jka+Oo+tu0qcDksppeqMrOwC2jeKJyLUdfgdiwvgjbNh52a45A0Y+T4kNg1EiEqpekKT6wBxhQj3DT+B/1zejWWbCzn3P9/v/RpTKaXUsXN7DD/nFtasvvXi16G0EEZNhg4XgC9nclRKKfycXIvIMBFZJSJrReS+arZHiMj73u3zvTN8ISLNRaRYRJZ6H2P9GWcgndOlMR/d3A9XiHDJKz/w0eIcp0NSSqlabX3ebnaXVhy5v3VFGcwbCy0HQaMugQhNKVUP+S25FhEX8CIwHDgRuFxETjxgt+uAHcaY1sAzwBNVtq0zxnT1Pm7yV5xO6NA4gc9uP5keTZP484dZPPLZL1S4tR+2UkodiyzvzIxdM48weczPH8Lu36DfHQGISilVX/mz5bo3sNYYs94YUwa8B5x3wD7nAW96X08GThWpH9/RJceE89Z1vbm2f3Ne/34DV01YQP7uUqfDUkqpWicru4DYiFBapsYeeidjYO5/oEEnaDUkcMEppeodfybXTYDsKss53nXV7mOMqQAKgRTvthYi8qOIzBaRAdVdQERGi8giEVmUl5fn2+gDIMwVwt/O6cD/XdyZRRt3MPipWYz/dj2lFTqro1JK1VRWTgGdMxIICTlM28yaryBvhZ0gpn604SilHBKsAxq3AE2NMd2APwETRST+wJ2MMeOMMT2NMT3T0tICHqSvXNIzk89uO5muTZP45+crOO3p2fzvp80YY5wOTSmlglpJuZsVW3Yeub/13OchPgM6XhiQuJRS9Zc/k+tcILPKcoZ3XbX7iEgokADkG2NKjTH5AMaYxcA6oK0fY3Vcu4ZxvPWH3rz5h97EhIdy28QfufDluSzeqDM7KqXUoazYspNytzl8pZDcJbDhWzjpZnCFBSw2pVT95M+K+QuBNiLSAptEXwaMPGCfKcDVwA/AxcA3xhgjImnAdmOMW0RaAm2A9X6MNWgMbJvGya1T+WhxDk9NX8VFL//A8I4NuXfYCTRP1el4lVKqqsqSpl0P13I993mIiIfuVwUkJqXqg/LycnJycigpqdvzdkRGRpKRkUFYWM0/mPstuTbGVIjIbcCXgAuYYIz5RUQeBRYZY6YArwFvi8haYDs2AQc4BXhURMoBD3CTMabeNOG6QoQRvTI5u0sjXp3zK6/MWceMFb9z5UnNuX1Ia5Jiwp0OUSmlgkJWTiHpcRE0TIisfoftv8LyT21f68iDehcqpY5RTk4OcXFxNG/enLpai8IYQ35+Pjk5ObRo0aLGx/l1rldjzFRg6gHrHq7yugS4pJrjPgI+8mdstUF0eCh3ntaGy3tn8vRXq3lj7q9MXpzN7UPacFW/ZkeeiUwppeq4rOyCw/e3nvcSiAv63BywmJSqD0pKSup0Yg0gIqSkpHC0RTOCdUCjqiI9PpIxF3Vm2p2n0L1ZEo9NtYMeP8vSQY9KqfqrsKic9dv2HLpLSNF2+PEd6DwC4hsFNDal6oO6nFhXOpb3qMl1LdKuYRxvXNubt6+zgx5vn/QjF7w0l4Ub6k2PGaWU2uun3AKAQw9mXPgalBfZLiFKqTqloKCAl1566aiPO/PMMykoKPB9QFVocl0LDWiTxud3DODJizuzpbCYS8b+wE1vL+bXbXucDk0ppQKmcjBjp4xqZmYsL4EFr0CbMyC9fWADU0r53aGS64qKisMeN3XqVBITE/0UleXXPtfKf1whwoiemZzduRHjv/2VsbPtoMfLemdyVd/mtG0Q53SISinlV1k5hbRMiyEhqppR/FmTYE+eTnWuVB113333sW7dOrp27UpYWBiRkZEkJSWxcuVKVq9ezfnnn092djYlJSXceeedjB49GoDmzZuzaNEidu/ezfDhwzn55JOZO3cuTZo04dNPPyUqKuq4Y9PkupaLDg/ljlPbcFnvTJ75ag0fLMzhnXmb6NksiZF9mnJmp0ZEhunAR6VU3WKMYWl2AQNapx680eOBH16Axt2g+cmBD06peuaRz35h+eadPj3niY3j+ds5HQ65fcyYMSxbtoylS5cya9YszjrrLJYtW7a3qseECRNITk6muLiYXr16cdFFF5GSkrLfOdasWcOkSZN49dVXGTFiBB999BGjRo067ti1W0gdkR4XyeMXduKH+4fw1zNPIH9PGX/6IIs+//qaRz9bztqtu5wOUSmlfOa3nSXk7SqtvlLIqqmQv9a2WteDAVdKKejdu/d+5fKef/55unTpwkknnUR2djZr1qw56JgWLVrQtWtXAHr06MGGDRt8Eou2XNcxKbERjD6lFTcMaMkP6/KZuGATb8/bwITvf6V3i2RG9m7KsI4NtTVbKVWrVfa3rja5nvs8JDaF9ucGNCal6qvDtTAHSkzMvon2Zs2axYwZM/jhhx+Ijo5m0KBB1U52ExERsfe1y+WiuLjYJ7Focl1HiQj9WqfSr3Uq23aXMnlxDpMWbOKu95eS9FkYF3XP4PI+TWmVFut0qEopddSWZhcS5hLaNzpgfMmm+ZA9H4Y/CS79E6dUXRUXF8euXdV/K19YWEhSUhLR0dGsXLmSefPmBTQ2vfPUA6mxEdw0sBWjB7Rk7rp8Ji3YxBtzNzD+u185qWUyl3tbs3VSGqXqNxEZBjyHnVV3vDFmzAHbmwETgDTsrLqjjDE53m1NgfFAJmCAM40xG/wVa1Z2Ae0bxR9835r7PEQlQbfj7zeplApeKSkp9O/fn44dOxIVFUWDBg32bhs2bBhjx46lffv2tGvXjpNOOimgsWlyXY+EhAgnt0nl5Dap5O0q5cPF2UxasIk731tKckw4F/fI4PLeTWmRGnPkkyml6hQRcQEvAqcDOcBCEZlijFleZbengLeMMW+KyBDgceBK77a3gMeMMV+JSCzg8Vesbo/h59xCLujWZP8N29bCys/hlLshXO9jStV1EydOrHZ9REQE06ZNq3ZbZb/q1NRUli1btnf93Xff7bO4NLmup9LiIrhlUGtuOqUV363dxsT5m3jtu18ZN2c9fVok069VKt2aJtIlM7H6MldKqbqmN7DWGLMeQETeA84DqibXJwJ/8r6eCXzi3fdEINQY8xWAMWa3PwNdn7eb3aUVB/e3/uE/4AqH3qP9eXmllDosTa7ruZAQ4ZS2aZzSNo2tO0v4cHEOny7N5dmvV1M5s3rLtBi6ZibSLTORrplJnNAojjCXFppRqo5pAmRXWc4B+hywTxZwIbbryAVAnIikAG2BAhH5L9ACmAHcZ4xxH3gRERkNjAZo2rTpMQW61DuYsWtmlcljdufB0knQ9XKITT+m8yqllC9ocq32So+P5NbBrbl1cGt2lpTzc04hS7ML+HHTDuaszuO/S3IBiAgNoVOTBLpmJtK1aSJdMxNpkhiFaMkrpeq6u4EXROQaYA6QC7ixf0sGAN2ATcD7wDXAaweewBgzDhgH0LNnT3MsQWTlFBAbEUrL1CoDsheMA3cZ9NWpzpVSztLkWlUrPjKM/q1T6e+doMEYQ86OYpZmF+x9vD1vI+O/+xWwgya7ZibSralt4e6cmUhshP56KVWL5GIHI1bK8K7byxizGdtyjbdf9UXGmAIRyQGWVulS8glwEtUk176QlV1I54wEQkK8H+jL9sDCV+GEsyC1tT8uqZRSNabZj6oRESEzOZrM5GjO6dIYgLIKDyt/22mT7U024Z6x4nfv/tAyNYZOTRLo2CSBTk0SOLFxPHGR2n9bqSC1EGgjIi2wSfVlwMiqO4hIKrDdGOMB7sdWDqk8NlFE0owxecAQYJE/giwpd7Niy05uOKXlvpU/vgvFO6CftlorpZynybU6ZuGhIXTOSKRzRiJX9bXrCorKyMopZOmmAn7OLWTe+u18snTz3mNapsbQsUkCHZvEe58TiNeEWynHGWMqROQ24EtsKb4JxphfRORRYJExZgowCHhcRAy2W8it3mPdInI38LXY/mGLgVf9EeeKLTup8Bi6ZCTaFe4KO9V5Rm9oGthyW0opVR1NrpVPJUaHM7BtGgPbpu1dl7erlGWbC1mWU8jPuYUs2rCdKVn7Eu5mKdF7W7c7NraJd2J0uBPhK1WvGWOmAlMPWPdwldeTgcmHOPYroLNfA2TfzIxdKyuFrJgCBRth6GP+vrRSKogUFBQwceJEbrnllqM+9tlnn2X06NFER0f7ITJNrlUApMVFMLhdOoPb7RvBn7+7lGWbd7Ist5BluYVkZRfw+U9b9m7PTI6iY+METmgYT/PUaJqlxNA8JVqTbqXquaycQhrER9AwIRKMsZPGJLeCdmc6HZpSKoAKCgp46aWXjjm5HjVqlCbXqm5JiY04qIV7x54yftm8k5+9CffPuYVMW/bbfsclRIXRLGVfsl353DQlmrTYCK1YolQdl5VdsK9LyIbvYPOPcPYzEKIzzCpVn9x3332sW7eOrl27cvrpp5Oens4HH3xAaWkpF1xwAY888gh79uxhxIgR5OTk4Ha7eeihh/j999/ZvHkzgwcPJjU1lZkzZ/o8Nk2uVdBIignfO4NkpZJyN9nbi9iQX8TG/D1syN/Dxvwib0v3ZjxVCnlFh7v2S7abp8TQLCWajMRooiNcRIW5iAxz4QrRBFyp2qiwqJz12/ZwUY8Mu2Lu8xCdCl0udzYwpeq7affBbz/79pwNO8HwMYfcPGbMGJYtW8bSpUuZPn06kydPZsGCBRhjOPfcc5kzZw55eXk0btyYzz//HIDCwkISEhJ4+umnmTlzJqmpqYc8//HQ5FoFtcgwF20axNGmQdxB28oqPOQWFLMhfw+b8ov2Jt6rft/FjBW/U+6uvoRuuCuEyLAQosL3JdyRYfZ15bqIsBC77F0XGeYiLTaCdO/X0Q3jI0mICtOWcqUC6KfcAgDbcr11BayZDoMfgLAoR+NSSjlr+vTpTJ8+nW7dugGwe/du1qxZw4ABA/jzn//Mvffey9lnn82AAQMCEo8m16rWCg8NoUVqDC1SYw7a5vYYthQWszG/iM0FxZSUuykud1Nc5qG43G2Xy9yUVNjnynVbd5Xb9eWefceUu/fOVllVRGgIDeIjaRAfQYN4m3A3iI+kQUIkDeJsEt4gPpLIMP26Wilf6N8qla//PJDGCVEw9W8QGgW9rnc6LKXUYVqYA8EYw/3338+NN9540LYlS5YwdepUHnzwQU499VQefvjhas7gW5pcqzrJFSJkJEWTkXT8gxWMMZSUe9i2u5Tfdpbw+84SfissYeuuUn4rtMvLcguZseJ3Sso9Bx2fEBVGw/hI0uMjiApzUeEx9uH27H12ewzlbmOfPXa5wm2o8HiocBvKK/fxGCJcIWQkR9M0OYrMJNsFJjM5mqbJ0WQkRRERqsm8qptCQoRWabGwcwv89D70uAaik50OSynlgLi4OHbt2gXA0KFDeeihh7jiiiuIjY0lNzeXsLAwKioqSE5OZtSoUSQmJjJ+/Pj9jtVuIUo5RESICnftnUTnUIwx7CypYOvOEm8SXro3Ef/dm5Tn7Sol1CWEhoQQGiKEuoTo8FDvOu/6yteufftU7u9yCaXlHnJ2FPHrtj3MXp23X0IvAg3jI/cm25WPyuXU2HDtyqJqv/ljwbih761OR6KUckhKSgr9+/enY8eODB8+nJEjR9K3r510IzY2lnfeeYe1a9dyzz33EBISQlhYGC+//DIAo0ePZtiwYTRu3NgvAxrFVPd9dy3Us2dPs2iRXyYEUypoGWPI211K9vYiNm0vYlN+MZu2F+1d/m1nyX77R4W5aJocTePESFwhIRhj8BiDx4Dxns9jDB4PeIypso69+1FlOSk6nMaJkTROjKJxYhRNvM+NErQ7zNESkcXGmJ5OxxFIx3TfLt0FT3eAVoNhxJv+CUwpdUQrVqygffv2TocRENW918Pds7XlWqlaTERIj4skPS6SHs0O/nq8pNxNzo7ifcm397GlsBiPB0JCIEQEEUGAELHLdp13OSSkyvK+Z4Dte8qYtSqPvN2lB/VLT40Nt0l3QpQ3+Y7cm3w3TozSVnR1bBa/CaWF0P8OpyNRSqlqaXKtVB0WGeaidXosrdNj/XqdsgoPv+8sIbegmM3eR25BCZsLilmXt5s5a/IoKnPvd0x4aAiNEyJJjgknPDSEMFcI4S77HBYaQphLiPCur3yEu2Tv9vC9z4IrJAS3x/Zh39d/3bO3H/vBfdsr+7t7t3mX8X5AMN4XVT8wmAO2HbTd+3xio3j+eHpbX/+IFYC7HOa9DM1OhiY9nI5GKaWqpcm1Uuq4hYeGHLZPujGGncUV+5LvwmLv6xK27yml3G3YXVpBWYWHcrdNive99i67PZRVHDxgtKb267/u7dfuqtLP3RUie1vkASpfVW1cF+/aQzW4iwhpcRHHHKM6gvWzYWcOnP2005EopdQh+TW5FpFhwHOACxhvjBlzwPYI4C2gB5APXGqM2eDddj9wHeAG7jDGfOnPWJVS/iMiJESHkRAdxomN44/5PMbYluZyt4fyCptwl7ttRRWXSwirJmEOc4UQImgXlLqgzWlw4xxo0MnpSJRS2HtyXb+3HsvYRL8l1yLiAl4ETgdygIUiMsUYs7zKbtcBO4wxrUXkMuAJ4FIRORG4DOgANAZmiEhbY8z+3ysrpeoVESHM2zWEcKejUY5o1MXpCJRSQGRkJPn5+aSkpNTZBNsYQ35+PpGRkUd1nD9brnsDa40x6wFE5D3gPKBqcn0e8Hfv68nAC2L/hc4D3jPGlAK/isha7/l+8GO8SimllFKqBjIyMsjJySEvL8/pUPwqMjKSjIyMozrGn8l1EyC7ynIO0OdQ+xhjKkSkEEjxrp93wLFN/BeqUkoppZSqqbCwMFq0aOF0GEEpxOkAjoeIjBaRRSKyqK5/clJKKaWUUsHPn8l1LpBZZTnDu67afUQkFEjADmysybEYY8YZY3oaY3qmpaX5MHSllFJKKaWOnj+T64VAGxFpISLh2AGKUw7YZwpwtff1xcA3xg7LnAJcJiIRItICaAMs8GOsSimllFJKHTe/Tn8uImcCz2JL8U0wxjwmIo8Ci4wxU0QkEngb6AZsBy6rMgDyAeAPQAVwlzFm2hGulQdsPIYwU4Ftx3CcPwRLLMESB2gs1QmWOCB4YgmWOODYY2lmjKlXX8HVgft2sMQBwRNLsMQBwRNLsMQBGkt1fH7P9mtyXRuIyKJDzQ0faMESS7DEARpLMMcBwRNLsMQBwRVLXRUsP+NgiQOCJ5ZgiQOCJ5ZgiQM0lkDFUasHNCqllFJKKRVMNLlWSimllFLKRzS5hnFOB1BFsMQSLHGAxlKdYIkDgieWYIkDgiuWuipYfsbBEgcETyzBEgcETyzBEgdoLNXxeRz1vs+1UkoppZRSvqIt10oppZRSSvlIvU6uRWSYiKwSkbUicp9DMWSKyEwRWS4iv4jInU7EcUBMLhH5UUT+53AciSIyWURWisgKEenrUBx/9P7bLBORSd4SkoG69gQR2Soiy6qsSxaRr0Rkjfc5ycFY/s/77/OTiHwsIolOxFFl259FxIhIqr/jOFwsInK79+fyi4g8GYhY6oNguGd74wiq+7besw+KQ+/Zh45F79kBuGfX2+RaRFzAi8Bw4ETgchE50YFQKoA/G2NOBE4CbnUojqruBFY4HAPAc8AXxpgTgC44EJOINAHuAHoaYzpia7ZfFsAQ3gCGHbDuPuBrY0wb4GvvslOxfAV0NMZ0BlYD9zsUByKSCZwBbApADIeMRUQGA+cBXYwxHYCnAhhPnRVE92wIvvu23rO99J59xFj0nh2Ae3a9Ta6B3sBaY8x6Y0wZ8B72hxtQxpgtxpgl3te7sDejJoGOo5KIZABnAeOdisEbRwJwCvAagDGmzBhT4FA4oUCUiIQC0cDmQF3YGDMHO8FSVecBb3pfvwmc71QsxpjpxpgK7+I8IMOJOLyeAf4CBGwgySFiuRkYY4wp9e6zNVDx1HFBcc+G4Lpv6z27WnrPPkQses8OzD27PifXTYDsKss5OJjUAohIc+xslfMdDONZ7C+7x8EYAFoAecDr3q87x4tITKCDMMbkYj/FbgK2AIXGmOmBjuMADYwxW7yvfwMaOBlMFX8ADjuTqr+IyHlArjEmy4nrH6AtMEBE5ovIbBHp5XRAdUTQ3bMhKO7bz6L37L30nn1U9J5t+fyeXZ+T66AiIrHAR9ip3nc6FMPZwFZjzGInrn+AUKA78LIxphuwh8B9lbaXt2/cedg/HI2BGBEZFeg4DsXYcj+Ol/wRkQewX5W/68C1o4G/Ag8H+tqHEAokY7sL3AN8ICLibEjKH5y+b+s9+2B6z64ZvWfvx+f37PqcXOcCmVWWM7zrAk5EwrA36HeNMf91Igav/sC5IrIB+5XrEBF5x6FYcoAcY0xla9Bk7I070E4DfjXG5BljyoH/Av0ciKOq30WkEYD32dFuByJyDXA2cIVxprZnK+wf0izv724GsEREGjoQC9jf3f8aawG2RTEgg3XquKC5Z0PQ3Lf1nn0wvWcfgd6zD+Lze3Z9Tq4XAm1EpIWIhGMHPEwJdBDeT0evASuMMU8H+vpVGWPuN8ZkGGOaY38e3xhjHPnEb4z5DcgWkXbeVacCyx0IZRNwkohEe/+tTsX5gUNTgKu9r68GPnUqEBEZhv1K+lxjTJETMRhjfjbGpBtjmnt/d3OA7t7fISd8AgwGEJG2QDiwzaFY6pKguGdD8Ny39Z5dLb1nH4bes6v1Cb6+Zxtj6u0DOBM7WnYd8IBDMZyM/YroJ2Cp93FmEPxsBgH/cziGrsAi78/mEyDJoTgeAVYCy4C3gYgAXnsStt9gOfYGdB2Qgh1xvgaYASQ7GMtabD/Yyt/dsU7EccD2DUCqgz+TcOAd7+/LEmBIoH5f6vojGO7Z3jiC7r6t9+z94tB79qFj0Xt2AO7ZOkOjUkoppZRSPlKfu4UopZRSSinlU5pcK6WUUkop5SOaXCullFJKKeUjmlwrpZRSSinlI5pcK6WUUkop5SOaXKs6S0TcIrK0ysNns4WJSHMRWXYU+8eIyAzv6+9EJNRXsSilVF2g92xVV+gvi6rLio0xXZ0Owqsv8IN3at49xpgKpwNSSqkgo/dsVSdoy7Wqd0Rkg4g8KSI/i8gCEWntXd9cRL4RkZ9E5GsRaepd30BEPhaRLO+jcipdl4i8KiK/iMh0EYmq5lqtRGQptkD9SGAx0MXbKpMemHeslFK1l96zVW2jybWqy6IO+Irx0irbCo0xnYAXgGe96/4DvGmM6Qy8CzzvXf88MNsY0wXoDvziXd8GeNEY0wEoAC46MABjzDpvS8xioDfwJnZ2qq7GmK2+e6tKKVXr6T1b1Qk6Q6Oqs0RktzEmtpr1G7DTm64XkTDgN2NMiohsAxoZY8q967cYY1JFJA/IMMaUVjlHc+ArY0wb7/K9QJgx5p+HiGWhMaaXiHwE3GmMyfH1+1VKqdpM79mqrtCWa1VfmUO8PhqlVV67qWYMg4iM9Q6iaeP9qnEY8D8R+eMxXlMppeojvWerWkOTa1VfXVrl+Qfv67nAZd7XVwDfel9/DdwMICIuEUmo6UWMMTcBjwD/AM4HPvd+vfjMcUWvlFL1i96zVa2h1UJUXRblbXmo9IUxprK0U5KI/IRtybjcu+524HURuQfIA671rr8TGCci12FbO24GthxFHAOBt4ABwOxjeSNKKVUP6D1b1Qna51rVO97+ez2NMducjkUppdTh6T1b1TbaLUQppZRSSikf0ZZrpZRSSimlfERbrpVSSimllPIRTa6VUkoppZTyEU2ulVJKKaWU8hFNrpVSSimllPIRTa6VUkoppZTyEU2ulVJKKaWU8pH/B4qUxFWIekgiAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 864x216 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axs = plt.subplots(1,2, figsize=(12,3))\n",
    "\n",
    "axs[0].plot(metrics[\"train_loss\"], label=\"train\")\n",
    "axs[0].plot(metrics[\"test_loss\"], label=\"test\")\n",
    "axs[0].legend()\n",
    "axs[0].set_xlabel(\"Epoch #\")\n",
    "axs[0].set_ylabel(\"Loss\")\n",
    "\n",
    "\n",
    "axs[1].plot(metrics[\"train_accuracy\"], label=\"train\")\n",
    "axs[1].plot(metrics[\"test_accuracy\"], label=\"test\")\n",
    "axs[1].legend()\n",
    "axs[1].set_xlabel(\"Epoch #\")\n",
    "axs[1].set_ylabel(\"Accuracy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ae58e9-826f-437f-b497-d4174885d461",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (RL)",
   "language": "python",
   "name": "py-rl"
  },
  "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.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}