{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:10:11.069094Z",
     "start_time": "2020-12-29T23:10:10.987566Z"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T00:26:09.816972Z",
     "start_time": "2020-12-30T00:26:09.761925Z"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from htools import *\n",
    "from incendio.layers import mish\n",
    "from incendio.utils import get_einops_image, plot_img_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T00:26:12.699223Z",
     "start_time": "2020-12-30T00:26:12.659006Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current directory: /Users/hmamin/incendio\n"
     ]
    }
   ],
   "source": [
    "cd_root()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T00:26:13.994392Z",
     "start_time": "2020-12-30T00:26:13.954643Z"
    }
   },
   "outputs": [],
   "source": [
    "bs = 2\n",
    "n_seq = 4\n",
    "n_emb = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T00:26:14.289000Z",
     "start_time": "2020-12-30T00:26:14.218655Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 8])"
      ]
     },
     "execution_count": 133,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(bs, n_seq, n_emb)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:19:21.446113Z",
     "start_time": "2020-12-29T23:19:21.408650Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 24])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lin = nn.Linear(n_emb, n_emb*3)\n",
    "qkv = lin(x)\n",
    "qkv.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:19:22.355886Z",
     "start_time": "2020-12-29T23:19:22.322446Z"
    }
   },
   "outputs": [],
   "source": [
    "q, k, v = qkv.chunk(3, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:19:36.854022Z",
     "start_time": "2020-12-29T23:19:36.804119Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[torch.Size([2, 4, 8]), torch.Size([2, 4, 8]), torch.Size([2, 4, 8])]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "smap(q, k, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:19:43.300421Z",
     "start_time": "2020-12-29T23:19:43.264827Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 4])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "qk = q @ k.transpose(-2, -1)\n",
    "qk.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:19:51.708470Z",
     "start_time": "2020-12-29T23:19:51.670169Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 4])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w = F.softmax(qk/np.sqrt(k.shape[-1]), dim=-1)\n",
    "w.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:19:56.970782Z",
     "start_time": "2020-12-29T23:19:56.935129Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 4, 4]), torch.Size([2, 4, 8]))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w.shape, v.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:19:57.612273Z",
     "start_time": "2020-12-29T23:19:57.576171Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.3194, 0.2145, 0.2537, 0.2123],\n",
       "         [0.1481, 0.3938, 0.1637, 0.2944],\n",
       "         [0.3235, 0.2150, 0.2943, 0.1672],\n",
       "         [0.2114, 0.3958, 0.1450, 0.2478]],\n",
       "\n",
       "        [[0.1403, 0.1413, 0.4371, 0.2814],\n",
       "         [0.1405, 0.1811, 0.4477, 0.2307],\n",
       "         [0.2538, 0.2546, 0.2017, 0.2899],\n",
       "         [0.1906, 0.2334, 0.2060, 0.3700]]], grad_fn=<SoftmaxBackward>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:20:06.853377Z",
     "start_time": "2020-12-29T23:20:06.819172Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0000, 1.0000, 1.0000, 1.0000],\n",
       "        [1.0000, 1.0000, 1.0000, 1.0000]], grad_fn=<SumBackward1>)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w.sum(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:20:08.551042Z",
     "start_time": "2020-12-29T23:20:08.488036Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.9252, -0.7189,  0.5745, -0.0615, -0.8207, -0.2445,  0.1443,\n",
       "           0.2584],\n",
       "         [-0.1945,  0.0850,  0.6434,  0.5749,  0.2454, -0.1575, -0.2515,\n",
       "           0.9941],\n",
       "         [-0.3747, -0.3643,  0.0776,  0.1407, -0.7063, -0.2824,  0.0551,\n",
       "           0.2374],\n",
       "         [-0.4754,  0.6410,  0.0406, -0.0262,  0.4806, -1.2904, -0.1688,\n",
       "           0.8364]],\n",
       "\n",
       "        [[-0.8540,  0.4700,  0.0350,  0.1249,  0.2180, -0.8239,  0.5342,\n",
       "           0.9115],\n",
       "         [-0.6555, -0.1389,  0.0239, -0.5375, -0.3968, -1.3929, -0.4080,\n",
       "           0.5622],\n",
       "         [-0.0428, -0.5596,  0.8501,  0.3117, -0.2734,  0.3237,  0.3952,\n",
       "           0.7681],\n",
       "         [-0.4201,  0.8392,  0.4082,  0.9275,  0.6644,  0.3775,  0.9765,\n",
       "           1.1330]]], grad_fn=<SplitBackward>)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-29T23:20:10.436580Z",
     "start_time": "2020-12-29T23:20:10.389198Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 8])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = w @ v\n",
    "res.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 369,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-01T02:01:27.203174Z",
     "start_time": "2021-01-01T02:01:26.740342Z"
    },
    "run_control": {
     "marked": false
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    }
   ],
   "source": [
    "class Projector(nn.Module):\n",
    "    \"\"\"Project input into multiple spaces. Used in DotProductAttention to\n",
    "    generate queries/keys/values.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_in, n_out_single=None, spaces=3):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_in: int\n",
    "            Size of input feature dimension, where input is (bs, n_in) or\n",
    "            (bs, seq_len, n_in). If the latter, this ONLY transforms the last \n",
    "            dimension. If you want to take multiple dimensions of information \n",
    "            into account simultaneously, you can flatten the input prior to \n",
    "            passing it in.\n",
    "        n_out_single: int or None\n",
    "            This determines the size of the feature dimension in each new \n",
    "            space. By default, this will be the same as n_in.\n",
    "        spaces: int\n",
    "            Number of spaces to project the input into. Default is 3 because\n",
    "            we commonly use this to generate queries, keys, and values for\n",
    "            attention computations.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.spaces = spaces\n",
    "        self.n_in = n_in\n",
    "        self.n_out_single = n_out_single or n_in\n",
    "        self.spaces = spaces\n",
    "        self.n_out = self.n_out_single * self.spaces\n",
    "        self.fc = nn.Linear(self.n_in, self.n_out)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.Tensor\n",
    "            Shape (bs, n_in) or (bs, seq_len, n_in).\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        tuple[torch.Tensor]: Tuple of `spaces` tensors where each tensor has\n",
    "        shape (bs, n_out_single) or (bs, seq_len, n_out_single), depending on\n",
    "        the input shape.\n",
    "        \"\"\"\n",
    "        return self.fc(x).chunk(self.spaces, dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Detour: testing fc projector vs conv projector\n",
    "\n",
    "NOTE: nn.Linear only acts on last dimension so if we pass in input of shape (bs, seq_len, emb_dim), we're not really using the seq dimension here. Maybe refresh memory of 1d conv to see if that would take both into account in the way I want. Otherwise we could do some reshaping before feeding to linear layer.\n",
    "\n",
    "UPDATE: Now that I think about this, maybe this behavior is desirable. At this stage we're not looking to take info from other time steps into account: we're just looking to get 3 different representations of EACH time step separately, and attention will be used to blend this information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:31:57.218929Z",
     "start_time": "2020-12-30T23:31:57.157916Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Object loaded from /Users/hmamin/.incendio/einops.zip.\n",
      "Object loaded from /Users/hmamin/.incendio/einops.zip.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([6, 3, 96, 96]), torch.Size([6, 96, 96]))"
      ]
     },
     "execution_count": 201,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img = get_einops_image()\n",
    "img_bw = get_einops_image(color=False)\n",
    "img.shape, img_bw.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:31:57.692113Z",
     "start_time": "2020-12-30T23:31:57.402621Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19100 characters]\n",
      "text/plain": [
       "<Figure size 288x864 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_img_batch(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:31:58.290456Z",
     "start_time": "2020-12-30T23:31:58.040068Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 18800 characters]\n",
      "text/plain": [
       "<Figure size 288x864 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_img_batch(img_bw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:32:00.835440Z",
     "start_time": "2020-12-30T23:32:00.788818Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([18, 96, 96])"
      ]
     },
     "execution_count": 204,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "proj = Projector(*img_bw.shape[1:], 3)\n",
    "with torch.no_grad():\n",
    "    proj.fc.weight.data.fill_(2)\n",
    "    \n",
    "qkv = torch.cat(proj(img_bw), dim=0)\n",
    "qkv.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 221,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:36:25.133780Z",
     "start_time": "2020-12-30T23:36:25.053928Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n"
     ]
    }
   ],
   "source": [
    "# Because we filled weights with constant, q/k/v are all equivalent for a \n",
    "# single sample. Precision is low for some reason (e.g. tolerance of 1e-3\n",
    "# finds differences).\n",
    "for i in range(6):\n",
    "    print(i)\n",
    "    assert torch.isclose(qkv[i], qkv[i+6], 1e-2, 1e-2).float().mean() == 1\n",
    "    assert torch.isclose(qkv[i], qkv[i+12], 1e-2, 1e-2).float().mean() == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:32:03.086589Z",
     "start_time": "2020-12-30T23:32:02.456461Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 51908 characters]\n",
      "text/plain": [
       "<Figure size 288x2592 with 18 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_img_batch(qkv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:23:24.889430Z",
     "start_time": "2020-12-30T23:23:24.727965Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6, 192, 96])"
      ]
     },
     "execution_count": 195,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = nn.Conv1d(img_bw.shape[-1], img_bw.shape[-1]*2, \n",
    "                 kernel_size=3, padding=1)\n",
    "with torch.no_grad():\n",
    "    conv.weight.data.fill_(2)\n",
    "# bw_out = conv(img_bw.permute(0, 2, 1))\n",
    "bw_out = conv(img_bw)\n",
    "bw_out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:23:26.582048Z",
     "start_time": "2020-12-30T23:23:26.346622Z"
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9772 characters]\n",
      "text/plain": [
       "<Figure size 288x864 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_img_batch(bw_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:18:07.330843Z",
     "start_time": "2020-12-30T23:18:07.271102Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4608"
      ]
     },
     "execution_count": 181,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bw_out.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-30T23:18:55.765636Z",
     "start_time": "2020-12-30T23:18:55.365871Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([6, 2304])\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 42836 characters]\n",
      "text/plain": [
       "<Figure size 288x864 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fc = nn.Linear(img_bw.shape[-1] * img_bw.shape[-2], \n",
    "               img_bw.shape[-1] * img_bw.shape[-2] // 4)\n",
    "with torch.no_grad():\n",
    "    fc.weight.data.fill_(2)\n",
    "bw_out = fc(img_bw.view(img_bw.shape[0], -1))\n",
    "plot_img_batch(\n",
    "    bw_out.view(bw_out.shape[0], int(np.sqrt(bw_out.shape[1])), -1)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 370,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-01T02:13:51.982472Z",
     "start_time": "2021-01-01T02:13:51.920392Z"
    }
   },
   "outputs": [],
   "source": [
    "class DotProductAttention(nn.Module):\n",
    "    \"\"\"GPT2-style attention block. This was mostly an intuition-building\n",
    "    exercise - in practice, Huggingface provides layers that should probably\n",
    "    be used instead.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_in, n_out=None, nf=None, n_heads=12,\n",
    "                 temperature='auto', p1=0.1, p2=0.1, return_attn=False):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_in: int\n",
    "            Last dimension of input, usually embedding dimension.\n",
    "        n_out: int or None\n",
    "            Size of output vectors. By default, this will be the same as the \n",
    "            input.\n",
    "        nf: int or None\n",
    "            Size (\"nf = number of features\") of queries/keys/values. \n",
    "            By default, this will be the same as n_in. Must be divisible by\n",
    "            n_heads.\n",
    "        n_heads: int\n",
    "            Number of attention heads to use. nf must be divisible\n",
    "            by this as each projected vector will be divided evenly among\n",
    "            each head.\n",
    "        temperature: str or float\n",
    "            If str, must be \"auto\", meaning softmax inputs will be scaled by\n",
    "            sqrt(n_proj_single). You can also specify a float, where values\n",
    "            <1 sharpen the distribution (usually not what we want here) and\n",
    "            values greater than one soften it (allowing attention head to \n",
    "            route more information from multiple neurons rather than almost \n",
    "            all from one).\n",
    "        p1: float\n",
    "            Value in (0.0, 1.0) setting the dropout probability on the \n",
    "            attention weights.\n",
    "        p2: float\n",
    "            Value in (0.0, 1.0) setting dropout probability following the \n",
    "            output layer.\n",
    "        return_attn: bool\n",
    "            If True, the `forward` method will return a tuple of \n",
    "            (output, attention_weights) tensors. If False (the default), just\n",
    "            return the output tensor.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        nf = nf or n_in\n",
    "        n_out = n_out or n_in\n",
    "        assert nf % n_heads == 0, \\\n",
    "            'n_proj_single must be divisible by n_heads'\n",
    "        \n",
    "        self.proj_in = Projector(n_in, nf, spaces=3)\n",
    "        # Reshape so hidden dimension is split equally between each head.\n",
    "        self.head_splitter = Rearrange('bs seq (heads f) -> bs heads seq f', \n",
    "                                       heads=n_heads)\n",
    "        self.soft = SmoothSoftmax(temperature)\n",
    "        self.drop_attn = nn.Dropout(p1)\n",
    "        # Concatenate output of each head.\n",
    "        self.head_merger = Rearrange('bs heads seq f -> bs seq (heads f)')\n",
    "        self.fc_out = nn.Linear(nf, n_out)\n",
    "        self.drop_out = nn.Dropout(p2)\n",
    "    \n",
    "        # Non-layer attributes.\n",
    "        self.n_heads = n_heads\n",
    "        self.temperature = temperature\n",
    "        self.p1 = p1\n",
    "        self.p2 = p2\n",
    "        self.return_attn = return_attn\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.Tensor\n",
    "            Shape (bs, seq_len, n_in). n_in will usually be the sum of\n",
    "            embedding dimensions (word and positional). For other problems\n",
    "            (e.g. web browsing sequence classificaiton), this might include\n",
    "            other features about the page at time step T.\n",
    "        \"\"\"\n",
    "        q, k, v = map(self.head_splitter, self.proj_in(x))\n",
    "        scores = q @ k.transpose(-2, -1)\n",
    "        weights = self.drop_attn(self.soft(scores))\n",
    "        x = weights @ v\n",
    "        x = self.head_merger(x)\n",
    "        x = self.drop_out(self.fc_out(x))\n",
    "        return (x, weights) if self.return_attn else x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 364,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-31T00:33:59.100527Z",
     "start_time": "2020-12-31T00:33:59.060325Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 6, 768])\n",
      "torch.Size([2, 12, 6, 64])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 768])"
      ]
     },
     "execution_count": 364,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(x.shape)\n",
    "z = attn.head_splitter(x)\n",
    "print(z.shape)\n",
    "merger = Rearrange('bs heads seq f -> bs seq (heads f)')\n",
    "merger(z).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 365,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-31T00:34:03.247634Z",
     "start_time": "2020-12-31T00:34:03.154730Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 6, 768])\n",
      "torch.Size([2, 6, 768])\n"
     ]
    }
   ],
   "source": [
    "bs = 2\n",
    "seq_len = 6\n",
    "n_emb = 768\n",
    "n_heads = 12\n",
    "x = torch.randn(bs, seq_len, n_emb)\n",
    "print(x.shape)\n",
    "\n",
    "attn = DotProductAttention(n_emb, n_heads=n_heads)\n",
    "res = attn(x)\n",
    "print(res.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 366,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-31T00:34:17.616465Z",
     "start_time": "2020-12-31T00:34:17.533985Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 768])"
      ]
     },
     "execution_count": 366,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn = DotProductAttention(n_emb, n_heads=3)\n",
    "res = attn(x)\n",
    "res.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 371,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-01T02:22:35.926151Z",
     "start_time": "2021-01-01T02:22:35.867699Z"
    }
   },
   "outputs": [],
   "source": [
    "class FanForward(nn.Module):\n",
    "    \"\"\"Fan out to a larger dimension and back in to the original space.\n",
    "    Not fully clear on why this is necessary or useful but we apparently need\n",
    "    it for our GPT2-esque decoder.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_in, n_hid=None, scalar=4, p=0.1, act=mish):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_in: int\n",
    "            Number of input features. x will typically have shape \n",
    "            (bs, seq_len, n_in).\n",
    "        n_hid: int or None\n",
    "            If int, this will be the size of the output of the first linear \n",
    "            layer. If None, hidden_scalar will be used to determine this.\n",
    "        scalar: int\n",
    "            In practice, we often want our hidden layer to produce outputs\n",
    "            of n_in*n where n is some semi-arbitrary scalar. If n_hid is None,\n",
    "            this scalar will be used to compute n_hid from n_in.\n",
    "        p: float\n",
    "            Dropout probability. Value in (0.0, 1.0). This is applied after \n",
    "            the output layer but not the input layer for some reason.\n",
    "        act: nn.Module or function\n",
    "            Callable activation function. This will be used after both the\n",
    "            first and second linear layers.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        n_hid = ifnone(n_hid, n_in * scalar)\n",
    "        self.fc1 = nn.Linear(n_in, n_hid)\n",
    "        self.fc2 = nn.Linear(n_hid, n_in)\n",
    "        self.act = act\n",
    "        self.drop = nn.Dropout(p)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.Tensor\n",
    "            Shape (bs, n_in) or (bs, seq_len, n_in).\n",
    "            \n",
    "        Output\n",
    "        ------\n",
    "        torch.Tensor: Same shape as input tensor.\n",
    "        \"\"\"\n",
    "        x = self.act(self.fc1(x))\n",
    "        return self.drop(self.act(self.fc2(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 351,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-31T00:29:08.069337Z",
     "start_time": "2020-12-31T00:29:08.026510Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 24])"
      ]
     },
     "execution_count": 351,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff = FanForward(24, 2)\n",
    "ff(res).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 373,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:26:20.791136Z",
     "start_time": "2021-01-03T04:26:20.626242Z"
    }
   },
   "outputs": [],
   "source": [
    "class TransformerDecoder(nn.Module):\n",
    "    \"\"\"GPT2-esque decoder network. This was mostly an intuition-building\n",
    "    exercise - in practice, it's probably best to just use something from \n",
    "    Huggingface.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_in=768, n_heads=12):\n",
    "        super().__init__()\n",
    "        self.attn = DotProductAttention(n_in, n_heads=n_heads)\n",
    "        self.ff = FanForward(n_in)\n",
    "        self.norm_1 = nn.LayerNorm(n_in)\n",
    "        self.norm_2 = nn.LayerNorm(n_in)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x += self.attn(self.norm_1(x))\n",
    "        return x + self.ff(self.norm_2(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 374,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:26:21.800707Z",
     "start_time": "2021-01-03T04:26:21.748664Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FanForward(\n",
       "  (fc1): Linear(in_features=5, out_features=6, bias=True)\n",
       "  (fc2): Linear(in_features=6, out_features=5, bias=True)\n",
       "  (drop): Dropout(p=0.1, inplace=False)\n",
       ")"
      ]
     },
     "execution_count": 374,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FanForward(5, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 375,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:26:28.638510Z",
     "start_time": "2021-01-03T04:26:28.577177Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 768])"
      ]
     },
     "execution_count": 375,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ff = FanForward(n_emb, 4)\n",
    "ff(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 376,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:26:39.371921Z",
     "start_time": "2021-01-03T04:26:39.194732Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/hmamin/anaconda3/lib/python3.7/site-packages/ipykernel/ipkernel.py:287: DeprecationWarning: `should_run_async` will not call `transform_cell` automatically in the future. Please pass the result to `transformed_cell` argument and any exception that happen during thetransform in `preprocessing_exc_tuple` in IPython 7.17 and above.\n",
      "  and should_run_async(code)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TransformerDecoder(\n",
       "  (attn): DotProductAttention(\n",
       "    (proj_in): Projector(\n",
       "      (fc): Linear(in_features=768, out_features=2304, bias=True)\n",
       "    )\n",
       "    (head_splitter): Rearrange('bs seq (heads f) -> bs heads seq f', heads=12)\n",
       "    (soft): SmoothSoftmax(\n",
       "      (act): Softmax(dim=-1)\n",
       "    )\n",
       "    (drop_attn): Dropout(p=0.1, inplace=False)\n",
       "    (head_merger): Rearrange('bs heads seq f -> bs seq (heads f)')\n",
       "    (fc_out): Linear(in_features=768, out_features=768, bias=True)\n",
       "    (drop_out): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (ff): FanForward(\n",
       "    (fc1): Linear(in_features=768, out_features=3072, bias=True)\n",
       "    (fc2): Linear(in_features=3072, out_features=768, bias=True)\n",
       "    (drop): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (norm_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       "  (norm_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
       ")"
      ]
     },
     "execution_count": 376,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoder = TransformerDecoder()\n",
    "decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 379,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:27:01.519901Z",
     "start_time": "2021-01-03T04:27:01.447053Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 768])"
      ]
     },
     "execution_count": 379,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 378,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-01-03T04:26:57.143765Z",
     "start_time": "2021-01-03T04:26:57.085474Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 768])"
      ]
     },
     "execution_count": 378,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "decoder(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T21:04:09.920843Z",
     "start_time": "2020-07-05T21:04:09.868279Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers.modeling_gpt2 import Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T21:06:51.382375Z",
     "start_time": "2020-07-05T21:06:51.336974Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 166,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = Args(n_head=1, attn_pdrop=.1, resid_pdrop=.1, n_ctx=n_emb*2, \n",
    "              n_layer=1, n_embd=n_emb, layer_norm_epsilon=1e-3,\n",
    "              output_attentions=False)\n",
    "attn_gpt = Attention(n_emb, n_emb*2, config)\n",
    "res_gpt = attn_gpt(x)\n",
    "len(res_gpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T21:07:20.548943Z",
     "start_time": "2020-07-05T21:07:20.507601Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 3])"
      ]
     },
     "execution_count": 169,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res_gpt[-2].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T21:10:28.107648Z",
     "start_time": "2020-07-05T21:10:28.040112Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 9])"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn_gpt.c_attn(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T04:45:14.610064Z",
     "start_time": "2020-07-05T04:45:14.554130Z"
    }
   },
   "outputs": [],
   "source": [
    "# FROM BLOG POST\n",
    "class Conv1D(nn.Module):\n",
    "    def __init__(self, nx, nf):\n",
    "        super().__init__()\n",
    "        self.nf = nf\n",
    "        w = torch.empty(nx, nf)\n",
    "        nn.init.normal_(w, std=0.02)\n",
    "        self.weight = nn.Parameter(w)\n",
    "        self.bias = nn.Parameter(torch.zeros(nf))\n",
    "\n",
    "    def forward(self, x):\n",
    "        size_out = x.size()[:-1] + (self.nf,)\n",
    "        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)\n",
    "        x = x.view(*size_out)\n",
    "        return x\n",
    "    \n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, dropout, d_model=768, nx=768*4):\n",
    "        super().__init__()\n",
    "        self.c_fc    = Conv1D(d_model, nx)\n",
    "        self.c_proj  = Conv1D(nx, d_model)\n",
    "        self.act     = F.gelu\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.act(self.c_fc(x))\n",
    "        print(x.shape)\n",
    "        return self.dropout(self.c_proj(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T04:45:15.335366Z",
     "start_time": "2020-07-05T04:45:15.268366Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 4, 12])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 3])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FeedForward(0, n_emb, n_emb*4)(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T04:37:07.515736Z",
     "start_time": "2020-07-05T04:37:07.447864Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 4, 4]), torch.Size([2, 4, 3]))"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w.shape, v.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T04:17:32.277305Z",
     "start_time": "2020-07-05T04:17:32.240425Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 3, 3])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = torch.randn(bs, n_seq, n_emb, n_emb)\n",
    "z.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T04:18:00.547709Z",
     "start_time": "2020-07-05T04:18:00.487208Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 4, 3, 5])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lin2 = nn.Linear(3, 5)\n",
    "lin2(z).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T04:18:08.820952Z",
     "start_time": "2020-07-05T04:18:08.775091Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[-0.5565, -0.2874, -0.5044],\n",
       "        [ 0.3084,  0.2368, -0.2929],\n",
       "        [-0.1206,  0.0814, -0.2280],\n",
       "        [ 0.1191,  0.0834, -0.2181],\n",
       "        [-0.0737,  0.3648,  0.3930]], requires_grad=True)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lin2.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-05T04:20:00.134476Z",
     "start_time": "2020-07-05T04:20:00.095313Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z.dim()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
