reformer_LSH_Self_Attention.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hVHY_ZYxWDph"
   },
   "source": [
    "# Reformer Efficient Attention: Ungraded Lab\n",
    "The videos describe two 'reforms' made to the Transformer to make it more memory and compute efficient. The *Reversible Layers* reduce memory and *Locality Sensitive Hashing(LSH)* reduces the cost of the Dot Product attention for large input sizes. This ungraded lab will look more closely at LSH and how it is used in the Reformer model.\n",
    "\n",
    "Specifically, the notebook has 3 goals\n",
    "* review dot-product self attention for reference\n",
    "* examine LSH based self attention\n",
    "* extend our understanding and familiarity with Trax infrastructure\n",
    "\n",
    "## Outline\n",
    "- [Part 1:  Trax Efficient Attention classes](#1)\n",
    "- [Part 2:  Full Dot Product Self Attention](#2)\n",
    "    - [2.1  Description](#2.1)\n",
    "        - [2.1.1  our_softmax](#2.1.1)\n",
    "    - [2.2  our simple attend](#2.2)\n",
    "    - [2.3  Class OurSelfAttention](#2.3)\n",
    "- [Part 3:  Trax LSHSelfAttention](#3)\n",
    "    - [3.1  Description](#3.1)\n",
    "    - [3.2  our_hash_vectors](#3.2)\n",
    "    - [3.3  Sorting Buckets](#3.3)\n",
    "    - [3.4  Chunked dot product attention](#3.4)\n",
    "    - [3.5  OurLSHSelfAttention](#3.5)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "1nhevIw88Vm2"
   },
   "source": [
    "<a name=\"1\"></a>\n",
    "## Part 1.0  Trax Efficient Attention classes\n",
    "Trax is similar to other popular NN development platforms such as Keras (now integrated into Tensorflow) and Pytorch in that it uses 'layers' as a useful level of abstraction. Layers are often represented as *classes*. We're going to improve our understanding of Trax by locally extending the classes used in the attention layers. We will extend only the 'forward' functions and utilize the existing attention layers as parent classes. The original code can be found at [github:trax/layers/Research/Efficient_attention](https://github.com/google/trax/blob/v1.3.4/trax/layers/research/efficient_attention.py). This link references release 1.3.4 but note that this is under the 'research' directory as this is an area of active research. When accessing the code on Github for review on this assignment, be sure you select the 1.3.4 release tag, the master copy may have new changes.:\n",
    "<img src = \"C4W4_LN2_image11.PNG\" height=\"250\" width=\"250\">\n",
    "<center><b>Figure 1: Reference Tag 1.3.4 on github</b></center>\n",
    "\n",
    "\n",
    "\n",
    "While Trax uses classes liberally, we have not built many classes in the course so far. Let's spend a few moments reviewing the classes we will be using.\n",
    "<img src = \"C4W4_LN2_image1.PNG\" height=\"788\" width=\"1561\">\n",
    "\n",
    "<center><b>Figure 2: Classes from Trax/layers/Research/Efficient_Attention.py that we will be utilizing.</b></center>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting on the right in the diagram below you see EfficientAttentionBase. The parent to this class is the base.layer which has the routines used by all layers. EfficientAttentionBase leaves many routines to be overridden by child classes - but it has an important feature in the *Forward* routine. It supports a `use_reference_code` capability that selects implementations that limit some of the complexities to provide a more easily understood version of the algorithms. In particular, it implements a nested loop that treats each *'example, head'* independently. This simplifies our work as we need only worry about matrix operations on one *'example, head'* at a time. This loop calls *forward_unbatched*, which is the child process that we will be overriding.\n",
    "\n",
    "On the top left are the outlines of the two child classes we will be using. The SelfAttention layer is a 'traditional' implementation of the dot product attention. We will be implementing the *forward_unbatched* version of this to highlight the differences between this and the LSH implementation.\n",
    "\n",
    "Below that is the LSHSelfAttention. This is the routine used in the Reformer architecture. We will override the *forward_unbatched* section of this and some of the utility functions it uses to explore its implementation in more detail.\n",
    "\n",
    "The code we will be working with is from the Trax source, and as such has implementation details that will make it a bit harder to follow. However, it will allow use of the results along with the rest of the Trax infrastructure. I will try to briefly describe these as they arise. The [Trax documentation](https://trax-ml.readthedocs.io/en/latest/) can also be referenced."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"1.2\"></a>\n",
    "## Part 1.2  Trax Details\n",
    "The goal in this notebook is to override a few routines in the Trax classes with our own versions. To maintain their functionality in a full Trax environment, many of the details we might ignore in example version of routines will be maintained in this code. Here are some of the considerations that may impact our code:\n",
    "* Trax operates with multiple back-end libraries, we will see special cases that will utilize unique features.\n",
    "* 'Fancy' numpy indexing is not supported in all backend environments and must be emulated in other ways.\n",
    "* Some operations don't have gradients for backprop and must be ignored or include forced re-evaluation.\n",
    "\n",
    "Here are some of the functions we may see:\n",
    "* Abstracted as `fastmath`, Trax supports multiple backend's such as [Jax](https://github.com/google/jax) and [Tensorflow2](https://github.com/tensorflow/tensorflow)\n",
    "* [tie_in](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.tie_in.html): Some non-numeric operations must be invoked during backpropagation. Normally, the gradient compute graph would determine invocation but these functions are not included. To force re-evaluation, they are 'tied' to other numeric operations using tie_in.\n",
    "* [stop_gradient](https://trax-ml.readthedocs.io/en/latest/trax.fastmath.html): Some operations are intentionally excluded from backprop gradient calculations by setting their gradients to zero.\n",
    "* Below we will execute `from trax.fastmath import numpy as np `, this uses accelerated forms of numpy functions. This is, however a *subset* of numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BGLyG_n_4WhH",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:tokens_length=568 inputs_length=512 targets_length=114 noise_density=0.15 mean_noise_span_length=3.0 \n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import trax\n",
    "from trax import layers as tl  # core building block\n",
    "import jax\n",
    "from trax import fastmath  # uses jax, offers numpy on steroids\n",
    "\n",
    "# fastmath.use_backend('tensorflow-numpy')\n",
    "import functools\n",
    "from trax.fastmath import numpy as np  # note, using fastmath subset of numpy!\n",
    "from trax.layers import (\n",
    "    tie_in,\n",
    "    length_normalized,\n",
    "    apply_broadcasted_dropout,\n",
    "    look_adjacent,\n",
    "    permute_via_gather,\n",
    "    permute_via_sort,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"2\"></a>\n",
    "## Part 2  Full Dot-Product Self Attention\n",
    "<a name=\"1.2\"></a>\n",
    "### Part 2.1 Description\n",
    "<img src = \"C4W4_LN2_image2.PNG\" height=\"200\" width=\"600\">\n",
    "\n",
    "<center><b>Figure 3: Project datapath and primary data structures and where they are implemented</b></center>\n",
    "\n",
    "The diagram above shows many of the familiar data structures and operations related to attention and describes the routines in which they are implemented. We will start by working on *our_simple_attend* or our simpler version of the original *attend* function. We will review the steps in performing dot-product attention with more focus on the details of the operations and their significance. This is useful when comparing to LSH attention. Note we will be discussing a single example/head unless otherwise specified.\n",
    "\n",
    "<img src = \"C4W4_LN2_image3.PNG\" height=\"250\" width=\"700\">\n",
    "\n",
    "<center><b>Figure 4: dot-product of Query and Key</b></center>\n",
    "\n",
    "The *attend* function receives *Query* and *Key*. As a reminder, they are produced by a matrix multiply of all the inputs with a single set of weights. We will describe the inputs as *embeddings* assuming an NLP application, however, this is not required. This matrix multiply very much like a convolutional network where a set of weights (a filter) slide across the input vectors leaving behind a map of the similarity of the input to the filter. In this case, the filters are the weight matrices $W^Q$ and $W^K$. The resulting maps are Q and K. Q and K have the dimensions of (n_seq, n_q) where n_seq is the number input embeddings and n_q or n_k is the selected size of the Q or K vectors. Note the shading of Q and K, this reflects the fact that each entry is associated with a particular input embedding. You will note later in the code that K is optional. Apparently, similar results can be achieved using Query alone saving the compute and storage associated with K. In that case, the dot-product in *attend* is matmul(q,q). Note the resulting dot-product (*Dot*) entries describe a complete (n_seq,n_seq) map of the similarity of all entries of q vs all entries of k. This is reflected in the notation in the dot-product boxes of $w_n$,$w_m$ representing word_n, word_m. Note that each row of *Dot* describes the relationship of an input embedding, say $w_0$, with every other input.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In some applications some values are masked. This can be used, for example to exclude results that occur later in time (causal) or to mask padding or other inputs.\n",
    "<img src = \"C4W4_LN2_image4.PNG\" height=\"300\" width=\"900\">\n",
    "\n",
    "<center><b>Figure 5: Masking</b></center>\n",
    "\n",
    "\n",
    "The routine below *mask_self_attention* implements a flexible masking capability. The masking is controlled by the information in q_info and kv_info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kDAWq_xMX6C1"
   },
   "outputs": [],
   "source": [
    "def mask_self_attention(\n",
    "    dots, q_info, kv_info, causal=True, exclude_self=True, masked=False\n",
    "):\n",
    "    \"\"\"Performs masking for self-attention.\"\"\"\n",
    "    if causal:\n",
    "        mask = fastmath.lt(q_info, kv_info).astype(np.float32)\n",
    "        dots = dots - 1e9 * mask\n",
    "    if exclude_self:\n",
    "        mask = np.equal(q_info, kv_info).astype(np.float32)\n",
    "        dots = dots - 1e5 * mask\n",
    "    if masked:\n",
    "        zeros_like_kv_info = tie_in(kv_info, np.zeros_like(kv_info))\n",
    "        mask = fastmath.lt(kv_info, zeros_like_kv_info).astype(np.float32)\n",
    "        dots = dots - 1e9 * mask\n",
    "    return dots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A SoftMax is applied per row of the *Dot* matrix to scale the values in the row between 0 and 1.\n",
    "<img src = \"C4W4_LN2_image5.PNG\" height=\"300\" width=\"900\">\n",
    "\n",
    "<center><b>Figure 6: SoftMax per row of Dot</b></center>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"2.1.1\"></a>\n",
    "### Part 2.1.1 our_softmax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code uses a separable form of the softmax calculation. Recall the softmax:\n",
    "$$ softmax(x_i)=\\frac{\\exp(x_i)}{\\sum_j \\exp(x_j)}\\tag{1}$$\n",
    "This can be alternately implemented as:\n",
    "$$ logsumexp(x)=\\log{({\\sum_j \\exp(x_j)})}\\tag{2}$$\n",
    "$$ softmax(x_i)=\\exp({x_i - logsumexp(x)})\\tag{3}$$\n",
    "The work below will maintain a copy of the logsumexp allowing the softmax to be completed in sections. You will see how this is useful later in the LSHSelfAttention class.\n",
    "We'll create a routine to implement that here with the addition of a passthrough. The matrix operations we will be working on below are easier to follow if we can maintain integer values. So, for tests, we will skip the softmax in some cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def our_softmax(x, passthrough=False):\n",
    "    \"\"\" softmax with passthrough\"\"\"\n",
    "    logsumexp = fastmath.logsumexp(x, axis=-1, keepdims=True)\n",
    "    o = np.exp(x - logsumexp)\n",
    "    if passthrough:\n",
    "        return (x, np.zeros_like(logsumexp))\n",
    "    else:\n",
    "        return (o, logsumexp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check our implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0320586  0.08714432 0.2368828  0.6439142 ]\n",
      "[0.0320586  0.08714431 0.23688279 0.64391416]\n",
      "[4.44019]\n"
     ]
    }
   ],
   "source": [
    "## compare softmax(a) using both methods\n",
    "a = np.array([1.0, 2.0, 3.0, 4.0])\n",
    "sma = np.exp(a) / sum(np.exp(a))\n",
    "print(sma)\n",
    "sma2, a_logsumexp = our_softmax(a)\n",
    "print(sma2)\n",
    "print(a_logsumexp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The purpose of the dot-product is to 'focus attention' on some of the inputs. Dot now has entries appropriately scaled to enhance some values and reduce others. These are now applied to the $V$ entries.\n",
    "<img src = \"C4W4_LN2_image6.PNG\" height=\"300\" width=\"900\">\n",
    "\n",
    "<center><b>Figure 7: Applying Attention to $V$</b></center>\n",
    "\n",
    "$V$ is of size (n_seq,n_v). Note the shading in the diagram. This is to draw attention to the operation of the matrix multiplication. This is detailed below.\n",
    "\n",
    "<img src = \"C4W4_LN2_image7.PNG\" height=\"300\" width=\"600\"/>\n",
    "\n",
    "<center><b>Figure 7: The Matrix Multiply applies attention to the values of V</b></center>\n",
    "\n",
    "$V$ is formed by a matrix multiply of the input embedding with the weight matrix $W^v$ whose values were set by backpropagation. The row entries of $V$ are then related to the corresponding input embedding. The matrix multiply weights first column of V, representing a section of each of the input embeddings, with the first row of Dot, representing the similarity of $W_0$ and each word of the input embedding and deposits the value in $Z$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"2.2\"></a>\n",
    "### Part 2.2  our_simple_attend\n",
    "In this section we'll work on an implementation of *attend* whose operations you can see in figure 3. It is a slightly simplified version of the routine in [efficient_attention.py](https://github.com/google/trax/blob/v1.3.4/trax/layers/research/efficient_attention.py). We will fill in a few lines of code. The main goal is to become familiar with the routine. You have implemented similar functionality in a previous assignment.\n",
    "\n",
    "**Instructions**\n",
    "**Step 1:** matrix multiply (np.matmul) q and the k 'transpose' kr.\n",
    "**Step 2:** use our_softmax() to perform a softmax on masked output of the dot product, dots.\n",
    "**Step 3:** matrix multiply (np.matmul) dots and v."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UGvWTWQ4ExAx"
   },
   "outputs": [],
   "source": [
    "def our_simple_attend(\n",
    "    q,\n",
    "    k=None,\n",
    "    v=None,\n",
    "    mask_fn=None,\n",
    "    q_info=None,\n",
    "    kv_info=None,\n",
    "    dropout=0.0,\n",
    "    rng=None,\n",
    "    verbose=False,\n",
    "    passthrough=False,\n",
    "):\n",
    "    \"\"\"Dot-product attention,  with masking, without optional chunking and/or.\n",
    "\n",
    "  Args:\n",
    "    q: Query vectors, shape [q_len, d_qk]\n",
    "    k: Key vectors, shape [kv_len, d_qk]; or None\n",
    "    v: Value vectors, shape [kv_len, d_v]\n",
    "    mask_fn: a function reference that implements masking (e.g. mask_self_attention)\n",
    "    q_info: Query-associated metadata for masking\n",
    "    kv_info: Key-associated metadata for masking\n",
    "    dropout: Dropout rate\n",
    "    rng: RNG for dropout\n",
    "\n",
    "  Returns:\n",
    "    A tuple (output, dots_logsumexp). The output has shape [q_len, d_v], and\n",
    "    dots_logsumexp has shape [q_len]. The logsumexp of the attention\n",
    "    probabilities is useful for combining multiple rounds of attention (as in\n",
    "    LSH attention).\n",
    "  \"\"\"\n",
    "    assert v is not None\n",
    "    share_qk = k is None\n",
    "    if share_qk:\n",
    "        k = q\n",
    "        if kv_info is None:\n",
    "            kv_info = q_info\n",
    "\n",
    "    if share_qk:\n",
    "        k = length_normalized(k)\n",
    "    k = k / np.sqrt(k.shape[-1])\n",
    "\n",
    "    # Dot-product attention.\n",
    "    kr = np.swapaxes(k, -1, -2)  # note the fancy transpose for later..\n",
    "\n",
    "    ## Step 1  ##\n",
    "    dots = np.matmul(q, kr )\n",
    "    if verbose:\n",
    "        print(\"Our attend dots\", dots.shape)\n",
    "\n",
    "    # Masking\n",
    "    if mask_fn is not None:\n",
    "        dots = mask_fn(dots, q_info[..., :, None], kv_info[..., None, :])\n",
    "\n",
    "    # Softmax.\n",
    "    # dots_logsumexp = fastmath.logsumexp(dots, axis=-1, keepdims=True)  #original\n",
    "    # dots = np.exp(dots - dots_logsumexp)  #original\n",
    "    ## Step 2  ##\n",
    "    # replace with our_softmax()\n",
    "    dots, dots_logsumexp = our_softmax(dots, passthrough=passthrough)\n",
    "    if verbose:\n",
    "        print(\"Our attend dots post softmax\", dots.shape, dots_logsumexp.shape)\n",
    "\n",
    "    if dropout > 0.0:\n",
    "        assert rng is not our_softmax(dots, passthrough=passthrough)\n",
    "        # Dropout is broadcast across the bin dimension\n",
    "        dropout_shape = (dots.shape[-2], dots.shape[-1])\n",
    "        keep_prob = tie_in(dots, 1.0 - dropout)\n",
    "        keep = fastmath.random.bernoulli(rng, keep_prob, dropout_shape)\n",
    "        multiplier = keep.astype(dots.dtype) / tie_in(keep, keep_prob)\n",
    "        dots = dots * multiplier\n",
    "\n",
    "    ## Step 3  ##\n",
    "    # The softmax normalizer (dots_logsumexp) is used by multi-round LSH attn.\n",
    "    out = np.matmul(dots, v)\n",
    "    if verbose:\n",
    "        print(\"Our attend out1\", out.shape)\n",
    "    out = np.reshape(out, (-1, out.shape[-1]))\n",
    "    if verbose:\n",
    "        print(\"Our attend out2\", out.shape)\n",
    "    dots_logsumexp = np.reshape(dots_logsumexp, (-1,))\n",
    "    return out, dots_logsumexp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 202
    },
    "colab_type": "code",
    "id": "g1861d--trPg",
    "outputId": "58a8974e-e3c8-4ec7-92a0-530df96d6d71"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Our attend dots (8, 8)\n",
      "Our attend dots post softmax (8, 8) (8, 1)\n",
      "Our attend out1 (8, 4)\n",
      "Our attend out2 (8, 4)\n",
      "[[0.5606324  0.7290605  0.5251243  0.47101074]\n",
      " [0.5713517  0.71991956 0.5033342  0.46975708]\n",
      " [0.5622886  0.7288458  0.52172124 0.46318397]\n",
      " [0.5568317  0.72234154 0.542236   0.4699722 ]\n",
      " [0.56504494 0.72274375 0.5204978  0.47231334]\n",
      " [0.56175965 0.7216782  0.53293145 0.48003793]\n",
      " [0.56753993 0.72232544 0.5141734  0.46625748]\n",
      " [0.57100445 0.70785505 0.5325362  0.4590797 ]] \n",
      " [2.6512175 2.1914332 2.6630518 2.7792363 2.4583826 2.5421977 2.4145055\n",
      " 2.5111294]\n"
     ]
    }
   ],
   "source": [
    "seq_len = 8\n",
    "emb_len = 5\n",
    "d_qk = 3\n",
    "d_v = 4\n",
    "with fastmath.use_backend(\"jax\"):  # specify the backend for consistency\n",
    "    rng_attend = fastmath.random.get_prng(1)\n",
    "    q = k = jax.random.uniform(rng_attend, (seq_len, d_qk), dtype=np.float32)\n",
    "    v = jax.random.uniform(rng_attend, (seq_len, d_v), dtype=np.float32)\n",
    "    o, logits = our_simple_attend(\n",
    "        q,\n",
    "        k,\n",
    "        v,\n",
    "        mask_fn=None,\n",
    "        q_info=None,\n",
    "        kv_info=None,\n",
    "        dropout=0.0,\n",
    "        rng=rng_attend,\n",
    "        verbose=True,\n",
    "    )\n",
    "print(o, \"\\n\", logits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "a61Hs5AgH_fu"
   },
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\"><b> Expected Output </b></font>\n",
    "</summary>\n",
    "\n",
    "**Expected Output**\n",
    "```\n",
    "Our attend dots (8, 8)\n",
    "Our attend dots post softmax (8, 8) (8, 1)\n",
    "Our attend out1 (8, 4)\n",
    "Our attend out2 (8, 4)\n",
    "[[0.5606324  0.7290605  0.5251243  0.47101074]\n",
    " [0.5713517  0.71991956 0.5033342  0.46975708]\n",
    " [0.5622886  0.7288458  0.52172124 0.46318397]\n",
    " [0.5568317  0.72234154 0.542236   0.4699722 ]\n",
    " [0.56504494 0.72274375 0.5204978  0.47231334]\n",
    " [0.56175965 0.7216782  0.53293145 0.48003793]\n",
    " [0.56753993 0.72232544 0.5141734  0.46625748]\n",
    " [0.57100445 0.70785505 0.5325362  0.4590797 ]]\n",
    " [2.6512175 2.1914332 2.6630518 2.7792363 2.4583826 2.5421977 2.4145055\n",
    " 2.5111294]```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\"><b> completed code for reference </b></font>\n",
    "</summary>\n",
    "    This notebook is ungraded, so for reference, the completed code follows:\n",
    "\n",
    "```\n",
    "def our_simple_attend(\n",
    "    q, k=None, v=None,\n",
    "    mask_fn=None, q_info=None, kv_info=None,\n",
    "    dropout=0.0, rng=None, verbose=False, passthrough=False\n",
    "    ):\n",
    "  \"\"\"Dot-product attention,  with masking, without optional chunking and/or.\n",
    "\n",
    "  Args:\n",
    "    q: Query vectors, shape [q_len, d_qk]\n",
    "    k: Key vectors, shape [kv_len, d_qk]; or None\n",
    "    v: Value vectors, shape [kv_len, d_v]\n",
    "    mask_fn: a function reference that implements masking (e.g. mask_self_attention)\n",
    "    q_info: Query-associated metadata for masking\n",
    "    kv_info: Key-associated metadata for masking\n",
    "    dropout: Dropout rate\n",
    "    rng: RNG for dropout\n",
    "\n",
    "  Returns:\n",
    "    A tuple (output, dots_logsumexp). The output has shape [q_len, d_v], and\n",
    "    dots_logsumexp has shape [q_len]. The logsumexp of the attention\n",
    "    probabilities is useful for combining multiple rounds of attention (as in\n",
    "    LSH attention).\n",
    "  \"\"\"\n",
    "  assert v is not None\n",
    "  share_qk = (k is None)\n",
    "  if share_qk:\n",
    "    k = q\n",
    "    if kv_info is None:\n",
    "      kv_info = q_info\n",
    "\n",
    "  if share_qk:\n",
    "    k = length_normalized(k)\n",
    "  k = k / np.sqrt(k.shape[-1])\n",
    "\n",
    "  # Dot-product attention.\n",
    "  kr = np.swapaxes(k, -1, -2)  #note the fancy transpose for later..\n",
    "\n",
    "## Step 1  ##\n",
    "  dots = np.matmul(q, kr )\n",
    "  if verbose: print(\"Our attend dots\", dots.shape)\n",
    "\n",
    "  # Masking\n",
    "  if mask_fn is not None:\n",
    "    dots = mask_fn(dots, q_info[..., :, None], kv_info[..., None, :])\n",
    "\n",
    "  # Softmax.\n",
    "  #dots_logsumexp = fastmath.logsumexp(dots, axis=-1, keepdims=True)  #original\n",
    "  #dots = np.exp(dots - dots_logsumexp)  #original\n",
    "## Step 2  ##\n",
    "  #replace with our_softmax()\n",
    "  dots, dots_logsumexp = our_softmax(dots, passthrough=passthrough)\n",
    "  if verbose: print(\"Our attend dots post softmax\", dots.shape, dots_logsumexp.shape)\n",
    "\n",
    "  if dropout > 0.0:\n",
    "    assert rng is not None\n",
    "    # Dropout is broadcast across the bin dimension\n",
    "    dropout_shape = (dots.shape[-2], dots.shape[-1])\n",
    "    keep_prob = tie_in(dots, 1.0 - dropout)\n",
    "    keep = fastmath.random.bernoulli(rng, keep_prob, dropout_shape)\n",
    "    multiplier = keep.astype(dots.dtype) / tie_in(keep, keep_prob)\n",
    "    dots = dots * multiplier\n",
    "\n",
    "## Step 3  ##\n",
    "# The softmax normalizer (dots_logsumexp) is used by multi-round LSH attn.\n",
    "  out = np.matmul(dots, v)\n",
    "  if verbose: print(\"Our attend out1\", out.shape)\n",
    "  out = np.reshape(out, (-1, out.shape[-1]))\n",
    "  if verbose: print(\"Our attend out2\", out.shape)\n",
    "  dots_logsumexp = np.reshape(dots_logsumexp, (-1,))\n",
    "  return out, dots_logsumexp\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"2.3\"></a>\n",
    "## Part 2.3  Class OurSelfAttention\n",
    "Here we create our own self attention layer by creating a class `OurSelfAttention`. The parent class will be the tl.SelfAttention layer in Trax. We will only override the `forward_unbatched` routine.\n",
    "We're not asking you to modify anything in this routine. There are some comments to draw your attention to a few lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Wix7oM-i5NV3"
   },
   "outputs": [],
   "source": [
    "class OurSelfAttention(tl.SelfAttention):\n",
    "    \"\"\"Our self-attention. Just the Forward Function.\"\"\"\n",
    "\n",
    "    def forward_unbatched(\n",
    "        self, x, mask=None, *, weights, state, rng, update_state, verbose=False\n",
    "    ):\n",
    "        print(\"ourSelfAttention:forward_unbatched\")\n",
    "        del update_state\n",
    "        attend_rng, output_rng = fastmath.random.split(rng)\n",
    "        if self.bias:\n",
    "            if self.share_qk:\n",
    "                w_q, w_v, w_o, b_q, b_v = weights\n",
    "            else:\n",
    "                w_q, w_k, w_v, w_o, b_q, b_k, b_v = weights\n",
    "        else:\n",
    "            if self.share_qk:\n",
    "                w_q, w_v, w_o = weights\n",
    "            else:\n",
    "                w_q, w_k, w_v, w_o = weights\n",
    "\n",
    "        print(\"x.shape,w_q.shape\", x.shape, w_q.shape)\n",
    "        q = np.matmul(x, w_q)\n",
    "        k = None\n",
    "        if not self.share_qk:\n",
    "            k = np.matmul(x, w_k)\n",
    "        v = np.matmul(x, w_v)\n",
    "\n",
    "        if self.bias:\n",
    "            q = q + b_q\n",
    "            if not self.share_qk:\n",
    "                k = k + b_k\n",
    "            v = v + b_v\n",
    "\n",
    "        mask_fn = functools.partial(\n",
    "            mask_self_attention,\n",
    "            causal=self.causal,\n",
    "            exclude_self=self.share_qk,\n",
    "            masked=self.masked,\n",
    "        )\n",
    "        q_info = kv_info = tie_in(x, np.arange(q.shape[-2], dtype=np.int32))\n",
    "\n",
    "        assert (mask is not None) == self.masked\n",
    "        if self.masked:\n",
    "            # mask is a boolean array (True means \"is valid token\")\n",
    "            ones_like_mask = tie_in(x, np.ones_like(mask, dtype=np.int32))\n",
    "            kv_info = kv_info * np.where(mask, ones_like_mask, -ones_like_mask)\n",
    "\n",
    "        # Notice, we are callout our vesion of attend\n",
    "        o, _ = our_simple_attend(\n",
    "            q,\n",
    "            k,\n",
    "            v,\n",
    "            mask_fn=mask_fn,\n",
    "            q_info=q_info,\n",
    "            kv_info=kv_info,\n",
    "            dropout=self.attention_dropout,\n",
    "            rng=attend_rng,\n",
    "            verbose=True,\n",
    "        )\n",
    "\n",
    "        # Notice, wo weight matrix applied to output of attend in forward_unbatched\n",
    "        out = np.matmul(o, w_o)\n",
    "        out = apply_broadcasted_dropout(out, self.output_dropout, output_rng)\n",
    "        return out, state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "e1Krh9er8qTj",
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "causal = False\n",
    "masked = False\n",
    "mask = None\n",
    "attention_dropout = 0.0\n",
    "n_heads = 3\n",
    "d_qk = 3\n",
    "d_v = 4\n",
    "seq_len = 8\n",
    "emb_len = 5\n",
    "batch_size = 1\n",
    "\n",
    "osa = OurSelfAttention(\n",
    "    n_heads=n_heads,\n",
    "    d_qk=d_qk,\n",
    "    d_v=d_v,\n",
    "    causal=causal,\n",
    "    use_reference_code=True,\n",
    "    attention_dropout=attention_dropout,\n",
    "    mode=\"train\",\n",
    ")\n",
    "\n",
    "rng_osa = fastmath.random.get_prng(1)\n",
    "x = jax.random.uniform(\n",
    "    jax.random.PRNGKey(0), (batch_size, seq_len, emb_len), dtype=np.float32\n",
    ")\n",
    "_, _ = osa.init(tl.shapes.signature(x), rng=rng_osa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 790
    },
    "colab_type": "code",
    "id": "BTfRVNkL85i6",
    "outputId": "8a321eb9-09b8-4431-ecad-2290ea2310a3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ourSelfAttention:forward_unbatched\n",
      "x.shape,w_q.shape (8, 5) (5, 3)\n",
      "Our attend dots (8, 8)\n",
      "Our attend dots post softmax (8, 8) (8, 1)\n",
      "Our attend out1 (8, 4)\n",
      "Our attend out2 (8, 4)\n",
      "ourSelfAttention:forward_unbatched\n",
      "x.shape,w_q.shape (8, 5) (5, 3)\n",
      "Our attend dots (8, 8)\n",
      "Our attend dots post softmax (8, 8) (8, 1)\n",
      "Our attend out1 (8, 4)\n",
      "Our attend out2 (8, 4)\n",
      "ourSelfAttention:forward_unbatched\n",
      "x.shape,w_q.shape (8, 5) (5, 3)\n",
      "Our attend dots (8, 8)\n",
      "Our attend dots post softmax (8, 8) (8, 1)\n",
      "Our attend out1 (8, 4)\n",
      "Our attend out2 (8, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray([[[ 6.70414209e-01, -1.04319841e-01, -5.33822298e-01,\n",
       "                1.92711830e-01, -4.54187393e-05],\n",
       "              [ 6.64090097e-01, -1.01875424e-01, -5.35733163e-01,\n",
       "                1.88311756e-01, -6.30629063e-03],\n",
       "              [ 6.73380017e-01, -1.06952369e-01, -5.31989932e-01,\n",
       "                1.90056816e-01,  1.30271912e-03],\n",
       "              [ 6.84564888e-01, -1.13240272e-01, -5.50182462e-01,\n",
       "                1.95673436e-01,  5.47635555e-03],\n",
       "              [ 6.81435883e-01, -1.11068964e-01, -5.32343209e-01,\n",
       "                1.91912338e-01,  5.69400191e-03],\n",
       "              [ 6.80724978e-01, -1.08496904e-01, -5.34994125e-01,\n",
       "                1.96332246e-01,  5.89773059e-03],\n",
       "              [ 6.80933356e-01, -1.14087075e-01, -5.18659890e-01,\n",
       "                1.90674081e-01,  1.14096403e-02],\n",
       "              [ 6.80265009e-01, -1.09031796e-01, -5.38248718e-01,\n",
       "                1.94203183e-01,  4.23943996e-03]]], dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "osa(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Km8I7HMZ62Cs"
   },
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\"><b> Expected Output </b></font>\n",
    "</summary>\n",
    "\n",
    "**Expected Output**\n",
    "Notice a few things:\n",
    "* the w_q (and w_k) matrices are applied to each row or each embedding on the input. This is similar to the filter operation in convolution\n",
    "* forward_unbatched is called 3 times. This is because we have 3 heads in this example.\n",
    "\n",
    "```\n",
    "ourSelfAttention:forward_unbatched\n",
    "x.shape,w_q.shape (8, 5) (5, 3)\n",
    "Our attend dots (8, 8)\n",
    "Our attend dots post softmax (8, 8) (8, 1)\n",
    "Our attend out1 (8, 4)\n",
    "Our attend out2 (8, 4)\n",
    "ourSelfAttention:forward_unbatched\n",
    "x.shape,w_q.shape (8, 5) (5, 3)\n",
    "Our attend dots (8, 8)\n",
    "Our attend dots post softmax (8, 8) (8, 1)\n",
    "Our attend out1 (8, 4)\n",
    "Our attend out2 (8, 4)\n",
    "ourSelfAttention:forward_unbatched\n",
    "x.shape,w_q.shape (8, 5) (5, 3)\n",
    "Our attend dots (8, 8)\n",
    "Our attend dots post softmax (8, 8) (8, 1)\n",
    "Our attend out1 (8, 4)\n",
    "Our attend out2 (8, 4)\n",
    "DeviceArray([[[ 6.70414209e-01, -1.04319841e-01, -5.33822298e-01,\n",
    "                1.92711830e-01, -4.54187393e-05],\n",
    "              [ 6.64090097e-01, -1.01875424e-01, -5.35733163e-01,\n",
    "                1.88311756e-01, -6.30629063e-03],\n",
    "              [ 6.73380017e-01, -1.06952369e-01, -5.31989932e-01,\n",
    "                1.90056816e-01,  1.30271912e-03],\n",
    "              [ 6.84564888e-01, -1.13240272e-01, -5.50182462e-01,\n",
    "                1.95673436e-01,  5.47635555e-03],\n",
    "              [ 6.81435883e-01, -1.11068964e-01, -5.32343209e-01,\n",
    "                1.91912338e-01,  5.69400191e-03],\n",
    "              [ 6.80724978e-01, -1.08496904e-01, -5.34994125e-01,\n",
    "                1.96332246e-01,  5.89773059e-03],\n",
    "              [ 6.80933356e-01, -1.14087075e-01, -5.18659890e-01,\n",
    "                1.90674081e-01,  1.14096403e-02],\n",
    "              [ 6.80265009e-01, -1.09031796e-01, -5.38248718e-01,\n",
    "                1.94203183e-01,  4.23943996e-03]]], dtype=float32)\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"3\"></a>\n",
    "## Part 3.0 Trax LSHSelfAttention\n",
    "<a name=\"3.1\"></a>\n",
    "## Part 3.1 Description\n",
    "The larger the matrix multiply in the previous section is, the more context can be taken into account when making the next decision. However, the self attention dot product grows as the size of the input squared. For example, if one wished to have an input size of 1024, that would result in $1024^2$ or over a million dot products for each head! As a result, there has been significant research related to reducing the compute requirements. One such approach is Locality Sensitive Hashing(LSH) Self Attention.\n",
    "\n",
    "You may recall, earlier in the course you utilized LSH to find similar tweets without resorting to calculating cosine similarity for each pair of embeddings. We will use a similar approach here. It may be best described with an example.\n",
    "<img src = \"C4W4_LN2_image8.PNG\" height=\"400\" width=\"750\">\n",
    "\n",
    "<center><b>Figure 9: Example of LSH Self Attention</b></center>\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LSH Self attention uses Queries only, no Keys. Attention then generates a metric of the similarity of each value of Q relative to all the other values in Q. An earlier assignment demonstrated that values which hash to the same bucket are likely to be similar. Further, multiple random hashes can improve the chances of finding entries which are similar. This is the approach taken here, though the hash is implemented a bit differently. The values of Q are hashed into buckets using a randomly generated set of hash vectors.  Multiple sets of hash vectors are used, generating multiple hash tables. In the figure above, we have 3 hash tables with 4 buckets in each table. Notionally,  following the hash, the values of Q have been replicated 3 times and distributed to their appropriate bucket in each of the 3 tables. To find similarity then, one generates dot-products only between members of the buckets. The result of this operation provides information on which entries are similar. As the operation has been distributed over multiple hash tables, these results need to be combined to form a complete picture and this can be used to generate a reduced dot-product attention array. Its clear that because we do not do a compare of every value vs every other value, the size of *Dots* will be reduced.\n",
    "\n",
    "The challenge in this approach is getting it to operate efficiently. You may recall from the earlier assignments the buckets were lists of entries and had varying length. This will operate poorly on a vector processing machine such as a GPU or TPU. Ideally, operations are done in large blocks with uniform sizes. While it is straightforward to implement the hash algorithm this way, it is challenging to managed buckets and variable sized dot-products. This will be discussed further below. For now, we will examine and implement the hash function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"3.2\"></a>\n",
    "## Part 3.2 our_hash_vectors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*our_hash_vectors*, is a reimplementation of Trax *hashvector*. It takes in an array of vectors, hashes the entries and returns and array assigning each input vector to n_hash buckets. Hashing is described as creating *random rotations*, see [Practical and Optimal LSH for Angular Distance](https://arxiv.org/pdf/1509.02897.pdf).\n",
    "\n",
    "<img src = \"C4W4_LN2_image9.PNG\" height=\"400\" width=\"750\">\n",
    "<img src = \"C4W4_LN2_image10.PNG\" height=\"400\" width=\"750\">\n",
    "<center><b>Figure 10: Processing steps in our_hash_vectors </b></center>\n",
    "\n",
    "Note, in the diagram, sizes relate to our expected input $Q$ while our_hash_vectors is written assuming a generic input vector\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Instructions**\n",
    "**Step 1**\n",
    "create an array of random normal vectors which will be our hash vectors. Each vector will be hashed into a hash table and into `rot_size//2` buckets. We use `rot_size//2` to reduce computation. Later in the routine we will form the negative rotations with a simple negation and concatenate to get a full `rot_size` number of rotations.\n",
    "    * use fastmath.random.normal and create an array of random vectors of shape (vec.shape[-1],n_hashes, rot_size//2)\n",
    "\n",
    "**Step 2** In this step we simply do the matrix multiply. `jax` has an accelerated version of [einsum](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html). Here we will utilize more conventional routines.\n",
    "\n",
    "**Step 2x**\n",
    "    * 2a: np.reshape random_rotations into a 2 dimensional array ([-1, n_hashes * (rot_size // 2)])\n",
    "    * 2b: np.dot vecs and random_rotations forming our rotated_vecs\n",
    "    * 2c: back to 3 dimension with np.reshape  [-1, n_hashes, rot_size//2]\n",
    "    * 2d: prepare for concatenating by swapping dimensions np.transpose  (1, 0, 2)\n",
    "**Step 3** Here we concatenate our rotation vectors getting a fullrot_size number of buckets (note, n_buckets = rotsize)\n",
    "    * use np.concatenate, [rotated_vecs, -rotated_vecs], axis=-1\n",
    "**Step 4** **This is the exciting step!** You have no doubt been wondering how we will turn these vectors into bucket indexes. By performing np.argmax over the rotations for a given entry, you get the index to the best match! We will use this as a bucket index.\n",
    "    * np.argmax(...).astype(np.int32); be sure to use the correct axis!\n",
    "**Step 5** In this style of hashing, items which land in bucket 0 of hash table 0 are not necessarily similar to those landing in bucket 0 of hash table 1, so we keep them separate. We do this by offsetting the bucket numbers by 'n_buckets'.\n",
    "* add buckets and offsets and reshape into a one dimensional array\n",
    "This will return a 1D array of size n_hashes * vec.shape[0]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Vzf4dDSHnFQQ"
   },
   "outputs": [],
   "source": [
    "def our_hash_vectors(vecs, rng, n_buckets, n_hashes, mask=None, verbose=False):\n",
    "    \"\"\" \n",
    "  Args:\n",
    "    vecs: tensor of at least 2 dimension, \n",
    "    rng: random number generator\n",
    "    n_buckets: number of buckets in each hash table\n",
    "    n_hashes: the number of hash tables\n",
    "    mask: None indicating no mask or a 1D boolean array of length vecs.shape[0], containing the location of padding value\n",
    "    verbose: controls prints for debug\n",
    "  Returns:\n",
    "    A vector of size n_hashes * vecs.shape[0] containing the buckets associated with each input vector per hash table.\n",
    "    \n",
    "    \"\"\"\n",
    "\n",
    "    # check for even, integer bucket sizes\n",
    "    assert isinstance(n_buckets, int) and n_buckets % 2 == 0\n",
    "\n",
    "    rng = fastmath.stop_gradient(tie_in(vecs, rng))\n",
    "    rot_size = n_buckets\n",
    "    ### Start Code Here\n",
    "\n",
    "    ### Step 1 ###\n",
    "    rotations_shape = (vecs.shape[-1], n_hashes, rot_size // 2)\n",
    "    random_rotations = fastmath.random.normal(rng, rotations_shape).astype(np.float32)\n",
    "    if verbose:\n",
    "        print(\"random.rotations.shape\", random_rotations.shape)\n",
    "\n",
    "    ### Step 2 ###\n",
    "    if fastmath.backend_name() == \"jax\":\n",
    "        rotated_vecs = np.einsum(\"tf,fhb->htb\", vecs, random_rotations)\n",
    "        print(\"using jax\")\n",
    "    else:\n",
    "        # Step 2a\n",
    "        random_rotations = np.reshape(random_rotations,\n",
    "                                    [-1, n_hashes * (rot_size // 2)])\n",
    "        if verbose:\n",
    "            print(\"random_rotations reshaped\", random_rotations.shape)\n",
    "        # Step 2b\n",
    "        rotated_vecs = np.dot(vecs, random_rotations)\n",
    "        if verbose:\n",
    "            print(\"rotated_vecs1\", rotated_vecs.shape)\n",
    "        # Step 2c\n",
    "        rotated_vecs = np.reshape(rotated_vecs, [-1, n_hashes, rot_size//2])\n",
    "        if verbose:\n",
    "            print(\"rotated_vecs2\", rotated_vecs.shape)\n",
    "        # Step 2d\n",
    "        rotated_vecs = np.transpose(rotated_vecs, (1, 0, 2))\n",
    "        if verbose:\n",
    "            print(\"rotated_vecs3\", rotated_vecs.shape)\n",
    "\n",
    "    ### Step 3 ###\n",
    "    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)\n",
    "    if verbose:\n",
    "        print(\"rotated_vecs.shape\", rotated_vecs.shape)\n",
    "    ### Step 4 ###\n",
    "    buckets = np.argmax(rotated_vecs, axis=-1).astype(np.int32)\n",
    "    if verbose:\n",
    "        print(\"buckets.shape\", buckets.shape)\n",
    "    if verbose:\n",
    "        print(\"buckets\", buckets)\n",
    "\n",
    "    if mask is not None:\n",
    "        n_buckets += 1  # Create an extra bucket for padding tokens only\n",
    "        buckets = np.where(mask[None, :], buckets, n_buckets - 1)\n",
    "\n",
    "    # buckets is now (n_hashes, seqlen). Next we add offsets so that\n",
    "    # bucket numbers from different hashing rounds don't overlap.\n",
    "    offsets = tie_in(buckets, np.arange(n_hashes, dtype=np.int32))\n",
    "    offsets = np.reshape(offsets * n_buckets, (-1, 1))\n",
    "    ### Step 5 ###\n",
    "    buckets = np.reshape(buckets + offsets, (-1,))\n",
    "    if verbose:\n",
    "        print(\"buckets with offsets\", buckets.shape, \"\\n\", buckets)\n",
    "    ### End Code Here\n",
    "    return buckets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 403
    },
    "colab_type": "code",
    "id": "cPGNaVpAi8wM",
    "outputId": "a5a6a956-30b7-4de7-a5c2-65011a9d3816"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "random.rotations.shape (5, 3, 2)\n",
      "random_rotations reshaped (5, 6)\n",
      "rotated_vecs1 (8, 6)\n",
      "rotated_vecs2 (8, 3, 2)\n",
      "rotated_vecs3 (3, 8, 2)\n",
      "rotated_vecs.shape (3, 8, 4)\n",
      "buckets.shape (3, 8)\n",
      "buckets ndarray<tf.Tensor(\n",
      "[[3 3 3 3 3 3 3 3]\n",
      " [3 3 3 3 3 3 3 3]\n",
      " [3 3 3 3 3 3 3 3]], shape=(3, 8), dtype=int32)>\n",
      "buckets with offsets (24,) \n",
      " ndarray<tf.Tensor([ 3  3  3  3  3  3  3  3  7  7  7  7  7  7  7  7 11 11 11 11 11 11 11 11], shape=(24,), dtype=int32)>\n",
      "ohv shape (24,) \n",
      "ohv ndarray<tf.Tensor([ 3  3  3  3  3  3  3  3  7  7  7  7  7  7  7  7 11 11 11 11 11 11 11 11], shape=(24,), dtype=int32)>\n",
      "using jax\n",
      "ohv shape (24,) \n",
      "ohv [ 3  3  3  3  3  3  3  3  5  5  5  5  5  5  5  5 11 11 11 11 11 11 11 11]\n"
     ]
    }
   ],
   "source": [
    "# example code. Note for reference, the sizes in this example match the values in the diagram above.\n",
    "ohv_q = np.ones((8, 5))  # (seq_len=8, n_q=5)\n",
    "ohv_n_buckets = 4  # even number\n",
    "ohv_n_hashes = 3\n",
    "with fastmath.use_backend(\"tf\"):\n",
    "    ohv_rng = fastmath.random.get_prng(1)\n",
    "    ohv = our_hash_vectors(\n",
    "        ohv_q, ohv_rng, ohv_n_buckets, ohv_n_hashes, mask=None, verbose=True\n",
    "    )\n",
    "    print(\"ohv shape\", ohv.shape, \"\\nohv\", ohv)  # (ohv_n_hashes * ohv_n_buckets)\n",
    "# note the random number generators do not produce the same results with different backends\n",
    "with fastmath.use_backend(\"jax\"):\n",
    "    ohv_rng = fastmath.random.get_prng(1)\n",
    "    ohv = our_hash_vectors(ohv_q, ohv_rng, ohv_n_buckets, ohv_n_hashes, mask=None)\n",
    "    print(\"ohv shape\", ohv.shape, \"\\nohv\", ohv)  # (ohv_n_hashes * ohv_n_buckets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XAQqr1_XkCf6"
   },
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\"><b> Expected Output </b></font>\n",
    "</summary>\n",
    "\n",
    "**Expected Values**\n",
    "```\n",
    "random.rotations.shape (5, 3, 2)\n",
    "random_rotations reshaped (5, 6)\n",
    "rotated_vecs1 (8, 6)\n",
    "rotated_vecs2 (8, 3, 2)\n",
    "rotated_vecs3 (3, 8, 2)\n",
    "rotated_vecs.shape (3, 8, 4)\n",
    "buckets.shape (3, 8)\n",
    "buckets ndarray<tf.Tensor(\n",
    "[[3 3 3 3 3 3 3 3]\n",
    " [3 3 3 3 3 3 3 3]\n",
    " [3 3 3 3 3 3 3 3]], shape=(3, 8), dtype=int32)>\n",
    "buckets with offsets (24,)\n",
    " ndarray<tf.Tensor([ 3  3  3  3  3  3  3  3  7  7  7  7  7  7  7  7 11 11 11 11 11 11 11 11], shape=(24,), dtype=int32)>\n",
    "ohv shape (24,)\n",
    "ohv ndarray<tf.Tensor([ 3  3  3  3  3  3  3  3  7  7  7  7  7  7  7  7 11 11 11 11 11 11 11 11], shape=(24,), dtype=int32)>\n",
    "using jax\n",
    "ohv shape (24,)\n",
    "ohv [ 3  3  3  3  3  3  3  3  5  5  5  5  5  5  5  5 11 11 11 11 11 11 11 11]```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\" ><b>Completed code for reference </b></font>\n",
    "</summary>\n",
    "\n",
    "```\n",
    "# since this notebook is ungraded the completed code is provided here for reference\n",
    "\n",
    "def our_hash_vectors(vecs, rng, n_buckets, n_hashes, mask=None, verbose=False):\n",
    "    \"\"\"\n",
    "  Args:\n",
    "    vecs: tensor of at least 2 dimension,\n",
    "    rng: random number generator\n",
    "    n_buckets: number of buckets in each hash table\n",
    "    n_hashes: the number of hash tables\n",
    "    mask: None indicating no mask or a 1D boolean array of length vecs.shape[0], containing the location of padding value\n",
    "    verbose: controls prints for debug\n",
    "  Returns:\n",
    "    A vector of size n_hashes * vecs.shape[0] containing the buckets associated with each input vector per hash table.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # check for even, integer bucket sizes\n",
    "    assert isinstance(n_buckets, int) and n_buckets % 2 == 0\n",
    "\n",
    "    rng = fastmath.stop_gradient(tie_in(vecs, rng))\n",
    "    rot_size = n_buckets\n",
    "    ### Start Code Here\n",
    "\n",
    "    ### Step 1 ###\n",
    "    rotations_shape = (vecs.shape[-1], n_hashes, rot_size // 2)\n",
    "    random_rotations = fastmath.random.normal(rng, rotations_shape).astype(\n",
    "        np.float32)\n",
    "    if verbose: print(\"random.rotations.shape\", random_rotations.shape)\n",
    "\n",
    "    ### Step 2 ###\n",
    "    if fastmath.backend_name() == 'jax':\n",
    "      rotated_vecs = np.einsum('tf,fhb->htb', vecs, random_rotations)\n",
    "      if verbose: print(\"using jax\")\n",
    "    else:\n",
    "      #Step 2a\n",
    "      random_rotations = np.reshape(random_rotations,\n",
    "                                    [-1, n_hashes * (rot_size // 2)])\n",
    "      if verbose: print(\"random_rotations reshaped\", random_rotations.shape)\n",
    "      #Step 2b\n",
    "      rotated_vecs = np.dot(vecs, random_rotations)\n",
    "      if verbose: print(\"rotated_vecs1\", rotated_vecs.shape)\n",
    "      #Step 2c\n",
    "      rotated_vecs = np.reshape(rotated_vecs, [-1, n_hashes, rot_size//2])\n",
    "      if verbose: print(\"rotated_vecs2\", rotated_vecs.shape)\n",
    "      #Step 2d\n",
    "      rotated_vecs = np.transpose(rotated_vecs, (1, 0, 2))\n",
    "      if verbose: print(\"rotated_vecs3\", rotated_vecs.shape)\n",
    "\n",
    "    ### Step 3 ###\n",
    "    rotated_vecs = np.concatenate([rotated_vecs, -rotated_vecs], axis=-1)\n",
    "    if verbose: print(\"rotated_vecs.shape\", rotated_vecs.shape)\n",
    "    ### Step 4 ###\n",
    "    buckets = np.argmax(rotated_vecs, axis=-1).astype(np.int32)\n",
    "    if verbose: print(\"buckets.shape\", buckets.shape)\n",
    "    if verbose: print(\"buckets\", buckets)\n",
    "\n",
    "    if mask is not None:\n",
    "      n_buckets += 1  # Create an extra bucket for padding tokens only\n",
    "      buckets = np.where(mask[None, :], buckets, n_buckets - 1)\n",
    "\n",
    "    # buckets is now (n_hashes, seqlen). Next we add offsets so that\n",
    "    # bucket numbers from different hashing rounds don't overlap.\n",
    "    offsets = tie_in(buckets, np.arange(n_hashes, dtype=np.int32))\n",
    "    offsets = np.reshape(offsets * n_buckets, (-1, 1))\n",
    "    ### Step 5 ###\n",
    "    buckets = np.reshape(buckets + offsets, (-1,))\n",
    "    if verbose: print(\"buckets with offsets\", buckets.shape, \"\\n\", buckets)\n",
    "    return buckets```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"3.3\"></a>\n",
    "## Part 3.3 Sorting Buckets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! Now that we have a hash function, we can work on sorting our buckets and performing our matrix operations.\n",
    " We'll walk through this algorithm in small steps:\n",
    "* sort_buckets - we'll perform the sort\n",
    "* softmax\n",
    "* dotandv - do the matrix math to form the dotproduct and output\n",
    "These routines will demonstrate a simplified version of the algorithm. We won't address masking and variable bucket sizes but will consider how they would be handled.\n",
    "\n",
    "**sort_buckets**\n",
    "\n",
    "At this point, we have called the hash function and were returned the associated buckets. For example, if we started with\n",
    "`q[n_seq,n_q]`, with `n_hash = 2; n_buckets = 4; n_seq = 8`\n",
    "we might be returned:\n",
    "`bucket = [0,1,2,3,0,1,2,3, 4,5,6,7,4,5,6,7]  `\n",
    "Note that it is n_hash\\*n_seq long and that the bucket values for each hash have been offset by n_hash so the numbers do not overlap. Going forward, we going to sort this array of buckets to group together members of the same (hash,bucket) pair.\n",
    "\n",
    "**Instructions**\n",
    "**Step 1** Our goal is to sort $q$ rather than the bucket list, so we will need to track the association of the buckets to their elements in $q$.\n",
    "* using np.arange, create `ticker`, just a sequence of numbers (0..n_hashed * seqlen) associating members of q with their bucket.\n",
    "\n",
    "**Step 2** This step is provided to you as it is a bit difficult to describe. We want to disambiguate elements that map to the same bucket. When a sorting routine encounters a situation where multiple entries have the same value, it can correctly choose any entry to go first. This makes testing ambiguous. This prevents that. We multiply all the buckets by `seqlen` and then add `ticker % seqlen`\n",
    "\n",
    "**Step 3** Here we are! Ready to sort. This is the exciting part.\n",
    "* Utilize [fastmath.sort_key_val](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.sort_key_val.html#jax.lax.sort_key_val) and sort `buckets_and_t` and `ticker`.\n",
    "\n",
    "**Step 4** We need to be able to undo the sort at the end to get things back into their correct locations\n",
    "* sort `sticker` and `ticker` to for the reverse map\n",
    "\n",
    "**Step 5** create our sorted q and sorted v\n",
    "* use [np.take](https://numpy.org/doc/stable/reference/generated/numpy.take.html) and `st` to grab correct values in `q` for the sorted values, `sq`. Use axis=0.\n",
    "\n",
    "Use the example code below the routine to check and help debug your results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sort_buckets(buckets, q, v, n_buckets, n_hashes, seqlen, verbose=True):\n",
    "    \"\"\" \n",
    "  Args:\n",
    "    buckets: tensor of at least 2 dimension, \n",
    "    n_buckets: number of buckets in each hash table\n",
    "    n_hashes: the number of hash tables    \n",
    "    \"\"\"\n",
    "    if verbose:\n",
    "        print(\"---sort_buckets--\")\n",
    "    ## Step 1\n",
    "    ticker = np.arange(n_hashes * seqlen)\n",
    "    if verbose:\n",
    "        print(\"ticker\", ticker.shape, ticker)\n",
    "    ## Step 2\n",
    "    buckets_and_t = seqlen * buckets + (ticker % seqlen)  # provided\n",
    "    if verbose:\n",
    "        print(\"buckets_and_t\", buckets_and_t.shape, buckets_and_t)\n",
    "\n",
    "    # Hash-based sort (\"s\" at the start of variable names means \"sorted\")\n",
    "    # Step 3\n",
    "    sbuckets_and_t, sticker = fastmath.sort_key_val(buckets_and_t, ticker, dimension=-1)\n",
    "    if verbose:\n",
    "        print(\"sbuckets_and_t\", sbuckets_and_t.shape, sbuckets_and_t)\n",
    "    if verbose:\n",
    "        print(\"sticker\", sticker.shape, sticker)\n",
    "    # Step 4\n",
    "    _, undo_sort = fastmath.sort_key_val(sticker, ticker, dimension=-1)\n",
    "    if verbose:\n",
    "        print(\"undo_sort\", undo_sort.shape, undo_sort)\n",
    "\n",
    "    # Step 5\n",
    "    st = sticker % seqlen  # provided\n",
    "    sq = np.take(q, st, axis=0)\n",
    "    sv = np.take(v, st, axis=0)\n",
    "    return sq, sv, sticker, undo_sort"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "q\n",
      " [[0. 0. 0.]\n",
      " [1. 1. 1.]\n",
      " [2. 2. 2.]\n",
      " [3. 3. 3.]\n",
      " [0. 0. 0.]\n",
      " [1. 1. 1.]\n",
      " [2. 2. 2.]\n",
      " [3. 3. 3.]]\n",
      "t_buckets:  [0 1 2 3 0 1 2 3 4 5 6 7 4 5 6 7]\n",
      "---sort_buckets--\n",
      "ticker (16,) [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15]\n",
      "buckets_and_t (16,) [ 0  9 18 27  4 13 22 31 32 41 50 59 36 45 54 63]\n",
      "sbuckets_and_t (16,) [ 0  4  9 13 18 22 27 31 32 36 41 45 50 54 59 63]\n",
      "sticker (16,) [ 0  4  1  5  2  6  3  7  8 12  9 13 10 14 11 15]\n",
      "undo_sort (16,) [ 0  2  4  6  1  3  5  7  8 10 12 14  9 11 13 15]\n",
      "sq.shape (16, 3) sv.shape (16, 5)\n",
      "sq\n",
      " [[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [1. 1. 1.]\n",
      " [1. 1. 1.]\n",
      " [2. 2. 2.]\n",
      " [2. 2. 2.]\n",
      " [3. 3. 3.]\n",
      " [3. 3. 3.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [1. 1. 1.]\n",
      " [1. 1. 1.]\n",
      " [2. 2. 2.]\n",
      " [2. 2. 2.]\n",
      " [3. 3. 3.]\n",
      " [3. 3. 3.]]\n"
     ]
    }
   ],
   "source": [
    "t_n_hashes = 2\n",
    "t_n_buckets = 4\n",
    "t_n_seq = t_seqlen = 8\n",
    "t_n_q = 3\n",
    "n_v = 5\n",
    "\n",
    "t_q = (np.array([(j % t_n_buckets) for j in range(t_n_seq)]) * np.ones((t_n_q, 1))).T\n",
    "t_v = np.ones((t_n_seq, n_v))\n",
    "t_buckets = np.array(\n",
    "    [\n",
    "        (j % t_n_buckets) + t_n_buckets * i\n",
    "        for i in range(t_n_hashes)\n",
    "        for j in range(t_n_seq)\n",
    "    ]\n",
    ")\n",
    "print(\"q\\n\", t_q)\n",
    "print(\"t_buckets: \", t_buckets)\n",
    "\n",
    "t_sq, t_sv, t_sticker, t_undo_sort = sort_buckets(\n",
    "    t_buckets, t_q, t_v, t_n_buckets, t_n_hashes, t_seqlen, verbose=True\n",
    ")\n",
    "\n",
    "print(\"sq.shape\", t_sq.shape, \"sv.shape\", t_sv.shape)\n",
    "print(\"sq\\n\", t_sq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\"><b> Expected Output </b></font>\n",
    "</summary>\n",
    "\n",
    "**Expected Values**\n",
    "```\n",
    "q\n",
    " [[0. 0. 0.]\n",
    " [1. 1. 1.]\n",
    " [2. 2. 2.]\n",
    " [3. 3. 3.]\n",
    " [0. 0. 0.]\n",
    " [1. 1. 1.]\n",
    " [2. 2. 2.]\n",
    " [3. 3. 3.]]\n",
    "t_buckets:  [0 1 2 3 0 1 2 3 4 5 6 7 4 5 6 7]\n",
    "---sort_buckets--\n",
    "ticker (16,) [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15]\n",
    "buckets_and_t (16,) [ 0  9 18 27  4 13 22 31 32 41 50 59 36 45 54 63]\n",
    "sbuckets_and_t (16,) [ 0  4  9 13 18 22 27 31 32 36 41 45 50 54 59 63]\n",
    "sticker (16,) [ 0  4  1  5  2  6  3  7  8 12  9 13 10 14 11 15]\n",
    "undo_sort (16,) [ 0  2  4  6  1  3  5  7  8 10 12 14  9 11 13 15]\n",
    "sq.shape (16, 3) sv.shape (16, 5)\n",
    "sq\n",
    " [[0. 0. 0.]\n",
    " [0. 0. 0.]\n",
    " [1. 1. 1.]\n",
    " [1. 1. 1.]\n",
    " [2. 2. 2.]\n",
    " [2. 2. 2.]\n",
    " [3. 3. 3.]\n",
    " [3. 3. 3.]\n",
    " [0. 0. 0.]\n",
    " [0. 0. 0.]\n",
    " [1. 1. 1.]\n",
    " [1. 1. 1.]\n",
    " [2. 2. 2.]\n",
    " [2. 2. 2.]\n",
    " [3. 3. 3.]\n",
    " [3. 3. 3.]]\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\" ><b>Completed code for reference </b></font>\n",
    "</summary>\n",
    "\n",
    "```\n",
    "# since this notebook is ungraded the completed code is provided here for reference\n",
    "def sort_buckets(buckets, q, v, n_buckets, n_hashes, seqlen, verbose=True):\n",
    "    \"\"\"\n",
    "  Args:\n",
    "    buckets: tensor of at least 2 dimension,\n",
    "    n_buckets: number of buckets in each hash table\n",
    "    n_hashes: the number of hash tables\n",
    "    \"\"\"\n",
    "    if verbose: print(\"---sort_buckets--\")\n",
    "    ## Step 1\n",
    "    ticker = np.arange(n_hashes * seqlen)\n",
    "    if verbose: print(\"ticker\",ticker.shape, ticker)\n",
    "    ## Step 2\n",
    "    buckets_and_t = seqlen * buckets + (ticker % seqlen)\n",
    "    if verbose: print(\"buckets_and_t\",buckets_and_t.shape, buckets_and_t)\n",
    "\n",
    "    # Hash-based sort (\"s\" at the start of variable names means \"sorted\")\n",
    "    #Step 3\n",
    "    sbuckets_and_t, sticker = fastmath.sort_key_val(\n",
    "    buckets_and_t, ticker, dimension=-1)\n",
    "    if verbose: print(\"sbuckets_and_t\",sbuckets_and_t.shape, sbuckets_and_t)\n",
    "    if verbose: print(\"sticker\",sticker.shape, sticker)\n",
    "    #Step 4\n",
    "    _, undo_sort = fastmath.sort_key_val(sticker, ticker, dimension=-1)\n",
    "    if verbose: print(\"undo_sort\",undo_sort.shape, undo_sort)\n",
    "\n",
    "    #Step 4\n",
    "    st = (sticker % seqlen)\n",
    "    sq = np.take(q, st, axis=0)\n",
    "    sv = np.take(v, st, axis=0)\n",
    "    return sq, sv, sticker, undo_sort\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"3.4\"></a>\n",
    "## Part 3.4 Chunked dot product attention"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's create the dot product attention. We have sorted $Q$ so that elements that the hash has determined are likely to be similar are adjacent to each other. We now want to perform the dot-product within those limited regions - in 'chunks'.\n",
    "\n",
    "<img src = \"C4W4_LN2_image12.PNG\" height=\"400\" width=\"750\">\n",
    "<center><b>Figure 11: Performing dot product in 'chunks' </b></center>\n",
    "\n",
    "\n",
    "The example we have been working on is shown above, with sequences of 8, 2 hashes, 4 buckets and, conveniently, the content of Q was such that when sorted, there were 2 entries in each bucket. If we reshape Q into a (8,2,n_q), we can use numpy matmul to perform the operation. Numpy [matmul](https://numpy.org/doc/stable/reference/generated/numpy.matmul.html) will treat the inputs as a stack of matrices residing in the last two indexes. This will allow us to matrix multiply Q with itself in *chunks* and later can also be used to perform the matrix multiply with v.\n",
    "\n",
    "We will perform a softmax on the output of the dot product of Q and Q, but in this case, there is a bit more to the story. Recall the output of the hash had multiple hash tables. We will perform softmax on those separately and then must combine them. This is where the form of softmax we defined at the top of the notebook comes into play. The routines below will utilize the logsumexp values that the `our_softmax` routine calculates.\n",
    "\n",
    "There is a good deal of [reshaping](https://numpy.org/doc/stable/reference/generated/numpy.reshape.html) to get things into the right formats. The code has many print statements that match the expected values below. You can use those to check your work as you go along. If you don't do a lot of 3-dimensional matrix multiplications in your daily life, it might be worthwhile to open a spare cell and practice a few simple examples to get the hang of it! Here is one to start with:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(8, 2, 3)\n"
     ]
    }
   ],
   "source": [
    "a = np.arange(16 * 3).reshape((16, 3))\n",
    "chunksize = 2\n",
    "ar = np.reshape(\n",
    "    a, (-1, chunksize, a.shape[-1])\n",
    ")  # the -1 usage is very handy, see numpy reshape\n",
    "print(ar.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "**Instructions**\n",
    "**Step 1** Reshaping Q\n",
    "* np.reshape `sq` (sorted q) to be 3 dimensions. The middle dimension is the size of the 'chunk' specified by `kv_chunk_len`\n",
    "* np.swapaxes to perform a 'transpose' on the reshaped `sq`, *but only on the last two dimension*\n",
    "* np.matmul the two values.\n",
    "\n",
    "**Step 2**\n",
    "* use our_softmax to perform the softmax on the dot product. Don't forget `passthrough`\n",
    "\n",
    "**Step 3**\n",
    "* np.reshape `sv`. Like `sq`, the middle dimension is the size of the 'chunk' specified by `kv_chunk_len`\n",
    "* np.matmul dotlike and the reshaped `sv`\n",
    "* np.reshape so to a two dimensional array with the last dimension stays the same (`so.shape[-1]`)\n",
    "* `logits` also needs reshaping, we'll do that.\n",
    "\n",
    "**Step 4** Now we can undo the sort.\n",
    "* use [np.take](https://numpy.org/doc/stable/reference/generated/numpy.take.html) and `undo_sort` and axis = 0 to unsort so\n",
    "* do the same with `slogits`.\n",
    "\n",
    "**Step 5** This step combines the results of multiple hashes. Recall, the softmax was only over the values in one hash, this extends it to all the hashes. Read through it, the code is provided. Note this is taking place *after* the matrix multiply with v while the softmax output is used before the multiply. How does this achieve the correct result?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dotandv(\n",
    "    sq, sv, undo_sort, kv_chunk_len, n_hashes, seqlen, passthrough, verbose=False\n",
    "):\n",
    "    # Step 1\n",
    "    rsq = np.reshape(sq,(-1, kv_chunk_len, sq.shape[-1]))\n",
    "    rsqt = np.swapaxes(rsq, -1, -2)\n",
    "    if verbose:\n",
    "        print(\"rsq.shape,rsqt.shape: \", rsq.shape, rsqt.shape)\n",
    "    dotlike = np.matmul(rsq, rsqt)\n",
    "    if verbose:\n",
    "        print(\"dotlike\\n\", dotlike)\n",
    "\n",
    "    # Step 2\n",
    "    dotlike, slogits = our_softmax(dotlike, passthrough)\n",
    "    if verbose:\n",
    "        print(\"dotlike post softmax\\n\", dotlike)\n",
    "\n",
    "    # Step 3\n",
    "    vr = np.reshape(sv, (-1, kv_chunk_len, sv.shape[-1]))\n",
    "    if verbose:\n",
    "        print(\"dotlike.shape, vr.shape:\", dotlike.shape, vr.shape)\n",
    "    so = np.matmul(dotlike, vr)\n",
    "    if verbose: print(\"so.shape:\", so.shape)\n",
    "    so = np.reshape(so, (-1, so.shape[-1]))\n",
    "    slogits = np.reshape(slogits, (-1,))  # provided\n",
    "    if verbose: print(\"so.shape,slogits.shape\", so.shape, slogits.shape)\n",
    "\n",
    "    #Step 4\n",
    "    o = np.take(so, undo_sort, axis=0)\n",
    "    logits = np.take(slogits, undo_sort, axis=0)\n",
    "    if verbose: print(\"o.shape,o\", o.shape, o)\n",
    "    if verbose: print(\"logits.shape, logits\", logits.shape, logits)\n",
    "\n",
    "    #Step 5 (Provided)\n",
    "    if n_hashes > 1:\n",
    "      o = np.reshape(o, (n_hashes, seqlen, o.shape[-1]))\n",
    "      logits = np.reshape(logits, (n_hashes, seqlen, 1))\n",
    "      probs = np.exp(logits - fastmath.logsumexp(logits, axis=0, keepdims=True))\n",
    "      o = np.sum(o * probs, axis=0)\n",
    "\n",
    "    return(o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rsq.shape,rsqt.shape:  (8, 2, 3) (8, 3, 2)\n",
      "dotlike\n",
      " [[[ 0.  0.]\n",
      "  [ 0.  0.]]\n",
      "\n",
      " [[ 3.  3.]\n",
      "  [ 3.  3.]]\n",
      "\n",
      " [[12. 12.]\n",
      "  [12. 12.]]\n",
      "\n",
      " [[27. 27.]\n",
      "  [27. 27.]]\n",
      "\n",
      " [[ 0.  0.]\n",
      "  [ 0.  0.]]\n",
      "\n",
      " [[ 3.  3.]\n",
      "  [ 3.  3.]]\n",
      "\n",
      " [[12. 12.]\n",
      "  [12. 12.]]\n",
      "\n",
      " [[27. 27.]\n",
      "  [27. 27.]]]\n",
      "dotlike post softmax\n",
      " [[[ 0.  0.]\n",
      "  [ 0.  0.]]\n",
      "\n",
      " [[ 3.  3.]\n",
      "  [ 3.  3.]]\n",
      "\n",
      " [[12. 12.]\n",
      "  [12. 12.]]\n",
      "\n",
      " [[27. 27.]\n",
      "  [27. 27.]]\n",
      "\n",
      " [[ 0.  0.]\n",
      "  [ 0.  0.]]\n",
      "\n",
      " [[ 3.  3.]\n",
      "  [ 3.  3.]]\n",
      "\n",
      " [[12. 12.]\n",
      "  [12. 12.]]\n",
      "\n",
      " [[27. 27.]\n",
      "  [27. 27.]]]\n",
      "dotlike.shape, vr.shape: (8, 2, 2) (8, 2, 5)\n",
      "so.shape: (8, 2, 5)\n",
      "so.shape,slogits.shape (16, 5) (16,)\n",
      "o.shape,o (16, 5) [[ 0.  0.  0.  0.  0.]\n",
      " [ 6.  6.  6.  6.  6.]\n",
      " [24. 24. 24. 24. 24.]\n",
      " [54. 54. 54. 54. 54.]\n",
      " [ 0.  0.  0.  0.  0.]\n",
      " [ 6.  6.  6.  6.  6.]\n",
      " [24. 24. 24. 24. 24.]\n",
      " [54. 54. 54. 54. 54.]\n",
      " [ 0.  0.  0.  0.  0.]\n",
      " [ 6.  6.  6.  6.  6.]\n",
      " [24. 24. 24. 24. 24.]\n",
      " [54. 54. 54. 54. 54.]\n",
      " [ 0.  0.  0.  0.  0.]\n",
      " [ 6.  6.  6.  6.  6.]\n",
      " [24. 24. 24. 24. 24.]\n",
      " [54. 54. 54. 54. 54.]]\n",
      "logits.shape, logits (16,) [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
      "out\n",
      " [[ 0.  0.  0.  0.  0.]\n",
      " [ 6.  6.  6.  6.  6.]\n",
      " [24. 24. 24. 24. 24.]\n",
      " [54. 54. 54. 54. 54.]\n",
      " [ 0.  0.  0.  0.  0.]\n",
      " [ 6.  6.  6.  6.  6.]\n",
      " [24. 24. 24. 24. 24.]\n",
      " [54. 54. 54. 54. 54.]]\n",
      "\n",
      "-----With softmax enabled----\n",
      "\n",
      "rsq.shape,rsqt.shape:  (8, 2, 3) (8, 3, 2)\n",
      "dotlike\n",
      " [[[ 0.  0.]\n",
      "  [ 0.  0.]]\n",
      "\n",
      " [[ 3.  3.]\n",
      "  [ 3.  3.]]\n",
      "\n",
      " [[12. 12.]\n",
      "  [12. 12.]]\n",
      "\n",
      " [[27. 27.]\n",
      "  [27. 27.]]\n",
      "\n",
      " [[ 0.  0.]\n",
      "  [ 0.  0.]]\n",
      "\n",
      " [[ 3.  3.]\n",
      "  [ 3.  3.]]\n",
      "\n",
      " [[12. 12.]\n",
      "  [12. 12.]]\n",
      "\n",
      " [[27. 27.]\n",
      "  [27. 27.]]]\n",
      "dotlike post softmax\n",
      " [[[0.5        0.5       ]\n",
      "  [0.5        0.5       ]]\n",
      "\n",
      " [[0.5        0.5       ]\n",
      "  [0.5        0.5       ]]\n",
      "\n",
      " [[0.49999976 0.49999976]\n",
      "  [0.49999976 0.49999976]]\n",
      "\n",
      " [[0.49999976 0.49999976]\n",
      "  [0.49999976 0.49999976]]\n",
      "\n",
      " [[0.5        0.5       ]\n",
      "  [0.5        0.5       ]]\n",
      "\n",
      " [[0.5        0.5       ]\n",
      "  [0.5        0.5       ]]\n",
      "\n",
      " [[0.49999976 0.49999976]\n",
      "  [0.49999976 0.49999976]]\n",
      "\n",
      " [[0.49999976 0.49999976]\n",
      "  [0.49999976 0.49999976]]]\n",
      "dotlike.shape, vr.shape: (8, 2, 2) (8, 2, 5)\n",
      "so.shape: (8, 2, 5)\n",
      "so.shape,slogits.shape (16, 5) (16,)\n",
      "o.shape,o (16, 5) [[1.        1.        1.        1.        1.       ]\n",
      " [1.        1.        1.        1.        1.       ]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
      " [1.        1.        1.        1.        1.       ]\n",
      " [1.        1.        1.        1.        1.       ]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
      " [1.        1.        1.        1.        1.       ]\n",
      " [1.        1.        1.        1.        1.       ]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
      " [1.        1.        1.        1.        1.       ]\n",
      " [1.        1.        1.        1.        1.       ]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
      " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]]\n",
      "logits.shape, logits (16,) [ 0.6931472  3.6931472 12.693148  27.693148   0.6931472  3.6931472\n",
      " 12.693148  27.693148   0.6931472  3.6931472 12.693148  27.693148\n",
      "  0.6931472  3.6931472 12.693148  27.693148 ]\n",
      "out\n",
      " [[1.         1.         1.         1.         1.        ]\n",
      " [1.         1.         1.         1.         1.        ]\n",
      " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]\n",
      " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]\n",
      " [1.         1.         1.         1.         1.        ]\n",
      " [1.         1.         1.         1.         1.        ]\n",
      " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]\n",
      " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]]\n"
     ]
    }
   ],
   "source": [
    "t_kv_chunk_len = 2\n",
    "out = dotandv(\n",
    "    t_sq,\n",
    "    t_sv,\n",
    "    t_undo_sort,\n",
    "    t_kv_chunk_len,\n",
    "    t_n_hashes,\n",
    "    t_seqlen,\n",
    "    passthrough=True,\n",
    "    verbose=True,\n",
    ")\n",
    "print(\"out\\n\", out)\n",
    "print(\"\\n-----With softmax enabled----\\n\")\n",
    "out = dotandv(\n",
    "    t_sq,\n",
    "    t_sv,\n",
    "    t_undo_sort,\n",
    "    t_kv_chunk_len,\n",
    "    t_n_hashes,\n",
    "    t_seqlen,\n",
    "    passthrough=False,\n",
    "    verbose=True,\n",
    ")\n",
    "print(\"out\\n\", out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\"><b> Expected Output </b></font>\n",
    "</summary>\n",
    "\n",
    "**Expected Values**\n",
    "```\n",
    "rsq.shape,rsqt.shape:  (8, 2, 3) (8, 3, 2)\n",
    "dotlike\n",
    " [[[ 0.  0.]\n",
    "  [ 0.  0.]]\n",
    "\n",
    " [[ 3.  3.]\n",
    "  [ 3.  3.]]\n",
    "\n",
    " [[12. 12.]\n",
    "  [12. 12.]]\n",
    "\n",
    " [[27. 27.]\n",
    "  [27. 27.]]\n",
    "\n",
    " [[ 0.  0.]\n",
    "  [ 0.  0.]]\n",
    "\n",
    " [[ 3.  3.]\n",
    "  [ 3.  3.]]\n",
    "\n",
    " [[12. 12.]\n",
    "  [12. 12.]]\n",
    "\n",
    " [[27. 27.]\n",
    "  [27. 27.]]]\n",
    "dotlike post softmax\n",
    " [[[ 0.  0.]\n",
    "  [ 0.  0.]]\n",
    "\n",
    " [[ 3.  3.]\n",
    "  [ 3.  3.]]\n",
    "\n",
    " [[12. 12.]\n",
    "  [12. 12.]]\n",
    "\n",
    " [[27. 27.]\n",
    "  [27. 27.]]\n",
    "\n",
    " [[ 0.  0.]\n",
    "  [ 0.  0.]]\n",
    "\n",
    " [[ 3.  3.]\n",
    "  [ 3.  3.]]\n",
    "\n",
    " [[12. 12.]\n",
    "  [12. 12.]]\n",
    "\n",
    " [[27. 27.]\n",
    "  [27. 27.]]]\n",
    "dotlike.shape, vr.shape: (8, 2, 2) (8, 2, 5)\n",
    "so.shape: (8, 2, 5)\n",
    "so.shape,slogits.shape (16, 5) (16,)\n",
    "o.shape,o (16, 5) [[ 0.  0.  0.  0.  0.]\n",
    " [ 6.  6.  6.  6.  6.]\n",
    " [24. 24. 24. 24. 24.]\n",
    " [54. 54. 54. 54. 54.]\n",
    " [ 0.  0.  0.  0.  0.]\n",
    " [ 6.  6.  6.  6.  6.]\n",
    " [24. 24. 24. 24. 24.]\n",
    " [54. 54. 54. 54. 54.]\n",
    " [ 0.  0.  0.  0.  0.]\n",
    " [ 6.  6.  6.  6.  6.]\n",
    " [24. 24. 24. 24. 24.]\n",
    " [54. 54. 54. 54. 54.]\n",
    " [ 0.  0.  0.  0.  0.]\n",
    " [ 6.  6.  6.  6.  6.]\n",
    " [24. 24. 24. 24. 24.]\n",
    " [54. 54. 54. 54. 54.]]\n",
    "logits.shape, logits (16,) [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
    "out\n",
    " [[ 0.  0.  0.  0.  0.]\n",
    " [ 6.  6.  6.  6.  6.]\n",
    " [24. 24. 24. 24. 24.]\n",
    " [54. 54. 54. 54. 54.]\n",
    " [ 0.  0.  0.  0.  0.]\n",
    " [ 6.  6.  6.  6.  6.]\n",
    " [24. 24. 24. 24. 24.]\n",
    " [54. 54. 54. 54. 54.]]\n",
    "\n",
    "-----With softmax enabled----\n",
    "\n",
    "rsq.shape,rsqt.shape:  (8, 2, 3) (8, 3, 2)\n",
    "dotlike\n",
    " [[[ 0.  0.]\n",
    "  [ 0.  0.]]\n",
    "\n",
    " [[ 3.  3.]\n",
    "  [ 3.  3.]]\n",
    "\n",
    " [[12. 12.]\n",
    "  [12. 12.]]\n",
    "\n",
    " [[27. 27.]\n",
    "  [27. 27.]]\n",
    "\n",
    " [[ 0.  0.]\n",
    "  [ 0.  0.]]\n",
    "\n",
    " [[ 3.  3.]\n",
    "  [ 3.  3.]]\n",
    "\n",
    " [[12. 12.]\n",
    "  [12. 12.]]\n",
    "\n",
    " [[27. 27.]\n",
    "  [27. 27.]]]\n",
    "dotlike post softmax\n",
    " [[[0.5        0.5       ]\n",
    "  [0.5        0.5       ]]\n",
    "\n",
    " [[0.5        0.5       ]\n",
    "  [0.5        0.5       ]]\n",
    "\n",
    " [[0.49999976 0.49999976]\n",
    "  [0.49999976 0.49999976]]\n",
    "\n",
    " [[0.49999976 0.49999976]\n",
    "  [0.49999976 0.49999976]]\n",
    "\n",
    " [[0.5        0.5       ]\n",
    "  [0.5        0.5       ]]\n",
    "\n",
    " [[0.5        0.5       ]\n",
    "  [0.5        0.5       ]]\n",
    "\n",
    " [[0.49999976 0.49999976]\n",
    "  [0.49999976 0.49999976]]\n",
    "\n",
    " [[0.49999976 0.49999976]\n",
    "  [0.49999976 0.49999976]]]\n",
    "dotlike.shape, vr.shape: (8, 2, 2) (8, 2, 5)\n",
    "so.shape: (8, 2, 5)\n",
    "so.shape,slogits.shape (16, 5) (16,)\n",
    "o.shape,o (16, 5) [[1.        1.        1.        1.        1.       ]\n",
    " [1.        1.        1.        1.        1.       ]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
    " [1.        1.        1.        1.        1.       ]\n",
    " [1.        1.        1.        1.        1.       ]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
    " [1.        1.        1.        1.        1.       ]\n",
    " [1.        1.        1.        1.        1.       ]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
    " [1.        1.        1.        1.        1.       ]\n",
    " [1.        1.        1.        1.        1.       ]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]\n",
    " [0.9999995 0.9999995 0.9999995 0.9999995 0.9999995]]\n",
    "logits.shape, logits (16,) [ 0.6931472  3.6931472 12.693148  27.693148   0.6931472  3.6931472\n",
    " 12.693148  27.693148   0.6931472  3.6931472 12.693148  27.693148\n",
    "  0.6931472  3.6931472 12.693148  27.693148 ]\n",
    "out\n",
    " [[1.         1.         1.         1.         1.        ]\n",
    " [1.         1.         1.         1.         1.        ]\n",
    " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]\n",
    " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]\n",
    " [1.         1.         1.         1.         1.        ]\n",
    " [1.         1.         1.         1.         1.        ]\n",
    " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]\n",
    " [0.99999905 0.99999905 0.99999905 0.99999905 0.99999905]]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\" ><b>Completed code for reference </b></font>\n",
    "</summary>\n",
    "\n",
    "```\n",
    "# since this notebook is ungraded the completed code is provided here for reference\n",
    "def dotandv(sq, sv, undo_sort, kv_chunk_len, n_hashes, seqlen, passthrough, verbose=False ):\n",
    "    # Step 1\n",
    "    rsq = np.reshape(sq,(-1, kv_chunk_len, sq.shape[-1]))\n",
    "    rsqt =  np.swapaxes(rsq, -1, -2)\n",
    "    if verbose: print(\"rsq.shape,rsqt.shape: \", rsq.shape,rsqt.shape)\n",
    "    dotlike = np.matmul(rsq, rsqt)\n",
    "    if verbose: print(\"dotlike\\n\", dotlike)\n",
    "\n",
    "    #Step 2\n",
    "    dotlike, slogits = our_softmax(dotlike, passthrough)\n",
    "    if verbose: print(\"dotlike post softmax\\n\", dotlike)\n",
    "\n",
    "    #Step 3\n",
    "    vr = np.reshape(sv, (-1, kv_chunk_len, sv.shape[-1]))\n",
    "    if verbose:  print(\"dotlike.shape, vr.shape:\", dotlike.shape, vr.shape)\n",
    "    so = np.matmul(dotlike, vr)\n",
    "    if verbose: print(\"so.shape:\", so.shape)\n",
    "    so = np.reshape(so, (-1, so.shape[-1]))\n",
    "    slogits = np.reshape(slogits, (-1,))  # provided\n",
    "    if verbose: print(\"so.shape,slogits.shape\", so.shape, slogits.shape)\n",
    "\n",
    "    #Step 4\n",
    "    o = np.take(so, undo_sort, axis=0)\n",
    "    logits = np.take(slogits, undo_sort, axis=0)\n",
    "    if verbose: print(\"o.shape,o\", o.shape, o)\n",
    "    if verbose: print(\"logits.shape, logits\", logits.shape, logits)\n",
    "\n",
    "    #Step 5 (Provided)\n",
    "    if n_hashes > 1:\n",
    "      o = np.reshape(o, (n_hashes, seqlen, o.shape[-1]))\n",
    "      logits = np.reshape(logits, (n_hashes, seqlen, 1))\n",
    "      probs = np.exp(logits - fastmath.logsumexp(logits, axis=0, keepdims=True))\n",
    "      o = np.sum(o * probs, axis=0)\n",
    "\n",
    "    return(o)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! You have now done examples code for most of the operation that are unique to the LSH version of self-attention. I'm sure at this point you are wondering what happens if the number of entries in a bucket is not evenly distributed the way our example is. It is possible, for example for all of the `seqlen` entries to land in one bucket. Further, since the buckets are not aligned, our 'chunks' may be misaligned with the start of the bucket. The implementation addresses this by attending to adjacent chunks as was described in the lecture:\n",
    "\n",
    "<img src = \"C4W4_LN2_image13.PNG\" height=\"400\" width=\"750\">\n",
    "<center><b>Figure 12: Misaligned Access, looking before and after </b></center>\n",
    "\n",
    "Hopefully, having implemented parts of this, you will appreciate this diagram more fully.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"3.5\"></a>\n",
    "## Part 3.5 OurLSHSelfAttention\n",
    "\n",
    "You can examine the full implementations below. Area's we did not 'attend to' in our implementations above include variable bucket sizes and masking. We will instantiate a layer of the full implementation below. We tried to use the same variable names above to make it easier to decipher the full version. Note that some of the functionality we implemented in our routines is split between `attend` and `forward_unbatched`. We've inserted our version of hash below, but use the original version of `attend`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# original version from trax 1.3.4\n",
    "def attend(\n",
    "    q,\n",
    "    k=None,\n",
    "    v=None,\n",
    "    q_chunk_len=None,\n",
    "    kv_chunk_len=None,\n",
    "    n_chunks_before=0,\n",
    "    n_chunks_after=0,\n",
    "    mask_fn=None,\n",
    "    q_info=None,\n",
    "    kv_info=None,\n",
    "    dropout=0.0,\n",
    "    rng=None,\n",
    "):\n",
    "    \"\"\"Dot-product attention, with optional chunking and/or masking.\n",
    "\n",
    "  Args:\n",
    "    q: Query vectors, shape [q_len, d_qk]\n",
    "    k: Key vectors, shape [kv_len, d_qk]; or None\n",
    "    v: Value vectors, shape [kv_len, d_v]\n",
    "    q_chunk_len: Set to non-zero to enable chunking for query vectors\n",
    "    kv_chunk_len: Set to non-zero to enable chunking for key/value vectors\n",
    "    n_chunks_before: Number of adjacent previous chunks to attend to\n",
    "    n_chunks_after: Number of adjacent subsequent chunks to attend to\n",
    "    mask_fn: TODO(kitaev) doc\n",
    "    q_info: Query-associated metadata for masking\n",
    "    kv_info: Key-associated metadata for masking\n",
    "    dropout: Dropout rate\n",
    "    rng: RNG for dropout\n",
    "\n",
    "  Returns:\n",
    "    A tuple (output, dots_logsumexp). The output has shape [q_len, d_v], and\n",
    "    dots_logsumexp has shape [q_len]. The logsumexp of the attention\n",
    "    probabilities is useful for combining multiple rounds of attention (as in\n",
    "    LSH attention).\n",
    "  \"\"\"\n",
    "    assert v is not None\n",
    "    share_qk = k is None\n",
    "\n",
    "    if q_info is None:\n",
    "        q_info = np.arange(q.shape[-2], dtype=np.int32)\n",
    "\n",
    "    if kv_info is None and not share_qk:\n",
    "        kv_info = np.arange(v.shape[-2], dtype=np.int32)\n",
    "\n",
    "    # Split q/k/v into chunks along the time axis, if desired.\n",
    "    if q_chunk_len is not None:\n",
    "        q = np.reshape(q, (-1, q_chunk_len, q.shape[-1]))\n",
    "        q_info = np.reshape(q_info, (-1, q_chunk_len))\n",
    "\n",
    "    if share_qk:\n",
    "        assert kv_chunk_len is None or kv_chunk_len == q_chunk_len\n",
    "        k = q\n",
    "        kv_chunk_len = q_chunk_len\n",
    "        if kv_info is None:\n",
    "            kv_info = q_info\n",
    "        elif kv_chunk_len is not None:\n",
    "            # kv_info is not None, but reshape as required.\n",
    "            kv_info = np.reshape(kv_info, (-1, kv_chunk_len))\n",
    "    elif kv_chunk_len is not None:\n",
    "        k = np.reshape(k, (-1, kv_chunk_len, k.shape[-1]))\n",
    "        kv_info = np.reshape(kv_info, (-1, kv_chunk_len))\n",
    "\n",
    "    if kv_chunk_len is not None:\n",
    "        v = np.reshape(v, (-1, kv_chunk_len, v.shape[-1]))\n",
    "\n",
    "    if share_qk:\n",
    "        k = length_normalized(k)\n",
    "    k = k / np.sqrt(k.shape[-1])\n",
    "\n",
    "    # Optionally include adjacent chunks.\n",
    "    if q_chunk_len is not None or kv_chunk_len is not None:\n",
    "        assert q_chunk_len is not None and kv_chunk_len is not None\n",
    "    else:\n",
    "        assert n_chunks_before == 0 and n_chunks_after == 0\n",
    "\n",
    "    k = look_adjacent(k, n_chunks_before, n_chunks_after)\n",
    "    v = look_adjacent(v, n_chunks_before, n_chunks_after)\n",
    "    kv_info = look_adjacent(kv_info, n_chunks_before, n_chunks_after)\n",
    "\n",
    "    # Dot-product attention.\n",
    "    dots = np.matmul(q, np.swapaxes(k, -1, -2))\n",
    "\n",
    "    # Masking\n",
    "    if mask_fn is not None:\n",
    "        dots = mask_fn(dots, q_info[..., :, None], kv_info[..., None, :])\n",
    "\n",
    "    # Softmax.\n",
    "    dots_logsumexp = fastmath.logsumexp(dots, axis=-1, keepdims=True)\n",
    "    dots = np.exp(dots - dots_logsumexp)\n",
    "\n",
    "    if dropout > 0.0:\n",
    "        assert rng is not None\n",
    "        # Dropout is broadcast across the bin dimension\n",
    "        dropout_shape = (dots.shape[-2], dots.shape[-1])\n",
    "        #\n",
    "        keep_prob = tie_in(dots, 1.0 - dropout)\n",
    "        keep = fastmath.random.bernoulli(rng, keep_prob, dropout_shape)\n",
    "        multiplier = keep.astype(dots.dtype) / tie_in(keep, keep_prob)\n",
    "        dots = dots * multiplier\n",
    "\n",
    "    # The softmax normalizer (dots_logsumexp) is used by multi-round LSH attn.\n",
    "    out = np.matmul(dots, v)\n",
    "    out = np.reshape(out, (-1, out.shape[-1]))\n",
    "    dots_logsumexp = np.reshape(dots_logsumexp, (-1,))\n",
    "    return out, dots_logsumexp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ihFoYYBGKFVu"
   },
   "outputs": [],
   "source": [
    "class OurLSHSelfAttention(tl.LSHSelfAttention):\n",
    "    \"\"\"Our simplified LSH self-attention \"\"\"\n",
    "\n",
    "    def forward_unbatched(self, x, mask=None, *, weights, state, rng, update_state):\n",
    "        attend_rng, output_rng = fastmath.random.split(rng)\n",
    "        w_q, w_v, w_o = weights\n",
    "\n",
    "        q = np.matmul(x, w_q)\n",
    "        v = np.matmul(x, w_v)\n",
    "\n",
    "        if update_state:\n",
    "            _, old_hash_rng = state\n",
    "            hash_rng, hash_subrng = fastmath.random.split(old_hash_rng)\n",
    "            #      buckets = self.hash_vectors(q, hash_subrng, mask)  #  original\n",
    "            ## use our version of hash\n",
    "            buckets = our_hash_vectors(\n",
    "                q, hash_subrng, self.n_buckets, self.n_hashes, mask=mask\n",
    "            )\n",
    "            s_buckets = buckets\n",
    "            if self._max_length_for_buckets:\n",
    "                length = self.n_hashes * self._max_length_for_buckets\n",
    "                if buckets.shape[0] < length:\n",
    "                    s_buckets = np.concatenate(\n",
    "                        [buckets, np.zeros(length - buckets.shape[0], dtype=np.int32)],\n",
    "                        axis=0,\n",
    "                    )\n",
    "            state = (s_buckets, hash_rng)\n",
    "        else:\n",
    "            buckets, _ = state\n",
    "            if self._max_length_for_buckets:\n",
    "                buckets = buckets[: self.n_hashes * x.shape[0]]\n",
    "\n",
    "        seqlen = x.shape[0]\n",
    "        assert int(buckets.shape[0]) == self.n_hashes * seqlen\n",
    "\n",
    "        ticker = tie_in(x, np.arange(self.n_hashes * seqlen, dtype=np.int32))\n",
    "        buckets_and_t = seqlen * buckets + (ticker % seqlen)\n",
    "        buckets_and_t = fastmath.stop_gradient(buckets_and_t)\n",
    "\n",
    "        # Hash-based sort (\"s\" at the start of variable names means \"sorted\")\n",
    "        sbuckets_and_t, sticker = fastmath.sort_key_val(\n",
    "            buckets_and_t, ticker, dimension=-1\n",
    "        )\n",
    "        _, undo_sort = fastmath.sort_key_val(sticker, ticker, dimension=-1)\n",
    "        sbuckets_and_t = fastmath.stop_gradient(sbuckets_and_t)\n",
    "        sticker = fastmath.stop_gradient(sticker)\n",
    "        undo_sort = fastmath.stop_gradient(undo_sort)\n",
    "\n",
    "        st = sticker % seqlen\n",
    "        sq = np.take(q, st, axis=0)\n",
    "        sv = np.take(v, st, axis=0)\n",
    "\n",
    "        mask_fn = functools.partial(\n",
    "            mask_self_attention,\n",
    "            causal=self.causal,\n",
    "            exclude_self=True,\n",
    "            masked=self.masked,\n",
    "        )\n",
    "        q_info = st\n",
    "\n",
    "        assert (mask is not None) == self.masked\n",
    "        kv_info = None\n",
    "        if self.masked:\n",
    "            # mask is a boolean array (True means \"is valid token\")\n",
    "            smask = np.take(mask, st, axis=0)\n",
    "            ones_like_mask = tie_in(x, np.ones_like(smask, dtype=np.int32))\n",
    "            kv_info = q_info * np.where(smask, ones_like_mask, -ones_like_mask)\n",
    "\n",
    "        ## use original version of attend (could use ours but lacks masks and masking)\n",
    "        so, slogits = attend(\n",
    "            sq,\n",
    "            k=None,\n",
    "            v=sv,\n",
    "            q_chunk_len=self.chunk_len,\n",
    "            n_chunks_before=self.n_chunks_before,\n",
    "            n_chunks_after=self.n_chunks_after,\n",
    "            mask_fn=mask_fn,\n",
    "            q_info=q_info,\n",
    "            kv_info=kv_info,\n",
    "            dropout=self.attention_dropout,\n",
    "            rng=attend_rng,\n",
    "        )\n",
    "\n",
    "        # np.take(so, undo_sort, axis=0); np.take(slogits, undo_sort, axis=0) would\n",
    "        # also work, but these helpers include performance optimizations for TPU.\n",
    "        o = permute_via_gather(so, undo_sort, sticker, axis=0)\n",
    "        logits = permute_via_sort(slogits, sticker, buckets_and_t, axis=-1)\n",
    "\n",
    "        if self.n_hashes > 1:\n",
    "            o = np.reshape(o, (self.n_hashes, seqlen, o.shape[-1]))\n",
    "            logits = np.reshape(logits, (self.n_hashes, seqlen, 1))\n",
    "            probs = np.exp(logits - fastmath.logsumexp(logits, axis=0, keepdims=True))\n",
    "            o = np.sum(o * probs, axis=0)\n",
    "\n",
    "        assert o.shape == (seqlen, w_v.shape[-1])\n",
    "        out = np.matmul(o, w_o)\n",
    "        out = apply_broadcasted_dropout(out, self.output_dropout, output_rng)\n",
    "        return out, state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QG3yCwWV3zJd",
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Here we're going to try out our LSHSelfAttention\n",
    "n_heads = 3\n",
    "causal = False\n",
    "masked = False\n",
    "mask = None\n",
    "chunk_len = 8\n",
    "n_chunks_before = 0\n",
    "n_chunks_after = 0\n",
    "attention_dropout = 0.0\n",
    "n_hashes = 5\n",
    "n_buckets = 4\n",
    "seq_len = 8\n",
    "emb_len = 5\n",
    "al = OurLSHSelfAttention(\n",
    "    n_heads=n_heads,\n",
    "    d_qk=3,\n",
    "    d_v=4,\n",
    "    causal=causal,\n",
    "    chunk_len=8,\n",
    "    n_chunks_before=n_chunks_before,\n",
    "    n_chunks_after=n_chunks_after,\n",
    "    n_hashes=n_hashes,\n",
    "    n_buckets=n_buckets,\n",
    "    use_reference_code=True,\n",
    "    attention_dropout=attention_dropout,\n",
    "    mode=\"train\",\n",
    ")\n",
    "\n",
    "x = jax.random.uniform(jax.random.PRNGKey(0), (1, seq_len, emb_len), dtype=np.float32)\n",
    "al_osa = fastmath.random.get_prng(1)\n",
    "_, _ = al.init(tl.shapes.signature(x), rng=al_osa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "TzHug40iMe3S"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using jax\n",
      "using jax\n",
      "using jax\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray([[[ 6.6842824e-01, -1.1364323e-01, -5.4430610e-01,\n",
       "                2.1126242e-01, -1.0988623e-02],\n",
       "              [ 7.0949769e-01, -1.5455185e-01, -5.9923315e-01,\n",
       "                2.2719440e-01,  1.3833776e-02],\n",
       "              [ 7.1442688e-01, -1.2046628e-01, -5.3956544e-01,\n",
       "                1.7320301e-01, -1.6552269e-02],\n",
       "              [ 6.7178929e-01, -7.6611102e-02, -5.9399861e-01,\n",
       "                2.1236290e-01,  7.9482794e-04],\n",
       "              [ 7.1518433e-01, -1.1359170e-01, -5.7821894e-01,\n",
       "                2.1304411e-01,  3.0598268e-02],\n",
       "              [ 6.8235350e-01, -9.3979925e-02, -5.5341840e-01,\n",
       "                2.1608177e-01, -6.6673756e-04],\n",
       "              [ 6.1286640e-01, -8.1027031e-02, -4.8148823e-01,\n",
       "                1.9373313e-01,  3.1555295e-02],\n",
       "              [ 7.2203505e-01, -1.0199660e-01, -5.5215168e-01,\n",
       "                1.7872262e-01, -2.2289157e-02]]], dtype=float32)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "al(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "c5IBhMGjmg0z"
   },
   "source": [
    "<details>\n",
    "<summary>\n",
    "    <font size=\"3\"><b> Expected Output </b></font>\n",
    "</summary>\n",
    "\n",
    "**Expected Values**\n",
    "```\n",
    "using jax\n",
    "using jax\n",
    "using jax\n",
    "DeviceArray([[[ 6.6842824e-01, -1.1364323e-01, -5.4430610e-01,\n",
    "                2.1126242e-01, -1.0988623e-02],\n",
    "              [ 7.0949769e-01, -1.5455185e-01, -5.9923315e-01,\n",
    "                2.2719440e-01,  1.3833776e-02],\n",
    "              [ 7.1442688e-01, -1.2046628e-01, -5.3956544e-01,\n",
    "                1.7320301e-01, -1.6552269e-02],\n",
    "              [ 6.7178929e-01, -7.6611102e-02, -5.9399861e-01,\n",
    "                2.1236290e-01,  7.9482794e-04],\n",
    "              [ 7.1518433e-01, -1.1359170e-01, -5.7821894e-01,\n",
    "                2.1304411e-01,  3.0598268e-02],\n",
    "              [ 6.8235350e-01, -9.3979925e-02, -5.5341840e-01,\n",
    "                2.1608177e-01, -6.6673756e-04],\n",
    "              [ 6.1286640e-01, -8.1027031e-02, -4.8148823e-01,\n",
    "                1.9373313e-01,  3.1555295e-02],\n",
    "              [ 7.2203505e-01, -1.0199660e-01, -5.5215168e-01,\n",
    "                1.7872262e-01, -2.2289157e-02]]], dtype=float32)```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Congratuations!** you have created a custom layer and have become familiar with LSHSelfAttention."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "C4W4_LN2.ipynb",
   "provenance": []
  },
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}