solution_and_note.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from primitive import arange, where, tensor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [ones](https://numpy.org/doc/stable/reference/generated/numpy.ones.html)\n",
    "To think in tensor, replace traditional iterative programming primitives with basic tensor ops. \n",
    "- `arange` and `broadcast` to replace `for loop`. \n",
    "- `where` to replace `if-else`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1., 1., 1., 1., 1.])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.ones(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v1\n",
    "# def ones(n: int):\n",
    "#     return arange(n) * 0 + 1\n",
    "\n",
    "# v2\n",
    "def ones(n: int):\n",
    "    return where(arange(n) == 1, 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 1, 1, 1, 1])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ones(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "assert (ones(5) == np.ones(5)).all(), f\"Mismatch: {ones(5)} != {np.ones(5)}\"\n",
    "assert (ones(0) == np.ones(0)).all(), f\"Mismatch: {ones(0)} != {np.ones(0)}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [sum](https://numpy.org/doc/stable/reference/generated/numpy.sum.html)\n",
    "Reduce along a dim, think about `@`.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0, 1, 2, 3, 4]), 10)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.arange(5), np.sum(np.arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v1\n",
    "def sum(a_i):\n",
    "    i = a_i.shape[0]\n",
    "    return a_i @ ones(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "assert sum(np.array([1, 2, 3, 4])) == 10, f\"Sum mismatch: {sum(np.array([1, 2, 3, 4]))} != 10\"\n",
    "assert sum(np.array([])) == 0, f\"Sum mismatch: {sum(np.array([]))} != 0\"\n",
    "assert sum(np.array([-1, 1])) == 0, f\"Sum mismatch: {sum(np.array([-1, 1]))} != 0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [outer](https://numpy.org/doc/stable/reference/generated/numpy.outer.html)\n",
    "`a[:, None]`'s shape is `(i, 1)`. It transform row vector to column vector, introducing one dummy dimension. It is used frequently in the series, usually with `broadcast`. Get used to this pattern. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0., 0., 0., 0.],\n",
       "       [1., 1., 1., 1., 1.],\n",
       "       [2., 2., 2., 2., 2.]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.outer(np.arange(3), np.ones(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def outer(a_i, b_j):\n",
    "    return a_i[:, None] * b_j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 0, 0, 0, 0],\n",
       "       [1, 1, 1, 1, 1],\n",
       "       [2, 2, 2, 2, 2]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outer(arange(3), ones(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "assert np.array_equal(outer(np.array([1, 2]), np.array([1, 2])), np.outer(np.array([1, 2]), np.array([1, 2]))), \"Outer product mismatch\"\n",
    "assert np.array_equal(outer(np.array([0, 1]), np.array([0, 1])), np.outer(np.array([0, 1]), np.array([0, 1]))), \"Outer product mismatch\"\n",
    "assert np.array_equal(outer(np.array([-1, 1]), np.array([-1, 1])), np.outer(np.array([-1, 1]), np.array([-1, 1]))), \"Outer product mismatch\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [diag](https://numpy.org/doc/stable/reference/generated/numpy.diag.html)\n",
    "Think about indexing for retrival puzzle. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[0, 1, 2],\n",
       "        [3, 4, 5],\n",
       "        [6, 7, 8]]),\n",
       " array([0, 4, 8]))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "a_33 = np.arange(9).reshape((3, 3))\n",
    "a_33, np.diag(a_33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diag(a_ii):  # input has to be square matrix\n",
    "    i = a_ii.shape[0]\n",
    "    return a_ii[arange(i), arange(i)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 4, 8])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "diag(a_33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "assert np.array_equal(diag(np.array([[1, 2], [3, 4]])), np.diag(np.array([[1, 2], [3, 4]]))), \"Diagonal extraction mismatch\"\n",
    "assert np.array_equal(diag(np.array([[0, 1], [1, 0]])), np.diag(np.array([[0, 1], [1, 0]]))), \"Diagonal extraction mismatch\"\n",
    "assert np.array_equal(diag(np.array([[-1, -2], [-3, -4]])), np.diag(np.array([[-1, -2], [-3, -4]]))), \"Diagonal extraction mismatch\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [eye](https://numpy.org/doc/stable/reference/generated/numpy.eye.html)\n",
    "Semantically, it's similar to `outer`. Add a dim plus simple op to expand to 2d, then use `where` to cast `bool` to `int` for final result. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 0., 0.],\n",
       "       [0., 1., 0.],\n",
       "       [0., 0., 1.]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.eye(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eye(n: int):\n",
    "    return where(arange(n)[:, None] == arange(n), 1, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 0, 0],\n",
       "       [0, 1, 0],\n",
       "       [0, 0, 1]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eye(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "assert np.array_equal(eye(3), np.eye(3)), \"Identity matrix generation mismatch\"\n",
    "assert np.array_equal(eye(5), np.eye(5)), \"Identity matrix generation mismatch\"\n",
    "assert np.array_equal(eye(1), np.eye(1)), \"Identity matrix generation mismatch\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [triu](https://numpy.org/doc/stable/reference/generated/numpy.triu.html)\n",
    "It is the key to `attention` mechanism. Semantically the same trick of `eye`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[0, 1, 2],\n",
       "        [3, 4, 5],\n",
       "        [6, 7, 8]]),\n",
       " array([[0, 1, 2],\n",
       "        [0, 4, 5],\n",
       "        [0, 0, 8]]))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "a_33 = np.arange(9).reshape((3, 3))\n",
    "a_33, np.triu(a_33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "def triu(a_ii):\n",
    "    i = a_ii.shape[0]\n",
    "    cond = arange(i)[:, None] <= arange(i)\n",
    "    return where(cond, a_ii, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 2],\n",
       "       [0, 4, 5],\n",
       "       [0, 0, 8]])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "triu(a_33)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "assert np.array_equal(triu(np.arange(9).reshape((3, 3))), np.triu(np.arange(9).reshape((3, 3)))), \"Upper triangle matrix generation mismatch\"\n",
    "assert np.array_equal(triu(np.arange(4).reshape((2, 2))), np.triu(np.arange(4).reshape((2, 2)))), \"Upper triangle matrix generation mismatch\"\n",
    "assert np.array_equal(triu(np.zeros((3, 3))), np.triu(np.zeros((3, 3)))), \"Upper triangle matrix generation mismatch\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [cumsum](https://numpy.org/doc/stable/reference/generated/numpy.cumsum.html)\n",
    "- Again, key to attention and why autoregressive transformer training is so efficient and parallelizable. \n",
    "- V1 is direct. V2 is what really happen in the transformer attention module. \n",
    "\n",
    "Naive pseudo code of attention: \n",
    "```python\n",
    "def attention(q_btc, k_btc, v_btc):\n",
    "    attn = softmax(q @ rearrange(k, 'b t c -> b c t'))  # (b, t, c) @ (b, c, t) -> (b, t, t)\n",
    "    attn_mask = where(tril, attn, float('-inf'))\n",
    "    attn += attn_mask\n",
    "    return attn @ v  # (b, t, t) @ (b, t, c) -> (b, t, c)\n",
    "```\n",
    "`tril` here because each token in `t` dim has `c` channels. Attention is applying autoregressive weighted sum per channel. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.cumsum(np.ones(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v1, vector @ matrix\n",
    "# def cumsum(a_i):\n",
    "#     i = a_i.shape[0]\n",
    "#     return a_i @ triu(ones(i)[:, None] * ones(i))\n",
    "\n",
    "# v2, matrix @ vector\n",
    "def cumsum(a_i):\n",
    "    i = a_i.shape[0]\n",
    "    tril = where(arange(i)[:, None] >= arange(i), 1, 0)\n",
    "    return tril @ a_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cumsum(ones(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "assert np.array_equal(cumsum(np.arange(10)), np.cumsum(np.arange(10))), \"Cumulative sum mismatch\"\n",
    "assert np.array_equal(cumsum(np.zeros(10)), np.cumsum(np.zeros(10))), \"Cumulative sum mismatch for zeros\"\n",
    "assert np.array_equal(cumsum(np.ones(10)), np.cumsum(np.ones(10))), \"Cumulative sum mismatch for ones\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [diff](https://numpy.org/doc/stable/reference/generated/numpy.diff.htm)\n",
    "This off by one index is how to construct label for autorepressive training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 1, 1, 1])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.diff(np.arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "def diff(a_i):\n",
    "    i = a_i.shape[0]\n",
    "    return a_i[arange(i)[1:]] - a_i[arange(i)[:-1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 1, 1, 1])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "diff(arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "random_array = np.random.randint(0, 10, size=5)\n",
    "assert np.array_equal(diff(random_array), np.diff(random_array)), \"Difference mismatch for random array\"\n",
    "assert np.array_equal(diff(np.zeros(5)), np.diff(np.zeros(5))), \"Difference mismatch for zeros\"\n",
    "assert np.array_equal(diff(np.ones(5)), np.diff(np.ones(5))), \"Difference mismatch for ones\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [vstack](https://numpy.org/doc/stable/reference/generated/numpy.vstack.html)\n",
    "Don't have to fully develop the condition, think about broadcast. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 1., 1., 1., 1.],\n",
       "       [0., 1., 2., 3., 4.]])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.vstack((np.ones(5), np.arange(5)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vstack(a_i, b_i):\n",
    "    return where((arange(2) == 0)[:, None], a_i, b_i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 1, 1, 1, 1],\n",
       "       [0, 1, 2, 3, 4]])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vstack(ones(5), arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "length1 = np.random.randint(1, 10)\n",
    "assert np.array_equal(vstack(ones(length1), arange(length1)), np.vstack((ones(length1), arange(length1)))), \"vstack mismatch for ones and arange with random length\"\n",
    "\n",
    "length2 = np.random.randint(1, 10)\n",
    "assert np.array_equal(vstack(arange(length2), arange(length2)), np.vstack((arange(length2), arange(length2)))), \"vstack mismatch for two aranges with same random length\"\n",
    "\n",
    "length3 = np.random.randint(1, 10)\n",
    "assert np.array_equal(vstack(arange(length3), ones(length3)), np.vstack((arange(length3), ones(length3)))), \"vstack mismatch for arange and ones with random length\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [roll](https://numpy.org/doc/stable/reference/generated/numpy.roll.html)\n",
    "- The first implementation used direct `where` to deal with index i-1, which works but it's iterative programming thinking. \n",
    "- Whenever you want to deal with special case, the tranditonal if-else kicks in, think again in tensor form.\n",
    "- Remember the goal of these exercises is not to implement functions, but to twist the brain enough to learn the pattern of tensor ops. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 2, 3, 4, 0])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.roll(np.arange(5), shift=-1)  # negative, shift left"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v1\n",
    "# def roll(a_i):\n",
    "#     i = a_i.shape[0]\n",
    "#     return a_i[where(arange(i) == i - 1, 0, arange(i) + 1)]\n",
    "\n",
    "\n",
    "# v2\n",
    "def roll(a_i):\n",
    "    i = a_i.shape[0]\n",
    "    return a_i[(arange(i) + 1) % i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 2, 3, 4, 0])"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "roll(arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "# fmt: off\n",
    "test_arrays = [\n",
    "    (np.random.randint(1, 10), arange, \"arange with random length\"),\n",
    "    (np.random.randint(1, 10), lambda x: np.random.randint(0, 100, size=x), \"random array with random length\"),\n",
    "    (10, lambda x: np.linspace(0, 1, x), \"linspace array with fixed length\")\n",
    "]\n",
    "\n",
    "for length, array_func, description in test_arrays:\n",
    "    test_array = array_func(length)\n",
    "    assert np.array_equal(roll(test_array), np.roll(test_array, shift=-1)), f\"roll mismatch for {description}: {length}, roll: {roll(test_array)}, np.roll: {np.roll(test_array, shift=1)}\"\n",
    "# fmt: on"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [flip](https://numpy.org/doc/stable/reference/generated/numpy.flip.html)\n",
    "For each position, original index and the flipped index sums to `i - 1`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4, 3, 2, 1, 0])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.flip(np.arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flip(a_i):\n",
    "    i = a_i.shape[0]\n",
    "    return a_i[i - 1 - arange(i)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4, 3, 2, 1, 0])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flip(arange(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "# fmt: off\n",
    "test_arrays_flip = [\n",
    "    (arange(5), \"arange with fixed length\"),\n",
    "    (np.random.randint(0, 100, size=10), \"random array with fixed length\"),\n",
    "    (np.array([0, -1, -2, -3, 0, 1, 2, 3, 0]), \"array with negative numbers and zeros\")\n",
    "]\n",
    "\n",
    "for test_array, description in test_arrays_flip:\n",
    "    assert np.array_equal(flip(test_array), np.flip(test_array)), f\"flip mismatch for {description}: {test_array}, flip: {flip(test_array)}, np.flip: {np.flip(test_array)}\"\n",
    "# fmt: on"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [compress](https://numpy.org/doc/stable/reference/generated/numpy.compress.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 2])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.compress(condition=tensor([True, False, True, False]), a=np.arange(4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compress(condition_i, a_i):\n",
    "    return a_i[condition_i]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 2])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compress(tensor([True, False, True, False]), arange(4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "# Test 1: Alternating condition\n",
    "condition_1 = tensor([True, False, True, False])\n",
    "test_array_1 = arange(4)\n",
    "assert np.array_equal(compress(condition_1, test_array_1), np.compress(condition_1, test_array_1)), \"Test 1 failed: compress mismatch with np.compress for alternating condition\"\n",
    "\n",
    "# Test 2: All true except last\n",
    "condition_2 = tensor([True, True, True, False])\n",
    "test_array_2 = arange(4)\n",
    "assert np.array_equal(compress(condition_2, test_array_2), np.compress(condition_2, test_array_2)), \"Test 2 failed: compress mismatch with np.compress for all true except last\"\n",
    "\n",
    "# Test 3: All false except last two\n",
    "condition_3 = tensor([False, False, True, True])\n",
    "test_array_3 = arange(4)\n",
    "assert np.array_equal(compress(condition_3, test_array_3), np.compress(condition_3, test_array_3)), \"Test 3 failed: compress mismatch with np.compress for all false except last two\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## pad to\n",
    "- Pad `[0, 1, 2]` to len 8 -> `[0, 1, 2, 0, 0, 0, 0, 0]`\n",
    "- Think about dimension change as [linear transformations](https://www.youtube.com/watch?v=kYB8IZa5AuE&list=PLZHQObOWTQDPD3MizzM2xVFitgF8hE_ab&index=3).\n",
    "- Padding is applying a matrix to the source vector. \n",
    "- Padding is adding dimensions, new dim's basis vectors are all 0s. \n",
    "- Projection from low dim to high dim in this way hsa no info loss or distortion. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "# v1, vector @ matrix, easily interpretible wrt basis vectors\n",
    "def pad_to(a_i, j: int):\n",
    "    \"\"\"pad to len j, where j >= i\"\"\"\n",
    "    i = a_i.shape[0]\n",
    "    basis_vector_mapping = where(arange(i)[:, None] == arange(j), 1, 0)\n",
    "    return a_i @ basis_vector_mapping\n",
    "\n",
    "# v2, matrix @ vector, this form is bmm friendly\n",
    "# https://pytorch.org/docs/stable/generated/torch.bmm.html\n",
    "# def pad_to(a_i, j: int):\n",
    "#     i = a_i.shape[0]\n",
    "#     return where(arange(j)[:, None] == arange(i), 1, 0) @ a_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 2, 0, 0, 0, 0, 0])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pad_to(arange(3), 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 1, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 1, 0, 0, 0, 0, 0]])"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# see that xyz basis vector, first 3 columns, didn't change at all. \n",
    "# basis vectors of padded dims are all zero, since there is no info mapping from source dim to padded dim. \n",
    "# padded dimensions are dummy dimensions, which have no entropy. \n",
    "where(arange(3)[:, None] == arange(8), 1, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "padded_array = pad_to(arange(3), 8)\n",
    "assert padded_array.shape[0] == 8, f\"Length mismatch for pad to len 8: expected 8, got {padded_array.shape[0]}\"\n",
    "assert np.array_equal(padded_array[:3], arange(3)), \"Original elements altered in pad to len 8\"\n",
    "assert np.all(padded_array[3:] == 0), \"Padding not zeros in pad to len 8\"\n",
    "\n",
    "padded_array = pad_to(arange(5), 10)\n",
    "assert padded_array.shape[0] == 10, f\"Length mismatch for pad to len 10: expected 10, got {padded_array.shape[0]}\"\n",
    "assert np.array_equal(padded_array[:5], arange(5)), \"Original elements altered in pad to len 10\"\n",
    "assert np.all(padded_array[5:] == 0), \"Padding not zeros in pad to len 10\"\n",
    "\n",
    "padded_array = pad_to(np.zeros(4), 6)\n",
    "assert padded_array.shape[0] == 6, f\"Length mismatch for pad zeros to len 6: expected 6, got {padded_array.shape[0]}\"\n",
    "assert np.array_equal(padded_array[:4], np.zeros(4)), \"Original elements altered in pad zeros to len 6\"\n",
    "assert np.all(padded_array[4:] == 0), \"Padding not zeros in pad zeros to len 6\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [sequence mask](https://www.tensorflow.org/api_docs/python/tf/sequence_mask)\n",
    "\n",
    "```python\n",
    "tf.sequence_mask(\n",
    "    lengths,\n",
    "    maxlen=None,\n",
    ")\n",
    "\n",
    "tf.sequence_mask([1, 3, 2], 5)  # [[True, False, False, False, False],\n",
    "                                #  [True, True, True, False, False],\n",
    "                                #  [True, True, False, False, False]]\n",
    "```\n",
    "\n",
    "`sequence_mask` and `pad_to` are useful primitives to deal with `padding` and `packing`. Check out the [BlockDiagonalMask](https://facebookresearch.github.io/xformers/components/ops.html#module-xformers.ops) in `xformers`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "def seq_mask(len_B: list[int], maxlen: int):\n",
    "    \"\"\"assume maxlen >= len of longest seq in the len_B\n",
    "    B: batch size\n",
    "    maxlen: pad the batch to this len\n",
    "    \"\"\"\n",
    "    return len_B[:, None] > arange(maxlen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ True, False, False, False, False],\n",
       "       [ True,  True,  True, False, False],\n",
       "       [ True,  True, False, False, False]])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seq_mask(tensor([1, 3, 2]), 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "result_1 = seq_mask(tensor([1, 3, 2]), 5)\n",
    "expected_1 = np.array([[True, False, False, False, False],\n",
    "                        [True, True, True, False, False],\n",
    "                        [True, True, False, False, False]])\n",
    "assert np.array_equal(result_1, expected_1), f\"Test 1 failed: expected {expected_1}, got {result_1}\"\n",
    "\n",
    "result_2 = seq_mask(tensor([0, 4, 1]), 4)\n",
    "expected_2 = np.array([[False, False, False, False],\n",
    "                        [True, True, True, True],\n",
    "                        [True, False, False, False]])\n",
    "assert np.array_equal(result_2, expected_2), f\"Test 2 failed: expected {expected_2}, got {result_2}\"\n",
    "\n",
    "result_3 = seq_mask(tensor([5]), 5)\n",
    "expected_3 = np.array([[True, True, True, True, True]])\n",
    "assert np.array_equal(result_3, expected_3), f\"Test 3 failed: expected {expected_3}, got {result_3}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [bincount](https://numpy.org/doc/stable/reference/generated/numpy.bincount.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 3, 1, 1, 0, 0, 0, 1])"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.bincount(np.array([0, 1, 1, 3, 2, 1, 7]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bincount(a_i):\n",
    "    max = a_i.max()\n",
    "    i = a_i.shape[0]\n",
    "\n",
    "    counting = arange(max + 1)[:, None] == a_i  # (max, i)\n",
    "    reduce = where(counting, 1, 0) @ ones(i)\n",
    "\n",
    "    return reduce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 3, 1, 1, 0, 0, 0, 1])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bincount(tensor([0, 1, 1, 3, 2, 1, 7]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "my_result_1 = bincount(tensor([0, 1, 1, 3, 2, 1, 7]))\n",
    "np_result_1 = np.bincount(np.array([0, 1, 1, 3, 2, 1, 7]))\n",
    "assert np.array_equal(my_result_1, np_result_1), f\"Test 1 failed: expected {np_result_1}, got {my_result_1}\"\n",
    "\n",
    "my_result_2 = bincount(tensor([2, 3, 3, 2, 4]))\n",
    "np_result_2 = np.bincount(np.array([2, 3, 3, 2, 4]))\n",
    "assert np.array_equal(my_result_2, np_result_2), f\"Test 2 failed: expected {np_result_2}, got {my_result_2}\"\n",
    "\n",
    "my_result_3 = bincount(tensor([0, 0, 0]))\n",
    "np_result_3 = np.bincount(np.array([0, 0, 0]))\n",
    "assert np.array_equal(my_result_3, np_result_3), f\"Test 3 failed: expected {np_result_3}, got {my_result_3}\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [scatter add](https://pytorch-scatter.readthedocs.io/en/1.3.0/functions/add.html)\n",
    "![](https://raw.githubusercontent.com/rusty1s/pytorch_scatter/master/docs/source/_figures/add.svg?sanitize=true)\n",
    "\n",
    "The idea is related to `bincount` and `sequence mask`. Expand and reduce."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scatter_add(a_i, index_i):\n",
    "    n = index_i.max() + 1  # number of attributions\n",
    "\n",
    "    expand = arange(n)[:, None] == index_i  # (n, i)\n",
    "    reduce = where(expand, 1, 0) @ a_i  # (n, i) @ (i,) = (n,)\n",
    "\n",
    "    return reduce\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([8, 7, 5, 4])"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = tensor([5, 1, 7, 2, 3, 2, 1, 3])\n",
    "index = tensor([0, 0, 1, 0, 2, 2, 3, 3])\n",
    "\n",
    "scatter_add(a, index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "result_1 = scatter_add(tensor([5, 1, 7, 2, 3, 2, 1, 3]), tensor([0, 0, 1, 0, 2, 2, 3, 3]))\n",
    "expected_1 = np.array([8, 7, 5, 4])\n",
    "assert np.array_equal(result_1, expected_1), f\"Test 1 failed: expected {expected_1}, got {result_1}\"\n",
    "\n",
    "result_2 = scatter_add(tensor([2, 4, 1, 3]), tensor([1, 1, 0, 1]))\n",
    "expected_2 = np.array([1, 9])\n",
    "assert np.array_equal(result_2, expected_2), f\"Test 2 failed: expected {expected_2}, got {result_2}\"\n",
    "\n",
    "result_3 = scatter_add(tensor([0, 0, 0, 0]), tensor([0, 1, 2, 3]))\n",
    "expected_3 = np.array([0, 0, 0, 0])\n",
    "assert np.array_equal(result_3, expected_3), f\"Test 3 failed: expected {expected_3}, got {result_3}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [flattern](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.flatten.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[ 0,  1,  2,  3],\n",
       "        [ 4,  5,  6,  7],\n",
       "        [ 8,  9, 10, 11]]),\n",
       " array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11]))"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "a = np.arange(12).reshape((3, 4))\n",
    "a, a.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flatten(a_ij):\n",
    "    i, j = a_ij.shape\n",
    "    return a_ij[arange(i * j) // j, arange(i * j) % j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flatten(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "a_test_1 = np.array([[1, 2], [3, 4]])\n",
    "expected_1 = a_test_1.flatten()\n",
    "result_1 = flatten(a_test_1)\n",
    "assert np.array_equal(result_1, expected_1), f\"Test 1 failed: expected {expected_1}, got {result_1}\"\n",
    "\n",
    "a_test_2 = np.array([[0, 0, 0], [0, 0, 0]])\n",
    "expected_2 = a_test_2.flatten()\n",
    "result_2 = flatten(a_test_2)\n",
    "assert np.array_equal(result_2, expected_2), f\"Test 2 failed: expected {expected_2}, got {result_2}\"\n",
    "\n",
    "a_test_3 = np.array([[9]])\n",
    "expected_3 = a_test_3.flatten()\n",
    "result_3 = flatten(a_test_3)\n",
    "assert np.array_equal(result_3, expected_3), f\"Test 3 failed: expected {expected_3}, got {result_3}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [linspace](https://numpy.org/doc/stable/reference/generated/numpy.linspace.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.  , 2.25, 2.5 , 2.75, 3.  ])"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.linspace(2, 3, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linspace(start, stop, n: int):\n",
    "    \"\"\"inclusive for both start and stop\"\"\"\n",
    "    step = (stop - start) / (n - 1)\n",
    "    return start + step * arange(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.  , 2.25, 2.5 , 2.75, 3.  ])"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "linspace(2, 3, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "test_1_start, test_1_stop, test_1_n = -10, 0, 5\n",
    "expected_1 = np.linspace(test_1_start, test_1_stop, test_1_n)\n",
    "result_1 = linspace(test_1_start, test_1_stop, test_1_n)\n",
    "assert np.allclose(result_1, expected_1), f\"Test 1 failed: expected {expected_1}, got {result_1}\"\n",
    "\n",
    "test_2_start, test_2_stop, test_2_n = 0, 100, 20\n",
    "expected_2 = np.linspace(test_2_start, test_2_stop, test_2_n)\n",
    "result_2 = linspace(test_2_start, test_2_stop, test_2_n)\n",
    "assert np.allclose(result_2, expected_2), f\"Test 2 failed: expected {expected_2}, got {result_2}\"\n",
    "\n",
    "test_3_start, test_3_stop, test_3_n = -50, 50, 2\n",
    "expected_3 = np.linspace(test_3_start, test_3_stop, test_3_n)\n",
    "result_3 = linspace(test_3_start, test_3_stop, test_3_n)\n",
    "assert np.allclose(result_3, expected_3), f\"Test 3 failed: expected {expected_3}, got {result_3}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [heaviside](https://numpy.org/doc/stable/reference/generated/numpy.heaviside.html)\n",
    "```python\n",
    "                      0   if x1 < 0\n",
    "heaviside(x1, x2) =  x2   if x1 == 0\n",
    "                      1   if x1 > 0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0., 1., 1.])"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.heaviside([-1.5, 0, 2.0], 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "def heaviside(a_i, b):\n",
    "    return where(a_i == 0, b, where(a_i > 0, 1, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 1])"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "heaviside(tensor([-1.5, 0, 2.0]), 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "test_1_input = np.array([-2.5, 0, 3.5])\n",
    "expected_1 = np.heaviside(test_1_input, 0.5)\n",
    "result_1 = heaviside(test_1_input, 0.5)\n",
    "assert np.allclose(result_1, expected_1), f\"Test 1 failed: expected {expected_1}, got {result_1}\"\n",
    "\n",
    "test_2_input = np.array([-100, -50, 0, 50, 100])\n",
    "expected_2 = np.heaviside(test_2_input, 0.8)\n",
    "result_2 = heaviside(test_2_input, 0.8)\n",
    "assert np.allclose(result_2, expected_2), f\"Test 2 failed: expected {expected_2}, got {result_2}\"\n",
    "\n",
    "test_3_input = np.array([0])\n",
    "expected_3 = np.heaviside(test_3_input, 0)\n",
    "result_3 = heaviside(test_3_input, 0)\n",
    "assert np.allclose(result_3, expected_3), f\"Test 3 failed: expected {expected_3}, got {result_3}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [repeat(1d)](https://numpy.org/doc/stable/reference/generated/numpy.tile.html)\n",
    "\n",
    "The effect of `np.tile`, but only expand on `dim=0`. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 2],\n",
       "       [0, 1, 2],\n",
       "       [0, 1, 2],\n",
       "       [0, 1, 2],\n",
       "       [0, 1, 2]])"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# example\n",
    "np.tile(np.arange(3), (5, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "def repeat(a_i, n: int):\n",
    "    return ones(n)[:, None] * a_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 2],\n",
       "       [0, 1, 2],\n",
       "       [0, 1, 2],\n",
       "       [0, 1, 2],\n",
       "       [0, 1, 2]])"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "repeat(arange(3), 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "test_1_input = arange(4)\n",
    "n_1 = 3\n",
    "expected_1 = np.tile(test_1_input, (n_1, 1))\n",
    "result_1 = repeat(test_1_input, n_1)\n",
    "assert np.array_equal(result_1, expected_1), f\"Test 1 failed: expected {expected_1}, got {result_1}\"\n",
    "\n",
    "test_2_input = tensor([5])\n",
    "n_2 = 10\n",
    "expected_2 = np.tile(test_2_input, (n_2, 1))\n",
    "result_2 = repeat(test_2_input, n_2)\n",
    "assert np.array_equal(result_2, expected_2), f\"Test 2 failed: expected {expected_2}, got {result_2}\"\n",
    "\n",
    "test_3_input = linspace(0, 1, 5)\n",
    "n_3 = 2\n",
    "expected_3 = np.tile(test_3_input, (n_3, 1))\n",
    "result_3 = repeat(test_3_input, n_3)\n",
    "assert np.array_equal(result_3, expected_3), f\"Test 3 failed: expected {expected_3}, got {result_3}\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [bucketize](https://pytorch.org/docs/stable/generated/torch.bucketize.html)\n",
    "\n",
    "Returns the **indices** of the buckets to which each value in the input belongs.\n",
    "\n",
    "```python\n",
    "a = torch.tensor([3, 6, 9])\n",
    "boundary = torch.tensor([1, 3, 5, 7, 9])\n",
    "\n",
    "torch.bucketize(a, boundary) -> tensor([1, 3, 4])\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bucketize(a_i, boundary_j):\n",
    "    \"\"\"boundary has to be in monotonically increasing order.\"\"\"\n",
    "    j = boundary_j.shape[0]\n",
    "    return where(a_i[:, None] > boundary_j, 1, 0) @ ones(j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1, 3, 4])"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = tensor([3, 6, 9])\n",
    "boundary = tensor([1, 3, 5, 7, 9])\n",
    "\n",
    "bucketize(a, boundary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "test_1_input = tensor([0, 5, 10])\n",
    "boundary_1 = tensor([0, 3, 6, 9])\n",
    "expected_1 = tensor([0, 2, 4])\n",
    "result_1 = bucketize(test_1_input, boundary_1)\n",
    "assert np.array_equal(result_1, expected_1), f\"Test 1 failed: expected {expected_1}, got {result_1}\"\n",
    "\n",
    "test_2_input = tensor([-5, 0, 5, 10])\n",
    "boundary_2 = tensor([-10, -5, 0, 5, 10])\n",
    "expected_2 = tensor([1, 2, 3, 4])\n",
    "result_2 = bucketize(test_2_input, boundary_2)\n",
    "assert np.array_equal(result_2, expected_2), f\"Test 2 failed: expected {expected_2}, got {result_2}\"\n",
    "\n",
    "test_3_input = tensor([1, 2, 3, 4, 5])\n",
    "boundary_3 = tensor([1, 3, 5])\n",
    "expected_3 = tensor([0, 1, 1, 2, 2])\n",
    "result_3 = bucketize(test_3_input, boundary_3)\n",
    "assert np.array_equal(result_3, expected_3), f\"Test 3 failed: expected {expected_3}, got {result_3}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lab",
   "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}