notebooks/notes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notebook that helped me figure out topics as I created my LLM, might remove since it's not really polished.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "parent_directory = os.path.dirname(os.getcwd())\n",
    "sys.path.append(parent_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attention\n",
    "Given an input sequence of three tokens, with embedding dimension of 2:\n",
    "$$\n",
    "X = \\begin{bmatrix}\n",
    "x_{1,1} & x_{1,2} \\\\\n",
    "x_{2,1} & x_{2,2} \\\\\n",
    "x_{3,1} & x_{3,2}\n",
    "\\end{bmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And Q, K, and V learnable weights:\n",
    "$$\n",
    "W_Q = \\begin{bmatrix}\n",
    "w^{q}_{1,1} & w^{q}_{1,2} \\\\\n",
    "w^{q}_{2,1} & w^{q}_{2,2}\n",
    "\\end{bmatrix}\n",
    "\\quad\n",
    "W_K = \\begin{bmatrix}\n",
    "w^{k}_{1,1} & w^{k}_{1,2} \\\\\n",
    "w^{k}_{2,1} & w^{k}_{2,2}\n",
    "\\end{bmatrix}\n",
    "\\quad\n",
    "W_V = \\begin{bmatrix}\n",
    "w^{v}_{1,1} & w^{v}_{1,2} \\\\\n",
    "w^{v}_{2,1} & w^{v}_{2,2}\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "Q = X \\times W_Q\n",
    "$$\n",
    "\n",
    "$$\n",
    "Q = \\begin{bmatrix}\n",
    "x_{1,1} & x_{1,2} \\\\\n",
    "x_{2,1} & x_{2,2} \\\\\n",
    "x_{3,1} & x_{3,2}\n",
    "\\end{bmatrix}\n",
    "\\times\n",
    "\\begin{bmatrix}\n",
    "w^{q}_{1,1} & w^{q}_{1,2} \\\\\n",
    "w^{q}_{2,1} & w^{q}_{2,2}\n",
    "\\end{bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "x_{1,1}w^{q}_{1,1} + x_{1,2}w^{q}_{2,1} & x_{1,1}w^{q}_{1,2} + x_{1,2}w^{q}_{2,2} \\\\\n",
    "x_{2,1}w^{q}_{1,1} + x_{2,2}w^{q}_{2,1} & x_{2,1}w^{q}_{1,2} + x_{2,2}w^{q}_{2,2} \\\\\n",
    "x_{3,1}w^{q}_{1,1} + x_{3,2}w^{q}_{2,1} & x_{3,1}w^{q}_{1,2} + x_{3,2}w^{q}_{2,2}\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "$$\n",
    "Q = \\begin{bmatrix}\n",
    "q_{1,1} & q_{1,2} \\\\\n",
    "q_{2,1} & q_{2,2} \\\\\n",
    "q_{3,1} & q_{3,2}\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "K = X \\times W_K \\quad = \\quad \\begin{bmatrix}\n",
    "k_{1,1} & k_{1,2} \\\\\n",
    "k_{2,1} & k_{2,2} \\\\\n",
    "k_{3,1} & k_{3,2}\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "$$\n",
    "V = X \\times W_V \\quad = \\quad \\begin{bmatrix}\n",
    "v_{1,1} & v_{1,2} \\\\\n",
    "v_{2,1} & v_{2,2} \\\\\n",
    "v_{3,1} & v_{3,2}\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Attention Score:\n",
    "$$\n",
    "A\\_S = Q \\times K^T\n",
    "$$\n",
    "\n",
    "$$\n",
    "A\\_S = \\begin{bmatrix}\n",
    "q_{1,1} & q_{1,2} \\\\\n",
    "q_{2,1} & q_{2,2} \\\\\n",
    "q_{3,1} & q_{3,2}\n",
    "\\end{bmatrix}\n",
    "\\times\n",
    "\\begin{bmatrix}\n",
    "k_{1,1} & k_{2,1} & k_{3,1} \\\\\n",
    "k_{1,2} & k_{2,2} & k_{3,2}\n",
    "\\end{bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "q_{1,1}k_{1,1} + q_{1,2}k_{1,2} & q_{1,1}k_{2,1} + q_{1,2}k_{2,2} & q_{1,1}k_{3,1} + q_{1,2}k_{3,2} \\\\\n",
    "q_{2,1}k_{1,1} + q_{2,2}k_{1,2} & q_{2,1}k_{2,1} + q_{2,2}k_{2,2} & q_{2,1}k_{3,1} + q_{2,2}k_{3,2} \\\\\n",
    "q_{3,1}k_{1,1} + q_{3,2}k_{1,2} & q_{3,1}k_{2,1} + q_{3,2}k_{2,2} & q_{3,1}k_{3,1} + q_{3,2}k_{3,2}\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "$$\n",
    "A\\_S = \\begin{bmatrix}\n",
    "s_{1,1} & s_{1,2} & s_{1,3} \\\\\n",
    "s_{2,1} & s_{2,2} & s_{2,3} \\\\\n",
    "s_{3,1} & s_{3,2} & s_{3,3}\n",
    "\\end{bmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scale and Normalize attention scores:  \n",
    "(softmax is applied across the features dimension)\n",
    "$$\n",
    "\\text{Scaled, Normalized } A\\_S = \\text{softmax}\\left(\\frac{1}{\\sqrt{2}} \\begin{bmatrix}\n",
    "s_{1,1} & s_{1,2} & s_{1,3} \\\\\n",
    "s_{2,1} & s_{2,2} & s_{2,3} \\\\\n",
    "s_{3,1} & s_{3,2} & s_{3,3}\n",
    "\\end{bmatrix}\\right)\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Attention Weights\n",
    "$$\n",
    "\\text{Attention Weights} = \\text{Scaled, Normalized } A\\_S\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\text{Attention Weights} = \\begin{bmatrix}\n",
    "\\alpha_{1,1} & \\alpha_{1,2} & \\alpha_{1,3} \\\\\n",
    "\\alpha_{2,1} & \\alpha_{2,2} & \\alpha_{2,3} \\\\\n",
    "\\alpha_{3,1} & \\alpha_{3,2} & \\alpha_{3,3}\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Attention Output:\n",
    "$$\n",
    "\\text{Attention Output} = \\text{Attention Weights} \\times V\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\text{Attention Output} = \\begin{bmatrix}\n",
    "\\alpha_{1,1} & \\alpha_{1,2} & \\alpha_{1,3} \\\\\n",
    "\\alpha_{2,1} & \\alpha_{2,2} & \\alpha_{2,3} \\\\\n",
    "\\alpha_{3,1} & \\alpha_{3,2} & \\alpha_{3,3}\n",
    "\\end{bmatrix}\n",
    "\\times\n",
    "\\begin{bmatrix}\n",
    "v_{1,1} & v_{1,2} \\\\\n",
    "v_{2,1} & v_{2,2} \\\\\n",
    "v_{3,1} & v_{3,2}\n",
    "\\end{bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "\\alpha_{1,1}v_{1,1} + \\alpha_{1,2}v_{2,1} + \\alpha_{1,3}v_{3,1} & \\alpha_{1,1}v_{1,2} + \\alpha_{1,2}v_{2,2} + \\alpha_{1,3}v_{3,2} \\\\\n",
    "\\alpha_{2,1}v_{1,1} + \\alpha_{2,2}v_{2,1} + \\alpha_{2,3}v_{3,1} & \\alpha_{2,1}v_{1,2} + \\alpha_{2,2}v_{2,2} + \\alpha_{2,3}v_{3,2} \\\\\n",
    "\\alpha_{3,1}v_{1,1} + \\alpha_{3,2}v_{2,1} + \\alpha_{3,3}v_{3,1} & \\alpha_{3,1}v_{1,2} + \\alpha_{3,2}v_{2,2} + \\alpha_{3,3}v_{3,2}\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\text{Attention Output} = O = \\begin{bmatrix}\n",
    "o_{1,1} & o_{1,2} \\\\\n",
    "o_{2,1} & o_{2,2} \\\\\n",
    "o_{3,1} & o_{3,2}\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multi-head attention\n",
    "$$\n",
    "\\text{MultiHead}(Q, K, V) = \\text{Concat}(\\text{head}_1, \\dots, \\text{head}_h) W^O\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\text{head}_i = \\text{Attention}(Q W_Q^{(i)}, K W_K^{(i)}, V W_V^{(i)}) = \\text{softmax}\\left(\\frac{(Q W_Q^{(i)}) (K W_K^{(i)})^T}{\\sqrt{d_k}}\\right) (V W_V^{(i)})\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "example of mulit-head attention\n",
    "$$\n",
    "\\text{MultiHead}(Q, K, V) = \\text{Concat}(\\text{head}_1, \\text{head}_2, \\text{head}_3) W^O\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\text{Concat}(\\text{head}_1, \\text{head}_2, \\text{head}_3) =\n",
    "\\begin{bmatrix}\n",
    "\\text{head}_1 & \\text{head}_2 & \\text{head}_3\n",
    "\\end{bmatrix}\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "h_{1,1} & h_{1,2} & h_{1,3} & h_{1,4} & h_{1,5} & h_{1,6} \\\\\n",
    "h_{2,1} & h_{2,2} & h_{2,3} & h_{2,4} & h_{2,5} & h_{2,6} \\\\\n",
    "h_{3,1} & h_{3,2} & h_{3,3} & h_{3,4} & h_{3,5} & h_{3,6}\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\text{Attention Output} = \\begin{bmatrix}\n",
    "h_{1,1} & h_{1,2} & h_{1,3} & h_{1,4} & h_{1,5} & h_{1,6} \\\\\n",
    "h_{2,1} & h_{2,2} & h_{2,3} & h_{2,4} & h_{2,5} & h_{2,6} \\\\\n",
    "h_{3,1} & h_{3,2} & h_{3,3} & h_{3,4} & h_{3,5} & h_{3,6}\n",
    "\\end{bmatrix}\n",
    "\\times\n",
    "\\begin{bmatrix}\n",
    "w_{1,1} & w_{1,2} & \\dots & w_{1,d_{model}} \\\\\n",
    "w_{2,1} & w_{2,2} & \\dots & w_{2,d_{model}} \\\\\n",
    "\\vdots  & \\vdots  & \\ddots & \\vdots \\\\\n",
    "w_{6,1} & w_{6,2} & \\dots & w_{6,d_{model}}\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "$$\n",
    "=\n",
    "\\begin{bmatrix}\n",
    "o_{1,1} & o_{1,2} & \\dots & o_{1,d_{model}} \\\\\n",
    "o_{2,1} & o_{2,2} & \\dots & o_{2,d_{model}} \\\\\n",
    "o_{3,1} & o_{3,2} & \\dots & o_{3,d_{model}}\n",
    "\\end{bmatrix}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rotary Positional Embedding (RoPE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "RoPE (rotary positional embedding) -- spelled out.\n",
    "Rotate token embedding pairs before attending. This means that the positional embedding will no longer be a learned parameter (or fixed sinusoidal), it is just a transformation of a token's embedding. This encodes positional information intrinsically within the token's embedding.\n",
    "\n",
    "The rotation angle `theta_p` is set to `p * ω_k`, where `p` is the token's position in the sequence and `w_k` is a frequency associated with each pair of dimensions.\n",
    "`w_k` decreases exponentially with increasing dimension index, meaning lower-indexed pairs of dimensions rotate more than higher-indexed ones.\n",
    "This design allows RoPE to capture relative positional information across different dimensions in a smooth, continuous manner.\n",
    "'''\n",
    "\n",
    "import math\n",
    "import torch\n",
    "\n",
    "def rope_1d_slow(p, embedding):\n",
    "    # p: index-position of token in sequence\n",
    "    # embedding: embedding representation of token\n",
    "\n",
    "    dim = len(embedding)\n",
    "    assert dim % 2 == 0\n",
    "\n",
    "    def w_k(k):\n",
    "        return 1 / (10_000 ** (2*k/dim))\n",
    "    \n",
    "    def theta_p(k):\n",
    "        # k: index-position of embedding-pair for tokem embedding\n",
    "        return p * w_k(k)\n",
    "    \n",
    "    transformation = torch.zeros(dim//2, 2)\n",
    "\n",
    "    for i in range(0, len(embedding), 2):\n",
    "        x_i_0 = embedding[i]\n",
    "        x_i_1 = embedding[i+1]\n",
    "        theta_p_val = theta_p(i//2)\n",
    "        cos = math.cos(theta_p_val)\n",
    "        sin = math.sin(theta_p_val)\n",
    "        rotation = torch.tensor([\n",
    "            [cos, -1*sin],\n",
    "            [sin, cos]\n",
    "        ])\n",
    "        x = torch.tensor([x_i_0, x_i_1])\n",
    "        transformation[i//2, :] = torch.matmul(rotation, x)\n",
    "\n",
    "    return transformation.view(-1)\n",
    "\n",
    "def rope_1d_faster(p, embedding):\n",
    "    dim = len(embedding)\n",
    "    assert dim % 2 == 0\n",
    "\n",
    "    indices = torch.arange(0, dim // 2)\n",
    "    w_k = 1 / (10_000 ** (2 * indices / dim))\n",
    "    cosine = torch.cos(p * w_k)\n",
    "    cosine = cosine.repeat_interleave(2)\n",
    "    sine = torch.sin(p * w_k)\n",
    "    sine = sine.repeat_interleave(2)\n",
    "    \n",
    "    x = embedding\n",
    "    x_shifted = torch.empty(len(embedding), dtype=embedding.dtype)\n",
    "    x_shifted[1::2], x_shifted[::2] = embedding[::2], -embedding[1::2]\n",
    "\n",
    "    return x * cosine + x_shifted * sine\n",
    "\n",
    "def rope_multi_head_broadcasting(att_comp):\n",
    "    # att_comp: attention component, either Q or K.\n",
    "    # Perform RoPE on each attention head of att_comp\n",
    "    # expected shape (Batch, Sequence, Head, Head Dimension).\n",
    "    # Use broadcasting along the Batch and Sequence dimension to save memory.\n",
    "    B, S, H, HD = att_comp.shape\n",
    "    assert HD % 2 == 0\n",
    "\n",
    "    ind_emb = torch.arange(0, HD // 2) # embedding-pairs indices\n",
    "    w_k = 1 / (10_000 ** (2 * ind_emb / HD))\n",
    "\n",
    "    # Expand w_k for each token in Sequence and perform `p * w_k`\n",
    "    w_k = w_k.view(1, -1)\n",
    "    w_k = w_k.repeat_interleave(S, 0)\n",
    "    ind_tok = torch.arange(S).unsqueeze(1).repeat(1, w_k.shape[1]) # token indices\n",
    "    w_k = w_k * ind_tok\n",
    "\n",
    "    # Get cosine rot. values\n",
    "    cosine = torch.cos(w_k)\n",
    "    cosine = cosine.repeat_interleave(2, -1) # repeat for embedding pairs\n",
    "    # reshape to B=1, S, H=1, HD\n",
    "    cosine = cosine.view(1, S, 1, HD)\n",
    "\n",
    "    # Get sine rot. values\n",
    "    sine = torch.sin(w_k)\n",
    "    sine = sine.repeat_interleave(2, -1) # repeat for embedding pairs\n",
    "    sine = sine.repeat_interleave(H, 0) # repeat for all attn heads\n",
    "    # reshape to B, S, H, HD\n",
    "    sine = sine.view(1, S, H, HD)\n",
    "    sine = sine.repeat_interleave(B, 0)\n",
    "\n",
    "    x = att_comp\n",
    "    x_shifted = torch.empty(att_comp.shape, dtype=att_comp.dtype)\n",
    "    x_shifted[...,1::2], x_shifted[...,::2] = att_comp[...,::2], -att_comp[...,1::2]\n",
    "\n",
    "    return x * cosine + x_shifted * sine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[ 0.0000,  0.1000,  0.2000,  0.3000,  0.4000,  0.5000,  0.6000,\n",
      "            0.7000],\n",
      "          [ 0.0000,  0.1000,  0.2000,  0.3000,  0.4000,  0.5000,  0.6000,\n",
      "            0.7000]],\n",
      "\n",
      "         [[-0.0841,  0.0540,  0.1691,  0.3185,  0.3950,  0.5040,  0.5993,\n",
      "            0.7006],\n",
      "          [-0.0841,  0.0540,  0.1691,  0.3185,  0.3950,  0.5040,  0.5993,\n",
      "            0.7006]],\n",
      "\n",
      "         [[-0.0909, -0.0416,  0.1364,  0.3338,  0.3899,  0.5079,  0.5986,\n",
      "            0.7012],\n",
      "          [-0.0909, -0.0416,  0.1364,  0.3338,  0.3899,  0.5079,  0.5986,\n",
      "            0.7012]]],\n",
      "\n",
      "\n",
      "        [[[ 0.0000,  0.1000,  0.2000,  0.3000,  0.4000,  0.5000,  0.6000,\n",
      "            0.7000],\n",
      "          [ 0.0000,  0.1000,  0.2000,  0.3000,  0.4000,  0.5000,  0.6000,\n",
      "            0.7000]],\n",
      "\n",
      "         [[-0.0841,  0.0540,  0.1691,  0.3185,  0.3950,  0.5040,  0.5993,\n",
      "            0.7006],\n",
      "          [-0.0841,  0.0540,  0.1691,  0.3185,  0.3950,  0.5040,  0.5993,\n",
      "            0.7006]],\n",
      "\n",
      "         [[-0.0909, -0.0416,  0.1364,  0.3338,  0.3899,  0.5079,  0.5986,\n",
      "            0.7012],\n",
      "          [-0.0909, -0.0416,  0.1364,  0.3338,  0.3899,  0.5079,  0.5986,\n",
      "            0.7012]]]])\n"
     ]
    }
   ],
   "source": [
    "# Full RoPE implementation\n",
    "B, S, H, HD = 2, 3, 2, 8\n",
    "q = torch.empty(B, S, H, HD)\n",
    "q[:,:,:,:] = torch.tensor([0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])\n",
    "rope_multi_head_broadcasting_res = rope_multi_head_broadcasting(q)\n",
    "print(rope_multi_head_broadcasting_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.0909, -0.0416,  0.1364,  0.3338,  0.3899,  0.5079,  0.5986,  0.7012])\n",
      "tensor([-0.0909, -0.0416,  0.1364,  0.3338,  0.3899,  0.5079,  0.5986,  0.7012])\n"
     ]
    }
   ],
   "source": [
    "# Simpler RoPE implementations\n",
    "rope_1d_slow_res = rope_1d_slow(2, [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])\n",
    "print(rope_1d_slow(2, [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7]))\n",
    "print(rope_1d_faster(2, torch.tensor([0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n"
     ]
    }
   ],
   "source": [
    "# sanity check that naive implementation\n",
    "print(torch.equal(rope_multi_head_broadcasting_res[0,2,0,:], rope_1d_slow_res))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000],\n",
      "          [-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000]],\n",
      "\n",
      "         [[-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000],\n",
      "          [-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000]],\n",
      "\n",
      "         [[-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000],\n",
      "          [-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000]]],\n",
      "\n",
      "\n",
      "        [[[-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000],\n",
      "          [-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000]],\n",
      "\n",
      "         [[-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000],\n",
      "          [-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000]],\n",
      "\n",
      "         [[-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000],\n",
      "          [-0.4000, -0.5000, -0.6000, -0.7000,  0.0000,  0.1000,  0.2000,\n",
      "            0.3000]]]])\n"
     ]
    }
   ],
   "source": [
    "def rotate_half(x):\n",
    "    x1, x2 = x.chunk(2, dim=-1)\n",
    "    return torch.cat((-x2, x1), dim=-1)\n",
    "\n",
    "B, S, H, HD = 2, 3, 2, 8\n",
    "q = torch.empty(B, S, H, HD)\n",
    "q[:,:,:,:] = torch.tensor([0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7])\n",
    "q = rotate_half(q)\n",
    "print(q)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}