bracket-closing.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Understanding bracket closing in GPT-Neo\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/SamAdamDay/mechanistic-interpretability-projects/blob/main/bracket-closing.ipynb)\n",
    "\n",
    "The goal of this notebook is to explore the phenomenon of bracket closing in the [GPT-Neo 125M model](https://www.eleuther.ai/artifacts/gpt-neo), whereby it can correctly match open parentheses `([{<` with their corresponding closing versions `)]}>`.\n",
    "\n",
    "This is [Problem 2.13](https://www.alignmentforum.org/s/yivyHaCAmMJ3CqSyj/p/XNjRwEX9kxbpzWFWd#block71) in Neel Nanda's [200 Concrete Open Problems in Mechanistic Interpretability](https://www.alignmentforum.org/posts/LbrPTJ4fmABEdEnLf/200-concrete-open-problems-in-mechanistic-interpretability). The first goal is to figure out how the model determines whether an opening or closing bracket is more appropriate, and the second is to figure out how it knows the correct kind: `(`, `[`, `{` or `<`.\n",
    "\n",
    "I'm using the [TransformerLens library](https://github.com/neelnanda-io/TransformerLens), and a lot of this notebook is copied from Neel's [Exploratory Analysis notebook](https://neelnanda.io/exploratory-analysis-demo). See that notebook for more details on the techniques used.\n",
    "\n",
    "This notebook lives in my [mechanistic interpretability GitHub repository](https://github.com/SamAdamDay/mechanistic-interpretability-projects).\n",
    "\n",
    "There are a few inconsistencies in this notebook, where I changed my approach a bit in light of my findings."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Main findings\n",
    "\n",
    "- GPT-Neo is very good at predicting a closing bracket over an opening bracket when there's an imbalance, and is robust to most of the things I threw at it to try and break it.\n",
    "- I've thought a fair bit about what the task actually is.\n",
    "    + Does always predicting `)` over `(` when brackets are imbalanced correspond to a model's preference to generate text which is overall bracket-balanced?\n",
    "    + Should we include tokens like '],' when seeing if the model wants to generate `[`?\n",
    "    + I explore these questions in more detail with reference to some of the experimental results (see [Discussion](#discussion) and [Reflections](#reflections)).\n",
    "- I develop two metrics for evaluating the model performance on the task.\n",
    "- I use the direct logit attribution method to find the parts of the model which most strongly directly affect its performance.\n",
    "    * The first part is the layer-5 attention. \n",
    "        + Many heads are responsible for the improvement seen here.\n",
    "        + The principal one, head 1, has an attention pattern which sends information from each scoping symbol (brackets and quotation marks) to the subsequent tokens until the next scoping symbol.\n",
    "        + In particular, I've found that the model includes quotation marks together with brackets when computing scopes.\n",
    "    * The second is the layer-8 MLP \n",
    "    * The third part is head 7 in attention layer 9.\n",
    "        + The attention pattern does something like sending information from opening brackets depending on the parity of the current bracket depth.\n",
    "    "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVELOPMENT_MODE = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running as a Jupyter notebook - intended for development only!\n",
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_36307/2873378927.py:13: DeprecationWarning: `magic(...)` is deprecated since IPython 0.13 (warning added in 8.1), use run_line_magic(magic_name, parameter_s).\n",
      "  ipython.magic(\"load_ext autoreload\")\n",
      "/tmp/ipykernel_36307/2873378927.py:14: DeprecationWarning: `magic(...)` is deprecated since IPython 0.13 (warning added in 8.1), use run_line_magic(magic_name, parameter_s).\n",
      "  ipython.magic(\"autoreload 2\")\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    import google.colab\n",
    "    IN_COLAB = True\n",
    "    print(\"Running as a Colab notebook\")\n",
    "    %pip install git+https://github.com/SamAdamDay/mechanistic-interpretability-projects.git\n",
    "except:\n",
    "    IN_COLAB = False\n",
    "    print(\"Running as a Jupyter notebook - intended for development only!\")\n",
    "    from IPython import get_ipython\n",
    "\n",
    "    ipython = get_ipython()\n",
    "    # Code to automatically update the HookedTransformer code as its edited without restarting the kernel\n",
    "    ipython.magic(\"load_ext autoreload\")\n",
    "    ipython.magic(\"autoreload 2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using renderer: colab\n"
     ]
    }
   ],
   "source": [
    "import plotly.io as pio\n",
    "\n",
    "if IN_COLAB or DEVELOPMENT_MODE:\n",
    "    pio.renderers.default = \"colab\"\n",
    "else:\n",
    "    pio.renderers.default = \"notebook_connected\"\n",
    "print(f\"Using renderer: {pio.renderers.default}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import things"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from pathlib import Path\n",
    "from typing import List, Union, Optional\n",
    "from functools import partial\n",
    "import copy\n",
    "import itertools\n",
    "import dataclasses\n",
    "from collections import OrderedDict\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import einops\n",
    "\n",
    "from fancy_einsum import einsum\n",
    "\n",
    "import tqdm.auto as tqdm\n",
    "\n",
    "import plotly.express as px\n",
    "\n",
    "from jaxtyping import Float, Int\n",
    "\n",
    "from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer\n",
    "import datasets\n",
    "\n",
    "from IPython.display import HTML\n",
    "\n",
    "import circuitsvis as cv\n",
    "from circuitsvis.utils.render import RenderedHTML\n",
    "\n",
    "import transformer_lens\n",
    "import transformer_lens.utils as utils\n",
    "from transformer_lens.hook_points import (\n",
    "    HookedRootModule,\n",
    "    HookPoint,\n",
    ")  # Hooking utilities\n",
    "from transformer_lens import (\n",
    "    HookedTransformer,\n",
    "    HookedTransformerConfig,\n",
    "    FactoredMatrix,\n",
    "    ActivationCache,\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Turn automatic differentiation off"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch.autograd.grad_mode.set_grad_enabled at 0x7f18e41a2a50>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.set_grad_enabled(False)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Torch device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sam/.virtualenvs/mech-interp/lib/python3.11/site-packages/torch/cuda/__init__.py:88: UserWarning:\n",
      "\n",
      "CUDA initialization: CUDA unknown error - this may be due to an incorrectly set up environment, e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the available devices to be zero. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:109.)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(tensor, renderer=None, **kwargs):\n",
    "    px.imshow(\n",
    "        utils.to_numpy(tensor),\n",
    "        color_continuous_midpoint=0.0,\n",
    "        color_continuous_scale=\"RdBu\",\n",
    "        **kwargs\n",
    "    ).show(renderer)\n",
    "\n",
    "\n",
    "def line(tensor, renderer=None, **kwargs):\n",
    "    px.line(y=utils.to_numpy(tensor), **kwargs).show(renderer)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Task specification\n",
    "\n",
    "The basic task is as follows.\n",
    "\n",
    "**Task.** Given a string $s$ containing some brackets, determine: (1) if an opening or closing bracket is more appropriate and (2) which type of bracket is most appropriate.\n",
    "\n",
    "We'll be using the GPT-Neo 125M model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained model gpt-neo-125M into HookedTransformer\n"
     ]
    }
   ],
   "source": [
    "model = HookedTransformer.from_pretrained(\n",
    "    \"gpt-neo-125M\",\n",
    "    center_unembed=True,\n",
    "    center_writing_weights=True,\n",
    "    fold_ln=True,\n",
    "    refactor_factored_attn_matrices=True,\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll look at the brackets `([{<)]}>`. We want to match all tokens which either begin with a bracket, or a space followed by the bracket. We want to match for example the token ').' for the bracket `)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "brackets: (('(', '[', '{', '<'), (')', ']', '}', '>'))\n",
      "bracket_pairs: (('(', ')'), ('[', ']'), ('{', '}'), ('<', '>'))\n",
      "brackets_flat: ('(', '[', '{', '<', ')', ']', '}', '>')\n",
      "( tokens: ['(', ' (', '()', ' (@', ' (\"', ' ($', '(\"', ' ()', '();', \"('\", ' (+', ' (%)', ' (-', ' ();', ' ((', '({', '($', ' (#', \" ('\", '((', ' (.', ' (*', '().', ' (!', '(),', ' (£', '([', ' ().', '(_', '())', ' ([', ' (),', ' (~', '(-', ' (?,', ' ())', '():', '());', ' (&', ' (−', ' (%', ' ({', '(\\\\', ' (<', ' ());', '(&', '(){', ' (_', ' (>', ' ($)', ' (=', '(*', ' (/']\n",
      "[ tokens: ['[', ' [', '[/', ' [\"', '[\"', ' [[', ' []', \"['\", '[]', ' […]', ' [];', ' [-', ' [+', ' [...]', '[_', '[[', ' [*', ' [*]', \" ['\", ' [/', ' [+]', ' [(', ' [|', ' [&']\n",
      "{ tokens: ['{', ' {', '{\"', ' {\"', ' {{', ' {}', '{{', '{\\\\', ' {\\\\', ' {:', ' {*']\n",
      "< tokens: ['<', ' <', '</', ' </', ' <<', '<<', ' <=', ' <-', ' <[', ' <@', ' <!--', ' <+', '<?']\n",
      ") tokens: [')', ').', '),', ' )', ');', '):', '))', ' );', ')(', ' ).', ' ),', ')-', ')|', ' ):', ' ))', ')]', ')\"', '));', ')\\\\', ')?', '){', ')/', ').\"', ')))', ')].', ')...', '),\"', ')*', ')—', ' ));', ')).', ')!', \")'\", '))))', ')=(', ')</', ')),', ')}', ')[', ')\",', ').[', ')--', ' )))', ')=', ')+', ' )]']\n",
      "] tokens: [']', ' ]', '].', '],', ']:', '][', '];', ']]', '])', '](', ' ],', '],\"', ' ].', ']=', ' ];', '].\"', ']\"', ' ])', ']).', ']);', '],[', ' ][', '][/', ']-', ']),', ']+', ']}', \"]'\"]\n",
      "} tokens: ['}', ' }', '},', '},{\"', ' },', '},\"', '}}', ' });', ' };', '}{', '};', '});', '}.', '})', ' })', '}\\\\', ' }}', '}\"', '}:', '}}}']\n",
      "> tokens: ['>', ' >', '>>', '><', ' >>', '></', ' >>>', '>>>>', ' >=', '>,', '>\"', '>.', ' ><', '>:', '>>>', '>>>>>>>>', '>(', '>>\\\\', '>[', '>]', '>)']\n",
      "bracket_tokens_flat: torch.Size([216])\n",
      "open_bracket_tokens: torch.Size([101])\n",
      "closed_bracket_tokens: torch.Size([115])\n",
      "bracket_tokens_sizes: [53, 24, 11, 13, 46, 28, 20, 21]\n"
     ]
    }
   ],
   "source": [
    "brackets = (tuple(\"([{<\"), tuple(\")]}>\"))\n",
    "bracket_pairs = tuple(zip(*brackets))\n",
    "brackets_flat = brackets[0] + brackets[1]\n",
    "num_brackets = len(brackets_flat)\n",
    "\n",
    "# Get all (non-endoftext tokens) tokens which either start with a bracket or a\n",
    "# space followed by a bracket\n",
    "bracket_tokens = OrderedDict([(bracket, []) for bracket in brackets_flat])\n",
    "bracket_token_strs = OrderedDict([(bracket, []) for bracket in brackets_flat])\n",
    "all_tokens = model.to_str_tokens(np.arange(model.cfg.d_vocab - 1), prepend_bos=False)\n",
    "for i, token_str in enumerate(all_tokens):\n",
    "    for bracket in brackets_flat:\n",
    "        if token_str.startswith(bracket) or token_str.startswith(\" \" + bracket):\n",
    "            bracket_tokens[bracket].append(i)\n",
    "            bracket_token_strs[bracket].append(token_str)\n",
    "for bracket, tokens in bracket_tokens.items():\n",
    "    bracket_tokens[bracket] = torch.tensor(tokens)\n",
    "\n",
    "# Flatten the dict of tokens, and record the sizes of each list\n",
    "bracket_tokens_flat = torch.cat(list(bracket_tokens.values()))\n",
    "bracket_tokens_sizes = [tokens.shape[0] for tokens in bracket_tokens.values()]\n",
    "\n",
    "# Select the open and closed bracket tokens\n",
    "num_open_bracket_tokens = sum(bracket_tokens_sizes[:num_brackets // 2])\n",
    "open_bracket_tokens = bracket_tokens_flat[:num_open_bracket_tokens]\n",
    "closed_bracket_tokens = bracket_tokens_flat[num_open_bracket_tokens:]\n",
    "\n",
    "print(\"brackets:\", brackets)\n",
    "print(\"bracket_pairs:\", bracket_pairs)\n",
    "print(\"brackets_flat:\", brackets_flat)\n",
    "for bracket, token_strs in bracket_token_strs.items():\n",
    "    print(f\"{bracket} tokens:\", token_strs)\n",
    "print(\"bracket_tokens_flat:\", bracket_tokens_flat.shape)\n",
    "print(\"open_bracket_tokens:\", open_bracket_tokens.shape)\n",
    "print(\"closed_bracket_tokens:\", closed_bracket_tokens.shape)\n",
    "print(\"bracket_tokens_sizes:\", bracket_tokens_sizes)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring model capability\n",
    "\n",
    "How good is GPT-Neo at closing brackets? In this section I explore its capabilities and try to break it. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I will explore the following variations on the string $s$.\n",
    "- Whether the brackets are balanced or not.\n",
    "- The type of brackets used.\n",
    "- Whether we mix different types.\n",
    "- The complexity of the bracket structure. This can be thought of as a tree, and we can consider varying both its depth and breadth.\n",
    "- The complexity of the rest of the string.\n",
    "- Whether $s$ looks like real code. I'll look at the following ways this could fail.\n",
    "    * It's actually natural language.\n",
    "    * It's like a programming language but has syntax errors.\n",
    "    * It's valid syntax but the symbol names are gibberish/unnatural.\n",
    "    * It consists only of brackets."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploratory prompts\n",
    "\n",
    "I will test the following prompts, to see what the model does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_exploratory_prompts: 72\n"
     ]
    }
   ],
   "source": [
    "exploratory_prompts = [\n",
    "    \"def line(tensor, renderer=None\",\n",
    "    \"def line(tensor, renderer=None)\",\n",
    "    \"exploratory_prompts = ['test'\",\n",
    "    \"exploratory_prompts = ['test']\",\n",
    "    \"array[0\",\n",
    "    \"array[0]\",\n",
    "    \"exploratory_dict = {'test': 'four'\",\n",
    "    \"exploratory_dict = {'test': 'four'}\",\n",
    "    \"<template\",\n",
    "    \"<template>\",\n",
    "    \"def sieve(num, prime_list = [2, 3]\",\n",
    "    \"def sieve(num, prime_list = [2, 3])\",\n",
    "    \"exploratory_dict = {'test': [3, 5]\",\n",
    "    \"exploratory_dict = {'test': [3, 5]}\",\n",
    "    \"exploratory_dict = {'test': get_test()\",\n",
    "    \"exploratory_dict = {'test': get_test()}\",\n",
    "    \"html_to_markdown('<s>'\",\n",
    "    \"html_to_markdown('<s>')\",\n",
    "    \"<table id='name()'\",\n",
    "    \"<table id='name()'>\",\n",
    "    \"load_model(build_structure()\",\n",
    "    \"load_model(build_structure())\",\n",
    "    \"load_model(build_structure(), get_hyperparameters()\",\n",
    "    \"load_model(build_structure(), get_hyperparameters())\",\n",
    "    \"load_model(build_structure(), get_hyperparameters(), (True, False)\",\n",
    "    \"load_model(build_structure(), get_hyperparameters(), (True, False))\",\n",
    "    \"load_model(build_structure(), get_hyperparameters(True), (True, False), get_extra_config(x)\",\n",
    "    \"load_model(build_structure(), get_hyperparameters(True), (True, False), get_extra_config(x))\",\n",
    "    \"x.detach().cpu().to_numpy(\",\n",
    "    \"x.detach().cpu().to_numpy()\",\n",
    "    \"enumerate(list(zip([1,3,65], [1, 2, 3])\",\n",
    "    \"enumerate(list(zip([1,3,65], [1, 2, 3]\",\n",
    "    \"enumerate(list(zip([1,3,65], [1, 2, 3\",\n",
    "    \"enumerate(list(zip([1,3,65], [1, 2, 3]))\",\n",
    "    \"enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4))\",\n",
    "    \"enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4)))\",\n",
    "    \"zip(enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4)))), [3, 4]\",\n",
    "    \"zip(enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4)))), [3, 4])\",\n",
    "    \"zip(enumerate(list(zip([1,3,65], list(np.zeros(4)))), [3, 4], {1: 2}.items()\",\n",
    "    \"zip(enumerate(list(zip([1,3,65], list(np.zeros(4)))), [3, 4], {1: 2}.items())\",\n",
    "    \"list())))).append(x\",\n",
    "    \"list())))).append(x)\",\n",
    "    \"px.imshow(to_numpy(tensor), ccm=0.0, ccs='RdBu', labels={'x':xaxis, 'y':yaxis}, **kwargs\",\n",
    "    \"px.imshow(to_numpy(tensor), ccm=0.0, ccs='RdBu', labels={'x':xaxis, 'y':yaxis}, **kwargs)\",\n",
    "    \"In the course our our analysis (which was long\",\n",
    "    \"In the course our our analysis (which was long)\",\n",
    "    \"He was eating a apple [sic\",\n",
    "    \"He was eating a apple [sic]\",\n",
    "    \"In the course our our analysis (which was long (though not too long\",\n",
    "    \"In the course our our analysis (which was long (though not too long))\",\n",
    "    \"def sieve(,num prime_list = 2[, 3]\",\n",
    "    \"def sieve(,num prime_list = 2[, 3])\",\n",
    "    \"defn line(tensor, renderer===None\",\n",
    "    \"defn line(tensor, renderer===None)\",\n",
    "    \"exploratory_dict = {'test': [3,} 5]\",\n",
    "    \"exploratory_dict = {'test': [3,} 5\",\n",
    "    \"exploratory_prompts = ['test'(]\",\n",
    "    \"exploratory_prompts = ['test'(])\",\n",
    "    \"def safasfd(oubefwef, vcewfec=afuasvfs\",\n",
    "    \"def safasfd(oubefwef, vcewfec=afuasvfs)\",\n",
    "    \"asdjhvauyrfsac = ['asdasdasd'\",\n",
    "    \"asdjhvauyrfsac = ['asdasdasd']\",\n",
    "    \"dfc = {'sdasd': 'casdasd'\",\n",
    "    \"dfc = {'sdasd': 'casdasd'}\",\n",
    "    \"<bwevzcxc\",\n",
    "    \"<bwevzcxc>\",\n",
    "    \"([]\",\n",
    "    \"([])\",\n",
    "    \"([({},[{()}])])\",\n",
    "    \"([({},[{()}])]\",\n",
    "    \"([({},[{()}])\",\n",
    "    \"([({},[{()}]\",\n",
    "]\n",
    "num_exploratory_prompts = len(exploratory_prompts)\n",
    "print(\"num_exploratory_prompts:\", num_exploratory_prompts)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's convert the prompts to padded tokens, keeping track of each unpadded length, so we can find the next predicted token for each one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "exp_prompt_tokens: torch.Size([72, 47])\n"
     ]
    }
   ],
   "source": [
    "# Compute the token length of each prompt, so we know where the next-token\n",
    "# prediction will be\n",
    "exp_prompt_token_lengths = []\n",
    "for prompt in exploratory_prompts:\n",
    "    prompt_tokens = model.to_tokens(prompt)\n",
    "    exp_prompt_token_lengths.append(prompt_tokens.shape[1])\n",
    "\n",
    "# Convert all the prompts to tokens, padding to make them the same length\n",
    "exp_prompt_tokens = model.to_tokens(exploratory_prompts)\n",
    "exp_prompt_tokens.to(device)\n",
    "\n",
    "print(\"exp_prompt_tokens:\", exp_prompt_tokens.shape)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Investigating performance\n",
    "\n",
    "The following function computes runs the model, looks at the predictions for the next tokens for each prompt, and computes the probability that it each possible bracket (including spaces), conditioned on that it actually is a bracket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_bracket_scores(\n",
    "    prompt_tokens: Float[torch.Tensor, \"batch pos\"],\n",
    "    prompt_token_lengths: Optional[list[int]] = None,\n",
    ") -> Float[torch.Tensor, \"batch n_brackets\"]:\n",
    "    \"\"\"Computes the conditional prob that the next token is each bracket\n",
    "\n",
    "    Conditioned on the next token actually being a bracket\n",
    "    \"\"\"\n",
    "\n",
    "    num_prompts = prompt_tokens.shape[0]\n",
    "\n",
    "    all_logits = model(prompt_tokens, return_type=\"logits\")  # batch pos d_vocab\n",
    "\n",
    "    d_vocab = all_logits.shape[2]\n",
    "\n",
    "    # Select the last token from each\n",
    "    if prompt_token_lengths is None:\n",
    "        logits = all_logits[:, prompt_tokens.shape[1] - 1, :]  # batch d_vocab\n",
    "    else:\n",
    "        indices = torch.tensor(prompt_token_lengths, device=device) - 1  # batch\n",
    "        indices = indices.reshape((num_prompts, 1, 1))  # batch 1 1\n",
    "        indices = indices.repeat((1, 1, d_vocab))  # batch 1 d_vocab\n",
    "        logits = torch.gather(all_logits, 1, indices).squeeze()  # batch d_vocab\n",
    "\n",
    "    probs = F.softmax(logits, dim=1)  # batch d_vocab\n",
    "\n",
    "    # Compute the probability for each bracket and spaced bracket, conditioned\n",
    "    # on the fact that it is a bracket\n",
    "    cond_probs = probs[:, bracket_tokens_flat]  # batch (2 n_bracket_tokens)\n",
    "    cond_probs = F.normalize(cond_probs, p=1.0, dim=1)\n",
    "\n",
    "    # Combine the conditional probabilities for each bracket\n",
    "    cond_probs_combined = torch.zeros((num_prompts, num_brackets))\n",
    "    index = 0\n",
    "    for i, size in enumerate(bracket_tokens_sizes):\n",
    "        cond_probs_combined[:, i] = cond_probs[:, index : index + size].sum(dim=1)\n",
    "        index += size\n",
    "\n",
    "    return cond_probs_combined\n",
    "\n",
    "\n",
    "bracket_scores = compute_bracket_scores(exp_prompt_tokens, exp_prompt_token_lengths)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's display these in a nice chart. I break it up into two since there are a lot of prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def display_bracket_scores(\n",
    "    prompts: list[str],\n",
    "    bracket_scores: Float[torch.Tensor, \"batch n_brackets\"],\n",
    "    height_scale: int = 30,\n",
    "):\n",
    "    \"\"\"Display the bracket scores nicely\"\"\"\n",
    "    num_prompts = len(prompts)\n",
    "    fig = px.imshow(\n",
    "        utils.to_numpy(bracket_scores),\n",
    "        color_continuous_scale=\"blues\",\n",
    "        labels=dict(x=\"Bracket\", color=\"Conditional Probability\"),\n",
    "        x=brackets_flat,\n",
    "        y=prompts,\n",
    "        height=height_scale * num_prompts,\n",
    "    )\n",
    "    for ix, bracket in enumerate(brackets_flat):\n",
    "        for iy in range(num_prompts):\n",
    "            fig.add_annotation(\n",
    "                x=ix,\n",
    "                y=iy,\n",
    "                text=bracket,\n",
    "                showarrow=False,\n",
    "                font_color=\"orange\",\n",
    "            )\n",
    "    fig.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<html>\n",
       "<head><meta charset=\"utf-8\" /></head>\n",
       "<body>\n",
       "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
       "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.20.0.min.js\"></script>                <div id=\"85b96044-2597-4bf4-aef5-68e3b55237b6\" class=\"plotly-graph-div\" style=\"height:1080px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"85b96044-2597-4bf4-aef5-68e3b55237b6\")) {                    Plotly.newPlot(                        \"85b96044-2597-4bf4-aef5-68e3b55237b6\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"(\",\"[\",\"{\",\"<\",\")\",\"]\",\"}\",\">\"],\"y\":[\"def line(tensor, renderer=None\",\"def line(tensor, renderer=None)\",\"exploratory_prompts = ['test'\",\"exploratory_prompts = ['test']\",\"array[0\",\"array[0]\",\"exploratory_dict = {'test': 'four'\",\"exploratory_dict = {'test': 'four'}\",\"<template\",\"<template>\",\"def sieve(num, prime_list = [2, 3]\",\"def sieve(num, prime_list = [2, 3])\",\"exploratory_dict = {'test': [3, 5]\",\"exploratory_dict = {'test': [3, 5]}\",\"exploratory_dict = {'test': get_test()\",\"exploratory_dict = {'test': get_test()}\",\"html_to_markdown('<s>'\",\"html_to_markdown('<s>')\",\"<table id='name()'\",\"<table id='name()'>\",\"load_model(build_structure()\",\"load_model(build_structure())\",\"load_model(build_structure(), get_hyperparameters()\",\"load_model(build_structure(), get_hyperparameters())\",\"load_model(build_structure(), get_hyperparameters(), (True, False)\",\"load_model(build_structure(), get_hyperparameters(), (True, False))\",\"load_model(build_structure(), get_hyperparameters(True), (True, False), get_extra_config(x)\",\"load_model(build_structure(), get_hyperparameters(True), (True, False), get_extra_config(x))\",\"x.detach().cpu().to_numpy(\",\"x.detach().cpu().to_numpy()\",\"enumerate(list(zip([1,3,65], [1, 2, 3])\",\"enumerate(list(zip([1,3,65], [1, 2, 3]\",\"enumerate(list(zip([1,3,65], [1, 2, 3\",\"enumerate(list(zip([1,3,65], [1, 2, 3]))\",\"enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4))\",\"enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4)))\"],\"z\":[[1.609617356734816e-05,3.3054224331863225e-05,1.992840594766676e-07,2.785497201784892e-07,0.9999461770057678,1.3260228115541395e-06,2.568428271843004e-06,7.152826242418087e-07],[0.3173849582672119,0.34839096665382385,0.22829656302928925,0.008658126927912235,0.08334029465913773,0.0021563128102570772,5.715771476388909e-05,0.01171552762389183],[5.701024201698601e-05,8.528495527571067e-05,1.4170408348945784e-06,0.00010017160093411803,3.256875788792968e-05,0.9981484413146973,0.0013124402612447739,0.00026260322192683816],[0.06506045162677765,0.06276267766952515,0.0041067167185246944,0.01613546721637249,0.47414323687553406,0.25735458731651306,0.10328751057386398,0.01714983582496643],[7.933303277241066e-05,0.00012013951345579699,1.8481894585420378e-05,0.00019387408974580467,0.00020359770860522985,0.9987819790840149,0.0004479825438465923,0.0001547060237498954],[0.04054002836346626,0.0996202751994133,0.09309915453195572,0.1915598064661026,0.229003444314003,0.2099682241678238,0.032061509788036346,0.10414756834506989],[7.024689693935215e-05,9.727911674417555e-05,9.942711585608777e-06,4.211832128930837e-05,4.4795464759772585e-07,7.518713391618803e-05,0.9996981620788574,6.701509391859872e-06],[0.03930394724011421,0.5143156051635742,0.026576373726129532,0.010578923858702183,0.21011267602443695,0.1617252230644226,0.031852636486291885,0.005534815602004528],[6.201586074894294e-05,3.1350384233519435e-05,4.58669592262595e-06,3.2986961741698906e-05,1.0613031236061943e-06,2.1040624176293932e-07,2.2042199177008115e-08,0.999867856502533],[0.003987165633589029,0.0003280879172962159,0.32636889815330505,0.663361132144928,0.0021191260311752558,0.00011482479749247432,0.0021040833089500666,0.0016169665614143014],[0.0003370226768311113,0.00044635715312324464,5.43029818800278e-05,1.3440484281090903e-06,0.9991406798362732,9.687974852568004e-06,9.131572937803867e-07,9.590323315933347e-06],[0.03845832869410515,0.2305075079202652,0.6141166687011719,0.0030709197744727135,0.10273788869380951,0.004210366867482662,0.00012686557602137327,0.006771921180188656],[0.00017551102791912854,0.000326345965731889,0.00013521083747036755,7.793560507707298e-05,4.511602583079366e-06,0.00018504710169509053,0.9990638494491577,3.189227936672978e-05],[0.022005023434758186,0.09190785884857178,0.005961531773209572,0.02382291853427887,0.20523318648338318,0.35139229893684387,0.29584819078445435,0.0038288270588964224],[0.0006291572353802621,0.007054484449326992,5.253767085378058e-05,0.0001529658038634807,6.04801016379497e-06,8.591034566052258e-05,0.9920012354850769,1.763384534569923e-05],[0.046468764543533325,0.386040061712265,0.018732581287622452,0.006070784758776426,0.3840903043746948,0.1094626635313034,0.04558943212032318,0.003545647719874978],[0.00863292533904314,0.026185426861047745,0.05960044637322426,0.02563023567199707,0.8597300052642822,0.009833105839788914,0.0050544473342597485,0.005333516281098127],[0.00745142949745059,0.22554203867912292,0.4050668478012085,0.07125060260295868,0.23563911020755768,0.02786274626851082,0.0151916379109025,0.011995525099337101],[0.0003436957485973835,0.0017240454908460379,0.0004237227258272469,0.004584697540849447,0.00022299525153357536,0.000546713883522898,0.00011198189167771488,0.9920421838760376],[0.0020694041159003973,0.004159319680184126,0.05741678550839424,0.9188671112060547,0.005957946181297302,0.0008211928652599454,0.010118766687810421,0.0005898329545743763],[0.011611465364694595,0.04259537160396576,0.0920693427324295,0.08050640672445297,0.7701142430305481,2.899942046497017e-05,3.718504012795165e-05,0.003037415212020278],[0.0008271806873381138,0.010552724823355675,0.9398097991943359,0.023127006366848946,0.012579532340168953,0.002497964771464467,0.0032037030905485153,0.007402563001960516],[0.00959315150976181,0.20443037152290344,0.007975746877491474,0.007019079755991697,0.7678603529930115,2.105999374180101e-05,9.49143577599898e-05,0.0030053460504859686],[0.002306845737621188,0.0952053964138031,0.6377246975898743,0.035850707441568375,0.12021268159151077,0.03836676850914955,0.04752829298377037,0.02280428633093834],[0.1416509598493576,0.08432875573635101,0.2721467614173889,0.023991448804736137,0.4356921315193176,0.0022873838897794485,0.0022391004022210836,0.03766360878944397],[0.06479150801897049,0.04926690831780434,0.16305406391620636,0.026185180991888046,0.6122965812683105,0.027922168374061584,0.012721704319119453,0.04376155510544777],[0.02259320765733719,0.06590501219034195,0.07726345211267471,0.014713450334966183,0.7394962310791016,0.010649065487086773,0.024793215095996857,0.044586002826690674],[0.01581173576414585,0.06648525595664978,0.19624970853328705,0.03211677446961403,0.5858823657035828,0.052803441882133484,0.00741581991314888,0.04323483631014824],[0.06164143979549408,0.7988715171813965,0.05018547177314758,0.014927104115486145,0.0549805574119091,0.0037797843106091022,0.00688927061855793,0.008724731393158436],[0.04242953658103943,0.22278039157390594,0.0026780515909194946,0.02065380848944187,0.6336902976036072,0.025079194456338882,0.0239628404378891,0.028725886717438698],[0.00025303184520453215,0.013660309836268425,0.0006482136668637395,0.0002263985079480335,0.9685782194137573,0.016530495136976242,2.914366814366076e-05,7.431350968545303e-05],[6.615332677029073e-05,6.781219417462125e-05,1.5903136954875663e-05,1.122779940487817e-05,0.9994860291481018,0.00034838810097426176,1.1610555930019473e-06,3.5140078580297995e-06],[2.929209040303249e-05,0.0001359581801807508,2.0037684862472815e-06,3.505746917653596e-06,0.00029827607795596123,0.9992111325263977,9.633033187128603e-05,0.00022338244889397174],[0.0015211787540465593,0.031201614066958427,0.0010750481160357594,0.0004477985785342753,0.9641165733337402,0.0011762977810576558,8.912254270398989e-05,0.0003722446854226291],[0.00027245012461207807,0.022021664306521416,2.7592808692133985e-05,0.00024337841023225337,0.965874433517456,0.011371328495442867,4.840717156184837e-05,0.00014080012624617666],[0.00019971512665506452,0.015324669890105724,1.8489130525267683e-05,0.00027502820012159646,0.9801182746887207,0.0035331544931977987,0.0002563177840784192,0.00027446672902442515]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Bracket: %{x}<br>y: %{y}<br>Conditional Probability: %{z}<extra></extra>\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Bracket\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\"},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Conditional Probability\"}},\"colorscale\":[[0.0,\"rgb(247,251,255)\"],[0.125,\"rgb(222,235,247)\"],[0.25,\"rgb(198,219,239)\"],[0.375,\"rgb(158,202,225)\"],[0.5,\"rgb(107,174,214)\"],[0.625,\"rgb(66,146,198)\"],[0.75,\"rgb(33,113,181)\"],[0.875,\"rgb(8,81,156)\"],[1.0,\"rgb(8,48,107)\"]]},\"margin\":{\"t\":60},\"height\":1080,\"annotations\":[{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":35}]},                        {\"responsive\": true}                    ).then(function(){\n",
       "                            \n",
       "var gd = document.getElementById('85b96044-2597-4bf4-aef5-68e3b55237b6');\n",
       "var x = new MutationObserver(function (mutations, observer) {{\n",
       "        var display = window.getComputedStyle(gd).display;\n",
       "        if (!display || display === 'none') {{\n",
       "            console.log([gd, 'removed!']);\n",
       "            Plotly.purge(gd);\n",
       "            observer.disconnect();\n",
       "        }}\n",
       "}});\n",
       "\n",
       "// Listen for the removal of the full notebook cells\n",
       "var notebookContainer = gd.closest('#notebook-container');\n",
       "if (notebookContainer) {{\n",
       "    x.observe(notebookContainer, {childList: true});\n",
       "}}\n",
       "\n",
       "// Listen for the clearing of the current output cell\n",
       "var outputEl = gd.closest('.output');\n",
       "if (outputEl) {{\n",
       "    x.observe(outputEl, {childList: true});\n",
       "}}\n",
       "\n",
       "                        })                };                            </script>        </div>\n",
       "</body>\n",
       "</html>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_bracket_scores(exploratory_prompts[:num_exploratory_prompts // 2], bracket_scores[:num_exploratory_prompts // 2, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<html>\n",
       "<head><meta charset=\"utf-8\" /></head>\n",
       "<body>\n",
       "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
       "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.20.0.min.js\"></script>                <div id=\"790ebab3-28bd-4ed7-95ee-cf00a1625414\" class=\"plotly-graph-div\" style=\"height:1080px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"790ebab3-28bd-4ed7-95ee-cf00a1625414\")) {                    Plotly.newPlot(                        \"790ebab3-28bd-4ed7-95ee-cf00a1625414\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"(\",\"[\",\"{\",\"<\",\")\",\"]\",\"}\",\">\"],\"y\":[\"zip(enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4)))), [3, 4]\",\"zip(enumerate(list(zip([1,3,65], [1, 2, 3], [4, 3], [1], list(np.zeros(4)))), [3, 4])\",\"zip(enumerate(list(zip([1,3,65], list(np.zeros(4)))), [3, 4], {1: 2}.items()\",\"zip(enumerate(list(zip([1,3,65], list(np.zeros(4)))), [3, 4], {1: 2}.items())\",\"list())))).append(x\",\"list())))).append(x)\",\"px.imshow(to_numpy(tensor), ccm=0.0, ccs='RdBu', labels={'x':xaxis, 'y':yaxis}, **kwargs\",\"px.imshow(to_numpy(tensor), ccm=0.0, ccs='RdBu', labels={'x':xaxis, 'y':yaxis}, **kwargs)\",\"In the course our our analysis (which was long\",\"In the course our our analysis (which was long)\",\"He was eating a apple [sic\",\"He was eating a apple [sic]\",\"In the course our our analysis (which was long (though not too long\",\"In the course our our analysis (which was long (though not too long))\",\"def sieve(,num prime_list = 2[, 3]\",\"def sieve(,num prime_list = 2[, 3])\",\"defn line(tensor, renderer===None\",\"defn line(tensor, renderer===None)\",\"exploratory_dict = {'test': [3,} 5]\",\"exploratory_dict = {'test': [3,} 5\",\"exploratory_prompts = ['test'(]\",\"exploratory_prompts = ['test'(])\",\"def safasfd(oubefwef, vcewfec=afuasvfs\",\"def safasfd(oubefwef, vcewfec=afuasvfs)\",\"asdjhvauyrfsac = ['asdasdasd'\",\"asdjhvauyrfsac = ['asdasdasd']\",\"dfc = {'sdasd': 'casdasd'\",\"dfc = {'sdasd': 'casdasd'}\",\"<bwevzcxc\",\"<bwevzcxc>\",\"([]\",\"([])\",\"([({},[{()}])])\",\"([({},[{()}])]\",\"([({},[{()}])\",\"([({},[{()}]\"],\"z\":[[3.783132342505269e-05,5.1831357268383726e-05,3.6239005112292944e-06,1.3839829989592545e-05,0.9997289180755615,0.00014232848479878157,4.293161509849597e-06,1.7146552636404522e-05],[0.00012034500105073676,0.006113966461271048,0.0001571957691339776,0.00028586047119461,0.9904594421386719,0.0018653444712981582,0.00024657801259309053,0.0007510727737098932],[0.00011730697588063776,0.01523576769977808,1.5473602616111748e-05,5.639057053485885e-05,0.9842443466186523,0.00015284647815860808,0.00015316430653911084,2.4948129066615365e-05],[0.00015281829109881073,0.1519942283630371,0.0007371343090198934,0.0008431801106780767,0.8418084383010864,0.001006837235763669,0.0017511456971988082,0.0017064801650121808],[0.014876144006848335,0.07230646908283234,9.023873280966654e-05,0.00043510922114364803,0.9108191132545471,0.0009662336087785661,0.00014890721649862826,0.0003581260098144412],[0.0051384675316512585,0.012794174253940582,0.03802643343806267,0.039977528154850006,0.1775882989168167,0.4542897343635559,0.24731911718845367,0.024866294115781784],[0.00038896402111276984,0.002319883555173874,0.00012180537305539474,3.3715277822921053e-06,0.9966431856155396,9.020037396112457e-05,0.00041413403232581913,1.8808779714163393e-05],[0.010102849453687668,0.2312985062599182,0.012104828841984272,0.019382573664188385,0.12173610180616379,0.24417459964752197,0.3023645281791687,0.058835748583078384],[0.0065517486073076725,0.0008548697805963457,1.844005601014942e-05,1.8493101379135624e-05,0.9919562339782715,0.00011712627747328952,0.000348947593010962,0.00013441653572954237],[0.6288301348686218,0.312997043132782,0.002632362535223365,0.0043196901679039,0.023539774119853973,0.017418742179870605,0.007827507331967354,0.0024350425228476524],[6.686874257866293e-05,3.308882878627628e-05,2.900109166148468e-07,3.884119905706029e-06,0.0015944107435643673,0.9982820153236389,1.1906608051504008e-05,7.566005479020532e-06],[0.25786975026130676,0.6813380122184753,0.011278100311756134,0.006678693927824497,0.009518916718661785,0.02544407732784748,0.00020725687500089407,0.007665226701647043],[0.0008377173799090087,0.00018905229808297008,3.536929398251232e-07,1.1715917935362086e-05,0.9989267587661743,1.564390186103992e-05,2.6314271508454112e-06,1.6034118743846193e-05],[0.43765321373939514,0.3432966470718384,0.005631020292639732,0.006646838039159775,0.1296047866344452,0.05732715129852295,0.014347088523209095,0.005493395496159792],[0.0010961915832012892,0.0008884978014975786,0.0014319627080112696,0.00014321375056169927,0.9958769679069519,7.854343130020425e-05,1.7474278138251975e-05,0.0004673725343309343],[0.04543733596801758,0.18676896393299103,0.6103681325912476,0.002208962570875883,0.14834395051002502,0.0028436484280973673,8.011127647478133e-05,0.003949080128222704],[0.000407856801757589,0.00012212026922497898,2.238086744910106e-05,2.8089953048038296e-05,0.9993014931678772,3.0249711926444434e-05,5.189376315684058e-05,3.591752101783641e-05],[0.020307833328843117,0.0180380679666996,0.946365475654602,0.0024651954881846905,0.00954336579889059,0.00018066668417304754,0.00018273454043082893,0.002917018486186862],[0.000651473761536181,0.00218518846668303,0.0007920117932371795,0.0004919193452224135,0.00036502169677987695,0.0014075201470404863,0.993802011013031,0.0003049311926588416],[0.00040118215838447213,0.001543401274830103,0.0003487435169517994,6.373901851475239e-05,0.0019522365182638168,0.9482523202896118,0.04738938808441162,4.917227488476783e-05],[0.019063137471675873,0.02085445262491703,0.004997895564883947,0.01302140299230814,0.8924995064735413,0.021639592945575714,0.004554110113531351,0.023370057344436646],[0.03222385793924332,0.2542784810066223,0.12594890594482422,0.01861872524023056,0.09022998064756393,0.4255065619945526,0.04021997004747391,0.01297358050942421],[0.10262304544448853,0.009719982743263245,0.005174960941076279,0.00107120955362916,0.8768436312675476,0.0006661756779067218,0.0023837503977119923,0.0015173490392044187],[0.3586027920246124,0.10367269068956375,0.4815979599952698,0.009928400628268719,0.03084094077348709,0.0017966675804927945,0.0004982450045645237,0.013063127174973488],[0.0005793988239020109,0.0002459466049913317,0.0002185478515457362,0.000270027230726555,0.0013606586726382375,0.9937530755996704,0.002253646496683359,0.0013188573066145182],[0.04938633367419243,0.02634953148663044,0.018402474001049995,0.01031226385384798,0.7248654961585999,0.03882462903857231,0.11274494975805283,0.019114237278699875],[0.0002476288937032223,0.00014439245569519699,9.689645958133042e-05,5.863437763764523e-05,2.6521552172198426e-06,0.0002103250881191343,0.9992278218269348,1.1652907232928555e-05],[0.03885957598686218,0.2546241879463196,0.04334506019949913,0.04490484297275543,0.2025327831506729,0.3134497106075287,0.09001661092042923,0.012267201207578182],[0.004141136072576046,0.004385257605463266,0.010346543043851852,0.06446536630392075,0.0029017720371484756,0.007811599411070347,0.007348623126745224,0.8985996842384338],[0.04773643985390663,0.09959903359413147,0.10708749294281006,0.704691469669342,0.0067106555216014385,0.003139645094051957,0.011323846876621246,0.019711676985025406],[0.4754627048969269,0.02873150072991848,0.06368735432624817,0.032147862017154694,0.3862423598766327,0.0028335172683000565,0.002544278046116233,0.008349926210939884],[0.19811204075813293,0.15874968469142914,0.16583722829818726,0.03626523166894913,0.1950610876083374,0.12608762085437775,0.09477539360523224,0.025111518800258636],[0.04672913998365402,0.0735243707895279,0.057244788855314255,0.015802543610334396,0.31973886489868164,0.26840636134147644,0.19851729273796082,0.020036181434988976],[0.0169257503002882,0.0010254014050588012,0.010664409026503563,0.0015906929038465023,0.9408483505249023,0.010732650756835938,0.010102471336722374,0.00811032485216856],[0.004795885179191828,0.013731338083744049,0.008551940321922302,0.0017096914816647768,0.049973804503679276,0.9123013019561768,0.008360871113836765,0.0005751546705141664],[0.007474340032786131,0.0005488938186317682,0.00526469387114048,0.00045276148011907935,0.9615497589111328,0.0156177943572402,0.008211473003029823,0.0008801835356280208]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Bracket: %{x}<br>y: %{y}<br>Conditional Probability: %{z}<extra></extra>\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Bracket\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\"},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Conditional Probability\"}},\"colorscale\":[[0.0,\"rgb(247,251,255)\"],[0.125,\"rgb(222,235,247)\"],[0.25,\"rgb(198,219,239)\"],[0.375,\"rgb(158,202,225)\"],[0.5,\"rgb(107,174,214)\"],[0.625,\"rgb(66,146,198)\"],[0.75,\"rgb(33,113,181)\"],[0.875,\"rgb(8,81,156)\"],[1.0,\"rgb(8,48,107)\"]]},\"margin\":{\"t\":60},\"height\":1080,\"annotations\":[{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":35},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":16},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":17},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":18},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":19},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":20},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":21},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":22},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":23},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":24},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":25},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":26},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":27},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":28},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":29},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":30},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":31},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":32},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":33},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":34},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":35}]},                        {\"responsive\": true}                    ).then(function(){\n",
       "                            \n",
       "var gd = document.getElementById('790ebab3-28bd-4ed7-95ee-cf00a1625414');\n",
       "var x = new MutationObserver(function (mutations, observer) {{\n",
       "        var display = window.getComputedStyle(gd).display;\n",
       "        if (!display || display === 'none') {{\n",
       "            console.log([gd, 'removed!']);\n",
       "            Plotly.purge(gd);\n",
       "            observer.disconnect();\n",
       "        }}\n",
       "}});\n",
       "\n",
       "// Listen for the removal of the full notebook cells\n",
       "var notebookContainer = gd.closest('#notebook-container');\n",
       "if (notebookContainer) {{\n",
       "    x.observe(notebookContainer, {childList: true});\n",
       "}}\n",
       "\n",
       "// Listen for the clearing of the current output cell\n",
       "var outputEl = gd.closest('.output');\n",
       "if (outputEl) {{\n",
       "    x.observe(outputEl, {childList: true});\n",
       "}}\n",
       "\n",
       "                        })                };                            </script>        </div>\n",
       "</body>\n",
       "</html>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_bracket_scores(exploratory_prompts[num_exploratory_prompts // 2:], bracket_scores[num_exploratory_prompts // 2:, :])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discussion\n",
    "\n",
    "- The model seems to do very well at predicting the correct closing bracket, and is robust to most things I've thrown at it.\n",
    "- One important case which occurs a few times though is when I'm looking to have the model predict `)` but it actually predicts `[`, with `)` being the second-most likely next bracket. I investigate this a little more below.\n",
    "- The fact that the model predicts `(` on `list())))).append(x` indicates that it is not confused by lots of closing brackets.\n",
    "- The model struggles a bit on the last prompts made purely of brackets and commas.\n",
    "- When the brackets are balanced, the model outputs vary a lot. Usually it predicts an opening bracket, though often spreading the probability over several types. Other times the probability is spread over both opening and closing brackets. And sometimes it predicts a closing bracket. \n",
    "    * It's not entirely clear what the model *should* predict in these cases. Oftentimes any kind of bracket would be inappropriate."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's think about the case where the model predicts `[` instead of `)`. I wouldn't say this is always wrong: it would be plausible to see a `[` in most of these cases, as a way of indexing some object, though usually this would be a bit weird.\n",
    "\n",
    "This observation raises a larger question about what exactly the task *is* and what the metric should be. The intuition is that a good model should be able to keep track of the open and closed brackets, and should prefer generating text which is *eventually* bracket-balanced. However, in the shorter term this may involve opening new brackets (after all, we wouldn't want to the model to be biased towards immediately closing all brackets it creates). I can think of the following ways of approaching this.\n",
    "1. The most direct way is to simply let the model continue to generate tokens, with the aim of seeing if the whole generated text is bracket-balanced and from there trying to understand how the model has done this. This would be a substantial undertaking, and beyond the scope of this small exploration.\n",
    "2. Another option is to focus only on the bracket type we care about. In this case, we'd only compare the prediction for `(` with `)`, and ignore the comparison with `[` and `]`. Of course there may still be instances where opening with `(` is a reasonable choice for the model to make so this doesn't completely eliminate the problem.\n",
    "3. The simplest way is to focus on clear-cut examples, where the only reasonable bracket is a closing one. I will go with this direction here, since it isolates more cleanly exactly what we want to investigate, which hopefully also makes the model behaviour more evident."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hypotheses\n",
    "\n",
    "How is GPT-Neo able to determine whether to close a bracket? Before I get my hands dirty with the model weights, I'm going briefly elaborate my thoughts for what might be going on."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- A basic component for this capability might be a 'bracket-counting' head. In this head certain tokens (perhaps those where some kind of bracket is likely as the next token, or just all tokens) attend to all the previous brackets. The opening bracket value vectors lie in an opposite direction to the closing bracket value vectors. This way, when we take the weighted sum, its projection onto the line determined by these opposing directions counts the value:\n",
    "```\n",
    "    (Number of opening brackets) - (Number of closing brackets)\n",
    "```\n",
    "- A simple way the transformer could use a bracket-counting head is by predicting an closing bracket if this number is positive and it is likely that the next token is some kind of bracket.\n",
    "- Such a simple head doesn't explain:\n",
    "    1. Why the model doesn't get confused by `list())))).append(x` (note that at `x` this count will be negative).\n",
    "    2. How the model can determine *which* bracket is appropriate.\n",
    "- Intuitively, in order to the deal with the first problem, the model needs some way of 'resetting' the count when it encounters the second `(`.\n",
    "- Here is one way this could be accomplished. There is a second bracket-counting head on a later layer, which works the same way except for the following modification. Any opening bracket which has a negative count from the first bracket-counting head gets the value vector which is the normal opening-bracket vector multiplied by the negative of the bracket count, plus one. This means that opening brackets appearing after a negatively balanced string reset the count, and counting can proceed as normal.\n",
    "- I can't think of a way to simplify this to a single head. Intuitively, the head which determines the final count already needs to have access to the bracket count computation, in order to determine when to reset. Perhaps there's a way to do it which doesn't involve counting.\n",
    "- To deal with the second token, the model needs some way of keeping track of the type of the most recent unclosed bracket. I haven't thought of a way this could work."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Experimental setup\n",
    "\n",
    "Here I define the prompts which I will be testing, and the metric used to quantify model performance on them. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference prompts\n",
    "\n",
    "I choose the prompts according to following criteria.\n",
    "1. They should have the same number of tokens.\n",
    "2. The next token, if it is a bracket, should be clearly a closing one.\n",
    "3. Each should have a corrupted version, which has the same number of tokens, differs only slightly, but after which the model predicts something different (ideally an opening bracket).\n",
    "4. There should be a variety of kinds of prompt."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The motivations for these are as follows.\n",
    "1. This makes working with the next predicted token across all prompts simultaneously easier.\n",
    "2. Cases where there are more opening than closing brackets are more clear-cut.\n",
    "3. Later I would like to use activation patching as an interpretability tool. This requires a corrupted version.\n",
    "4. We want to find a mechanism by which the model robustly accomplishes the task, rather than one which might be specific to one kind of prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompts\n",
      "14 <|endoftext|>#def# line#(#data#_#t#ensor#,# rend#erer#='#four#'\n",
      "14 <|endoftext|>#model#.#dat#as#et#.#data# =# table#['#respond#er#'\n",
      "14 <|endoftext|>#{#b#(),# c#(),# d#(),# e#(),# f#(#x#)\n",
      "14 <|endoftext|>#<#template# id#='#named#_#car#riage#'# name#='#time#'\n",
      "14 <|endoftext|>#rend#erer#(#new#_#document#,# True#).#begin#('#small#'\n",
      "14 <|endoftext|>#[#factor#(#x#_#new#),# test#(),# inspect#(#p#)\n",
      "14 <|endoftext|>#spec#ification# =# {#'#<#xml#>#':# more#(#True#)\n",
      "14 <|endoftext|>#<#html# style#='#b#.#blue# {#color#:# blue#}#'\n",
      "\n",
      "Corrupted prompts\n",
      "14 <|endoftext|>#def# line#(#data#_#t#ensor#,# rend#erer#='#four#')\n",
      "14 <|endoftext|>#model#.#dat#as#et#.#data# =# table#['#respond#er#']\n",
      "14 <|endoftext|>#{#b#(),# c#(),# d#(),# e#(),# f#(#x#)}\n",
      "14 <|endoftext|>#<#template# id#='#named#_#car#riage#'# name#='#time#'>\n",
      "14 <|endoftext|>#rend#erer#(#new#_#document#,# True#).#begin#('#small#')\n",
      "14 <|endoftext|>#[#factor#(#x#_#new#),# test#(),# inspect#(#p#)]\n",
      "14 <|endoftext|>#spec#ification# =# {#'#<#xml#>#':# more#(#True#)}\n",
      "14 <|endoftext|>#<#html# style#='#b#.#blue# {#color#:# blue#}#'>\n"
     ]
    }
   ],
   "source": [
    "# The regular prompts and their answers\n",
    "prompts = [\n",
    "    \"def line(data_tensor, renderer='four'\",\n",
    "    \"model.dataset.data = table['responder'\",\n",
    "    \"{b(), c(), d(), e(), f(x)\",\n",
    "    \"<template id='named_carriage' name='time'\",\n",
    "    \"renderer(new_document, True).begin('small'\",\n",
    "    \"[factor(x_new), test(), inspect(p)\",\n",
    "    \"specification = {'<xml>': more(True)\",\n",
    "    \"<html style='b.blue {color: blue}'\",\n",
    "]\n",
    "answers_openness = [0] * len(prompts) # 1 if opening bracket\n",
    "answer_symbols = list(\")]}>)]}>\")\n",
    "\n",
    "# The corrupted prompts and their answers\n",
    "# Note: there aren't clear answers to what the exact symbol should be\n",
    "corrupted_prompts = [\n",
    "    \"def line(data_tensor, renderer='four')\",\n",
    "    \"model.dataset.data = table['responder']\",\n",
    "    \"{b(), c(), d(), e(), f(x)}\",\n",
    "    \"<template id='named_carriage' name='time'>\",\n",
    "    \"renderer(new_document, True).begin('small')\",\n",
    "    \"[factor(x_new), test(), inspect(p)]\",\n",
    "    \"specification = {'<xml>': more(True)}\",\n",
    "    \"<html style='b.blue {color: blue}'>\",\n",
    "]\n",
    "corrupted_answers_openness = [1] * len(corrupted_prompts) # 1 if opening bracket\n",
    "\n",
    "# Combine the non-corrupted and corrupted\n",
    "all_prompts = prompts + corrupted_prompts\n",
    "all_answers_openness = answers_openness + corrupted_answers_openness\n",
    "\n",
    "print (\"Prompts\")\n",
    "for prompt in prompts:\n",
    "    prompt_as_tokens = model.to_str_tokens(prompt)\n",
    "    print(len(prompt_as_tokens), \"#\".join(prompt_as_tokens))\n",
    "\n",
    "print()\n",
    "print (\"Corrupted prompts\")\n",
    "for prompt in corrupted_prompts:\n",
    "    prompt_as_tokens = model.to_str_tokens(prompt)\n",
    "    print(len(prompt_as_tokens), \"#\".join(prompt_as_tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "prompt_tokens: torch.Size([8, 14])\n",
      "corrupted_prompt_tokens: torch.Size([8, 14])\n",
      "all_prompt_tokens: torch.Size([16, 14])\n"
     ]
    }
   ],
   "source": [
    "# Convert all the prompts to tokens, padding to make them the same length\n",
    "all_prompt_tokens = model.to_tokens(all_prompts).to(device)\n",
    "prompt_tokens = all_prompt_tokens[:len(prompts), :]\n",
    "corrupted_prompt_tokens = all_prompt_tokens[len(prompts):, :]\n",
    "\n",
    "print(\"prompt_tokens:\", prompt_tokens.shape)\n",
    "print(\"corrupted_prompt_tokens:\", corrupted_prompt_tokens.shape)\n",
    "print(\"all_prompt_tokens:\", all_prompt_tokens.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualise the model performance on the prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<html>\n",
       "<head><meta charset=\"utf-8\" /></head>\n",
       "<body>\n",
       "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
       "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.20.0.min.js\"></script>                <div id=\"ecbcf89c-e81c-4099-96bf-024c5e999634\" class=\"plotly-graph-div\" style=\"height:640px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"ecbcf89c-e81c-4099-96bf-024c5e999634\")) {                    Plotly.newPlot(                        \"ecbcf89c-e81c-4099-96bf-024c5e999634\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"(\",\"[\",\"{\",\"<\",\")\",\"]\",\"}\",\">\"],\"y\":[\"def line(data_tensor, renderer='four'\",\"model.dataset.data = table['responder'\",\"{b(), c(), d(), e(), f(x)\",\"<template id='named_carriage' name='time'\",\"renderer(new_document, True).begin('small'\",\"[factor(x_new), test(), inspect(p)\",\"specification = {'<xml>': more(True)\",\"<html style='b.blue {color: blue}'\",\"def line(data_tensor, renderer='four')\",\"model.dataset.data = table['responder']\",\"{b(), c(), d(), e(), f(x)}\",\"<template id='named_carriage' name='time'>\",\"renderer(new_document, True).begin('small')\",\"[factor(x_new), test(), inspect(p)]\",\"specification = {'<xml>': more(True)}\",\"<html style='b.blue {color: blue}'>\"],\"z\":[[4.0845909097697586e-05,0.000583446875680238,1.2041942682117224e-05,8.239389899244998e-06,0.9991393089294434,2.1074263258924475e-06,3.2315720090991817e-06,0.00021068696514703333],[1.1182510206708685e-05,1.6669526303303428e-05,1.987591531360522e-06,2.4423228751402348e-05,1.910684659378603e-06,0.9997237324714661,0.00019499423797242343,2.5355644538649358e-05],[0.002138939220458269,0.0007400010945275426,0.0054753185249865055,0.019648006185889244,0.00027858209796249866,0.0003157307510264218,0.9585424661636353,0.012861157767474651],[0.0010331079829484224,0.004348312970250845,0.004267917014658451,0.0045164539478719234,0.0006564196082763374,0.00045295903692021966,0.00025719363475218415,0.9844675660133362],[0.006689833011478186,0.0032863502856343985,0.0007558726356364787,0.0017512092599645257,0.9740374684333801,0.0019086762331426144,0.0014541835989803076,0.010116149671375751],[0.0016786952037364244,0.0022457106970250607,0.007600426208227873,0.004531071521341801,0.013487800024449825,0.960982620716095,0.001907364116050303,0.0075660827569663525],[0.0012869347119703889,0.001971310004591942,0.014029732905328274,0.003349706530570984,4.134665687161032e-06,7.910265412647277e-05,0.978911280632019,0.0003677530912682414],[0.0002617583959363401,0.000597211706917733,0.0006201667711138725,0.002838436746969819,2.3110736947273836e-05,5.83438049943652e-05,0.00028315838426351547,0.995317816734314],[0.022926494479179382,0.3047207295894623,0.6372870206832886,0.0005036972579546273,0.03238394483923912,0.00036158208968117833,7.979639485711232e-05,0.0017372311558574438],[0.10054793953895569,0.4974648356437683,0.03890956938266754,0.04471215233206749,0.19859470427036285,0.046642571687698364,0.06460931897163391,0.008518622256815434],[0.1256423443555832,0.09267544746398926,0.19473883509635925,0.12798269093036652,0.23578301072120667,0.08953504264354706,0.0554659478366375,0.07817614078521729],[0.0032398421317338943,0.005509546957910061,0.7559093236923218,0.23360438644886017,0.000465273653389886,0.00020666010095737875,0.000719101692084223,0.0003460742882452905],[0.049847111105918884,0.08086181432008743,0.21990390121936798,0.07810601592063904,0.27875152230262756,0.05998514965176582,0.2074015885591507,0.02514275349676609],[0.1636640578508377,0.009415828622877598,0.20234617590904236,0.03084702417254448,0.34799134731292725,0.1194072961807251,0.06983835250139236,0.05649033933877945],[0.018467478454113007,0.016953861340880394,0.04200029745697975,0.2677575945854187,0.10177432000637054,0.09709315001964569,0.4207594692707062,0.03519400954246521],[0.006267456337809563,0.0010978813515976071,0.14616267383098602,0.8405607342720032,0.002107215579599142,0.0001774320990080014,0.002876672660931945,0.000750187668018043]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Bracket: %{x}<br>y: %{y}<br>Conditional Probability: %{z}<extra></extra>\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Bracket\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\"},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Conditional Probability\"}},\"colorscale\":[[0.0,\"rgb(247,251,255)\"],[0.125,\"rgb(222,235,247)\"],[0.25,\"rgb(198,219,239)\"],[0.375,\"rgb(158,202,225)\"],[0.5,\"rgb(107,174,214)\"],[0.625,\"rgb(66,146,198)\"],[0.75,\"rgb(33,113,181)\"],[0.875,\"rgb(8,81,156)\"],[1.0,\"rgb(8,48,107)\"]]},\"margin\":{\"t\":60},\"height\":640,\"annotations\":[{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"(\",\"x\":0,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"[\",\"x\":1,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"{\",\"x\":2,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"<\",\"x\":3,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\")\",\"x\":4,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"]\",\"x\":5,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\"}\",\"x\":6,\"y\":15},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":0},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":1},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":2},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":3},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":4},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":5},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":6},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":7},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":8},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":9},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":10},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":11},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":12},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":13},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":14},{\"font\":{\"color\":\"orange\"},\"showarrow\":false,\"text\":\">\",\"x\":7,\"y\":15}]},                        {\"responsive\": true}                    ).then(function(){\n",
       "                            \n",
       "var gd = document.getElementById('ecbcf89c-e81c-4099-96bf-024c5e999634');\n",
       "var x = new MutationObserver(function (mutations, observer) {{\n",
       "        var display = window.getComputedStyle(gd).display;\n",
       "        if (!display || display === 'none') {{\n",
       "            console.log([gd, 'removed!']);\n",
       "            Plotly.purge(gd);\n",
       "            observer.disconnect();\n",
       "        }}\n",
       "}});\n",
       "\n",
       "// Listen for the removal of the full notebook cells\n",
       "var notebookContainer = gd.closest('#notebook-container');\n",
       "if (notebookContainer) {{\n",
       "    x.observe(notebookContainer, {childList: true});\n",
       "}}\n",
       "\n",
       "// Listen for the clearing of the current output cell\n",
       "var outputEl = gd.closest('.output');\n",
       "if (outputEl) {{\n",
       "    x.observe(outputEl, {childList: true});\n",
       "}}\n",
       "\n",
       "                        })                };                            </script>        </div>\n",
       "</body>\n",
       "</html>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bracket_scores = compute_bracket_scores(all_prompt_tokens)\n",
    "display_bracket_scores(all_prompts, bracket_scores, height_scale=40)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metrics\n",
    "\n",
    "I now define the metrics used to evaluate the model performance. The first metric measures the success at predicting openness or closedness of the bracket. The second measures in addition how well the model predicts the actual token. Both use average logit difference."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the first metric, we take the sum of the logits for open bracket tokens and find the difference with the sum of the logits for closed bracket tokens. When we expect the answer to be an open bracket, the metric is the first of these quantities take away the second. When we expect a closed bracket, it's the other way around."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def openness_metric(\n",
    "    logits: torch.Tensor, answers_openness: list, per_prompt=False\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"Computes the average difference between the open and closed logits\"\"\"\n",
    "\n",
    "    # Turn the answer openness into a sign tensor\n",
    "    answers_openness = torch.tensor(answers_openness, device=device)\n",
    "    answers_openness_sign = torch.sign(2 * answers_openness - 1)\n",
    "\n",
    "    # Select the final open and closed bracket logits\n",
    "    open_bracket_logits = logits[:, -1, open_bracket_tokens]\n",
    "    closed_bracket_logits = logits[:, -1, closed_bracket_tokens]\n",
    "\n",
    "    # Sum up the logits for open and closed brackets\n",
    "    open_bracket_logits_sum = open_bracket_logits.sum(dim=-1)\n",
    "    closed_bracket_logits_sum = closed_bracket_logits.sum(dim=-1)\n",
    "\n",
    "    # Compute the difference signed by the answer openness\n",
    "    bracket_logit_diff = open_bracket_logits_sum - closed_bracket_logits_sum\n",
    "    bracket_logit_diff = bracket_logit_diff * answers_openness_sign\n",
    "\n",
    "    if per_prompt:\n",
    "        return bracket_logit_diff\n",
    "    else:\n",
    "        return bracket_logit_diff.mean()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test on the reference prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 101.9118,  332.5992,    6.9850,  223.4595,  227.4698,  295.8047,\n",
      "        -162.3903,  132.8040])\n"
     ]
    }
   ],
   "source": [
    "logits = model(prompt_tokens, return_type=\"logits\")\n",
    "openness_metrics = openness_metric(logits, answers_openness, per_prompt=True)\n",
    "\n",
    "print(openness_metrics)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's a fair bit of variance in the metric for the reference prompts, even though the conditional probabilities all clearly favour one option. This is because:\n",
    "1. We're looking at logits not probabilities (i.e. they are not 'normalised' by the softmax).\n",
    "2. Earlier we consider the conditional probability, which has to sum to one over all brackets. It could be that the model predicts a non-bracket higher than any bracket."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next I define the metric for how well the model predicts the correct bracket. There are several ways of doing this. Here I take the sum of the logits corresponding to the correct bracket, and take away the mean of the sum of the logits for the rest of the brackets.\n",
    "\n",
    "The motivation for this is as follows. We want the metric to be linear in the logits, because this makes later analysis easier. During training the optimiser tries to minimise the cross entropy loss of the softmax of the logits. If $\\{x_i\\}$ is the set of all logits, and $x_{\\text{true}}$ is the logit for the true next token, this corresponds to maximising:\n",
    "$$\n",
    "    x_{\\text{true}} - \\log \\left(\\sum_i \\exp(x_i) \\right)\n",
    "$$\n",
    "If we want to focus on just getting the correct bracket, we can see this as maximising:\n",
    "$$\n",
    "    x_{\\text{true}} - \\log \\left(\\sum_{i \\in B} \\exp(x_i) \\right)\n",
    "$$\n",
    "where $B$ is the set of logits corresponding to brackets.\n",
    "\n",
    "How do approximate this with a linear function? In general, logsumexp is not very linear, but approximating it with the mean seems ok for the purposes of making a metric.\n",
    "\n",
    "Now the above is a bit of a simplification, since there isn't a 'true' next bracket token, because there are many tokens whose string representation starts with the same token. So instead we combine all logits whose tokens begin with the same bracket, and thus arrive at our metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bracket_symbol_metric(\n",
    "    logits: torch.Tensor, answer_symbols: list, per_prompt=False\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"Compute the difference from the answer bracket logit to all others\"\"\"\n",
    "\n",
    "    batch_size = logits.shape[0]\n",
    "\n",
    "    # Sum the logits corresponding to each bracket\n",
    "    sum_per_bracket = torch.zeros((batch_size, num_brackets))\n",
    "    for i, tokens in enumerate(bracket_tokens.values()):\n",
    "        sum_per_bracket[:, i] = logits[:, -1, tokens].sum(dim=-1)\n",
    "\n",
    "    # Turn the answers_symbol list into a tensor for indexing `sum_per_bracket`\n",
    "    answer_symbol_indices = [brackets_flat.index(bracket) for bracket in answer_symbols]\n",
    "    answer_symbol_indices = torch.tensor(answer_symbol_indices)\n",
    "    answer_symbol_indices = answer_symbol_indices.reshape((batch_size, 1))\n",
    "\n",
    "    # Compute the logits difference from the answer to the sum of the other\n",
    "    # brackets\n",
    "    answer_logits = sum_per_bracket.gather(dim=-1, index=answer_symbol_indices)\n",
    "    answer_logits = answer_logits.squeeze()\n",
    "    logit_diff = 2 * answer_logits - sum_per_bracket.mean(dim=-1)\n",
    "\n",
    "    if per_prompt:\n",
    "        return logit_diff\n",
    "    else:\n",
    "        return logit_diff.mean()\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check this metric the with reference prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([585.3561, 648.4344, 364.9975, 245.7553, 678.2637, 472.1611, 374.2183,\n",
       "        320.5639])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bracket_symbol_metric(logits, answer_symbols, per_prompt=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's compare with some random incorrect answers, to make sure the metric is doing what we want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-36.8596,  89.3942, 149.0051, 199.3244, -38.3681, 345.0772, 142.3415,\n",
       "         77.5222])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bracket_symbol_metric(logits, list(\"<[>){)[]\"), per_prompt=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Direct logit attribution\n",
    "\n",
    "In this section I investigate the model using the 'direct logit attribution' method, which looks at how different parts of the model directly affect the output logits.\n",
    "\n",
    "Much of this section is copied directly from the [Exploratory Analysis notebook](https://neelnanda.io/exploratory-analysis-demo). See the 'Direct Logit Attribution' section in that notebook for more details on the techniques."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a first pass, I will make the following simplifications.\n",
    "1. I will focus solely on the task determining if the next bracket should be opening or closing.\n",
    "2. Rather than comparing all tokens beginning with a bracket across all bracket types, I will fix a bracket type per prompt and compare only the tokens corresponding to the opening and closing versions. This is to be able to talk about residual directions, looking at the logit difference between the two possible tokens."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first add some wrong answers then tokenise everything."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "answer_wrong_symbols = list(\"([{<([{<\")\n",
    "\n",
    "# Tokenise everything\n",
    "answer_tokens = [model.to_single_token(b) for b in answer_symbols]\n",
    "answer_wrong_tokens = [model.to_single_token(b) for b in answer_wrong_symbols]\n",
    "answer_tokens = torch.tensor(answer_tokens)\n",
    "answer_wrong_tokens = torch.tensor(answer_wrong_tokens)\n",
    "answer_both_tokens = torch.stack((answer_tokens, answer_wrong_tokens)).T"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's run the model and cache the intermediate activations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "original_logits, cache = model.run_with_cache(prompt_tokens)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now compute the directions in the residual stream corresponding to moving from the wrong answer to the right one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer residual directions shape: torch.Size([8, 2, 768])\n",
      "Logit difference directions shape: torch.Size([8, 768])\n"
     ]
    }
   ],
   "source": [
    "answer_residual_directions = model.tokens_to_residual_directions(answer_both_tokens)\n",
    "print(\"Answer residual directions shape:\", answer_residual_directions.shape)\n",
    "logit_diff_directions = (\n",
    "    answer_residual_directions[:, 0] - answer_residual_directions[:, 1]\n",
    ")\n",
    "print(\"Logit difference directions shape:\", logit_diff_directions.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's test to see if this works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final residual stream shape: torch.Size([8, 14, 768])\n",
      "Calculated average logit diff: 3.5070555210113525\n",
      "Original logit difference: 3.8240890502929688\n"
     ]
    }
   ],
   "source": [
    "# cache syntax - resid_post is the residual stream at the end of the layer, -1\n",
    "# gets the final layer. The general syntax is [activation_name, layer_index,\n",
    "# sub_layer_type].\n",
    "final_residual_stream = cache[\"resid_post\", -1]\n",
    "print(\"Final residual stream shape:\", final_residual_stream.shape)\n",
    "final_token_residual_stream = final_residual_stream[:, -1, :]\n",
    "# Apply LayerNorm scaling. pos_slice is the subset of the positions we take -\n",
    "# here the final token of each prompt\n",
    "scaled_final_token_residual_stream = cache.apply_ln_to_stack(\n",
    "    final_token_residual_stream, layer=-1, pos_slice=-1\n",
    ")\n",
    "\n",
    "# Get the original logit difference\n",
    "final_logits = original_logits[:, -1, :]\n",
    "answer_logits = final_logits.gather(dim=-1, index=answer_both_tokens.to(device))\n",
    "original_average_logit_diff = answer_logits[:, 0] - answer_logits[:, 1]\n",
    "original_average_logit_diff = original_average_logit_diff.mean()\n",
    "\n",
    "average_logit_diff = einsum(\n",
    "    \"batch d_model, batch d_model -> \",\n",
    "    scaled_final_token_residual_stream,\n",
    "    logit_diff_directions,\n",
    ") / len(prompts)\n",
    "print(\"Calculated average logit diff:\", average_logit_diff.item())\n",
    "print(\"Original logit difference:\", original_average_logit_diff.item())\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logit lens"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we look at the residual stream after each layer calculate the logit difference from that. This gives an idea of when the model starts being able to do the task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def residual_stack_to_logit_diff(\n",
    "    residual_stack: Float[torch.Tensor, \"components batch d_model\"],\n",
    "    cache: ActivationCache,\n",
    ") -> float:\n",
    "    scaled_residual_stack = cache.apply_ln_to_stack(\n",
    "        residual_stack, layer=-1, pos_slice=-1\n",
    "    )\n",
    "    return einsum(\n",
    "        \"... batch d_model, batch d_model -> ...\",\n",
    "        scaled_residual_stack,\n",
    "        logit_diff_directions,\n",
    "    ) / len(prompts)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([25, 8, 768])\n",
      "torch.Size([25])\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<html>\n",
       "<head><meta charset=\"utf-8\" /></head>\n",
       "<body>\n",
       "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
       "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.20.0.min.js\"></script>                <div id=\"827e33d7-d9f8-4310-9a68-72aa82ac3e91\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"827e33d7-d9f8-4310-9a68-72aa82ac3e91\")) {                    Plotly.newPlot(                        \"827e33d7-d9f8-4310-9a68-72aa82ac3e91\",                        [{\"hovertemplate\":\"<b>%{hovertext}</b><br><br>x=%{x}<br>y=%{y}<extra></extra>\",\"hovertext\":[\"0_pre\",\"0_mid\",\"1_pre\",\"1_mid\",\"2_pre\",\"2_mid\",\"3_pre\",\"3_mid\",\"4_pre\",\"4_mid\",\"5_pre\",\"5_mid\",\"6_pre\",\"6_mid\",\"7_pre\",\"7_mid\",\"8_pre\",\"8_mid\",\"9_pre\",\"9_mid\",\"10_pre\",\"10_mid\",\"11_pre\",\"11_mid\",\"final_post\"],\"legendgroup\":\"\",\"line\":{\"color\":\"#636efa\",\"dash\":\"solid\"},\"marker\":{\"symbol\":\"circle\"},\"mode\":\"lines\",\"name\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"x\":[0.0,0.5,1.0,1.5,2.0,2.5,3.0,3.5,4.0,4.5,5.0,5.5,6.0,6.5,7.0,7.5,8.0,8.5,9.0,9.5,10.0,10.5,11.0,11.5,12.0],\"xaxis\":\"x\",\"y\":[0.1326475888490677,-1.0357563495635986,-2.0593316555023193,-2.940925121307373,-3.6660237312316895,-7.584017753601074,-7.596715927124023,-6.966047286987305,-7.235817909240723,-6.579507350921631,-6.8386125564575195,1.0659483671188354,0.37148529291152954,0.8216102123260498,0.8474677801132202,1.4055311679840088,1.328263759613037,1.4906032085418701,3.0765366554260254,6.096038818359375,5.654181003570557,4.404540061950684,5.445803165435791,3.6348345279693604,3.5070555210113525],\"yaxis\":\"y\",\"type\":\"scatter\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"x\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"y\"}},\"legend\":{\"tracegroupgap\":0},\"title\":{\"text\":\"Logit Difference From Accumulate Residual Stream\"}},                        {\"responsive\": true}                    ).then(function(){\n",
       "                            \n",
       "var gd = document.getElementById('827e33d7-d9f8-4310-9a68-72aa82ac3e91');\n",
       "var x = new MutationObserver(function (mutations, observer) {{\n",
       "        var display = window.getComputedStyle(gd).display;\n",
       "        if (!display || display === 'none') {{\n",
       "            console.log([gd, 'removed!']);\n",
       "            Plotly.purge(gd);\n",
       "            observer.disconnect();\n",
       "        }}\n",
       "}});\n",
       "\n",
       "// Listen for the removal of the full notebook cells\n",
       "var notebookContainer = gd.closest('#notebook-container');\n",
       "if (notebookContainer) {{\n",
       "    x.observe(notebookContainer, {childList: true});\n",
       "}}\n",
       "\n",
       "// Listen for the clearing of the current output cell\n",
       "var outputEl = gd.closest('.output');\n",
       "if (outputEl) {{\n",
       "    x.observe(outputEl, {childList: true});\n",
       "}}\n",
       "\n",
       "                        })                };                            </script>        </div>\n",
       "</body>\n",
       "</html>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "accumulated_residual, labels = cache.accumulated_resid(\n",
    "    layer=-1, incl_mid=True, pos_slice=-1, return_labels=True\n",
    ")\n",
    "print(accumulated_residual.shape)\n",
    "logit_lens_logit_diffs = residual_stack_to_logit_diff(accumulated_residual, cache)\n",
    "print(logit_lens_logit_diffs.shape)\n",
    "line(\n",
    "    logit_lens_logit_diffs,\n",
    "    x=np.arange(model.cfg.n_layers * 2 + 1) / 2,\n",
    "    hover_name=labels,\n",
    "    title=\"Logit Difference From Accumulate Residual Stream\",\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interestingly, for the first six layers the model actually gets worse at the task. In layer 5 the model performance jumps back up to around baseline performance, and stays there until layers 8 and 9 where it achieves best performance. Afterwards the performance decreases a little bit.\n",
    "\n",
    "Here are my provisional thoughts on what might be happening.\n",
    "- It could be that solving the task requires intermediate computation steps, and during these steps the model predicts the wrong token (at least on the prompts on which we're testing).\n",
    "- Alternatively, the initial dip in performance might be unrelated to the task. Perhaps the model doesn't try to figure out bracket balance until the later layers. Earlier on it might be doing other things with the logit directions; in other words there's some superposition going on, and the different superposed features are computed at different stages of the model.\n",
    "- The final decrease in performance might be because the sample of prompts is not representative enough. Perhaps in order to get the best performance across all bracket matching tasks (weighted by the data distribution), the optimiser decided to reduce performance on the current set of prompts in favour of others. In other words, while we might see decreasing performance on these prompts in the last layers, on others might still be low at layer 9 and continue increasing.\n",
    "- Another possibility is that the current metric for performance is too crude: it only focuses on the difference between two tokens. Could it be for example that in the later layers the model tries to decide which of the tokens beginning with the correct bracket is most appropriate? "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer attribution\n",
    "\n",
    "We now repeat the above analysis but per layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<html>\n",
       "<head><meta charset=\"utf-8\" /></head>\n",
       "<body>\n",
       "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
       "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.20.0.min.js\"></script>                <div id=\"b74daba2-f2bb-4f79-94b5-d72e0cb168a1\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"b74daba2-f2bb-4f79-94b5-d72e0cb168a1\")) {                    Plotly.newPlot(                        \"b74daba2-f2bb-4f79-94b5-d72e0cb168a1\",                        [{\"hovertemplate\":\"<b>%{hovertext}</b><br><br>x=%{x}<br>y=%{y}<extra></extra>\",\"hovertext\":[\"embed\",\"pos_embed\",\"0_attn_out\",\"0_mlp_out\",\"1_attn_out\",\"1_mlp_out\",\"2_attn_out\",\"2_mlp_out\",\"3_attn_out\",\"3_mlp_out\",\"4_attn_out\",\"4_mlp_out\",\"5_attn_out\",\"5_mlp_out\",\"6_attn_out\",\"6_mlp_out\",\"7_attn_out\",\"7_mlp_out\",\"8_attn_out\",\"8_mlp_out\",\"9_attn_out\",\"9_mlp_out\",\"10_attn_out\",\"10_mlp_out\",\"11_attn_out\",\"11_mlp_out\"],\"legendgroup\":\"\",\"line\":{\"color\":\"#636efa\",\"dash\":\"solid\"},\"marker\":{\"symbol\":\"circle\"},\"mode\":\"lines\",\"name\":\"\",\"orientation\":\"v\",\"showlegend\":false,\"x\":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25],\"xaxis\":\"x\",\"y\":[0.1380680352449417,-0.0054205008782446384,-1.1684043407440186,-1.023576021194458,-0.8815919160842896,-0.7250982522964478,-3.9179940223693848,-0.012698769569396973,0.6306684017181396,-0.2697727680206299,0.6563105583190918,-0.25910326838493347,7.904561996459961,-0.6944642066955566,0.4501255452632904,0.025857575237751007,0.558063268661499,-0.07726666331291199,0.1623375117778778,1.5859346389770508,3.0195040702819824,-0.4418591856956482,-1.2496426105499268,1.0412615537643433,-1.8109644651412964,-0.12777948379516602],\"yaxis\":\"y\",\"type\":\"scatter\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"x\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"title\":{\"text\":\"y\"}},\"legend\":{\"tracegroupgap\":0},\"title\":{\"text\":\"Logit Difference From Each Layer\"}},                        {\"responsive\": true}                    ).then(function(){\n",
       "                            \n",
       "var gd = document.getElementById('b74daba2-f2bb-4f79-94b5-d72e0cb168a1');\n",
       "var x = new MutationObserver(function (mutations, observer) {{\n",
       "        var display = window.getComputedStyle(gd).display;\n",
       "        if (!display || display === 'none') {{\n",
       "            console.log([gd, 'removed!']);\n",
       "            Plotly.purge(gd);\n",
       "            observer.disconnect();\n",
       "        }}\n",
       "}});\n",
       "\n",
       "// Listen for the removal of the full notebook cells\n",
       "var notebookContainer = gd.closest('#notebook-container');\n",
       "if (notebookContainer) {{\n",
       "    x.observe(notebookContainer, {childList: true});\n",
       "}}\n",
       "\n",
       "// Listen for the clearing of the current output cell\n",
       "var outputEl = gd.closest('.output');\n",
       "if (outputEl) {{\n",
       "    x.observe(outputEl, {childList: true});\n",
       "}}\n",
       "\n",
       "                        })                };                            </script>        </div>\n",
       "</body>\n",
       "</html>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "per_layer_residual, labels = cache.decompose_resid(\n",
    "    layer=-1, pos_slice=-1, return_labels=True\n",
    ")\n",
    "per_layer_logit_diffs = residual_stack_to_logit_diff(per_layer_residual, cache)\n",
    "line(per_layer_logit_diffs, hover_name=labels, title=\"Logit Difference From Each Layer\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Layer yields a big improvement in performance, as observed before. We can now see that this is entirely due to the attention layer.\n",
    "- In fact, up to layer 7 almost all of the performance changes (up and down) can be attributed to attention.\n",
    "- The biggest decrease in performance comes from the layer-2 attention.\n",
    "- However, the MLP layers do play a role in later layers.\n",
    "- The ultimate gain in performance in layers 8 and 9 can be attributed to the layer-8 MLP and the layer-9 attention.\n",
    "    * A tentative conclusion might be that the main way the model solves the task is by an MLP followed by an attention layer.\n",
    "- The final decrease in performance can mostly be attributed to the attention layers."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Head attribution\n",
    "\n",
    "Let's break things down further by looking at the individual heads in the attention layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tried to stack head results when they weren't cached. Computing head results now\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<html>\n",
       "<head><meta charset=\"utf-8\" /></head>\n",
       "<body>\n",
       "    <div>            <script src=\"https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/MathJax.js?config=TeX-AMS-MML_SVG\"></script><script type=\"text/javascript\">if (window.MathJax && window.MathJax.Hub && window.MathJax.Hub.Config) {window.MathJax.Hub.Config({SVG: {font: \"STIX-Web\"}});}</script>                <script type=\"text/javascript\">window.PlotlyConfig = {MathJaxConfig: 'local'};</script>\n",
       "        <script charset=\"utf-8\" src=\"https://cdn.plot.ly/plotly-2.20.0.min.js\"></script>                <div id=\"79f3796f-dba0-40ca-9210-23de87d97122\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"79f3796f-dba0-40ca-9210-23de87d97122\")) {                    Plotly.newPlot(                        \"79f3796f-dba0-40ca-9210-23de87d97122\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-0.12933094799518585,-0.6222943067550659,-0.4824504554271698,-0.44885796308517456,-0.032292768359184265,0.19880644977092743,-0.6575671434402466,-0.25205954909324646,-0.08363711833953857,-0.803784966468811,0.3485720753669739,-0.39176297187805176],[-0.8944721817970276,-0.09146858751773834,-0.048817411065101624,-0.03221786767244339,0.14971162378787994,0.23544767498970032,-0.5945754647254944,0.048346612602472305,0.025732897222042084,-0.08966928720474243,0.04102242738008499,0.09965965896844864],[-0.00497126579284668,-0.22112548351287842,-2.1787075996398926,0.36066383123397827,-0.22513234615325928,-0.10903863608837128,-0.24295371770858765,0.20368219912052155,-1.0950706005096436,-0.2724267244338989,0.010421235114336014,0.1451311558485031],[-0.2924911379814148,0.264126181602478,-0.017994003370404243,0.1904219388961792,0.04398790001869202,-0.05595690757036209,0.005974318832159042,-0.04355607181787491,-0.05246491730213165,-0.09018054604530334,0.22531819343566895,-0.05855701118707657],[0.15699885785579681,0.0057273805141448975,0.029716406017541885,0.06105861812829971,-0.04578716307878494,-0.0630362331867218,0.11294683814048767,0.16936665773391724,-0.05473984032869339,0.11588207632303238,-0.023634813725948334,-0.10929717123508453],[0.33663809299468994,1.9193885326385498,1.1729332208633423,1.2352879047393799,0.26866450905799866,0.2698501646518707,0.1985558569431305,-0.3415093719959259,0.06970397382974625,0.8207395076751709,0.28354546427726746,0.9034391641616821],[-0.17713510990142822,-0.03257058560848236,2.4880430698394775,0.0828595831990242,-1.209618330001831,0.38764286041259766,-1.1624231338500977,0.09608856588602066,0.15122757852077484,-0.13170620799064636,-0.01975012570619583,-0.06334562599658966],[0.0437442883849144,0.27756351232528687,0.0010266434401273727,0.47497791051864624,-0.11010842025279999,0.19258910417556763,0.11238044500350952,0.47061672806739807,0.09391490370035172,0.045786526054143906,0.09935048222541809,-0.25966876745224],[-0.05819308012723923,-0.06336840987205505,-0.10717733204364777,-0.039399147033691406,-0.18735401332378387,0.15486344695091248,0.0799737200140953,0.01947592943906784,0.041385650634765625,-0.05886710435152054,0.21845388412475586,-0.17910712957382202],[-0.1365678608417511,-0.24948696792125702,-0.20118993520736694,0.07397547364234924,0.12050814926624298,-0.18652862310409546,0.011613108217716217,2.8222508430480957,-0.0074574947357177734,0.07528018951416016,0.5913159847259521,0.2526225447654724],[-0.21369516849517822,0.07717274874448776,-0.3487119972705841,-0.07537136971950531,-0.29706454277038574,0.02698347717523575,-0.17574633657932281,-0.2279108762741089,-0.33434104919433594,-0.012749925255775452,-0.43710964918136597,-0.09722991287708282],[0.15608692169189453,-0.5279218554496765,0.07175523042678833,0.012573085725307465,-1.6704082489013672,-0.10854470729827881,-0.1292581856250763,-0.17145642638206482,0.19786059856414795,-0.22796568274497986,0.11236831545829773,0.0481741800904274]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}<br>Layer: %{y}<br>color: %{z}<extra></extra>\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Logit Difference From Each Head\"}},                        {\"responsive\": true}                    ).then(function(){\n",
       "                            \n",
       "var gd = document.getElementById('79f3796f-dba0-40ca-9210-23de87d97122');\n",
       "var x = new MutationObserver(function (mutations, observer) {{\n",
       "        var display = window.getComputedStyle(gd).display;\n",
       "        if (!display || display === 'none') {{\n",
       "            console.log([gd, 'removed!']);\n",
       "            Plotly.purge(gd);\n",
       "            observer.disconnect();\n",
       "        }}\n",
       "}});\n",
       "\n",
       "// Listen for the removal of the full notebook cells\n",
       "var notebookContainer = gd.closest('#notebook-container');\n",
       "if (notebookContainer) {{\n",
       "    x.observe(notebookContainer, {childList: true});\n",
       "}}\n",
       "\n",
       "// Listen for the clearing of the current output cell\n",
       "var outputEl = gd.closest('.output');\n",
       "if (outputEl) {{\n",
       "    x.observe(outputEl, {childList: true});\n",
       "}}\n",
       "\n",
       "                        })                };                            </script>        </div>\n",
       "</body>\n",
       "</html>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "per_head_residual, labels = cache.stack_head_results(\n",
    "    layer=-1, pos_slice=-1, return_labels=True\n",
    ")\n",
    "per_head_logit_diffs = residual_stack_to_logit_diff(per_head_residual, cache)\n",
    "per_head_logit_diffs = einops.rearrange(\n",
    "    per_head_logit_diffs,\n",
    "    \"(layer head_index) -> layer head_index\",\n",
    "    layer=model.cfg.n_layers,\n",
    "    head_index=model.cfg.n_heads,\n",
    ")\n",
    "imshow(\n",
    "    per_head_logit_diffs,\n",
    "    labels={\"x\": \"Head\", \"y\": \"Layer\"},\n",
    "    title=\"Logit Difference From Each Head\",\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- There's quite a lot going on here, with most heads contributing something.\n",
    "- The layer-2 drop in performance can mostly be attributed to L2H2.\n",
    "- The layer-2 recovery can be attributed to many heads, in order: L5H1, L5H3, L5H2, L5H11, L5H9. There are other heads which contribute to a lesser degree.\n",
    "- Even though the total gain from the layer-6 attention is small, it actually has one head which contributes a lot (L6H2), which is counterbalanced by two heads which detract (L6H4 and L6H6).\n",
    "- The performance gain in layer 9 is almost all down to L9H7.\n",
    "- The main culprit for latter performance loss is L11H4."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention analysis\n",
    "\n",
    "Let's now zoom in on the important heads, and see what they're doing in terms of moving information about."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualise_attention_patterns(\n",
    "    head_indices: Union[List[int], int, Float[torch.Tensor, \"heads\"]],\n",
    "    batch_index: int = 0,\n",
    "    visualisation_type: str = \"attention_heads\",\n",
    ") -> RenderedHTML:\n",
    "    \"\"\"Visualise selected attention heads\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    head_indices: int or list or torch.Tensor\n",
    "        Indices of the heads to visualise, as integers in the range [0,\n",
    "        n_layers * n_heads)\n",
    "    batch_index: int, default=0\n",
    "        Which prompt to look at\n",
    "    visualisation_type: str, default=\"attention_heads\"\n",
    "        The type of visualisation to make. Either \"attention_heads\" or\n",
    "        \"attention_patterns\". The later can't label the heads.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    html: RenderedHTML\n",
    "        The HTML object built.\n",
    "    \"\"\"\n",
    "\n",
    "    if visualisation_type not in (\"attention_heads\", \"attention_patterns\"):\n",
    "        raise ValueError(\n",
    "            \"Parameter 'visualisation_type' must be one of 'attention_heads' \"\n",
    "            \"or 'attention_patterns'\"\n",
    "        )\n",
    "\n",
    "    if isinstance(head_indices, int):\n",
    "        head_indices = [head_indices]\n",
    "    elif isinstance(head_indices, list) or isinstance(head_indices, torch.Tensor):\n",
    "        head_indices = utils.to_numpy(head_indices)\n",
    "\n",
    "    labels = []\n",
    "    patterns = []\n",
    "    for head_index in head_indices:\n",
    "        layer = head_index // model.cfg.n_heads\n",
    "        head_index = head_index % model.cfg.n_heads\n",
    "        # Get the attention patterns for the head.\n",
    "        # Attention patterns have shape [batch, head_index, query_pos, key_pos]\n",
    "        patterns.append(cache[\"attn\", layer][batch_index, head_index])\n",
    "        labels.append(f\"L{layer}H{head_index}\")\n",
    "    str_tokens = model.to_str_tokens(prompt_tokens[batch_index])\n",
    "    patterns = torch.stack(patterns, dim=0)\n",
    "\n",
    "    # Plot the attention patterns\n",
    "    if visualisation_type == \"attention_heads\":\n",
    "        return cv.attention.attention_heads(\n",
    "            attention=patterns, tokens=str_tokens, attention_head_names=labels\n",
    "        )\n",
    "    elif visualisation_type == \"attention_patterns\":\n",
    "        return cv.attention.attention_patterns(attention=patterns, tokens=str_tokens)\n",
    "    else:\n",
    "        raise ValueError\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We look at the attention patterns of the top 3 positive and top 3 negative heads. The `circuitsvis` package has two different visualisers for attention patterns, each with their pros and cons. I use both below. We can look at different prompts by varying `batch_index`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_k = 3\n",
    "top_positive_logit_attr_heads = torch.topk(per_head_logit_diffs.flatten(), k=top_k).indices\n",
    "top_negative_logit_attr_heads = torch.topk(-per_head_logit_diffs.flatten(), k=top_k).indices\n",
    "top_logit_attr_heads = torch.cat((top_positive_logit_attr_heads, top_negative_logit_attr_heads))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The top 3 positive and top 3 negative heads. Prompt 0\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-8c6362e6-335b\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, AttentionHeads } from \"https://unpkg.com/circuitsvis@1.39.1/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-8c6362e6-335b\",\n",
       "      AttentionHeads,\n",
       "      {\"attention\": [[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.880813479423523, 0.11918650567531586, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8866000175476074, 0.041837844997644424, 0.0715620368719101, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6065509915351868, 0.010962745174765587, 0.22979700565338135, 0.15268921852111816, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.20068036019802094, 0.011121049523353577, 0.01270347274839878, 0.7394205927848816, 0.03607449308037758, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6909565329551697, 0.0030536274425685406, 0.024777865037322044, 0.15370815992355347, 0.08442375063896179, 0.04308013617992401, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2006145417690277, 0.009903253987431526, 0.038164082914590836, 0.6281787753105164, 0.08458612859249115, 0.019058970734477043, 0.019494274631142616, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.17909038066864014, 0.01399916410446167, 0.021557670086622238, 0.5900844931602478, 0.05053335800766945, 0.02113083004951477, 0.012256230227649212, 0.11134788393974304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2426910698413849, 0.004641556181013584, 0.08246645331382751, 0.05032024160027504, 0.15283311903476715, 0.025794463232159615, 0.05596958473324776, 0.30032670497894287, 0.08495678007602692, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6304879784584045, 0.0034488930832594633, 0.0019310996867716312, 0.1427067071199417, 0.008115739561617374, 0.015357663854956627, 0.012509833090007305, 0.031667932868003845, 0.14121681451797485, 0.01255732774734497, 0.0, 0.0, 0.0, 0.0], [0.12943795323371887, 0.009073421359062195, 0.0023400415666401386, 0.6571571230888367, 0.005963231436908245, 0.007471141871064901, 0.012556728906929493, 0.02838405780494213, 0.12203390151262283, 0.013591778464615345, 0.011990723200142384, 0.0, 0.0, 0.0], [0.6771817803382874, 0.0016069903504103422, 0.07132568210363388, 0.031885724514722824, 0.00990009680390358, 0.0033402314875274897, 0.003915388602763414, 0.05777247995138168, 0.011799772270023823, 0.032717153429985046, 0.01631552167236805, 0.08223918825387955, 0.0, 0.0], [0.33163970708847046, 0.002229334320873022, 0.005863015074282885, 0.1292608231306076, 0.0015864258166402578, 0.003054470755159855, 0.004276403225958347, 0.021820753812789917, 0.06715165823698044, 0.023694880306720734, 0.007324681617319584, 0.3905360698699951, 0.011561781167984009, 0.0], [0.045904893428087234, 0.01685239188373089, 0.004483601078391075, 0.7136902809143066, 0.002983863465487957, 0.0007091029547154903, 0.00306225405074656, 0.01964682899415493, 0.09558134526014328, 0.006645140238106251, 0.0034323367290198803, 0.05660869553685188, 8.052153862081468e-05, 0.030318718403577805]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9062837958335876, 0.09371623396873474, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.828139066696167, 0.13903765380382538, 0.03282339125871658, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6529006958007812, 0.19007615745067596, 0.0321105495095253, 0.12491263449192047, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6157164573669434, 0.08498001098632812, 0.05196775123476982, 0.15143358707427979, 0.09590217471122742, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.534752368927002, 0.0877794548869133, 0.040105126798152924, 0.13253909349441528, 0.11982043087482452, 0.08500358462333679, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4263680875301361, 0.0921923890709877, 0.04208358749747276, 0.12282796949148178, 0.10456105321645737, 0.08813812583684921, 0.12382875382900238, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4970209300518036, 0.05227413401007652, 0.032834243029356, 0.038823019713163376, 0.07813955098390579, 0.08493656665086746, 0.08154695481061935, 0.1344246119260788, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4189852178096771, 0.04504605755209923, 0.014605037868022919, 0.020590944215655327, 0.062069956213235855, 0.023343056440353394, 0.04156344383955002, 0.10190647095441818, 0.2718897759914398, 0.0, 0.0, 0.0, 0.0, 0.0], [0.40449240803718567, 0.027698826044797897, 0.009495067410171032, 0.02591489441692829, 0.05592276155948639, 0.05252164602279663, 0.05948542803525925, 0.06758347898721695, 0.28797072172164917, 0.008914805948734283, 0.0, 0.0, 0.0, 0.0], [0.4861564338207245, 0.019685767590999603, 0.011181261390447617, 0.0340634286403656, 0.047057006508111954, 0.038313377648591995, 0.04315846785902977, 0.037123050540685654, 0.23647885024547577, 0.009700099006295204, 0.03708221763372421, 0.0, 0.0, 0.0], [0.3521716296672821, 0.009464077651500702, 0.007490101270377636, 0.010845880955457687, 0.018343951553106308, 0.02628241293132305, 0.018422529101371765, 0.017087260261178017, 0.0654035434126854, 0.0060485126450657845, 0.021186089143157005, 0.4472539722919464, 0.0, 0.0], [0.6491422057151794, 0.017789626494050026, 0.007113546598702669, 0.03054223209619522, 0.01250661350786686, 0.020662635564804077, 0.018974920734763145, 0.02402409166097641, 0.05400443822145462, 0.012453031726181507, 0.0306521225720644, 0.09129032492637634, 0.030844250693917274, 0.0], [0.3954005539417267, 0.024430355057120323, 0.013967443257570267, 0.0171529408544302, 0.018088042736053467, 0.010854436084628105, 0.02308524027466774, 0.02904055267572403, 0.10933482646942139, 0.018913419917225838, 0.05505087599158287, 0.1602170318365097, 0.013321652077138424, 0.11114262789487839]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8271328806877136, 0.17286713421344757, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4016323685646057, 0.19811268150806427, 0.4002549350261688, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11043158918619156, 0.02414160594344139, 0.23269100487232208, 0.6327357292175293, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12678006291389465, 0.02629311941564083, 0.07796094566583633, 0.6229124665260315, 0.14605344831943512, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10085837543010712, 0.01737082563340664, 0.020829033106565475, 0.6493512988090515, 0.1368805468082428, 0.07470981776714325, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06161661073565483, 0.015364686027169228, 0.04669824242591858, 0.5932425260543823, 0.13312405347824097, 0.05430753901600838, 0.09564629197120667, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.035465970635414124, 0.010724017396569252, 0.02340884879231453, 0.3901020884513855, 0.13724792003631592, 0.19031432271003723, 0.16456522047519684, 0.04817154258489609, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.017150431871414185, 0.0024521027226001024, 0.017884371802210808, 0.7197827100753784, 0.0227514635771513, 0.008119339123368263, 0.01574013940989971, 0.0032569707836955786, 0.19286251068115234, 0.0, 0.0, 0.0, 0.0, 0.0], [0.020057933405041695, 0.007940851151943207, 0.012289399281144142, 0.44327953457832336, 0.016624685376882553, 0.009680046699941158, 0.010975202545523643, 0.008294289000332355, 0.3392893373966217, 0.13156874477863312, 0.0, 0.0, 0.0, 0.0], [0.026629742234945297, 0.001512529095634818, 0.006891750264912844, 0.32978832721710205, 0.002915039425715804, 0.003407103708013892, 0.00906258262693882, 0.0011415474582463503, 0.3987167477607727, 0.1365651786327362, 0.08336936682462692, 0.0, 0.0, 0.0], [0.0026538583915680647, 0.0003491208190098405, 0.0006576338200829923, 0.047059316188097, 0.0021381531842052937, 0.0007629544124938548, 0.0020055798813700676, 0.0006289319135248661, 0.06153493374586105, 0.024671513587236404, 0.04335427284240723, 0.8141837120056152, 0.0, 0.0], [0.0016588016878813505, 0.0009802269050851464, 0.0005441831890493631, 0.0007990115555003285, 0.0013185753487050533, 0.0005097362445667386, 0.001671497942879796, 0.0016786481719464064, 0.011077039875090122, 0.006781391799449921, 0.00901136826723814, 0.7218093276023865, 0.24216017127037048, 0.0], [0.04869362339377403, 0.002771251369267702, 0.011116416193544865, 0.07579924166202545, 0.004150938708335161, 0.0008619367145001888, 0.002991914050653577, 0.0022288537584245205, 0.24542565643787384, 0.02463546022772789, 0.030237557366490364, 0.20632369816303253, 0.0030813817866146564, 0.3416821360588074]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9136444330215454, 0.0863555446267128, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8263124823570251, 0.07250402122735977, 0.10118347406387329, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7650721669197083, 0.07281088829040527, 0.09706933051347733, 0.06504761427640915, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6993734240531921, 0.06561672687530518, 0.09219622611999512, 0.06813137233257294, 0.07468223571777344, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6717727780342102, 0.06651146709918976, 0.08390183746814728, 0.0662839412689209, 0.0708223283290863, 0.04070763662457466, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6427624225616455, 0.06329774856567383, 0.08380358666181564, 0.060755111277103424, 0.06873518228530884, 0.037290722131729126, 0.04335525631904602, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6156875491142273, 0.05519231781363487, 0.07588260620832443, 0.05121631920337677, 0.06373441964387894, 0.03430108726024628, 0.04193982109427452, 0.06204581260681152, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5806622505187988, 0.054447874426841736, 0.07623729854822159, 0.0494152270257473, 0.061665259301662445, 0.030378330498933792, 0.038375746458768845, 0.057961605489254, 0.05085642635822296, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5314738750457764, 0.04943231865763664, 0.06743114441633224, 0.0478399433195591, 0.05720377340912819, 0.033642955124378204, 0.03790438920259476, 0.055076926946640015, 0.05096087232232094, 0.06903379410505295, 0.0, 0.0, 0.0, 0.0], [0.49908578395843506, 0.04876786097884178, 0.0650462955236435, 0.04369903728365898, 0.05270344391465187, 0.029971566051244736, 0.03549615293741226, 0.051128026098012924, 0.04734424874186516, 0.06642548739910126, 0.060332104563713074, 0.0, 0.0, 0.0], [0.4636313319206238, 0.04477505385875702, 0.06128520146012306, 0.04428810253739357, 0.05117896571755409, 0.02639702893793583, 0.031872496008872986, 0.04749530553817749, 0.04682840779423714, 0.061201758682727814, 0.05287045240402222, 0.06817581504583359, 0.0, 0.0], [0.42335471510887146, 0.04269725829362869, 0.05866531282663345, 0.04051211476325989, 0.05088456720113754, 0.02540493942797184, 0.031905610114336014, 0.048182517290115356, 0.04080118611454964, 0.06125872954726219, 0.05457599088549614, 0.06249537318944931, 0.05926157906651497, 0.0], [0.42228785157203674, 0.03822512924671173, 0.056164562702178955, 0.03373495861887932, 0.044957708567380905, 0.02277299202978611, 0.0300452820956707, 0.04248708486557007, 0.04043055698275566, 0.056272298097610474, 0.050744082778692245, 0.057307954877614975, 0.05850312486290932, 0.04606639966368675]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [5.4443644330604e-06, 0.9999945163726807, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.000340827158652246, 1.522992351965513e-05, 0.9996439218521118, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.5033668887554086e-07, 1.5539581843526662e-09, 1.0828672429852304e-06, 0.9999988079071045, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.4860424016660545e-05, 3.924570137314731e-06, 4.7554352931911126e-05, 0.0029133642092347145, 0.9970203042030334, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.1183689469239653e-08, 5.399687363905059e-10, 4.442081191768921e-08, 0.004807529039680958, 1.7774240404833108e-05, 0.9951747059822083, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.2336930012679659e-05, 1.1335694694025733e-07, 1.7473822708780062e-06, 8.646557398606092e-05, 0.002848046598955989, 0.001072101527824998, 0.9959790706634521, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.4322378067154204e-06, 3.87321415473707e-06, 2.9090304451528937e-05, 0.00010180847311858088, 0.008485228754580021, 0.0004340687009971589, 0.005096187349408865, 0.9858474135398865, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [6.682396302970517e-11, 4.753054838602111e-08, 2.8264020457413608e-08, 1.485501229581132e-06, 7.235502721414377e-07, 4.88603109261021e-06, 9.3033901293893e-07, 1.1520458429004066e-05, 0.9999804496765137, 0.0, 0.0, 0.0, 0.0, 0.0], [4.509112841333263e-06, 1.4262801641962142e-06, 3.9936694520292804e-05, 6.5301328504574485e-06, 3.416110121179372e-05, 8.898381929611787e-05, 1.3975942238175776e-05, 0.0004578701627906412, 3.156268212478608e-05, 0.9993211030960083, 0.0, 0.0, 0.0, 0.0], [1.2304453775868751e-05, 8.846908713167068e-06, 1.1869107765960507e-05, 1.8301088857697323e-05, 0.0006726715946570039, 2.9870310754631646e-05, 9.385061275679618e-05, 0.0032646076288074255, 0.004700236953794956, 0.0007520715589635074, 0.9904352426528931, 0.0, 0.0, 0.0], [3.518028179883004e-08, 1.1920349256289153e-11, 3.274404969388911e-09, 2.6973480089509394e-06, 4.262755426420739e-11, 3.6191207186675456e-07, 3.968440159818343e-11, 2.6712736467260356e-09, 1.1683687350938499e-10, 6.754234060935005e-09, 1.0540309736128961e-09, 0.9999969005584717, 0.0, 0.0], [1.216761802425026e-07, 2.3296263407246443e-06, 1.0410583684006269e-07, 1.0372890528742573e-06, 2.3239222457505093e-07, 1.1366762009856757e-05, 3.483045407648433e-09, 5.737806077377172e-06, 2.5247034045605687e-06, 1.4229512999008875e-05, 3.3851924854388926e-06, 2.753440276137553e-05, 0.9999314546585083, 0.0], [0.0021027533803135157, 5.3797426517121494e-05, 2.642461004143115e-05, 5.549227716983296e-05, 0.0004349422815721482, 5.564442108152434e-05, 0.0003221925289835781, 0.0019075508462265134, 0.012244115583598614, 0.00019615399651229382, 0.005641697905957699, 7.855122021283023e-06, 5.750841592089273e-05, 0.9768938422203064]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0009072376415133476, 0.9990928173065186, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.848416806955356e-05, 0.00033618693123571575, 0.9996353387832642, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.08687606825697e-06, 2.3963496005308116e-06, 4.65728362541995e-06, 0.9999918937683105, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.0430768674705178e-05, 1.9211251128581353e-05, 4.7870503294689115e-06, 0.0004905093810521066, 0.9994551539421082, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.5894426016748184e-06, 1.5198468190646963e-06, 7.012368996583973e-08, 3.681573798530735e-05, 0.013727720826864243, 0.98623126745224, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [9.139445319306105e-05, 0.00011409086437197402, 2.19574917537102e-06, 0.00022433913545683026, 0.0043365610763430595, 0.0023445438127964735, 0.9928868412971497, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.246472988394089e-05, 5.881246579519939e-06, 8.339918053934525e-07, 3.120941255474463e-05, 0.0001950529112946242, 7.835734140826389e-05, 0.003378617577254772, 0.9962775111198425, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0792935427161865e-05, 1.9718265320989303e-05, 3.5675466278917156e-06, 0.0003161633212585002, 0.0006945938221178949, 1.7135267626144923e-05, 0.0010499282507225871, 0.0006411561626009643, 0.9972468614578247, 0.0, 0.0, 0.0, 0.0, 0.0], [1.969202912732726e-06, 9.843746511251084e-07, 7.95048526924802e-06, 1.2612697446456878e-06, 7.5254224611853715e-06, 2.0916523624237016e-08, 4.208237669445225e-07, 9.048753213392047e-07, 2.441770448058378e-06, 0.999976634979248, 0.0, 0.0, 0.0, 0.0], [3.388111508684233e-06, 1.532968667561363e-07, 8.007928045117296e-07, 6.9148195507295895e-06, 9.635970172894304e-07, 8.262756701071794e-09, 2.7631315901999187e-08, 1.2104433722015528e-07, 8.786131729721092e-06, 0.010999530553817749, 0.9889793395996094, 0.0, 0.0, 0.0], [3.4670763398025883e-06, 3.2853522498044185e-06, 1.7884923408928444e-06, 1.271074688702356e-05, 1.981398440875637e-07, 6.105980165926894e-10, 1.331544474680868e-08, 1.059634935351994e-09, 6.316852250165539e-07, 0.0005560617428272963, 0.000964729639235884, 0.9984571933746338, 0.0, 0.0], [0.0001142337714554742, 4.442620138433995e-06, 6.048322120477678e-06, 2.263612077513244e-05, 2.3252282232988364e-07, 7.76467601326658e-09, 5.912871081648063e-09, 1.0200583488995107e-07, 1.1434232192186755e-06, 7.391762665065471e-06, 4.037147149915654e-08, 3.6513690702122403e-06, 0.9998400211334229, 0.0], [0.00011470724712125957, 1.807276362342236e-06, 1.0345141845391481e-06, 6.004875831422396e-05, 3.8280772969301324e-06, 1.865960825853108e-07, 3.2268536642732215e-07, 3.506449672840972e-07, 4.888435341854347e-06, 7.464206532858952e-07, 1.698989535725559e-07, 3.743731667782413e-06, 4.9968955863732845e-05, 0.9997581839561462]]], \"attentionHeadNames\": [\"L9H7\", \"L6H2\", \"L5H1\", \"L2H2\", \"L11H4\", \"L6H4\"], \"tokens\": [\"<|endoftext|>\", \"def\", \" line\", \"(\", \"data\", \"_\", \"t\", \"ensor\", \",\", \" rend\", \"erer\", \"='\", \"four\", \"'\"]}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7f17f2d95f90>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch_index = 0\n",
    "print(f\"The top {top_k} positive and top {top_k} negative heads. Prompt {batch_index}\")\n",
    "visualise_attention_patterns(top_logit_attr_heads, batch_index=batch_index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The same heads with a different visualiser\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-5ccb13c5-c62c\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, AttentionPatterns } from \"https://unpkg.com/circuitsvis@1.39.1/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-5ccb13c5-c62c\",\n",
       "      AttentionPatterns,\n",
       "      {\"tokens\": [\"<|endoftext|>\", \"def\", \" line\", \"(\", \"data\", \"_\", \"t\", \"ensor\", \",\", \" rend\", \"erer\", \"='\", \"four\", \"'\"], \"attention\": [[[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.880813479423523, 0.11918650567531586, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8866000175476074, 0.041837844997644424, 0.0715620368719101, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6065509915351868, 0.010962745174765587, 0.22979700565338135, 0.15268921852111816, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.20068036019802094, 0.011121049523353577, 0.01270347274839878, 0.7394205927848816, 0.03607449308037758, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6909565329551697, 0.0030536274425685406, 0.024777865037322044, 0.15370815992355347, 0.08442375063896179, 0.04308013617992401, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2006145417690277, 0.009903253987431526, 0.038164082914590836, 0.6281787753105164, 0.08458612859249115, 0.019058970734477043, 0.019494274631142616, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.17909038066864014, 0.01399916410446167, 0.021557670086622238, 0.5900844931602478, 0.05053335800766945, 0.02113083004951477, 0.012256230227649212, 0.11134788393974304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.2426910698413849, 0.004641556181013584, 0.08246645331382751, 0.05032024160027504, 0.15283311903476715, 0.025794463232159615, 0.05596958473324776, 0.30032670497894287, 0.08495678007602692, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6304879784584045, 0.0034488930832594633, 0.0019310996867716312, 0.1427067071199417, 0.008115739561617374, 0.015357663854956627, 0.012509833090007305, 0.031667932868003845, 0.14121681451797485, 0.01255732774734497, 0.0, 0.0, 0.0, 0.0], [0.12943795323371887, 0.009073421359062195, 0.0023400415666401386, 0.6571571230888367, 0.005963231436908245, 0.007471141871064901, 0.012556728906929493, 0.02838405780494213, 0.12203390151262283, 0.013591778464615345, 0.011990723200142384, 0.0, 0.0, 0.0], [0.6771817803382874, 0.0016069903504103422, 0.07132568210363388, 0.031885724514722824, 0.00990009680390358, 0.0033402314875274897, 0.003915388602763414, 0.05777247995138168, 0.011799772270023823, 0.032717153429985046, 0.01631552167236805, 0.08223918825387955, 0.0, 0.0], [0.33163970708847046, 0.002229334320873022, 0.005863015074282885, 0.1292608231306076, 0.0015864258166402578, 0.003054470755159855, 0.004276403225958347, 0.021820753812789917, 0.06715165823698044, 0.023694880306720734, 0.007324681617319584, 0.3905360698699951, 0.011561781167984009, 0.0], [0.045904893428087234, 0.01685239188373089, 0.004483601078391075, 0.7136902809143066, 0.002983863465487957, 0.0007091029547154903, 0.00306225405074656, 0.01964682899415493, 0.09558134526014328, 0.006645140238106251, 0.0034323367290198803, 0.05660869553685188, 8.052153862081468e-05, 0.030318718403577805]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9062837958335876, 0.09371623396873474, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.828139066696167, 0.13903765380382538, 0.03282339125871658, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6529006958007812, 0.19007615745067596, 0.0321105495095253, 0.12491263449192047, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6157164573669434, 0.08498001098632812, 0.05196775123476982, 0.15143358707427979, 0.09590217471122742, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.534752368927002, 0.0877794548869133, 0.040105126798152924, 0.13253909349441528, 0.11982043087482452, 0.08500358462333679, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4263680875301361, 0.0921923890709877, 0.04208358749747276, 0.12282796949148178, 0.10456105321645737, 0.08813812583684921, 0.12382875382900238, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4970209300518036, 0.05227413401007652, 0.032834243029356, 0.038823019713163376, 0.07813955098390579, 0.08493656665086746, 0.08154695481061935, 0.1344246119260788, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4189852178096771, 0.04504605755209923, 0.014605037868022919, 0.020590944215655327, 0.062069956213235855, 0.023343056440353394, 0.04156344383955002, 0.10190647095441818, 0.2718897759914398, 0.0, 0.0, 0.0, 0.0, 0.0], [0.40449240803718567, 0.027698826044797897, 0.009495067410171032, 0.02591489441692829, 0.05592276155948639, 0.05252164602279663, 0.05948542803525925, 0.06758347898721695, 0.28797072172164917, 0.008914805948734283, 0.0, 0.0, 0.0, 0.0], [0.4861564338207245, 0.019685767590999603, 0.011181261390447617, 0.0340634286403656, 0.047057006508111954, 0.038313377648591995, 0.04315846785902977, 0.037123050540685654, 0.23647885024547577, 0.009700099006295204, 0.03708221763372421, 0.0, 0.0, 0.0], [0.3521716296672821, 0.009464077651500702, 0.007490101270377636, 0.010845880955457687, 0.018343951553106308, 0.02628241293132305, 0.018422529101371765, 0.017087260261178017, 0.0654035434126854, 0.0060485126450657845, 0.021186089143157005, 0.4472539722919464, 0.0, 0.0], [0.6491422057151794, 0.017789626494050026, 0.007113546598702669, 0.03054223209619522, 0.01250661350786686, 0.020662635564804077, 0.018974920734763145, 0.02402409166097641, 0.05400443822145462, 0.012453031726181507, 0.0306521225720644, 0.09129032492637634, 0.030844250693917274, 0.0], [0.3954005539417267, 0.024430355057120323, 0.013967443257570267, 0.0171529408544302, 0.018088042736053467, 0.010854436084628105, 0.02308524027466774, 0.02904055267572403, 0.10933482646942139, 0.018913419917225838, 0.05505087599158287, 0.1602170318365097, 0.013321652077138424, 0.11114262789487839]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8271328806877136, 0.17286713421344757, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.4016323685646057, 0.19811268150806427, 0.4002549350261688, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.11043158918619156, 0.02414160594344139, 0.23269100487232208, 0.6327357292175293, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.12678006291389465, 0.02629311941564083, 0.07796094566583633, 0.6229124665260315, 0.14605344831943512, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.10085837543010712, 0.01737082563340664, 0.020829033106565475, 0.6493512988090515, 0.1368805468082428, 0.07470981776714325, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.06161661073565483, 0.015364686027169228, 0.04669824242591858, 0.5932425260543823, 0.13312405347824097, 0.05430753901600838, 0.09564629197120667, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.035465970635414124, 0.010724017396569252, 0.02340884879231453, 0.3901020884513855, 0.13724792003631592, 0.19031432271003723, 0.16456522047519684, 0.04817154258489609, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.017150431871414185, 0.0024521027226001024, 0.017884371802210808, 0.7197827100753784, 0.0227514635771513, 0.008119339123368263, 0.01574013940989971, 0.0032569707836955786, 0.19286251068115234, 0.0, 0.0, 0.0, 0.0, 0.0], [0.020057933405041695, 0.007940851151943207, 0.012289399281144142, 0.44327953457832336, 0.016624685376882553, 0.009680046699941158, 0.010975202545523643, 0.008294289000332355, 0.3392893373966217, 0.13156874477863312, 0.0, 0.0, 0.0, 0.0], [0.026629742234945297, 0.001512529095634818, 0.006891750264912844, 0.32978832721710205, 0.002915039425715804, 0.003407103708013892, 0.00906258262693882, 0.0011415474582463503, 0.3987167477607727, 0.1365651786327362, 0.08336936682462692, 0.0, 0.0, 0.0], [0.0026538583915680647, 0.0003491208190098405, 0.0006576338200829923, 0.047059316188097, 0.0021381531842052937, 0.0007629544124938548, 0.0020055798813700676, 0.0006289319135248661, 0.06153493374586105, 0.024671513587236404, 0.04335427284240723, 0.8141837120056152, 0.0, 0.0], [0.0016588016878813505, 0.0009802269050851464, 0.0005441831890493631, 0.0007990115555003285, 0.0013185753487050533, 0.0005097362445667386, 0.001671497942879796, 0.0016786481719464064, 0.011077039875090122, 0.006781391799449921, 0.00901136826723814, 0.7218093276023865, 0.24216017127037048, 0.0], [0.04869362339377403, 0.002771251369267702, 0.011116416193544865, 0.07579924166202545, 0.004150938708335161, 0.0008619367145001888, 0.002991914050653577, 0.0022288537584245205, 0.24542565643787384, 0.02463546022772789, 0.030237557366490364, 0.20632369816303253, 0.0030813817866146564, 0.3416821360588074]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.9136444330215454, 0.0863555446267128, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.8263124823570251, 0.07250402122735977, 0.10118347406387329, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.7650721669197083, 0.07281088829040527, 0.09706933051347733, 0.06504761427640915, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6993734240531921, 0.06561672687530518, 0.09219622611999512, 0.06813137233257294, 0.07468223571777344, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6717727780342102, 0.06651146709918976, 0.08390183746814728, 0.0662839412689209, 0.0708223283290863, 0.04070763662457466, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6427624225616455, 0.06329774856567383, 0.08380358666181564, 0.060755111277103424, 0.06873518228530884, 0.037290722131729126, 0.04335525631904602, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.6156875491142273, 0.05519231781363487, 0.07588260620832443, 0.05121631920337677, 0.06373441964387894, 0.03430108726024628, 0.04193982109427452, 0.06204581260681152, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5806622505187988, 0.054447874426841736, 0.07623729854822159, 0.0494152270257473, 0.061665259301662445, 0.030378330498933792, 0.038375746458768845, 0.057961605489254, 0.05085642635822296, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5314738750457764, 0.04943231865763664, 0.06743114441633224, 0.0478399433195591, 0.05720377340912819, 0.033642955124378204, 0.03790438920259476, 0.055076926946640015, 0.05096087232232094, 0.06903379410505295, 0.0, 0.0, 0.0, 0.0], [0.49908578395843506, 0.04876786097884178, 0.0650462955236435, 0.04369903728365898, 0.05270344391465187, 0.029971566051244736, 0.03549615293741226, 0.051128026098012924, 0.04734424874186516, 0.06642548739910126, 0.060332104563713074, 0.0, 0.0, 0.0], [0.4636313319206238, 0.04477505385875702, 0.06128520146012306, 0.04428810253739357, 0.05117896571755409, 0.02639702893793583, 0.031872496008872986, 0.04749530553817749, 0.04682840779423714, 0.061201758682727814, 0.05287045240402222, 0.06817581504583359, 0.0, 0.0], [0.42335471510887146, 0.04269725829362869, 0.05866531282663345, 0.04051211476325989, 0.05088456720113754, 0.02540493942797184, 0.031905610114336014, 0.048182517290115356, 0.04080118611454964, 0.06125872954726219, 0.05457599088549614, 0.06249537318944931, 0.05926157906651497, 0.0], [0.42228785157203674, 0.03822512924671173, 0.056164562702178955, 0.03373495861887932, 0.044957708567380905, 0.02277299202978611, 0.0300452820956707, 0.04248708486557007, 0.04043055698275566, 0.056272298097610474, 0.050744082778692245, 0.057307954877614975, 0.05850312486290932, 0.04606639966368675]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [5.4443644330604e-06, 0.9999945163726807, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.000340827158652246, 1.522992351965513e-05, 0.9996439218521118, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.5033668887554086e-07, 1.5539581843526662e-09, 1.0828672429852304e-06, 0.9999988079071045, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.4860424016660545e-05, 3.924570137314731e-06, 4.7554352931911126e-05, 0.0029133642092347145, 0.9970203042030334, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.1183689469239653e-08, 5.399687363905059e-10, 4.442081191768921e-08, 0.004807529039680958, 1.7774240404833108e-05, 0.9951747059822083, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.2336930012679659e-05, 1.1335694694025733e-07, 1.7473822708780062e-06, 8.646557398606092e-05, 0.002848046598955989, 0.001072101527824998, 0.9959790706634521, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.4322378067154204e-06, 3.87321415473707e-06, 2.9090304451528937e-05, 0.00010180847311858088, 0.008485228754580021, 0.0004340687009971589, 0.005096187349408865, 0.9858474135398865, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [6.682396302970517e-11, 4.753054838602111e-08, 2.8264020457413608e-08, 1.485501229581132e-06, 7.235502721414377e-07, 4.88603109261021e-06, 9.3033901293893e-07, 1.1520458429004066e-05, 0.9999804496765137, 0.0, 0.0, 0.0, 0.0, 0.0], [4.509112841333263e-06, 1.4262801641962142e-06, 3.9936694520292804e-05, 6.5301328504574485e-06, 3.416110121179372e-05, 8.898381929611787e-05, 1.3975942238175776e-05, 0.0004578701627906412, 3.156268212478608e-05, 0.9993211030960083, 0.0, 0.0, 0.0, 0.0], [1.2304453775868751e-05, 8.846908713167068e-06, 1.1869107765960507e-05, 1.8301088857697323e-05, 0.0006726715946570039, 2.9870310754631646e-05, 9.385061275679618e-05, 0.0032646076288074255, 0.004700236953794956, 0.0007520715589635074, 0.9904352426528931, 0.0, 0.0, 0.0], [3.518028179883004e-08, 1.1920349256289153e-11, 3.274404969388911e-09, 2.6973480089509394e-06, 4.262755426420739e-11, 3.6191207186675456e-07, 3.968440159818343e-11, 2.6712736467260356e-09, 1.1683687350938499e-10, 6.754234060935005e-09, 1.0540309736128961e-09, 0.9999969005584717, 0.0, 0.0], [1.216761802425026e-07, 2.3296263407246443e-06, 1.0410583684006269e-07, 1.0372890528742573e-06, 2.3239222457505093e-07, 1.1366762009856757e-05, 3.483045407648433e-09, 5.737806077377172e-06, 2.5247034045605687e-06, 1.4229512999008875e-05, 3.3851924854388926e-06, 2.753440276137553e-05, 0.9999314546585083, 0.0], [0.0021027533803135157, 5.3797426517121494e-05, 2.642461004143115e-05, 5.549227716983296e-05, 0.0004349422815721482, 5.564442108152434e-05, 0.0003221925289835781, 0.0019075508462265134, 0.012244115583598614, 0.00019615399651229382, 0.005641697905957699, 7.855122021283023e-06, 5.750841592089273e-05, 0.9768938422203064]], [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0009072376415133476, 0.9990928173065186, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.848416806955356e-05, 0.00033618693123571575, 0.9996353387832642, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.08687606825697e-06, 2.3963496005308116e-06, 4.65728362541995e-06, 0.9999918937683105, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.0430768674705178e-05, 1.9211251128581353e-05, 4.7870503294689115e-06, 0.0004905093810521066, 0.9994551539421082, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [2.5894426016748184e-06, 1.5198468190646963e-06, 7.012368996583973e-08, 3.681573798530735e-05, 0.013727720826864243, 0.98623126745224, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [9.139445319306105e-05, 0.00011409086437197402, 2.19574917537102e-06, 0.00022433913545683026, 0.0043365610763430595, 0.0023445438127964735, 0.9928868412971497, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [3.246472988394089e-05, 5.881246579519939e-06, 8.339918053934525e-07, 3.120941255474463e-05, 0.0001950529112946242, 7.835734140826389e-05, 0.003378617577254772, 0.9962775111198425, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0792935427161865e-05, 1.9718265320989303e-05, 3.5675466278917156e-06, 0.0003161633212585002, 0.0006945938221178949, 1.7135267626144923e-05, 0.0010499282507225871, 0.0006411561626009643, 0.9972468614578247, 0.0, 0.0, 0.0, 0.0, 0.0], [1.969202912732726e-06, 9.843746511251084e-07, 7.95048526924802e-06, 1.2612697446456878e-06, 7.5254224611853715e-06, 2.0916523624237016e-08, 4.208237669445225e-07, 9.048753213392047e-07, 2.441770448058378e-06, 0.999976634979248, 0.0, 0.0, 0.0, 0.0], [3.388111508684233e-06, 1.532968667561363e-07, 8.007928045117296e-07, 6.9148195507295895e-06, 9.635970172894304e-07, 8.262756701071794e-09, 2.7631315901999187e-08, 1.2104433722015528e-07, 8.786131729721092e-06, 0.010999530553817749, 0.9889793395996094, 0.0, 0.0, 0.0], [3.4670763398025883e-06, 3.2853522498044185e-06, 1.7884923408928444e-06, 1.271074688702356e-05, 1.981398440875637e-07, 6.105980165926894e-10, 1.331544474680868e-08, 1.059634935351994e-09, 6.316852250165539e-07, 0.0005560617428272963, 0.000964729639235884, 0.9984571933746338, 0.0, 0.0], [0.0001142337714554742, 4.442620138433995e-06, 6.048322120477678e-06, 2.263612077513244e-05, 2.3252282232988364e-07, 7.76467601326658e-09, 5.912871081648063e-09, 1.0200583488995107e-07, 1.1434232192186755e-06, 7.391762665065471e-06, 4.037147149915654e-08, 3.6513690702122403e-06, 0.9998400211334229, 0.0], [0.00011470724712125957, 1.807276362342236e-06, 1.0345141845391481e-06, 6.004875831422396e-05, 3.8280772969301324e-06, 1.865960825853108e-07, 3.2268536642732215e-07, 3.506449672840972e-07, 4.888435341854347e-06, 7.464206532858952e-07, 1.698989535725559e-07, 3.743731667782413e-06, 4.9968955863732845e-05, 0.9997581839561462]]]}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7f17f2400910>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(f\"The same heads with a different visualiser\")\n",
    "visualise_attention_patterns(top_logit_attr_heads, visualisation_type=\"attention_patterns\", batch_index=batch_index)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- The attention pattern of L5H1 is the clearest to analyse.\n",
    "    * It works on all 'scoping symbols': opening and closing bracket *and quotation marks*.\n",
    "    * It sends from every scoping symbol, and the BOS token, to all the subsequent tokens until the next scoping symbol.\n",
    "    * It seems to ignore the '()' tokens, rather than counting them as new scoping symbols.\n",
    "- It's a bit less clear what L6H2 is doing.\n",
    "    * It sends from certain key symbols: the BOS token, `.`, `,`, `=`, `(`, `[` and `'`.\n",
    "    * It mostly sends to all subsequent tokens, though sometimes it pays attention to scope.\n",
    "- The attention patterns for L9H7 are more complex.\n",
    "    * What's clear is that it sends information from the BOS token, opening brackets *and quotation marks* to some tokens within their scope.\n",
    "    * Actually the condition of being inside the scope gets violated in one instance: in prompt 3 information is sent from the first quotation mark to tokens inside the scope of the third. I'm not sure if this is 'intentional' or 'accidental'. Testing more examples would help.\n",
    "    * It seems plausible that L9H7 is using the computation from L5H1 to compute these activation patterns.\n",
    "    * One striking feature is the way the attention pattern seems to alternate between opening brackets.\n",
    "    * The patterns are similar to what we'd get if we set tokens to receive information depending on the parity of the current bracket depth, though not quite.\n",
    "    * This might be accomplished by exploiting the fact that the attention weights are normalised, by putting making the key vectors of later opening brackets larger.\n",
    "    * Knowing the parity alone is not enough to solve the task, but the model is probably transferring important data according to this activation pattern which allows it to solve the task.\n",
    "- The activation patterns of the negative heads are not very interesting.\n",
    "    * I suppose their negative effects come from *what* gets transferred according to the patterns."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reflections\n",
    "\n",
    "- It would be nice to examine the issue of which tokens to use in the metric in more detail.\n",
    "    * It seems to make sense to include most of the tokens beginning with a bracket in the metric, since the model will consider the whole range of possibilities when thinking about the next token.\n",
    "    * There are some tokens which it might not make sense to include though, such as ')('.\n",
    "    * I didn't look at what effect including tokens beginning with a space has. Is it worth it?\n",
    "    * It's not clear how multi-token metrics translate to the logit directions used in direct logit attribution: there isn't a single direction of improvement.\n",
    "    * Could it be that the simple metric used above comparing just two tokens (e.g. `(` and`)`) is good enough? In other words, does the logit difference between the two tokens sufficiently characterise the model's belief about whether the next bracket should be opening or closing?\n",
    "- I think it's worth making clear and precise what the model behaviour actually *is* that we want to explain.\n",
    "    * As I said in the [discussion section above](#discussion), intuitively we'd like to say that the model can tell when brackets are balanced, and prefers generating text with balanced brackets.\n",
    "    * But does this really correspond to giving a higher probability to `(` than `)` at every point?\n",
    "- A bigger diversity of prompts would have been better.\n",
    "    * I feel this rather detracted from my ability to understand how the model was working.\n",
    "    * Some useful variations would have been:\n",
    "        + More prompts with an earlier closed parenthesis set, followed later by an opening bracket.\n",
    "        + Greater bracket depths.\n",
    "        + More examples of quotation marks, including double ones.\n",
    "        + Prompts where the brackets were unbalanced by more than one.\n",
    "    * I used a trick to get the reference prompts and corrupted prompts to look very similar and have the same number of characters: the string \"')\" is a single token, so ending the reference prompt with a quotation mark allowed me to add in a bracket to get a balanced corrupted prompt without increasing the number of tokens.\n",
    "        + This restricted the range of prompts I could write.\n",
    "        + I'm not sure how necessary it was.\n",
    "- In general there's a bit of inconsistency in this document, since I changed the approach several times as I discovered new things."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mech-interp",
   "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.11.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}