{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-MYot7DJh9kk"
      },
      "source": [
        "See code at https://github.com/google-research/vision_transformer/\n",
        "\n",
        "See papers at\n",
        "\n",
        "- Vision Transformer: https://arxiv.org/abs/2010.11929\n",
        "- MLP-Mixer: https://arxiv.org/abs/2105.01601\n",
        "- How to train your ViT: https://arxiv.org/abs/2106.10270\n",
        "- When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations: https://arxiv.org/abs/2106.01548\n",
        "\n",
        "This Colab allows you to run the [JAX](https://jax.readthedocs.org) implementation of the Vision Transformer.\n",
        "\n",
        "If you just want to load a pre-trained checkpoint from a large repository and\n",
        "directly use it for inference, you probably want to go [this Colab](https://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax_augreg.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sXhZm0kpPpH6"
      },
      "source": [
        "##### Copyright 2021 Google LLC."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cfe6jvTCo1yQ"
      },
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "KfmzfvFxPuk7",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iOVCm4CnP1Do"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cyD76dm5JaeW"
      },
      "source": [
        "### Setup\n",
        "\n",
        "Needs to be executed once in every VM.\n",
        "\n",
        "The cell below downloads the code from Github and install necessary dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 74
        },
        "id": "zZvI8OXt78sj",
        "outputId": "61904a04-1072-42cb-9f8d-a69ccf8710a2",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<h1 style=\"color:red\">CHANGES NOT PERSISTED</h1>"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "#@markdown Select whether you would like to store data in your personal drive.\n",
        "#@markdown\n",
        "#@markdown If you select **yes**, you will need to authorize Colab to access\n",
        "#@markdown your personal drive\n",
        "#@markdown\n",
        "#@markdown If you select **no**, then any changes you make will diappear when\n",
        "#@markdown this Colab's VM restarts after some time of inactivity...\n",
        "use_gdrive = 'no'  #@param [\"yes\", \"no\"]\n",
        "\n",
        "if use_gdrive == 'yes':\n",
        "  from google.colab import drive\n",
        "  drive.mount('/gdrive')\n",
        "  root = '/gdrive/My Drive/vision_transformer_colab'\n",
        "  import os\n",
        "  if not os.path.isdir(root):\n",
        "    os.mkdir(root)\n",
        "  os.chdir(root)\n",
        "  print(f'\\nChanged CWD to \"{root}\"')\n",
        "else:\n",
        "  from IPython import display\n",
        "  display.display(display.HTML(\n",
        "      '<h1 style=\"color:red\">CHANGES NOT PERSISTED</h1>'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "GeEy6gN71CDa",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "10bc6db8-2df4-4285-f986-a7e840cf6961"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'vision_transformer'...\n",
            "remote: Enumerating objects: 49, done.\u001b[K\n",
            "remote: Counting objects: 100% (49/49), done.\u001b[K\n",
            "remote: Compressing objects: 100% (44/44), done.\u001b[K\n",
            "remote: Total 49 (delta 6), reused 16 (delta 1), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (49/49), 1.86 MiB | 13.63 MiB/s, done.\n",
            "Resolving deltas: 100% (6/6), done.\n",
            "Already up to date.\n"
          ]
        }
      ],
      "source": [
        "# Clone repository and pull latest changes.\n",
        "![ -d vision_transformer ] || git clone --depth=1 https://github.com/google-research/vision_transformer\n",
        "!cd vision_transformer && git pull"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "sCN4d-GQJdU4",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "93aa7c2d-cca7-4fa6-f9d0-645f06219ea0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  Running command git clone --filter=blob:none --quiet https://github.com/google/flaxformer /tmp/pip-req-build-bpeaplsy\n"
          ]
        }
      ],
      "source": [
        "# Colab already includes most of the dependencies, so we only install the delta:\n",
        "!pip install einops>=0.3.0 ml-collections>=0.1.0 aqtp>=0.2.0 clu>=0.0.3 git+https://github.com/google/flaxformer tensorflow-text>=2.9.0\n",
        "\n",
        "#!pip install -qr vision_transformer/vit_jax/requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcLBTSXuNjK6"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "BcnlJF7FKTfD",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d4c2a531-59d4-4b06-d0cd-cc9313620f5f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "gs://vit_models/imagenet21k+imagenet2012/:\n",
            "377.57 MiB  2020-11-30T16:17:02Z  gs://vit_models/imagenet21k+imagenet2012/R50+ViT-B_16.npz\n",
            "330.29 MiB  2020-10-29T17:05:52Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_16-224.npz\n",
            " 331.4 MiB  2020-10-20T11:48:22Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_16.npz\n",
            "336.89 MiB  2020-10-20T11:47:36Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_32.npz\n",
            "334.78 MiB  2021-03-12T09:04:16Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_8.npz\n",
            "  1.13 GiB  2020-10-29T17:08:31Z  gs://vit_models/imagenet21k+imagenet2012/ViT-L_16-224.npz\n",
            "  1.14 GiB  2020-10-20T11:53:44Z  gs://vit_models/imagenet21k+imagenet2012/ViT-L_16.npz\n",
            "  1.14 GiB  2020-10-20T11:50:56Z  gs://vit_models/imagenet21k+imagenet2012/ViT-L_32.npz\n",
            "\n",
            "gs://vit_models/imagenet21k/:\n",
            "450.23 MiB  2021-01-20T14:12:43Z  gs://vit_models/imagenet21k/R26+ViT-B_32.npz\n",
            "439.85 MiB  2020-11-30T10:10:15Z  gs://vit_models/imagenet21k/R50+ViT-B_16.npz\n",
            "  1.31 GiB  2021-01-20T14:11:54Z  gs://vit_models/imagenet21k/R50+ViT-L_32.npz\n",
            "393.69 MiB  2020-10-22T21:38:39Z  gs://vit_models/imagenet21k/ViT-B_16.npz\n",
            "400.01 MiB  2020-11-02T08:30:56Z  gs://vit_models/imagenet21k/ViT-B_32.npz\n",
            "393.72 MiB  2021-03-10T13:28:28Z  gs://vit_models/imagenet21k/ViT-B_8.npz\n",
            "  2.46 GiB  2020-11-03T10:46:11Z  gs://vit_models/imagenet21k/ViT-H_14.npz\n",
            "  1.22 GiB  2020-11-09T14:39:51Z  gs://vit_models/imagenet21k/ViT-L_16.npz\n",
            "  1.23 GiB  2020-11-02T08:35:10Z  gs://vit_models/imagenet21k/ViT-L_32.npz\n",
            "TOTAL: 17 objects, 14306096550 bytes (13.32 GiB)\n",
            " 330.3 MiB  2021-07-13T19:39:09Z  gs://vit_models/sam/ViT-B_16.npz\n",
            "336.61 MiB  2021-07-13T19:39:10Z  gs://vit_models/sam/ViT-B_32.npz\n",
            "  1.13 GiB  2021-07-13T19:39:38Z  gs://vit_models/sam/ViT-L_16.npz\n",
            "  1.14 GiB  2021-07-13T19:39:38Z  gs://vit_models/sam/ViT-L_32.npz\n",
            "252.57 MiB  2022-02-18T06:43:52Z  gs://vit_models/sam/ViT-S_16.npz\n",
            "TOTAL: 5 objects, 3407859850 bytes (3.17 GiB)\n",
            "       6 B  2021-06-28T13:07:12Z  gs://mixer_models/sam_$folder$\n",
            "\n",
            "gs://mixer_models/gsam/:\n",
            "228.47 MiB  2022-07-21T05:47:43Z  gs://mixer_models/gsam/Mixer-B_16.npz\n",
            "230.04 MiB  2022-07-21T05:48:18Z  gs://mixer_models/gsam/Mixer-B_32.npz\n",
            " 70.71 MiB  2022-07-21T05:48:30Z  gs://mixer_models/gsam/Mixer-S_16.npz\n",
            " 72.91 MiB  2022-07-21T05:48:42Z  gs://mixer_models/gsam/Mixer-S_32.npz\n",
            " 78.79 MiB  2022-07-21T05:48:56Z  gs://mixer_models/gsam/Mixer-S_8.npz\n",
            "\n",
            "gs://mixer_models/imagenet1k/:\n",
            "228.47 MiB  2021-05-05T14:09:01Z  gs://mixer_models/imagenet1k/Mixer-B_16.npz\n",
            "794.29 MiB  2021-05-05T14:09:02Z  gs://mixer_models/imagenet1k/Mixer-L_16.npz\n",
            "\n",
            "gs://mixer_models/imagenet21k/:\n",
            "289.61 MiB  2021-05-05T14:09:11Z  gs://mixer_models/imagenet21k/Mixer-B_16.npz\n",
            "875.78 MiB  2021-05-05T14:09:12Z  gs://mixer_models/imagenet21k/Mixer-L_16.npz\n",
            "\n",
            "gs://mixer_models/sam/:\n",
            "228.47 MiB  2021-06-28T13:08:09Z  gs://mixer_models/sam/Mixer-B_16.npz\n",
            "230.04 MiB  2021-06-28T13:08:08Z  gs://mixer_models/sam/Mixer-B_32.npz\n",
            "TOTAL: 12 objects, 3489219014 bytes (3.25 GiB)\n"
          ]
        }
      ],
      "source": [
        "# Shows all available pre-trained models.\n",
        "!gsutil ls -lh gs://vit_models/imagenet*\n",
        "!gsutil ls -lh gs://vit_models/sam\n",
        "!gsutil ls -lh gs://mixer_models/*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "6ztOhq_fzZyO",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e626faa7-f1d8-4f05-90eb-589fd0383b06"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Copying gs://vit_models/imagenet21k/ViT-B_32.npz...\n",
            "/ [0 files][    0.0 B/400.0 MiB]                                                \r==> NOTE: You are downloading one or more large file(s), which would\n",
            "run significantly faster if you enabled sliced object downloads. This\n",
            "feature is enabled by default but requires that compiled crcmod be\n",
            "installed (see \"gsutil help crcmod\").\n",
            "\n",
            "/ [1 files][400.0 MiB/400.0 MiB]   33.6 MiB/s                                   \n",
            "Operation completed over 1 objects/400.0 MiB.                                    \n"
          ]
        }
      ],
      "source": [
        "# Download a pre-trained model.\n",
        "\n",
        "# Note: you can really choose any of the above, but this Colab has been tested\n",
        "# with the models of below selection...\n",
        "model_name = 'ViT-B_32'  #@param [\"ViT-B_32\", \"Mixer-B_16\"]\n",
        "\n",
        "if model_name.startswith('ViT'):\n",
        "  ![ -e \"$model_name\".npz ] || gsutil cp gs://vit_models/imagenet21k/\"$model_name\".npz .\n",
        "if model_name.startswith('Mixer'):\n",
        "  ![ -e \"$model_name\".npz ] || gsutil cp gs://mixer_models/imagenet21k/\"$model_name\".npz .\n",
        "\n",
        "import os\n",
        "assert os.path.exists(f'{model_name}.npz')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "4EzOChfJeVrU",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "77797aa3-7d73-4889-d286-4a8d6032eb53"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "No TPU detected. Can be changed under \"Runtime/Change runtime type\".\n"
          ]
        }
      ],
      "source": [
        "# Google Colab \"TPU\" runtimes are configured in \"2VM mode\", meaning that JAX\n",
        "# cannot see the TPUs because they're not directly attached. Instead we need to\n",
        "# setup JAX to communicate with a second machine that has the TPUs attached.\n",
        "import os\n",
        "if 'google.colab' in str(get_ipython()) and 'COLAB_TPU_ADDR' in os.environ:\n",
        "  import jax\n",
        "  import jax.tools.colab_tpu\n",
        "  jax.tools.colab_tpu.setup_tpu()\n",
        "  print('Connected to TPU.')\n",
        "else:\n",
        "  print('No TPU detected. Can be changed under \"Runtime/Change runtime type\".')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "igqZ6qYNeHWo",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c509f24a-bef4-4e3c-b649-80544ac1f0a0"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[cuda(id=0)]"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ],
      "source": [
        "from absl import logging\n",
        "import flax\n",
        "import jax\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import optax\n",
        "import tqdm\n",
        "\n",
        "logging.set_verbosity(logging.INFO)\n",
        "\n",
        "# Shows the number of available devices.\n",
        "# In a CPU/GPU runtime this will be a single device.\n",
        "# In a TPU runtime this will be 8 cores.\n",
        "jax.local_devices()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "id": "9TuMn31fNj0T",
        "outputId": "53ecf6de-7a53-4e6f-e6c8-8ea736ce39aa",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "      ((filepath) => {{\n",
              "        if (!google.colab.kernel.accessAllowed) {{\n",
              "          return;\n",
              "        }}\n",
              "        google.colab.files.view(filepath);\n",
              "      }})(\"/content/vision_transformer/vit_jax/configs/common.py\")"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "      ((filepath) => {{\n",
              "        if (!google.colab.kernel.accessAllowed) {{\n",
              "          return;\n",
              "        }}\n",
              "        google.colab.files.view(filepath);\n",
              "      }})(\"/content/vision_transformer/vit_jax/configs/models.py\")"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "      ((filepath) => {{\n",
              "        if (!google.colab.kernel.accessAllowed) {{\n",
              "          return;\n",
              "        }}\n",
              "        google.colab.files.view(filepath);\n",
              "      }})(\"/content/vision_transformer/vit_jax/checkpoint.py\")"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "      ((filepath) => {{\n",
              "        if (!google.colab.kernel.accessAllowed) {{\n",
              "          return;\n",
              "        }}\n",
              "        google.colab.files.view(filepath);\n",
              "      }})(\"/content/vision_transformer/vit_jax/input_pipeline.py\")"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "      ((filepath) => {{\n",
              "        if (!google.colab.kernel.accessAllowed) {{\n",
              "          return;\n",
              "        }}\n",
              "        google.colab.files.view(filepath);\n",
              "      }})(\"/content/vision_transformer/vit_jax/models.py\")"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "      ((filepath) => {{\n",
              "        if (!google.colab.kernel.accessAllowed) {{\n",
              "          return;\n",
              "        }}\n",
              "        google.colab.files.view(filepath);\n",
              "      }})(\"/content/vision_transformer/vit_jax/train.py\")"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "# Open some code files in a split editor on the right.\n",
        "# You can open more files in the file tab on the left.\n",
        "from google.colab import files\n",
        "files.view('vision_transformer/vit_jax/configs/common.py')\n",
        "files.view('vision_transformer/vit_jax/configs/models.py')\n",
        "files.view('vision_transformer/vit_jax/checkpoint.py')\n",
        "files.view('vision_transformer/vit_jax/input_pipeline.py')\n",
        "files.view('vision_transformer/vit_jax/models.py')\n",
        "files.view('vision_transformer/vit_jax/train.py')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "sjN0_b-YbaHu",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Import files from repository.\n",
        "# Updating the files in the editor on the right will immediately update the\n",
        "# modules by re-importing them.\n",
        "\n",
        "import sys\n",
        "if './vision_transformer' not in sys.path:\n",
        "  sys.path.append('./vision_transformer')\n",
        "\n",
        "%load_ext autoreload\n",
        "%autoreload 2\n",
        "\n",
        "from vit_jax import checkpoint\n",
        "from vit_jax import input_pipeline\n",
        "from vit_jax import utils\n",
        "from vit_jax import models\n",
        "from vit_jax import train\n",
        "from vit_jax.configs import common as common_config\n",
        "from vit_jax.configs import models as models_config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "GojydzsXgknd",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Helper functions for images.\n",
        "\n",
        "labelnames = dict(\n",
        "  # https://www.cs.toronto.edu/~kriz/cifar.html\n",
        "  cifar10=('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'),\n",
        "  # https://www.cs.toronto.edu/~kriz/cifar.html\n",
        "  cifar100=('apple', 'aquarium_fish', 'baby', 'bear', 'beaver', 'bed', 'bee', 'beetle', 'bicycle', 'bottle', 'bowl', 'boy', 'bridge', 'bus', 'butterfly', 'camel', 'can', 'castle', 'caterpillar', 'cattle', 'chair', 'chimpanzee', 'clock', 'cloud', 'cockroach', 'couch', 'crab', 'crocodile', 'cup', 'dinosaur', 'dolphin', 'elephant', 'flatfish', 'forest', 'fox', 'girl', 'hamster', 'house', 'kangaroo', 'computer_keyboard', 'lamp', 'lawn_mower', 'leopard', 'lion', 'lizard', 'lobster', 'man', 'maple_tree', 'motorcycle', 'mountain', 'mouse', 'mushroom', 'oak_tree', 'orange', 'orchid', 'otter', 'palm_tree', 'pear', 'pickup_truck', 'pine_tree', 'plain', 'plate', 'poppy', 'porcupine', 'possum', 'rabbit', 'raccoon', 'ray', 'road', 'rocket', 'rose', 'sea', 'seal', 'shark', 'shrew', 'skunk', 'skyscraper', 'snail', 'snake', 'spider', 'squirrel', 'streetcar', 'sunflower', 'sweet_pepper', 'table', 'tank', 'telephone', 'television', 'tiger', 'tractor', 'train', 'trout', 'tulip', 'turtle', 'wardrobe', 'whale', 'willow_tree', 'wolf', 'woman', 'worm')\n",
        ")\n",
        "def make_label_getter(dataset):\n",
        "  \"\"\"Returns a function converting label indices to names.\"\"\"\n",
        "  def getter(label):\n",
        "    if dataset in labelnames:\n",
        "      return labelnames[dataset][label]\n",
        "    return f'label={label}'\n",
        "  return getter\n",
        "\n",
        "def show_img(img, ax=None, title=None):\n",
        "  \"\"\"Shows a single image.\"\"\"\n",
        "  if ax is None:\n",
        "    ax = plt.gca()\n",
        "  ax.imshow(img[...])\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",
        "    img = (img + 1) / 2  # Denormalize\n",
        "    show_img(img, axs[i // n][i % n], title)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QZfK1vIIMmFz"
      },
      "source": [
        "### Load dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "TSAVpYtP5VaE",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "dataset = 'cifar10'\n",
        "batch_size = 512\n",
        "config = common_config.with_dataset(common_config.get_config(), dataset)\n",
        "config.batch = batch_size\n",
        "config.pp.crop = 224"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "ruzdzpsMNhGm",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 445,
          "referenced_widgets": [
            "88b7391eac6346b6848f9a0b42983bb3",
            "ff4ce5fb8b4d42ff9567995085be3fe0",
            "60530198d0044832b19d2a465455a810",
            "9b9d565e72c34a4890404eec2ac7d377",
            "2a23cdf1a8a847f7a2ce8f22b5b77bb1",
            "9a15e456d4f54167afe81f560918064d",
            "afc6d112cf5646788bdf31894a232222",
            "1633effb0fe24526a419fc2411804816",
            "d140d24e2cc64a2fab959d8bbd57d7bc",
            "beb0f1d2984b4c91a17ad9ff3cfd7dc0",
            "2c679fdde70f4e1394f89f7d9edd62ba",
            "a2006ee171e4402fbb2f28e3ab4cafa5",
            "8d83fafc13f641fdb5db20bc7a74c906",
            "54fd4ba8000e4a6cafacbdac591e4542",
            "7c70d6032fa0415a87917de014492c5e",
            "1b90aa4ddccd44249938f46bf1c57d24",
            "898f31eb6de6488186fafd9999c608eb",
            "ea6d78f0e27a409cac736eea6f0bebd3",
            "4a1c284e25914060b0b439dc74d4798e",
            "33d74a0d0c554135926b6eca6e32e8c4",
            "8bbe1a3515bd45eaa246328f496427c8",
            "9b66f4cbff664452a4d03cb0716518a8",
            "95ac0ffea8274bf0b9eb5fd5214c7f69",
            "740c230a0bdb49d8a60f0e4a5c0ab28e",
            "170456fc38604ec398fba752c90cd417",
            "521b5953067f4cc1bb10003591f997b9",
            "ca515be22fdf4098b695334ee7ffe7b2",
            "5bff4f12df1549ba94a1d1e5e169c80c",
            "adfc4aaac7c045059d12c1e1cf37f4c2",
            "5fcecb05c8c24e9285aff1536b185bea",
            "3c662c7be9964edfa0628cf8effd3d8f",
            "1b1f1eb77092490ca6c5bbac60627ad6",
            "e3f6b487be0b44f180113151c6392c92",
            "3eff2841b3f84abea71febf38c1dee8e",
            "7d6dc9752873434dab3c1bc09bd2d7a4",
            "7e8e8e1d697749b9b07fbb8da3f0c3b7",
            "9923a1e8acfd41cd93879de31055406f",
            "77f9a36bd6384955871d3b7f3deddb9d",
            "b2ea0941ad0646f092d7bd5113dd1725",
            "c05389b29b0a40c7af3495f1de9d4897",
            "46da21bfa2934247945d2c3f0f121d2c",
            "fc2d97654d744874aecfa15f4bb2c2e3",
            "45a4c68ca2694c23acf5933739375939",
            "c0076d262b684281b18db3c82810fcc2",
            "9c224504cbcb483091f698e7189d0c90",
            "203ec9c61ec94b53bcc86ac8952bb780",
            "eea538f2b1604587a51fee35ce16f5be",
            "fd9cd18bf458434e933ae970f80ef6dc",
            "58decf174a574a45a6fb2d7685598410",
            "07fcf33390e14a7dbb62523c80cc8c91",
            "0029395e502e4505b51434c8211221c9",
            "52b495f2c2d04372aded6bb4dbf1f5c6",
            "26febddaef2a4db795f2b3319cd5a8ef",
            "928ef89206014b9282a9c8c36ae04fd8",
            "64f81fce9df04a098557fa8d68fbf8fe",
            "858d201e6ca6454c9c60c6e3cc0daa1b",
            "cc9699e15bfc4573986c18cf5005b36e",
            "50ccc5e2256a48a38a2f1996956d5f73",
            "83214684d4af410f82a3830ffb8d3b48",
            "63f0ce0b2bf546458aaa35efdb3a2f07",
            "ce720edde76b4613828a4bf4b97d4106",
            "cc1e72ca287243fe8196c332dbb27423",
            "3ef22f6fa51947bc90c1e6e3b2b4a8ef",
            "d9f1aee8f1fb44b88b6d83d5788d1ea0",
            "d39bca2d23be4adab61bd277d6c80e52",
            "c1c497d8574c493b8d90366be94f8e1b",
            "3f61424d60d9474cb7f6cdaa16e757d2",
            "e0d5e0d2c66d41f7b2b392b12476076d",
            "d8105d46b03c4063ad00fc7ceedd6ce4",
            "69e4d51c37994ee39143f10beca47200",
            "7069223e26324cbcbed7b44258ee73c3",
            "cc8544485c584c43ae0825a7b96510ae",
            "f22bce220a7049cfa81c561835c64a6e",
            "6c90f244d5c34a56b53ac492ba92988c",
            "98c875a0af064305b540b33d341f85e9",
            "1d27837081434b88aa2526d2c7bbfe53",
            "530d6ca4cc82481d8c1376bbed324d04",
            "cd8b4d532cf94190ae89032bee8d1d64",
            "60149eabd64d4d6ebf1f5659478e8b24",
            "674cce53dac842fb8d8f30d58270f91c",
            "d226d17a362246d28eacee5d3084c377",
            "dc237fc9345f400394c01e1ac11dc048",
            "31b648ce84b04155913adff950edf4e2",
            "8e986d355825423ea01808765ef7385a",
            "1af824e5c2c548b69ace8d5bbc207c16",
            "410e720c1cc945ae8f56362bdb96879b",
            "082e18037c274d219da99bd58a4c93f1",
            "2570ced1bb0646d08f323d54b22bd7f8"
          ]
        },
        "outputId": "98f18937-a064-4401-d820-e5d61dcd91b7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: cifar10/3.0.2\n",
            "INFO:absl:Load dataset info from /tmp/tmpyg6qqyl7tfds\n",
            "INFO:absl:For 'cifar10/3.0.2': fields info.[citation, splits, supervised_keys, module_name] differ on disk and in the code. Keeping the one from code.\n",
            "INFO:absl:Generating dataset cifar10 (/root/tensorflow_datasets/cifar10/3.0.2)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading and preparing dataset 162.17 MiB (download: 162.17 MiB, generated: 132.40 MiB, total: 294.58 MiB) to /root/tensorflow_datasets/cifar10/3.0.2...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Dl Completed...: 0 url [00:00, ? url/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "88b7391eac6346b6848f9a0b42983bb3"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Dl Size...: 0 MiB [00:00, ? MiB/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a2006ee171e4402fbb2f28e3ab4cafa5"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Extraction completed...: 0 file [00:00, ? file/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "95ac0ffea8274bf0b9eb5fd5214c7f69"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Downloading https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz into /root/tensorflow_datasets/downloads/cs.toronto.edu_kriz_cifar-10-binaryODHPtIjLh3oLcXirEISTO7dkzyKjRCuol6lV8Wc6C7s.tar.gz.tmp.7757bd27a0a24767a589a744255c6d2d...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Generating splits...:   0%|          | 0/2 [00:00<?, ? splits/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "3eff2841b3f84abea71febf38c1dee8e"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Generating train examples...:   0%|          | 0/50000 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "9c224504cbcb483091f698e7189d0c90"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Shuffling /root/tensorflow_datasets/cifar10/3.0.2.incompleteHVPR1P/cifar10-train.tfrecord*...:   0%|          …"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "858d201e6ca6454c9c60c6e3cc0daa1b"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Done writing /root/tensorflow_datasets/cifar10/3.0.2.incompleteHVPR1P/cifar10-train.tfrecord*. Number of examples: 50000 (shards: [50000])\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Generating test examples...:   0%|          | 0/10000 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "3f61424d60d9474cb7f6cdaa16e757d2"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Shuffling /root/tensorflow_datasets/cifar10/3.0.2.incompleteHVPR1P/cifar10-test.tfrecord*...:   0%|          |…"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "cd8b4d532cf94190ae89032bee8d1d64"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Done writing /root/tensorflow_datasets/cifar10/3.0.2.incompleteHVPR1P/cifar10-test.tfrecord*. Number of examples: 10000 (shards: [10000])\n",
            "INFO:absl:Creating a tf.data.Dataset reading 1 files located in folders: /root/tensorflow_datasets/cifar10/3.0.2.\n",
            "INFO:absl:Constructing tf.data.Dataset cifar10 for split train[:98%], from /root/tensorflow_datasets/cifar10/3.0.2\n",
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Dataset cifar10 downloaded and prepared to /root/tensorflow_datasets/cifar10/3.0.2. Subsequent calls will reuse this data.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2\n",
            "INFO:absl:Reusing dataset cifar10 (/root/tensorflow_datasets/cifar10/3.0.2)\n",
            "INFO:absl:Creating a tf.data.Dataset reading 1 files located in folders: /root/tensorflow_datasets/cifar10/3.0.2.\n",
            "INFO:absl:Constructing tf.data.Dataset cifar10 for split test, from /root/tensorflow_datasets/cifar10/3.0.2\n",
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2\n",
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2\n"
          ]
        }
      ],
      "source": [
        "# For details about setting up datasets, see input_pipeline.py on the right.\n",
        "ds_train = input_pipeline.get_data_from_tfds(config=config, mode='train')\n",
        "ds_test = input_pipeline.get_data_from_tfds(config=config, mode='test')\n",
        "num_classes = input_pipeline.get_dataset_info(dataset, 'train')['num_classes']\n",
        "del config  # Only needed to instantiate datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "7c-LfxOJdj8_",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "528e0adf-ad31-4cb3-cd07-5462fe322a9e"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1, 512, 224, 224, 3)"
            ]
          },
          "metadata": {},
          "execution_count": 14
        }
      ],
      "source": [
        "# Fetch a batch of test images for illustration purposes.\n",
        "batch = next(iter(ds_test.as_numpy_iterator()))\n",
        "# Note the shape : [num_local_devices, local_batch_size, h, w, c]\n",
        "batch['image'].shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "rL0jQRBCgeJA",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 752
        },
        "outputId": "c31ee9f9-0d7a-4b53-ad70-f33fc4aa52fb"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 900x900 with 9 Axes>"
            ],
            "image/png": "[encoded data, 478040 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# Show some images with their labels.\n",
        "images, labels = batch['image'][0][:9], batch['label'][0][:9]\n",
        "titles = map(make_label_getter(dataset), labels.argmax(axis=1))\n",
        "show_img_grid(images, titles)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "jFqi3h7yMEsB",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 752
        },
        "outputId": "e9efe5f9-81d1-498a-d999-26cbc0c6e6c1"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 900x900 with 9 Axes>"
            ],
            "image/png": "[encoded data, 385540 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "# Same as above, but with train images.\n",
        "# Note how images are cropped/scaled differently.\n",
        "# Check out input_pipeline.get_data() in the editor at your right to see how the\n",
        "# images are preprocessed differently.\n",
        "batch = next(iter(ds_train.as_numpy_iterator()))\n",
        "images, labels = batch['image'][0][:9], batch['label'][0][:9]\n",
        "titles = map(make_label_getter(dataset), labels.argmax(axis=1))\n",
        "show_img_grid(images, titles)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ehzbRTSN20E5"
      },
      "source": [
        "### Load pre-trained"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "8kIeQWuyf4Ga",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e2f7f856-e79c-40a2-8871-84ce6ab5b4f6"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "classifier: token\n",
              "hidden_size: 768\n",
              "model_name: ViT-B_32\n",
              "patches:\n",
              "  size: !!python/tuple\n",
              "  - 32\n",
              "  - 32\n",
              "representation_size: null\n",
              "transformer:\n",
              "  attention_dropout_rate: 0.0\n",
              "  dropout_rate: 0.0\n",
              "  mlp_dim: 3072\n",
              "  num_heads: 12\n",
              "  num_layers: 12"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ],
      "source": [
        "model_config = models_config.MODEL_CONFIGS[model_name]\n",
        "model_config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "DMKr-4nK3DlT",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Load model definition & initialize random parameters.\n",
        "# This also compiles the model to XLA (takes some minutes the first time).\n",
        "if model_name.startswith('Mixer'):\n",
        "  model = models.MlpMixer(num_classes=num_classes, **model_config)\n",
        "else:\n",
        "  model = models.VisionTransformer(num_classes=num_classes, **model_config)\n",
        "variables = jax.jit(lambda: model.init(\n",
        "    jax.random.PRNGKey(0),\n",
        "    # Discard the \"num_local_devices\" dimension of the batch for initialization.\n",
        "    batch['image'][0, :1],\n",
        "    train=False,\n",
        "), backend='cpu')()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zIXjOEDkvAWM",
        "outputId": "5564f304-bc07-4550-83a2-aeae39bd18f6",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Inspect extra keys:\n",
            "{'pre_logits/kernel', 'pre_logits/bias'}\n",
            "INFO:absl:load_pretrained: drop-head variant\n"
          ]
        }
      ],
      "source": [
        "# Load and convert pretrained checkpoint.\n",
        "# This involves loading the actual pre-trained model results, but then also also\n",
        "# modifying the parameters a bit, e.g. changing the final layers, and resizing\n",
        "# the positional embeddings.\n",
        "# For details, refer to the code and to the methods of the paper.\n",
        "params = checkpoint.load_pretrained(\n",
        "    pretrained_path=f'{model_name}.npz',\n",
        "    init_params=variables['params'],\n",
        "    model_config=model_config,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aQVKzhaR8o-J"
      },
      "source": [
        "### Evaluate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "WB6ywRTY-LOa",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e6f101ee-14fa-44b1-8fd6-6f6351526568"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "params.cls: ArrayImpl (10,)\n",
            "params_repl.cls: ArrayImpl (1, 10)\n"
          ]
        }
      ],
      "source": [
        "# So far, all our data is in the host memory. Let's now replicate the arrays\n",
        "# into the devices.\n",
        "# This will make every array in the pytree params become a ShardedDeviceArray\n",
        "# that has the same data replicated across all local devices.\n",
        "# For TPU it replicates the params in every core.\n",
        "# For a single GPU this simply moves the data onto the device.\n",
        "# For CPU it simply creates a copy.\n",
        "params_repl = flax.jax_utils.replicate(params)\n",
        "print('params.cls:', type(params['head']['bias']).__name__,\n",
        "      params['head']['bias'].shape)\n",
        "print('params_repl.cls:', type(params_repl['head']['bias']).__name__,\n",
        "      params_repl['head']['bias'].shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "_unNxEZAK0Cu",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Then map the call to our model's forward pass onto all available devices.\n",
        "vit_apply_repl = jax.pmap(lambda params, inputs: model.apply(\n",
        "    dict(params=params), inputs, train=False))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "ZgjFBUQ88p4z",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "def get_accuracy(params_repl):\n",
        "  \"\"\"Returns accuracy evaluated on the test set.\"\"\"\n",
        "  good = total = 0\n",
        "  steps = input_pipeline.get_dataset_info(dataset, 'test')['num_examples'] // batch_size\n",
        "  for _, batch in zip(tqdm.trange(steps), ds_test.as_numpy_iterator()):\n",
        "    predicted = vit_apply_repl(params_repl, batch['image'])\n",
        "    is_same = predicted.argmax(axis=-1) == batch['label'].argmax(axis=-1)\n",
        "    good += is_same.sum()\n",
        "    total += len(is_same.flatten())\n",
        "  return good / total"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "3qc7j0lv-F6-",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ea5dab65-4082-4abb-c662-2f9b88dbf82d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2\n",
            "100%|██████████| 19/19 [01:05<00:00,  3.44s/it]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(0.10063734, dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ],
      "source": [
        "# Random performance without fine-tuning.\n",
        "get_accuracy(params_repl)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HxMdU_e5NeoT"
      },
      "source": [
        "### Fine-tune"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "MI62dexw8mGo",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# 100 Steps take approximately 15 minutes in the TPU runtime.\n",
        "total_steps = 100\n",
        "warmup_steps = 5\n",
        "decay_type = 'cosine'\n",
        "grad_norm_clip = 1\n",
        "# This controls in how many forward passes the batch is split. 8 works well with\n",
        "# a TPU runtime that has 8 devices. 64 should work on a GPU. You can of course\n",
        "# also adjust the batch_size above, but that would require you to adjust the\n",
        "# learning rate accordingly.\n",
        "accum_steps = 8\n",
        "base_lr = 0.03"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "vzlfREb1ZHBY",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Check out train.make_update_fn in the editor on the right side for details.\n",
        "lr_fn = utils.create_learning_rate_schedule(total_steps, base_lr, decay_type, warmup_steps)\n",
        "# We use a momentum optimizer that uses half precision for state to save\n",
        "# memory. It als implements the gradient clipping.\n",
        "tx = optax.chain(\n",
        "    optax.clip_by_global_norm(grad_norm_clip),\n",
        "    optax.sgd(\n",
        "        learning_rate=lr_fn,\n",
        "        momentum=0.9,\n",
        "        accumulator_dtype='bfloat16',\n",
        "    ),\n",
        ")\n",
        "update_fn_repl = train.make_update_fn(\n",
        "    apply_fn=model.apply, accum_steps=accum_steps, tx=tx)\n",
        "opt_state = tx.init(params)\n",
        "opt_state_repl = flax.jax_utils.replicate(opt_state)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "RTU7OmgjHb-G",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Initialize PRNGs for dropout.\n",
        "update_rng_repl = flax.jax_utils.replicate(jax.random.PRNGKey(0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "zKn4IfUWHWPk",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 878
        },
        "outputId": "ec5c7f51-6182-470e-df7f-32b70f9fee26"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 100/100 [09:22<00:00,  5.62s/it]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[<matplotlib.lines.Line2D at 0x7e412274c850>]"
            ]
          },
          "metadata": {},
          "execution_count": 27
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 29540 characters]\n"
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "[encoded data, 28796 characters]\n"
          },
          "metadata": {}
        }
      ],
      "source": [
        "losses = []\n",
        "lrs = []\n",
        "# Completes in ~20 min on the TPU runtime.\n",
        "for step, batch in zip(\n",
        "    tqdm.trange(1, total_steps + 1),\n",
        "    ds_train.as_numpy_iterator(),\n",
        "):\n",
        "\n",
        "  params_repl, opt_state_repl, loss_repl, update_rng_repl = update_fn_repl(\n",
        "      params_repl, opt_state_repl, batch, update_rng_repl)\n",
        "  losses.append(loss_repl[0])\n",
        "  lrs.append(lr_fn(step))\n",
        "\n",
        "plt.plot(losses)\n",
        "plt.figure()\n",
        "plt.plot(lrs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "jJhKAMhMI2D6",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d6a9b68f-dde2-411f-bfa0-7722592dcb45"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2\n",
            "100%|██████████| 19/19 [00:30<00:00,  1.59s/it]\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array(0.9771793, dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ],
      "source": [
        "# Should be ~96.7% for Mixer-B/16 or 97.7% for ViT-B/32 on CIFAR10 (both @224)\n",
        "get_accuracy(params_repl)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-wIdj_qnbIM"
      },
      "source": [
        "### Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "id": "BrOSyiBsE-N-",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ea99a96a-b8d2-4edd-c823-29c603d909db"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Copying gs://vit_models/imagenet21k+imagenet2012/ViT-B_32.npz...\n",
            "/ [0 files][    0.0 B/336.9 MiB]                                                \r==> NOTE: You are downloading one or more large file(s), which would\n",
            "run significantly faster if you enabled sliced object downloads. This\n",
            "feature is enabled by default but requires that compiled crcmod be\n",
            "installed (see \"gsutil help crcmod\").\n",
            "\n",
            "/ [1 files][336.9 MiB/336.9 MiB]   35.3 MiB/s                                   \n",
            "Operation completed over 1 objects/336.9 MiB.                                    \n"
          ]
        }
      ],
      "source": [
        "# Download a pre-trained model.\n",
        "\n",
        "if model_name.startswith('Mixer'):\n",
        "  # Download model trained on imagenet2012\n",
        "  ![ -e \"$model_name\"_imagenet2012.npz ] || gsutil cp gs://mixer_models/imagenet1k/\"$model_name\".npz \"$model_name\"_imagenet2012.npz\n",
        "  model = models.MlpMixer(num_classes=1000, **model_config)\n",
        "else:\n",
        "  # Download model pre-trained on imagenet21k and fine-tuned on imagenet2012.\n",
        "  ![ -e \"$model_name\"_imagenet2012.npz ] || gsutil cp gs://vit_models/imagenet21k+imagenet2012/\"$model_name\".npz \"$model_name\"_imagenet2012.npz\n",
        "  model = models.VisionTransformer(num_classes=1000, **model_config)\n",
        "\n",
        "import os\n",
        "assert os.path.exists(f'{model_name}_imagenet2012.npz')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "id": "TE7BoGkyoY7X",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Load and convert pretrained checkpoint.\n",
        "params = checkpoint.load(f'{model_name}_imagenet2012.npz')\n",
        "params['pre_logits'] = {}  # Need to restore empty leaf for Flax."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "id": "p3TTypN0gkwO",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "af5863fa-2b6c-43c4-d547-56c8a4ff4641"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2024-02-06 10:52:41--  https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt\n",
            "Resolving storage.googleapis.com (storage.googleapis.com)... 172.217.203.207, 172.253.123.207, 142.250.98.207, ...\n",
            "Connecting to storage.googleapis.com (storage.googleapis.com)|172.217.203.207|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 21675 (21K) [text/plain]\n",
            "Saving to: ‘ilsvrc2012_wordnet_lemmas.txt’\n",
            "\n",
            "\r          ilsvrc201   0%[                    ]       0  --.-KB/s               \rilsvrc2012_wordnet_ 100%[===================>]  21.17K  --.-KB/s    in 0s      \n",
            "\n",
            "2024-02-06 10:52:41 (59.2 MB/s) - ‘ilsvrc2012_wordnet_lemmas.txt’ saved [21675/21675]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Get imagenet labels.\n",
        "!wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt\n",
        "imagenet_labels = dict(enumerate(open('ilsvrc2012_wordnet_lemmas.txt')))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "id": "IQwomPOopDlr",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 679
        },
        "outputId": "d10ce619-290d-4df3-d2bb-449bbe56f281"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2024-02-06 10:52:41--  https://picsum.photos/384\n",
            "Resolving picsum.photos (picsum.photos)... 104.26.5.30, 172.67.74.163, 104.26.4.30, ...\n",
            "Connecting to picsum.photos (picsum.photos)|104.26.5.30|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://fastly.picsum.photos/id/997/384/384.jpg?hmac=VWE41G8Zk15BPtMQxAa0HPtjkgN80R-eoJVXHwwyE4k [following]\n",
            "--2024-02-06 10:52:42--  https://fastly.picsum.photos/id/997/384/384.jpg?hmac=VWE41G8Zk15BPtMQxAa0HPtjkgN80R-eoJVXHwwyE4k\n",
            "Resolving fastly.picsum.photos (fastly.picsum.photos)... 151.101.1.91, 151.101.65.91, 151.101.129.91, ...\n",
            "Connecting to fastly.picsum.photos (fastly.picsum.photos)|151.101.1.91|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 21922 (21K) [image/jpeg]\n",
            "Saving to: ‘picsum.jpg’\n",
            "\n",
            "picsum.jpg          100%[===================>]  21.41K  --.-KB/s    in 0.001s  \n",
            "\n",
            "2024-02-06 10:52:42 (14.1 MB/s) - ‘picsum.jpg’ saved [21922/21922]\n",
            "\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=384x384>"
            ],
            "image/png": "[encoded data, 264208 characters]\n"
          },
          "metadata": {},
          "execution_count": 32
        }
      ],
      "source": [
        "# Get a random picture with the correct dimensions.\n",
        "resolution = 224 if model_name.startswith('Mixer') else 384\n",
        "!wget https://picsum.photos/$resolution -O picsum.jpg\n",
        "import PIL\n",
        "img = PIL.Image.open('picsum.jpg')\n",
        "img"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "id": "wrvwNAGJshzb",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# Predict on a batch with a single item (note very efficient TPU usage...)\n",
        "logits, = model.apply(dict(params=params), (np.array(img) / 128 - 1)[None, ...], train=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "id": "64hwCdaehs42",
        "vscode": {
          "languageId": "python"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "94bbda94-db37-40be-dce9-0d314319e310"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.52548 : lakeside, lakeshore\n",
            "0.16299 : promontory, headland, head, foreland\n",
            "0.12497 : volcano\n",
            "0.04307 : valley, vale\n",
            "0.02912 : cliff, drop, drop-off\n",
            "0.02730 : seashore, coast, seacoast, sea-coast\n",
            "0.01116 : dam, dike, dyke\n",
            "0.00892 : geyser\n",
            "0.00597 : sandbar, sand_bar\n",
            "0.00570 : alp\n"
          ]
        }
      ],
      "source": [
        "preds = np.array(jax.nn.softmax(logits))\n",
        "for idx in preds.argsort()[:-11:-1]:\n",
        "  print(f'{preds[idx]:.5f} : {imagenet_labels[idx]}', end='')"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "gpuClass": "standard",
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "88b7391eac6346b6848f9a0b42983bb3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_ff4ce5fb8b4d42ff9567995085be3fe0",
              "IPY_MODEL_60530198d0044832b19d2a465455a810",
              "IPY_MODEL_9b9d565e72c34a4890404eec2ac7d377"
            ],
            "layout": "IPY_MODEL_2a23cdf1a8a847f7a2ce8f22b5b77bb1"
          }
        },
        "ff4ce5fb8b4d42ff9567995085be3fe0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_9a15e456d4f54167afe81f560918064d",
            "placeholder": "​",
            "style": "IPY_MODEL_afc6d112cf5646788bdf31894a232222",
            "value": "Dl Completed...: 100%"
          }
        },
        "60530198d0044832b19d2a465455a810": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1633effb0fe24526a419fc2411804816",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d140d24e2cc64a2fab959d8bbd57d7bc",
            "value": 1
          }
        },
        "9b9d565e72c34a4890404eec2ac7d377": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_beb0f1d2984b4c91a17ad9ff3cfd7dc0",
            "placeholder": "​",
            "style": "IPY_MODEL_2c679fdde70f4e1394f89f7d9edd62ba",
            "value": " 1/1 [00:05&lt;00:00,  3.86s/ url]"
          }
        },
        "2a23cdf1a8a847f7a2ce8f22b5b77bb1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9a15e456d4f54167afe81f560918064d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "afc6d112cf5646788bdf31894a232222": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "1633effb0fe24526a419fc2411804816": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": "20px"
          }
        },
        "d140d24e2cc64a2fab959d8bbd57d7bc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "beb0f1d2984b4c91a17ad9ff3cfd7dc0": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2c679fdde70f4e1394f89f7d9edd62ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "a2006ee171e4402fbb2f28e3ab4cafa5": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_8d83fafc13f641fdb5db20bc7a74c906",
              "IPY_MODEL_54fd4ba8000e4a6cafacbdac591e4542",
              "IPY_MODEL_7c70d6032fa0415a87917de014492c5e"
            ],
            "layout": "IPY_MODEL_1b90aa4ddccd44249938f46bf1c57d24"
          }
        },
        "8d83fafc13f641fdb5db20bc7a74c906": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_898f31eb6de6488186fafd9999c608eb",
            "placeholder": "​",
            "style": "IPY_MODEL_ea6d78f0e27a409cac736eea6f0bebd3",
            "value": "Dl Size...: 100%"
          }
        },
        "54fd4ba8000e4a6cafacbdac591e4542": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4a1c284e25914060b0b439dc74d4798e",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_33d74a0d0c554135926b6eca6e32e8c4",
            "value": 1
          }
        },
        "7c70d6032fa0415a87917de014492c5e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8bbe1a3515bd45eaa246328f496427c8",
            "placeholder": "​",
            "style": "IPY_MODEL_9b66f4cbff664452a4d03cb0716518a8",
            "value": " 162/162 [00:05&lt;00:00, 44.27 MiB/s]"
          }
        },
        "1b90aa4ddccd44249938f46bf1c57d24": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "898f31eb6de6488186fafd9999c608eb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ea6d78f0e27a409cac736eea6f0bebd3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "4a1c284e25914060b0b439dc74d4798e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": "20px"
          }
        },
        "33d74a0d0c554135926b6eca6e32e8c4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "8bbe1a3515bd45eaa246328f496427c8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9b66f4cbff664452a4d03cb0716518a8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "95ac0ffea8274bf0b9eb5fd5214c7f69": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_740c230a0bdb49d8a60f0e4a5c0ab28e",
              "IPY_MODEL_170456fc38604ec398fba752c90cd417",
              "IPY_MODEL_521b5953067f4cc1bb10003591f997b9"
            ],
            "layout": "IPY_MODEL_ca515be22fdf4098b695334ee7ffe7b2"
          }
        },
        "740c230a0bdb49d8a60f0e4a5c0ab28e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5bff4f12df1549ba94a1d1e5e169c80c",
            "placeholder": "​",
            "style": "IPY_MODEL_adfc4aaac7c045059d12c1e1cf37f4c2",
            "value": "Extraction completed...: 100%"
          }
        },
        "170456fc38604ec398fba752c90cd417": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5fcecb05c8c24e9285aff1536b185bea",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_3c662c7be9964edfa0628cf8effd3d8f",
            "value": 1
          }
        },
        "521b5953067f4cc1bb10003591f997b9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1b1f1eb77092490ca6c5bbac60627ad6",
            "placeholder": "​",
            "style": "IPY_MODEL_e3f6b487be0b44f180113151c6392c92",
            "value": " 8/8 [00:05&lt;00:00,  5.56s/ file]"
          }
        },
        "ca515be22fdf4098b695334ee7ffe7b2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5bff4f12df1549ba94a1d1e5e169c80c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "adfc4aaac7c045059d12c1e1cf37f4c2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "5fcecb05c8c24e9285aff1536b185bea": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": "20px"
          }
        },
        "3c662c7be9964edfa0628cf8effd3d8f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "1b1f1eb77092490ca6c5bbac60627ad6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e3f6b487be0b44f180113151c6392c92": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "3eff2841b3f84abea71febf38c1dee8e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_7d6dc9752873434dab3c1bc09bd2d7a4",
              "IPY_MODEL_7e8e8e1d697749b9b07fbb8da3f0c3b7",
              "IPY_MODEL_9923a1e8acfd41cd93879de31055406f"
            ],
            "layout": "IPY_MODEL_77f9a36bd6384955871d3b7f3deddb9d"
          }
        },
        "7d6dc9752873434dab3c1bc09bd2d7a4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b2ea0941ad0646f092d7bd5113dd1725",
            "placeholder": "​",
            "style": "IPY_MODEL_c05389b29b0a40c7af3495f1de9d4897",
            "value": "Generating splits...: 100%"
          }
        },
        "7e8e8e1d697749b9b07fbb8da3f0c3b7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_46da21bfa2934247945d2c3f0f121d2c",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_fc2d97654d744874aecfa15f4bb2c2e3",
            "value": 2
          }
        },
        "9923a1e8acfd41cd93879de31055406f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_45a4c68ca2694c23acf5933739375939",
            "placeholder": "​",
            "style": "IPY_MODEL_c0076d262b684281b18db3c82810fcc2",
            "value": " 2/2 [00:55&lt;00:00, 24.37s/ splits]"
          }
        },
        "77f9a36bd6384955871d3b7f3deddb9d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": "hidden",
            "width": null
          }
        },
        "b2ea0941ad0646f092d7bd5113dd1725": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c05389b29b0a40c7af3495f1de9d4897": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "46da21bfa2934247945d2c3f0f121d2c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fc2d97654d744874aecfa15f4bb2c2e3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "45a4c68ca2694c23acf5933739375939": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c0076d262b684281b18db3c82810fcc2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "9c224504cbcb483091f698e7189d0c90": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_203ec9c61ec94b53bcc86ac8952bb780",
              "IPY_MODEL_eea538f2b1604587a51fee35ce16f5be",
              "IPY_MODEL_fd9cd18bf458434e933ae970f80ef6dc"
            ],
            "layout": "IPY_MODEL_58decf174a574a45a6fb2d7685598410"
          }
        },
        "203ec9c61ec94b53bcc86ac8952bb780": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_07fcf33390e14a7dbb62523c80cc8c91",
            "placeholder": "​",
            "style": "IPY_MODEL_0029395e502e4505b51434c8211221c9",
            "value": "Generating train examples...: 100%"
          }
        },
        "eea538f2b1604587a51fee35ce16f5be": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_52b495f2c2d04372aded6bb4dbf1f5c6",
            "max": 50000,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_26febddaef2a4db795f2b3319cd5a8ef",
            "value": 50000
          }
        },
        "fd9cd18bf458434e933ae970f80ef6dc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_928ef89206014b9282a9c8c36ae04fd8",
            "placeholder": "​",
            "style": "IPY_MODEL_64f81fce9df04a098557fa8d68fbf8fe",
            "value": " 49794/50000 [00:45&lt;00:00, 1104.41 examples/s]"
          }
        },
        "58decf174a574a45a6fb2d7685598410": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": "hidden",
            "width": null
          }
        },
        "07fcf33390e14a7dbb62523c80cc8c91": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0029395e502e4505b51434c8211221c9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "52b495f2c2d04372aded6bb4dbf1f5c6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "26febddaef2a4db795f2b3319cd5a8ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "928ef89206014b9282a9c8c36ae04fd8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "64f81fce9df04a098557fa8d68fbf8fe": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "858d201e6ca6454c9c60c6e3cc0daa1b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_cc9699e15bfc4573986c18cf5005b36e",
              "IPY_MODEL_50ccc5e2256a48a38a2f1996956d5f73",
              "IPY_MODEL_83214684d4af410f82a3830ffb8d3b48"
            ],
            "layout": "IPY_MODEL_63f0ce0b2bf546458aaa35efdb3a2f07"
          }
        },
        "cc9699e15bfc4573986c18cf5005b36e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ce720edde76b4613828a4bf4b97d4106",
            "placeholder": "​",
            "style": "IPY_MODEL_cc1e72ca287243fe8196c332dbb27423",
            "value": "Shuffling /root/tensorflow_datasets/cifar10/3.0.2.incompleteHVPR1P/cifar10-train.tfrecord*...:  98%"
          }
        },
        "50ccc5e2256a48a38a2f1996956d5f73": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3ef22f6fa51947bc90c1e6e3b2b4a8ef",
            "max": 50000,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d9f1aee8f1fb44b88b6d83d5788d1ea0",
            "value": 50000
          }
        },
        "83214684d4af410f82a3830ffb8d3b48": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d39bca2d23be4adab61bd277d6c80e52",
            "placeholder": "​",
            "style": "IPY_MODEL_c1c497d8574c493b8d90366be94f8e1b",
            "value": " 49183/50000 [00:00&lt;00:00, 109403.50 examples/s]"
          }
        },
        "63f0ce0b2bf546458aaa35efdb3a2f07": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": "hidden",
            "width": null
          }
        },
        "ce720edde76b4613828a4bf4b97d4106": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "cc1e72ca287243fe8196c332dbb27423": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "3ef22f6fa51947bc90c1e6e3b2b4a8ef": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d9f1aee8f1fb44b88b6d83d5788d1ea0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "d39bca2d23be4adab61bd277d6c80e52": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c1c497d8574c493b8d90366be94f8e1b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "3f61424d60d9474cb7f6cdaa16e757d2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e0d5e0d2c66d41f7b2b392b12476076d",
              "IPY_MODEL_d8105d46b03c4063ad00fc7ceedd6ce4",
              "IPY_MODEL_69e4d51c37994ee39143f10beca47200"
            ],
            "layout": "IPY_MODEL_7069223e26324cbcbed7b44258ee73c3"
          }
        },
        "e0d5e0d2c66d41f7b2b392b12476076d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_cc8544485c584c43ae0825a7b96510ae",
            "placeholder": "​",
            "style": "IPY_MODEL_f22bce220a7049cfa81c561835c64a6e",
            "value": "Generating test examples...:  98%"
          }
        },
        "d8105d46b03c4063ad00fc7ceedd6ce4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6c90f244d5c34a56b53ac492ba92988c",
            "max": 10000,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_98c875a0af064305b540b33d341f85e9",
            "value": 10000
          }
        },
        "69e4d51c37994ee39143f10beca47200": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1d27837081434b88aa2526d2c7bbfe53",
            "placeholder": "​",
            "style": "IPY_MODEL_530d6ca4cc82481d8c1376bbed324d04",
            "value": " 9847/10000 [00:09&lt;00:00, 1120.42 examples/s]"
          }
        },
        "7069223e26324cbcbed7b44258ee73c3": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": "hidden",
            "width": null
          }
        },
        "cc8544485c584c43ae0825a7b96510ae": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f22bce220a7049cfa81c561835c64a6e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "6c90f244d5c34a56b53ac492ba92988c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "98c875a0af064305b540b33d341f85e9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "1d27837081434b88aa2526d2c7bbfe53": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "530d6ca4cc82481d8c1376bbed324d04": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "cd8b4d532cf94190ae89032bee8d1d64": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_60149eabd64d4d6ebf1f5659478e8b24",
              "IPY_MODEL_674cce53dac842fb8d8f30d58270f91c",
              "IPY_MODEL_d226d17a362246d28eacee5d3084c377"
            ],
            "layout": "IPY_MODEL_dc237fc9345f400394c01e1ac11dc048"
          }
        },
        "60149eabd64d4d6ebf1f5659478e8b24": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_31b648ce84b04155913adff950edf4e2",
            "placeholder": "​",
            "style": "IPY_MODEL_8e986d355825423ea01808765ef7385a",
            "value": "Shuffling /root/tensorflow_datasets/cifar10/3.0.2.incompleteHVPR1P/cifar10-test.tfrecord*...:   0%"
          }
        },
        "674cce53dac842fb8d8f30d58270f91c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1af824e5c2c548b69ace8d5bbc207c16",
            "max": 10000,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_410e720c1cc945ae8f56362bdb96879b",
            "value": 10000
          }
        },
        "d226d17a362246d28eacee5d3084c377": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_082e18037c274d219da99bd58a4c93f1",
            "placeholder": "​",
            "style": "IPY_MODEL_2570ced1bb0646d08f323d54b22bd7f8",
            "value": " 0/10000 [00:00&lt;?, ? examples/s]"
          }
        },
        "dc237fc9345f400394c01e1ac11dc048": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": "hidden",
            "width": null
          }
        },
        "31b648ce84b04155913adff950edf4e2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8e986d355825423ea01808765ef7385a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "1af824e5c2c548b69ace8d5bbc207c16": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "410e720c1cc945ae8f56362bdb96879b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "082e18037c274d219da99bd58a4c93f1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2570ced1bb0646d08f323d54b22bd7f8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}