{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Layers\n",
    "\n",
    "> Custom activations, layers, and layer blocks are contained in this module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from fastai.layers import trunc_normal_\n",
    "import numpy as np\n",
    "from operator import add, truediv, sub\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from htools import add_docstring\n",
    "from incendio.data import probabilistic_hash_tensor\n",
    "from incendio.utils import concat, weighted_avg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Used for testing only.\n",
    "from collections import defaultdict, Counter\n",
    "from itertools import chain\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from htools import assert_raises, InvalidArgumentError\n",
    "from incendio.data import probabilistic_hash_item\n",
    "import pandas_htools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class GRelu(nn.Module):\n",
    "    \"\"\"Generic ReLU.\"\"\"\n",
    "\n",
    "    def __init__(self, leak=0.0, max=float('inf'), sub=0.0):\n",
    "        super().__init__()\n",
    "        self.leak = leak\n",
    "        self.max = max\n",
    "        self.sub = sub\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"Check which operations are necessary to save computation.\"\"\"\n",
    "        x = F.leaky_relu(x, self.leak) if self.leak else F.relu(x)\n",
    "        if self.sub:\n",
    "            x -= self.sub\n",
    "        if self.max:\n",
    "            x = torch.clamp(x, max=self.max)\n",
    "        return x\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f'GReLU(leak={self.leak}, max={self.max}, sub={self.sub})'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "JRelu = GRelu(leak=.1, sub=.4, max=6.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Mish(nn.Module):\n",
    "    \"\"\"OOP form of mish activation.\n",
    "\n",
    "    Mish: A Self Regularized Non-Monotonic Neural Activation Function\n",
    "    https://arxiv.org/pdf/1908.08681v1.pdf\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x * torch.tanh(F.softplus(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "def mish(x):\n",
    "    \"\"\"Functional form of mish activation.\n",
    "\n",
    "    Mish: A Self Regularized Non-Monotonic Neural Activation Function\n",
    "    https://arxiv.org/pdf/1908.08681v1.pdf\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x: torch.Tensor[float]\n",
    "        Input tensor.\n",
    "    Returns\n",
    "    -------\n",
    "    torch.Tensor[float]: Tensor of same shape as input x.\n",
    "    \"\"\"\n",
    "    return x * torch.tanh(F.softplus(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_activations(z, a, mode='scatter', **kwargs):\n",
    "    \"\"\"Plot an input tensor and its corresponding activations.  Both tensors\n",
    "    will be flattened for plotting.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    z: tf.Tensor\n",
    "        Tensor containing values to plot on the x axis (we can often think of\n",
    "        this as the output of a linear layer, where z=f(x) and a=mish(z)).\n",
    "    a: tf.Tensor\n",
    "        Tensor containing values to plot on y axis.\n",
    "    mode: str\n",
    "        'scatter' for scatter plot or 'plot' for line plot.\n",
    "    kwargs: Values to be passed to the matplotlib plotting function, such as \n",
    "        's' when in 'scatter' mode or 'lw' in 'plot' mode.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    None\n",
    "    \"\"\"\n",
    "    plt_func = getattr(plt, mode)\n",
    "    kwargs = kwargs or {}\n",
    "    if mode == 'scatter' and not kwargs:\n",
    "        kwargs = {'s': .75}\n",
    "    plt_func(z.numpy().flatten(), a.numpy().flatten(), **kwargs)\n",
    "    plt.axvline(0, lw=.5, alpha=.5)\n",
    "    plt.axhline(0, lw=.5, alpha=.5)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.arange(-5, 5, .05)\n",
    "a = mish(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10384 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_activations(x, a, 'plot')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer Blocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ConvBlock(nn.Module):\n",
    "    \"\"\"Create a convolutional block optionally followed by a batch norm layer.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, c_in, c_out, kernel_size=3, norm=True, activation=JRelu,\n",
    "                 **kwargs):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        -----------\n",
    "        c_in: int\n",
    "            # of input channels.\n",
    "        c_out: int\n",
    "            # of output channels.\n",
    "        kernel_size: int\n",
    "            Size of kernel in conv2d layer. An integer argument will be used\n",
    "            as both the height and width.\n",
    "        norm: bool\n",
    "            If True, include a batch norm layer after the conv layer. If False,\n",
    "            no norm layer will be used. Note that batch norm has learnable\n",
    "            affine parameters which remove the need for a bias in the preceding\n",
    "            conv layer. When batch norm is not used, however, the conv layer\n",
    "            will include a bias term.\n",
    "        activation: nn.Module\n",
    "            Activation function to use at the end of the convolutional block.\n",
    "            (In some cases such as our ResBlock implementation, we pass in None\n",
    "            so that an extra addition can be performed before the final\n",
    "            activation.) Do not use the functional form here as it will be\n",
    "            added to a sequential object.\n",
    "        kwargs: any\n",
    "            Additional keyword args are passed to Conv2d. Useful kwargs include\n",
    "            stride, and padding (see pytorch docs for nn.Conv2d).\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.norm = norm\n",
    "        layers = [nn.Conv2d(c_in, c_out, kernel_size, bias=not norm, **kwargs)]\n",
    "        if norm:\n",
    "            layers.append(nn.BatchNorm2d(c_out))\n",
    "        if activation is not None:\n",
    "            layers.append(activation)\n",
    "        self.block = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.block(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvBlock(\n",
       "  (block): Sequential(\n",
       "    (0): Conv2d(3, 5, kernel_size=(3, 3), stride=(1, 1))\n",
       "    (1): GReLU(leak=0.1, max=6.0, sub=0.4)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = ConvBlock(3, 5, norm=False)\n",
    "conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5, 2, 2])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.rand(2, 3, 4, 4)\n",
    "conv(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class ResBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, c_in, activation=JRelu, f=3, stride=1, pad=1,\n",
    "                 skip_size=2, norm=True):\n",
    "        \"\"\"Residual block using 2D convolutional layers. Note that f,\n",
    "        stride, and pad must be selected such that the height and width of \n",
    "        the input remain the same.\n",
    "\n",
    "        Parameters\n",
    "        -----------\n",
    "        c_in: int\n",
    "            # of input channels.\n",
    "        activation: callable\n",
    "            Activation function to use.\n",
    "        f: int\n",
    "            Size of filter (f x f) used in convolution. Default 3.\n",
    "        stride: int\n",
    "            # of pixels the filter moves between each convolution. Default 1.\n",
    "        pad: int\n",
    "            Pixel padding around the input. Default 1.\n",
    "        skip_size: int\n",
    "            Number of conv blocks inside the skip connection (default 2).\n",
    "            ResNet paper notes that skipping a single layer did not show\n",
    "            noticeable improvements.\n",
    "        norm: bool\n",
    "            Specifies whether to include a batch norm layer after each conv\n",
    "            layer.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.skip_size = skip_size\n",
    "        self.layers = nn.ModuleList([\n",
    "            ConvBlock(c_in, c_in, norm=norm, activation=None, kernel_size=f,\n",
    "                      stride=stride, padding=pad)\n",
    "            for i in range(skip_size)\n",
    "        ])\n",
    "        self.activation = activation\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_out = x\n",
    "        for i, layer in enumerate(self.layers):\n",
    "            x_out = layer(x_out)\n",
    "\n",
    "            # Final activation must be applied after addition.\n",
    "            if i != self.skip_size - 1:\n",
    "                x_out = self.activation(x_out)\n",
    "\n",
    "        return self.activation(x + x_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (layers): ModuleList(\n",
       "    (0): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (1): BatchNorm2d(4, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (activation): GReLU(leak=0.1, max=6.0, sub=0.4)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ResBlock(4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ResBlock(\n",
       "  (layers): ModuleList(\n",
       "    (0): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "    )\n",
       "    (1): ConvBlock(\n",
       "      (block): Sequential(\n",
       "        (0): Conv2d(4, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (activation): GReLU(leak=0.1, max=6.0, sub=0.4)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ResBlock(4, norm=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@add_docstring(nn.Conv2d)\n",
    "class ReflectionPaddedConv2d(nn.Module):\n",
    "    \"\"\"Conv2d only allows padding_mode of `zeros` or `circular`. This\n",
    "    layer is a quick way for us to use reflection padding.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, in_channels, out_channels, padding=1, \n",
    "                 kernel_size=3, **kwargs):\n",
    "        \"\"\"Do not specify a padding mode.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        if 'padding_mode' in kwargs: \n",
    "            raise InvalidArgumentError('Remove `padding_mode` from arguments.')\n",
    "        self.reflect = nn.ReflectionPad2d(padding)\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,\n",
    "                              padding=0)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.reflect(x)\n",
    "        return self.conv(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_img(img):\n",
    "    plt.imshow(img.permute(1, 2, 0) / 255)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ReflectionPaddedConv2d(\n",
       "  (reflect): ReflectionPad2d((2, 2, 2, 2))\n",
       "  (conv): Conv2d(3, 3, kernel_size=(1, 1), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rconv = ReflectionPaddedConv2d(3, 3, kernel_size=1, padding=2)\n",
    "rconv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 5016 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = torch.randint(255, (1, 3, 3, 3)).float()\n",
    "show_img(x[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4180 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "x2 = rconv.reflect(x)\n",
    "show_img(x2[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got InvalidArgumentError(Remove `padding_mode` from arguments.).\n"
     ]
    }
   ],
   "source": [
    "# Tests\n",
    "assert nn.Conv2d.__doc__ in ReflectionPaddedConv2d.__doc__\n",
    "\n",
    "with assert_raises(InvalidArgumentError):\n",
    "    ReflectionPaddedConv2d(3, 3, padding_mode='zeros')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Dropin(nn.Module):\n",
    "    \"\"\"Additive dropout. This injects small amounts of noise into a model\n",
    "    in the form of randomly generated floats from a zero-centered\n",
    "    gaussian distribution (variance can be adjusted). This does nothing \n",
    "    in eval mode. Unlike Dropout, this does not scale weights during \n",
    "    training since it does not bias them in any direction.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, scale=.5):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        scale: float\n",
    "            Used to scale the magnitude of the random noise. Keep in mind \n",
    "            that the scalar term is square rooted, so the relationship\n",
    "            will not be linear. Relatively large values (e.g. 1.0) will have\n",
    "            a stronger regularizing effect, while small values (e.g. 0.1)\n",
    "            will have a slight regularizing effect. There is no max value\n",
    "            enforced, so it's up to the user to select a reasonable value.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.scale = scale\n",
    "        \n",
    "    def forward(self, x):\n",
    "        if not self.training:\n",
    "            return x\n",
    "        \n",
    "        # Storing noise allows us to run diagnostics.\n",
    "        self.noise = torch.randn_like(x) * np.sqrt(self.scale / x.shape[-1])\n",
    "        return x + self.noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.drop = Dropin()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.drop(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net()\n",
    "x = torch.randn(8, 128, 128, 3)\n",
    "assert np.corrcoef(net(x).flatten(), x.flatten())[0][1] > .9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.eval()\n",
    "assert torch.eq(net(x), x).all()\n",
    "assert not net.drop.training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_activation_stats(scale=1.0, trials=10_000):\n",
    "    act_stats = defaultdict(list)\n",
    "    noise_stats = defaultdict(list)\n",
    "    \n",
    "    drop = Dropin(scale)\n",
    "    for _ in range(trials):\n",
    "        x = torch.randn(3, 4, dtype=torch.float)\n",
    "        z = drop(x)\n",
    "        noise = drop.noise\n",
    "        noise_stats['mean'].append(noise.mean())\n",
    "        noise_stats['std'].append(noise.std())\n",
    "        noise_stats['act_corr'].append(\n",
    "            np.corrcoef(z.flatten(), noise.flatten())[0][1]\n",
    "        )\n",
    "        \n",
    "        act_stats['mean'].append(z.mean())\n",
    "        act_stats['std'].append(z.std())\n",
    "        act_stats['x_corr'].append(\n",
    "            np.corrcoef(z.flatten(), x.flatten())[0][1]\n",
    "        )\n",
    "\n",
    "    return pd.DataFrame(dict(\n",
    "        act={k: np.mean(v).round(4) for k, v in act_stats.items()}, \n",
    "        noise={k: np.mean(v).round(4) for k, v in noise_stats.items()}\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 10\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.0132</td>\n",
       "      <td>0.0094</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.8189</td>\n",
       "      <td>1.5192</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.5324</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.8304</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 1\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>-0.0141</td>\n",
       "      <td>0.0034</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0921</td>\n",
       "      <td>0.4870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.8855</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.4282</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.75\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>-0.0015</td>\n",
       "      <td>0.0022</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0633</td>\n",
       "      <td>0.4240</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9100</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.3899</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.5\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.0107</td>\n",
       "      <td>0.0008</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0558</td>\n",
       "      <td>0.3442</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9409</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.3235</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.25\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>0.0098</td>\n",
       "      <td>-0.0057</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>1.0013</td>\n",
       "      <td>0.2461</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9667</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.2298</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " 0.1\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>act</th>\n",
       "      <th>noise</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>mean</th>\n",
       "      <td>-0.0057</td>\n",
       "      <td>-0.0014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>std</th>\n",
       "      <td>0.9969</td>\n",
       "      <td>0.1533</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>x_corr</th>\n",
       "      <td>0.9868</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>act_corr</th>\n",
       "      <td>NaN</td>\n",
       "      <td>0.1394</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for scale in [10, 1, .75, .5, .25, .1]:\n",
    "    print('\\n', scale)\n",
    "    simulate_activation_stats(scale, 1_000).pprint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "@add_docstring(nn.Embedding)\n",
    "def PaddedEmbedding(num_embeddings, embedding_dim, padding_idx=None, **kwargs):\n",
    "    \"\"\"Patched version of Fastai `embedding` that allows us to specify a row of\n",
    "    zeros for a padding token.\n",
    "    \"\"\"\n",
    "    emb = nn.Embedding(num_embeddings, embedding_dim, padding_idx, **kwargs)\n",
    "    with torch.no_grad():\n",
    "        trunc_normal_(emb.weight, std=.01)\n",
    "        if padding_idx is not None:\n",
    "            torch.zero_(emb.weight[0])\n",
    "    return emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0145, -0.0067, -0.0038],\n",
       "        [-0.0176, -0.0179,  0.0010],\n",
       "        [-0.0042, -0.0020,  0.0027]], requires_grad=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PaddedEmbedding(4, 3, 0).weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class BloomEmbedding(nn.Module):\n",
    "    \"\"\"Bloom Embedding layer for memory-efficient word representations.\n",
    "    Each word is encoded by a combination of rows of the embedding\n",
    "    matrix. The number of rows can therefore be far lower than the number\n",
    "    of words in our vocabulary while still providing unique representations.\n",
    "    The reduction in rows allows us to use memory in other ways: a larger\n",
    "    embedding dimension, more or larger layers after the embedding,\n",
    "    larger batch sizes, etc.\n",
    "    \n",
    "    Note that if hashing is done in the Dataset, we could use a simple\n",
    "    nn.EmbeddingBag to achieve the same thing. Many users have reported \n",
    "    poor performance with this layer though (especially on CPU, but in some\n",
    "    cases on GPU) so I stick with the standard Embedding. We also bake in\n",
    "    the truncated normal intialization provided by fastai, with a slight tweak\n",
    "    to allow a row for padding.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n_emb=251, emb_dim=100, n_hashes=4, padding_idx=0,\n",
    "                 pre_hashed=False):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_emb: int\n",
    "            Number of rows to create in the embedding matrix. A prime\n",
    "            number is recommended. Lower numbers will be more \n",
    "            memory-efficient but increase the chances of collisions.\n",
    "        emb_dim: int\n",
    "            Size of each embedding. If emb_dim=100, each word will\n",
    "            be represented by a 100-dimensional vector.\n",
    "        n_hashes: int\n",
    "            This determines the number of hashes that will be taken\n",
    "            for each word index, and as a result, the number of rows\n",
    "            that will be summed to create each unique representation.\n",
    "            The higher the number, the lower the chances of a collision.\n",
    "        padding_idx: int or None\n",
    "            If an integer is provided, this will set aside the corresponding\n",
    "            row in the embedding matrix as a vector of zeros. If None, no\n",
    "            padding vector will be allocated.\n",
    "        pre_hashed: bool\n",
    "            Pass in True if the input tensor will already be hashed by the time \n",
    "            it enters this layer (you may prefer pre-compute the hashes in the\n",
    "            Dataset to save computation time during training). In this\n",
    "            scenario, the layer is a simple embedding bag with mode \"sum\". \n",
    "            Pass in False if the inputs will be word indices that have not yet\n",
    "            been hashed. In this case, hashing will be done inside the \n",
    "            `forward` call.\n",
    "            \n",
    "        Suggested values for a vocab size of ~30,000:\n",
    "        \n",
    "        | n_emb | n_hashes | unique combos |\n",
    "        |-------|----------|---------------|\n",
    "        | 127   | 5        | 29,998        |\n",
    "        | 251   | 4        | 29,996        |\n",
    "        | 997   | 3        | 29,997        |\n",
    "        | 5,003 | 2        | 29,969        |\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.n_emb = n_emb\n",
    "        self.emb = PaddedEmbedding(n_emb, emb_dim, padding_idx=padding_idx)\n",
    "        self.n_hashes = n_hashes\n",
    "        self.pad_idx = padding_idx\n",
    "        self.pre_hashed = pre_hashed\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.LongTensor\n",
    "            Input tensor of word indices (bs x seq_len) if pre_hashed is False.\n",
    "            Hashed indices (bs x seq_len x n_hashes) if pre_hashed is False.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        torch.FloatTensor: Words encoded with combination of embeddings.\n",
    "            (bs x seq_len x emb_dim)\n",
    "        \"\"\"\n",
    "        if not self.pre_hashed:\n",
    "            # (bs, seq_len) -> hash -> (bs, seq_len, n_hashes)\n",
    "            hashed = probabilistic_hash_tensor(x, \n",
    "                                               self.n_emb,\n",
    "                                               self.n_hashes,\n",
    "                                               self.pad_idx)\n",
    "        # (bs, seq_len, n_hashes, emb_dim) -> sum -> (bs, seq_len, emb_dim)\n",
    "        return self.emb(hashed).sum(-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Data(Dataset):\n",
    "    \n",
    "    def __init__(self, sentences, labels, seq_len):\n",
    "        x = [s.split(' ') for s in sentences]\n",
    "        self.w2i = self.make_w2i(x)\n",
    "        self.seq_len = seq_len\n",
    "        self.x = self.encode(x)\n",
    "        self.y = torch.tensor(labels)\n",
    "        \n",
    "    def __getitem__(self, i):\n",
    "        return self.x[i], self.y[i]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "    \n",
    "    def make_w2i(self, tok_rows):\n",
    "        return {k: i for i, (k, v) in \n",
    "                enumerate(Counter(chain(*tok_rows)).most_common(), 1)}\n",
    "    \n",
    "    def encode(self, tok_rows):\n",
    "        enc = np.zeros((len(tok_rows), self.seq_len), dtype=int)\n",
    "        for i, row in enumerate(tok_rows):\n",
    "            trunc = [self.w2i.get(w, 0) for w in row[:self.seq_len]]\n",
    "            enc[i, :len(trunc)] = trunc\n",
    "        return torch.tensor(enc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sents = [\n",
    "    'I walked to the store so I hope it is not closed.',\n",
    "    'The theater is closed today and the sky is grey.',\n",
    "    'His dog is brown while hers is grey.'\n",
    "]\n",
    "labels = [0, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([13, 14,  1, 15, 16, 17,  3, 18,  1,  4]), tensor(1))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = Data(sents, labels, 10)\n",
    "ds[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "         [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "         [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]]), tensor([0, 1, 1]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dl = DataLoader(ds, batch_size=3)\n",
    "x, y = next(iter(dl))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "         [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "         [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]]), tensor([0, 1, 1]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x, y = next(iter(dl))\n",
    "x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0054, -0.0142, -0.0027, -0.0052],\n",
       "        [ 0.0097,  0.0003, -0.0098, -0.0082],\n",
       "        [ 0.0049,  0.0084,  0.0028, -0.0158],\n",
       "        [ 0.0035,  0.0056, -0.0106, -0.0135],\n",
       "        [ 0.0088, -0.0058,  0.0105, -0.0072],\n",
       "        [-0.0003, -0.0004, -0.0012,  0.0142],\n",
       "        [ 0.0089, -0.0114, -0.0001,  0.0037],\n",
       "        [ 0.0089,  0.0007,  0.0076,  0.0034],\n",
       "        [-0.0039,  0.0119,  0.0090, -0.0077],\n",
       "        [ 0.0040, -0.0154,  0.0072, -0.0024]], requires_grad=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "be = BloomEmbedding(11, 4)\n",
    "be.emb.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1],\n",
       "        [13, 14,  1, 15, 16, 17,  3, 18,  1,  4],\n",
       "        [19, 20,  1, 21, 22, 23,  1,  4,  0,  0]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 10, 4])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# (bs x seq_len) -> (bs -> seq_len -> emb_size)\n",
    "y = be(x)\n",
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0337, -0.0130, -0.0147, -0.0183],\n",
       "        [ 0.0075, -0.0316,  0.0120,  0.0236],\n",
       "        [ 0.0304, -0.0328,  0.0388, -0.0240],\n",
       "        [ 0.0144,  0.0004, -0.0022,  0.0019],\n",
       "        [ 0.0084,  0.0117,  0.0089, -0.0284],\n",
       "        [ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0177, -0.0087,  0.0344, -0.0139],\n",
       "        [ 0.0272, -0.0108, -0.0036,  0.0130],\n",
       "        [ 0.0035, -0.0162,  0.0048,  0.0260]], grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we show by step how to get from x to y. This is meant to demonstrate the basic mechanism, not to show how PyTorch actually implements this under the hood. Let's look at a single row of x, corresponding to 1 sentence where each word is mapped to its index in the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we hash each item."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[8, 2, 7, 8],\n",
       " [2, 8, 1, 2],\n",
       " [6, 6, 10, 10],\n",
       " [10, 5, 5, 5],\n",
       " [6, 9, 7, 2],\n",
       " [5, 9, 4, 0],\n",
       " [8, 2, 7, 8],\n",
       " [5, 10, 8, 9],\n",
       " [7, 8, 6, 2],\n",
       " [6, 10, 6, 0]]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hashed = [probabilistic_hash_item(i.item(), 11, int, 4) for i in x[0]]\n",
    "hashed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then use each row of hashed integers to index into the embedding weight matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 4, 4])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0089,  0.0007,  0.0076,  0.0034],\n",
       "         [ 0.0097,  0.0003, -0.0098, -0.0082],\n",
       "         [ 0.0089, -0.0114, -0.0001,  0.0037],\n",
       "         [ 0.0089,  0.0007,  0.0076,  0.0034]],\n",
       "\n",
       "        [[ 0.0097,  0.0003, -0.0098, -0.0082],\n",
       "         [ 0.0089,  0.0007,  0.0076,  0.0034],\n",
       "         [ 0.0054, -0.0142, -0.0027, -0.0052],\n",
       "         [ 0.0097,  0.0003, -0.0098, -0.0082]]], grad_fn=<SliceBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output = []\n",
    "for row in hashed:\n",
    "    row_out = be.emb.weight[row]\n",
    "    output.append(row_out)\n",
    "output = torch.stack(output)\n",
    "print(output.shape)\n",
    "output[:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we sum up the embedding rows. Above, each word is represented by four rows of the embedding matrix. After summing, we get a single vector for each word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0337, -0.0130, -0.0147, -0.0183],\n",
       "        [ 0.0075, -0.0316,  0.0120,  0.0236],\n",
       "        [ 0.0304, -0.0328,  0.0388, -0.0240],\n",
       "        [ 0.0144,  0.0004, -0.0022,  0.0019],\n",
       "        [ 0.0084,  0.0117,  0.0089, -0.0284],\n",
       "        [ 0.0363, -0.0098,  0.0052,  0.0023],\n",
       "        [ 0.0177, -0.0087,  0.0344, -0.0139],\n",
       "        [ 0.0272, -0.0108, -0.0036,  0.0130],\n",
       "        [ 0.0035, -0.0162,  0.0048,  0.0260]], grad_fn=<SumBackward1>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output = output.sum(-2)\n",
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " Notice that the values now match the output of our embedding layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert torch.isclose(output, y[0]).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TESTING - Axial Embedding\n",
    "\n",
    "TODO\n",
    "- port to torch\n",
    "- use better initialization scheme (maybe can use modified embedding layer from incendio)\n",
    "- idea: is there a way to combine this with bloom embeddings?\n",
    "- idea: what if we used more than 2 emb matrices? Could we reduce params further this way? And are there pros/cons to doing this? Seems like by having more matrices, we could use smaller segments for each one which might be beneficial. With two matrices, that means many words are sharing large segments of embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AxialEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_dim, emb_dim, pad_idx=None):\n",
    "        super().__init__()\n",
    "        self.v = self._decompose_mult(vocab_dim)\n",
    "        self.e = self._decompose_add(emb_dim)\n",
    "        self.emb = nn.ModuleList(PaddedEmbedding(self.v, self.e, pad_idx) \n",
    "                                 for _ in range(2))\n",
    "    \n",
    "    def _decompose_mult(self, dim, ratio):\n",
    "        return int(np.ceil(np.sqrt(dim)))\n",
    "    \n",
    "    def _decompose_add(self, dim, ratio):\n",
    "        return int(np.ceil(dim / 2))\n",
    "    \n",
    "    def forward(self, idx):\n",
    "        return torch.cat([self.emb[0](idx%self.v), self.emb[1](idx//self.v)], \n",
    "                         dim=-1)\n",
    "\n",
    "class MultiAxialEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_dim, emb_dim, n_blocks=1, pad_idx=None):\n",
    "        super().__init__()\n",
    "        self.n_blocks = n_blocks\n",
    "        self.v = self._decompose_mult(vocab_dim)\n",
    "        self.e = self._decompose_add(emb_dim)\n",
    "        self.emb = nn.ModuleList(PaddedEmbedding(self.v, self.e, pad_idx) \n",
    "                                  for _ in range(n_blocks))\n",
    "    \n",
    "    def _decompose_mult(self, dim):\n",
    "        return int(np.ceil(dim ** (1 / self.n_blocks)))\n",
    "    \n",
    "    def _decompose_add(self, dim):\n",
    "        return int(np.ceil(dim // self.n_blocks))\n",
    "    \n",
    "    def forward(self, idx):\n",
    "        res = [e(i%self.v) for i, e in zip(idx, self.emb)]\n",
    "        return torch.cat(res, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PreHashedAxialEncoder(AxialEncoding2):\n",
    "    \n",
    "    def __init__(self, vocab_dim, emb_dim, n_blocks=1, pad_idx=None):\n",
    "        super().__init__(vocab_dim, emb_dim, n_blocks, pad_idx)\n",
    "    \n",
    "    def forward(self, idx):\n",
    "        return torch.cat([e(idx.squeeze()) for e, idx in \n",
    "                          zip(self.emb, torch.chunk(xhash, xhash.shape[0], -1))],\n",
    "                         dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([28834, 11824, 24136, 14813, 13900, 13288, 14173,  6981, 17569, 29101,\n",
      "        25553, 15089])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AxialEncoding(\n",
       "  (emb): ModuleList(\n",
       "    (0): Embedding(122, 50)\n",
       "    (1): Embedding(122, 50)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "emb_dim = 100\n",
    "bs = 12\n",
    "\n",
    "ax = AxialEncoding(vocab_size, emb_dim, 2)\n",
    "x = torch.randint(0, vocab_size, (bs,))\n",
    "print(x)\n",
    "ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "index out of range: Tried to access index 236 out of table with 121 rows. at ../aten/src/TH/generic/THTensorEvenMoreMath.cpp:418",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-340-ba5c0491a78f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0max\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mres\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mres\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdetach\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mround\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    531\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 532\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    533\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-333-7cb3c2c9e65d>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, idx)\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0;31m#         return torch.cat([self.emb1(idx%self.v1), self.emb2(idx//self.v1)],\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0;31m#                          dim=-1)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m         return torch.cat([self.emb[0](idx%self.v1), self.emb[1](idx//self.v1)], \n\u001b[0m\u001b[1;32m     24\u001b[0m                          dim=-1)\n\u001b[1;32m     25\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m    530\u001b[0m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    531\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 532\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    533\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    534\u001b[0m             \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/torch/nn/modules/sparse.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    112\u001b[0m         return F.embedding(\n\u001b[1;32m    113\u001b[0m             \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpadding_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax_norm\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 114\u001b[0;31m             self.norm_type, self.scale_grad_by_freq, self.sparse)\n\u001b[0m\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    116\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mextra_repr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.pyenv/versions/3.7.4/envs/main/lib/python3.7/site-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36membedding\u001b[0;34m(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)\u001b[0m\n\u001b[1;32m   1482\u001b[0m         \u001b[0;31m# remove once script supports set_grad_enabled\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1483\u001b[0m         \u001b[0m_no_grad_embedding_renorm_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_norm\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnorm_type\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1484\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0membedding\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpadding_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mscale_grad_by_freq\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msparse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1485\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1486\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: index out of range: Tried to access index 236 out of table with 121 rows. at ../aten/src/TH/generic/THTensorEvenMoreMath.cpp:418"
     ]
    }
   ],
   "source": [
    "res = ax(x)\n",
    "print(res.shape)\n",
    "res.detach().numpy().round(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "18400"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "3 mil -> 18,400\n",
    "122*50+246*50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0114, -0.0137, -0.0056,  ..., -0.0089, -0.0061,  0.0051],\n",
       "        [-0.0114, -0.0137, -0.0056,  ..., -0.0032,  0.0157, -0.0045],\n",
       "        [-0.0078,  0.0031,  0.0065,  ..., -0.0135, -0.0141,  0.0131],\n",
       "        ...,\n",
       "        [-0.0114, -0.0137, -0.0056,  ..., -0.0105,  0.0051,  0.0001],\n",
       "        [-0.0114, -0.0137, -0.0056,  ..., -0.0035,  0.0041,  0.0061],\n",
       "        [-0.0114, -0.0137, -0.0056,  ..., -0.0032, -0.0003,  0.0062]],\n",
       "       grad_fn=<CatBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[29804, 14878],\n",
      "        [28725, 25398],\n",
      "        [ 7542,  1851],\n",
      "        [20846, 12400],\n",
      "        [ 1678, 19067],\n",
      "        [ 5784, 25606],\n",
      "        [  445, 19821],\n",
      "        [13784,  8734],\n",
      "        [  668, 10067],\n",
      "        [27693,  4930],\n",
      "        [11528, 13906],\n",
      "        [10509, 13890]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AxialEncoding2(\n",
       "  (emb): ModuleList(\n",
       "    (0): Embedding(14, 25)\n",
       "    (1): Embedding(14, 25)\n",
       "    (2): Embedding(14, 25)\n",
       "    (3): Embedding(14, 25)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "emb_dim = 100\n",
    "bs = 12\n",
    "\n",
    "ax = AxialEncoding2(vocab_size, emb_dim, 4)\n",
    "x = torch.randint(0, vocab_size, (bs, 2))\n",
    "print(x)\n",
    "ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([12, 2, 4])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0,  1, 12, 10],\n",
       "         [ 9,  8,  3,  0]],\n",
       "\n",
       "        [[ 1, 10, 12, 12],\n",
       "         [ 6, 10,  9,  0]],\n",
       "\n",
       "        [[ 8, 10,  2,  9],\n",
       "         [12,  4, 12,  7]],\n",
       "\n",
       "        [[ 5, 11,  1, 10],\n",
       "         [ 4, 13, 11, 12]],\n",
       "\n",
       "        [[ 4,  5, 12,  1],\n",
       "         [ 7,  1,  0,  2]],\n",
       "\n",
       "        [[ 9,  6,  9, 11],\n",
       "         [ 1,  8,  7,  5]],\n",
       "\n",
       "        [[ 8, 13,  0,  5],\n",
       "         [ 2,  6,  3,  1]],\n",
       "\n",
       "        [[ 6, 10,  2,  4],\n",
       "         [ 0,  9,  9,  7]],\n",
       "\n",
       "        [[ 1,  5,  1,  1],\n",
       "         [12, 11, 13, 12]],\n",
       "\n",
       "        [[ 5, 12,  2, 13],\n",
       "         [11, 12,  3,  3]],\n",
       "\n",
       "        [[11,  5,  3, 12],\n",
       "         [ 3,  2,  6,  0]],\n",
       "\n",
       "        [[12,  6, 13,  3],\n",
       "         [ 1,  6,  0, 13]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xhash = probabilistic_hash_tensor(x, 14, 4)\n",
    "print(xhash.shape)\n",
    "xhash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(tensor([[ 3,  7, 12, 13],\n",
       "          [ 2, 11,  5,  5],\n",
       "          [ 1, 10,  3, 10],\n",
       "          [ 1, 11,  6,  3],\n",
       "          [ 5, 10, 12,  3],\n",
       "          [ 6,  2,  8,  7],\n",
       "          [ 9,  1,  5,  9],\n",
       "          [11,  0,  8, 12],\n",
       "          [ 0,  8,  0,  0],\n",
       "          [ 8, 10,  4,  3],\n",
       "          [12,  7,  6,  5],\n",
       "          [ 1,  9,  0,  4]]), 'a')]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(zip(xhash, ['a', 'b', 'c', 'd']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 12, 1])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.split(xhash, xhash.shape[0], -1)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 12, 1])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.chunk(xhash, xhash.shape[1], -1)[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 7920, 26623],\n",
      "        [16646, 20293],\n",
      "        [17905, 14433],\n",
      "        [28334, 17367],\n",
      "        [ 2305, 23681],\n",
      "        [19342, 10013],\n",
      "        [ 5689,  7792],\n",
      "        [25094, 18898],\n",
      "        [10175,  4162],\n",
      "        [13035, 29493],\n",
      "        [13181, 20727],\n",
      "        [  387, 14492]])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PreHashedAxialEncoder(\n",
       "  (emb): ModuleList(\n",
       "    (0): Embedding(14, 25)\n",
       "    (1): Embedding(14, 25)\n",
       "    (2): Embedding(14, 25)\n",
       "    (3): Embedding(14, 25)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab_size = 30_000\n",
    "emb_dim = 100\n",
    "bs = 12\n",
    "\n",
    "ax = PreHashedAxialEncoder(vocab_size, emb_dim, 4)\n",
    "x = torch.randint(0, vocab_size, (bs, 2))\n",
    "print(x)\n",
    "ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([12, 2, 4])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[0, 6, 2, 0],\n",
       "         [4, 6, 5, 4]],\n",
       "\n",
       "        [[6, 7, 5, 3],\n",
       "         [6, 0, 3, 0]],\n",
       "\n",
       "        [[2, 7, 1, 7],\n",
       "         [2, 5, 3, 4]],\n",
       "\n",
       "        [[3, 6, 0, 7],\n",
       "         [7, 1, 4, 5]],\n",
       "\n",
       "        [[3, 4, 4, 6],\n",
       "         [6, 5, 2, 4]],\n",
       "\n",
       "        [[6, 4, 6, 2],\n",
       "         [3, 2, 2, 7]],\n",
       "\n",
       "        [[7, 1, 2, 0],\n",
       "         [7, 1, 5, 7]],\n",
       "\n",
       "        [[2, 2, 1, 1],\n",
       "         [2, 3, 0, 0]],\n",
       "\n",
       "        [[3, 7, 1, 3],\n",
       "         [7, 4, 0, 0]],\n",
       "\n",
       "        [[3, 0, 2, 3],\n",
       "         [6, 2, 3, 3]],\n",
       "\n",
       "        [[2, 1, 7, 7],\n",
       "         [5, 5, 3, 7]],\n",
       "\n",
       "        [[0, 6, 5, 6],\n",
       "         [1, 7, 1, 0]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xhash = probabilistic_hash_tensor(x, ax.v, len(ax.emb))\n",
    "print(xhash.shape)\n",
    "xhash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([12, 2, 8])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res = ax(xhash)\n",
    "res.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([12, 8])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res[:, 0, :].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0.005,  0.005, -0.006,  0.003, -0.009, -0.007,  0.006,  0.017],\n",
       "        [-0.004,  0.004, -0.006,  0.003, -0.005, -0.007, -0.   ,  0.003]],\n",
       "\n",
       "       [[ 0.008, -0.014,  0.   , -0.006, -0.005, -0.007,  0.012, -0.008],\n",
       "        [ 0.008, -0.014, -0.001, -0.003, -0.001,  0.018,  0.006,  0.017]],\n",
       "\n",
       "       [[-0.005,  0.006,  0.   , -0.006, -0.006, -0.012, -0.018,  0.005],\n",
       "        [-0.005,  0.006,  0.009, -0.004, -0.001,  0.018, -0.   ,  0.003]],\n",
       "\n",
       "       [[ 0.012,  0.018, -0.006,  0.003,  0.007, -0.009, -0.018,  0.005],\n",
       "        [-0.001, -0.003, -0.003, -0.004,  0.001, -0.002,  0.01 , -0.006]],\n",
       "\n",
       "       ...,\n",
       "\n",
       "       [[ 0.012,  0.018,  0.   , -0.006, -0.006, -0.012,  0.012, -0.008],\n",
       "        [-0.001, -0.003,  0.015,  0.01 ,  0.007, -0.009,  0.006,  0.017]],\n",
       "\n",
       "       [[ 0.012,  0.018, -0.001, -0.003, -0.009, -0.007,  0.012, -0.008],\n",
       "        [ 0.008, -0.014,  0.002, -0.006, -0.001,  0.018,  0.012, -0.008]],\n",
       "\n",
       "       [[-0.005,  0.006, -0.003, -0.004,  0.006, -0.006, -0.018,  0.005],\n",
       "        [-0.009, -0.003,  0.009, -0.004, -0.001,  0.018, -0.018,  0.005]],\n",
       "\n",
       "       [[ 0.005,  0.005, -0.006,  0.003, -0.005, -0.007,  0.013,  0.018],\n",
       "        [-0.006, -0.013,  0.   , -0.006, -0.006, -0.012,  0.006,  0.017]]], dtype=float32)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res.detach().numpy().round(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cat pros\n",
    "\n",
    "-smaller emb matrices\n",
    "\n",
    "# Cat cons\n",
    "\n",
    "-different items will share large blocks of emb vectors. Maybe this works if we have some pre-defined knowledge of what words should be indexed close together. Sorting by frequency does this somewhat, but maybe would work better to first cluster by w2vec, sort so similar words are close together. Maybe this is why this is good for positional encodings: adjacent time steps share some (but not all) of embs\n",
    "\n",
    "# Add pros\n",
    "\n",
    "-words won't necessarily end up sharing embs w/ other words even if most of their rows are the same - the 1 unshared row could learn to be very different\n",
    "\n",
    "# Add cons\n",
    "\n",
    "-larger emb matrices\n",
    "-each row serves multiple purposes - maybe \"pulled in multiple directions\" by diff words"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## END TESTING - Axial Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LinearSkipBlock(nn.Module):\n",
    "    \"\"\"This lets us easily create residual block equivalents with linear\n",
    "    layers.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, layer_dims, op, activation=mish):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        x_dim: int\n",
    "            Size of input tensor.\n",
    "        layer_dims: Iterable[int]\n",
    "            Size of each layer. The length of this list will be the skip size\n",
    "            (2 is probably a reasonable starting point).\n",
    "        op: function\n",
    "            This will be called on the input x and the processed x in the\n",
    "            forward method. This is a concatenation for dense blocks and an\n",
    "            addition for residual blocks, but any operation is possible.\n",
    "        activation: callable\n",
    "            Activation function or callable class. This will be applied after\n",
    "            each layer. The final activation is applied after the `op` function.\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.skip_size = len(layer_dims)\n",
    "        self.activation = activation\n",
    "        self.layers = nn.ModuleList([nn.Linear(d_in, d_out) for d_in, d_out\n",
    "                                     in zip([x_dim]+list(layer_dims), layer_dims)])\n",
    "        self.op = op\n",
    "            \n",
    "    def forward(self, x):\n",
    "        out = x\n",
    "        for i, layer in enumerate(self.layers, 1):\n",
    "            out = layer(out)\n",
    "            if i < self.skip_size: out = self.activation(out)\n",
    "        return self.activation(self.op(x, out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LinearResBlock(LinearSkipBlock):\n",
    "    \"\"\"Equivalent of ResNet block with linear layers.\"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, hidden_dims, activation=mish):\n",
    "        if hidden_dims[-1] != x_dim:\n",
    "            raise InvalidArgumentError(\n",
    "                'Last hidden dimension must match input dimension.'\n",
    "            )\n",
    "        super().__init__(x_dim, hidden_dims, add, activation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LinearDenseBlock(LinearSkipBlock):\n",
    "    \"\"\"Equivalent of DenseNet block with linear layers.\"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, hidden_dims, activation=mish):\n",
    "        super().__init__(x_dim, hidden_dims, concat, activation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class WeightedLinearResBlock(LinearSkipBlock):\n",
    "    \"\"\"Like a LinearResBlock but takes a weighted average of the input and output\n",
    "    rather than adding them. Addition gives them equal weight and we may want to \n",
    "    weight the output more heavily.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, x_dim, hidden_dims, weights=(.25, .75), activation=mish):\n",
    "        super().__init__(x_dim, hidden_dims, \n",
    "                         partial(weighted_avg, weights=list(weights)), activation)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
