homeworks/hw1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5rUp2xWkdGFu"
   },
   "source": [
    "# Homework 1: Autoregressive models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B0w70vkwZ-tw"
   },
   "source": [
    "## Task 1: Theory (5pt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ivV4DY1SPEIL"
   },
   "source": [
    "### Problem 1: f-divergence (1.5pt)\n",
    "\n",
    "In the Lecture 1 we have discussed that it is possible to formulate the generative modelling task as the divergence minimization task.\n",
    "\n",
    "\n",
    "$$\n",
    "\t\\min_{\\boldsymbol{\\theta}} D(\\pi || p),\n",
    "$$\n",
    "\n",
    "Let consider the family of f-divergencies:\n",
    "$$\n",
    "\tD_f(\\pi || p) = \\mathbb{E}_{p(\\mathbf{x})}  f\\left( \\frac{\\pi(\\mathbf{x})}{p(\\mathbf{x})} \\right)  = \\int p(\\mathbf{x}) f\\left( \\frac{\\pi(\\mathbf{x})}{p(\\mathbf{x})} \\right) d \\mathbf{x},\n",
    "$$\n",
    "where $f: \\mathbb{R}_+ \\rightarrow \\mathbb{R}$ is a convex, lower semicontinuous function satisfying $f(1) = 0$.\n",
    "\n",
    "It is easy to show that any function $D_f$ satisties the divergence definition.\n",
    "\n",
    "Your task here is to show that forward and reverse KL lie in the class of f-divergence. \n",
    "Find the functions $f$ that correspond to forward and KL divergences (they should satisfy the conditions on function $f$ described above)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x1C_RPljZ-t3"
   },
   "source": [
    "```\n",
    "your solution for problem 1\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c7wm4AOrZ-t4"
   },
   "source": [
    "### Problem 2: Curse of dimensionality (1.5pt)\n",
    "\n",
    "The main problem of the generative modelling is the curse of dimensionality. Let try to get some intuition about it.\n",
    "\n",
    "Let consider a sphere of radius $r = 1$ in $\\mathbb{R}^m$. Let say that it has the volume $V_m(r)$. Our goal is to find the fraction of the volume of the sphere that lies between radius $r = 1 - \\epsilon$ and $r = 1$. Our geometric intuition is that this fraction is small. But the magic happens with $m$ goes to infinity. Basically, the volume of a high dimensional sphere is mostly in this small fraction!\n",
    "\n",
    "1. Find the expression of the volume of a shpere of radius $r$ in $m$ dimensions.\n",
    "\n",
    "2. Find the required fraction $\\Delta = \\frac{V_m(1) - V_m(1 - \\epsilon)}{V_m(1)}$.\n",
    "\n",
    "3. Prove that, for large $m$, the fraction tends to 1 even for small values of $\\epsilon$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sfyWlBzOZ-t4"
   },
   "source": [
    "```\n",
    "your solution for problem 2\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Problem 3: Expressivity of normalizing flows (2pt)\n",
    "\n",
    "This problem considers the following question. Are normalizing flows able to represent **any** distribution $\\pi(\\mathbf{x})$, even if the base distribution $p(\\mathbf{u})$ is restricted to be simple?\n",
    "\n",
    "Let $\\pi(\\mathbf{x})$ is an absolutely continuous probability distribution supported everywhere on $\\mathbb{R}^m$ (i.e. $\\pi(\\mathbf{x}) > 0$ for all $\\mathbf{x} \\in \\mathbb{R}^m$). Additionally we suppose, that the pdf $\\pi(\\mathbf{x})$ is continuously differentiable on  $\\mathbb{R}^m$.  Our ultimate goal is to show, that there exists diffeomorphism (invertible continuously differentiable map):\n",
    "$$\n",
    "    \\mathbf{u} = F(\\mathbf{x}) \\quad F : \\mathbb{R}^m \\rightarrow \\mathbb{U}.\n",
    "$$\n",
    "Here $\\mathbb{U} = [0, 1]^m$ is a hypercube, which turns $\\pi(\\mathbf{x})$ into uniform distribution $p(\\mathbf{u}) = U[0, 1]^m$ on the hypercube $\\mathbb{U}$ ($p(\\mathbf{u}) = 1 \\,,\\, \\mathbf{u} \\in \\mathbb{U}$). Here we have to think about openness and closeness of $\\mathbb{U}$ for formal math correctness, but we omit it in this task. \n",
    "\n",
    "<details> \n",
    "  <summary> <i>Math correctness comment</i> </summary>\n",
    "  \n",
    "   Strictly speaking, the diffeomorphism $F$ maps $\\mathbb{R}^m$ to the open cube $\\text{int}\\left(\\mathbb{U}\\right) = (0, 1)^m$, sinse the reverse mapping $F^{-1} : \\text{int}\\left(\\mathbb{U}\\right) \\rightarrow \\mathbb{R}^m$ is continuous $\\Rightarrow$ the preimage of the open set $\\mathbb{R}^m$ with respect to $F^{-1}$ should be <b>open</b>: $(F^{-1})^{-1}(\\mathbb{R}^m) = F(\\mathbb{R}^m) = \\text{an open set}$\n",
    "   \n",
    "</details>\n",
    "\n",
    "So, if such function exists, it means that **there exists normalizing flow model from base uniform distribution to any(!) target distribution.**\n",
    "\n",
    "---------\n",
    "\n",
    "1. Consider the autoregressive decomposition of $\\pi(\\mathbf{x})$:\n",
    "\n",
    "$$\n",
    "    \\pi(\\mathbf{x}) = \\prod\\limits_{j = 1}^{m} \\pi(x_j|\\mathbf{x}_{1:j - 1}).\n",
    "$$\n",
    "\n",
    "2. Treat each component $\\pi(x_j|\\mathbf{x}_{1:j - 1})$ in the decomposition above separately. Consider the transformations:\n",
    "\n",
    "$$\n",
    "    \\mathbf{x} \\rightarrow u_j = F_j(x_j, \\mathbf{x}_{1:j - 1}) = \\int\\limits_{- \\infty}^{x_j} \\pi(x_j'|\\mathbf{x}_{1:j - 1}) d x_j'.\n",
    "$$\n",
    "\n",
    "   Here $F_j(x_j, \\mathbf{x}_{1:j - 1})$ is the cumulative distribution function of $j$-th conditional $x_j$ (given $\\mathbf{x}_{1:j - 1}$) Note, that $u_j \\in [0, 1]$.\n",
    "\n",
    "3. Define the transform $F : \\mathbb{R}^m \\rightarrow \\mathbb{U}$ as follows:\n",
    "    \n",
    "$$\n",
    "    \\mathbf{x} \\rightarrow \\mathbf{u} = \\begin{bmatrix} F_1(x_1) \\\\ F_2(x_2,x_1) \\\\ \\dots \\\\ F_m(x_m, \\mathbf{x}_{1:m-1})\\end{bmatrix}.\n",
    "$$\n",
    "    \n",
    "Given the properties of $\\pi$ ($\\pi(\\mathbf{x}) > 0$ and continuously differentiable) it is easy to show, that $F$ is continuously differentiable. We omit the details since they are just boring mathematical calculations.\n",
    "\n",
    "1) Prove, that $\\det \\mathbf{J}_F (\\mathbf{x}) = \\pi(x)$ (it follows, that the function $F : \\mathbb{R}^m \\rightarrow \\mathbb{U}$ is invertible).\n",
    "\n",
    "2) Prove, that $\\mathbf{u}$ is a uniformly distributed ($p(\\mathbf{u}) = U[0,1]^m)$.\n",
    "\n",
    "3) Let $\\pi(\\mathbf{x})$ and $\\mu(\\mathbf{y})$ are absolutely continuous probability distributions supported everywhere on $\\mathbb{R}^m$, whose pdfs are continuously differentiable. Prove that there exists a diffeomorphism $G : \\mathbb{R}^m \\rightarrow \\mathbb{R}^m$ which turns $\\pi$ into $\\mu$, i.e.:\n",
    "\n",
    "$$\\mathbf{y} = G(\\mathbf{x}), \\quad \\text{where } \\mathbf{x} \\sim \\pi(\\mathbf{x}), \\, \\mathbf{y} \\sim \\mu(\\mathbf{y}).$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "your solution for problem 3\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d2LH7vKPPVUt"
   },
   "source": [
    "Now it time to move on to practical part of homework.\n",
    "\n",
    "In our course we will use a small util [package](https://github.com/r-isachenko/dgm_utils) with some usefull functions for loading and visualizing the images and training curves. In each homework there will be a cell with installing this package. Please read carefully the sources of the functions from this package. It could help you to solve the tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7wXq7VJ_SRAl",
    "outputId": "058adcba-b263-4460-e419-7befd58c286d",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "COMMIT_HASH = \"9c3ef18c821268297f31d45e36c3b0f2bb538a80\"\n",
    "!if [ -d dgm_utils ]; then rm -Rf dgm_utils; fi\n",
    "!git clone https://github.com/r-isachenko/dgm_utils.git\n",
    "%cd dgm_utils\n",
    "!git checkout {COMMIT_HASH}\n",
    "!pip install ./\n",
    "%cd ./..\n",
    "!rm -Rf dgm_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m-44Dxk6SRAp"
   },
   "outputs": [],
   "source": [
    "from dgm_utils import train_model\n",
    "from dgm_utils import show_samples, visualize_images, load_dataset\n",
    "from dgm_utils import BaseModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zk6rWePvdGFv"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from typing import Optional, Literal\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.utils.data as data\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    DEVICE = \"cuda\"\n",
    "    print('GPU found :)') \n",
    "else: \n",
    "    DEVICE = \"cpu\"\n",
    "    print('GPU not found :(')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 523
    },
    "id": "PiOBcSvx_PN8",
    "outputId": "186f6cf9-47ee-475a-ff31-985ab44908b8"
   },
   "outputs": [],
   "source": [
    "train_data, test_data = load_dataset(\"mnist\", flatten=False, binarize=True)\n",
    "visualize_images(train_data, \"MNIST samples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KHouHarf_Hs-"
   },
   "source": [
    "## Task 2: PixelCNN receptive field and autocompletion on MNIST (5pt)\n",
    "\n",
    "In this task, you will **analyze PixelCNN**, focusing on its sampling capabilities and receptive field. See the [lecture 1](https://github.com/r-isachenko/2025-DGM-AIMasters-course/blob/main/lectures/lecture1/Lecture1.pdf), [seminar 1](https://github.com/r-isachenko/2025-DGM-AIMasters-course/blob/main/seminars/seminar1/seminar1_Part2_PixelCNN.ipynb) and [paper](https://arxiv.org/abs/1601.06759) for details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "elXdTtB7_6Xl"
   },
   "source": [
    "Masked Convolution Layer is the basic building block of PixelCNN model. Look carefully at this code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ybALnsgO_PRW"
   },
   "outputs": [],
   "source": [
    "class MaskedConv2d(nn.Conv2d):\n",
    "    def __init__(\n",
    "        self, mask_type: Literal['A', 'B'], in_channels: int, out_channels: int, kernel_size: int = 5\n",
    "    ) -> None:\n",
    "        assert mask_type in [\"A\", \"B\"]\n",
    "        super().__init__(\n",
    "            in_channels=in_channels,\n",
    "            out_channels=out_channels,\n",
    "            kernel_size=kernel_size,\n",
    "            padding=kernel_size // 2,\n",
    "        )\n",
    "        # register_buffer() registers the tensor as a non-trainable buffer:\n",
    "        # 1. It stays on the same device as the model (moves with .to(device)).\n",
    "        # 2. It is not updated during backpropagation (not a learnable parameter).\n",
    "        self.register_buffer(\"mask\", torch.zeros_like(self.weight))\n",
    "        self.create_mask(mask_type)\n",
    "\n",
    "    def forward(self, input: torch.Tensor) -> torch.Tensor:\n",
    "        return F.conv2d(input, self.weight * self.mask, self.bias, padding=self.padding)\n",
    "\n",
    "    def create_mask(self, mask_type: Literal['A', 'B']) -> None:\n",
    "        # NOTE: try to understand the logic about mask_type\n",
    "        k = self.kernel_size[0]\n",
    "        self.mask[:, :, : k // 2] = 1\n",
    "        self.mask[:, :, k // 2, : k // 2] = 1\n",
    "        if mask_type == \"B\":\n",
    "            self.mask[:, :, k // 2, k // 2] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lPw0osOrAJFm"
   },
   "source": [
    "[Layer Normalization](https://arxiv.org/abs/1607.06450) stabilizes training. We normalize only over the channel dimension because `nn.LayerNorm([C, H, W])` would normalize each pixel position independently, disrupting the feature map structure essential for PixelCNN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "npUNrimpAGUZ"
   },
   "outputs": [],
   "source": [
    "class LayerNorm(nn.LayerNorm):\n",
    "    def __init__(self, n_filters: int) -> None:\n",
    "        super().__init__(n_filters)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = x.permute(0, 2, 3, 1).contiguous()\n",
    "        x = super().forward(x)\n",
    "        return x.permute(0, 3, 1, 2).contiguous()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "11_PstPkAXig"
   },
   "source": [
    "Now we are ready to define the main PixelCNN model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xC6E-hrHALYf"
   },
   "outputs": [],
   "source": [
    "class PixelCNN(BaseModel):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_shape: tuple[int, int],\n",
    "        n_filters: int = 64,\n",
    "        kernel_size: int = 7,\n",
    "        n_layers: int = 5,\n",
    "        use_layer_norm: bool = True,\n",
    "    ) -> None:\n",
    "\n",
    "        super().__init__()\n",
    "        self.input_shape = input_shape\n",
    "\n",
    "        # we apply the sequence of MaskedConv2d -> LayerNorm (it is optional) -> ReLU\n",
    "        # the last layer should be MaskedConv2d (not ReLU)\n",
    "        # Note 1: the first conv layer should be of type 'A'\n",
    "        # Note 2: final output_dim in MaskedConv2d must be 2\n",
    "        model: list[nn.Module] = [MaskedConv2d(\"A\", 1, n_filters, kernel_size=kernel_size)]\n",
    "\n",
    "        for _ in range(n_layers):\n",
    "            if use_layer_norm:\n",
    "                model.append(LayerNorm(n_filters))\n",
    "            model.append(nn.ReLU())\n",
    "            model.append(MaskedConv2d(\"B\", n_filters, n_filters, kernel_size=kernel_size))\n",
    "\n",
    "        model.append(nn.ReLU())\n",
    "        model.append(MaskedConv2d(\"B\", in_channels=n_filters, out_channels=2, kernel_size=1))\n",
    "        self.net = nn.Sequential(*model)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        batch_size = x.shape[0]\n",
    "        out = (x.float() - 0.5) / 0.5\n",
    "        out = self.net(out)\n",
    "        return out.view(batch_size, 2, 1, *self.input_shape)\n",
    "\n",
    "    def loss(self, x: torch.Tensor) -> dict:\n",
    "        # our loss is just cross entropy\n",
    "        total_loss = F.cross_entropy(self(x), x.long())\n",
    "        return {\"total_loss\": total_loss}\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def sample(self, n: int) -> np.ndarray:\n",
    "        # read carefully the sampling process\n",
    "        # here you see the sequential process of sampling\n",
    "        samples = torch.zeros(n, 1, *self.input_shape).to(self.device)\n",
    "        for r in range(self.input_shape[0]):\n",
    "            for c in range(self.input_shape[1]):\n",
    "                logits = self(samples)[:, :, :, r, c]\n",
    "                probs = F.softmax(logits, dim=1).squeeze(-1)\n",
    "                samples[:, 0, r, c] = torch.multinomial(probs, num_samples=1).squeeze(-1)\n",
    "        return samples.cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68,
     "referenced_widgets": [
      "cbddb1ae1a1d4124b972851b77bc24b9",
      "f519525bb7b640bebd08160d3710008c",
      "4a84d8a794ff4aeb8376f85e45384a89",
      "72ca46babd314fa88a8e59b2aa41f7af",
      "1c55fe8da4c848478801552a1aa3d420",
      "5425bd677e334fa385ff37ddf7eee1d5",
      "7045dbbf8309480bb4a697edb4f4d127",
      "6bb8abfa64ab4e1280cf3efe38b3f7a8",
      "67b3bc3b3359480f863649780ec1cf13",
      "8c11c45785494f2593bf113240f97f46",
      "0773888b28204289b0c9eb77c975ea91"
     ]
    },
    "id": "9FI5foNQBlr5",
    "outputId": "5c19070e-fe92-4c97-f365-d479dd87beb1"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters\n",
    "EPOCHS = \n",
    "BATCH_SIZE = \n",
    "LR = \n",
    "N_LAYERS = \n",
    "N_FILTERS = \n",
    "USE_LAYER_NORM = \n",
    "# ====\n",
    "\n",
    "pixel_cnn = PixelCNN(\n",
    "    input_shape=(28, 28),\n",
    "    n_filters=N_FILTERS,\n",
    "    kernel_size=7,\n",
    "    n_layers=N_LAYERS,\n",
    "    use_layer_norm=USE_LAYER_NORM,\n",
    ")\n",
    "\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = data.DataLoader(test_data, batch_size=BATCH_SIZE)\n",
    "\n",
    "# choose any optimizer/scheduler as you want\n",
    "optimizer = torch.optim.Adam(pixel_cnn.parameters(), lr=LR)\n",
    "\n",
    "train_model(\n",
    "    pixel_cnn,\n",
    "    train_loader,\n",
    "    test_loader,\n",
    "    epochs=EPOCHS,\n",
    "    optimizer=optimizer,\n",
    "    device=DEVICE,\n",
    "    n_samples=16, # Remember the infernce is slow, so, you can lower the value \n",
    "    visualize_samples=True # or you can turn off smapling by setting `visualize_samples=False`\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bdMY__h8CwE5"
   },
   "source": [
    "Now we sample the new images from the model. Notice that the model is sequential and **sampling is really slow** (it is a drawback of all AR models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 523
    },
    "id": "skwg5JlRCt0s",
    "outputId": "f019771c-8247-4ddc-fc61-211117ce3268"
   },
   "outputs": [],
   "source": [
    "samples = pixel_cnn.sample(25)\n",
    "show_samples(samples, title=\"PixelCNN samples\", nrow=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z8eXfFdWzoVu"
   },
   "source": [
    "### Image autocompletion\n",
    "\n",
    "One more feature of autoregressive models is that they can be **easily adapted for image autocompletion**. As autoregressive models predict pixels one by one, we can set the first pixels to predefined values and check how the model completes the image.\n",
    "\n",
    "For implementing this, we just need to skip the iterations in the sampling loop that already have a value unequals to -1.\n",
    "We redefine the sample method in our PixelCNN class to allow it to take the init of the sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2X24nc_eqET-"
   },
   "outputs": [],
   "source": [
    "class PixelCNNAutoComplete(PixelCNN):\n",
    "    @torch.no_grad()\n",
    "    def sample(self, n: int, init: Optional[torch.Tensor] = None) -> np.ndarray:\n",
    "        # ====\n",
    "        # your code\n",
    "        # this method almost the same as the method of the base PixelCNN model\n",
    "        # but now if init is given, this tensor will be used as a starting image.\n",
    "        # NOTE: fill the pixels only with value equals to -1 in the input tensor.\n",
    "        \n",
    "        # ===="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To skip training the model again, we simply load the parameters from an already trained one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "481551c56615463f913d7c88850917c2",
      "439e35ef3240410d9634834725f97411",
      "47c79f4511e5463fbff01e4ee2ba0f03",
      "5d7677ac892b49e68cdd6a366a86e184",
      "62d49bf2ec954f4eaa1ed42d1ea1b329",
      "0aa50015efba4e22b0f7fea2c25c20be",
      "bd5ac497e2654ad9a9763a2e7ec96326",
      "c343c766d3f74e2d936bb6b12e91dabb",
      "23c4bb2b6af14243b554f8d0ace59475",
      "1c99231bb4ee40249e67cc233833faee",
      "224de606324240b89a65ec0a537910f2"
     ]
    },
    "id": "qFdMAilywnI-",
    "outputId": "2f50c251-b653-4420-d880-410e5c9b85e7"
   },
   "outputs": [],
   "source": [
    "pixel_cnn_ac = PixelCNNAutoComplete(\n",
    "    input_shape=(28, 28),\n",
    "    n_filters=N_FILTERS,\n",
    "    kernel_size=7,\n",
    "    n_layers=N_LAYERS,\n",
    "    use_layer_norm=USE_LAYER_NORM,\n",
    ")\n",
    "\n",
    "pixel_cnn_ac.load_state_dict(pixel_cnn.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Rd9rpg2I0paF"
   },
   "source": [
    "We randomly take images from the training set, mask the lower half of the image (set -1's), and let the model autocomplete it. We do this several times for each image to see the diversity of samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "yR0zivWlwsYR",
    "outputId": "62b5bc0c-79fc-4c98-b420-02b6ff0fc42a"
   },
   "outputs": [],
   "source": [
    "def autocomplete_image(\n",
    "    original_image: np.ndarray, \n",
    "    model: PixelCNNAutoComplete, \n",
    "    n_samples: int\n",
    ") -> None:\n",
    "    # Remove lower half of the image\n",
    "    original_image = np.repeat(original_image[None, :, :, :], n_samples, axis=0)\n",
    "    masked_image = original_image.copy()\n",
    "    masked_image[:, :, masked_image.shape[2] // 2 :, :] = -1\n",
    "    \n",
    "    # Generate completions\n",
    "    generated_image = model.sample(n_samples, torch.tensor(masked_image).to(model.device))  \n",
    "    images_to_show = np.concat([np.maximum(masked_image, 0), generated_image, original_image], axis=0)\n",
    "    show_samples(images_to_show, title=\"PixelCNN autocompletion\", nrow=n_samples)\n",
    "\n",
    "for i in range(1, 4):\n",
    "    autocomplete_image(test_data[i], pixel_cnn_ac, n_samples=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6LK1ttmcqfvW"
   },
   "source": [
    "### Receptive field\n",
    "\n",
    "Next, let's visualize the model's receptive field. It **can be empirically measured** by backpropagating an arbitrary loss from a specific output pixel back to the input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "J1ekd7rGDBTi"
   },
   "outputs": [],
   "source": [
    "def plot_receptive_field(model: nn.Module, model_name: str) -> None:\n",
    "    # ====\n",
    "    # your code\n",
    "    # 1) create tensor with zeros and set required_grad to True.\n",
    "    # 2) apply model to the tensor\n",
    "    # 3) apply backward() to the center pixel of model output\n",
    "    # 4) take the gradient with respect to input\n",
    "    # 5) binary receptive field is an indicator map in which we stay 1's if absolute gradient more than 1e-8\n",
    "    # 6) weighted receptive field is the normalized absolute gradient (values lies in [0, 1])\n",
    "\n",
    "    # ====\n",
    "\n",
    "    # we stack the maps to get RGB image\n",
    "    binary_map = np.stack([binary_map, binary_map, binary_map], axis=-1)\n",
    "    weighted_map = np.stack([weighted_map, weighted_map, weighted_map], axis=-1)\n",
    "\n",
    "    # center point will be red\n",
    "    binary_map[x_center, y_center] = [1, 0, 0]\n",
    "    weighted_map[x_center, y_center] = [1, 0, 0]\n",
    "\n",
    "    fig, ax = plt.subplots(1, 2, figsize=(10, 6))\n",
    "    ax[0].imshow(weighted_map, vmin=0.0, vmax=1.0)\n",
    "    ax[1].imshow(binary_map, vmin=0.0, vmax=1.0)\n",
    "\n",
    "    ax[0].set_title(f\"Weighted receptive field for {model_name}\")\n",
    "    ax[1].set_title(f\"Binary receptive field for {model_name}\")\n",
    "\n",
    "    ax[0].axis(\"off\")\n",
    "    ax[1].axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "ruOurWdFDIYP",
    "outputId": "eb22b324-aee8-4c1a-a957-325fc1ee40bd"
   },
   "outputs": [],
   "source": [
    "for n_layers in [1, 3, 5, 6]:\n",
    "    model = PixelCNN(\n",
    "        input_shape=(28, 28),\n",
    "        n_filters=32,\n",
    "        kernel_size=5,\n",
    "        n_layers=n_layers,\n",
    "        use_layer_norm=True\n",
    "    )\n",
    "    model = model.to(DEVICE)\n",
    "    plot_receptive_field(model, model_name=f\"PixelCNN {n_layers} layers\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Njl3VAmruh_v"
   },
   "source": [
    "You can see that as the number of convolutional layers increases, the **receptive field expands**. However, it **remains relatively limited**. Even with a well-trained model, generated samples **lack long-range dependencies**. Moreover, you may notice a strange **blind spot** in the binary receptive field plot on the right side. This is a known issue with the PixelCNN model. Please try to understand why it happens.\n",
    "\n",
    "One way to solve this problem is a [GatedPixelCNN](https://arxiv.org/pdf/1606.05328.pdf) model (see paper, if you are interested in)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DajrWO8YJyYB"
   },
   "source": [
    "## Task 3: ImageGPT on MNIST (5pt)\n",
    "\n",
    "In this task you will try to implement the Image Transformer net for an autoregressive generation MNIST images. See the [blog](https://openai.com/blog/image-gpt/) and [paper](https://cdn.openai.com/papers/Generative_Pretraining_from_Pixels_V2.pdf) for details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Wz6to33lJydk"
   },
   "source": [
    "### Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bl_g0XJDX8Py"
   },
   "source": [
    "Let start with **the multihead attention block**. Each head of the multihead attention computes an embedding\n",
    "$$\n",
    "    \\text{Attention}(\\mathbf{Q}, \\mathbf{K}, \\mathbf{V}) = \\text{softmax}\\left(\\frac{\\mathbf{Q}\\mathbf{K}^T}{\\sqrt{d_k}}\\right) \\mathbf{V},\n",
    "$$\n",
    "where $\\mathbf{Q}, \\mathbf{K}, \\mathbf{V}$ - query/key/value matrices obtained using Linear projection of $\\mathbf{h}^l$.\n",
    "\n",
    "**To make the model autoregressive** we will introduce the following changes:\n",
    "\n",
    "1. We will apply the upper triangular mask to the matrix of attention logits ($\\mathbf{Q}\\mathbf{K}^T$). Masked values are made close to minus infinity so they will turn zero after softmax.\n",
    "2. During training we will add \"start of sequence\" token to the input tensor and pop the last pixel.\n",
    "\n",
    "Note: we will use raster order to identify which pixels come first (as we have done in PixelCNN). For each pixel the predicted probabality is conditioned on all the previous pixels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lv9pX1eyQ09M"
   },
   "outputs": [],
   "source": [
    "class MultiheadAttention(nn.MultiheadAttention):\n",
    "    def __init__(self, embed_dim: int, num_heads: int) -> None:\n",
    "        super().__init__(embed_dim, num_heads)\n",
    "\n",
    "    def get_attention_mask(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        # ====\n",
    "        # your code\n",
    "        # define attention mask, it should contain\n",
    "        # - zeros under and on the main diagonal\n",
    "        # - minus Inf above the main diagonal\n",
    "\n",
    "        # ====\n",
    "        return attention_mask\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        attn_mask = self.get_attention_mask(x)\n",
    "        return super().forward(x, x, x, attn_mask=attn_mask, need_weights=False)[0]\n",
    "\n",
    "\n",
    "def test_attention_mask() -> None:\n",
    "    x = torch.zeros(2, 4, 16)  # (pixel_num, batch_size, emb_dim)\n",
    "    mask = np.array([[0.0, -np.inf], [0.0, 0.0]])\n",
    "    layer = MultiheadAttention(16, 8)\n",
    "    attention_mask = layer.get_attention_mask(x)\n",
    "    assert attention_mask.size() == (x.size(0), x.size(0))\n",
    "    assert np.allclose(attention_mask.numpy(), mask)\n",
    "    out = layer(x)\n",
    "    assert x.size() == out.size()\n",
    "\n",
    "\n",
    "test_attention_mask()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "a3DcBLp3T86-"
   },
   "source": [
    "Now we will define **the decoder block** wich transformes the tensor as follows:\n",
    "\\begin{align}\n",
    "    \\mathbf{n}^l &= \\text{layer\\_norm}_1\\left(\\mathbf{h}^l\\right), \\\\\n",
    "    \\mathbf{a}^l &= \\mathbf{h}^l + \\text{multihead\\_attention}\\left(\\mathbf{n}^l\\right), \\\\\n",
    "    \\mathbf{h}^{l+1} &= \\mathbf{a}^l + \\text{MLP}\\left(\\text{layer\\_norm}_2\\left(a^l\\right)\\right). \\\\\n",
    "\\end{align}\n",
    "\n",
    "The $l^{\\text{th}}$ block receives a tensor $\\mathbf{h}^l$ with shape `(pixel_num, batch_size, emb_dim)` where `pixel_num` is a total number of pixels $(28*28)$ and `emb_dim` is a hyperparameter for the size of the embeddings. We will use just 2 linear layers with ReLU activation for MLP. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yBxhyjd7h3Ww"
   },
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, embed_dim: int, num_heads: int) -> None:\n",
    "        super().__init__()\n",
    "        assert embed_dim % num_heads == 0\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # define multihead attention (https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html)\n",
    "        # define LayerNorm (https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html)\n",
    "        # (here we do not use previous class for LayerNorm because we do not need to change order of tensor dimensions)\n",
    "        # define MLP - 2 linear layers with ReLU\n",
    "        # (You could choose the latent dimensionality of MLP as you like. For example, double the embed_dim)\n",
    "\n",
    "        # ====\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        # ====\n",
    "        # your code\n",
    "        # here you have to implement formulas that described above.\n",
    "\n",
    "        # ====\n",
    "        return x\n",
    "\n",
    "\n",
    "def test_decoder_block() -> None:\n",
    "    block = TransformerBlock(embed_dim=12, num_heads=4)\n",
    "    x = torch.zeros(4, 28, 12)\n",
    "    assert x.shape == block(x).shape\n",
    "\n",
    "\n",
    "test_decoder_block()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's construct the Image Transformer. It consists of $L$ sequentially applied decoder blocks.\n",
    "\n",
    "The initial embedding $h^1$ is the sum of the token embedding of the input batch and a trainable positional embedding, with a \"start of sequence\" token prepended. \n",
    "\n",
    "After passing through the last decoder block, we apply `nn.LayerNorm` followed by a `nn.Linear` layer to produce logits of shape `(pixel_num, batch_size, 1)`:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "  n^L &= \\text{layer\\_norm}\\left(h^L\\right) \\\\\n",
    "  \\text{logits} &= \\text{linear}\\left(n^L\\right)\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2yXwW9fvVJN1"
   },
   "outputs": [],
   "source": [
    "class ImageGPT(BaseModel):\n",
    "    def __init__(\n",
    "        self, input_shape: tuple[int, int], embed_dim: int, num_heads: int, num_layers: int\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.embed_dim = embed_dim\n",
    "        self.input_shape = input_shape\n",
    "        self.criterion = nn.BCEWithLogitsLoss()\n",
    "\n",
    "        # \"start of sequence\" token (we initialize it from Normal distribution)\n",
    "        self.sos = nn.Parameter(torch.zeros(embed_dim))\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) define token_embeddings (we will have 2 embeddings in total, because our images are binary)\n",
    "        # 2) define position_embeddings (we will use nn.Embedding)\n",
    "\n",
    "        # ====\n",
    "\n",
    "        self.layers = nn.ModuleList()\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) add decoder blocks to self.layers list\n",
    "        # 2) define last LayerNorm\n",
    "        # 3) define final Linear layer (without bias)\n",
    "\n",
    "        # ====\n",
    "\n",
    "    def add_sos_token(self, embeddings: torch.Tensor) -> torch.Tensor:\n",
    "        batch_size = embeddings.size(1)\n",
    "        # ====\n",
    "        # your code\n",
    "        # prepend sos (start of sequence) token\n",
    "        # 1) repeat sos token batch_size times (make it of size (1, batch_size, emd_size))\n",
    "        # 2) drop last embedding from embeddings\n",
    "        # 3) concat repeated sos token to embeddings (after dropping)\n",
    "\n",
    "        # ====\n",
    "        return embeddings\n",
    "\n",
    "    def add_pos_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:\n",
    "        length = embeddings.size(0)\n",
    "        # ====\n",
    "        # your code\n",
    "        # add positional embeddings\n",
    "        # 1) define tensor with positions (just torch.arange) of size (length, 1)\n",
    "        # 2) add position embeddings to initial embeddings\n",
    "\n",
    "        # ====\n",
    "        return embeddings\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        x = x.long()\n",
    "        x = x.reshape(x.size(0), -1)  # (batch_size, length)\n",
    "        x = x.permute(1, 0)\n",
    "\n",
    "        embeddings = self.token_embeddings(x)  # (length, batch_size, emb_size)\n",
    "        embeddings = self.add_sos_token(embeddings)\n",
    "        embeddings = self.add_pos_embeddings(embeddings)\n",
    "\n",
    "        # ====\n",
    "        # your code\n",
    "        # 1) apply all decoder layers\n",
    "        # 2) apply final LayerNorm and Linear layer\n",
    "\n",
    "        # ====\n",
    "\n",
    "        return logits.permute(1, 0, 2)  # (length, batch_size, emb_size) -> (batch_size, length, emb_size)\n",
    "\n",
    "    def loss(self, x: torch.Tensor) -> dict:\n",
    "        logits = self(x)\n",
    "        loss = self.criterion(logits.reshape(-1), x.reshape(-1).float())\n",
    "        return {\"total_loss\": loss}\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def sample(self, n_samples: int) -> np.ndarray:\n",
    "        # read sampling carefully\n",
    "        seq_len = self.input_shape[0] * self.input_shape[1]\n",
    "        samples = torch.zeros(n_samples, seq_len).long().to(self.device)\n",
    "        for i in range(seq_len):\n",
    "            logits = self(samples)\n",
    "            dist = torch.distributions.Bernoulli(logits=logits[:, i, 0])\n",
    "            samples[:, i] = dist.sample()\n",
    "        samples = samples.reshape(n_samples, 1, *self.input_shape)\n",
    "        return samples.cpu().numpy()\n",
    "\n",
    "\n",
    "def test_image_gpt() -> None:\n",
    "    image_gpt = ImageGPT(input_shape=(2, 2), embed_dim=12, num_heads=4, num_layers=2)\n",
    "    x = torch.LongTensor([[0, 1, 0, 0], [0, 1, 1, 1]])\n",
    "    assert image_gpt(x).shape == torch.Size([2, 4, 1])\n",
    "    assert image_gpt.loss(x)[\"total_loss\"].requires_grad == True\n",
    "    assert image_gpt.sample(1).shape == torch.Size([1, 1, 2, 2])\n",
    "    img = torch.randint(2, size=(1, 1, 2, 2))\n",
    "\n",
    "\n",
    "test_image_gpt()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NptSgGtRXyca"
   },
   "source": [
    "### Training\n",
    "\n",
    "Now we are ready to train our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "b58b14e1ff08487aa8572cc79d242928",
      "cfdf6e891231495cadf4636063eeda53",
      "893cb4a9d8f94d7c9a6f90d46fa8a15e",
      "1ea2d22fcc7c42e796b8da898a9233f0",
      "0944b22a17c949f681afe1995f7f684c",
      "31c04490292b4a1e9c8623ad7ba19e8f",
      "6b834161adcc40548ab3e795e324565a",
      "5fc400b74c5241eb9b986a0268adfc9f",
      "2586235a37db44598e414f55d9c73883",
      "3d63d10d58464437adb4cc7ab05ea593",
      "9055399b0d844bfe92b95fc15942bb01"
     ]
    },
    "id": "yXy7j5HoXJtE",
    "outputId": "bed89288-e650-463c-8270-50f1d6bd50c7"
   },
   "outputs": [],
   "source": [
    "# ====\n",
    "# your code\n",
    "# choose these parameters\n",
    "EPOCHS = \n",
    "BATCH_SIZE = \n",
    "LR = \n",
    "\n",
    "EMB_DIM = \n",
    "NUM_HEADS = \n",
    "NUM_LAYERS = \n",
    "# ====\n",
    "\n",
    "image_gpt = ImageGPT((28, 28), EMB_DIM, NUM_HEADS, NUM_LAYERS)\n",
    "\n",
    "train_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_loader = data.DataLoader(test_data, batch_size=BATCH_SIZE)\n",
    "\n",
    "# choose any optimizer/scheduler as you want\n",
    "optimizer = torch.optim.Adam(image_gpt.parameters(), lr=LR)\n",
    "\n",
    "train_model(\n",
    "    image_gpt,\n",
    "    train_loader,\n",
    "    test_loader,\n",
    "    epochs=EPOCHS,\n",
    "    optimizer=optimizer,\n",
    "    device=DEVICE,\n",
    "    n_samples=16, # Remember the infernce is slow, so, you can lower the value \n",
    "    visualize_samples=True # or you can turn off smapling by setting `visualize_samples=False`\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hz5t-srXk2-i"
   },
   "source": [
    "Let sample from our model. You probably get better samples than PixelCNN samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 523
    },
    "id": "jUC3JBReYCau",
    "outputId": "30fe5fc3-5f71-457c-910b-36dd8a5e8c44"
   },
   "outputs": [],
   "source": [
    "samples = image_gpt.sample(25)\n",
    "show_samples(samples, title=\"ImageGPT samples\", nrow=5)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "784c2e2ba9cdfc1ce1e8c565791a35aa78e810f3990b00899de93c9f5ea5b088"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0773888b28204289b0c9eb77c975ea91": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "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": ""
     }
    },
    "0944b22a17c949f681afe1995f7f684c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "0aa50015efba4e22b0f7fea2c25c20be": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "1c55fe8da4c848478801552a1aa3d420": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "1c99231bb4ee40249e67cc233833faee": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "1ea2d22fcc7c42e796b8da898a9233f0": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "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_3d63d10d58464437adb4cc7ab05ea593",
      "placeholder": "​",
      "style": "IPY_MODEL_9055399b0d844bfe92b95fc15942bb01",
      "value": " 4/4 [18:33&lt;00:00, 278.10s/it]"
     }
    },
    "224de606324240b89a65ec0a537910f2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "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": ""
     }
    },
    "23c4bb2b6af14243b554f8d0ace59475": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "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": ""
     }
    },
    "2586235a37db44598e414f55d9c73883": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "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": ""
     }
    },
    "31c04490292b4a1e9c8623ad7ba19e8f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "3d63d10d58464437adb4cc7ab05ea593": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "439e35ef3240410d9634834725f97411": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "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_0aa50015efba4e22b0f7fea2c25c20be",
      "placeholder": "​",
      "style": "IPY_MODEL_bd5ac497e2654ad9a9763a2e7ec96326",
      "value": "100%"
     }
    },
    "47c79f4511e5463fbff01e4ee2ba0f03": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "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_c343c766d3f74e2d936bb6b12e91dabb",
      "max": 4,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_23c4bb2b6af14243b554f8d0ace59475",
      "value": 4
     }
    },
    "481551c56615463f913d7c88850917c2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "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_439e35ef3240410d9634834725f97411",
       "IPY_MODEL_47c79f4511e5463fbff01e4ee2ba0f03",
       "IPY_MODEL_5d7677ac892b49e68cdd6a366a86e184"
      ],
      "layout": "IPY_MODEL_62d49bf2ec954f4eaa1ed42d1ea1b329"
     }
    },
    "4a84d8a794ff4aeb8376f85e45384a89": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "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_6bb8abfa64ab4e1280cf3efe38b3f7a8",
      "max": 4,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_67b3bc3b3359480f863649780ec1cf13",
      "value": 4
     }
    },
    "5425bd677e334fa385ff37ddf7eee1d5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "5d7677ac892b49e68cdd6a366a86e184": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "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_1c99231bb4ee40249e67cc233833faee",
      "placeholder": "​",
      "style": "IPY_MODEL_224de606324240b89a65ec0a537910f2",
      "value": " 4/4 [03:05&lt;00:00, 46.37s/it]"
     }
    },
    "5fc400b74c5241eb9b986a0268adfc9f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "62d49bf2ec954f4eaa1ed42d1ea1b329": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "67b3bc3b3359480f863649780ec1cf13": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "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": ""
     }
    },
    "6b834161adcc40548ab3e795e324565a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "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": ""
     }
    },
    "6bb8abfa64ab4e1280cf3efe38b3f7a8": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "7045dbbf8309480bb4a697edb4f4d127": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "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": ""
     }
    },
    "72ca46babd314fa88a8e59b2aa41f7af": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "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_8c11c45785494f2593bf113240f97f46",
      "placeholder": "​",
      "style": "IPY_MODEL_0773888b28204289b0c9eb77c975ea91",
      "value": " 4/4 [03:05&lt;00:00, 46.39s/it]"
     }
    },
    "893cb4a9d8f94d7c9a6f90d46fa8a15e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "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_5fc400b74c5241eb9b986a0268adfc9f",
      "max": 4,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_2586235a37db44598e414f55d9c73883",
      "value": 4
     }
    },
    "8c11c45785494f2593bf113240f97f46": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "9055399b0d844bfe92b95fc15942bb01": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "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": ""
     }
    },
    "b58b14e1ff08487aa8572cc79d242928": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "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_cfdf6e891231495cadf4636063eeda53",
       "IPY_MODEL_893cb4a9d8f94d7c9a6f90d46fa8a15e",
       "IPY_MODEL_1ea2d22fcc7c42e796b8da898a9233f0"
      ],
      "layout": "IPY_MODEL_0944b22a17c949f681afe1995f7f684c"
     }
    },
    "bd5ac497e2654ad9a9763a2e7ec96326": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "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": ""
     }
    },
    "c343c766d3f74e2d936bb6b12e91dabb": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "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
     }
    },
    "cbddb1ae1a1d4124b972851b77bc24b9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "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_f519525bb7b640bebd08160d3710008c",
       "IPY_MODEL_4a84d8a794ff4aeb8376f85e45384a89",
       "IPY_MODEL_72ca46babd314fa88a8e59b2aa41f7af"
      ],
      "layout": "IPY_MODEL_1c55fe8da4c848478801552a1aa3d420"
     }
    },
    "cfdf6e891231495cadf4636063eeda53": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "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_31c04490292b4a1e9c8623ad7ba19e8f",
      "placeholder": "​",
      "style": "IPY_MODEL_6b834161adcc40548ab3e795e324565a",
      "value": "100%"
     }
    },
    "f519525bb7b640bebd08160d3710008c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "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_5425bd677e334fa385ff37ddf7eee1d5",
      "placeholder": "​",
      "style": "IPY_MODEL_7045dbbf8309480bb4a697edb4f4d127",
      "value": "100%"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}