makemore_part4_backprop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## makemore: becoming a backprop ninja\n",
    "\n",
    "swole doge style"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# there no change change in the first several cells from last lecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt # for making figures\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32033\n",
      "15\n",
      "['emma', 'olivia', 'ava', 'isabella', 'sophia', 'charlotte', 'mia', 'amelia']\n"
     ]
    }
   ],
   "source": [
    "# read in all the words\n",
    "words = open('names.txt', 'r').read().splitlines()\n",
    "print(len(words))\n",
    "print(max(len(w) for w in words))\n",
    "print(words[:8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{1: 'a', 2: 'b', 3: 'c', 4: 'd', 5: 'e', 6: 'f', 7: 'g', 8: 'h', 9: 'i', 10: 'j', 11: 'k', 12: 'l', 13: 'm', 14: 'n', 15: 'o', 16: 'p', 17: 'q', 18: 'r', 19: 's', 20: 't', 21: 'u', 22: 'v', 23: 'w', 24: 'x', 25: 'y', 26: 'z', 0: '.'}\n",
      "27\n"
     ]
    }
   ],
   "source": [
    "# build the vocabulary of characters and mappings to/from integers\n",
    "chars = sorted(list(set(''.join(words))))\n",
    "stoi = {s:i+1 for i,s in enumerate(chars)}\n",
    "stoi['.'] = 0\n",
    "itos = {i:s for s,i in stoi.items()}\n",
    "vocab_size = len(itos)\n",
    "print(itos)\n",
    "print(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([182625, 3]) torch.Size([182625])\n",
      "torch.Size([22655, 3]) torch.Size([22655])\n",
      "torch.Size([22866, 3]) torch.Size([22866])\n"
     ]
    }
   ],
   "source": [
    "# build the dataset\n",
    "block_size = 3 # context length: how many characters do we take to predict the next one?\n",
    "\n",
    "def build_dataset(words):  \n",
    "  X, Y = [], []\n",
    "  \n",
    "  for w in words:\n",
    "    context = [0] * block_size\n",
    "    for ch in w + '.':\n",
    "      ix = stoi[ch]\n",
    "      X.append(context)\n",
    "      Y.append(ix)\n",
    "      context = context[1:] + [ix] # crop and append\n",
    "\n",
    "  X = torch.tensor(X)\n",
    "  Y = torch.tensor(Y)\n",
    "  print(X.shape, Y.shape)\n",
    "  return X, Y\n",
    "\n",
    "import random\n",
    "random.seed(42)\n",
    "random.shuffle(words)\n",
    "n1 = int(0.8*len(words))\n",
    "n2 = int(0.9*len(words))\n",
    "\n",
    "Xtr,  Ytr  = build_dataset(words[:n1])     # 80%\n",
    "Xdev, Ydev = build_dataset(words[n1:n2])   # 10%\n",
    "Xte,  Yte  = build_dataset(words[n2:])     # 10%"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ok biolerplate done, now we get to the action:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# utility function we will use later when comparing manual gradients to PyTorch gradients\n",
    "def cmp(s, dt, t):\n",
    "  ex = torch.all(dt == t.grad).item()\n",
    "  app = torch.allclose(dt, t.grad)\n",
    "  maxdiff = (dt - t.grad).abs().max().item()\n",
    "  print(f'{s:15s} | exact: {str(ex):5s} | approximate: {str(app):5s} | maxdiff: {maxdiff}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4137\n"
     ]
    }
   ],
   "source": [
    "n_embd = 10 # the dimensionality of the character embedding vectors\n",
    "n_hidden = 64 # the number of neurons in the hidden layer of the MLP\n",
    "\n",
    "g = torch.Generator().manual_seed(2147483647) # for reproducibility\n",
    "C  = torch.randn((vocab_size, n_embd),            generator=g)\n",
    "# Layer 1\n",
    "W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * (5/3)/((n_embd * block_size)**0.5)\n",
    "b1 = torch.randn(n_hidden,                        generator=g) * 0.1 # using b1 just for fun, it's useless because of BN\n",
    "# Layer 2\n",
    "W2 = torch.randn((n_hidden, vocab_size),          generator=g) * 0.1\n",
    "b2 = torch.randn(vocab_size,                      generator=g) * 0.1\n",
    "# BatchNorm parameters\n",
    "bngain = torch.randn((1, n_hidden))*0.1 + 1.0\n",
    "bnbias = torch.randn((1, n_hidden))*0.1\n",
    "\n",
    "# Note: I am initializating many of these parameters in non-standard ways\n",
    "# because sometimes initializating with e.g. all zeros could mask an incorrect\n",
    "# implementation of the backward pass.\n",
    "\n",
    "parameters = [C, W1, b1, W2, b2, bngain, bnbias]\n",
    "print(sum(p.nelement() for p in parameters)) # number of parameters in total\n",
    "for p in parameters:\n",
    "  p.requires_grad = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "n = batch_size # a shorter variable also, for convenience\n",
    "# construct a minibatch\n",
    "ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)\n",
    "Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(3.3470, grad_fn=<NegBackward0>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# forward pass, \"chunkated\" into smaller steps that are possible to backward one at a time\n",
    "\n",
    "emb = C[Xb] # embed the characters into vectors\n",
    "embcat = emb.view(emb.shape[0], -1) # concatenate the vectors\n",
    "# Linear layer 1\n",
    "hprebn = embcat @ W1 + b1 # hidden layer pre-activation\n",
    "# BatchNorm layer\n",
    "bnmeani = 1/n*hprebn.sum(0, keepdim=True)\n",
    "bndiff = hprebn - bnmeani\n",
    "bndiff2 = bndiff**2\n",
    "bnvar = 1/(n-1)*(bndiff2).sum(0, keepdim=True) # note: Bessel's correction (dividing by n-1, not n)\n",
    "bnvar_inv = (bnvar + 1e-5)**-0.5\n",
    "bnraw = bndiff * bnvar_inv\n",
    "hpreact = bngain * bnraw + bnbias\n",
    "# Non-linearity\n",
    "h = torch.tanh(hpreact) # hidden layer\n",
    "# Linear layer 2\n",
    "logits = h @ W2 + b2 # output layer\n",
    "# cross entropy loss (same as F.cross_entropy(logits, Yb))\n",
    "logit_maxes = logits.max(1, keepdim=True).values\n",
    "norm_logits = logits - logit_maxes # subtract max for numerical stability\n",
    "counts = norm_logits.exp()\n",
    "counts_sum = counts.sum(1, keepdims=True)\n",
    "counts_sum_inv = counts_sum**-1 # if I use (1.0 / counts_sum) instead then I can't get backprop to be bit exact...\n",
    "probs = counts * counts_sum_inv\n",
    "logprobs = probs.log()\n",
    "loss = -logprobs[range(n), Yb].mean()\n",
    "\n",
    "# PyTorch backward pass\n",
    "for p in parameters:\n",
    "  p.grad = None\n",
    "for t in [logprobs, probs, counts, counts_sum, counts_sum_inv, # afaik there is no cleaner way\n",
    "          norm_logits, logit_maxes, logits, h, hpreact, bnraw,\n",
    "         bnvar_inv, bnvar, bndiff2, bndiff, hprebn, bnmeani,\n",
    "         embcat, emb]:\n",
    "  t.retain_grad()\n",
    "loss.backward()\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logprobs        | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "probs           | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "counts_sum_inv  | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "counts_sum      | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "counts          | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "norm_logits     | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "logit_maxes     | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "logits          | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "h               | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "W2              | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "b2              | exact: True  | approximate: True  | maxdiff: 0.0\n",
      "hpreact         | exact: False | approximate: True  | maxdiff: 4.656612873077393e-10\n",
      "bngain          | exact: False | approximate: True  | maxdiff: 3.725290298461914e-09\n",
      "bnbias          | exact: False | approximate: True  | maxdiff: 3.725290298461914e-09\n",
      "bnraw           | exact: False | approximate: True  | maxdiff: 6.984919309616089e-10\n",
      "bnvar_inv       | exact: False | approximate: True  | maxdiff: 3.725290298461914e-09\n",
      "bnvar           | exact: False | approximate: True  | maxdiff: 9.313225746154785e-10\n",
      "bndiff2         | exact: False | approximate: True  | maxdiff: 2.9103830456733704e-11\n",
      "bndiff          | exact: False | approximate: True  | maxdiff: 4.656612873077393e-10\n",
      "bnmeani         | exact: False | approximate: True  | maxdiff: 1.862645149230957e-09\n",
      "hprebn          | exact: False | approximate: True  | maxdiff: 4.656612873077393e-10\n",
      "embcat          | exact: False | approximate: True  | maxdiff: 1.862645149230957e-09\n",
      "W1              | exact: False | approximate: True  | maxdiff: 4.6566128730773926e-09\n",
      "b1              | exact: False | approximate: True  | maxdiff: 1.3969838619232178e-09\n",
      "emb             | exact: False | approximate: True  | maxdiff: 1.862645149230957e-09\n",
      "C               | exact: False | approximate: True  | maxdiff: 5.587935447692871e-09\n"
     ]
    }
   ],
   "source": [
    "# Exercise 1: backprop through the whole thing manually, \n",
    "# backpropagating through exactly all of the variables \n",
    "# as they are defined in the forward pass above, one by one\n",
    "\n",
    "dlogprobs = torch.zeros_like(logprobs)\n",
    "dlogprobs[range(n), Yb] = -1.0/n\n",
    "dprobs = (1.0 / probs) * dlogprobs\n",
    "dcounts_sum_inv = (counts * dprobs).sum(1, keepdim=True)\n",
    "dcounts = counts_sum_inv * dprobs\n",
    "dcounts_sum = (-counts_sum**-2) * dcounts_sum_inv\n",
    "dcounts += torch.ones_like(counts) * dcounts_sum\n",
    "dnorm_logits = counts * dcounts\n",
    "dlogits = dnorm_logits.clone()\n",
    "dlogit_maxes = (-dnorm_logits).sum(1, keepdim=True)\n",
    "dlogits += F.one_hot(logits.max(1).indices, num_classes=logits.shape[1]) * dlogit_maxes\n",
    "dh = dlogits @ W2.T\n",
    "dW2 = h.T @ dlogits\n",
    "db2 = dlogits.sum(0)\n",
    "dhpreact = (1.0 - h**2) * dh\n",
    "dbngain = (bnraw * dhpreact).sum(0, keepdim=True)\n",
    "dbnraw = bngain * dhpreact\n",
    "dbnbias = dhpreact.sum(0, keepdim=True)\n",
    "dbndiff = bnvar_inv * dbnraw\n",
    "dbnvar_inv = (bndiff * dbnraw).sum(0, keepdim=True)\n",
    "dbnvar = (-0.5*(bnvar + 1e-5)**-1.5) * dbnvar_inv\n",
    "dbndiff2 = (1.0/(n-1))*torch.ones_like(bndiff2) * dbnvar\n",
    "dbndiff += (2*bndiff) * dbndiff2\n",
    "dhprebn = dbndiff.clone()\n",
    "dbnmeani = (-dbndiff).sum(0)\n",
    "dhprebn += 1.0/n * (torch.ones_like(hprebn) * dbnmeani)\n",
    "dembcat = dhprebn @ W1.T\n",
    "dW1 = embcat.T @ dhprebn\n",
    "db1 = dhprebn.sum(0)\n",
    "demb = dembcat.view(emb.shape)\n",
    "dC = torch.zeros_like(C)\n",
    "for k in range(Xb.shape[0]):\n",
    "  for j in range(Xb.shape[1]):\n",
    "    ix = Xb[k,j]\n",
    "    dC[ix] += demb[k,j]\n",
    "    \n",
    "cmp('logprobs', dlogprobs, logprobs)\n",
    "cmp('probs', dprobs, probs)\n",
    "cmp('counts_sum_inv', dcounts_sum_inv, counts_sum_inv)\n",
    "cmp('counts_sum', dcounts_sum, counts_sum)\n",
    "cmp('counts', dcounts, counts)\n",
    "cmp('norm_logits', dnorm_logits, norm_logits)\n",
    "cmp('logit_maxes', dlogit_maxes, logit_maxes)\n",
    "cmp('logits', dlogits, logits)\n",
    "cmp('h', dh, h)\n",
    "cmp('W2', dW2, W2)\n",
    "cmp('b2', db2, b2)\n",
    "cmp('hpreact', dhpreact, hpreact)\n",
    "cmp('bngain', dbngain, bngain)\n",
    "cmp('bnbias', dbnbias, bnbias)\n",
    "cmp('bnraw', dbnraw, bnraw)\n",
    "cmp('bnvar_inv', dbnvar_inv, bnvar_inv)\n",
    "cmp('bnvar', dbnvar, bnvar)\n",
    "cmp('bndiff2', dbndiff2, bndiff2)\n",
    "cmp('bndiff', dbndiff, bndiff)\n",
    "cmp('bnmeani', dbnmeani, bnmeani)\n",
    "cmp('hprebn', dhprebn, hprebn)\n",
    "cmp('embcat', dembcat, embcat)\n",
    "cmp('W1', dW1, W1)\n",
    "cmp('b1', db1, b1)\n",
    "cmp('emb', demb, emb)\n",
    "cmp('C', dC, C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.347013235092163 diff: -2.384185791015625e-07\n"
     ]
    }
   ],
   "source": [
    "# Exercise 2: backprop through cross_entropy but all in one go\n",
    "# to complete this challenge look at the mathematical expression of the loss,\n",
    "# take the derivative, simplify the expression, and just write it out\n",
    "\n",
    "# forward pass\n",
    "\n",
    "# before:\n",
    "# logit_maxes = logits.max(1, keepdim=True).values\n",
    "# norm_logits = logits - logit_maxes # subtract max for numerical stability\n",
    "# counts = norm_logits.exp()\n",
    "# counts_sum = counts.sum(1, keepdims=True)\n",
    "# counts_sum_inv = counts_sum**-1 # if I use (1.0 / counts_sum) instead then I can't get backprop to be bit exact...\n",
    "# probs = counts * counts_sum_inv\n",
    "# logprobs = probs.log()\n",
    "# loss = -logprobs[range(n), Yb].mean()\n",
    "\n",
    "# now:\n",
    "loss_fast = F.cross_entropy(logits, Yb)\n",
    "print(loss_fast.item(), 'diff:', (loss_fast - loss).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "logits          | exact: False | approximate: True  | maxdiff: 5.820766091346741e-09\n"
     ]
    }
   ],
   "source": [
    "# backward pass\n",
    "\n",
    "dlogits = F.softmax(logits, 1)\n",
    "dlogits[range(n), Yb] -= 1\n",
    "dlogits /= n\n",
    "\n",
    "cmp('logits', dlogits, logits) # I can only get approximate to be true, my maxdiff is 6e-9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([32, 27]), torch.Size([32]))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logits.shape, Yb.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0724, 0.0869, 0.0204, 0.0468, 0.0190, 0.0792, 0.0238, 0.0353, 0.0184,\n",
       "        0.0300, 0.0399, 0.0340, 0.0374, 0.0301, 0.0380, 0.0140, 0.0093, 0.0192,\n",
       "        0.0149, 0.0549, 0.0497, 0.0216, 0.0254, 0.0703, 0.0589, 0.0267, 0.0234],\n",
       "       grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.softmax(logits, 1)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0724,  0.0869,  0.0204,  0.0468,  0.0190,  0.0792,  0.0238,  0.0353,\n",
       "        -0.9816,  0.0300,  0.0399,  0.0340,  0.0374,  0.0301,  0.0380,  0.0140,\n",
       "         0.0093,  0.0192,  0.0149,  0.0549,  0.0497,  0.0216,  0.0254,  0.0703,\n",
       "         0.0589,  0.0267,  0.0234], grad_fn=<MulBackward0>)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dlogits[0] * n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-2.3283e-10, grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dlogits[0].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f2cfeb068f0>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAATMAAAFgCAYAAADXQp4HAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8g+/7EAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAkIElEQVR4nO3de0xUd/o/8DcqM4LAsKjcVqR4qXfY1q2UtHVtZVU2abTSxF6S1cZodLFZZbtt2PS+m9C1Seu2ofpPV9Ok1q5J1bSbtWlpwXQX3ZWVRaqOQGm1EXCrgRlALsL5/dGfs04Vznvw8J3x4/uVTKLD4+d85pzD45k5z+eZKMuyLIiI3ORGhXsCIiJOUDITESMomYmIEZTMRMQISmYiYgQlMxExgpKZiBhByUxEjDAm3BP4oYGBAZw7dw7x8fGIiooK93REJIwsy4Lf70d6ejpGjRr62iviktm5c+eQkZER7mmISAQ5e/YsJk2aNGTMiCWzsrIyvPrqq2hpaUFOTg7efPNNLFiwwPbfxcfHAwCqq6sRFxc3ZOzo0aNtx/P5fNR8o6Ojqbi+vj7bGLt5X9HR0WEbw7xGAMjJyaHiampqbGPs/ge8wsmVcOxYzNX65cuXHd0msz/YscaOHUvFMeP19vZSYzH7LCYmhhprYGCAiuvp6bGNYV5jR0cHFi5cGMgLQxmRZPb++++juLgYO3bsQG5uLrZt24alS5fC6/UiOTl5yH97ZcfHxcXZvgDmF509ydhkxpxAzI4HuJOMTWbsW3Jmbjd7MmP+wwllm04mMzZp3OzJzOVy2caEcv4wr2FEbgC89tprWLduHZ544gnMnj0bO3bsQGxsLP785z+PxOZERJxPZr29vaiurkZ+fv7/NjJqFPLz81FVVXVNfE9PD3w+X9BDRCRUjiez7777Dv39/UhJSQl6PiUlBS0tLdfEl5aWwuPxBB768F9EhiPsdWYlJSVob28PPM6ePRvuKYnITcjxGwATJkzA6NGj0draGvR8a2srUlNTr4l3u91wu91OT0NEbjGOX5m5XC7Mnz8f5eXlgecGBgZQXl6OvLw8pzcnIgJghEoziouLsXr1avz0pz/FggULsG3bNnR2duKJJ54Yic2JiIxMMlu1ahX++9//4vnnn0dLSwt+8pOf4ODBg9fcFBjK5cuXbYsfmeLIpKQkantMASvA1X11dXVRYzF1NmzNl9frpeKYOiG25q6/v5+Kc7LQdfbs2bYxJ0+epMZia6aYOPY4sfuMqZVjawuZbbL7oru7m4pjfk+YbYaypHHEVgBs2rQJmzZtGqnhRUSChP1upoiIE5TMRMQISmYiYgQlMxExgpKZiBhByUxEjKBkJiJGiLi22Vf09PTYFm8yBXVskZ+TjfrYhopsQzwnx2I6gDIxAF8oysSxhbp1dXW2MdOnT6fGOnHiBBU3Zoz9rwl7/iQmJlJxTBE324SSmb+TY7HjsecPS1dmImIEJTMRMYKSmYgYQclMRIygZCYiRlAyExEjKJmJiBGUzETECEpmImKEiF0BMHr0aNtqYydbQLPteZmqZbaCnmln7OTKBIDbZ2yVN9sC2kljx461jfn666+psdjjxOwz9jj5/X4qjlm5wp6zTKvxL7/8khqLPc+Yb1xj9hl7LgK6MhMRQyiZiYgRlMxExAhKZiJiBCUzETGCkpmIGEHJTESMoGQmIkaI2KLZnJwc25jTp0/bxly+fJnaHlv06GShH1N0yhamMsWkgLMtoNk4priTPU7MWJmZmdRY9fX1VBxTAMpiC12ZYu/e3l5qLKY9OFMYDPDnNjM3tqCdpSszETGCkpmIGEHJTESMoGQmIkZQMhMRIyiZiYgRlMxExAhKZiJiBCUzETFCxK4A+M9//oP4+PgbHoetWHaybfalS5eosRhsZT9bDe5kq3F2dYKT22TivvnmG2osFtNem62gnzNnDhXHrG4ZPXo0NRbzO8CeP2xcQkKCbQzTGjwUjl+Zvfjii4iKigp6zJw50+nNiIgEGZErszlz5uDTTz/930ZC+FICEZHhGJEsM2bMGKSmpo7E0CIi1zUiNwDq6+uRnp6OKVOm4PHHH8eZM2cGje3p6YHP5wt6iIiEyvFklpubi127duHgwYPYvn07mpqacN999w36fYGlpaXweDyBR0ZGhtNTEpFbQJTFNqUapra2NmRmZuK1117D2rVrr/l5T09P0N0in8+HjIwM3c38/1wuFxUXjruZfX19VByzz9jjxMyNnRe7zxjhuJvJ/uo6eTeT5dTdTL/fj3nz5qG9vd12zBH/ZD4xMRG33347Ghoarvtzt9vtaPM7Ebk1jXjRbEdHBxobG5GWljbSmxKRW5jjyeypp55CZWUlvv76a/zjH//AQw89hNGjR+PRRx91elMiIgGOv8389ttv8eijj+LChQuYOHEi7r33Xhw+fBgTJ04MaZzo6Gjbz0e6urpsx2HfwnZ0dFBxzOcP7OcnMTExtjHs5yLs51zTp0+3jamrq6PGYj/nYvaHkxXocXFx1Fjs6grmM1BmlQDAfRYGcMedPebMWOxnxuwxH+yG39WYz1LZVSbACCSzPXv2OD2kiIgtLTQXESMomYmIEZTMRMQISmYiYgQlMxExgpKZiBhByUxEjBCxXRP7+/ttC+aYtsHsom+2qPfixYu2MU4uDmcX23d2dlJxtbW1tjFMMSPg7EJztoA1MzPTNsbr9VJjOdljgS06ZY8nU8TNbpM5TmwL7suXL1NxTLE6UxDLnouArsxExBBKZiJiBCUzETGCkpmIGEHJTESMoGQmIkZQMhMRIyiZiYgRlMxExAgRuwIgKirKtsKZqYBm2+62tbVRcUwF9KxZs6ixGhsbqTgG+zpDqai24+TX+LFtp0+dOkXFMdh9wVTHh9Le2SnsShNmbuxqAjaOWXnDtP0OZZWGrsxExAhKZiJiBCUzETGCkpmIGEHJTESMoGQmIkZQMhMRIyiZiYgRlMxExAgRuwKgr6/Ptnf5tGnTbMdhq+zZ3uZM1TLbg57pzc722U9ISKDimEp79nsT2BUADLYHPVOBzlaph6OC3ufzUXFMD32/30+NxXy/gtPHnPk9YfZrKCsrdGUmIkZQMhMRIyiZiYgRlMxExAhKZiJiBCUzETGCkpmIGEHJTESMELFFswMDA7YFc0wLZbY1Mhs3MDBgG8O2+mUKApntAUBHRwcVx7xOdl+whcZM0WZvby81FjO3SZMmUWO1tLRQcUxBL1tMyhanTp061TamtraWGqurq8s2hj3mbHGwU4XG7PaAYVyZHTp0CA8++CDS09MRFRWF/fv3B/3csiw8//zzSEtLQ0xMDPLz81FfXx/qZkREQhJyMuvs7EROTg7Kysqu+/OtW7fijTfewI4dO3DkyBGMGzcOS5cuRXd39w1PVkRkMCG/zSwoKEBBQcF1f2ZZFrZt24Znn30Wy5cvBwC88847SElJwf79+/HII4/c2GxFRAbh6A2ApqYmtLS0ID8/P/Ccx+NBbm4uqqqqrvtvenp64PP5gh4iIqFyNJld+UA1JSUl6PmUlJRBP2wtLS2Fx+MJPDIyMpyckojcIsJemlFSUoL29vbA4+zZs+GekojchBxNZqmpqQCA1tbWoOdbW1sDP/sht9uNhISEoIeISKgcTWZZWVlITU1FeXl54Dmfz4cjR44gLy/PyU2JiAQJ+W5mR0cHGhoaAn9vampCTU0NkpKSMHnyZGzevBl/+MMfMH36dGRlZeG5555Deno6VqxY4eS8RUSChJzMjh49ivvvvz/w9+LiYgDA6tWrsWvXLjz99NPo7OzE+vXr0dbWhnvvvRcHDx6kqsCvFhUVZVuVzFQts1Xqg5Wb/NBf//pX25iYmBhqLKZtM9s22MlVB+w22epsplU3OxazUuCrr76ixmJbdTNxbGU/e25cfcEwGPbcZuKcbFsOcCsimGPJntfAMJLZokWLhtxAVFQUXn75Zbz88suhDi0iMmxhv5spIuIEJTMRMYKSmYgYQclMRIygZCYiRlAyExEjKJmJiBEitm02YF8wxxQDssW6TDEs4GwBJbMOlS2MnDNnDhV38uRJ2xi2aJZtFc1g24MzRZvsMY+OjqbimOJOdiymgBjg9i1bwJqYmGgbc/HiRWostrjWqbHYdt6ArsxExBBKZiJiBCUzETGCkpmIGEHJTESMoGQmIkZQMhMRIyiZiYgRlMxExAgRvQLArsKZqQ5mK4jZamqmUj0uLo4aq7Oz0zaGrcavra2l4pg2xE5WeQPffwOXne7ubmqsnJwc2xiv10uNxa7UYIwbN46Ka29vp+KYfdbV1UWN1dfXZxvDtHB3GvM7pxUAInLLUTITESMomYmIEZTMRMQISmYiYgQlMxExgpKZiBhByUxEjKBkJiJGiNgVAC6Xy7YqmalsZvq3A1zFNcD1cGer2ZkK6NjYWGostoc+swKAxVZnT5061TbmxIkT1Fh1dXW2Mcx5AfD7gjk3Ojo6qLHY7ydg5saes+z3SDDYFSkM5vwPZXu6MhMRIyiZiYgRlMxExAhKZiJiBCUzETGCkpmIGEHJTESMoGQmIkaI2KLZO+64w7aorrGx0XYctoCSKYZlsS2UmUJLJwtwASA6OtqxsVj19fW2MWwLaKalN1tAzOwLgDsGMTEx1Fhsq+4xY+x/NdmCUqa4mS3AZfct8/vEFAaHUuQd8pXZoUOH8OCDDyI9PR1RUVHYv39/0M/XrFmDqKiooMeyZctC3YyISEhCTmadnZ3IyclBWVnZoDHLli1Dc3Nz4PHee+/d0CRFROyE/DazoKAABQUFQ8a43W6kpqYOe1IiIqEakRsAFRUVSE5OxowZM7Bx40ZcuHBh0Nienh74fL6gh4hIqBxPZsuWLcM777yD8vJy/PGPf0RlZSUKCgoG/bCytLQUHo8n8MjIyHB6SiJyC3D8buYjjzwS+PO8efOQnZ2NqVOnoqKiAosXL74mvqSkBMXFxYG/+3w+JTQRCdmI15lNmTIFEyZMQENDw3V/7na7kZCQEPQQEQnViCezb7/9FhcuXEBaWtpIb0pEbmEhv83s6OgIuspqampCTU0NkpKSkJSUhJdeegmFhYVITU1FY2Mjnn76aUybNg1Lly51dOIiIleLskLso1xRUYH777//mudXr16N7du3Y8WKFTh27Bja2tqQnp6OJUuW4Pe//z1SUlKo8X0+HzweD2pqahAfHz9kLFON7PF4qO12dnZScUxltpNV0mxranabDLYaPCsri4pjVmowlf2hxDHYY85wcgUGwLW6ZlcAONmCm90mc5yY3yW/34/Zs2ejvb3d9iOokK/MFi1aNOTO+fjjj0MdUkTkhmmhuYgYQclMRIygZCYiRlAyExEjKJmJiBGUzETECEpmImIEJTMRMULEfgdAbm6ubVX1119/bTuO0z30me8UYBdVMNtkv0+ArWZnVgowldkAcPLkSSqOqWZnv6uBqaB3sjIecLYfP7uCgdkfTlbts/vf6e+HcJKuzETECEpmImIEJTMRMYKSmYgYQclMRIygZCYiRlAyExEjKJmJiBEitmi2qqoKcXFxQ8a0t7fbjjN27Fhqe0wLa4AremQLKJlvomKLftkCSqZoli3AZVtAM9i23729vbYx7DGPjY11bJtse3P2PGP2LVv0a9d+HgAuXrxIjcUW/TKF0tOmTbONCaWrv67MRMQISmYiYgQlMxExgpKZiBhByUxEjKBkJiJGUDITESMomYmIEZTMRMQIEbsCALBv0cu08GUry1lspTeDqW52ss0yAEyfPt025vTp09RYbAtlpu00W+l96dIl2xh2XzBV6gC3ooM9TsyqDwDo6uqyjWHPbWZFB7tqgl3dwszN6/Xaxvj9ftxxxx3UNnVlJiJGUDITESMomYmIEZTMRMQISmYiYgQlMxExgpKZiBhByUxEjBCxRbNut9u2FbSTBZRsC2imGJAtrGUKI9liTKYwFQBOnTplG+N0q3Enx2Lag7MtxP1+PxXHFAezBcTMOQs426qbKXRli5bZ8zE7O9s25sSJE45tD9CVmYgYIqRkVlpairvuugvx8fFITk7GihUrrlmS0N3djaKiIowfPx5xcXEoLCxEa2uro5MWEfmhkJJZZWUlioqKcPjwYXzyySfo6+vDkiVLgtZ+bdmyBR9++CH27t2LyspKnDt3DitXrnR84iIiVwvpM7ODBw8G/X3Xrl1ITk5GdXU1Fi5ciPb2drz99tvYvXs3HnjgAQDAzp07MWvWLBw+fBh33323czMXEbnKDX1mduV7K5OSkgAA1dXV6OvrQ35+fiBm5syZmDx5Mqqqqq47Rk9PD3w+X9BDRCRUw05mAwMD2Lx5M+655x7MnTsXANDS0gKXy4XExMSg2JSUFLS0tFx3nNLSUng8nsAjIyNjuFMSkVvYsJNZUVER6urqsGfPnhuaQElJCdrb2wOPs2fP3tB4InJrGlad2aZNm/DRRx/h0KFDmDRpUuD51NRU9Pb2oq2tLejqrLW1Fampqdcdi6knExGxE9KVmWVZ2LRpE/bt24fPPvsMWVlZQT+fP38+oqOjUV5eHnjO6/XizJkzyMvLc2bGIiLXEdKVWVFREXbv3o0DBw4gPj4+8DmYx+NBTEwMPB4P1q5di+LiYiQlJSEhIQFPPvkk8vLyQr6Teeedd9pWVX/11Ve247Btfp2MY6vxmdUJbGU5W0HPVHqz7aTZts3d3d22MWw1OzN/pno+FEwVOrvP4uPjqThmn7GcPGdZx48fd3Q8RkivYPv27QCARYsWBT2/c+dOrFmzBgDw+uuvY9SoUSgsLERPTw+WLl2Kt956y5HJiogMJqRkxvyvOHbsWJSVlaGsrGzYkxIRCZXWZoqIEZTMRMQISmYiYgQlMxExgpKZiBhByUxEjKBkJiJGiNjvAKiqqkJcXNyQMenp6bbjfPPNN9T22Ap6phqc7fOekJBgG8NWgrN9+5mqfXZfOFk1zq4mYKr72X1hd36Fsk32OySutM2yw6xXZvv2M+fZxYsXqbFC6clvhznm7HkB6MpMRAyhZCYiRlAyExEjKJmJiBGUzETECEpmImIEJTMRMYKSmYgYIWKLZl0ulyNfdMK0pg4FMye26JRptcwWDTpZ9OtkYSTAFXey7cFdLteNTidkzDnE7jMni4NZTHEz27acLUhm5s+081bRrIjccpTMRMQISmYiYgQlMxExgpKZiBhByUxEjKBkJiJGUDITESMomYmIESJ2BUB/f79thXBLS4vtOH6/n9oeu9qAqbRnx2JaYs+cOZMa69SpU1QcU3WdmJhIjcW2Wmaqy5l5AdwKAHY1BBvHrGBgK9XZlQLM/mCr9s+fP28bc9tttzk2FsDtM+b3JJSVELoyExEjKJmJiBGUzETECEpmImIEJTMRMYKSmYgYQclMRIygZCYiRlAyExEjROwKALfbbVsh3NHRYTsOU4kM8JXGTAU303OdjTt9+jQ1Fvs6marx9vZ2aiy2HzzT35/d/8z3JrD7gj1OTHV/dnY2NVZNTQ0Vx5xn7OuMj4+3jWltbaXGio6OpuKYFQzMChh2lQYQ4pVZaWkp7rrrLsTHxyM5ORkrVqyA1+sNilm0aBGioqKCHhs2bAhlMyIiIQspmVVWVqKoqAiHDx/GJ598gr6+PixZsgSdnZ1BcevWrUNzc3PgsXXrVkcnLSLyQyG9zTx48GDQ33ft2oXk5GRUV1dj4cKFgedjY2ORmprqzAxFRAg3dAPgymcrSUlJQc+/++67mDBhAubOnYuSkhJ0dXUNOkZPTw98Pl/QQ0QkVMO+ATAwMIDNmzfjnnvuwdy5cwPPP/bYY8jMzER6ejpqa2vxzDPPwOv14oMPPrjuOKWlpXjppZeGOw0REQA3kMyKiopQV1eHL774Iuj59evXB/48b948pKWlYfHixWhsbMTUqVOvGaekpATFxcWBv/t8PmRkZAx3WiJyixpWMtu0aRM++ugjHDp0CJMmTRoyNjc3FwDQ0NBw3WTGlGCIiNgJKZlZloUnn3wS+/btQ0VFBbKysmz/zZW6mrS0tGFNUESEEVIyKyoqwu7du3HgwAHEx8cH2lZ7PB7ExMSgsbERu3fvxi9+8QuMHz8etbW12LJlCxYuXEgXFV7R29trW0jJFA0yBZsA37aZuYpkW3UnJCTYxly6dIkai23bPGvWLNuY48ePU2OxLaDZOAZzPNljzhaAMoWb7D5jW10z5yO7X5nzrLm5mRqLPc+Y+TO/v2xhMBBiMtu+fTuA7wtjr7Zz506sWbMGLpcLn376KbZt24bOzk5kZGSgsLAQzz77bCibEREJWchvM4eSkZGBysrKG5qQiMhwaKG5iBhByUxEjKBkJiJGUDITESMomYmIEZTMRMQISmYiYoSIbZvd399vW0XMVHq7XC5qe5mZmVRcY2OjbQxbgT5Ua6Qr2IprtrK8vr7eNoZtVcyummBfA4Npdc1W9jMtuNnx2GPO7luPx2Mb09bWRo114cIF2xi20r6vr4+KY1YnxMTE2MawxwjQlZmIGELJTESMoGQmIkZQMhMRIyiZiYgRlMxExAhKZiJiBCUzETFCxBbNxsTE2BbVMQWI3d3d1Pa8Xi8Vx2BbhJ84ccI2hm2N7GShK1t0yhbNOtVCmY1ji3SZok0A6OzstI1hv5SHLW7u6OiwjWHPDWafxcbGUmOxRehMQS9znOxa519NV2YiYgQlMxExgpKZiBhByUxEjKBkJiJGUDITESMomYmIEZTMRMQISmYiYoSIXQHQ1dVlW+HMVDYzbZbZsQCu6rq2tpYai5kbWwHNtFkGgEmTJtnGnD59mhqLxVS9s22nGWPHjqXimLblLLadNPs6mX3GrsBgzll2X7DnI3MMmBUA7O8voCszETGEkpmIGEHJTESMoGQmIkZQMhMRIyiZiYgRlMxExAhKZiJiBCUzETFCxK4AWLBggW21dENDg+04bGU228PdyR76zNzYinGmTz0AnDx50rFtshXozOoKtjc+u00G+zqdnL+TffvZfcF8P0R8fDw1Fvs6/X6/bQyz/0M53iFdmW3fvh3Z2dlISEhAQkIC8vLy8Le//S3w8+7ubhQVFWH8+PGIi4tDYWEhWltbQ9mEiMiwhJTMJk2ahFdeeQXV1dU4evQoHnjgASxfvhxffvklAGDLli348MMPsXfvXlRWVuLcuXNYuXLliExcRORqURa7wnoQSUlJePXVV/Hwww9j4sSJ2L17Nx5++GEAwKlTpzBr1ixUVVXh7rvvpsbz+XzweDwYM2bMTfs2k8V+JZqTYzm5uDdS32ayb/NZTu4zJ99msou+mbHGjRtHjfV//TbT7/fjjjvuQHt7OxISEoaeGzWz6+jv78eePXvQ2dmJvLw8VFdXo6+vD/n5+YGYmTNnYvLkyaiqqhp0nJ6eHvh8vqCHiEioQk5mx48fR1xcHNxuNzZs2IB9+/Zh9uzZaGlpgcvlQmJiYlB8SkoKWlpaBh2vtLQUHo8n8MjIyAj5RYiIhJzMZsyYgZqaGhw5cgQbN27E6tWrqW/mHkxJSQna29sDj7Nnzw57LBG5dYVcmuFyuTBt2jQAwPz58/Gvf/0Lf/rTn7Bq1Sr09vaira0t6OqstbUVqampg47ndrvpz6tERAZzw0WzAwMD6Onpwfz58xEdHY3y8vLAz7xeL86cOYO8vLwb3YyIyJBCujIrKSlBQUEBJk+eDL/fj927d6OiogIff/wxPB4P1q5di+LiYiQlJSEhIQFPPvkk8vLy6DuZV6upqbEt5GPu5sTGxlLbY4tOmeJCdizmLpnTxaTM/mCKLAF+boyYmBgqjp0bg50/c9ft9ttvp8Y6fvw4Fcfsj8uXL1NjMXcq2XOWxdzdZeYfSrFFSMns/Pnz+OUvf4nm5mZ4PB5kZ2fj448/xs9//nMAwOuvv45Ro0ahsLAQPT09WLp0Kd56661QNiEiMiwhJbO33357yJ+PHTsWZWVlKCsru6FJiYiESgvNRcQISmYiYgQlMxExgpKZiBhByUxEjKBkJiJGiLhOs1eK5Do6OmxjmaJZtpi0q6uLimOwBYhMQSDbDZVtAcTsD7YwlZ0bgy0ADUfRLLNv2eJO5rwGuP1x6dIlaixm/t3d3dRYLKeKZq/sL+p35Ub7mTnt22+/VecMEQly9uxZTJo0aciYiEtmAwMDOHfuHOLj4wP/8/t8PmRkZODs2bO2DdoikeYffjf7a7hV529ZFvx+P9LT022vpCPubeaoUaMGzcBXvnvgZqX5h9/N/hpuxfl7PB4qTjcARMQISmYiYoSbIpm53W688MILN20TR80//G7216D524u4GwAiIsNxU1yZiYjYUTITESMomYmIEZTMRMQIN0UyKysrw2233YaxY8ciNzcX//znP8M9JcqLL76IqKiooMfMmTPDPa1BHTp0CA8++CDS09MRFRWF/fv3B/3csiw8//zzSEtLQ0xMDPLz81FfXx+eyV6H3fzXrFlzzfFYtmxZeCZ7HaWlpbjrrrsQHx+P5ORkrFixAl6vNyimu7sbRUVFGD9+POLi4lBYWIjW1tYwzTgYM/9FixZdcww2bNjgyPYjPpm9//77KC4uxgsvvIB///vfyMnJwdKlS3H+/PlwT40yZ84cNDc3Bx5ffPFFuKc0qM7OTuTk5Az6HQ5bt27FG2+8gR07duDIkSMYN24cli5d6vgi5eGymz8ALFu2LOh4vPfee/+HMxxaZWUlioqKcPjwYXzyySfo6+vDkiVLghoXbNmyBR9++CH27t2LyspKnDt3DitXrgzjrP+HmT8ArFu3LugYbN261ZkJWBFuwYIFVlFRUeDv/f39Vnp6ulVaWhrGWXFeeOEFKycnJ9zTGBYA1r59+wJ/HxgYsFJTU61XX3018FxbW5vldrut9957LwwzHNoP529ZlrV69Wpr+fLlYZnPcJw/f94CYFVWVlqW9f3+jo6Otvbu3RuIOXnypAXAqqqqCtc0B/XD+VuWZf3sZz+zfv3rX4/I9iL6yqy3txfV1dXIz88PPDdq1Cjk5+ejqqoqjDPj1dfXIz09HVOmTMHjjz+OM2fOhHtKw9LU1ISWlpagY+HxeJCbm3vTHAsAqKioQHJyMmbMmIGNGzfiwoUL4Z7SoNrb2wEASUlJAIDq6mr09fUFHYOZM2di8uTJEXkMfjj/K959911MmDABc+fORUlJiWPttyJuofnVvvvuO/T39yMlJSXo+ZSUFJw6dSpMs+Ll5uZi165dmDFjBpqbm/HSSy/hvvvuQ11dHfVlwpGkpaUFAK57LK78LNItW7YMK1euRFZWFhobG/G73/0OBQUFqKqqwujRo8M9vSADAwPYvHkz7rnnHsydOxfA98fA5XIhMTExKDYSj8H15g8Ajz32GDIzM5Geno7a2lo888wz8Hq9+OCDD254mxGdzG52BQUFgT9nZ2cjNzcXmZmZ+Mtf/oK1a9eGcWa3pkceeSTw53nz5iE7OxtTp05FRUUFFi9eHMaZXauoqAh1dXUR/RnrUAab//r16wN/njdvHtLS0rB48WI0NjZi6tSpN7TNiH6bOWHCBIwePfqauzWtra1ITU0N06yGLzExEbfffjsaGhrCPZWQXdnfphwLAJgyZQomTJgQccdj06ZN+Oijj/D5558HtcNKTU1Fb28v2traguIj7RgMNv/ryc3NBQBHjkFEJzOXy4X58+ejvLw88NzAwADKy8uRl5cXxpkNT0dHBxobG5GWlhbuqYQsKysLqampQcfC5/PhyJEjN+WxAL7vanzhwoWIOR6WZWHTpk3Yt28fPvvsM2RlZQX9fP78+YiOjg46Bl6vF2fOnImIY2A3/+upqakBAGeOwYjcVnDQnj17LLfbbe3atcs6ceKEtX79eisxMdFqaWkJ99Rs/eY3v7EqKiqspqYm6+9//7uVn59vTZgwwTp//ny4p3Zdfr/fOnbsmHXs2DELgPXaa69Zx44ds7755hvLsizrlVdesRITE60DBw5YtbW11vLly62srCzr0qVLYZ7594aav9/vt5566imrqqrKampqsj799FPrzjvvtKZPn251d3eHe+qWZVnWxo0bLY/HY1VUVFjNzc2BR1dXVyBmw4YN1uTJk63PPvvMOnr0qJWXl2fl5eWFcdb/Yzf/hoYG6+WXX7aOHj1qNTU1WQcOHLCmTJliLVy40JHtR3wysyzLevPNN63JkydbLpfLWrBggXX48OFwT4myatUqKy0tzXK5XNaPf/xja9WqVVZDQ0O4pzWozz//3AJwzWP16tWWZX1fnvHcc89ZKSkpltvtthYvXmx5vd7wTvoqQ82/q6vLWrJkiTVx4kQrOjrayszMtNatWxdR/yleb+4ArJ07dwZiLl26ZP3qV7+yfvSjH1mxsbHWQw89ZDU3N4dv0lexm/+ZM2eshQsXWklJSZbb7bamTZtm/fa3v7Xa29sd2b5aAImIESL6MzMREZaSmYgYQclMRIygZCYiRlAyExEjKJmJiBGUzETECEpmImIEJTMRMYKSmYgYQclMRIygZCYiRvh/4udmDx3GLhsAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(4, 4))\n",
    "plt.imshow(dlogits.detach(), cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max diff: tensor(4.7684e-07, grad_fn=<MaxBackward1>)\n"
     ]
    }
   ],
   "source": [
    "# Exercise 3: backprop through batchnorm but all in one go\n",
    "# to complete this challenge look at the mathematical expression of the output of batchnorm,\n",
    "# take the derivative w.r.t. its input, simplify the expression, and just write it out\n",
    "\n",
    "# forward pass\n",
    "\n",
    "# before:\n",
    "# bnmeani = 1/n*hprebn.sum(0, keepdim=True)\n",
    "# bndiff = hprebn - bnmeani\n",
    "# bndiff2 = bndiff**2\n",
    "# bnvar = 1/(n-1)*(bndiff2).sum(0, keepdim=True) # note: Bessel's correction (dividing by n-1, not n)\n",
    "# bnvar_inv = (bnvar + 1e-5)**-0.5\n",
    "# bnraw = bndiff * bnvar_inv\n",
    "# hpreact = bngain * bnraw + bnbias\n",
    "\n",
    "# now:\n",
    "hpreact_fast = bngain * (hprebn - hprebn.mean(0, keepdim=True)) / torch.sqrt(hprebn.var(0, keepdim=True, unbiased=True) + 1e-5) + bnbias\n",
    "print('max diff:', (hpreact_fast - hpreact).abs().max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hprebn          | exact: False | approximate: True  | maxdiff: 9.313225746154785e-10\n"
     ]
    }
   ],
   "source": [
    "# backward pass\n",
    "\n",
    "# before we had:\n",
    "# dbnraw = bngain * dhpreact\n",
    "# dbndiff = bnvar_inv * dbnraw\n",
    "# dbnvar_inv = (bndiff * dbnraw).sum(0, keepdim=True)\n",
    "# dbnvar = (-0.5*(bnvar + 1e-5)**-1.5) * dbnvar_inv\n",
    "# dbndiff2 = (1.0/(n-1))*torch.ones_like(bndiff2) * dbnvar\n",
    "# dbndiff += (2*bndiff) * dbndiff2\n",
    "# dhprebn = dbndiff.clone()\n",
    "# dbnmeani = (-dbndiff).sum(0)\n",
    "# dhprebn += 1.0/n * (torch.ones_like(hprebn) * dbnmeani)\n",
    "\n",
    "# calculate dhprebn given dhpreact (i.e. backprop through the batchnorm)\n",
    "# (you'll also need to use some of the variables from the forward pass up above)\n",
    "\n",
    "dhprebn = bngain*bnvar_inv/n * (n*dhpreact - dhpreact.sum(0) - n/(n-1)*bnraw*(dhpreact*bnraw).sum(0))\n",
    "\n",
    "cmp('hprebn', dhprebn, hprebn) # I can only get approximate to be true, my maxdiff is 9e-10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([32, 64]),\n",
       " torch.Size([1, 64]),\n",
       " torch.Size([1, 64]),\n",
       " torch.Size([32, 64]),\n",
       " torch.Size([64]))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dhprebn.shape, bngain.shape, bnvar_inv.shape, dbnraw.shape, dbnraw.sum(0).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12297\n",
      "      0/ 200000: 3.7557\n",
      "  10000/ 200000: 2.1838\n",
      "  20000/ 200000: 2.4098\n",
      "  30000/ 200000: 2.4363\n",
      "  40000/ 200000: 2.0216\n",
      "  50000/ 200000: 2.4138\n",
      "  60000/ 200000: 2.3818\n",
      "  70000/ 200000: 2.0321\n",
      "  80000/ 200000: 2.4064\n",
      "  90000/ 200000: 2.0996\n",
      " 100000/ 200000: 1.9639\n",
      " 110000/ 200000: 2.3385\n",
      " 120000/ 200000: 2.0382\n",
      " 130000/ 200000: 2.4536\n",
      " 140000/ 200000: 2.3317\n",
      " 150000/ 200000: 2.2015\n",
      " 160000/ 200000: 1.9416\n",
      " 170000/ 200000: 1.8330\n",
      " 180000/ 200000: 1.9010\n",
      " 190000/ 200000: 1.8763\n"
     ]
    }
   ],
   "source": [
    "# Exercise 4: putting it all together!\n",
    "# Train the MLP neural net with your own backward pass\n",
    "\n",
    "# init\n",
    "n_embd = 10 # the dimensionality of the character embedding vectors\n",
    "n_hidden = 200 # the number of neurons in the hidden layer of the MLP\n",
    "\n",
    "g = torch.Generator().manual_seed(2147483647) # for reproducibility\n",
    "C  = torch.randn((vocab_size, n_embd),            generator=g)\n",
    "# Layer 1\n",
    "W1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * (5/3)/((n_embd * block_size)**0.5)\n",
    "b1 = torch.randn(n_hidden,                        generator=g) * 0.1\n",
    "# Layer 2\n",
    "W2 = torch.randn((n_hidden, vocab_size),          generator=g) * 0.1\n",
    "b2 = torch.randn(vocab_size,                      generator=g) * 0.1\n",
    "# BatchNorm parameters\n",
    "bngain = torch.randn((1, n_hidden))*0.1 + 1.0\n",
    "bnbias = torch.randn((1, n_hidden))*0.1\n",
    "\n",
    "parameters = [C, W1, b1, W2, b2, bngain, bnbias]\n",
    "print(sum(p.nelement() for p in parameters)) # number of parameters in total\n",
    "for p in parameters:\n",
    "  p.requires_grad = True\n",
    "\n",
    "# same optimization as last time\n",
    "max_steps = 200000\n",
    "batch_size = 32\n",
    "n = batch_size # convenience\n",
    "lossi = []\n",
    "\n",
    "# use this context manager for efficiency once your backward pass is written (TODO)\n",
    "with torch.no_grad():\n",
    "\n",
    "  # kick off optimization\n",
    "  for i in range(max_steps):\n",
    "\n",
    "    # minibatch construct\n",
    "    ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)\n",
    "    Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y\n",
    "\n",
    "    # forward pass\n",
    "    emb = C[Xb] # embed the characters into vectors\n",
    "    embcat = emb.view(emb.shape[0], -1) # concatenate the vectors\n",
    "    # Linear layer\n",
    "    hprebn = embcat @ W1 + b1 # hidden layer pre-activation\n",
    "    # BatchNorm layer\n",
    "    # -------------------------------------------------------------\n",
    "    bnmean = hprebn.mean(0, keepdim=True)\n",
    "    bnvar = hprebn.var(0, keepdim=True, unbiased=True)\n",
    "    bnvar_inv = (bnvar + 1e-5)**-0.5\n",
    "    bnraw = (hprebn - bnmean) * bnvar_inv\n",
    "    hpreact = bngain * bnraw + bnbias\n",
    "    # -------------------------------------------------------------\n",
    "    # Non-linearity\n",
    "    h = torch.tanh(hpreact) # hidden layer\n",
    "    logits = h @ W2 + b2 # output layer\n",
    "    loss = F.cross_entropy(logits, Yb) # loss function\n",
    "\n",
    "    # backward pass\n",
    "    for p in parameters:\n",
    "      p.grad = None\n",
    "    #loss.backward() # use this for correctness comparisons, delete it later!\n",
    "\n",
    "    # manual backprop! #swole_doge_meme\n",
    "    # -----------------\n",
    "    dlogits = F.softmax(logits, 1)\n",
    "    dlogits[range(n), Yb] -= 1\n",
    "    dlogits /= n\n",
    "    # 2nd layer backprop\n",
    "    dh = dlogits @ W2.T\n",
    "    dW2 = h.T @ dlogits\n",
    "    db2 = dlogits.sum(0)\n",
    "    # tanh\n",
    "    dhpreact = (1.0 - h**2) * dh\n",
    "    # batchnorm backprop\n",
    "    dbngain = (bnraw * dhpreact).sum(0, keepdim=True)\n",
    "    dbnbias = dhpreact.sum(0, keepdim=True)\n",
    "    dhprebn = bngain*bnvar_inv/n * (n*dhpreact - dhpreact.sum(0) - n/(n-1)*bnraw*(dhpreact*bnraw).sum(0))\n",
    "    # 1st layer\n",
    "    dembcat = dhprebn @ W1.T\n",
    "    dW1 = embcat.T @ dhprebn\n",
    "    db1 = dhprebn.sum(0)\n",
    "    # embedding\n",
    "    demb = dembcat.view(emb.shape)\n",
    "    dC = torch.zeros_like(C)\n",
    "    for k in range(Xb.shape[0]):\n",
    "      for j in range(Xb.shape[1]):\n",
    "        ix = Xb[k,j]\n",
    "        dC[ix] += demb[k,j]\n",
    "    grads = [dC, dW1, db1, dW2, db2, dbngain, dbnbias]\n",
    "    # -----------------\n",
    "\n",
    "    # update\n",
    "    lr = 0.1 if i < 100000 else 0.01 # step learning rate decay\n",
    "    for p, grad in zip(parameters, grads):\n",
    "      #p.data += -lr * p.grad # old way of cheems doge (using PyTorch grad from .backward())\n",
    "      p.data += -lr * grad # new way of swole doge TODO: enable\n",
    "\n",
    "    # track stats\n",
    "    if i % 10000 == 0: # print every once in a while\n",
    "      print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')\n",
    "    lossi.append(loss.log10().item())\n",
    "\n",
    "  #   if i >= 100: # TODO: delete early breaking when you're ready to train the full net\n",
    "  #     break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# useful for checking your gradients\n",
    "# for p,g in zip(parameters, grads):\n",
    "#   cmp(str(tuple(p.shape)), g, p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calibrate the batch norm at the end of training\n",
    "\n",
    "with torch.no_grad():\n",
    "  # pass the training set through\n",
    "  emb = C[Xtr]\n",
    "  embcat = emb.view(emb.shape[0], -1)\n",
    "  hpreact = embcat @ W1 + b1\n",
    "  # measure the mean/std over the entire training set\n",
    "  bnmean = hpreact.mean(0, keepdim=True)\n",
    "  bnvar = hpreact.var(0, keepdim=True, unbiased=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train 2.0711452960968018\n",
      "val 2.1081483364105225\n"
     ]
    }
   ],
   "source": [
    "# evaluate train and val loss\n",
    "\n",
    "@torch.no_grad() # this decorator disables gradient tracking\n",
    "def split_loss(split):\n",
    "  x,y = {\n",
    "    'train': (Xtr, Ytr),\n",
    "    'val': (Xdev, Ydev),\n",
    "    'test': (Xte, Yte),\n",
    "  }[split]\n",
    "  emb = C[x] # (N, block_size, n_embd)\n",
    "  embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size * n_embd)\n",
    "  hpreact = embcat @ W1 + b1\n",
    "  hpreact = bngain * (hpreact - bnmean) * (bnvar + 1e-5)**-0.5 + bnbias\n",
    "  h = torch.tanh(hpreact) # (N, n_hidden)\n",
    "  logits = h @ W2 + b2 # (N, vocab_size)\n",
    "  loss = F.cross_entropy(logits, y)\n",
    "  print(split, loss.item())\n",
    "\n",
    "split_loss('train')\n",
    "split_loss('val')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# I achieved:\n",
    "# train 2.0718822479248047\n",
    "# val 2.1162495613098145"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mora.\n",
      "mayah.\n",
      "seel.\n",
      "ndheyla.\n",
      "ren.\n",
      "rugendraegustered.\n",
      "elin.\n",
      "shi.\n",
      "jen.\n",
      "eden.\n",
      "sana.\n",
      "arleigh.\n",
      "malaia.\n",
      "noshubergihi.\n",
      "jest.\n",
      "jair.\n",
      "jennex.\n",
      "teron.\n",
      "ubrence.\n",
      "ryyah.\n"
     ]
    }
   ],
   "source": [
    "# sample from the model\n",
    "g = torch.Generator().manual_seed(2147483647 + 10)\n",
    "\n",
    "for _ in range(20):\n",
    "    \n",
    "    out = []\n",
    "    context = [0] * block_size # initialize with all ...\n",
    "    while True:\n",
    "      # ------------\n",
    "      # forward pass:\n",
    "      # Embedding\n",
    "      emb = C[torch.tensor([context])] # (1,block_size,d)      \n",
    "      embcat = emb.view(emb.shape[0], -1) # concat into (N, block_size * n_embd)\n",
    "      hpreact = embcat @ W1 + b1\n",
    "      hpreact = bngain * (hpreact - bnmean) * (bnvar + 1e-5)**-0.5 + bnbias\n",
    "      h = torch.tanh(hpreact) # (N, n_hidden)\n",
    "      logits = h @ W2 + b2 # (N, vocab_size)\n",
    "      # ------------\n",
    "      # Sample\n",
    "      probs = F.softmax(logits, dim=1)\n",
    "      ix = torch.multinomial(probs, num_samples=1, generator=g).item()\n",
    "      context = context[1:] + [ix]\n",
    "      out.append(ix)\n",
    "      if ix == 0:\n",
    "        break\n",
    "    \n",
    "    print(''.join(itos[i] for i in out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}