steps.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x1d1a80bf270>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set numpy print mode to 2 decimal places and turn off scientific notation\n",
    "np.set_printoptions(precision=2, suppress=True)\n",
    "# Set torch mode to no grad\n",
    "torch.set_grad_enabled(False)\n",
    "# Set seed for reproducibility\n",
    "torch.manual_seed(79)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.61  0.23 -0.22  0.59]\n",
      " [-0.67  1.13  0.3   1.35]\n",
      " [-0.43 -1.38  1.01  1.45]\n",
      " [-0.03 -0.56 -0.61  2.2 ]\n",
      " [ 0.07  0.72  1.4   1.05]]\n"
     ]
    }
   ],
   "source": [
    "# We want to go through one layer of a transformer, starting from the input embeddings\n",
    "# Make random tensor with size (seq_len, emb_dim)\n",
    "seq_len = 5\n",
    "emb_dim = 4\n",
    "emb = torch.randn(seq_len, emb_dim)\n",
    "print(emb.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0. 0. 0. 0.]\n",
      " [1. 1. 1. 1.]\n",
      " [2. 2. 2. 2.]\n",
      " [3. 3. 3. 3.]\n",
      " [4. 4. 4. 4.]]\n"
     ]
    }
   ],
   "source": [
    "# Next is positional encoding\n",
    "# Make a tensor with size (seq_len, emb_dim) with columns with 1 to 5 and rows that are all the same\n",
    "pos = torch.arange(seq_len).unsqueeze(1).expand(seq_len, emb_dim).float()\n",
    "print(pos.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.61  0.23 -0.22  0.59]\n",
      " [ 0.33  2.13  1.3   2.35]\n",
      " [ 1.57  0.62  3.01  3.45]\n",
      " [ 2.97  2.44  2.39  5.2 ]\n",
      " [ 4.07  4.72  5.4   5.05]]\n"
     ]
    }
   ],
   "source": [
    "# Add the positional encoding to the embeddings\n",
    "emb_pos = emb + pos\n",
    "print(emb_pos.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 0.91 -0.22 -1.55  0.85]\n",
      " [-1.51  0.76 -0.29  1.03]\n",
      " [-0.52 -1.37  0.75  1.14]\n",
      " [-0.24 -0.71 -0.75  1.7 ]\n",
      " [-1.51 -0.19  1.21  0.49]]\n"
     ]
    }
   ],
   "source": [
    "# Next is layernorm\n",
    "# Make a layernorm module\n",
    "ln = nn.LayerNorm(emb_dim)\n",
    "emb_pos_ln = ln(emb_pos)\n",
    "print(emb_pos_ln.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wq1\n",
      "[[-0.74  0.7 ]\n",
      " [-1.3   1.73]\n",
      " [-0.86 -0.05]\n",
      " [-0.76 -0.09]]\n",
      "Wk1\n",
      "[[ 0.98 -0.18]\n",
      " [-1.44  0.51]\n",
      " [ 0.35  1.95]\n",
      " [ 1.3  -0.13]]\n",
      "Wv1\n",
      "[[-0.02 -0.41]\n",
      " [-0.56  0.38]\n",
      " [-0.82 -0.78]\n",
      " [ 0.23 -0.61]]\n",
      "Wq2\n",
      "[[ 0.2  -0.41]\n",
      " [ 0.62  1.26]\n",
      " [ 1.64  1.35]\n",
      " [-0.89  0.74]]\n",
      "Wk2\n",
      "[[-0.83  0.09]\n",
      " [-0.95 -2.35]\n",
      " [-0.59  1.18]\n",
      " [-2.3   2.21]]\n",
      "Wv2\n",
      "[[ 0.03 -0.8 ]\n",
      " [ 0.51  1.45]\n",
      " [-0.78 -0.59]\n",
      " [ 0.62 -1.37]]\n",
      "Wo\n",
      "[[ 0.12 -1.6  -0.37 -1.27]\n",
      " [ 0.22  1.31 -0.12 -0.3 ]\n",
      " [-0.19 -0.28  0.15  0.42]\n",
      " [-0.42 -0.23  1.29  0.22]]\n",
      "Q1\n",
      "[[ 0.3   0.25]\n",
      " [-0.41  0.18]\n",
      " [ 0.65 -2.87]\n",
      " [ 0.46 -1.51]\n",
      " [-0.06 -1.47]]\n",
      "K1\n",
      "[[ 1.77 -3.4 ]\n",
      " [-1.33 -0.03]\n",
      " [ 3.19  0.71]\n",
      " [ 2.72 -1.99]\n",
      " [-0.16  2.47]]\n",
      "V1\n",
      "[[ 1.56  0.23]\n",
      " [ 0.08  0.5 ]\n",
      " [ 0.41 -1.59]\n",
      " [ 1.39 -0.63]\n",
      " [-0.74 -0.69]]\n",
      "Q2\n",
      "[[-3.26 -2.1 ]\n",
      " [-1.21  1.95]\n",
      " [-0.73  0.35]\n",
      " [-3.22 -0.54]\n",
      " [ 1.14  2.37]]\n",
      "K2\n",
      "[[-1.59  0.64]\n",
      " [-1.68  0.02]\n",
      " [-1.33  6.56]\n",
      " [-2.58  4.5 ]\n",
      " [-0.41  2.81]]\n",
      "V2\n",
      "[[ 1.65 -1.29]\n",
      " [ 1.22  1.06]\n",
      " [-0.59 -3.58]\n",
      " [ 1.27 -2.71]\n",
      " [-0.77 -0.46]]\n"
     ]
    }
   ],
   "source": [
    "# Next up is multi-head attention\n",
    "# Doing it manually, create W_q, W_k, W_v, W_o\n",
    "# Make random tensors with size (emb_dim, emb_dim/2) for W_q, W_k, W_v, (emb_dim, emb_dim) for W_o\n",
    "W_q_1 = torch.randn(emb_dim, emb_dim//2)\n",
    "W_k_1 = torch.randn(emb_dim, emb_dim//2)\n",
    "W_v_1 = torch.randn(emb_dim, emb_dim//2)\n",
    "print(\"Wq1\")\n",
    "print(W_q_1.numpy())\n",
    "print(\"Wk1\")\n",
    "print(W_k_1.numpy())\n",
    "print(\"Wv1\")\n",
    "print(W_v_1.numpy())\n",
    "\n",
    "W_q_2 = torch.randn(emb_dim, emb_dim//2)\n",
    "W_k_2 = torch.randn(emb_dim, emb_dim//2)\n",
    "W_v_2 = torch.randn(emb_dim, emb_dim//2)\n",
    "print(\"Wq2\")\n",
    "print(W_q_2.numpy())\n",
    "print(\"Wk2\")\n",
    "print(W_k_2.numpy())\n",
    "print(\"Wv2\")\n",
    "print(W_v_2.numpy())\n",
    "\n",
    "W_o = torch.randn(emb_dim, emb_dim)\n",
    "print(\"Wo\")\n",
    "print(W_o.numpy())\n",
    "\n",
    "# Calculate Q, K, V\n",
    "Q1 = emb_pos_ln @ W_q_1\n",
    "K1 = emb_pos_ln @ W_k_1\n",
    "V1 = emb_pos_ln @ W_v_1\n",
    "print(\"Q1\")\n",
    "print(Q1.numpy())\n",
    "print(\"K1\")\n",
    "print(K1.numpy())\n",
    "print(\"V1\")\n",
    "print(V1.numpy())\n",
    "\n",
    "Q2 = emb_pos_ln @ W_q_2\n",
    "K2 = emb_pos_ln @ W_k_2\n",
    "V2 = emb_pos_ln @ W_v_2\n",
    "print(\"Q2\")\n",
    "print(Q2.numpy())\n",
    "print(\"K2\")\n",
    "print(K2.numpy())\n",
    "print(\"V2\")\n",
    "print(V2.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "K1^T\n",
      "[[ 1.77 -1.33  3.19  2.72 -0.16]\n",
      " [-3.4  -0.03  0.71 -1.99  2.47]]\n",
      "Scores1\n",
      "[[-0.33 -0.41  1.14  0.31  0.58]\n",
      " [-1.35  0.54 -1.18 -1.48  0.51]\n",
      " [10.91 -0.79  0.04  7.48 -7.18]\n",
      " [ 5.96 -0.57  0.39  4.26 -3.81]\n",
      " [ 4.92  0.11 -1.23  2.78 -3.63]]\n",
      "K2^T\n",
      "[[-1.59 -1.68 -1.33 -2.58 -0.41]\n",
      " [ 0.64  0.02  6.56  4.5   2.81]]\n",
      "Scores2\n",
      "[[ 3.84  5.42 -9.44 -1.03 -4.56]\n",
      " [ 3.18  2.08 14.44 11.92  5.98]\n",
      " [ 1.38  1.23  3.25  3.44  1.27]\n",
      " [ 4.78  5.39  0.76  5.9  -0.19]\n",
      " [-0.3  -1.86 14.02  7.71  6.18]]\n"
     ]
    }
   ],
   "source": [
    "# Attention scores\n",
    "scores1 = Q1 @ K1.T\n",
    "scores2 = Q2 @ K2.T\n",
    "print(\"K1^T\")\n",
    "print(K1.T.numpy())\n",
    "print(\"Scores1\")\n",
    "print(scores1.numpy())\n",
    "print(\"K2^T\")\n",
    "print(K2.T.numpy())\n",
    "print(\"Scores2\")\n",
    "print(scores2.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mask\n",
      "[[1. 0. 0. 0. 0.]\n",
      " [1. 1. 0. 0. 0.]\n",
      " [1. 1. 1. 0. 0.]\n",
      " [1. 1. 1. 1. 0.]\n",
      " [1. 1. 1. 1. 1.]]\n",
      "Scores1 masked\n",
      "[[-0.33  -inf  -inf  -inf  -inf]\n",
      " [-1.35  0.54  -inf  -inf  -inf]\n",
      " [10.91 -0.79  0.04  -inf  -inf]\n",
      " [ 5.96 -0.57  0.39  4.26  -inf]\n",
      " [ 4.92  0.11 -1.23  2.78 -3.63]]\n",
      "Scores2 masked\n",
      "[[ 3.84  -inf  -inf  -inf  -inf]\n",
      " [ 3.18  2.08  -inf  -inf  -inf]\n",
      " [ 1.38  1.23  3.25  -inf  -inf]\n",
      " [ 4.78  5.39  0.76  5.9   -inf]\n",
      " [-0.3  -1.86 14.02  7.71  6.18]]\n",
      "Scores1 softmax\n",
      "[[1.   0.   0.   0.   0.  ]\n",
      " [0.21 0.79 0.   0.   0.  ]\n",
      " [1.   0.   0.   0.   0.  ]\n",
      " [0.75 0.01 0.01 0.23 0.  ]\n",
      " [0.79 0.03 0.01 0.17 0.  ]]\n",
      "Scores2 softmax\n",
      "[[1.   0.   0.   0.   0.  ]\n",
      " [0.69 0.31 0.   0.   0.  ]\n",
      " [0.18 0.16 0.66 0.   0.  ]\n",
      " [0.21 0.32 0.01 0.46 0.  ]\n",
      " [0.   0.   0.98 0.01 0.  ]]\n"
     ]
    }
   ],
   "source": [
    "# Add causal mask\n",
    "mask = torch.tril(torch.ones(seq_len, seq_len))\n",
    "print(\"Mask\")\n",
    "print(mask.numpy())\n",
    "scores1 = scores1.masked_fill(mask==0, float('-inf'))\n",
    "scores2 = scores2.masked_fill(mask==0, float('-inf'))\n",
    "print(\"Scores1 masked\")\n",
    "print(scores1.numpy())\n",
    "print(\"Scores2 masked\")\n",
    "print(scores2.numpy())\n",
    "# Scale and softmax\n",
    "scores1 /= np.sqrt(emb_dim//2)\n",
    "scores2 /= np.sqrt(emb_dim//2)\n",
    "scores1 = torch.softmax(scores1, dim=-1)\n",
    "scores2 = torch.softmax(scores2, dim=-1)\n",
    "print(\"Scores1 softmax\")\n",
    "print(scores1.numpy())\n",
    "print(\"Scores2 softmax\")\n",
    "print(scores2.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Z1\n",
      "[[1.56 0.23]\n",
      " [0.39 0.44]\n",
      " [1.56 0.22]\n",
      " [1.5  0.01]\n",
      " [1.48 0.06]]\n",
      "Z2\n",
      "[[ 1.65 -1.29]\n",
      " [ 1.51 -0.55]\n",
      " [ 0.09 -2.43]\n",
      " [ 1.31 -1.22]\n",
      " [-0.57 -3.55]]\n"
     ]
    }
   ],
   "source": [
    "# Multiply by V\n",
    "Z1 = scores1 @ V1\n",
    "Z2 = scores2 @ V2\n",
    "print(\"Z1\")\n",
    "print(Z1.numpy())\n",
    "print(\"Z2\")\n",
    "print(Z2.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Z\n",
      "[[ 1.56  0.23  1.65 -1.29]\n",
      " [ 0.39  0.44  1.51 -0.55]\n",
      " [ 1.56  0.22  0.09 -2.43]\n",
      " [ 1.5   0.01  1.31 -1.22]\n",
      " [ 1.48  0.06 -0.57 -3.55]]\n"
     ]
    }
   ],
   "source": [
    "# Concatenate heads\n",
    "Z = torch.cat((Z1, Z2), dim=-1)\n",
    "print(\"Z\")\n",
    "print(Z.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output\n",
      "[[ 0.46 -2.37 -2.01 -1.63]\n",
      " [ 0.08 -0.34 -0.67 -0.11]\n",
      " [ 1.24 -1.68 -3.73 -2.54]\n",
      " [ 0.44 -2.47 -1.92 -1.61]\n",
      " [ 1.8  -1.31 -5.23 -2.91]]\n"
     ]
    }
   ],
   "source": [
    "# Multiply by W_o\n",
    "output = Z @ W_o\n",
    "print(\"Output\")\n",
    "print(output.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output with residual connection\n",
      "[[ 1.07 -2.15 -2.23 -1.04]\n",
      " [ 0.41  1.79  0.62  2.24]\n",
      " [ 2.81 -1.06 -0.72  0.91]\n",
      " [ 3.41 -0.03  0.47  3.59]\n",
      " [ 5.86  3.41  0.18  2.14]]\n"
     ]
    }
   ],
   "source": [
    "# Add residual connection\n",
    "output += emb_pos\n",
    "print(\"Output with residual connection\")\n",
    "print(output.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output with layernorm\n",
      "[[ 1.62 -0.8  -0.86  0.03]\n",
      " [-1.11  0.68 -0.83  1.27]\n",
      " [ 1.51 -1.01 -0.78  0.28]\n",
      " [ 0.94 -1.15 -0.84  1.05]\n",
      " [ 1.44  0.25 -1.32 -0.37]]\n"
     ]
    }
   ],
   "source": [
    "# Add another layernorm\n",
    "output = ln(output)\n",
    "print(\"Output with layernorm\")\n",
    "print(output.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wff1\n",
      "[[-1.15  0.36  0.31  0.01 -0.86  0.69  0.1   0.69]\n",
      " [-0.15  0.06  0.34  1.06  0.21 -0.52  1.14 -1.17]\n",
      " [ 0.17  1.28 -0.45  2.57  0.03  0.78 -0.81 -0.68]\n",
      " [ 1.76 -0.32  0.68  0.18  0.89  0.91  0.62 -0.91]]\n",
      "Wff2\n",
      "[[ 2.17 -0.94  0.74 -0.79]\n",
      " [ 1.43 -1.12  0.81  0.86]\n",
      " [-0.01 -2.6   0.54 -1.22]\n",
      " [-0.62  1.34 -0.15  0.3 ]\n",
      " [-0.35 -0.2  -0.17 -0.2 ]\n",
      " [-0.46  0.12  0.81  0.56]\n",
      " [ 2.05 -0.03 -0.42 -0.  ]\n",
      " [-3.42 -2.2  -0.29  1.93]]\n",
      "Bff1\n",
      "[-1.41 -0.15 -1.25 -0.58  0.26 -1.85  0.13 -0.  ]\n",
      "Bff2\n",
      "[ 0.25  0.52 -1.63  0.03]\n",
      "Wf1\n",
      "[[-1.83 -0.58  0.63 -3.03 -1.57  0.9  -0.03  2.6 ]\n",
      " [ 3.25 -1.84  1.13 -1.21  2.2  -0.62  2.12 -2.15]\n",
      " [-1.24 -0.62  0.66 -3.01 -1.3   1.21 -0.2   2.51]\n",
      " [ 0.78 -1.15  0.99 -3.17 -0.16  1.55  0.11  1.61]\n",
      " [-2.57 -1.04  0.87 -3.18 -1.56 -0.5   1.26  1.93]]\n",
      "Bf1\n",
      "[[-3.24 -0.73 -0.61 -3.61 -1.31 -0.95  0.09  2.6 ]\n",
      " [ 1.85 -1.98 -0.11 -1.79  2.45 -2.47  2.25 -2.15]\n",
      " [-2.65 -0.76 -0.59 -3.59 -1.04 -0.64 -0.07  2.5 ]\n",
      " [-0.63 -1.3  -0.26 -3.75  0.1  -0.3   0.23  1.61]\n",
      " [-3.97 -1.19 -0.38 -3.76 -1.3  -2.35  1.39  1.92]]\n",
      "Rf1\n",
      "[[0.   0.   0.   0.   0.   0.   0.09 2.6 ]\n",
      " [1.85 0.   0.   0.   2.45 0.   2.25 0.  ]\n",
      " [0.   0.   0.   0.   0.   0.   0.   2.5 ]\n",
      " [0.   0.   0.   0.   0.1  0.   0.23 1.61]\n",
      " [0.   0.   0.   0.   0.   0.   1.39 1.92]]\n",
      "Wf2\n",
      "[[-8.7  -5.72 -0.78  5.02]\n",
      " [ 7.75 -2.29 -0.01 -1.96]\n",
      " [-8.56 -5.5  -0.72  4.83]\n",
      " [-5.07 -3.57 -0.58  3.09]\n",
      " [-3.73 -4.27 -1.13  3.71]]\n",
      "Bf2\n",
      "[[-8.45 -5.2  -2.41  5.04]\n",
      " [ 8.   -1.77 -1.63 -1.94]\n",
      " [-8.31 -4.98 -2.34  4.86]\n",
      " [-4.81 -3.05 -2.2   3.11]\n",
      " [-3.48 -3.75 -2.76  3.74]]\n"
     ]
    }
   ],
   "source": [
    "# Now simulate the feedforward network\n",
    "# Wff1 (emb_dim, emb_dim*2), Wff2 (emb_dim*2, emb_dim)\n",
    "# Bff1 (emb_dim*2), Bff2 (emb_dim)\n",
    "W_ff1 = torch.randn(emb_dim, emb_dim*2)\n",
    "W_ff2 = torch.randn(emb_dim*2, emb_dim)\n",
    "B_ff1 = torch.randn(emb_dim*2)\n",
    "B_ff2 = torch.randn(emb_dim)\n",
    "print(\"Wff1\")\n",
    "print(W_ff1.numpy())\n",
    "print(\"Wff2\")\n",
    "print(W_ff2.numpy())\n",
    "print(\"Bff1\")\n",
    "print(B_ff1.numpy())\n",
    "print(\"Bff2\")\n",
    "print(B_ff2.numpy())\n",
    "\n",
    "# Calculate feedforward\n",
    "wf1 = output @ W_ff1\n",
    "bf1 = wf1 + B_ff1\n",
    "print(\"Wf1\")\n",
    "print(wf1.numpy())\n",
    "print(\"Bf1\")\n",
    "print(bf1.numpy())\n",
    "rf1 = nn.ReLU()(bf1)\n",
    "print(\"Rf1\")\n",
    "print(rf1.numpy())\n",
    "wf2 = rf1 @ W_ff2\n",
    "bf2 = wf2 + B_ff2\n",
    "print(\"Wf2\")\n",
    "print(wf2.numpy())\n",
    "print(\"Bf2\")\n",
    "print(bf2.numpy())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output\n",
      "[[ 0.12  1.54 -1.16 -0.5 ]\n",
      " [-0.06 -1.45  1.37  0.14]\n",
      " [ 0.12  1.54 -1.16 -0.5 ]\n",
      " [-0.09  1.62 -1.08 -0.45]\n",
      " [-0.32  1.69 -0.92 -0.44]]\n"
     ]
    }
   ],
   "source": [
    "# Multiply by some random (emb_dim, emb_dim) matrix to simulate the next layer\n",
    "# Then layernorm\n",
    "W = torch.randn(emb_dim, emb_dim)\n",
    "output = bf2 @ W\n",
    "output = ln(output)\n",
    "print(\"Output\")\n",
    "print(output.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wlin\n",
      "[[-0.69 -0.08  0.13  0.59 -1.39  0.96 -1.98  0.23  1.38 -0.09]\n",
      " [-0.32  2.06  0.49 -0.99 -1.92 -1.9  -0.25 -1.82 -0.75 -0.42]\n",
      " [ 0.2   0.29 -0.61  1.04  0.35 -0.05 -0.95 -0.49  0.22  1.01]\n",
      " [ 0.67 -0.37  0.84 -0.9  -1.14 -0.39 -0.97  1.77 -0.22  0.  ]]\n",
      "Logits\n",
      "[[-1.14  3.02  1.06 -2.22 -2.98 -2.56  0.97 -3.1  -1.14 -1.83]\n",
      " [ 0.87 -2.64 -1.43  2.7   3.2   2.56 -0.95  2.21  1.27  1.99]\n",
      " [-1.14  3.02  1.06 -2.23 -2.98 -2.56  0.97 -3.09 -1.14 -1.83]\n",
      " [-0.97  3.21  1.06 -2.38 -2.86 -2.93  1.24 -3.24 -1.48 -1.76]\n",
      " [-0.79  3.41  0.97 -2.43 -2.62 -3.29  1.52 -3.48 -1.82 -1.61]]\n"
     ]
    }
   ],
   "source": [
    "# Assume vocab size is 10\n",
    "vocab_size = 10\n",
    "# Output is (seq_len, emb_dim)\n",
    "# Make a linear layer with size (emb_dim, vocab_size)\n",
    "wlin = torch.randn(emb_dim, vocab_size)\n",
    "print(\"Wlin\")\n",
    "print(wlin.numpy())\n",
    "# Calculate logits\n",
    "logits = output @ wlin\n",
    "print(\"Logits\")\n",
    "print(logits.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Probs\n",
      "[[0.01 0.76 0.11 0.   0.   0.   0.1  0.   0.01 0.01]\n",
      " [0.03 0.   0.   0.2  0.32 0.17 0.01 0.12 0.05 0.1 ]\n",
      " [0.01 0.76 0.11 0.   0.   0.   0.1  0.   0.01 0.01]\n",
      " [0.01 0.77 0.09 0.   0.   0.   0.11 0.   0.01 0.01]\n",
      " [0.01 0.79 0.07 0.   0.   0.   0.12 0.   0.   0.01]]\n"
     ]
    }
   ],
   "source": [
    "# Calculate softmax\n",
    "probs = torch.softmax(logits, dim=-1)\n",
    "print(\"Probs\")\n",
    "print(probs.numpy())"
   ]
  }
 ],
 "metadata": {
  "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}