{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import abc, Counter\n",
    "from functools import lru_cache\n",
    "from itertools import chain\n",
    "import matplotlib.pyplot as plt\n",
    "import mmh3\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from htools import hdir, eprint, assert_raises, flatten"
   ]
  },
  {
   "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": [],
   "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": [
    {
     "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": [
       "torch.Size([3, 10])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "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": [
    "ds.x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'is': 1,\n",
       " 'I': 2,\n",
       " 'the': 3,\n",
       " 'grey.': 4,\n",
       " 'walked': 5,\n",
       " 'to': 6,\n",
       " 'store': 7,\n",
       " 'so': 8,\n",
       " 'hope': 9,\n",
       " 'it': 10,\n",
       " 'not': 11,\n",
       " 'closed.': 12,\n",
       " 'The': 13,\n",
       " 'theater': 14,\n",
       " 'closed': 15,\n",
       " 'today': 16,\n",
       " 'and': 17,\n",
       " 'sky': 18,\n",
       " 'His': 19,\n",
       " 'dog': 20,\n",
       " 'brown': 21,\n",
       " 'while': 22,\n",
       " 'hers': 23}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.w2i"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For now, just convert int to str and take hash. Another option that is meant for ints is Knuth's multiplicative method:\n",
    "\n",
    "hash(i) = i*2654435761 mod 2^32\n",
    "\n",
    "But we'd need to make this dependent on a random seed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def probabilistic_hash_item(x, n_buckets, mode=int, n_hashes=3):\n",
    "    \"\"\"Slightly hacky way to probabilistically hash an integer by\n",
    "    first converting it to a string.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: int\n",
    "        The integer or string to hash.\n",
    "    n_buckets: int\n",
    "        The number of buckets that items will be mapped to. Typically \n",
    "        this would occur outside the hashing function, but since \n",
    "        the intended use case is so narrow here it makes sense to me \n",
    "        to include it here.\n",
    "    mode: type\n",
    "        The type of input you want to hash. This is user-provided to prevent\n",
    "        accidents where we pass in a different item than intended and hash \n",
    "        the wrong thing. One of (int, str). When using this inside a\n",
    "        BloomEmbedding layer, this must be `int` because there are no\n",
    "        string tensors. When used inside a dataset or as a one-time\n",
    "        pre-processing step, you can choose either as long as you\n",
    "        pass in the appropriate inputs.\n",
    "    n_hashes: int\n",
    "        The number of times to hash x, each time with a different seed.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    list[int]: A list of integers with length `n_hashes`, where each integer\n",
    "        is in [0, n_buckets).\n",
    "    \"\"\"\n",
    "    # Check type to ensure we don't accidentally hash Tensor(5) instead of 5.\n",
    "    assert isinstance(x, mode), f'Input `x` must have type {mode}.'\n",
    "    return [mmh3.hash(str(x), i, signed=False) % n_buckets \n",
    "            for i in range(n_hashes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def probabilistic_hash_tensor(x_r2, n_buckets, n_hashes=3, pad_idx=0):\n",
    "    \"\"\"Hash a rank 2 LongTensor.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x_r2: torch.LongTensor\n",
    "        Rank 2 tensor of integers. Shape: (bs, seq_len)\n",
    "    n_buckets: int\n",
    "        Number of buckets to hash items into (i.e. the number of \n",
    "        rows in the embedding matrix). Typically a moderately large\n",
    "        prime number, like 251 or 997.\n",
    "    n_hashes: int\n",
    "        Number of hashes to take for each input index. This determines\n",
    "        the number of rows of the embedding matrix that will be summed\n",
    "        to get the representation for each word. Typically 2-5.\n",
    "    pad_idx: int or None\n",
    "        If you want to pad sequences with vectors of zeros, pass in an\n",
    "        integer (same as the `padding_idx` argument to nn.Embedding).\n",
    "        If None, no padding index will be used. The sequences must be\n",
    "        padded before passing them into this function.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.LongTensor: Tensor of indices where each row corresponds\n",
    "        to one of the input indices. Shape: (bs, seq_len, n_hashes)\n",
    "    \"\"\"\n",
    "    return torch.tensor(\n",
    "        [[probabilistic_hash_item(x.item(), n_buckets, int, n_hashes) \n",
    "          if x != pad_idx else [pad_idx]*n_hashes for x in row]\n",
    "         for row in x_r2]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "tensor([ 2,  5,  6,  3,  7,  8,  2,  9, 10,  1])\n",
      "[2, 5, 6, 3, 7, 8, 2, 9, 10, 1]\n",
      "\n",
      "1\n",
      "tensor([13, 14,  1, 15, 16, 17,  3, 18,  1,  4])\n",
      "[13, 14, 1, 15, 16, 17, 3, 18, 1, 4]\n",
      "\n",
      "2\n",
      "tensor([19, 20,  1, 21, 22, 23,  1,  4,  0,  0])\n",
      "[19, 20, 1, 21, 22, 23, 1, 4, 0, 0]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i, row in enumerate(x):\n",
    "    print(i)\n",
    "    print(row)\n",
    "    print([x.item() for x in row], end='\\n\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 8,  2,  7],\n",
       "         [ 2,  8,  1],\n",
       "         [ 6,  6, 10],\n",
       "         [10,  5,  5],\n",
       "         [ 6,  9,  7],\n",
       "         [ 5,  9,  4],\n",
       "         [ 8,  2,  7],\n",
       "         [ 5, 10,  8],\n",
       "         [ 7,  8,  6],\n",
       "         [ 6, 10,  6]],\n",
       "\n",
       "        [[ 2,  5,  1],\n",
       "         [ 9,  8,  8],\n",
       "         [ 6, 10,  6],\n",
       "         [ 8,  1, 10],\n",
       "         [ 6,  3,  9],\n",
       "         [ 2,  6,  9],\n",
       "         [10,  5,  5],\n",
       "         [10,  5,  0],\n",
       "         [ 6, 10,  6],\n",
       "         [ 4,  8,  6]],\n",
       "\n",
       "        [[ 2, 10,  6],\n",
       "         [ 1,  7,  8],\n",
       "         [ 6, 10,  6],\n",
       "         [ 2,  7,  1],\n",
       "         [ 9,  4,  8],\n",
       "         [ 5,  3,  3],\n",
       "         [ 6, 10,  6],\n",
       "         [ 4,  8,  6],\n",
       "         [ 0,  0,  0],\n",
       "         [ 0,  0,  0]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probabilistic_hash_tensor(x, 11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 8,  2,  7],\n",
       "         [ 2,  8,  1],\n",
       "         [ 6,  6, 10],\n",
       "         [10,  5,  5],\n",
       "         [ 6,  9,  7],\n",
       "         [ 5,  9,  4],\n",
       "         [ 8,  2,  7],\n",
       "         [ 5, 10,  8],\n",
       "         [ 7,  8,  6],\n",
       "         [ 6, 10,  6]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probabilistic_hash_tensor(x[0, None], 11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 2, 10,  6],\n",
       "         [ 1,  7,  8],\n",
       "         [ 6, 10,  6],\n",
       "         [ 2,  7,  1],\n",
       "         [ 9,  4,  8],\n",
       "         [ 5,  3,  3],\n",
       "         [ 6, 10,  6],\n",
       "         [ 4,  8,  6],\n",
       "         [ 0,  0,  0],\n",
       "         [ 0,  0,  0]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probabilistic_hash_tensor(x[2, None], 11)"
   ]
  },
  {
   "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": [
       "[[8, 2, 7],\n",
       " [2, 8, 1],\n",
       " [6, 6, 10],\n",
       " [10, 5, 5],\n",
       " [6, 9, 7],\n",
       " [5, 9, 4],\n",
       " [8, 2, 7],\n",
       " [5, 10, 8],\n",
       " [7, 8, 6],\n",
       " [6, 10, 6]]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[probabilistic_hash_item(n.item(), 11) for n in x[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[9, 8, 1]\n",
      "[2, 6, 9]\n",
      "[8, 2, 7]\n",
      "[2, 5, 7]\n",
      "[10, 10, 2]\n",
      "[0, 10, 4]\n",
      "[3, 8, 4]\n",
      "[6, 5, 0]\n",
      "[6, 10, 8]\n",
      "[4, 4, 0]\n",
      "[7, 1, 0]\n",
      "[10, 4, 6]\n"
     ]
    }
   ],
   "source": [
    "for i in range(0, 200, 17):\n",
    "    print(probabilistic_hash_item(i, 11))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 0: ('I', [0, 5, 5])\n",
      " 1: ('walked', [2, 5, 1])\n",
      " 2: ('to', [10, 4, 6])\n",
      " 3: ('the', [4, 1, 4])\n",
      " 4: ('store', [4, 6, 3])\n",
      " 5: ('so', [7, 8, 8])\n",
      " 6: ('I', [0, 5, 5])\n",
      " 7: ('hope', [1, 2, 7])\n",
      " 8: ('it', [3, 9, 0])\n",
      " 9: ('is', [3, 1, 3])\n",
      "10: ('not', [6, 10, 4])\n",
      "11: ('closed.', [3, 6, 10])\n",
      "\n",
      " 0: ('The', [1, 6, 9])\n",
      " 1: ('theater', [8, 10, 2])\n",
      " 2: ('is', [3, 1, 3])\n",
      " 3: ('closed', [5, 5, 0])\n",
      " 4: ('today', [3, 10, 8])\n",
      " 5: ('and', [7, 2, 4])\n",
      " 6: ('the', [4, 1, 4])\n",
      " 7: ('sky', [1, 2, 9])\n",
      " 8: ('is', [3, 1, 3])\n",
      " 9: ('grey.', [7, 6, 7])\n",
      "\n",
      " 0: ('His', [0, 10, 3])\n",
      " 1: ('dog', [8, 6, 6])\n",
      " 2: ('is', [3, 1, 3])\n",
      " 3: ('brown', [9, 8, 9])\n",
      " 4: ('while', [9, 2, 8])\n",
      " 5: ('hers', [0, 5, 4])\n",
      " 6: ('is', [3, 1, 3])\n",
      " 7: ('grey.', [7, 6, 7])\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for row in [s.split(' ') for s in sents]:\n",
    "    eprint(list(zip(row, (probabilistic_hash_item(word, 11, str) for word in row))))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "As expected, got AssertionError(torch tensor).\n"
     ]
    }
   ],
   "source": [
    "assert isinstance(np.array(sents), (list, np.ndarray)), 'np array'\n",
    "assert isinstance(sents, (list, np.ndarray)), 'list'\n",
    "with assert_raises(AssertionError):\n",
    "    assert isinstance(x, (list, np.ndarray)), 'torch tensor'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iterable tensor True\n",
      "Iterable list True\n",
      "Iterable array True\n",
      "\n",
      "Collection tensor True\n",
      "Collection list True\n",
      "Collection array True\n",
      "\n",
      "Container tensor True\n",
      "Container list True\n",
      "Container array True\n",
      "\n",
      "Sequence tensor False\n",
      "Sequence list True\n",
      "Sequence array False\n",
      "\n",
      "MutableSequence tensor False\n",
      "MutableSequence list True\n",
      "MutableSequence array False\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for t in (abc.Iterable, abc.Collection, abc.Container, abc.Sequence, abc.MutableSequence):\n",
    "    tname = t.__name__\n",
    "    print(tname, 'tensor', isinstance(x, t))\n",
    "    print(tname, 'list', isinstance(sents, t))\n",
    "    print(tname, 'array', isinstance(np.array(sents), t))\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# To do:\n",
    "- [x] maybe handle padding differently (i.e. just a row of zeros as usual. Should this happen in hash_int, hash_int_tensor, or BloomEmbedding?)\n",
    "- maybe use fastai embedding() instead of nn.Embedding? Check if the weight init method might help for this.\n",
    "- maybe use nn.embeddingbag? (note: people report this is much slower on CPU, and possibly slower on GPU even though it's supposed to be faster)\n",
    "- experiment with different numbers of embeddings and hashes. Try to find guidelines for what reasonable choices are to prevent collisions. \n",
    "    - Eventually, maybe better to let user input vocab size and choose prob of collision, then automatically select values for n_emb and n_hashes. \n",
    "- check to make sure indices are working correctly after switching hash_int_tensor to take only 2d tensors. Also consider if this is the preferred way to do this.\n",
    "- should we let user choose between mean and sum? Wonder if mean would be better bc we could try different values of n_hashes while still loading pre-trained embeddings (bc scale is standardized)? But that probably doesn't work bc the hashes will be different anyway."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.5849, -0.1108, -0.1991, -1.1982],\n",
       "        [-0.5891,  1.5808,  1.4062, -1.5162],\n",
       "        [-1.0055,  0.1712, -0.5717, -0.2766],\n",
       "        [ 1.6280, -0.1770, -1.4850, -2.0490],\n",
       "        [ 0.8724,  0.5324,  0.1913, -0.9653]], requires_grad=True)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "emb = nn.EmbeddingBag(5, 4, mode='sum')\n",
    "emb.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-2.1837,  3.3329,  2.2408, -3.3091],\n",
       "        [ 1.2073, -0.1165, -2.2557, -3.5238]], grad_fn=<EmbeddingBagBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.Tensor([[1, 2, 1],\n",
    "                  [3, 2, 0]]).long()\n",
    "emb(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "    \n",
    "    The reduction in rows allows us to use memory in other ways: \n",
    "    a larger embedding dimension, more or larger layers after the embedding,\n",
    "    or larger batch sizes.\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 `forward`\n",
    "            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 = nn.Embedding(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": [
    {
     "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.0890,  2.8199,  1.8165, -1.9656],\n",
       "        [ 1.2061,  0.3585, -0.5961, -1.2669],\n",
       "        [-1.5584, -2.4416,  1.1176, -0.4356],\n",
       "        [ 0.1442, -0.2951,  0.8433, -0.1157],\n",
       "        [-0.2319,  0.6422, -0.1494,  0.2830],\n",
       "        [ 0.7345,  2.0709,  1.4002,  0.0562],\n",
       "        [-0.3707,  0.5033,  0.6533, -1.4163],\n",
       "        [ 0.7275,  0.1103,  1.3854,  0.0924],\n",
       "        [ 0.3850,  0.0692,  0.1639,  0.3658],\n",
       "        [-1.2039, -1.9414, -0.2005, -0.4414]], 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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[9, 8, 1]\n",
      "[6, 10, 6]\n",
      "[8, 2, 7]\n",
      "[10, 5, 5]\n",
      "[4, 8, 6]\n",
      "[2, 8, 1]\n",
      "[6, 6, 10]\n",
      "[6, 9, 7]\n",
      "[5, 9, 4]\n",
      "[5, 10, 8]\n",
      "[7, 8, 6]\n",
      "[7, 9, 2]\n",
      "[1, 7, 3]\n",
      "[2, 5, 1]\n",
      "[9, 8, 8]\n",
      "[8, 1, 10]\n",
      "[6, 3, 9]\n",
      "[2, 6, 9]\n",
      "[10, 5, 0]\n",
      "[2, 10, 6]\n",
      "[1, 7, 8]\n",
      "[2, 7, 1]\n",
      "[9, 4, 8]\n",
      "[5, 3, 3]\n"
     ]
    }
   ],
   "source": [
    "for i in range(24):\n",
    "    print(probabilistic_hash_item(i, 11))"
   ]
  },
  {
   "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([[ 2.2904,  1.0825,  2.8281, -2.4985],\n",
       "        [ 3.2287,  3.6471,  2.0098, -4.4070],\n",
       "        [-0.9388,  0.2589,  2.3995, -0.7704],\n",
       "        [-1.8995, -0.0147, -0.6485,  0.4076],\n",
       "        [ 1.9549,  3.0019,  1.6213, -2.2612],\n",
       "        [ 0.2973,  0.4164,  0.8578,  0.5331],\n",
       "        [ 2.2904,  1.0825,  2.8281, -2.4985],\n",
       "        [-0.3233, -1.1196,  1.1995,  0.2997],\n",
       "        [ 2.2974,  3.0430,  2.8428, -2.5347],\n",
       "        [ 0.2651,  2.2004,  2.6000, -0.3290]], grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 2.2693,  4.1790,  0.4750, -4.2163],\n",
       "        [ 2.5744,  2.3608,  4.3350,  0.6067],\n",
       "        [ 0.2651,  2.2004,  2.6000, -0.3290],\n",
       "        [-1.9458, -1.4528,  4.1191, -2.7502],\n",
       "        [-0.4390, -0.3015,  2.6817, -0.0136],\n",
       "        [ 2.3256,  2.4986,  0.9680, -0.8449],\n",
       "        [-1.8995, -0.0147, -0.6485,  0.4076],\n",
       "        [-0.2297, -0.9407, -0.9459, -1.4253],\n",
       "        [ 0.2651,  2.2004,  2.6000, -0.3290],\n",
       "        [ 2.8123,  2.2446,  3.0328, -1.2341]], grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.4671, -1.4535,  0.4032, -2.0935],\n",
       "        [ 1.1803,  5.5044,  5.2554, -3.2333],\n",
       "        [ 0.2651,  2.2004,  2.6000, -0.3290],\n",
       "        [ 1.0686,  3.3866,  2.7170, -4.7645],\n",
       "        [ 1.9912,  1.9554,  3.7928,  0.3986],\n",
       "        [-2.6212, -4.1307,  3.4713, -0.4958],\n",
       "        [ 0.2651,  2.2004,  2.6000, -0.3290],\n",
       "        [ 2.8123,  2.2446,  3.0328, -1.2341],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000],\n",
       "        [ 0.0000,  0.0000,  0.0000,  0.0000]], grad_fn=<SelectBackward>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y[2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "is 1 [ 0.26507235  2.2003503   2.5999665  -0.32898915]\n",
      "I 2 [ 2.2903948  1.0824531  2.828073  -2.498523 ]\n",
      "the 3 [-1.8995289  -0.01470172 -0.6485326   0.40761688]\n",
      "grey. 4 [ 2.812286   2.24457    3.0328324 -1.2340608]\n",
      "walked 5 [ 3.2286701  3.6471388  2.009758  -4.4069753]\n",
      "to 6 [-0.93877876  0.2589332   2.3995147  -0.7704108 ]\n",
      "store 7 [ 1.9548693  3.0019045  1.6212814 -2.2612271]\n",
      "so 8 [0.29734743 0.41637933 0.85779124 0.5330522 ]\n",
      "hope 9 [-0.3232504  -1.1195885   1.1994998   0.29972154]\n",
      "it 10 [ 2.2973628  3.042996   2.8428445 -2.534657 ]\n",
      "not 11 [ 0.8497246  1.4343383  0.8743619 -3.7337801]\n",
      "closed. 12 [-3.043953  -1.0598472  3.3869708 -4.2589087]\n",
      "The 13 [ 2.2692842  4.179036   0.4749601 -4.2163124]\n",
      "theater 14 [2.5744486  2.3608148  4.334958   0.60669684]\n",
      "closed 15 [-1.9457765 -1.452824   4.1191187 -2.7502215]\n",
      "today 16 [-0.43895298 -0.30150518  2.6817129  -0.01360181]\n",
      "and 17 [ 2.325552    2.498587    0.9679917  -0.84489024]\n",
      "sky 18 [-0.22965312 -0.9407249  -0.9459034  -1.4252954 ]\n",
      "His 19 [-0.4671499  -1.4534968   0.40321434 -2.0935135 ]\n",
      "dog 20 [ 1.1802679  5.5044327  5.2554398 -3.2333224]\n",
      "brown 21 [ 1.0686436  3.3865533  2.716979  -4.764516 ]\n",
      "while 22 [1.9911952  1.9553655  3.792798   0.39860588]\n",
      "hers 23 [-2.621228   -4.1306973   3.471336   -0.49583316]\n"
     ]
    }
   ],
   "source": [
    "for w, i in ds.w2i.items():\n",
    "    print(w, i, be(torch.tensor([[i]])).detach().numpy().squeeze())\n",
    "#           .emb.weight[hash_int(i, be.n_emb)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[5, 3, 3, 8]]])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hashed = probabilistic_hash_tensor(torch.tensor([23]).unsqueeze(0), 11, 4)\n",
    "hashed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-2.6212, -4.1307,  3.4713, -0.4958]]], grad_fn=<SumBackward1>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "be.emb.weight[hashed].sum(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unique_combos(tups):\n",
    "    return len(set(tuple(sorted(x)) for x in tups))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hash_all_idx(vocab_size, n_buckets, n_hashes):\n",
    "    return [probabilistic_hash_item(i, n_buckets, int, n_hashes) \n",
    "            for i in range(vocab_size)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Buckets: 127 Hashes: 5 Unique combos: 29998 % unique: 0.99993\n",
      "Buckets: 251 Hashes: 4 Unique combos: 29996 % unique: 0.99987\n",
      "Buckets: 997 Hashes: 3 Unique combos: 29997 % unique: 0.9999\n",
      "Buckets: 5003 Hashes: 2 Unique combos: 29969 % unique: 0.99897\n"
     ]
    }
   ],
   "source": [
    "buckets2hashes = {127: 5,\n",
    "                  251: 4,\n",
    "                  997: 3,\n",
    "                  5_003: 2}\n",
    "for b, h in buckets2hashes.items():\n",
    "    tups = hash_all_idx(30_000, b,  h)\n",
    "    unique = unique_combos(tups)\n",
    "    print('Buckets:', b, 'Hashes:', h, 'Unique combos:', unique,\n",
    "          '% unique:', round(unique/30_000, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_n_buckets(vocab_size, hash_sizes, bucket_sizes):\n",
    "    for bs in bucket_sizes:\n",
    "        for hs in hash_sizes:\n",
    "            tups = hash_all_idx(vocab_size, bs, hs)\n",
    "            unique = unique_combos(tups)\n",
    "            print('buckets:', bs, \n",
    "                  'hashes:', hs, \n",
    "                  'unique:', round(unique/vocab_size, 4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "buckets: 5 hashes: 2 unique: 0.1875\n",
      "buckets: 5 hashes: 3 unique: 0.3625\n",
      "buckets: 5 hashes: 4 unique: 0.5375\n",
      "buckets: 5 hashes: 5 unique: 0.625\n",
      "buckets: 11 hashes: 2 unique: 0.625\n",
      "buckets: 11 hashes: 3 unique: 0.875\n",
      "buckets: 11 hashes: 4 unique: 0.975\n",
      "buckets: 11 hashes: 5 unique: 1.0\n",
      "buckets: 13 hashes: 2 unique: 0.675\n",
      "buckets: 13 hashes: 3 unique: 0.9\n",
      "buckets: 13 hashes: 4 unique: 1.0\n",
      "buckets: 13 hashes: 5 unique: 1.0\n",
      "buckets: 19 hashes: 2 unique: 0.85\n",
      "buckets: 19 hashes: 3 unique: 0.975\n",
      "buckets: 19 hashes: 4 unique: 1.0\n",
      "buckets: 19 hashes: 5 unique: 1.0\n",
      "buckets: 29 hashes: 2 unique: 0.9375\n",
      "buckets: 29 hashes: 3 unique: 0.9875\n",
      "buckets: 29 hashes: 4 unique: 1.0\n",
      "buckets: 29 hashes: 5 unique: 1.0\n",
      "buckets: 37 hashes: 2 unique: 0.925\n",
      "buckets: 37 hashes: 3 unique: 1.0\n",
      "buckets: 37 hashes: 4 unique: 1.0\n",
      "buckets: 37 hashes: 5 unique: 1.0\n"
     ]
    }
   ],
   "source": [
    "eval_n_buckets(80, range(2, 6), [5, 11, 13, 19, 29, 37])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 500])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randint(0, 30_000, (64, 500))\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "339 ms ± 2.82 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 127, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "343 ms ± 2.82 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 251, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "352 ms ± 4.19 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 997, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "346 ms ± 3.62 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 5_003, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "318 ms ± 1.54 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 251, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "303 ms ± 3.05 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 997, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "284 ms ± 2.39 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 5_003, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13412 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(range(2, 6), [284, 303, 318, 339], s=[127, 251, 997, 5_003])\n",
    "plt.xlabel('# of Hashes')\n",
    "plt.ylabel('Time (ms) \\nto encode \\na batch with \\n32,000 words',\n",
    "           rotation=0, labelpad=40)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Try memoizing hash func.\n",
    "\n",
    "Caching results does save computation but it increases memory usage, which is one of the benefits of Bloom Embeddings. However, this is probably still more memory efficient since each word in the cache is represented by only a few indices rather than a large embedding. A large embedding matrix also forces us to devote a lot of memory to unused gradients during training (if the layer is not frozen). Still, the fact that `probabilistic_hash_tensor` is not vectorized means that it may be best to pre-compute these and load the indices in the Dataset. Rather than caching, it may be easiest to create a w2hash dict similar to a standard w2index dict. Leaving this implementation lets the user determine what tradeoff they want between speed and memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@lru_cache(maxsize=30_000)\n",
    "def probabilistic_hash_item(x, n_buckets, mode=int, n_hashes=3):\n",
    "    \"\"\"Slightly hacky way to probabilistically hash an integer by\n",
    "    first converting it to a string.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x: int\n",
    "        The integer or string to hash.\n",
    "    n_buckets: int\n",
    "        The number of buckets that items will be mapped to. Typically \n",
    "        this would occur outside the hashing function, but since \n",
    "        the intended use case is so narrow here it makes sense to me \n",
    "        to include it here.\n",
    "    mode: type\n",
    "        The type of input you want to hash. This is user-provided to prevent\n",
    "        accidents where we pass in a different item than intended and hash \n",
    "        the wrong thing. One of (int, str). When using this inside a\n",
    "        BloomEmbedding layer, this must be `int` because there are no\n",
    "        string tensors. When used inside a dataset or as a one-time\n",
    "        pre-processing step, you can choose either as long as you\n",
    "        pass in the appropriate inputs.\n",
    "    n_hashes: int\n",
    "        The number of times to hash x, each time with a different seed.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    list[int]: A list of integers with length `n_hashes`, where each integer\n",
    "        is in [0, n_buckets).\n",
    "    \"\"\"\n",
    "    # Check type to ensure we don't accidentally hash Tensor(5) instead of 5.\n",
    "    assert isinstance(x, mode), f'Input `x` must have type {mode}.'\n",
    "    return [mmh3.hash(str(x), i, signed=False) % n_buckets \n",
    "            for i in range(n_hashes)]\n",
    "\n",
    "\n",
    "def probabilistic_hash_tensor(x_r2, n_buckets, n_hashes=3, pad_idx=0):\n",
    "    \"\"\"Hash a rank 2 LongTensor.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x_r2: torch.LongTensor\n",
    "        Rank 2 tensor of integers. Shape: (bs, seq_len)\n",
    "    n_buckets: int\n",
    "        Number of buckets to hash items into (i.e. the number of \n",
    "        rows in the embedding matrix). Typically a moderately large\n",
    "        prime number, like 251 or 997.\n",
    "    n_hashes: int\n",
    "        Number of hashes to take for each input index. This determines\n",
    "        the number of rows of the embedding matrix that will be summed\n",
    "        to get the representation for each word. Typically 2-5.\n",
    "    pad_idx: int or None\n",
    "        If you want to pad sequences with vectors of zeros, pass in an\n",
    "        integer (same as the `padding_idx` argument to nn.Embedding).\n",
    "        If None, no padding index will be used. The sequences must be\n",
    "        padded before passing them into this function.\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    torch.LongTensor: Tensor of indices where each row corresponds\n",
    "        to one of the input indices. Shape: (bs, seq_len, n_hashes)\n",
    "    \"\"\"\n",
    "    return torch.tensor(\n",
    "        [[probabilistic_hash_item(x.item(), n_buckets, int, n_hashes) \n",
    "          if x != pad_idx else [pad_idx]*n_hashes for x in row]\n",
    "         for row in x_r2]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "191 ms ± 12 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 5_003, 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "185 ms ± 8.76 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 5_003, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "195 ms ± 5.1 ms per loop (mean ± std. dev. of 5 runs, 5 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 5 -r 5\n",
    "hashed = probabilistic_hash_tensor(x, 5_003, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CacheInfo(hits=780342, misses=19633, maxsize=30000, currsize=19633)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probabilistic_hash_item.cache_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19634"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set(flatten(x.numpy())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "probabilistic_hash_item.cache_clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CacheInfo(hits=0, misses=0, maxsize=30000, currsize=0)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probabilistic_hash_item.cache_info()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
