NNsight_Walkthrough.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/JadenFiotto-Kaufman/nnsight/blob/docs_0.3/NNsight_v0_3.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Scfr942GwO4"
      },
      "source": [
        "<img src=\"https://nnsight.net/_static/images/nnsight_logo.svg\" alt=\"drawing\" width=\"200\"/>\n",
        "\n",
        "# **NNsight**\n",
        "\n",
        "## The API for a transparent science on black-box AI\n",
        "\n",
        "In this era of large-scale deep learning, the most interesting AI models are\n",
        "massive black boxes that are hard to run. Ordinary commercial inference service\n",
        "APIs let us interact with huge models, but they do not let us access model\n",
        "internals.\n",
        "\n",
        "The nnsight library is different: it provides full access to all the neural\n",
        "network internals. When used together with a remote service like the\n",
        "[National Deep Inference Fabric](https://thevisible.net/docs/NDIF-proposal.pdf)\n",
        "(NDIF), it makes possible to run complex experiments on huge open  models easily,\n",
        "with fully transparent access.\n",
        "\n",
        "Our team wants to enable entire labs and independent researchers alike, as we\n",
        "believe a large, passionate, and collaborative community will produce the next\n",
        "big insights on a profoundly important field.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N1OemD2VGyZx"
      },
      "source": [
        "# 1 First, let's start small"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WyLMmhrAKTNM"
      },
      "source": [
        "## Tracing Context\n",
        "\n",
        "To demonstrate the core functionality and syntax of nnsight, we'll define and\n",
        "use a tiny two layer neural network.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "R0RJijD0eXwf"
      },
      "outputs": [],
      "source": [
        "# Install nnsight\n",
        "!pip install nnsight\n",
        "!pip install --upgrade transformers torch\n",
        "\n",
        "from IPython.display import clear_output\n",
        "\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bgydw7i3HmIH"
      },
      "source": [
        "Our little model here is composed of two submodules – linear layers 'layer1' and 'layer2'. We specify the sizes of each of these modules, and create\n",
        "some complementary example input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "2pX2Wg8Ceo6N"
      },
      "outputs": [],
      "source": [
        "from collections import OrderedDict\n",
        "import torch\n",
        "\n",
        "input_size = 5\n",
        "hidden_dims = 10\n",
        "output_size = 2\n",
        "\n",
        "net = torch.nn.Sequential(\n",
        "    OrderedDict(\n",
        "        [\n",
        "            (\"layer1\", torch.nn.Linear(input_size, hidden_dims)),\n",
        "            (\"layer2\", torch.nn.Linear(hidden_dims, output_size)),\n",
        "        ]\n",
        "    )\n",
        ").requires_grad_(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tIPa2h2pJIwl"
      },
      "source": [
        "The core object of the nnsight package is `NNsight`. This wraps around a given\n",
        "pytorch model to enable the capabilities nnsight provides.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8H9R_ynTJI5y"
      },
      "outputs": [],
      "source": [
        "import nnsight\n",
        "from nnsight import NNsight\n",
        "\n",
        "tiny_model = NNsight(net)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2NfISlQ_Ilvp"
      },
      "source": [
        "Printing a Pytorch model shows a named hierarchy of modules which is very useful\n",
        "when accessing sub-components directly. NNsight reflect the same hierarchy and can be similarly printed.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bYtnbJHvlGZV",
        "outputId": "5eb7c572-451c-43c6-956f-af224ae1867c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sequential(\n",
            "  (layer1): Linear(in_features=5, out_features=10, bias=True)\n",
            "  (layer2): Linear(in_features=10, out_features=2, bias=True)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "print(tiny_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "djC5kyJWLuUH"
      },
      "source": [
        "Before we actually get to using the model we just created, let's talk about\n",
        "Python contexts.\n",
        "\n",
        "Python contexts define a scope using the `with` statement and are often used to\n",
        "create some object, or initiate some logic, that you later want to destroy or\n",
        "conclude.\n",
        "\n",
        "The most common application is opening files like the following example:\n",
        "\n",
        "```python\n",
        "with open('myfile.txt', 'r') as file:\n",
        "  text = file.read()\n",
        "```\n",
        "\n",
        "Python uses the `with` keyword to enter a context-like object. This object\n",
        "defines logic to be run at the start of the `with` block, as well as logic to be\n",
        "run when exiting. When using `with` for a file, entering the context opens the\n",
        "file and exiting the context closes it. Being within the context means we can\n",
        "read from the file. Simple enough! Now we can discuss how `nnsight` uses\n",
        "contexts to enable intuitive access into the internals of a neural network.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iNvuCOeyojcA"
      },
      "source": [
        "The main tool with `nnsight` is a context for tracing.\n",
        "\n",
        "We enter the tracing context by calling `model.trace(<input>)` on an `NNsight`\n",
        "model, which defines how we want to run the model. Inside the context, we will\n",
        "be able to customize how the neural network runs. The model is actually run upon\n",
        "exiting the tracing context.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "qEXQ4auPSL-m"
      },
      "outputs": [],
      "source": [
        "# random input\n",
        "input = torch.rand((1, input_size))\n",
        "\n",
        "with tiny_model.trace(input) as tracer:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UZQsHinjqicJ"
      },
      "source": [
        "But where's the output? To get that, we'll have to learn how to request it from\n",
        "within the tracing context.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xMfBpYzDPMoB"
      },
      "source": [
        "## Getting\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B5_aFwFRv0ax"
      },
      "source": [
        "Earlier, when we wrapped our little neural net with the `NNsight` class. This\n",
        "added a couple properties to each module in the model (including the root model\n",
        "itself). The two most important ones are `.input` and `.output`.\n",
        "\n",
        "```python\n",
        "model.input\n",
        "model.output\n",
        "```\n",
        "\n",
        "The names are self explanatory. They correspond to the inputs and outputs of\n",
        "their respective modules during a forward pass of the model. We can use these\n",
        "attributes inside the `with` block.\n",
        "\n",
        "However, it is important to understand that the model is not executed until the\n",
        "end of the tracing context. How can we access inputs and outputs before the\n",
        "model is run? The trick is deferred execution.\n",
        "\n",
        "`.input` and `.output` are Proxies for the eventual inputs and outputs of a\n",
        "module. In other words, when we access `model.output` what we are\n",
        "communicating to `nnsight` is, \"When you compute the output of `model`, please\n",
        "grab it for me and put the value into its corresponding Proxy object's `.value`\n",
        "attribute.\" Let's try it:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 303
        },
        "id": "cYe8-r9ptGaG",
        "outputId": "adf25efe-d2cf-4f65-9eaf-15b155762f4b"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "Accessing value before it's been set.",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[6], line 5\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m tiny_model\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;28minput\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m tracer:\n\u001b[1;32m      3\u001b[0m     output \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39moutput\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[43moutput\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m)\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Proxy.py:50\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     43\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvalue\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     44\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this proxy's node.\u001b[39;00m\n\u001b[1;32m     45\u001b[0m \n\u001b[1;32m     46\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     47\u001b[0m \u001b[38;5;124;03m        Any: The stored value of the proxy, populated during execution of the model.\u001b[39;00m\n\u001b[1;32m     48\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 50\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Node.py:182\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    172\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this node.\u001b[39;00m\n\u001b[1;32m    173\u001b[0m \n\u001b[1;32m    174\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    178\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or destroyed).\u001b[39;00m\n\u001b[1;32m    179\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    181\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdone():\n\u001b[0;32m--> 182\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAccessing value before it\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms been set.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    184\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_value\n",
            "\u001b[0;31mValueError\u001b[0m: Accessing value before it's been set."
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    output = tiny_model.output\n",
        "\n",
        "print(output.value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBz5qfwuR-6F"
      },
      "source": [
        "Oh no an error! \"Accessing Proxy value before it's been set.\"\n",
        "\n",
        "Why doesn't our `output` have a `value`?\n",
        "\n",
        "Proxy objects will only have their value at the end of a context if we call\n",
        "`.save()` on them. This helps to reduce memory costs. Adding `.save()` fixes the\n",
        "error:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y_bXRd5dvsBu",
        "outputId": "a8e8985c-2661-46a8-975b-d43f67cfc162"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[-0.2687, -0.3314]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    output = tiny_model.output.save()\n",
        "\n",
        "print(output.value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t5C0UZCwvrYn"
      },
      "source": [
        "Success! We now have the model output. We just completed out first\n",
        "intervention using `nnsight`.\n",
        "\n",
        "Each time we access a module's input or output, we create an _intervention_ in\n",
        "the neural network's forward pass. Collectively these requests form the\n",
        "_intervention graph_. We call the process of executing it alongside the model's\n",
        "normal computation graph, _interleaving_.\n",
        "\n",
        "<details>\n",
        "<summary>On Model output</summary>\n",
        "\n",
        "---\n",
        "\n",
        "If we don't need to access anything other than the final model output, we can\n",
        "call the tracing context with `trace=False` and not use it as a context. This could be especially useful for easy remote inference.\n",
        "\n",
        "```python\n",
        "  output = model.trace(<inputs>, trace=False)\n",
        "```\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n",
        "\n",
        "Just like we saved the output of the model as a whole, we can save the output of\n",
        "any of its submodules. We use normal Python attribute syntax. We can discover\n",
        "how to access them by name by printing out the model:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0WcVUSP-0CJi",
        "outputId": "7fd9d6ac-152d-4431-9a4d-0b60321638eb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sequential(\n",
            "  (layer1): Linear(in_features=5, out_features=10, bias=True)\n",
            "  (layer2): Linear(in_features=10, out_features=2, bias=True)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "print(tiny_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's access the output of the first layer (non-coincidentally named 'layer1'):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "akWr-cNqy-9O",
        "outputId": "68b9e071-c609-4b70-906d-2f2cd93b4495"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    l1_output = tiny_model.layer1.output.save()\n",
        "\n",
        "print(l1_output.value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "85A-aP_03ht6"
      },
      "source": [
        "Let's do the same for the input of layer2. While we're at it, let's also drop\n",
        "the `as tracer`, as we won't be needing the tracer object itself for a few\n",
        "sections:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3EHEN38N3nXR",
        "outputId": "5afd17a7-0bcc-4d1a-c131-7038bde9455d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    l2_input = tiny_model.layer2.input.save()\n",
        "\n",
        "print(l2_input.value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jk-U8zi33Gi-"
      },
      "source": [
        "<details>\n",
        "  <summary>On module inputs</summary>\n",
        "\n",
        "---\n",
        "\n",
        "Notice how the value for `l2_input`, is just a single tensor. By default, the `.input` attribute of a module will return the **first** tensor input to the module.\n",
        "\n",
        "We can also access the full input to a module by using the `.inputs` attribute which will return the values in the form of:\n",
        "\n",
        "      tuple(tuple(args), dictionary(kwargs))\n",
        "\n",
        "Where the first index of the tuple is itself a tuple of all positional\n",
        "arguments, and the second index is a dictionary of the keyword arguments.\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Until now we were saving the output of the model and its submodules within the `Trace` context to then print it after exiting the context. We will continuing doing this in the rest of the tutorial since it's a good practice to save the computation results for later analysis. \n",
        "\n",
        "However, we can also log the outputs of the model and its submodules within the `Trace` context. This is useful for debugging and understanding the model's behavior while saving memory. Let's see how to do this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 - out:  tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "  tracer.log(\"Layer 1 - out: \", tiny_model.layer1.output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p7Xo_PHyPr4p"
      },
      "source": [
        "## Functions, Methods, and Operations\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ehSofWbx5DSx"
      },
      "source": [
        "Now that we can access activations, we also want to do some post-processing on\n",
        "it. Let's find out which dimension of layer1's output has the highest value.\n",
        "\n",
        "We could do this by calling `torch.argmax(...)` after the tracing context or we\n",
        "can just leverage the fact that `nnsight` handles Pytorch functions and methods within\n",
        "the tracing context, by creating a Proxy request for it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c5XCiSZn2p3k",
        "outputId": "97d34dee-7bb5-4735-8c53-6aac6377af9e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(7)\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Note we don't need to call .save() on the output,\n",
        "    # as we're only using its value within the tracing context.\n",
        "    l1_output = tiny_model.layer1.output\n",
        "\n",
        "    # We do need to save the argmax tensor however, \n",
        "    # as we're using it outside the tracing context.\n",
        "    l1_amax = torch.argmax(l1_output, dim=1).save()\n",
        "\n",
        "print(l1_amax[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BGPUJvWq_bOp"
      },
      "source": [
        "Nice! That worked seamlessly, but hold on, how come we didn't need to call\n",
        "`.value[0]` on the result? In previous sections, we were just being explicit to\n",
        "get an understanding of Proxies and their value. In practice, however, `nnsight`\n",
        "knows that when outside of the tracing context we only care about the actual\n",
        "value, and so printing, indexing, and applying functions all immediately return\n",
        "and reflect the data in `.value`. So for the rest of the tutorial we won't use\n",
        "it.\n",
        "\n",
        "The same principles work for Pytorch methods and all operators as well:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FIcOYeEjFuln",
        "outputId": "d836abd3-b713-44c5-aedc-95bf85abf11a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(-1.3797)\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    value = (tiny_model.layer1.output.sum() + tiny_model.layer2.output.sum()).save()\n",
        "\n",
        "print(value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V0vNOJtJ2oFR"
      },
      "source": [
        "The code block above is saying to `nnsight`, \"Run the model with\n",
        "the given `input`. When the output of layer1 is computed, take its sum. Then do\n",
        "the same for layer2. Now that both of those are computed, add them and make sure\n",
        "not to delete this value as I wish to use it outside of the tracing context.\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Custom Functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Everything within the tracing context operates on the intervention graph. Therefore for `nnsight` to trace a  function it must also be a part of the intervention graph. \n",
        "\n",
        "Out-of-the-box `nnsight` supports Pytorch functions and methods, all operators, as well the `einops` library. We don't need to do anything special to use them. But what do we do if we want to use custom functions? How do we add them to the intervention graph?\n",
        "\n",
        "Enter `nnsight.apply()`. It allows us to add new functions to the intervention graph. Let's see how it works:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(-0.7796) tensor(-0.7796)\n"
          ]
        }
      ],
      "source": [
        "# Take a tensor and return the sum of its elements\n",
        "def tensor_sum(tensor):\n",
        "    flat = tensor.flatten()\n",
        "    total = 0\n",
        "    for element in flat:\n",
        "        total += element.item()\n",
        "    \n",
        "    return torch.tensor(total)\n",
        "\n",
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    # Specify the function name and its arguments (in a coma-separated form) to add to the intervention graph\n",
        "    custom_sum = nnsight.apply(tensor_sum, tiny_model.layer1.output).save()\n",
        "    sum = tiny_model.layer1.output.sum()\n",
        "    sum.save()\n",
        "    \n",
        "\n",
        "print(custom_sum, sum)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`nnsight.apply()` executes the function it wraps and returns its output as a Proxy object. We can then use this Proxy object as we would any other.\n",
        "\n",
        "The applications of `nnsight.apply` are wide. It can be used to wrap any custom function or functions from libraries that `nnsight` does not support out-of-the-box."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m_5qH5gHPOT_"
      },
      "source": [
        "## Setting\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rgju-b_IOLlq"
      },
      "source": [
        "Getting and analyzing the activations from various points in a model can be\n",
        "really insightful, and a number of ML techniques do exactly that. However, often\n",
        "times we not only want to view the computation of a model, but influence it as\n",
        "well.\n",
        "\n",
        "To demonstrate the effect of editing the flow of information through the model,\n",
        "let's set the first dimension of the first layer's output to 0. `NNsight` makes\n",
        "this really easy using '=' operator:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-6y2wzJqOz3a",
        "outputId": "6274e5c5-c5fe-4e07-87c4-a119cfd09742"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Before: tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n",
            "After: tensor([[ 0.0000,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Save the output before the edit to compare.\n",
        "    # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    # Access the 0th index of the hidden state dimension and set it to 0.\n",
        "    tiny_model.layer1.output[:, 0] = 0\n",
        "\n",
        "    # Save the output after to see our edit.\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oZz9SMs3Y_iS"
      },
      "source": [
        "Seems our change was reflected. Now the same for the last dimension:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 363
        },
        "id": "qwlqvHFcxld2",
        "outputId": "28ae7175-5af0-4bd0-f9f0-21680a4fd748"
      },
      "outputs": [
        {
          "ename": "IndexError",
          "evalue": "Above exception when execution Node: 'setitem_0' in Graph: '6063279136'",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Node.py:380\u001b[0m, in \u001b[0;36mNode.execute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    379\u001b[0m \u001b[38;5;66;03m# Call the target to get value.\u001b[39;00m\n\u001b[0;32m--> 380\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    382\u001b[0m \u001b[38;5;66;03m# Set value.\u001b[39;00m\n",
            "\u001b[0;31mIndexError\u001b[0m: index 10 is out of bounds for dimension 1 with size 10",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[16], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;43;01mwith\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtiny_model\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m      2\u001b[0m \n\u001b[1;32m      3\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# Save the output before the edit to compare.\u001b[39;49;00m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# Notice we apply .clone() before saving as the setting operation is in-place.\u001b[39;49;00m\n\u001b[1;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[43ml1_output_before\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mtiny_model\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlayer1\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mclone\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msave\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# Access the last index of the hidden state dimension and set it to 0.\u001b[39;49;00m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/Tracer.py:103\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     98\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minvoker\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__exit__\u001b[39m(\u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m    100\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m--> 103\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__exit__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mexc_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_val\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_tb\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/GraphBasedContext.py:218\u001b[0m, in \u001b[0;36mGraphBasedContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m    215\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    216\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n\u001b[0;32m--> 218\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/backends/LocalBackend.py:27\u001b[0m, in \u001b[0;36mLocalBackend.__call__\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, obj: LocalMixin):\n\u001b[0;32m---> 27\u001b[0m     \u001b[43mobj\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlocal_backend_execute\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/Tracer.py:147\u001b[0m, in \u001b[0;36mTracer.local_backend_execute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    143\u001b[0m     invoker_inputs \u001b[38;5;241m=\u001b[39m resolve_dependencies(invoker_inputs)\n\u001b[1;32m    145\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\u001b[38;5;241m.\u001b[39mexecute()\n\u001b[0;32m--> 147\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minterleave\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    148\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    149\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    150\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minvoker_inputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    151\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    152\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    154\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\n\u001b[1;32m    155\u001b[0m graph\u001b[38;5;241m.\u001b[39malive \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/models/NNsightModel.py:462\u001b[0m, in \u001b[0;36mNNsight.interleave\u001b[0;34m(self, fn, intervention_graph, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    454\u001b[0m intervention_handler \u001b[38;5;241m=\u001b[39m InterventionHandler(\n\u001b[1;32m    455\u001b[0m     intervention_graph, batch_groups, batch_size\n\u001b[1;32m    456\u001b[0m )\n\u001b[1;32m    458\u001b[0m module_paths \u001b[38;5;241m=\u001b[39m InterventionProtocol\u001b[38;5;241m.\u001b[39mget_interventions(\n\u001b[1;32m    459\u001b[0m     intervention_graph\n\u001b[1;32m    460\u001b[0m )\u001b[38;5;241m.\u001b[39mkeys()\n\u001b[0;32m--> 462\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mwith\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mHookHandler\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    463\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_model\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    464\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mlist\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mmodule_paths\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    465\u001b[0m \u001b[43m    \u001b[49m\u001b[43minput_hook\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mactivations\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodule_path\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mInterventionProtocol\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mintervene\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    466\u001b[0m \u001b[43m        \u001b[49m\u001b[43mactivations\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodule_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43minput\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mintervention_handler\u001b[49m\n\u001b[1;32m    467\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    468\u001b[0m \u001b[43m    \u001b[49m\u001b[43moutput_hook\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mactivations\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodule_path\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mInterventionProtocol\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mintervene\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    469\u001b[0m \u001b[43m        \u001b[49m\u001b[43mactivations\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodule_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43moutput\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mintervention_handler\u001b[49m\n\u001b[1;32m    470\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    471\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m    472\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mtry\u001b[39;49;00m\u001b[43m:\u001b[49m\n\u001b[1;32m    473\u001b[0m \u001b[43m        \u001b[49m\u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/intervention.py:574\u001b[0m, in \u001b[0;36mHookHandler.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m    571\u001b[0m     handle\u001b[38;5;241m.\u001b[39mremove()\n\u001b[1;32m    573\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(exc_val, \u001b[38;5;167;01mException\u001b[39;00m):\n\u001b[0;32m--> 574\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/models/NNsightModel.py:473\u001b[0m, in \u001b[0;36mNNsight.interleave\u001b[0;34m(self, fn, intervention_graph, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    462\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m HookHandler(\n\u001b[1;32m    463\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model,\n\u001b[1;32m    464\u001b[0m     \u001b[38;5;28mlist\u001b[39m(module_paths),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    470\u001b[0m     ),\n\u001b[1;32m    471\u001b[0m ):\n\u001b[1;32m    472\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 473\u001b[0m         \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    474\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m protocols\u001b[38;5;241m.\u001b[39mEarlyStopProtocol\u001b[38;5;241m.\u001b[39mEarlyStopException:\n\u001b[1;32m    475\u001b[0m         \u001b[38;5;66;03m# TODO: Log.\u001b[39;00m\n\u001b[1;32m    476\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m node \u001b[38;5;129;01min\u001b[39;00m intervention_graph\u001b[38;5;241m.\u001b[39mnodes\u001b[38;5;241m.\u001b[39mvalues():\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/models/NNsightModel.py:584\u001b[0m, in \u001b[0;36mNNsight._execute\u001b[0;34m(self, *prepared_inputs, **kwargs)\u001b[0m\n\u001b[1;32m    581\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m    582\u001b[0m     \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[0;32m--> 584\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_model\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    585\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mprepared_inputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    586\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    587\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/torch/nn/modules/module.py:1603\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1600\u001b[0m     bw_hook \u001b[38;5;241m=\u001b[39m hooks\u001b[38;5;241m.\u001b[39mBackwardHook(\u001b[38;5;28mself\u001b[39m, full_backward_hooks, backward_pre_hooks)\n\u001b[1;32m   1601\u001b[0m     args \u001b[38;5;241m=\u001b[39m bw_hook\u001b[38;5;241m.\u001b[39msetup_input_hook(args)\n\u001b[0;32m-> 1603\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1604\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks:\n\u001b[1;32m   1605\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m hook_id, hook \u001b[38;5;129;01min\u001b[39;00m (\n\u001b[1;32m   1606\u001b[0m         \u001b[38;5;241m*\u001b[39m_global_forward_hooks\u001b[38;5;241m.\u001b[39mitems(),\n\u001b[1;32m   1607\u001b[0m         \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks\u001b[38;5;241m.\u001b[39mitems(),\n\u001b[1;32m   1608\u001b[0m     ):\n\u001b[1;32m   1609\u001b[0m         \u001b[38;5;66;03m# mark that always called hook is run\u001b[39;00m\n",
            "File \u001b[0;32m/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/torch/nn/modules/container.py:219\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    218\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 219\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    220\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n",
            "File \u001b[0;32m/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/torch/nn/modules/module.py:1553\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1552\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1553\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/torch/nn/modules/module.py:1616\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1614\u001b[0m     hook_result \u001b[38;5;241m=\u001b[39m hook(\u001b[38;5;28mself\u001b[39m, args, kwargs, result)\n\u001b[1;32m   1615\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1616\u001b[0m     hook_result \u001b[38;5;241m=\u001b[39m \u001b[43mhook\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mresult\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1618\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m hook_result \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1619\u001b[0m     result \u001b[38;5;241m=\u001b[39m hook_result\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/intervention.py:559\u001b[0m, in \u001b[0;36mHookHandler.__enter__.<locals>.output_hook\u001b[0;34m(module, input, output, module_path)\u001b[0m\n\u001b[1;32m    558\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21moutput_hook\u001b[39m(module, \u001b[38;5;28minput\u001b[39m, output, module_path\u001b[38;5;241m=\u001b[39mmodule_path):\n\u001b[0;32m--> 559\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutput_hook\u001b[49m\u001b[43m(\u001b[49m\u001b[43moutput\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodule_path\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/models/NNsightModel.py:468\u001b[0m, in \u001b[0;36mNNsight.interleave.<locals>.<lambda>\u001b[0;34m(activations, module_path)\u001b[0m\n\u001b[1;32m    454\u001b[0m intervention_handler \u001b[38;5;241m=\u001b[39m InterventionHandler(\n\u001b[1;32m    455\u001b[0m     intervention_graph, batch_groups, batch_size\n\u001b[1;32m    456\u001b[0m )\n\u001b[1;32m    458\u001b[0m module_paths \u001b[38;5;241m=\u001b[39m InterventionProtocol\u001b[38;5;241m.\u001b[39mget_interventions(\n\u001b[1;32m    459\u001b[0m     intervention_graph\n\u001b[1;32m    460\u001b[0m )\u001b[38;5;241m.\u001b[39mkeys()\n\u001b[1;32m    462\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m HookHandler(\n\u001b[1;32m    463\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model,\n\u001b[1;32m    464\u001b[0m     \u001b[38;5;28mlist\u001b[39m(module_paths),\n\u001b[1;32m    465\u001b[0m     input_hook\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mlambda\u001b[39;00m activations, module_path: InterventionProtocol\u001b[38;5;241m.\u001b[39mintervene(\n\u001b[1;32m    466\u001b[0m         activations, module_path, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minput\u001b[39m\u001b[38;5;124m\"\u001b[39m, intervention_handler\n\u001b[1;32m    467\u001b[0m     ),\n\u001b[0;32m--> 468\u001b[0m     output_hook\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mlambda\u001b[39;00m activations, module_path: \u001b[43mInterventionProtocol\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mintervene\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    469\u001b[0m \u001b[43m        \u001b[49m\u001b[43mactivations\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodule_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43moutput\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mintervention_handler\u001b[49m\n\u001b[1;32m    470\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m    471\u001b[0m ):\n\u001b[1;32m    472\u001b[0m     \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    473\u001b[0m         fn(\u001b[38;5;241m*\u001b[39minputs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/intervention.py:449\u001b[0m, in \u001b[0;36mInterventionProtocol.intervene\u001b[0;34m(cls, activations, module_path, key, intervention_handler)\u001b[0m\n\u001b[1;32m    442\u001b[0m     value \u001b[38;5;241m=\u001b[39m util\u001b[38;5;241m.\u001b[39mapply(\n\u001b[1;32m    443\u001b[0m         activations,\n\u001b[1;32m    444\u001b[0m         narrow,\n\u001b[1;32m    445\u001b[0m         torch\u001b[38;5;241m.\u001b[39mTensor,\n\u001b[1;32m    446\u001b[0m     )\n\u001b[1;32m    448\u001b[0m \u001b[38;5;66;03m# Value injection.\u001b[39;00m\n\u001b[0;32m--> 449\u001b[0m \u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mset_value\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    451\u001b[0m \u001b[38;5;66;03m# Check if through the previous value injection, there was a 'swap' intervention.\u001b[39;00m\n\u001b[1;32m    452\u001b[0m \u001b[38;5;66;03m# This would mean we want to replace activations for this batch with some other ones.\u001b[39;00m\n\u001b[1;32m    453\u001b[0m value \u001b[38;5;241m=\u001b[39m protocols\u001b[38;5;241m.\u001b[39mSwapProtocol\u001b[38;5;241m.\u001b[39mget_swap(\n\u001b[1;32m    454\u001b[0m     intervention_handler\u001b[38;5;241m.\u001b[39mgraph, value\n\u001b[1;32m    455\u001b[0m )\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Node.py:410\u001b[0m, in \u001b[0;36mNode.set_value\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m    407\u001b[0m     listener\u001b[38;5;241m.\u001b[39mremaining_dependencies \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m    409\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m listener\u001b[38;5;241m.\u001b[39mfulfilled() \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\u001b[38;5;241m.\u001b[39msequential:\n\u001b[0;32m--> 410\u001b[0m         \u001b[43mlistener\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    412\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m dependency \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39marg_dependencies:\n\u001b[1;32m    413\u001b[0m     dependency\u001b[38;5;241m.\u001b[39mremaining_listeners \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Node.py:387\u001b[0m, in \u001b[0;36mNode.execute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    383\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mset_value(output)\n\u001b[1;32m    385\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m--> 387\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(e)(\n\u001b[1;32m    388\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAbove exception when execution Node: \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m in Graph: \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\u001b[38;5;241m.\u001b[39mid\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    389\u001b[0m     ) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m    391\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    392\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mremaining_dependencies \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n",
            "\u001b[0;31mIndexError\u001b[0m: Above exception when execution Node: 'setitem_0' in Graph: '6063279136'"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Save the output before the edit to compare.\n",
        "    # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    # Access the last index of the hidden state dimension and set it to 0.\n",
        "    tiny_model.layer1.output[:, hidden_dims] = 0\n",
        "\n",
        "    # Save the output after to see our edit.\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Oh no, we are getting an error. Looks like it's happening when we are setting the output. \n",
        "\n",
        "How can we find what went wrong more precisely? Is there an easy way to debug this?\n",
        "\n",
        "Come in \"Scanning\" and \"Validating\"! We can enable these features by setting the `scan=True` and `validate=True` flags in the `trace` method.\n",
        "\n",
        "\n",
        "Let's run this again and see what it can do for us:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {},
      "outputs": [
        {
          "ename": "IndexError",
          "evalue": "index 10 is out of bounds for dimension 1 with size 10",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[17], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# turn on scan and validate\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;43;01mwith\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtiny_model\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mscan\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalidate\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43ml1_output_before\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mtiny_model\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlayer1\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mclone\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msave\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      6\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# the error is happening here\u001b[39;49;00m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/Tracer.py:103\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     98\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minvoker\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__exit__\u001b[39m(\u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m    100\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m--> 103\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__exit__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mexc_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_val\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_tb\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/GraphBasedContext.py:216\u001b[0m, in \u001b[0;36mGraphBasedContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m    214\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\u001b[38;5;241m.\u001b[39malive \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m    215\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 216\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n\u001b[1;32m    218\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbackend(\u001b[38;5;28mself\u001b[39m)\n",
            "Cell \u001b[0;32mIn[17], line 7\u001b[0m\n\u001b[1;32m      4\u001b[0m     l1_output_before \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39mlayer1\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39mclone()\u001b[38;5;241m.\u001b[39msave()\n\u001b[1;32m      6\u001b[0m     \u001b[38;5;66;03m# the error is happening here\u001b[39;00m\n\u001b[0;32m----> 7\u001b[0m     \u001b[43mtiny_model\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlayer1\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhidden_dims\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m      9\u001b[0m     l1_output_after \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39mlayer1\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39msave()\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBefore:\u001b[39m\u001b[38;5;124m\"\u001b[39m, l1_output_before)\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Proxy.py:89\u001b[0m, in \u001b[0;36mProxy.__setitem__\u001b[0;34m(self, key, value)\u001b[0m\n\u001b[1;32m     88\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__setitem__\u001b[39m(\u001b[38;5;28mself\u001b[39m, key: Union[Proxy, Any], value: Union[Self, Any]) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m---> 89\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     90\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtarget\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43moperator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msetitem\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     91\u001b[0m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     92\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Node.py:270\u001b[0m, in \u001b[0;36mNode.create\u001b[0;34m(self, target, proxy_value, args, kwargs, name)\u001b[0m\n\u001b[1;32m    267\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m value\n\u001b[1;32m    269\u001b[0m \u001b[38;5;66;03m# Otherwise just create the Node on the Graph like normal.\u001b[39;00m\n\u001b[0;32m--> 270\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    271\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtarget\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    272\u001b[0m \u001b[43m    \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    273\u001b[0m \u001b[43m    \u001b[49m\u001b[43mproxy_value\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mproxy_value\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    274\u001b[0m \u001b[43m    \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    275\u001b[0m \u001b[43m    \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    276\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Graph.py:113\u001b[0m, in \u001b[0;36mGraph.create\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    106\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcreate\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Proxy:\n\u001b[1;32m    107\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Creates a Node directly on this `Graph` and returns its `Proxy`.\u001b[39;00m\n\u001b[1;32m    108\u001b[0m \n\u001b[1;32m    109\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m    110\u001b[0m \u001b[38;5;124;03m        Proxy: `Proxy` for newly created `Node`.\u001b[39;00m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 113\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mproxy_class(\u001b[43mNode\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m)\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Node.py:96\u001b[0m, in \u001b[0;36mNode.__init__\u001b[0;34m(self, target, graph, proxy_value, args, kwargs, name)\u001b[0m\n\u001b[1;32m     93\u001b[0m \u001b[38;5;66;03m# If theres an alive Graph, add it.\u001b[39;00m\n\u001b[1;32m     94\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattached():\n\u001b[0;32m---> 96\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Graph.py:131\u001b[0m, in \u001b[0;36mGraph.add\u001b[0;34m(self, node)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;66;03m# If we're validating and the user did not provide a proxy_value, execute the given target with meta proxy values to compute new proxy_value.\u001b[39;00m\n\u001b[1;32m    129\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mvalidate \u001b[38;5;129;01mand\u001b[39;00m node\u001b[38;5;241m.\u001b[39mproxy_value \u001b[38;5;129;01mis\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39m_empty:\n\u001b[0;32m--> 131\u001b[0m     node\u001b[38;5;241m.\u001b[39mproxy_value \u001b[38;5;241m=\u001b[39m \u001b[43mvalidate\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    133\u001b[0m \u001b[38;5;66;03m# Get name of target.\u001b[39;00m\n\u001b[1;32m    134\u001b[0m name \u001b[38;5;241m=\u001b[39m (\n\u001b[1;32m    135\u001b[0m     node\u001b[38;5;241m.\u001b[39mtarget\n\u001b[1;32m    136\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(node\u001b[38;5;241m.\u001b[39mtarget, \u001b[38;5;28mstr\u001b[39m)\n\u001b[1;32m    137\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m node\u001b[38;5;241m.\u001b[39mtarget\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\n\u001b[1;32m    138\u001b[0m )\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/util.py:20\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m FakeTensorMode(\n\u001b[1;32m     15\u001b[0m     allow_non_fake_inputs\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m     16\u001b[0m     shape_env\u001b[38;5;241m=\u001b[39mShapeEnv(assume_static_by_default\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m),\n\u001b[1;32m     17\u001b[0m ) \u001b[38;5;28;01mas\u001b[39;00m fake_mode:\n\u001b[1;32m     18\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m FakeCopyMode(fake_mode):\n\u001b[0;32m---> 20\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43;01mwith\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mGlobalTracingContext\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexit_global_tracing_context\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m     22\u001b[0m \u001b[43m            \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mNode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprepare_inputs\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mproxy\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m     24\u001b[0m \u001b[43m            \u001b[49m\u001b[38;5;28;43;01mreturn\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     25\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     26\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     27\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/GraphBasedContext.py:330\u001b[0m, in \u001b[0;36mGlobalTracingContext.GlobalTracingExit.__exit__\u001b[0;34m(self, exc_type, exc_val, traceback)\u001b[0m\n\u001b[1;32m    326\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mPATCHER\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__enter__\u001b[39m()\n\u001b[1;32m    328\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(exc_val, \u001b[38;5;167;01mBaseException\u001b[39;00m):\n\u001b[0;32m--> 330\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/util.py:24\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m GlobalTracingContext\u001b[38;5;241m.\u001b[39mexit_global_tracing_context():\n\u001b[1;32m     22\u001b[0m     args, kwargs \u001b[38;5;241m=\u001b[39m Node\u001b[38;5;241m.\u001b[39mprepare_inputs((args, kwargs), proxy\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m---> 24\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtarget\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     25\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     26\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     27\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/contexts/GraphBasedContext.py:312\u001b[0m, in \u001b[0;36mGlobalTracingContext.GlobalTracingTorchHandler.__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m    305\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m_VariableFunctionsClass\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01min\u001b[39;00m func\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__qualname__\u001b[39m:\n\u001b[1;32m    306\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m GlobalTracingContext\u001b[38;5;241m.\u001b[39mGLOBAL_TRACING_CONTEXT\u001b[38;5;241m.\u001b[39mapply(\n\u001b[1;32m    307\u001b[0m         func,\n\u001b[1;32m    308\u001b[0m         \u001b[38;5;241m*\u001b[39margs,\n\u001b[1;32m    309\u001b[0m         \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m    310\u001b[0m     )\n\u001b[0;32m--> 312\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
            "\u001b[0;31mIndexError\u001b[0m: index 10 is out of bounds for dimension 1 with size 10"
          ]
        }
      ],
      "source": [
        "# turn on scan and validate\n",
        "with tiny_model.trace(input, scan=True, validate=True):\n",
        "\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "    \n",
        "    # the error is happening here\n",
        "    tiny_model.layer1.output[:, hidden_dims] = 0\n",
        "    \n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JCNR_Jkpxr8l"
      },
      "source": [
        "Ah of course, we needed to index at `hidden_dims - 1` not `hidden_dims`. \n",
        "\n",
        "How was `nnsight` able to catch this error?\n",
        "\n",
        "Earlier when discussing contexts in Python, we learned some logic happens upon\n",
        "entering, and some logic happens upon exiting. We know the model is actually run\n",
        "on exit, but what happens on enter? \n",
        "\n",
        "\n",
        "If `scan` and `validate` are enabled, our input is run though the model, but under its own \"fake\" context. This means the input makes its way through all of the model operations, allowing `nnsight` to record the shapes and data types of module inputs and outputs! The operations are never executed using tensors with real values so it doesn't incur any memory costs. Then, when creating proxy requests like the setting one above, `nnsight` also attempts to execute the request on the \"fake\" values we recorded. Hence, it lets us know if our request is feasible before even running the model.\n",
        "\n",
        "\"Scanning\" is what we call running \"fake\" inputs throught the model to collect\n",
        "information like shapes and types. \"Validating\" is what we call trying to\n",
        "execute the intervention proxies with \"fake\" inputs to see if they work. \n",
        "\"Validating\" is dependent on \"Scanning\" to work correctly, so we need to run the scan of the model at least once to debug with validate.\n",
        "\n",
        "<details>\n",
        "<summary>A word of caution</summary>\n",
        "\n",
        "---\n",
        "\n",
        "Some pytorch operations and related libraries don't work well with fake tensors\n",
        "\n",
        "If you are doing anything in a loop where efficiency is important, you should keep scanning and validating off. It's best to use them only when debugging or when you are unsure if your intervention will work.\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n",
        "\n",
        "Let's try again with the correct indexing, and view the shape of the output\n",
        "before leaving the tracing context:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wf8ugJKB2mru",
        "outputId": "4eed6c19-cb8c-4deb-e0e1-25aa80959c64"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 output shape: torch.Size([1, 10])\n",
            "Before: tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n",
            "After: tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158,  0.0000]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Save the output before the edit to compare.\n",
        "    # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    print(f\"Layer 1 output shape: {tiny_model.layer1.output.shape}\")\n",
        "\n",
        "    # Access the last index of the hidden state dimension and set it to 0.\n",
        "    tiny_model.layer1.output[:, hidden_dims - 1] = 0\n",
        "\n",
        "    # Save the output after to see our edit.\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9DH7PeL13-WU"
      },
      "source": [
        "We can also just replace proxy inputs and outputs with tensors of the same shape\n",
        "and type. Let's use the shape information we have at our disposal to add noise\n",
        "to the output, and replace it with this new noised tensor:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AP5wRmz_4YE7",
        "outputId": "4c30dc1d-232d-4900-bfdd-0c93a9b23ad7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Before: tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n",
            "After: tensor([[ 0.2283,  0.3262, -0.8443, -0.5498, -0.3424, -0.1178,  0.4780,  0.5522,\n",
            "         -0.2394, -0.1264]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Save the output before the edit to compare.\n",
        "    # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    # Create random noise with variance of .001\n",
        "    noise = (0.001**0.5) * torch.randn(l1_output_before.shape)\n",
        "\n",
        "    # Add to original value and replace.\n",
        "    tiny_model.layer1.output = l1_output_before + noise\n",
        "\n",
        "    # Save the output after to see our edit.\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "There is also another way to check the shape of the input and outputs of model's modules. We can run `.scan` manually to get the module's dimensions before running the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "10\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.scan(input):\n",
        "\n",
        "    dim = tiny_model.layer1.output.shape[-1]\n",
        "\n",
        "print(dim)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s016CelFP8sx"
      },
      "source": [
        "## Gradients\n",
        "\n",
        "`NNsight` also lets us apply backpropagation and access gradients with respect to a\n",
        "loss. Like `.input` and `.output` on modules, `nnsight` exposes `.grad` on\n",
        "Proxies themselves (assuming they are proxies of tensors):\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "L-cYUNNyRDn1",
        "outputId": "d03dafa1-ef1e-40f1-de28-4dc2f77fe4db"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 output gradient: tensor([[-0.2777, -0.1917,  0.1359,  0.2426,  0.1477, -0.0748,  0.0050, -0.1204,\n",
            "          0.1260,  0.2847]])\n",
            "Layer 2 output gradient: tensor([[1., 1.]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # We need to explicitly have the tensor require grad\n",
        "    # as the model we defined earlier turned off requiring grad.\n",
        "    tiny_model.layer1.output.requires_grad = True\n",
        "\n",
        "    # We call .grad on a tensor Proxy to communicate we want to store its gradient.\n",
        "    # We need to call .save() since .grad is its own Proxy.\n",
        "    layer1_output_grad = tiny_model.layer1.output.grad.save()\n",
        "    layer2_output_grad = tiny_model.layer2.output.grad.save()\n",
        "\n",
        "    # Need a loss to propagate through the later modules in order to have a grad.\n",
        "    loss = tiny_model.output.sum()\n",
        "    loss.backward()\n",
        "\n",
        "print(\"Layer 1 output gradient:\", layer1_output_grad)\n",
        "print(\"Layer 2 output gradient:\", layer2_output_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1o0JaaYvWlHG"
      },
      "source": [
        "All of the features we learned previously, also apply to `.grad`. In other\n",
        "words, we can apply operations to and edit the gradients. Let's zero the grad of\n",
        "`layer1` and double the grad of `layer2`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9bFoaJpOWlRb",
        "outputId": "1c026534-9ffe-40e7-d016-e4565ddde8ac"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 output gradient: tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])\n",
            "Layer 2 output gradient: tensor([[2., 2.]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # We need to explicitly have the tensor require grad\n",
        "    # as the model we defined earlier turned off requiring grad.\n",
        "    tiny_model.layer1.output.requires_grad = True\n",
        "\n",
        "    tiny_model.layer1.output.grad[:] = 0\n",
        "    tiny_model.layer2.output.grad = tiny_model.layer2.output.grad * 2\n",
        "\n",
        "    layer1_output_grad = tiny_model.layer1.output.grad.save()\n",
        "    layer2_output_grad = tiny_model.layer2.output.grad.save()\n",
        "\n",
        "    # Need a loss to propagate through the later modules in order to have a grad.\n",
        "    loss = tiny_model.output.sum()\n",
        "    loss.backward()\n",
        "\n",
        "print(\"Layer 1 output gradient:\", layer1_output_grad)\n",
        "print(\"Layer 2 output gradient:\", layer2_output_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Early Stopping"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If we are only interested in a model's intermediate computations, we can halt a forward pass run at any module level, reducing runtime and conserving compute resources. One examples where this could be particularly useful would if we are working with SAEs - we can train an SAE on one layer and then stop the execution. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "L1 - Output:  tensor([[ 0.2341,  0.3416, -0.8637, -0.5382, -0.3792, -0.1253,  0.4137,  0.5758,\n",
            "         -0.3158, -0.1226]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "   l1_out = tiny_model.layer1.output.save()\n",
        "   tiny_model.layer1.output.stop()\n",
        "\n",
        "# get the output of the first layer and stop tracing\n",
        "print(\"L1 - Output: \", l1_out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Interventions within the tracing context do not necessarily execute in the order they are defined. Instead, their execution is tied to the module they are associated with.\n",
        "\n",
        "As a result, if the forward pass is terminated early any interventions linked to modules beyond that point will be skipped, even if they were defined earlier in the context.\n",
        "\n",
        "In the example below, the output of layer 2 _**cannot**_ be accessed since the model's execution was stopped at layer 1."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "L2 - Output:  "
          ]
        },
        {
          "ename": "ValueError",
          "evalue": "Accessing value before it's been set.",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[24], line 5\u001b[0m\n\u001b[1;32m      2\u001b[0m    l2_out \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39mlayer2\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39msave()\n\u001b[1;32m      3\u001b[0m    tiny_model\u001b[38;5;241m.\u001b[39mlayer1\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39mstop()\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;43mprint\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mL2 - Output: \u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43ml2_out\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Proxy.py:56\u001b[0m, in \u001b[0;36mProxy.__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__str__\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mstr\u001b[39m:\n\u001b[1;32m     54\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mattached():\n\u001b[0;32m---> 56\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mstr\u001b[39m(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m)\n\u001b[1;32m     58\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m): \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mproxy_value\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mif\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mproxy_value\u001b[38;5;250m \u001b[39m\u001b[38;5;129;01mis\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;129;01mnot\u001b[39;00m\u001b[38;5;250m \u001b[39minspect\u001b[38;5;241m.\u001b[39m_empty\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01melse\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Proxy.py:50\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     42\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     43\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvalue\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     44\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this proxy's node.\u001b[39;00m\n\u001b[1;32m     45\u001b[0m \n\u001b[1;32m     46\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     47\u001b[0m \u001b[38;5;124;03m        Any: The stored value of the proxy, populated during execution of the model.\u001b[39;00m\n\u001b[1;32m     48\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 50\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\n",
            "File \u001b[0;32m~/Projects/NDIF/nnsight/src/nnsight/tracing/Node.py:182\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    172\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this node.\u001b[39;00m\n\u001b[1;32m    173\u001b[0m \n\u001b[1;32m    174\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    178\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or destroyed).\u001b[39;00m\n\u001b[1;32m    179\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    181\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdone():\n\u001b[0;32m--> 182\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAccessing value before it\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms been set.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    184\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_value\n",
            "\u001b[0;31mValueError\u001b[0m: Accessing value before it's been set."
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "   l2_out = tiny_model.layer2.output.save()\n",
        "   tiny_model.layer1.output.stop()\n",
        "\n",
        "print(\"L2 - Output: \", l2_out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Conditional Interventions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Interventions can also be made conditional. \n",
        "\n",
        "Inside the tracing context we can specify a new - conditional - context. This context will only execute the interventions within it if the condition is met."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Random Integer  tensor([-5])  is Odd\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
        "\n",
        "  with tracer.cond(rand_int % 2 == 0):\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
        "\n",
        "  with tracer.cond(rand_int % 2 == 1):\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In the example above, we have two conditional contexts with mutually exclusive conditions, just like a usual `If`-`Else` statement."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Conditional contexts can also be nested, if we want our interventions to depend on more than one condition at a time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Rand Int  8  is Positive and Even\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "  non_rand_int = 8\n",
        "  \n",
        "  with tracer.cond(non_rand_int > 0):\n",
        "    with tracer.cond(non_rand_int % 2 == 0):\n",
        "      tracer.log(\"Rand Int \", non_rand_int, \" is Positive and Even\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gvQ1nZgYQDG3"
      },
      "source": [
        "# 2️ Bigger"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "miQgUY4NAmDQ"
      },
      "source": [
        "Now that we have the basics of `nnsight` under our belt, we can scale our model\n",
        "up and combine the techniques we've learned into more interesting experiments.\n",
        "\n",
        "The `NNsight` class is very bare bones. It wraps a pre-defined model and does no\n",
        "pre-processing on the inputs we enter. It's designed to be extended with more\n",
        "complex and powerful types of models and we're excited to see what can be done\n",
        "to leverage its features.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5TJDblHiQpp1"
      },
      "source": [
        "## LanguageModel\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9l9mZOY5HFH2"
      },
      "source": [
        "`LanguageModel` is a subclass of `NNsight`. While we could define and create a\n",
        "model to pass in directly, `LanguageModel` includes special support for\n",
        "Huggingface language models, including automatically loading models from a\n",
        "Huggingface ID, and loading the model together with the appropriate tokenizer.\n",
        "\n",
        "Here is how we can use `LanguageModel` to load `GPT-2`:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 769,
          "referenced_widgets": [
            "0a676b8f2dc34fed9731f19a03e4bfa1",
            "74358c341daa44faa41a92b0eb073615",
            "c40e17a994ed44d79ef54421f4772858",
            "431ffcb0dbb1419dbbb237a29842b905",
            "c28c185a508e459b95edb29afedd9c79",
            "7a0f3d2deeb14bbba7f0df208133a7e2",
            "94cdcf1fb0914dfe99c3c615e68b7ec8",
            "f3e24aa6619e4185a3e2f53f2e473e45",
            "b7be20bd738940658f710309729d8ccd",
            "562a287082eb4882af864062ae13c91b",
            "11e99924bdf84739b3cfc7d676b797e5",
            "60fe026e908741209637604edfeb0563",
            "e7a73cfc3fa04425afc08368532ef0be",
            "d839da1c31a34659988e3420ba520725",
            "96a3b8c765374a5e98eaf713dd5e85d8",
            "b49ce0d00e5b4df097b3a9fc4c0bd8e7",
            "b2748df9aff540df954f002310606a56",
            "f4caf2bc58704464b6ac7bf6da31b71d",
            "3731a5747724403e86ef1abaf4fb389b",
            "688bf676958a4050b73cc9597f5110de",
            "8776d04bfd3748b69e10708466e9c132",
            "56490c20aef04fe4a63ba7c990e71e77",
            "b8ad433cc3dd4a8689cb06655158d545",
            "e51620404d244a4e806eeaf3b96a474f",
            "e4568f72971749ac92cc5577cb9c4f9c",
            "1a1ada8b122d4502949c4c75861a82bf",
            "18d4c2302f424bc18f9215f2d5b83393",
            "f28f557210b7402b8b2d0fd83f038dc1",
            "bc8190c700e949ebb6f874c8f8055161",
            "ceebda1d07b545f5b605551c08e73be8",
            "6aa96b80e3b746b48f5c13bcafe1134e",
            "8de63414803f41c29e75242756cfd67f",
            "f693dfed91df492c986afbd4ff925f89",
            "cef2516aabc24f6d852079984afa8090",
            "95d303d83d3348a0adfdf6b5cb94c41c",
            "9373fcd054014760b23e3acc928da0a5",
            "7e6b41ec535648b190bfc655d9e430dd",
            "d0198fa98e5f412380f2ba3e9448aef3",
            "eb0c489b913744a28304e5334bb66532",
            "c80603555c634e338f16c0021f752420",
            "0fddebfb42e847e08890d09f30161984",
            "d07a0a49a2924092affc0e07bf0b5ef0",
            "beb4231b49024d8a9bb7ead19647d99d",
            "26d09aa072d64af0bf4df9a89309093a"
          ]
        },
        "id": "1OD2z7d3HQJU",
        "outputId": "7a53864c-6965-4c25-9668-4b0d56db9f6b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GPT2LMHeadModel(\n",
            "  (transformer): GPT2Model(\n",
            "    (wte): Embedding(50257, 768)\n",
            "    (wpe): Embedding(1024, 768)\n",
            "    (drop): Dropout(p=0.1, inplace=False)\n",
            "    (h): ModuleList(\n",
            "      (0-11): 12 x GPT2Block(\n",
            "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPT2SdpaAttention(\n",
            "          (c_attn): Conv1D()\n",
            "          (c_proj): Conv1D()\n",
            "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
            "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPT2MLP(\n",
            "          (c_fc): Conv1D()\n",
            "          (c_proj): Conv1D()\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "  )\n",
            "  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
            "  (generator): WrapperModule()\n",
            ")\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/opt/homebrew/anaconda3/envs/nnsight_local/lib/python3.12/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be deprecated in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
            "  warnings.warn(\n"
          ]
        }
      ],
      "source": [
        "from nnsight import LanguageModel\n",
        "\n",
        "llm = LanguageModel(\"openai-community/gpt2\", device_map=\"auto\")\n",
        "\n",
        "print(llm)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hw2VaAr_ezkj"
      },
      "source": [
        "<details>\n",
        "<summary>On Model Initialization</summary>\n",
        "\n",
        "---\n",
        "\n",
        "A few important things to note:\n",
        "\n",
        "Keyword arguments passed to the initialization of `LanguageModel` is forwarded\n",
        "to HuggingFace specific loading logic. In this case, `device_map` specifies\n",
        "which devices to use and its value `auto` indicates to evenly distribute it to\n",
        "all available GPUs (and CPU if no GPUs available). Other arguments can be found\n",
        "here:\n",
        "https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoModelForCausalLM\n",
        "\n",
        "When we initialize `LanguageModel`, we aren't yet loading the parameters of the\n",
        "model into memory. We are actually loading a 'meta' version of the model which\n",
        "doesn't take up any memory, but still allows us to view and trace actions on it.\n",
        "After exiting the first tracing context, the model is then fully loaded into\n",
        "memory. To load into memory on initialization, you can pass `dispatch=True` into\n",
        "`LanguageModel` like\n",
        "`LanguageModel('openai-community/gpt2', device_map=\"auto\", dispatch=True)`.\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n",
        "\n",
        "Let's put together some of the features we applied to the small model, but now\n",
        "on `GPT-2`. Unlike `NNsight`, `LanguageModel` does define logic to pre-process\n",
        "inputs upon entering the tracing context. This makes interacting with the model\n",
        "simpler without having to directly access the tokenizer.\n",
        "\n",
        "In the following example, we ablate the value coming from the last layer's MLP\n",
        "module and decode the logits to see what token the model predicts without\n",
        "influence from that particular module:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mLSapaMCgLNU",
        "outputId": "192b8998-ecca-4f25-ecf2-5196c0d2d672"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Token IDs: tensor([[ 262,   12,  417, 8765,   11,  257,  262, 3504,  338, 3576]],\n",
            "       device='mps:0')\n",
            "Prediction:  London\n"
          ]
        }
      ],
      "source": [
        "with llm.trace(\"The Eiffel Tower is in the city of\"):\n",
        "\n",
        "    # Access the last layer using h[-1] as it's a ModuleList\n",
        "    # Access the first index of .output as that's where the hidden states are.\n",
        "    llm.transformer.h[-1].mlp.output[0][:] = 0\n",
        "\n",
        "    # Logits come out of model.lm_head and we apply argmax to get the predicted token ids.\n",
        "    token_ids = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"\\nToken IDs:\", token_ids)\n",
        "\n",
        "# Apply the tokenizer to decode the ids into words after the tracing context.\n",
        "print(\"Prediction:\", llm.tokenizer.decode(token_ids[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X2-HiCQhkv-B"
      },
      "source": [
        "We just ran a little intervention on a much more complex model with a lot more\n",
        "parameters! An important piece of information we're missing though is what the\n",
        "prediction would look like without our ablation.\n",
        "\n",
        "Of course we could just run two tracing contexts and compare the outputs. This,\n",
        "however, would require two forward passes through the model. `NNsight` can do\n",
        "better than that.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vLjHnRyRPXjp"
      },
      "source": [
        "<a name=\"batching-id\"></a>\n",
        "\n",
        "## Batching\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Gc-Zr6NmEOn"
      },
      "source": [
        "It's time to bring back the `Tracer` object we dropped before. \n",
        "\n",
        "See, when we call `.trace(...)`, it's actually creating two different\n",
        "contexts behind the scenes. The second one is the invoker context. The invoker context defines\n",
        "the values of `.input` and `.output` Proxies. \n",
        "\n",
        "If we call `.trace(...)` with some input, the input is passed on to the invoker. \n",
        "Since there is only one input – only one invoker context is created.\n",
        "\n",
        "If we call `.trace()` without input then we can call `tracer.invoke(...)` to manually\n",
        "create the invoker context with our input. Now every subsequent time we call\n",
        "`.invoke(...)`, new interventions will only refer to the input in that\n",
        "particular invoke. When exiting the tracing context, the inputs from all of the\n",
        "invokers will be batched together, and they will be executed in one forward\n",
        "pass! So let's do the ablation experiment, and compute a 'control' output to\n",
        "compare to:\n",
        "\n",
        "<details>\n",
        "<summary>More on the invoker context</summary>\n",
        "\n",
        "---\n",
        "\n",
        "Note that when injecting data to only the relevant invoker interventions,\n",
        "`nnsight` tries, but can't guarantee, that it can narrow the data into the right\n",
        "batch idxs. So there are cases\n",
        "where all invokes will get all of the data. Specifically, if the input or output data is stored \n",
        "as an object that is not an arbitrary collection of tensors, it will be broadcasted to all invokes.\n",
        "\n",
        "Just like `.trace(...)` created a `Tracer` object, `.invoke(...)` creates an\n",
        "`Invoker` object. For `LaguageModel` models, the `Invoker` prepares the input by running a tokenizer on it.\n",
        "`Invoker` stores pre-processed inputs at `invoker.inputs`, which can be accessed to see information about our inputs.\n",
        "In case when we are passing a single input to `.trace(...)` directly, we can still access the invoker\n",
        "object at `tracer.invoker` without having to call `tracer.invoke(...)`.\n",
        "\n",
        "Keyword arguments given to `.invoke(..)` make its way to the input pre-processing.  \n",
        "`LanguageModel` has keyword arguments `max_length` and `truncation` used for tokenization, and they can be \n",
        "passed to the invoker. If we are calling a single-input `.trace(...)` and want to pass the keyword arguments, \n",
        "we can do so in the form of `invoker_args` that should be a dictionary of keyword arguments for the invoker. \n",
        "\n",
        "Here is an example to demonstrate everything we've described:\n",
        "\n",
        "**This snippet**\n",
        "\n",
        "```\n",
        "with llm.trace(\"hello\", invoker_args={\"max_length\":10}) as tracer:\n",
        "  invoker = tracer.invoker\n",
        "\n",
        "```\n",
        "  **does the same as**\n",
        "  \n",
        "\n",
        "```\n",
        "with llm.trace() as tracer:\n",
        "  with tracer.invoke(\"hello\", max_length=10) as invoker:\n",
        "    invoker = invoker\n",
        "```\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Kdcq4oCNmEua",
        "outputId": "e6d33c07-d22a-41af-b84a-91c9681c7524"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Original token IDs: tensor([[ 198,   12,  417, 8765,  318,  257,  262, 3504, 7372, 6342]],\n",
            "       device='mps:0')\n",
            "Modified token IDs: tensor([[ 262,   12,  417, 8765,   11,  257,  262, 3504,  338, 3576]],\n",
            "       device='mps:0')\n",
            "Original prediction:  Paris\n",
            "Modified prediction:  London\n"
          ]
        }
      ],
      "source": [
        "with llm.trace() as tracer:\n",
        "\n",
        "    with tracer.invoke(\"The Eiffel Tower is in the city of\"):\n",
        "\n",
        "        # Ablate the last MLP for only this batch.\n",
        "        llm.transformer.h[-1].mlp.output[0][:] = 0\n",
        "\n",
        "        # Get the output for only the intervened on batch.\n",
        "        token_ids_intervention = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "    with tracer.invoke(\"The Eiffel Tower is in the city of\"):\n",
        "\n",
        "        # Get the output for only the original batch.\n",
        "        token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"Original token IDs:\", token_ids_original)\n",
        "print(\"Modified token IDs:\", token_ids_intervention)\n",
        "\n",
        "print(\"Original prediction:\", llm.tokenizer.decode(token_ids_original[0][-1]))\n",
        "print(\"Modified prediction:\", llm.tokenizer.decode(token_ids_intervention[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8BsybgsK2Bbr"
      },
      "source": [
        "So it did end up affecting what the model predicted. That's pretty neat!\n",
        "\n",
        "Another cool thing with multiple invokes is that the Proxies can interact\n",
        "between them. Here we transfer the word token embeddings from a real prompt into\n",
        "another placeholder prompt. Therefore the latter prompt produces the output of\n",
        "the former prompt:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "syKTI_KhpvCY",
        "outputId": "b2238c5c-d2d8-47bf-b73a-730f1e03cf99"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Original prediction:  _\n",
            "Modified prediction:  Paris\n"
          ]
        }
      ],
      "source": [
        "with llm.trace() as tracer:\n",
        "\n",
        "    with tracer.invoke(\"The Eiffel Tower is in the city of\"):\n",
        "\n",
        "        embeddings = llm.transformer.wte.output\n",
        "\n",
        "    with tracer.invoke(\"_ _ _ _ _ _ _ _ _ _\"):\n",
        "\n",
        "        llm.transformer.wte.output = embeddings\n",
        "\n",
        "        token_ids_intervention = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "    with tracer.invoke(\"_ _ _ _ _ _ _ _ _ _\"):\n",
        "\n",
        "        token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"Original prediction:\", llm.tokenizer.decode(token_ids_original[0][-1]))\n",
        "print(\"Modified prediction:\", llm.tokenizer.decode(token_ids_intervention[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vvWA-CWqQtah"
      },
      "source": [
        "## .next()\n",
        "\n",
        "Some HuggingFace models define methods to generate multiple outputs at a time.\n",
        "`LanguageModel` wraps that functionality to provide the same tracing features by\n",
        "using `.generate(...)` instead of `.trace(...)`. This calls the underlying\n",
        "model's `.generate` method. It passes the output through a `.generator`\n",
        "module that we've added onto the model, allowing us to get the generate output\n",
        "at `.generator.output`.\n",
        "\n",
        "In a case like this, the underlying model is called more than once; the modules\n",
        "of said model produce more than one output. Which iteration should a given\n",
        "`module.output` refer to? That's where `Module.next()` comes in.\n",
        "\n",
        "Each module has a call idx associated with it and `.next()` simply increments\n",
        "that attribute. At the time of execution, data is injected into the intervention\n",
        "graph only at the iteration that matches the call idx.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yy5Z9NE1GkaN",
        "outputId": "0b4c1bd3-a6f6-4ba1-9cb8-6019c222311f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Prediction 1:   Paris\n",
            "Prediction 2:  ,\n",
            "Prediction 3:   and\n",
            "All token ids:  tensor([[ 464,  412,  733,  417, 8765,  318,  287,  262, 1748,  286, 6342,   11,\n",
            "          290]], device='mps:0')\n",
            "Full prediction:  ['The Eiffel Tower is in the city of Paris, and']\n"
          ]
        }
      ],
      "source": [
        "with llm.generate(\"The Eiffel Tower is in the city of\", max_new_tokens=3):\n",
        "\n",
        "    token_ids_1 = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "    token_ids_2 = llm.lm_head.next().output.argmax(dim=-1).save()\n",
        "\n",
        "    token_ids_3 = llm.lm_head.next().output.argmax(dim=-1).save()\n",
        "\n",
        "    output = llm.generator.output.save()\n",
        "\n",
        "print(\"\\nPrediction 1: \", llm.tokenizer.decode(token_ids_1[0][-1]))\n",
        "print(\"Prediction 2: \", llm.tokenizer.decode(token_ids_2[0][-1]))\n",
        "print(\"Prediction 3: \", llm.tokenizer.decode(token_ids_3[0][-1]))\n",
        "\n",
        "print(\"All token ids: \", output)\n",
        "\n",
        "print(\"Full prediction: \", llm.tokenizer.batch_decode(output.value))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Model Editing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "NNsight's model editing feature allows you to create persistently modified versions of a model with a use of `.edit()`. Unlike interventions in a tracing context, which are temporary, the **Editor** context enables you to make lasting changes to a model instance.\n",
        "\n",
        "This feature is useful for:\n",
        "* Creating modified model variants without altering the original\n",
        "* Applying changes that persist across multiple forward passes\n",
        "* Comparing interventions between original and edited models\n",
        "\n",
        "Let's explore how to use the **Editor** context to make a simple persistent change to a model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Original Prediction:   Rome\n",
            "Modified Prediction:   Paris\n"
          ]
        }
      ],
      "source": [
        "# we take the hidden states with the expected output \"Paris\"\n",
        "with llm.trace(\"The Eiffel Tower is located in the city of\") as tracer:\n",
        "    hs11 = llm.transformer.h[11].output[0][:, -1, :].save()\n",
        "\n",
        "# the edited model will now always predict \"Paris\" as the next token\n",
        "with llm.edit() as llm_edited:\n",
        "    llm.transformer.h[11].output[0][:, -1, :] = hs11\n",
        "\n",
        "# we demonstrate this by comparing the output of an unmodified model...\n",
        "with llm.trace(\"Vatican is located in the city of\") as tracer:\n",
        "    original_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        " \n",
        "# ...with the output of the edited model\n",
        "with llm_edited.trace(\"Vatican is located in the city of\") as tracer:\n",
        "    modified_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "\n",
        "print(\"\\nOriginal Prediction: \", llm.tokenizer.decode(original_tokens[0][-1]))\n",
        "print(\"Modified Prediction: \", llm.tokenizer.decode(modified_tokens[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Edits defined within an **Editor** context create a new, modified version of the model by default, preserving the original. This allows for safe experimentation with model changes. If you wish to modify the original model directly, you can set `inplace=True` when calling `.edit()`. \n",
        "\n",
        "Use this option cautiously, as in-place edits alter the base model for all the consequent model calls."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Modified In-place:   Paris\n"
          ]
        }
      ],
      "source": [
        "# we use the hidden state we saved above (hs11)\n",
        "with llm.edit(inplace=True) as llm_edited:\n",
        "    llm.transformer.h[11].output[0][:, -1, :] = hs11\n",
        "\n",
        "# we demonstrate this by comparing the output of an unmodified model...\n",
        "with llm.trace(\"Vatican is located in the city of\") as tracer:\n",
        "    modified_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"Modified In-place: \", llm.tokenizer.decode(modified_tokens[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If you've made in-place edits to your model and need to revert these changes, `.clear_edits()` can help. This method removes all edits applied to the model, effectively restoring it to its original state. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Edits cleared:   Rome\n"
          ]
        }
      ],
      "source": [
        "llm.clear_edits()\n",
        "\n",
        "with llm.trace(\"Vatican is located in the city of\"):\n",
        "    modified_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"Edits cleared: \", llm.tokenizer.decode(modified_tokens[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rpkX-LwBQZHo"
      },
      "source": [
        "# 3 I thought you said huge models?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uCZR65VmILEb"
      },
      "source": [
        "`NNsight` is only one part of our project to democratize access to AI internals.\n",
        "The other half is `NDIF` (National Deep Inference Fabric).\n",
        "\n",
        "The interaction between the two is fairly straightforward. The\n",
        "**intervention graph** we create via the tracing context can be encoded into a\n",
        "custom json format and sent via an http request to the `NDIF` servers. `NDIF`\n",
        "then decodes the **intervention graph** and **interleaves** it alongside the\n",
        "specified model.\n",
        "\n",
        "To see which models are currently being hosted, check out the following status\n",
        "page: https://nnsight.net/status/\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65Ks1LUvQaER"
      },
      "source": [
        "## Remote execution\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Aa-ZuVOFKS5k"
      },
      "source": [
        "In its current state, `NDIF` requires you to receive an API key. Therefore, to\n",
        "run the rest of this colab, you would need one of your own. To get one, simply\n",
        "go to https://login.ndif.us and sign up.\n",
        "\n",
        "With a valid API key, you then can configure `nnsight` by doing the following:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "id": "ax1NWoS9MJeZ"
      },
      "outputs": [],
      "source": [
        "from nnsight import CONFIG\n",
        "\n",
        "CONFIG.set_default_api_key(\"YOUR_API_KEY\") \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mvXLOh65MXmF"
      },
      "source": [
        "This only needs to be run once as it will save this api key as the default in a\n",
        "config file along with the `nnsight` installation.\n",
        "\n",
        "To amp things up a few levels, let's demonstrate using `nnsight`'s tracing\n",
        "context with one of the larger open source language models, `Llama-3.1-70b`!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "id": "Q1N04sJPZnJt"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# llama3.1 70b is a gated model and you need access via your huggingface token\n",
        "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\"\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "oqVzjNoyNGc7"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:11:21,150 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:11:21,184 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:11:21,206 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - RUNNING: Your job has started running.\n",
            "2024-08-30 07:11:21,398 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - COMPLETED: Your job has been completed.\n",
            "Downloading result:   0%|          | 0.00/9.48M [00:00<?, ?B/s]huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
            "To disable this warning, you can either:\n",
            "\t- Avoid using `tokenizers` before the fork if possible\n",
            "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
            "Downloading result: 100%|██████████| 9.48M/9.48M [00:02<00:00, 3.21MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(tensor([[[ 5.4688, -4.9062,  2.2344,  ..., -3.6875,  0.9609,  1.2578],\n",
            "         [ 1.5469, -0.6172, -1.4531,  ..., -1.1562, -0.1406, -2.1250],\n",
            "         [ 1.7812, -1.8906, -1.1875,  ...,  0.1680,  0.9609,  0.5625],\n",
            "         ...,\n",
            "         [ 0.9453, -0.3711,  1.3516,  ...,  1.3828, -0.7969, -1.9297],\n",
            "         [-0.8906,  0.3672,  0.2617,  ...,  2.4688, -0.4414, -0.6758],\n",
            "         [-1.6094,  1.0938,  1.7031,  ...,  1.8672, -1.1328, -0.5000]]],\n",
            "       dtype=torch.bfloat16), DynamicCache())\n",
            "tensor([[[ 6.3750,  8.6250, 13.0000,  ..., -4.1562, -4.1562, -4.1562],\n",
            "         [-2.8594, -2.2344, -3.0938,  ..., -8.6250, -8.6250, -8.6250],\n",
            "         [ 8.9375,  3.5938,  4.5000,  ..., -3.9375, -3.9375, -3.9375],\n",
            "         ...,\n",
            "         [ 3.5781,  3.4531,  0.0796,  ..., -6.5625, -6.5625, -6.5625],\n",
            "         [10.8750,  6.4062,  4.9375,  ..., -4.0000, -4.0000, -3.9844],\n",
            "         [ 7.2500,  6.1562,  3.5156,  ..., -4.7188, -4.7188, -4.7188]]])\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# We'll never actually load the parameters so no need to specify a device_map.\n",
        "llama = LanguageModel(\"meta-llama/Meta-Llama-3.1-70B\")\n",
        "\n",
        "# All we need to specify using NDIF vs executing locally is remote=True.\n",
        "with llama.trace(\"The Eiffel Tower is in the city of\", remote=True) as runner:\n",
        "\n",
        "    hidden_states = llama.model.layers[-1].output.save()\n",
        "\n",
        "    output = llama.output.save()\n",
        "\n",
        "print(hidden_states)\n",
        "\n",
        "print(output[\"logits\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IceVpnZcZ9F0"
      },
      "source": [
        "It really is as simple as `remote=True`. All of the techniques we went through\n",
        "in earlier sections work just the same when running locally or remotely."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sessions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "NDIF uses a queue to handle concurrent requests from multiple users. To optimize the execution of our experiments we can use the `session` context to efficiently package multiple interventions together as one single request to the server.\n",
        "\n",
        "This offers the following benefits: \n",
        "1) All interventions within a session will be executed one after another without additional wait in the queue\n",
        "2) All intermediate outputs of each intervention are stored on the server and can be accessed by other interventions in the same session without moving the data back and forth between NDIF and the local machine.\n",
        "\n",
        "Let's take a look:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:11:28,206 MainProcess nnsight_remote INFO     4a6576dd-b5fd-4f1f-9836-a619f8277057 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:11:28,206 MainProcess nnsight_remote INFO     4a6576dd-b5fd-4f1f-9836-a619f8277057 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:11:28,207 MainProcess nnsight_remote INFO     4a6576dd-b5fd-4f1f-9836-a619f8277057 - RUNNING: Your job has started running.\n",
            "2024-08-30 07:11:28,207 MainProcess nnsight_remote INFO     4a6576dd-b5fd-4f1f-9836-a619f8277057 - LOG: 80\n",
            "2024-08-30 07:11:28,416 MainProcess nnsight_remote INFO     4a6576dd-b5fd-4f1f-9836-a619f8277057 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 1.69k/1.69k [00:00<00:00, 6.30MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "T1 - Original Prediction:   Paris\n",
            "T2 - Modified Prediction:   Paris\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  with llama.trace(\"The Eiffel Tower is in the city of\") as t1:\n",
        "    # capture the hidden state from layer 11 at the last token\n",
        "    hs_79 = llama.model.layers[79].output[0][:, -1, :] # no .save()\n",
        "    t1_tokens_out = llama.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "  with llama.trace(\"Buckingham Palace is in the city of\") as t2:\n",
        "    llama.model.layers[1].output[0][:, -1, :] = hs_79[:]\n",
        "    t2_tokens_out = llama.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"\\nT1 - Original Prediction: \", llama.tokenizer.decode(t1_tokens_out[0][-1]))\n",
        "print(\"T2 - Modified Prediction: \", llama.tokenizer.decode(t2_tokens_out[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In the example above, we are interested in replacing the hidden state of a later layer with an earlier one. Since we are using a `session`, we don't have to save the hidden state from Tracer 1 to reference it in Tracer 2.\n",
        "\n",
        "It is important to note that all the traces defined within the `session` context are executed sequentially, strictly following the order of definition (i.e. `t2` being executed after `t1` and `t3` after `t2` etc.)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `session` context object has its own methods to log values and be terminated early."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:11:30,900 MainProcess nnsight_remote INFO     28ac8e47-fa48-45a1-acb0-3e17960e36b8 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:11:30,934 MainProcess nnsight_remote INFO     28ac8e47-fa48-45a1-acb0-3e17960e36b8 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:11:30,935 MainProcess nnsight_remote INFO     28ac8e47-fa48-45a1-acb0-3e17960e36b8 - RUNNING: Your job has started running.\n",
            "2024-08-30 07:11:30,951 MainProcess nnsight_remote INFO     28ac8e47-fa48-45a1-acb0-3e17960e36b8 - LOG: -- Early Stop --\n",
            "2024-08-30 07:11:30,953 MainProcess nnsight_remote INFO     28ac8e47-fa48-45a1-acb0-3e17960e36b8 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 4.11MB/s]\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  session.log(\"-- Early Stop --\")\n",
        "  session.exit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In addition to the benefits mentioned above, the `session` context also enables interesting experiments not possible with other `nnsight` tools - since every trace is run on its own model, it means that within one session we can run interventions between different models – for example, we can swap activations between vanilla and instruct versions of the Llama model and compare the outputs. And `session` can also be used to run experiments entirely locally! "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Looping"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We mention earlier that the `session` context enables multi-tracing execution. But how can we optimize a process that would require running an intervention graph in a loop? If we create a simple `for` loop with a **Tracer context** inside of it, this will result in creating a new intervention graph at each iteration, which is not scalable. \n",
        "\n",
        "We solve this problem the `nnsight` way via the **Iterator context**: an intervention loop that iteratively executes and updates a single intervention graph.\n",
        "\n",
        "Use a `session` to define the **Iterator context** and pass in a sequence of items that you want to loop over at each iteration:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:11:34,689 MainProcess nnsight_remote INFO     667c4310-9041-451d-99a5-f713f639abb8 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:11:34,708 MainProcess nnsight_remote INFO     667c4310-9041-451d-99a5-f713f639abb8 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:11:34,726 MainProcess nnsight_remote INFO     667c4310-9041-451d-99a5-f713f639abb8 - RUNNING: Your job has started running.\n",
            "2024-08-30 07:11:35,332 MainProcess nnsight_remote INFO     667c4310-9041-451d-99a5-f713f639abb8 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 6.49MB/s]\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  with session.iter([0, 1, 2]) as item:\n",
        "    # define intervention body here ...\n",
        "\n",
        "    with llama.trace(\"_\"):\n",
        "      # define interventions here ...\n",
        "      pass\n",
        "\n",
        "    with llama.trace(\"_\"):\n",
        "      # define interventions here ...\n",
        "      pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `Iterator` context extends all the `nnsight` graph-based functionalities, but also closely mimics the conventional `for` loop statement in Python, which allows it to support all kind of iterative operations with a use of `as item` syntax:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:11:36,315 MainProcess nnsight_remote INFO     8734ee75-b616-4270-9d4d-7cfabd6d63ae - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:11:36,334 MainProcess nnsight_remote INFO     8734ee75-b616-4270-9d4d-7cfabd6d63ae - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:11:36,342 MainProcess nnsight_remote INFO     8734ee75-b616-4270-9d4d-7cfabd6d63ae - RUNNING: Your job has started running.\n",
            "2024-08-30 07:11:36,354 MainProcess nnsight_remote INFO     8734ee75-b616-4270-9d4d-7cfabd6d63ae - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 1.06k/1.06k [00:00<00:00, 13.7MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "List:  [0, 1, 2]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "\n",
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  li = nnsight.list()\n",
        "  [li.append([num]) for num in range(0, 3)] # adding [0], [1], [2] to the list\n",
        "  li2 = nnsight.list().save()\n",
        "  \n",
        "  # You can create nested Iterator contexts\n",
        "  with session.iter(li) as item:\n",
        "    with session.iter(item) as item_2:\n",
        "      li2.append(item_2)\n",
        "\n",
        "print(\"\\nList: \", li2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Notice how we used the `nnsight.list()` method to create a list of lists to loop over. This type of method is what we call an **NNsight Built-in**. It is a special type of methods that serve as a wrapper around `nnsight.apply()` to provide a more user-friendly interface for adding common datatypes to the Intervention Graph.\n",
        "\n",
        "<details>\n",
        "<summary>A full list of NNsight Built-Ins</summary>\n",
        "\n",
        "\n",
        "`nnsight.bool()` creates a traceable Boolean\n",
        "\n",
        "`nnsight.bytes()` creates a traceable Bytes\n",
        "\n",
        "`nnsight.int()` creates a traceable Integer\n",
        "\n",
        "`nnsight.float()` creates a traceable Float\n",
        "\n",
        "`nnsight.str()` creates a traceable String\n",
        "\n",
        "`nnsight.comples()` creates a traceable Complex number\n",
        "\n",
        "`nnsight.bytearray()` creates a traceable Bytearray\n",
        "\n",
        "`nnsight.tuple()` creates a traceable Tuple\n",
        "\n",
        "`nnsight.list()` creates a traceable List\n",
        "\n",
        "`nnsight.set()` creates a traceable Set\n",
        "\n",
        "`nnsight.dict()` creates a traceable Dictionary\n",
        "\n",
        "</details>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can also expose the `iterator` context object via a `return_context` flag. You can then use it to `exit` out of the Iteration loop early and log the intermediate outputs within the loop:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:11:38,551 MainProcess nnsight_remote INFO     5f0b434d-178d-4807-913f-331f403eb0ea - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:11:38,571 MainProcess nnsight_remote INFO     5f0b434d-178d-4807-913f-331f403eb0ea - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:11:38,593 MainProcess nnsight_remote INFO     5f0b434d-178d-4807-913f-331f403eb0ea - RUNNING: Your job has started running.\n",
            "2024-08-30 07:11:38,594 MainProcess nnsight_remote INFO     5f0b434d-178d-4807-913f-331f403eb0ea - LOG: 0\n",
            "2024-08-30 07:11:38,610 MainProcess nnsight_remote INFO     5f0b434d-178d-4807-913f-331f403eb0ea - LOG: 1\n",
            "2024-08-30 07:11:38,611 MainProcess nnsight_remote INFO     5f0b434d-178d-4807-913f-331f403eb0ea - LOG: 2\n",
            "2024-08-30 07:11:38,630 MainProcess nnsight_remote INFO     5f0b434d-178d-4807-913f-331f403eb0ea - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 992/992 [00:00<00:00, 5.68MB/s]\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "  \n",
        "  with session.iter([0, 1, 2, 3], return_context=True) as (item, iterator):\n",
        "      \n",
        "      iterator.log(item)\n",
        "      \n",
        "      with iterator.cond(item == 2):\n",
        "        iterator.exit()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The **Iterator** context is a niece piece of functionality that allows you to define a bunch of basic code operations that can now be \"traceable\" by `nnsight`.\n",
        "\n",
        "But in what kind of experimental scenario would someone even need to use it?\n",
        "\n",
        "In the next section, we delve into a powerful use case of the `Iterator` context and see how it enables it!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Training a LoRA"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Here is an example of a task that uses everything we have covered in the last section - remote execution, **Session** context and iterative interventions. We're going to train a very simple LORA that, when applied, will make our model always predict \"Paris\" no matter what."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 43,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "from nnsight.envoy import Envoy # \n",
        "\n",
        "# We will define a LORA class.\n",
        "# The LORA class call method operations are simply traced like you would normally do in a .trace.\n",
        "class LORA(nn.Module):\n",
        "    def __init__(self, module: Envoy, dim: int, r: int) -> None:\n",
        "        \"\"\"Init.\n",
        "\n",
        "        Args:\n",
        "            module (Envoy): Which model Module we are adding the LORA to.\n",
        "            dim (int): Dimension of the layer we are adding to (This could potentially be auto populated if the user scanned first so we know the shape)\n",
        "            r (int): Inner dimension of the LORA\n",
        "        \"\"\"\n",
        "        super(LORA, self).__init__()\n",
        "        self.r = r\n",
        "        self.module = module\n",
        "        self.WA = torch.nn.Parameter(torch.randn(dim, self.r), requires_grad=True).save()\n",
        "        self.WB = torch.nn.Parameter(torch.zeros(self.r, dim), requires_grad=True).save()\n",
        "\n",
        "    # The Call method defines how to actually apply the LORA.\n",
        "    def __call__(self, alpha: float = 1.0):\n",
        "        \"\"\"Call.\n",
        "\n",
        "        Args:\n",
        "            alpha (float, optional): How much to apply the LORA. Can be altered after training for inference. Defaults to 1.0.\n",
        "        \"\"\"\n",
        "\n",
        "        # We apply WA to the first positional arg (the hidden states)\n",
        "        A_x = torch.matmul(self.module.input[0][0], self.WA)\n",
        "        BA_x = torch.matmul(A_x, self.WB)\n",
        "\n",
        "        # LORA is additive\n",
        "        h = BA_x + self.module.output\n",
        "\n",
        "        # Replace the output with our new one * alpha\n",
        "        # Could also have been self.module.output[:] = h * alpha, for in-place\n",
        "        self.module.output = h * alpha\n",
        "\n",
        "    def parameters(self):\n",
        "        # Some way to get all the parameters.\n",
        "        return [self.WA, self.WB]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's define all the variables to use in LORA training. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "metadata": {},
      "outputs": [],
      "source": [
        "# We need the token id of the correct answer.\n",
        "answer = \" Paris\"\n",
        "answer_token = llama.tokenizer.encode(answer)[1]\n",
        "# Inner LORA dimension\n",
        "lora_dim = 4\n",
        "# Module to train LORA on\n",
        "module = llama.model.layers[-1].mlp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can use the `.scan()` method to get the shape of the module without having to fully run the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "8192\n"
          ]
        }
      ],
      "source": [
        "with llama.scan(\" \"):\n",
        "    dim = module.output.shape[-1]\n",
        "\n",
        "print(dim)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "It's time to run the LORA training loop! We using the **Session** and the **Iterator** contexts to achieve this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:12:21,091 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:12:21,146 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:12:21,166 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - RUNNING: Your job has started running.\n",
            "2024-08-30 07:12:21,704 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.6289, -1.1172, -0.6719,  0.1816],\n",
            "        [-0.2715,  0.5547, -1.2812, -0.8086],\n",
            "        [ 1.0938, -0.5820, -0.7070, -1.1094],\n",
            "        ...,\n",
            "        [-0.2910, -0.6016,  0.6602, -0.4590],\n",
            "        [-2.7969, -0.3477, -1.3438, -1.1797],\n",
            "        [-1.0312, -1.0469, -0.7930,  0.4141]], requires_grad=True)\n",
            "2024-08-30 07:12:21,853 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.6094, -1.0859, -0.6523,  0.1758],\n",
            "        [-0.2637,  0.5391, -1.2422, -0.7852],\n",
            "        [ 1.0625, -0.5664, -0.6875, -1.0781],\n",
            "        ...,\n",
            "        [-0.2832, -0.5820,  0.6406, -0.4453],\n",
            "        [-2.7188, -0.3379, -1.3047, -1.1406],\n",
            "        [-1.0000, -1.0156, -0.7695,  0.4023]], requires_grad=True)\n",
            "2024-08-30 07:12:21,983 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.5898, -1.0547, -0.6328,  0.1709],\n",
            "        [-0.2559,  0.5234, -1.2031, -0.7617],\n",
            "        [ 1.0312, -0.5508, -0.6680, -1.0469],\n",
            "        ...,\n",
            "        [-0.2754, -0.5664,  0.6211, -0.4316],\n",
            "        [-2.6406, -0.3281, -1.2656, -1.1094],\n",
            "        [-0.9688, -0.9844, -0.7461,  0.3906]], requires_grad=True)\n",
            "2024-08-30 07:12:22,102 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.5703, -1.0234, -0.6133,  0.1660],\n",
            "        [-0.2480,  0.5078, -1.1641, -0.7383],\n",
            "        [ 1.0000, -0.5352, -0.6484, -1.0156],\n",
            "        ...,\n",
            "        [-0.2676, -0.5508,  0.6016, -0.4180],\n",
            "        [-2.5625, -0.3184, -1.2266, -1.0781],\n",
            "        [-0.9414, -0.9531, -0.7227,  0.3789]], requires_grad=True)\n",
            "2024-08-30 07:12:22,218 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.5547, -0.9922, -0.5938,  0.1611],\n",
            "        [-0.2402,  0.4922, -1.1328, -0.7148],\n",
            "        [ 0.9688, -0.5195, -0.6289, -0.9844],\n",
            "        ...,\n",
            "        [-0.2598, -0.5352,  0.5820, -0.4062],\n",
            "        [-2.4844, -0.3086, -1.1875, -1.0469],\n",
            "        [-0.9141, -0.9258, -0.6992,  0.3672]], requires_grad=True)\n",
            "2024-08-30 07:12:22,374 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.5391, -0.9609, -0.5742,  0.1562],\n",
            "        [-0.2334,  0.4766, -1.1016, -0.6953],\n",
            "        [ 0.9414, -0.5039, -0.6094, -0.9531],\n",
            "        ...,\n",
            "        [-0.2520, -0.5195,  0.5664, -0.3945],\n",
            "        [-2.4062, -0.2988, -1.1484, -1.0156],\n",
            "        [-0.8867, -0.8984, -0.6797,  0.3555]], requires_grad=True)\n",
            "2024-08-30 07:12:22,755 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.5234, -0.9336, -0.5586,  0.1514],\n",
            "        [-0.2266,  0.4629, -1.0703, -0.6758],\n",
            "        [ 0.9141, -0.4883, -0.5898, -0.9258],\n",
            "        ...,\n",
            "        [-0.2441, -0.5039,  0.5508, -0.3828],\n",
            "        [-2.3281, -0.2891, -1.1172, -0.9844],\n",
            "        [-0.8594, -0.8711, -0.6602,  0.3457]], requires_grad=True)\n",
            "2024-08-30 07:12:22,757 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.5078, -0.9062, -0.5430,  0.1465],\n",
            "        [-0.2197,  0.4492, -1.0391, -0.6562],\n",
            "        [ 0.8867, -0.4727, -0.5703, -0.8984],\n",
            "        ...,\n",
            "        [-0.2363, -0.4883,  0.5352, -0.3711],\n",
            "        [-2.2656, -0.2812, -1.0859, -0.9531],\n",
            "        [-0.8320, -0.8438, -0.6406,  0.3359]], requires_grad=True)\n",
            "2024-08-30 07:12:22,757 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.4922, -0.8789, -0.5273,  0.1426],\n",
            "        [-0.2129,  0.4355, -1.0078, -0.6367],\n",
            "        [ 0.8594, -0.4590, -0.5547, -0.8711],\n",
            "        ...,\n",
            "        [-0.2295, -0.4727,  0.5195, -0.3594],\n",
            "        [-2.2031, -0.2734, -1.0547, -0.9258],\n",
            "        [-0.8086, -0.8203, -0.6211,  0.3262]], requires_grad=True)\n",
            "2024-08-30 07:12:22,900 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - LOG: Parameter containing:\n",
            "tensor([[-0.4766, -0.8516, -0.5117,  0.1387],\n",
            "        [-0.2061,  0.4219, -0.9766, -0.6172],\n",
            "        [ 0.8320, -0.4453, -0.5391, -0.8438],\n",
            "        ...,\n",
            "        [-0.2227, -0.4590,  0.5039, -0.3477],\n",
            "        [-2.1406, -0.2656, -1.0234, -0.8984],\n",
            "        [-0.7852, -0.7969, -0.6016,  0.3164]], requires_grad=True)\n",
            "2024-08-30 07:12:22,901 MainProcess nnsight_remote INFO     3e78f88a-e620-4679-ac73-abeb4f14ce8e - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 133k/133k [00:00<00:00, 571kB/s]\n"
          ]
        }
      ],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "\n",
        "# The LORA object itself isn't transmitted to the server. Only the forward / call method. \n",
        "# The parameters are created remotely and never sent only retrieved \n",
        "with llama.session(remote=True) as session:\n",
        "\n",
        "    # Create dataset of 100 pairs of a blank prompt and the \" Paris \" id\n",
        "    dataset = [[\"_\", answer_token]] * 100\n",
        "\n",
        "    # Create a dataloader from it.\n",
        "    dataloader = DataLoader(dataset, batch_size=10)\n",
        "\n",
        "    # Create our LORA on the last mlp\n",
        "    lora = LORA(module, dim, lora_dim)\n",
        "\n",
        "    # Create an optimizer. Use the parameters from LORA\n",
        "    optimizer = torch.optim.AdamW(lora.parameters(), lr=3)\n",
        "\n",
        "    # Iterate over dataloader using .iter. \n",
        "    with session.iter(dataloader, return_context=True) as (batch, iterator):\n",
        "\n",
        "        prompt = batch[0]\n",
        "        correct_token = batch[1]\n",
        "\n",
        "        # Run .trace with prompt\n",
        "        with llama.trace(prompt) as tracer:\n",
        "\n",
        "            # Apply LORA to intervention graph just by calling it with .trace\n",
        "            lora()\n",
        "\n",
        "            # Get logits\n",
        "            logits = llama.lm_head.output\n",
        "\n",
        "            # Do cross entropy on last predicted token and correct_token\n",
        "            loss = torch.nn.functional.cross_entropy(logits[:, -1], batch[1])\n",
        "            # Call backward\n",
        "            loss.backward()\n",
        "\n",
        "        # Call methods on optimizer. Graphs that arent from .trace (so in this case session and iterator both have their own graph) are executed sequentially.\n",
        "        # The Graph of Iterator here will be:\n",
        "        # 1.) Index batch at 0 for prompt\n",
        "        # 2.) Index batch at 1 for correct_token\n",
        "        # 3.) Execute the .trace using the prompt\n",
        "        # 4.) Call .step() on optimizer\n",
        "        optimizer.step()\n",
        "        # 5.) Call .zero_grad() in optimizer\n",
        "        optimizer.zero_grad()\n",
        "        # 6.) Print out the lora WA weights to show they are indeed changing\n",
        "        iterator.log(lora.WA)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now `WA` and `WB` are optimized! So we generate with the lora just by calling `lora()` in the `.generate` and save the output to then de-tokenize it. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:12:41,410 MainProcess nnsight_remote INFO     c1c6e24c-9f3f-415b-8f90-de8404fc2e74 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:12:41,539 MainProcess nnsight_remote INFO     c1c6e24c-9f3f-415b-8f90-de8404fc2e74 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:12:41,572 MainProcess nnsight_remote INFO     c1c6e24c-9f3f-415b-8f90-de8404fc2e74 - RUNNING: Your job has started running.\n",
            "2024-08-30 07:12:41,695 MainProcess nnsight_remote INFO     c1c6e24c-9f3f-415b-8f90-de8404fc2e74 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 1.31k/1.31k [00:00<00:00, 7.88MB/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['<|begin_of_text|>Hello Paris']\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-08-30 07:12:42,203 MainProcess nnsight_remote INFO     88838ccd-230d-485b-8f57-8399453c2250 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 07:12:42,224 MainProcess nnsight_remote INFO     88838ccd-230d-485b-8f57-8399453c2250 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 07:12:42,231 MainProcess nnsight_remote INFO     88838ccd-230d-485b-8f57-8399453c2250 - RUNNING: Your job has started running.\n",
            "2024-08-30 07:12:42,350 MainProcess nnsight_remote INFO     88838ccd-230d-485b-8f57-8399453c2250 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 1.31k/1.31k [00:00<00:00, 2.93MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['<|begin_of_text|>Hello.']\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# With lora. Should produce \"Hello Paris\"\n",
        "with llama.generate(\"Hello\", remote=True) as generator:\n",
        "\n",
        "    lora()\n",
        "\n",
        "    out = llama.generator.output.save()\n",
        "\n",
        "print(llama.tokenizer.batch_decode(out.value))\n",
        "\n",
        "# Then without. Should produce \"Hello,\"\n",
        "with llama.generate(\"Hello\", remote=True) as generator:\n",
        "\n",
        "    out = llama.generator.output.save()\n",
        "\n",
        "print(llama.tokenizer.batch_decode(out.value))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P3zRm-7VRRov"
      },
      "source": [
        "# Getting Involved!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BnCc9xgjvxEP"
      },
      "source": [
        "Note that both `nnsight`, but especially `NDIF`, is in active development and\n",
        "therefore there may be caveats, changes, and errors to work through.v\n",
        "\n",
        "If you're interested in following updates to `nnsight`, contributing, giving\n",
        "feedback, or finding collaborators, please join the\n",
        "[NDIF discord](https://discord.gg/6uFJmCSwW7)!\n",
        "\n",
        "Our website [NNsight.net](https://nnsight.net/), has a bunch more tutorials\n",
        "detailing more complex interpretability techniques using `nnsight`. If you want\n",
        "to share any of the work you do using `nnsight`, let others know on either of\n",
        "the discords above and we might turn it into a tutorial on our website.\n",
        "\n",
        "You can also follow us on X/Twitter [@ndif_team](https://x.com/ndif_team).\n",
        "\n",
        "💟\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.12.5"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "0a676b8f2dc34fed9731f19a03e4bfa1": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_74358c341daa44faa41a92b0eb073615",
              "IPY_MODEL_c40e17a994ed44d79ef54421f4772858",
              "IPY_MODEL_431ffcb0dbb1419dbbb237a29842b905"
            ],
            "layout": "IPY_MODEL_c28c185a508e459b95edb29afedd9c79"
          }
        },
        "0fddebfb42e847e08890d09f30161984": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "11e99924bdf84739b3cfc7d676b797e5": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "18d4c2302f424bc18f9215f2d5b83393": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1a1ada8b122d4502949c4c75861a82bf": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8de63414803f41c29e75242756cfd67f",
            "placeholder": "​",
            "style": "IPY_MODEL_f693dfed91df492c986afbd4ff925f89",
            "value": " 456k/456k [00:00&lt;00:00, 1.87MB/s]"
          }
        },
        "26d09aa072d64af0bf4df9a89309093a": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "3731a5747724403e86ef1abaf4fb389b": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "431ffcb0dbb1419dbbb237a29842b905": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_562a287082eb4882af864062ae13c91b",
            "placeholder": "​",
            "style": "IPY_MODEL_11e99924bdf84739b3cfc7d676b797e5",
            "value": " 665/665 [00:00&lt;00:00, 12.8kB/s]"
          }
        },
        "562a287082eb4882af864062ae13c91b": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "56490c20aef04fe4a63ba7c990e71e77": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "60fe026e908741209637604edfeb0563": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e7a73cfc3fa04425afc08368532ef0be",
              "IPY_MODEL_d839da1c31a34659988e3420ba520725",
              "IPY_MODEL_96a3b8c765374a5e98eaf713dd5e85d8"
            ],
            "layout": "IPY_MODEL_b49ce0d00e5b4df097b3a9fc4c0bd8e7"
          }
        },
        "688bf676958a4050b73cc9597f5110de": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "6aa96b80e3b746b48f5c13bcafe1134e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "74358c341daa44faa41a92b0eb073615": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7a0f3d2deeb14bbba7f0df208133a7e2",
            "placeholder": "​",
            "style": "IPY_MODEL_94cdcf1fb0914dfe99c3c615e68b7ec8",
            "value": "config.json: 100%"
          }
        },
        "7a0f3d2deeb14bbba7f0df208133a7e2": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7e6b41ec535648b190bfc655d9e430dd": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_beb4231b49024d8a9bb7ead19647d99d",
            "placeholder": "​",
            "style": "IPY_MODEL_26d09aa072d64af0bf4df9a89309093a",
            "value": " 1.36M/1.36M [00:00&lt;00:00, 4.16MB/s]"
          }
        },
        "8776d04bfd3748b69e10708466e9c132": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "8de63414803f41c29e75242756cfd67f": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9373fcd054014760b23e3acc928da0a5": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0fddebfb42e847e08890d09f30161984",
            "max": 1355256,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_d07a0a49a2924092affc0e07bf0b5ef0",
            "value": 1355256
          }
        },
        "94cdcf1fb0914dfe99c3c615e68b7ec8": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "95d303d83d3348a0adfdf6b5cb94c41c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_eb0c489b913744a28304e5334bb66532",
            "placeholder": "​",
            "style": "IPY_MODEL_c80603555c634e338f16c0021f752420",
            "value": "tokenizer.json: 100%"
          }
        },
        "96a3b8c765374a5e98eaf713dd5e85d8": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_8776d04bfd3748b69e10708466e9c132",
            "placeholder": "​",
            "style": "IPY_MODEL_56490c20aef04fe4a63ba7c990e71e77",
            "value": " 1.04M/1.04M [00:00&lt;00:00, 3.16MB/s]"
          }
        },
        "b2748df9aff540df954f002310606a56": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b49ce0d00e5b4df097b3a9fc4c0bd8e7": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b7be20bd738940658f710309729d8ccd": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "b8ad433cc3dd4a8689cb06655158d545": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e51620404d244a4e806eeaf3b96a474f",
              "IPY_MODEL_e4568f72971749ac92cc5577cb9c4f9c",
              "IPY_MODEL_1a1ada8b122d4502949c4c75861a82bf"
            ],
            "layout": "IPY_MODEL_18d4c2302f424bc18f9215f2d5b83393"
          }
        },
        "bc8190c700e949ebb6f874c8f8055161": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "beb4231b49024d8a9bb7ead19647d99d": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c28c185a508e459b95edb29afedd9c79": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c40e17a994ed44d79ef54421f4772858": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f3e24aa6619e4185a3e2f53f2e473e45",
            "max": 665,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_b7be20bd738940658f710309729d8ccd",
            "value": 665
          }
        },
        "c80603555c634e338f16c0021f752420": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ceebda1d07b545f5b605551c08e73be8": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "cef2516aabc24f6d852079984afa8090": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_95d303d83d3348a0adfdf6b5cb94c41c",
              "IPY_MODEL_9373fcd054014760b23e3acc928da0a5",
              "IPY_MODEL_7e6b41ec535648b190bfc655d9e430dd"
            ],
            "layout": "IPY_MODEL_d0198fa98e5f412380f2ba3e9448aef3"
          }
        },
        "d0198fa98e5f412380f2ba3e9448aef3": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d07a0a49a2924092affc0e07bf0b5ef0": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "d839da1c31a34659988e3420ba520725": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3731a5747724403e86ef1abaf4fb389b",
            "max": 1042301,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_688bf676958a4050b73cc9597f5110de",
            "value": 1042301
          }
        },
        "e4568f72971749ac92cc5577cb9c4f9c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ceebda1d07b545f5b605551c08e73be8",
            "max": 456318,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_6aa96b80e3b746b48f5c13bcafe1134e",
            "value": 456318
          }
        },
        "e51620404d244a4e806eeaf3b96a474f": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f28f557210b7402b8b2d0fd83f038dc1",
            "placeholder": "​",
            "style": "IPY_MODEL_bc8190c700e949ebb6f874c8f8055161",
            "value": "merges.txt: 100%"
          }
        },
        "e7a73cfc3fa04425afc08368532ef0be": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b2748df9aff540df954f002310606a56",
            "placeholder": "​",
            "style": "IPY_MODEL_f4caf2bc58704464b6ac7bf6da31b71d",
            "value": "vocab.json: 100%"
          }
        },
        "eb0c489b913744a28304e5334bb66532": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f28f557210b7402b8b2d0fd83f038dc1": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f3e24aa6619e4185a3e2f53f2e473e45": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f4caf2bc58704464b6ac7bf6da31b71d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "f693dfed91df492c986afbd4ff925f89": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}