demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A cute little demo showing the simplest usage of minGPT. Configured to run fine on Macbook Air in like a minute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "from mingpt.utils import set_seed\n",
    "set_seed(3407)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "class SortDataset(Dataset):\n",
    "    \"\"\" \n",
    "    Dataset for the Sort problem. E.g. for problem length 6:\n",
    "    Input: 0 0 2 1 0 1 -> Output: 0 0 0 1 1 2\n",
    "    Which will feed into the transformer concatenated as:\n",
    "    input:  0 0 2 1 0 1 0 0 0 1 1\n",
    "    output: I I I I I 0 0 0 1 1 2\n",
    "    where I is \"ignore\", as the transformer is reading the input sequence\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, split, length=6, num_digits=3):\n",
    "        assert split in {'train', 'test'}\n",
    "        self.split = split\n",
    "        self.length = length\n",
    "        self.num_digits = num_digits\n",
    "    \n",
    "    def __len__(self):\n",
    "        return 10000 # ...\n",
    "    \n",
    "    def get_vocab_size(self):\n",
    "        return self.num_digits\n",
    "    \n",
    "    def get_block_size(self):\n",
    "        # the length of the sequence that will feed into transformer, \n",
    "        # containing concatenated input and the output, but -1 because\n",
    "        # the transformer starts making predictions at the last input element\n",
    "        return self.length * 2 - 1\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \n",
    "        # use rejection sampling to generate an input example from the desired split\n",
    "        while True:\n",
    "            # generate some random integers\n",
    "            inp = torch.randint(self.num_digits, size=(self.length,), dtype=torch.long)\n",
    "            # half of the time let's try to boost the number of examples that \n",
    "            # have a large number of repeats, as this is what the model seems to struggle\n",
    "            # with later in training, and they are kind of rate\n",
    "            if torch.rand(1).item() < 0.5:\n",
    "                if inp.unique().nelement() > self.length // 2:\n",
    "                    # too many unqiue digits, re-sample\n",
    "                    continue\n",
    "            # figure out if this generated example is train or test based on its hash\n",
    "            h = hash(pickle.dumps(inp.tolist()))\n",
    "            inp_split = 'test' if h % 4 == 0 else 'train' # designate 25% of examples as test\n",
    "            if inp_split == self.split:\n",
    "                break # ok\n",
    "        \n",
    "        # solve the task: i.e. sort\n",
    "        sol = torch.sort(inp)[0]\n",
    "\n",
    "        # concatenate the problem specification and the solution\n",
    "        cat = torch.cat((inp, sol), dim=0)\n",
    "\n",
    "        # the inputs to the transformer will be the offset sequence\n",
    "        x = cat[:-1].clone()\n",
    "        y = cat[1:].clone()\n",
    "        # we only want to predict at output locations, mask out the loss at the input locations\n",
    "        y[:self.length-1] = -1\n",
    "        return x, y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 -1\n",
      "0 -1\n",
      "1 -1\n",
      "0 -1\n",
      "0 -1\n",
      "0 0\n",
      "0 0\n",
      "0 0\n",
      "0 0\n",
      "0 1\n",
      "1 1\n"
     ]
    }
   ],
   "source": [
    "# print an example instance of the dataset\n",
    "train_dataset = SortDataset('train')\n",
    "test_dataset = SortDataset('test')\n",
    "x, y = train_dataset[0]\n",
    "for a, b in zip(x,y):\n",
    "    print(int(a),int(b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of parameters: 0.09M\n"
     ]
    }
   ],
   "source": [
    "# create a GPT instance\n",
    "from minbert.model import BERT\n",
    "\n",
    "model_config = BERT.get_default_config()\n",
    "model_config.model_type = 'gpt-nano'\n",
    "model_config.vocab_size = train_dataset.get_vocab_size()\n",
    "model_config.block_size = train_dataset.get_block_size()\n",
    "model = BERT(model_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running on device cpu\n"
     ]
    }
   ],
   "source": [
    "# create a Trainer object\n",
    "from minbert.trainer import Trainer\n",
    "\n",
    "train_config = Trainer.get_default_config()\n",
    "train_config.learning_rate = 5e-4 # the model we're using is so small that we can go a bit faster\n",
    "train_config.max_iters = 2000\n",
    "train_config.num_workers = 0\n",
    "train_config.device = 'cpu'\n",
    "trainer = Trainer(train_config, model, train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter_dt 0.00ms; iter 0: train loss 1.06215\n",
      "iter_dt 13.99ms; iter 100: train loss 0.21042\n",
      "iter_dt 14.83ms; iter 200: train loss 0.07224\n",
      "iter_dt 16.20ms; iter 300: train loss 0.03696\n",
      "iter_dt 14.11ms; iter 400: train loss 0.01761\n",
      "iter_dt 14.02ms; iter 500: train loss 0.01130\n",
      "iter_dt 13.89ms; iter 600: train loss 0.01612\n",
      "iter_dt 13.68ms; iter 700: train loss 0.03921\n",
      "iter_dt 13.82ms; iter 800: train loss 0.00338\n",
      "iter_dt 14.21ms; iter 900: train loss 0.01256\n",
      "iter_dt 14.15ms; iter 1000: train loss 0.01387\n",
      "iter_dt 14.06ms; iter 1100: train loss 0.00115\n",
      "iter_dt 14.27ms; iter 1200: train loss 0.00537\n",
      "iter_dt 18.30ms; iter 1300: train loss 0.00040\n",
      "iter_dt 14.06ms; iter 1400: train loss 0.03832\n",
      "iter_dt 13.67ms; iter 1500: train loss 0.00626\n",
      "iter_dt 13.93ms; iter 1600: train loss 0.02706\n",
      "iter_dt 14.04ms; iter 1700: train loss 0.00188\n",
      "iter_dt 13.99ms; iter 1800: train loss 0.00099\n",
      "iter_dt 13.93ms; iter 1900: train loss 0.01076\n"
     ]
    }
   ],
   "source": [
    "def batch_end_callback(trainer):\n",
    "    if trainer.iter_num % 100 == 0:\n",
    "        print(f\"iter_dt {trainer.iter_dt * 1000:.2f}ms; iter {trainer.iter_num}: train loss {trainer.loss.item():.5f}\")\n",
    "trainer.set_callback('on_batch_end', batch_end_callback)\n",
    "\n",
    "trainer.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now let's perform some evaluation\n",
    "model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train final score: 5000/5000 = 100.00% correct\n",
      "test final score: 5000/5000 = 100.00% correct\n"
     ]
    }
   ],
   "source": [
    "def eval_split(trainer, split, max_batches):\n",
    "    dataset = {'train':train_dataset, 'test':test_dataset}[split]\n",
    "    n = train_dataset.length # naugy direct access shrug\n",
    "    results = []\n",
    "    mistakes_printed_already = 0\n",
    "    loader = DataLoader(dataset, batch_size=100, num_workers=0, drop_last=False)\n",
    "    for b, (x, y) in enumerate(loader):\n",
    "        x = x.to(trainer.device)\n",
    "        y = y.to(trainer.device)\n",
    "        # isolate the input pattern alone\n",
    "        inp = x[:, :n]\n",
    "        sol = y[:, -n:]\n",
    "        # let the model sample the rest of the sequence\n",
    "        cat = model.generate(inp, n, do_sample=False) # using greedy argmax, not sampling\n",
    "        sol_candidate = cat[:, n:] # isolate the filled in sequence\n",
    "        # compare the predicted sequence to the true sequence\n",
    "        correct = (sol == sol_candidate).all(1).cpu() # Software 1.0 vs. Software 2.0 fight RIGHT on this line haha\n",
    "        for i in range(x.size(0)):\n",
    "            results.append(int(correct[i]))\n",
    "            if not correct[i] and mistakes_printed_already < 3: # only print up to 5 mistakes to get a sense\n",
    "                mistakes_printed_already += 1\n",
    "                print(\"GPT claims that %s sorted is %s but gt is %s\" % (inp[i].tolist(), sol_candidate[i].tolist(), sol[i].tolist()))\n",
    "        if max_batches is not None and b+1 >= max_batches:\n",
    "            break\n",
    "    rt = torch.tensor(results, dtype=torch.float)\n",
    "    print(\"%s final score: %d/%d = %.2f%% correct\" % (split, rt.sum(), len(results), 100*rt.mean()))\n",
    "    return rt.sum()\n",
    "\n",
    "# run a lot of examples from both train and test through the model and verify the output correctness\n",
    "with torch.no_grad():\n",
    "    train_score = eval_split(trainer, 'train', max_batches=50)\n",
    "    test_score  = eval_split(trainer, 'test',  max_batches=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input sequence  : [[0, 0, 2, 1, 0, 1]]\n",
      "predicted sorted: [[0, 0, 0, 1, 1, 2]]\n",
      "gt sort         : [0, 0, 0, 1, 1, 2]\n",
      "matches         : True\n"
     ]
    }
   ],
   "source": [
    "# let's run a random given sequence through the model as well\n",
    "n = train_dataset.length # naugy direct access shrug\n",
    "inp = torch.tensor([[0, 0, 2, 1, 0, 1]], dtype=torch.long).to(trainer.device)\n",
    "assert inp[0].nelement() == n\n",
    "with torch.no_grad():\n",
    "    cat = model.generate(inp, n, do_sample=False)\n",
    "sol = torch.sort(inp[0])[0]\n",
    "sol_candidate = cat[:, n:]\n",
    "print('input sequence  :', inp.tolist())\n",
    "print('predicted sorted:', sol_candidate.tolist())\n",
    "print('gt sort         :', sol.tolist())\n",
    "print('matches         :', bool((sol == sol_candidate).all()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[ -12552.3867, -116095.0234, -298504.4375,  ...,   -8466.6279,\n",
       "            -36166.3828,  -30076.0234],\n",
       "          [ -12553.0264, -116094.8281, -298504.7812,  ...,   -8466.7178,\n",
       "            -36164.1602,  -30074.7637],\n",
       "          [ -12553.8408, -116095.3438, -298504.5625,  ...,   -8466.0527,\n",
       "            -36166.1211,  -30074.7266],\n",
       "          ...,\n",
       "          [ -12554.3965, -116094.1172, -298504.4062,  ...,   -8464.6836,\n",
       "            -36164.7070,  -30074.0488],\n",
       "          [ -12553.8652, -116094.6250, -298504.7812,  ...,   -8467.8955,\n",
       "            -36165.2383,  -30075.2832],\n",
       "          [ -12554.0273, -116095.6016, -298504.6875,  ...,   -8465.4580,\n",
       "            -36165.7773,  -30075.0352]],\n",
       "\n",
       "         [[ -12554.1064, -116096.2734, -298505.6562,  ...,   -8467.5830,\n",
       "            -36163.6797,  -30075.2383],\n",
       "          [ -12555.8779, -116093.5781, -298503.8125,  ...,   -8467.3242,\n",
       "            -36165.5820,  -30074.4277],\n",
       "          [ -12552.3613, -116095.5547, -298504.2812,  ...,   -8464.6445,\n",
       "            -36163.6680,  -30075.6191],\n",
       "          ...,\n",
       "          [ -12552.1787, -116094.4453, -298505.6250,  ...,   -8465.5439,\n",
       "            -36164.8828,  -30074.2266],\n",
       "          [ -12555.9912, -116094.2031, -298504.5000,  ...,   -8466.9658,\n",
       "            -36163.3398,  -30074.3945],\n",
       "          [ -12554.9746, -116093.7422, -298505.2500,  ...,   -8465.8662,\n",
       "            -36165.6562,  -30073.1406]],\n",
       "\n",
       "         [[ -12552.6006, -116093.0391, -298505.4062,  ...,   -8466.1904,\n",
       "            -36164.8164,  -30073.0488],\n",
       "          [ -12555.1826, -116096.5000, -298505.4062,  ...,   -8466.9043,\n",
       "            -36164.1133,  -30074.2402],\n",
       "          [ -12553.4766, -116095.7344, -298505.5000,  ...,   -8467.6553,\n",
       "            -36164.5703,  -30072.4570],\n",
       "          ...,\n",
       "          [ -12553.8994, -116094.0938, -298505.8125,  ...,   -8465.9775,\n",
       "            -36166.9453,  -30074.0137],\n",
       "          [ -12552.8848, -116095.7188, -298505.6875,  ...,   -8465.8271,\n",
       "            -36166.0039,  -30075.3320],\n",
       "          [ -12553.8232, -116095.3203, -298505.4688,  ...,   -8464.9004,\n",
       "            -36165.5039,  -30074.2402]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ -12553.7529, -116096.6719, -298504.5000,  ...,   -8467.5195,\n",
       "            -36165.2734,  -30074.7656],\n",
       "          [ -12555.7275, -116095.0703, -298505.1875,  ...,   -8467.3125,\n",
       "            -36163.0586,  -30072.6152],\n",
       "          [ -12553.8896, -116095.1797, -298505.6250,  ...,   -8468.0762,\n",
       "            -36164.7500,  -30074.7695],\n",
       "          ...,\n",
       "          [ -12552.8242, -116095.1797, -298505.5312,  ...,   -8467.3447,\n",
       "            -36164.4102,  -30074.6523],\n",
       "          [ -12554.1797, -116094.7734, -298504.7812,  ...,   -8466.1221,\n",
       "            -36165.4883,  -30073.9824],\n",
       "          [ -12553.2939, -116094.7422, -298505.5625,  ...,   -8466.2871,\n",
       "            -36166.6094,  -30073.4473]],\n",
       "\n",
       "         [[ -12554.4883, -116094.4766, -298502.2812,  ...,   -8465.8262,\n",
       "            -36165.3945,  -30074.4727],\n",
       "          [ -12553.5361, -116093.0078, -298503.9688,  ...,   -8466.5996,\n",
       "            -36164.6914,  -30075.0039],\n",
       "          [ -12552.7568, -116094.7656, -298503.7812,  ...,   -8464.4580,\n",
       "            -36165.1484,  -30074.7676],\n",
       "          ...,\n",
       "          [ -12552.0830, -116095.6953, -298504.1250,  ...,   -8464.5713,\n",
       "            -36165.5664,  -30074.3320],\n",
       "          [ -12553.8975, -116094.5859, -298504.4688,  ...,   -8465.6777,\n",
       "            -36166.9727,  -30075.1973],\n",
       "          [ -12555.2197, -116093.9375, -298505.4688,  ...,   -8466.9512,\n",
       "            -36164.5156,  -30074.4199]],\n",
       "\n",
       "         [[ -12553.5537, -116094.6016, -298503.5625,  ...,   -8465.7637,\n",
       "            -36165.2500,  -30075.5840],\n",
       "          [ -12552.6094, -116095.3906, -298504.3750,  ...,   -8467.5859,\n",
       "            -36164.3984,  -30074.5352],\n",
       "          [ -12554.4785, -116097.4766, -298504.2812,  ...,   -8467.5049,\n",
       "            -36165.1797,  -30075.1992],\n",
       "          ...,\n",
       "          [ -12554.1875, -116092.9141, -298506.5000,  ...,   -8466.5498,\n",
       "            -36165.3672,  -30075.5391],\n",
       "          [ -12553.9434, -116094.8203, -298504.5312,  ...,   -8466.8965,\n",
       "            -36163.9688,  -30073.9180],\n",
       "          [ -12555.6904, -116093.9766, -298505.6250,  ...,   -8466.0908,\n",
       "            -36166.0664,  -30074.5391]]],\n",
       "\n",
       "\n",
       "        [[[-207005.2969,   84380.6016,  -65319.9531,  ...,  -51888.5156,\n",
       "            -58986.2148, -183961.8750],\n",
       "          [-207005.2188,   84379.5781,  -65318.9141,  ...,  -51888.2773,\n",
       "            -58985.8906, -183960.8125],\n",
       "          [-207005.6719,   84379.8281,  -65320.3867,  ...,  -51889.8047,\n",
       "            -58984.6406, -183960.1719],\n",
       "          ...,\n",
       "          [-207006.2656,   84381.1484,  -65322.1914,  ...,  -51887.0703,\n",
       "            -58986.6445, -183960.5938],\n",
       "          [-207005.8906,   84378.9297,  -65318.7227,  ...,  -51888.9414,\n",
       "            -58984.8789, -183962.0156],\n",
       "          [-207005.4219,   84379.6484,  -65322.0117,  ...,  -51887.6133,\n",
       "            -58984.0391, -183961.9375]],\n",
       "\n",
       "         [[-207005.5000,   84381.2891,  -65321.1328,  ...,  -51888.5156,\n",
       "            -58983.9961, -183961.1406],\n",
       "          [-207005.4062,   84380.5312,  -65319.9609,  ...,  -51888.4414,\n",
       "            -58985.1836, -183963.3281],\n",
       "          [-207007.4844,   84379.1172,  -65319.5078,  ...,  -51887.5000,\n",
       "            -58986.0234, -183962.2500],\n",
       "          ...,\n",
       "          [-207005.5469,   84380.5000,  -65321.1484,  ...,  -51886.0430,\n",
       "            -58983.9805, -183960.9531],\n",
       "          [-207006.9688,   84379.1484,  -65318.6797,  ...,  -51890.0039,\n",
       "            -58985.8828, -183959.1562],\n",
       "          [-207006.7812,   84378.0391,  -65320.3555,  ...,  -51889.0039,\n",
       "            -58985.3672, -183961.3750]],\n",
       "\n",
       "         [[-207005.5781,   84378.6719,  -65321.3164,  ...,  -51886.9688,\n",
       "            -58985.4961, -183961.6250],\n",
       "          [-207007.2188,   84378.4531,  -65320.6367,  ...,  -51888.7422,\n",
       "            -58985.1992, -183960.3750],\n",
       "          [-207004.5156,   84379.2344,  -65320.7461,  ...,  -51888.9219,\n",
       "            -58986.9844, -183961.9531],\n",
       "          ...,\n",
       "          [-207005.0000,   84378.3281,  -65320.0898,  ...,  -51886.7422,\n",
       "            -58985.6953, -183959.9062],\n",
       "          [-207004.5938,   84379.4531,  -65319.3516,  ...,  -51889.4023,\n",
       "            -58986.3750, -183961.8438],\n",
       "          [-207007.5469,   84380.4141,  -65319.7891,  ...,  -51887.2812,\n",
       "            -58985.5820, -183963.1875]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-207004.8594,   84380.5391,  -65320.6250,  ...,  -51885.9688,\n",
       "            -58984.5547, -183960.1875],\n",
       "          [-207006.4062,   84377.9219,  -65322.0312,  ...,  -51889.1914,\n",
       "            -58985.5234, -183961.3438],\n",
       "          [-207004.7188,   84378.9844,  -65321.1406,  ...,  -51887.6758,\n",
       "            -58984.2578, -183961.5469],\n",
       "          ...,\n",
       "          [-207006.5312,   84378.6484,  -65321.2734,  ...,  -51888.3711,\n",
       "            -58983.7930, -183962.1719],\n",
       "          [-207006.9844,   84379.6172,  -65318.6680,  ...,  -51889.7812,\n",
       "            -58983.5234, -183962.8594],\n",
       "          [-207005.9531,   84378.6953,  -65319.2070,  ...,  -51888.6289,\n",
       "            -58984.0312, -183961.1094]],\n",
       "\n",
       "         [[-207007.2500,   84379.9062,  -65320.9453,  ...,  -51888.3633,\n",
       "            -58984.6094, -183961.8906],\n",
       "          [-207007.5781,   84381.7188,  -65319.8984,  ...,  -51887.4023,\n",
       "            -58983.4883, -183961.1094],\n",
       "          [-207005.4219,   84381.5391,  -65320.8555,  ...,  -51887.6758,\n",
       "            -58985.3281, -183960.2031],\n",
       "          ...,\n",
       "          [-207007.2969,   84378.8281,  -65321.3945,  ...,  -51887.6836,\n",
       "            -58983.2500, -183962.7656],\n",
       "          [-207006.7344,   84381.5078,  -65318.2500,  ...,  -51890.3320,\n",
       "            -58986.1250, -183960.9062],\n",
       "          [-207003.3750,   84379.0547,  -65321.1562,  ...,  -51887.8789,\n",
       "            -58984.9219, -183962.6406]],\n",
       "\n",
       "         [[-207007.9688,   84379.0312,  -65320.2148,  ...,  -51887.2969,\n",
       "            -58985.2422, -183960.0312],\n",
       "          [-207003.7969,   84379.8750,  -65320.2773,  ...,  -51888.0352,\n",
       "            -58985.1367, -183961.7344],\n",
       "          [-207007.1250,   84379.8594,  -65321.2461,  ...,  -51888.0352,\n",
       "            -58984.7930, -183962.3125],\n",
       "          ...,\n",
       "          [-207007.6562,   84378.2109,  -65320.2266,  ...,  -51887.1484,\n",
       "            -58986.7656, -183961.0312],\n",
       "          [-207006.5781,   84380.1953,  -65320.5078,  ...,  -51887.7930,\n",
       "            -58983.1445, -183961.4219],\n",
       "          [-207006.0312,   84380.9141,  -65319.5547,  ...,  -51885.7383,\n",
       "            -58984.9102, -183961.3750]]],\n",
       "\n",
       "\n",
       "        [[[-238457.2188,  -47106.9531, -320729.9062,  ..., -118620.7969,\n",
       "            -97102.6797, -146419.0312],\n",
       "          [-238456.7969,  -47107.0742, -320731.0312,  ..., -118621.8672,\n",
       "            -97102.4766, -146419.0469],\n",
       "          [-238457.8281,  -47108.4258, -320731.5938,  ..., -118622.6797,\n",
       "            -97100.9062, -146418.3438],\n",
       "          ...,\n",
       "          [-238458.1406,  -47109.7930, -320729.9375,  ..., -118622.7891,\n",
       "            -97101.7344, -146417.7031],\n",
       "          [-238455.2031,  -47109.5664, -320731.5000,  ..., -118620.6250,\n",
       "            -97102.9844, -146419.2188],\n",
       "          [-238456.7969,  -47108.6836, -320730.5312,  ..., -118622.8281,\n",
       "            -97100.7578, -146419.0938]],\n",
       "\n",
       "         [[-238458.5312,  -47108.4375, -320731.0000,  ..., -118622.0078,\n",
       "            -97101.0469, -146419.9844],\n",
       "          [-238455.5625,  -47106.3750, -320730.3125,  ..., -118622.0391,\n",
       "            -97103.5391, -146420.7969],\n",
       "          [-238456.2344,  -47108.4883, -320732.7188,  ..., -118621.8516,\n",
       "            -97102.8203, -146418.5000],\n",
       "          ...,\n",
       "          [-238458.4062,  -47109.2148, -320732.1250,  ..., -118621.1953,\n",
       "            -97100.8438, -146420.4062],\n",
       "          [-238456.2969,  -47109.3828, -320729.7812,  ..., -118620.9844,\n",
       "            -97100.8828, -146418.3906],\n",
       "          [-238456.0000,  -47108.1641, -320730.8438,  ..., -118622.7812,\n",
       "            -97102.2578, -146420.4688]],\n",
       "\n",
       "         [[-238456.6562,  -47107.6016, -320731.3125,  ..., -118622.0781,\n",
       "            -97102.1328, -146419.6719],\n",
       "          [-238456.2969,  -47107.6445, -320731.8750,  ..., -118621.7422,\n",
       "            -97101.9844, -146418.8594],\n",
       "          [-238457.6250,  -47107.6445, -320730.3750,  ..., -118620.8594,\n",
       "            -97102.3359, -146419.0938],\n",
       "          ...,\n",
       "          [-238457.7656,  -47109.1953, -320731.0625,  ..., -118624.0156,\n",
       "            -97100.3203, -146419.3438],\n",
       "          [-238456.4844,  -47105.2109, -320729.7812,  ..., -118621.4922,\n",
       "            -97101.2266, -146420.1875],\n",
       "          [-238457.3125,  -47108.9688, -320730.2812,  ..., -118621.0859,\n",
       "            -97102.1484, -146420.5156]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-238456.6719,  -47110.6641, -320729.2812,  ..., -118621.0625,\n",
       "            -97101.1172, -146418.8125],\n",
       "          [-238457.0156,  -47109.0820, -320730.6250,  ..., -118621.3828,\n",
       "            -97102.6875, -146418.4688],\n",
       "          [-238459.0156,  -47107.7617, -320729.8750,  ..., -118623.7422,\n",
       "            -97101.8125, -146418.8438],\n",
       "          ...,\n",
       "          [-238457.4531,  -47108.9883, -320728.5938,  ..., -118621.7109,\n",
       "            -97101.4531, -146420.7656],\n",
       "          [-238455.1250,  -47107.4531, -320731.1875,  ..., -118621.1797,\n",
       "            -97103.0781, -146419.9844],\n",
       "          [-238456.0781,  -47108.5898, -320731.7500,  ..., -118621.7266,\n",
       "            -97100.3906, -146420.5312]],\n",
       "\n",
       "         [[-238456.3125,  -47109.8711, -320729.4062,  ..., -118619.5078,\n",
       "            -97100.9922, -146419.3906],\n",
       "          [-238456.8125,  -47107.3320, -320731.4062,  ..., -118622.0859,\n",
       "            -97101.9453, -146418.6562],\n",
       "          [-238456.5938,  -47108.9180, -320730.6875,  ..., -118622.2812,\n",
       "            -97103.9062, -146419.1250],\n",
       "          ...,\n",
       "          [-238457.1719,  -47108.8203, -320731.5312,  ..., -118621.7031,\n",
       "            -97103.2578, -146418.2969],\n",
       "          [-238456.7188,  -47109.2656, -320729.5938,  ..., -118622.9922,\n",
       "            -97102.1797, -146418.3125],\n",
       "          [-238456.2188,  -47108.6289, -320730.0000,  ..., -118622.1562,\n",
       "            -97100.5547, -146419.2344]],\n",
       "\n",
       "         [[-238457.8594,  -47109.7188, -320730.0000,  ..., -118622.2031,\n",
       "            -97102.4297, -146418.8594],\n",
       "          [-238457.2969,  -47109.9492, -320731.3750,  ..., -118623.5234,\n",
       "            -97100.5078, -146417.8906],\n",
       "          [-238457.2500,  -47108.6602, -320732.5000,  ..., -118623.4531,\n",
       "            -97101.4219, -146418.7969],\n",
       "          ...,\n",
       "          [-238455.9375,  -47109.8242, -320729.6875,  ..., -118622.3594,\n",
       "            -97101.5391, -146419.5156],\n",
       "          [-238457.6562,  -47108.9688, -320730.7188,  ..., -118622.8047,\n",
       "            -97102.1172, -146418.5156],\n",
       "          [-238456.0000,  -47110.3594, -320729.7188,  ..., -118621.9375,\n",
       "            -97100.5312, -146419.0938]]],\n",
       "\n",
       "\n",
       "        ...,\n",
       "\n",
       "\n",
       "        [[[-175370.3594, -162652.8438,  185812.4844,  ...,   -7561.1968,\n",
       "           -247513.8125,   87173.2031],\n",
       "          [-175371.5469, -162651.7500,  185814.3125,  ...,   -7559.2798,\n",
       "           -247514.7500,   87176.0703],\n",
       "          [-175369.0469, -162651.5625,  185814.8594,  ...,   -7561.2861,\n",
       "           -247513.0156,   87174.9141],\n",
       "          ...,\n",
       "          [-175369.0781, -162652.1250,  185815.4531,  ...,   -7560.4531,\n",
       "           -247513.4062,   87176.5625],\n",
       "          [-175369.6406, -162652.8281,  185815.0938,  ...,   -7560.6367,\n",
       "           -247512.7031,   87173.4375],\n",
       "          [-175369.7031, -162651.8125,  185814.4688,  ...,   -7560.6392,\n",
       "           -247514.0000,   87176.2031]],\n",
       "\n",
       "         [[-175370.0625, -162653.2812,  185815.4375,  ...,   -7559.9731,\n",
       "           -247511.2344,   87174.6875],\n",
       "          [-175369.4219, -162652.6250,  185813.6094,  ...,   -7561.8330,\n",
       "           -247513.4688,   87175.5859],\n",
       "          [-175370.5312, -162651.5312,  185814.4375,  ...,   -7560.1777,\n",
       "           -247512.3750,   87174.3125],\n",
       "          ...,\n",
       "          [-175368.3906, -162654.3438,  185815.9688,  ...,   -7560.8037,\n",
       "           -247513.4844,   87175.2500],\n",
       "          [-175370.8281, -162649.9688,  185815.6719,  ...,   -7560.4849,\n",
       "           -247513.3438,   87175.2578],\n",
       "          [-175369.9062, -162651.2188,  185816.0938,  ...,   -7561.9004,\n",
       "           -247513.3594,   87174.9844]],\n",
       "\n",
       "         [[-175370.7969, -162652.0469,  185815.5625,  ...,   -7562.2637,\n",
       "           -247510.4219,   87174.6562],\n",
       "          [-175368.5156, -162652.7344,  185814.2188,  ...,   -7560.1353,\n",
       "           -247512.8438,   87173.9844],\n",
       "          [-175369.9531, -162652.9219,  185814.7812,  ...,   -7558.6069,\n",
       "           -247514.6562,   87176.8672],\n",
       "          ...,\n",
       "          [-175367.4531, -162651.5469,  185815.2344,  ...,   -7559.6255,\n",
       "           -247513.1250,   87175.9688],\n",
       "          [-175369.5625, -162654.2344,  185814.7969,  ...,   -7562.3696,\n",
       "           -247513.3750,   87173.7891],\n",
       "          [-175369.5469, -162652.1719,  185814.9531,  ...,   -7559.5557,\n",
       "           -247512.3438,   87174.1328]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-175370.2812, -162653.8125,  185814.4688,  ...,   -7559.2080,\n",
       "           -247512.7188,   87175.0312],\n",
       "          [-175369.7188, -162651.9062,  185815.8594,  ...,   -7560.1519,\n",
       "           -247511.7344,   87174.6953],\n",
       "          [-175370.1562, -162653.3281,  185814.2188,  ...,   -7560.4219,\n",
       "           -247512.8594,   87173.7109],\n",
       "          ...,\n",
       "          [-175369.6719, -162652.5938,  185815.2344,  ...,   -7561.6572,\n",
       "           -247511.9844,   87174.9531],\n",
       "          [-175369.6719, -162652.7812,  185813.3281,  ...,   -7561.5044,\n",
       "           -247513.5469,   87174.4688],\n",
       "          [-175370.1875, -162651.5000,  185813.9844,  ...,   -7561.7310,\n",
       "           -247514.4844,   87177.2344]],\n",
       "\n",
       "         [[-175370.4375, -162653.7969,  185814.6562,  ...,   -7560.9585,\n",
       "           -247512.4062,   87175.1328],\n",
       "          [-175371.0156, -162654.7188,  185813.2656,  ...,   -7562.3804,\n",
       "           -247514.4688,   87175.3594],\n",
       "          [-175368.6094, -162652.0469,  185814.5625,  ...,   -7560.6040,\n",
       "           -247513.7812,   87175.4453],\n",
       "          ...,\n",
       "          [-175367.8594, -162653.1406,  185814.4375,  ...,   -7560.1367,\n",
       "           -247513.8750,   87173.8594],\n",
       "          [-175370.3906, -162653.0781,  185814.4531,  ...,   -7560.8047,\n",
       "           -247514.3125,   87173.1406],\n",
       "          [-175368.7656, -162653.6094,  185812.5312,  ...,   -7560.9404,\n",
       "           -247511.1094,   87175.0078]],\n",
       "\n",
       "         [[-175369.4531, -162653.8594,  185812.6875,  ...,   -7562.9180,\n",
       "           -247512.8438,   87176.5938],\n",
       "          [-175370.3125, -162651.6562,  185813.0312,  ...,   -7560.6045,\n",
       "           -247513.4688,   87175.1250],\n",
       "          [-175369.3125, -162651.5156,  185814.6875,  ...,   -7559.9204,\n",
       "           -247514.0781,   87176.3594],\n",
       "          ...,\n",
       "          [-175370.5000, -162651.4219,  185814.6406,  ...,   -7559.3979,\n",
       "           -247512.3438,   87173.5703],\n",
       "          [-175370.6719, -162652.1562,  185813.4688,  ...,   -7559.4844,\n",
       "           -247514.4219,   87174.0156],\n",
       "          [-175370.1719, -162652.6406,  185815.6094,  ...,   -7561.0811,\n",
       "           -247513.0156,   87174.0312]]],\n",
       "\n",
       "\n",
       "        [[[-156329.6875,   12027.3643,  -66127.4766,  ...,  -25045.9941,\n",
       "           -191375.6875, -138114.6719],\n",
       "          [-156329.9219,   12027.9307,  -66128.2109,  ...,  -25044.5840,\n",
       "           -191375.5781, -138113.8438],\n",
       "          [-156330.6406,   12027.2695,  -66128.7031,  ...,  -25045.9727,\n",
       "           -191374.9844, -138114.4531],\n",
       "          ...,\n",
       "          [-156331.4844,   12027.3809,  -66127.1484,  ...,  -25046.1816,\n",
       "           -191373.8281, -138114.9062],\n",
       "          [-156330.4531,   12028.7275,  -66129.4219,  ...,  -25045.8242,\n",
       "           -191375.5938, -138114.9375],\n",
       "          [-156331.5156,   12026.7803,  -66126.6562,  ...,  -25047.2930,\n",
       "           -191376.4062, -138115.0781]],\n",
       "\n",
       "         [[-156330.8281,   12026.4043,  -66129.8125,  ...,  -25044.7012,\n",
       "           -191375.5938, -138114.7969],\n",
       "          [-156331.0000,   12028.8828,  -66128.9297,  ...,  -25045.7793,\n",
       "           -191374.3438, -138115.6406],\n",
       "          [-156331.4531,   12027.8906,  -66128.6250,  ...,  -25045.1602,\n",
       "           -191373.9531, -138117.2344],\n",
       "          ...,\n",
       "          [-156330.8125,   12025.8682,  -66128.0000,  ...,  -25048.0469,\n",
       "           -191378.6875, -138114.6250],\n",
       "          [-156329.3906,   12025.8008,  -66128.0000,  ...,  -25045.4258,\n",
       "           -191374.1562, -138114.0156],\n",
       "          [-156328.9688,   12026.8730,  -66129.6562,  ...,  -25046.4902,\n",
       "           -191374.7344, -138115.8750]],\n",
       "\n",
       "         [[-156331.0781,   12026.9023,  -66130.2969,  ...,  -25047.3691,\n",
       "           -191376.1875, -138114.2500],\n",
       "          [-156331.3438,   12026.5674,  -66128.7578,  ...,  -25046.7773,\n",
       "           -191376.7812, -138115.1562],\n",
       "          [-156330.9375,   12026.6807,  -66128.6484,  ...,  -25045.3711,\n",
       "           -191374.3125, -138114.9062],\n",
       "          ...,\n",
       "          [-156330.8125,   12026.5850,  -66127.8672,  ...,  -25046.4102,\n",
       "           -191374.4219, -138117.2656],\n",
       "          [-156331.4062,   12027.2627,  -66130.4141,  ...,  -25045.8496,\n",
       "           -191373.2812, -138115.1094],\n",
       "          [-156331.1875,   12027.6611,  -66128.5547,  ...,  -25045.8809,\n",
       "           -191375.2656, -138116.0781]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-156331.4219,   12026.8154,  -66129.6484,  ...,  -25045.7559,\n",
       "           -191374.2031, -138115.9688],\n",
       "          [-156330.8906,   12025.4180,  -66128.7969,  ...,  -25045.3828,\n",
       "           -191375.1719, -138115.0156],\n",
       "          [-156330.0938,   12028.2246,  -66128.1172,  ...,  -25047.7539,\n",
       "           -191373.5312, -138114.7188],\n",
       "          ...,\n",
       "          [-156329.4531,   12027.0703,  -66129.2969,  ...,  -25045.4902,\n",
       "           -191375.5469, -138115.4688],\n",
       "          [-156330.7812,   12027.0195,  -66128.6953,  ...,  -25045.3945,\n",
       "           -191373.9219, -138114.4219],\n",
       "          [-156331.0781,   12028.7686,  -66128.1172,  ...,  -25047.0195,\n",
       "           -191374.2500, -138116.2500]],\n",
       "\n",
       "         [[-156331.0938,   12025.6680,  -66128.4844,  ...,  -25046.5977,\n",
       "           -191376.1562, -138117.3125],\n",
       "          [-156333.8906,   12025.0986,  -66129.2266,  ...,  -25046.8555,\n",
       "           -191374.7500, -138117.4688],\n",
       "          [-156330.4375,   12026.5723,  -66129.6172,  ...,  -25046.2188,\n",
       "           -191376.3750, -138115.2344],\n",
       "          ...,\n",
       "          [-156331.8125,   12025.0264,  -66130.2344,  ...,  -25044.9199,\n",
       "           -191375.0312, -138116.7812],\n",
       "          [-156330.5938,   12025.9580,  -66130.1484,  ...,  -25046.1973,\n",
       "           -191374.6562, -138116.0469],\n",
       "          [-156329.7031,   12026.2441,  -66128.0469,  ...,  -25046.0723,\n",
       "           -191374.7656, -138114.8281]],\n",
       "\n",
       "         [[-156330.9531,   12024.7998,  -66127.5625,  ...,  -25045.9902,\n",
       "           -191374.5312, -138114.5312],\n",
       "          [-156332.1250,   12024.8682,  -66129.0000,  ...,  -25045.2031,\n",
       "           -191375.2500, -138116.4688],\n",
       "          [-156329.7969,   12027.4414,  -66128.9844,  ...,  -25045.0488,\n",
       "           -191376.4219, -138115.7969],\n",
       "          ...,\n",
       "          [-156331.0469,   12026.3867,  -66128.7031,  ...,  -25045.1230,\n",
       "           -191376.5156, -138116.3125],\n",
       "          [-156332.7969,   12025.2441,  -66128.7812,  ...,  -25045.1641,\n",
       "           -191375.8281, -138116.0312],\n",
       "          [-156331.3906,   12027.1406,  -66128.4453,  ...,  -25046.1836,\n",
       "           -191374.0000, -138117.0156]]],\n",
       "\n",
       "\n",
       "        [[[-201364.2500, -146707.9375, -292038.9688,  ..., -272524.3438,\n",
       "            -68409.6797, -103326.9844],\n",
       "          [-201363.0000, -146709.5781, -292038.7812,  ..., -272524.9375,\n",
       "            -68409.8047, -103325.1484],\n",
       "          [-201362.7500, -146707.8438, -292038.5938,  ..., -272524.5938,\n",
       "            -68410.1172, -103327.9688],\n",
       "          ...,\n",
       "          [-201362.1562, -146708.6875, -292039.0000,  ..., -272524.2500,\n",
       "            -68411.7422, -103326.8828],\n",
       "          [-201362.9531, -146709.2500, -292039.1250,  ..., -272524.3125,\n",
       "            -68410.5625, -103326.9922],\n",
       "          [-201362.2969, -146708.6875, -292038.7500,  ..., -272525.3125,\n",
       "            -68410.6406, -103325.6719]],\n",
       "\n",
       "         [[-201362.5938, -146707.8750, -292039.2188,  ..., -272523.4688,\n",
       "            -68411.0703, -103326.0547],\n",
       "          [-201362.3281, -146708.1562, -292038.7188,  ..., -272524.4062,\n",
       "            -68409.5391, -103327.7266],\n",
       "          [-201359.4844, -146708.6562, -292037.6875,  ..., -272524.9375,\n",
       "            -68409.5703, -103325.4141],\n",
       "          ...,\n",
       "          [-201363.2500, -146706.7188, -292039.7188,  ..., -272525.5000,\n",
       "            -68410.1094, -103327.0156],\n",
       "          [-201361.8750, -146707.3438, -292039.0000,  ..., -272526.0625,\n",
       "            -68411.6016, -103327.0625],\n",
       "          [-201361.1250, -146707.5938, -292039.0938,  ..., -272527.0312,\n",
       "            -68409.5938, -103327.7969]],\n",
       "\n",
       "         [[-201362.8281, -146707.9219, -292038.5625,  ..., -272525.4375,\n",
       "            -68409.9062, -103326.9297],\n",
       "          [-201360.5625, -146709.2656, -292037.6250,  ..., -272525.3125,\n",
       "            -68412.0078, -103326.9609],\n",
       "          [-201361.6562, -146709.2969, -292039.0625,  ..., -272526.2500,\n",
       "            -68410.7812, -103326.1562],\n",
       "          ...,\n",
       "          [-201361.9375, -146707.7031, -292038.3125,  ..., -272525.6562,\n",
       "            -68410.0391, -103326.9375],\n",
       "          [-201362.0938, -146708.3438, -292038.6875,  ..., -272525.5625,\n",
       "            -68410.3125, -103324.5391],\n",
       "          [-201361.9688, -146710.9844, -292039.8438,  ..., -272526.7188,\n",
       "            -68411.5234, -103326.1094]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[-201361.4062, -146709.7656, -292039.4062,  ..., -272524.9375,\n",
       "            -68410.4375, -103326.5547],\n",
       "          [-201361.2500, -146710.7188, -292039.4062,  ..., -272524.5938,\n",
       "            -68408.8281, -103327.3203],\n",
       "          [-201360.5625, -146707.9062, -292038.4062,  ..., -272525.4375,\n",
       "            -68410.5703, -103326.5000],\n",
       "          ...,\n",
       "          [-201363.0469, -146708.0312, -292038.7500,  ..., -272524.5000,\n",
       "            -68409.3516, -103326.2812],\n",
       "          [-201360.6094, -146709.9062, -292038.5000,  ..., -272526.1875,\n",
       "            -68410.9922, -103328.3281],\n",
       "          [-201361.6719, -146707.8594, -292038.0938,  ..., -272525.3438,\n",
       "            -68411.0312, -103328.4062]],\n",
       "\n",
       "         [[-201361.2031, -146709.7656, -292037.6875,  ..., -272524.5312,\n",
       "            -68408.6641, -103326.4531],\n",
       "          [-201361.3125, -146706.3906, -292038.0938,  ..., -272525.5000,\n",
       "            -68410.4922, -103326.5000],\n",
       "          [-201361.9531, -146708.8750, -292037.9375,  ..., -272525.1875,\n",
       "            -68409.8828, -103327.4766],\n",
       "          ...,\n",
       "          [-201362.7188, -146707.8281, -292038.7188,  ..., -272525.5312,\n",
       "            -68410.3516, -103327.1484],\n",
       "          [-201361.7812, -146707.9531, -292039.0938,  ..., -272527.1562,\n",
       "            -68411.5234, -103327.6875],\n",
       "          [-201362.0781, -146707.4219, -292038.8438,  ..., -272525.3438,\n",
       "            -68410.6328, -103326.8672]],\n",
       "\n",
       "         [[-201362.0156, -146708.1094, -292037.1875,  ..., -272524.8438,\n",
       "            -68411.6719, -103324.6719],\n",
       "          [-201360.9844, -146706.9844, -292040.3750,  ..., -272527.0625,\n",
       "            -68409.1094, -103327.0078],\n",
       "          [-201361.0156, -146709.9531, -292040.6562,  ..., -272526.2188,\n",
       "            -68411.0234, -103326.0938],\n",
       "          ...,\n",
       "          [-201362.2344, -146708.2969, -292038.6250,  ..., -272524.9062,\n",
       "            -68411.4766, -103327.0781],\n",
       "          [-201362.2656, -146709.0312, -292037.7188,  ..., -272524.2812,\n",
       "            -68410.3203, -103326.3281],\n",
       "          [-201361.8750, -146708.4219, -292040.1250,  ..., -272524.1562,\n",
       "            -68411.5391, -103326.4141]]]])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B, T, C = 64, 8, 128\n",
    "nh = 16\n",
    "\n",
    "att = torch.randn((B, nh, T, T))\n",
    "mask = torch.randn((B, T))\n",
    "mask = mask[:, None, None, :]\n",
    "att -= 100000 * (1-mask)\n",
    "att"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-hyper-trees]",
   "language": "python",
   "name": "conda-env-.conda-hyper-trees-py"
  },
  "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.9.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "3ad933181bd8a04b432d3370b9dc3b0662ad032c4dfaa4e4f1596c548f763858"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}