challenges-for-true-pytorch-heroes-solutions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch Challenges\n",
    "\n",
    "These set of challenges are concerned about broadcasting, one of the key concepts when dealing with tensors.\n",
    "\n",
    "[Sasha Rush](https://twitter.com/srush_nlp) compiled a set of [16 Tensor mini-puzzles](https://github.com/srush/Tensor-Puzzles) that involve reasoning about broadcasting in a constrained setting: people are allowed to use only a single PyTorch function: `torch.arange`. Can you do it?\n",
    "\n",
    "Here, I've extended his list to 26 puzzles! \n",
    "\n",
    "**Rules**\n",
    "\n",
    "- Each puzzle needs to be solved in 1 line (<80 columns) of code.\n",
    "- You are allowed @, arithmetic, comparison, shape, any indexing (e.g. `a[:j], a[:, None], a[arange(10)]`), and previous puzzle functions.\n",
    "- To start off, we give you an implementation for the `torch.arange` function.\n",
    "\n",
    "**Anti-Rules**\n",
    "- Nothing else. No `.view, .sum, .take, .squeeze, .tensor`.\n",
    "- No cheating. Stackoverflow is great, but this is about first-principles.\n",
    "- Hint... these puzzles are mostly about [Broadcasting](https://pytorch.org/docs/master/notes/broadcasting.html). Make sure you understand this rule.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from spec import make_test, run_test, TT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### arange\n",
    "\n",
    "This is given for free! Think about it as a \"for-loop\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 2, 3, 4, 5])"
      ]
     },
     "execution_count": 170,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def arange(i: int):\n",
    "    return torch.arange(i)\n",
    "\n",
    "arange(6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### where"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0, -1,  2, -1])"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def where(q, a, b):\n",
    "    return q * a + (~q) * b\n",
    "\n",
    "where(arange(4) % 2 == 0, arange(4), -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 1, 1, 1])"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def ones(i: int):\n",
    "    return where(arange(i) >= 0, 1, 0)\n",
    "\n",
    "ones(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(6)"
      ]
     },
     "execution_count": 173,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def sum(a: torch.Tensor):\n",
    "    return ones(a.shape[0]) @ a\n",
    "\n",
    "sum(arange(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### outer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 0, 0],\n",
       "        [1, 1, 1],\n",
       "        [2, 2, 2],\n",
       "        [3, 3, 3]])"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def outer(a: torch.Tensor, b: torch.Tensor):\n",
    "    return a[:, None] * b[None, :]\n",
    "\n",
    "outer(arange(4), ones(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### diag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 2, 3])"
      ]
     },
     "execution_count": 175,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def diag(a: torch.Tensor):\n",
    "    return a[arange(a.shape[0]), arange(a.shape[0])]\n",
    "\n",
    "diag(outer(arange(4), ones(4)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### eye"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 0, 0, 0],\n",
       "        [0, 1, 0, 0],\n",
       "        [0, 0, 1, 0],\n",
       "        [0, 0, 0, 1]])"
      ]
     },
     "execution_count": 176,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def eye(j: int):\n",
    "    return (arange(j)[:, None] == arange(j)[None, :]) * 1\n",
    "\n",
    "eye(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### triu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 1, 1, 1],\n",
       "        [0, 1, 1, 1],\n",
       "        [0, 0, 1, 1],\n",
       "        [0, 0, 0, 1]])"
      ]
     },
     "execution_count": 177,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def triu(j: int):\n",
    "    return (arange(j)[:,None] <= arange(j))*1\n",
    "\n",
    "triu(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### cumsum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 3, 6])"
      ]
     },
     "execution_count": 178,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def cumsum(a: torch.Tensor):\n",
    "    return (outer(ones(a.shape[0]), a) @ triu(a.shape[0]))[0]\n",
    "\n",
    "cumsum(torch.arange(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 1, 1])"
      ]
     },
     "execution_count": 179,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def diff(a: torch.Tensor, i: int):\n",
    "    return a - a[where(arange(i) > 0, arange(i)-1, 0)] + (a*(arange(i) <= 0))\n",
    "\n",
    "diff(arange(4), 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### vstack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 1, 2, 3],\n",
       "        [1, 1, 1, 1]])"
      ]
     },
     "execution_count": 180,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def vstack(a: torch.Tensor, b: torch.Tensor):\n",
    "    return a * (1-arange(2)[:, None]) + b * arange(2)[:, None]\n",
    "\n",
    "vstack(arange(4), ones(4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### roll"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 2, 3, 0])"
      ]
     },
     "execution_count": 181,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def roll(a: torch.Tensor, i: int):\n",
    "    return a[(arange(i) + 1) * ((arange(i) + 1) < i)]\n",
    "\n",
    "roll(arange(4), 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### flip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([3, 2, 1, 0])"
      ]
     },
     "execution_count": 182,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def flip(a: torch.Tensor, i: int):\n",
    "    return a[i - arange(i) - 1]\n",
    "\n",
    "flip(arange(4), 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### compress"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 2, 0])"
      ]
     },
     "execution_count": 183,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def compress(g: torch.Tensor, v: torch.Tensor, i: int):\n",
    "    return sum(eye(i)[:sum(g*1)] * outer(v[g], ones(i)))\n",
    "\n",
    "compress(torch.tensor([False, True, True]), arange(3), 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### pad_to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 2, 0, 0])"
      ]
     },
     "execution_count": 184,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def pad_to(a: torch.Tensor, i: int, j: int):\n",
    "    return sum((arange(i)[:, None] == arange(j)[None, :]) * a[:, None])\n",
    "\n",
    "pad_to(arange(3), 3, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### sequence_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 185,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 1, 0],\n",
       "        [1, 1, 0],\n",
       "        [1, 0, 0],\n",
       "        [1, 1, 1]])"
      ]
     },
     "execution_count": 185,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def sequence_mask(values: torch.Tensor, length: torch.Tensor):\n",
    "    return values * (length[:, None] > arange(values.shape[-1])[None, :])\n",
    "\n",
    "sequence_mask(outer(ones(4), ones(3)), torch.tensor([2,2,1,3]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### bincount"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 3, 4, 2])"
      ]
     },
     "execution_count": 186,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def bincount(a: torch.Tensor, j: int):\n",
    "    return ones(len(a)) @ ((a[:, None] == arange(j)[None, :]) * 1)\n",
    "\n",
    "bincount(torch.tensor([2, 1, 3, 3, 1, 2, 2, 2, 1, 0]), 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### scatter_add"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([8, 7, 5, 4])"
      ]
     },
     "execution_count": 187,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def scatter_add(values: torch.Tensor, link: torch.Tensor, j: int):\n",
    "    return sum((link[:, None] == arange(j)[None, :]) * outer(values, ones(j)))\n",
    "\n",
    "scatter_add(torch.tensor([5,1,7,2,3,2,1,3]), torch.tensor([0,0,1,0,2,2,3,3]), 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### flatten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 189,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15])"
      ]
     },
     "execution_count": 189,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def flatten(a: torch.Tensor, i:int, j:int):\n",
    "    return a[outer(ones(i), ones(j)) == 1]\n",
    "\n",
    "flatten(arange(16).view(4, 4), 4, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### linspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0000, 0.1111, 0.2222, 0.3333, 0.4444, 0.5556, 0.6667, 0.7778, 0.8889,\n",
       "        1.0000])"
      ]
     },
     "execution_count": 190,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def linspace(i: float, j: float, n: int):\n",
    "    return i + (j - i) * arange(n) / max(1, (n - 1))\n",
    "\n",
    "linspace(0, 1, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### heaviside"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 1.0000, -2.6444,  0.0000])"
      ]
     },
     "execution_count": 191,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def heaviside(a: torch.Tensor, b: torch.Tensor):\n",
    "    return (a > 0) + (a == 0) * b\n",
    "\n",
    "heaviside(torch.tensor([1, 0, -2]), torch.randn(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### hstack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 192,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 1],\n",
       "        [1, 1],\n",
       "        [2, 1]])"
      ]
     },
     "execution_count": 192,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def hstack(a: torch.Tensor, b: torch.Tensor):\n",
    "    return a[:,None] * eye(2)[0] + b[:,None] * eye(2)[1]\n",
    "\n",
    "hstack(arange(3), ones(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### view (1d to 2d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 1],\n",
       "        [2, 3],\n",
       "        [4, 5]])"
      ]
     },
     "execution_count": 193,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def view(a: torch.Tensor, i: int, j: int):\n",
    "    return a[(j * arange(i)[:,None] + arange(j)[None]) % len(a)][:i, :j]\n",
    "\n",
    "view(arange(6), 3, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### repeat (1d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4])"
      ]
     },
     "execution_count": 194,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def repeat(a: torch.Tensor, d: int):\n",
    "    return (ones(d)[:, None] * a)[outer(ones(d), ones(len(a))) == 1]\n",
    "\n",
    "repeat(arange(5),  3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### repeat_interleave (1d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4])"
      ]
     },
     "execution_count": 195,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def repeat_interleave(a: torch.Tensor, d: int):\n",
    "    return (ones(d)[:, None] * a).T[outer(ones(len(a)), ones(d)) == 1]\n",
    "\n",
    "repeat_interleave(arange(5), 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([0, 1]),\n",
       " tensor([2, 3]),\n",
       " tensor([4, 5]),\n",
       " tensor([6, 7]),\n",
       " tensor([8, 9]),\n",
       " tensor([10, 11])]"
      ]
     },
     "execution_count": 198,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def chunk(a: torch.Tensor, c: int):\n",
    "    return list(view(a, c, len(a)//c))\n",
    "\n",
    "chunk(torch.arange(12), 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### nonzero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 0],\n",
       "        [1, 1],\n",
       "        [2, 2]])"
      ]
     },
     "execution_count": 200,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def nonzero(a: torch.Tensor, i: int, j: int):\n",
    "    return hstack(outer(arange(i),ones(j))[a!=0],outer(ones(i),arange(j))[a!=0])\n",
    "\n",
    "nonzero(eye(3), 3, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### bucketize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 3, 4])"
      ]
     },
     "execution_count": 201,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def bucketize(v: torch.Tensor, boundaries: torch.Tensor):\n",
    "    return sum((v[:,None] > boundaries[None, :]).T * 1)\n",
    "\n",
    "bucketize(torch.tensor([3, 6, 9]), torch.tensor([1, 3, 5, 7, 9]))"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "celltoolbar": "Raw Cell Format",
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}