PyTorchified Inside

inside.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inside class\n",
    "\n",
    "Encapsulating ideas from the blogpost into a single `Inside` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def start_end_idxs(sentence_length, span_length, device):\n",
    "    start = torch.arange(sentence_length - span_length + 1,\n",
    "                         device=device)\n",
    "    end = start + span_length - 1\n",
    "    return start, end\n",
    "\n",
    "def split_idxs(sentence_length, span_length, device):\n",
    "    start = torch.arange(sentence_length - span_length + 1,\n",
    "                         device=device)\n",
    "    end = start + span_length - 1\n",
    "    splits = torch.arange(span_length - 1, device=device)\n",
    "    split_idxs = start[:, None] + splits\n",
    "    return split_idxs\n",
    "\n",
    "class Inside(nn.Module):\n",
    "    def __init__(self, operator, hidden_size):\n",
    "        super(Inside, self).__init__()\n",
    "        self.op = operator\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "    def forward(self, X: torch.Tensor):\n",
    "        M = self.fill_table(X)\n",
    "        return M\n",
    "\n",
    "    def fill_table(self, X: torch.Tensor):\n",
    "        # Setup\n",
    "        sentence_length = X.size(0)\n",
    "        batch_size = X.size(1)\n",
    "        idxs = torch.arange(sentence_length,\n",
    "                            dtype=torch.long, device=X.device)\n",
    "        M = torch.zeros((sentence_length,\n",
    "                         sentence_length,\n",
    "                         batch_size, self.hidden_size))\n",
    "        # 1st level\n",
    "        M[idxs, idxs] = X\n",
    "        \n",
    "        # Fill the table\n",
    "        for span_length in range(2, sentence_length + 1):\n",
    "            start, end = start_end_idxs(sentence_length, span_length, \n",
    "                                        device=X.device)\n",
    "            split = split_idxs(sentence_length, span_length,\n",
    "                               device=X.device)\n",
    "            l_vals = M[start[:, None], split]\n",
    "            r_vals = M[split + 1, end[:, None]]\n",
    "            M[start, end] = self.op(l_vals, r_vals)\n",
    "        return M"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example `op` to be used within `Inside`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MeanOp(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MeanOp, self).__init__()\n",
    "    def forward(self, a, b):\n",
    "        return torch.mean(a + b, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 1, 1])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.,  1.,  3.,  6., 10., 15., 21., 28., 36., 45.],\n",
       "        [ 0.,  1.,  3.,  6., 10., 15., 21., 28., 36., 45.],\n",
       "        [ 0.,  0.,  2.,  5.,  9., 14., 20., 27., 35., 44.],\n",
       "        [ 0.,  0.,  0.,  3.,  7., 12., 18., 25., 33., 42.],\n",
       "        [ 0.,  0.,  0.,  0.,  4.,  9., 15., 22., 30., 39.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  5., 11., 18., 26., 35.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  6., 13., 21., 30.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  7., 15., 24.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  8., 17.],\n",
       "        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  9.]])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mean_inside = Inside(MeanOp(), 1)\n",
    "X = torch.arange(10, dtype=torch.float)[:, None, None]\n",
    "M = mean_inside(X)\n",
    "M[:, :, 0, 0]"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}