{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eo91KAWJ-HSJ"
      },
      "source": [
        "# NNsight 0.3 - User Guide"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aqLF_dP_Fb_B"
      },
      "source": [
        "## Set up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "sTJ7xkITsjuE"
      },
      "outputs": [],
      "source": [
        "from IPython.display import clear_output\n",
        "\n",
        "!pip install nnsight\n",
        "!pip install --upgrade transformers torch\n",
        "\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "ZDGNHRmxdQDw"
      },
      "outputs": [],
      "source": [
        "from google.colab import userdata\n",
        "from nnsight import CONFIG\n",
        "\n",
        "from nnsight.logger import remote_logger\n",
        "remote_logger.propagate = False\n",
        "\n",
        "CONFIG.set_default_api_key('422220a9817141e49c5add1868af07a5')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "F4d92-2eVJ1r"
      },
      "outputs": [],
      "source": [
        "from collections import OrderedDict\n",
        "from nnsight import NNsight\n",
        "import torch\n",
        "\n",
        "input_size = 5\n",
        "hidden_dims = 10\n",
        "output_size = 2\n",
        "\n",
        "torch.manual_seed(423)\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, hidden_dims))\n",
        "        ]\n",
        "    )\n",
        ").requires_grad_(False)\n",
        "\n",
        "input = torch.rand((1, input_size))\n",
        "input_2 = torch.rand((1, input_size))\n",
        "\n",
        "tiny_model = NNsight(net)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "ivKBOJ3RdBde",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5bcd5b20-28ff-402b-fe85-c62cc7d55974"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-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 depracted 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",
        "lm = LanguageModel(\"openai-community/gpt2\", dispatch=True)\n",
        "llm = LanguageModel(\"meta-llama/Meta-Llama-3.1-8B\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Breaking Changes"
      ],
      "metadata": {
        "id": "tn8kHURI7yCf"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## input/inputs"
      ],
      "metadata": {
        "id": "edM1XEFTP1Bc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Module input access has a syntactic change:\n",
        "\n",
        "- Old: `nnsight.Envoy.input`\n",
        "\n",
        "- New: `nnsight.Envoy.inputs`\n",
        "\n",
        "- Note: `nnsight.Envoy.input` now provides access to the first positional argument of the module's input."
      ],
      "metadata": {
        "id": "dqUXNzX7DnJI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with lm.trace(\"Hello World\"):\n",
        "  l2_ins = lm.transformer.h[2].inputs.save()\n",
        "  l2_in = lm.transformer.h[2].input.save()\n",
        "\n",
        "print(\"Inputs: \", l2_ins)\n",
        "print(\"First Positional Argument Input: \", l2_in)"
      ],
      "metadata": {
        "id": "nm3RAmLEP0e0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5e115d69-3c58-4f44-acea-1087b384090c"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "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"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Inputs:  ((tensor([[[ 1.0967, -1.8792,  1.0121,  ..., -1.0442, -0.5490, -1.1239],\n",
            "         [-0.2300,  0.5204,  0.4756,  ..., -1.6795, -0.5285,  0.4190]]],\n",
            "       grad_fn=<AddBackward0>),), {'layer_past': None, 'attention_mask': None, 'head_mask': None, 'encoder_hidden_states': None, 'encoder_attention_mask': None, 'use_cache': True, 'output_attentions': False})\n",
            "First Positional Argument Input:  tensor([[[ 1.0967, -1.8792,  1.0121,  ..., -1.0442, -0.5490, -1.1239],\n",
            "         [-0.2300,  0.5204,  0.4756,  ..., -1.6795, -0.5285,  0.4190]]],\n",
            "       grad_fn=<AddBackward0>)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## `scan` and `validate`"
      ],
      "metadata": {
        "id": "rV9C-4jiP5Nn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "`scan` and `validate` are now set to `False` by default in the `Tracer` context."
      ],
      "metadata": {
        "id": "UAOuVY4AgMZ_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# New Features"
      ],
      "metadata": {
        "id": "cN8E6k790GR1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Scanning"
      ],
      "metadata": {
        "id": "sIMmbTk5fH9J"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can scan a model without executing it to gather important insights. This is useful for looking at internal modules' shapes for example. You can pass a dummy input to the model, and it will not be executed. This is also means that you don't have to call `save()` on any variable."
      ],
      "metadata": {
        "id": "4ABUUMY0fmiO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with tiny_model.scan(torch.tensor([0, 0, 0, 0, 0])):\n",
        "  dim = tiny_model.layer2.input.shape\n",
        "\n",
        "print(dim)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mpj0ieE5fLgb",
        "outputId": "413611b8-5805-4148-fd2e-80aa400b599e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([10])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## nnsight builtins"
      ],
      "metadata": {
        "id": "p14OBYvg7fbK"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can now define multiple `Python` builtins to be traceable by the Intervention graph.\n",
        "\n",
        "Simply use the `nnsight` import to call constructors for these data structures."
      ],
      "metadata": {
        "id": "xM-5nNieSy44"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import nnsight\n",
        "\n",
        "with tiny_model.trace(input):\n",
        "  num = nnsight.int(5).save()\n",
        "  l = nnsight.list().save()\n",
        "  l.append(num)\n",
        "  d = nnsight.dict({\"five\": num}).save()\n",
        "\n",
        "print(\"Interger: \", num)\n",
        "print(\"List: \", l)\n",
        "print(\"Dictionary: \", d)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8ZwiHFmSSyfx",
        "outputId": "44091911-ef6b-48a1-f7c1-3dd031b50ae2"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Interger:  5\n",
            "List:  [5]\n",
            "Dictionary:  {'five': 5}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here is the complete list of supported `Python` builtins:\n",
        "- bool\n",
        "- bytes\n",
        "- int\n",
        "- float\n",
        "- str\n",
        "- complex\n",
        "- bytearray\n",
        "- tuple\n",
        "- list\n",
        "- set\n",
        "- dict"
      ],
      "metadata": {
        "id": "2G-Mw0r1T87b"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VV5ROMYuGTCF"
      },
      "source": [
        "## Proxy Update"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xU02sRDi7hdU"
      },
      "source": [
        "For literals created and traced by `nnsight`, there is no direct way of setting their values.\n",
        "\n",
        "Use our `.update()` method on Intervention Proxies to assign it a new value."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import nnsight\n",
        "\n",
        "with tiny_model.trace(input):\n",
        "  input_str = nnsight.str(\"I am a \").save()\n",
        "  input_str.update(input_str + \"Transformer\")\n",
        "\n",
        "print(\"Input: \", input_str)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2S5qlgZUUzoH",
        "outputId": "95ae0e94-4811-45d0-db30-57798ef662cd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input:  I am a Transformer\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This is also useful for calculating running sums and other statistics."
      ],
      "metadata": {
        "id": "QypCaxALVnzc"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BmtjvU4vhzMK"
      },
      "source": [
        "## Logging"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wOBzLbqZiT3y"
      },
      "source": [
        "We are probably all guilty, at least once, of trying to print an Intervention Proxy from within the tracing context to look at its value:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cIsQuKNNi2Ct",
        "outputId": "39e73bbb-09e6-4a59-c3a8-1a5a6ae5a111"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "InterventionProxy (InterventionProtocol_0): \n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "    print(tiny_model.layer1.output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tIP4OIrUkQ9y"
      },
      "source": [
        "The reason this does not print any actual value is because the model is only executed upon exiting the `Tracer` context, and thus, the proxies' values have not been populated yet."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AdUcR0TqkoL7"
      },
      "source": [
        "If you are still only interested in looking at some intermediate values without necessary saving them, you can call our logging feature which will be executed as an `nnsight` node during the model's execution and show you the actual values."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8IMHHTkwhyTm",
        "outputId": "fa179999-549a-420c-857f-f8287e2260dd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Layer 1 - out:  tensor([[ 7.2426e-01,  2.4406e-01, -5.3356e-01,  2.4396e-01, -4.4996e-04,\n",
            "         -1.3551e-01,  1.8728e-01,  7.5127e-01, -3.3102e-01, -8.0205e-01]])\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "\n",
        "with tiny_model.trace(input) as tracer:\n",
        "  nnsight.log(\"Layer 1 - out: \", tiny_model.layer1.output)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Tracing function calls"
      ],
      "metadata": {
        "id": "hVScJAjpSbZd"
      }
    },
    {
      "cell_type": "markdown",
      "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.\n",
        "\n",
        "For custom functions we can use `nnsight.apply()` to add them to the intervention graph."
      ],
      "metadata": {
        "id": "GUC1CpwuG-vw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import nnsight\n",
        "import torch\n",
        "\n",
        "# We define a simple custom function that sums all the elements of a tensor\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 lm.trace(\"The Eiffel Tower is in the city of\") 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, lm.transformer.h[0].output[0]).save()\n",
        "    sum = lm.transformer.h[0].output[0].sum().save()\n",
        "\n",
        "print(\"PyTorch sum: \", sum)\n",
        "print(\"Our sum: \", custom_sum)"
      ],
      "metadata": {
        "id": "yYozeK24So2l",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "cd5e7821-1294-46c0-dc84-6bc802d0aafe"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "PyTorch sum:  tensor(191.2440, grad_fn=<SumBackward0>)\n",
            "Our sum:  tensor(191.2442)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Early Stopping"
      ],
      "metadata": {
        "id": "UuKa3RZ1RLKW"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0RWqu5cA-clS"
      },
      "source": [
        "If you are only interested in a model's intermediate computations, you can halt a forward pass run at any module level, reducing runtime and conserving computational resources. This is particularly useful if you are working with SAEs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "85dDFQQSMgRf",
        "outputId": "efac2e7f-ffb1-4bf5-d439-2398373ed5eb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "L1 - Output:  tensor([[ 7.2426e-01,  2.4406e-01, -5.3356e-01,  2.4396e-01, -4.4996e-04,\n",
            "         -1.3551e-01,  1.8728e-01,  7.5127e-01, -3.3102e-01, -8.0205e-01]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "   l1_out = tiny_model.layer1.output.save()\n",
        "   tiny_model.layer1.output.stop()\n",
        "\n",
        "print(\"L1 - Output: \", l1_out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZY92I8U-fJ2"
      },
      "source": [
        "Interventions within the `Tracer` 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": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 339
        },
        "id": "57_IfUB6Zs9e",
        "outputId": "0525db2b-3738-4a46-b897-3b6abf11ad11"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "L2 - Output:  "
          ]
        },
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "Accessing value before it's been set.",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-12-bd7c79651e17>\u001b[0m in \u001b[0;36m<cell line: 5>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      3\u001b[0m    \u001b[0mtiny_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayer1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"L2 - Output: \"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0ml2_out\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Proxy.py\u001b[0m in \u001b[0;36m__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     54\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mattached\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 56\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     57\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     58\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0;34mf\"{type(self).__name__} ({self.node.name}): {self.node.proxy_value if self.node.proxy_value is not inspect._empty else ''}\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Proxy.py\u001b[0m in \u001b[0;36mvalue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     48\u001b[0m         \"\"\"\n\u001b[1;32m     49\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 50\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     51\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__str__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Node.py\u001b[0m in \u001b[0;36mvalue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    180\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    181\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 182\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Accessing value before it's been set.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    183\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    184\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_value\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\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": {
        "id": "pG0DiXO885Ts"
      },
      "source": [
        "## Conditional Interventions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lyDbzWwe8OPy"
      },
      "source": [
        "You can make interventions conditional!\n",
        "\n",
        "Create a Conditional context and pass it a value to be evaluated as a boolean. The context will wrap all the interventions that you wish to be dependent on the condition specified.\n",
        "\n",
        "Let's take a look at how you can do that:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "56L4_IdiiqCG",
        "outputId": "116be79d-2280-4b98-aa30-92341a06faa4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Random Integer  -1  is Odd\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "  rand_int = torch.randint(low=-10, high=10, size=(1,)).item()\n",
        "\n",
        "  with tracer.cond(rand_int % 2 == 0):\n",
        "    tracer.apply(print, \"Random Integer \", rand_int, \" is Even\")\n",
        "\n",
        "  with tracer.cond(rand_int % 2 == 1):\n",
        "    tracer.apply(print, \"Random Integer \", rand_int, \" is Odd\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sIfJKz2NkTcl"
      },
      "source": [
        "In the example above, we have two Conditional contexts with mutually exclusive conditions, mimicking a conventional `If`-`Else` statement."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LYhVDGjPtZgm"
      },
      "source": [
        "The condition passed to the Conditional context is evaluated directly by calling `bool()` on the proxy value, so be mindful of how your Intervention Proxy condition evaluates to boolean."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 495
        },
        "id": "h219j-m2ufFw",
        "outputId": "094f6c65-3780-4579-e873-ad2a0106f05a",
        "collapsed": true
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "Above exception when execution Node: 'ne_0' in Graph: '136573620465632'",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Node.py\u001b[0m in \u001b[0;36mexecute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    371\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 372\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtarget\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    373\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/protocols.py\u001b[0m in \u001b[0;36mexecute\u001b[0;34m(cls, node)\u001b[0m\n\u001b[1;32m    759\u001b[0m         \u001b[0mcond_value\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mNode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprepare_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 760\u001b[0;31m         \u001b[0;32mif\u001b[0m \u001b[0mcond_value\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    761\u001b[0m             \u001b[0;31m# cond_value is True\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: Boolean value of Tensor with more than one value is ambiguous",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Node.py\u001b[0m in \u001b[0;36mexecute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    382\u001b[0m                 \u001b[0;31m# Set value.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 383\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_value\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    384\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Node.py\u001b[0m in \u001b[0;36mset_value\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m    409\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mlistener\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfulfilled\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msequential\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 410\u001b[0;31m                 \u001b[0mlistener\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    411\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Node.py\u001b[0m in \u001b[0;36mexecute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    386\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 387\u001b[0;31m             raise type(e)(\n\u001b[0m\u001b[1;32m    388\u001b[0m                 \u001b[0;34mf\"Above exception when execution Node: '{self.name}' in Graph: '{self.graph.id}'\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: Above exception when execution Node: 'ConditionalProtocol_0' in Graph: '136573620465632'",
            "\nThe above exception was the direct cause of the following exception:\n",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-19-9ec87ef9235a>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32mwith\u001b[0m \u001b[0mtiny_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m   \u001b[0ml1_out\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtiny_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayer1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m   \u001b[0;32mwith\u001b[0m \u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcond\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ml1_out\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0mtracer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mapply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"Condition is True\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/contexts/Tracer.py\u001b[0m in \u001b[0;36m__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    102\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 103\u001b[0;31m         \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__exit__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexc_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexc_val\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexc_tb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    105\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0minvoke\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mInvoker\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/contexts/GraphBasedContext.py\u001b[0m in \u001b[0;36m__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m    216\u001b[0m             \u001b[0;32mraise\u001b[0m \u001b[0mexc_val\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 218\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    219\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    220\u001b[0m     \u001b[0;31m### BACKENDS ########\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/contexts/backends/LocalBackend.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, obj)\u001b[0m\n\u001b[1;32m     25\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mobj\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mLocalMixin\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m         \u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlocal_backend_execute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/contexts/Tracer.py\u001b[0m in \u001b[0;36mlocal_backend_execute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    145\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    146\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m         self.model.interleave(\n\u001b[0m\u001b[1;32m    148\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_execute\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/models/NNsightModel.py\u001b[0m in \u001b[0;36minterleave\u001b[0;34m(self, fn, intervention_graph, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    460\u001b[0m         ).keys()\n\u001b[1;32m    461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 462\u001b[0;31m         with HookHandler(\n\u001b[0m\u001b[1;32m    463\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_model\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    464\u001b[0m             \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule_paths\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/intervention.py\u001b[0m in \u001b[0;36m__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m    572\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    573\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexc_val\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mException\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 574\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0mexc_val\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    575\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    576\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/models/NNsightModel.py\u001b[0m in \u001b[0;36minterleave\u001b[0;34m(self, fn, intervention_graph, *inputs, **kwargs)\u001b[0m\n\u001b[1;32m    471\u001b[0m         ):\n\u001b[1;32m    472\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 473\u001b[0;31m                 \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    474\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mprotocols\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mEarlyStopProtocol\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mEarlyStopException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    475\u001b[0m                 \u001b[0;31m# TODO: Log.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/models/NNsightModel.py\u001b[0m in \u001b[0;36m_execute\u001b[0;34m(self, *prepared_inputs, **kwargs)\u001b[0m\n\u001b[1;32m    582\u001b[0m             \u001b[0;32mpass\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    583\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 584\u001b[0;31m         return self._model(\n\u001b[0m\u001b[1;32m    585\u001b[0m             \u001b[0;34m*\u001b[0m\u001b[0mprepared_inputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    586\u001b[0m             \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1552\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1553\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1554\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1555\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1601\u001b[0m                 \u001b[0margs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbw_hook\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msetup_input_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1602\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1603\u001b[0;31m             \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1604\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0m_global_forward_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1605\u001b[0m                 for hook_id, hook in (\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    217\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    218\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 219\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    220\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    221\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1551\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1552\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1553\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1554\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1555\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1614\u001b[0m                         \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1615\u001b[0m                     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1616\u001b[0;31m                         \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1617\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1618\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0mhook_result\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/intervention.py\u001b[0m in \u001b[0;36moutput_hook\u001b[0;34m(module, input, output, module_path)\u001b[0m\n\u001b[1;32m    557\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    558\u001b[0m                 \u001b[0;32mdef\u001b[0m \u001b[0moutput_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule_path\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodule_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 559\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    560\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    561\u001b[0m                 self.handles.append(\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/models/NNsightModel.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(activations, module_path)\u001b[0m\n\u001b[1;32m    466\u001b[0m                 \u001b[0mactivations\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"input\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mintervention_handler\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    467\u001b[0m             ),\n\u001b[0;32m--> 468\u001b[0;31m             output_hook=lambda activations, module_path: InterventionProtocol.intervene(\n\u001b[0m\u001b[1;32m    469\u001b[0m                 \u001b[0mactivations\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"output\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mintervention_handler\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    470\u001b[0m             ),\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/intervention.py\u001b[0m in \u001b[0;36mintervene\u001b[0;34m(cls, activations, module_path, key, intervention_handler)\u001b[0m\n\u001b[1;32m    447\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    448\u001b[0m                 \u001b[0;31m# Value injection.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 449\u001b[0;31m                 \u001b[0mnode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_value\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    450\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    451\u001b[0m                 \u001b[0;31m# Check if through the previous value injection, there was a 'swap' intervention.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Node.py\u001b[0m in \u001b[0;36mset_value\u001b[0;34m(self, value)\u001b[0m\n\u001b[1;32m    408\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    409\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mlistener\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfulfilled\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msequential\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 410\u001b[0;31m                 \u001b[0mlistener\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecute\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    411\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    412\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mdependency\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marg_dependencies\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/nnsight/tracing/Node.py\u001b[0m in \u001b[0;36mexecute\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    385\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    386\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 387\u001b[0;31m             raise type(e)(\n\u001b[0m\u001b[1;32m    388\u001b[0m                 \u001b[0;34mf\"Above exception when execution Node: '{self.name}' in Graph: '{self.graph.id}'\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    389\u001b[0m             ) from e\n",
            "\u001b[0;31mRuntimeError\u001b[0m: Above exception when execution Node: 'ne_0' in Graph: '136573620465632'"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "  l1_out = tiny_model.layer1.output\n",
        "  with tracer.cond(l1_out != 1):\n",
        "    tracer.apply(print, \"Condition is True\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OkhsHKecu_U-"
      },
      "source": [
        "The code above throws the **ERROR**: `Boolean value of Tensor with more than one value is ambiguous`, because the condition specified in the Conditional context cannot be handled properly.\n",
        "\n",
        "Instead, use something like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G6G7NovlvcNk",
        "outputId": "23e2a5e7-18ec-498b-baf8-e238f86e50a3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Condition is True\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "  l1_out = tiny_model.layer1.output\n",
        "  with tracer.cond(torch.all(l1_out != 1)):\n",
        "    tracer.apply(print, \"Condition is True\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6-jntBjMtNcC"
      },
      "source": [
        "Conditional contexts can also be nested, if you want your interventions to depend on more than one condition at a time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LvsSFzOu9-o7",
        "outputId": "a1f00a8e-6c7f-45de-d1c3-6b441b84997f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Rand Int  6  is Positive and Even\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "  rand_int = tracer.apply(int, 6)\n",
        "  with tracer.cond(rand_int > 0):\n",
        "    with tracer.cond(rand_int % 2 == 0):\n",
        "      tracer.apply(print, \"Rand Int \", rand_int, \" is Positive and Even\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bh7sOlVKNOjJ"
      },
      "source": [
        "## Model Editing"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can alter a model by setting default edits and interventions in an editing context, applied before each forward pass. This can be used to attach additional modules like SAEs.\n",
        "\n"
      ],
      "metadata": {
        "id": "N9AiWwkjsL-k"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with tiny_model.edit() as edited_model:\n",
        "  tiny_model.layer1.output[0][:] = 0\n",
        "\n",
        "with tiny_model.trace(input):\n",
        "  l1_out = tiny_model.layer1.output.save()\n",
        "\n",
        "with edited_model.trace(input):\n",
        "  l1_out_edited = edited_model.layer1.output.save()\n",
        "\n",
        "print(\"L1 - Out: \", l1_out)\n",
        "print(\"L1 - Out [edited]: \", l1_out_edited)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x9IWHelvYzS3",
        "outputId": "37690666-f80d-4256-fdb7-19282a5a9dad"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "L1 - Out:  tensor([[ 7.2426e-01,  2.4406e-01, -5.3356e-01,  2.4396e-01, -4.4996e-04,\n",
            "         -1.3551e-01,  1.8728e-01,  7.5127e-01, -3.3102e-01, -8.0205e-01]])\n",
            "L1 - Out [edited]:  tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's look at anotehr example"
      ],
      "metadata": {
        "id": "SmDdUhacZW6A"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "K7NkeurTCFIk",
        "outputId": "7cd2e840-c22e-43ef-a44c-e404586b3426"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "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"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Original output:  tensor([[[-23.8122, -24.2187, -27.4893,  ..., -30.6782, -29.9407, -24.9513],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6782, -29.9407, -24.9513],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6782, -29.9407, -24.9513],\n",
            "         ...,\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6781, -29.9407, -24.9512],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6781, -29.9407, -24.9512],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6781, -29.9407, -24.9512]]],\n",
            "       grad_fn=<UnsafeViewBackward0>)\n",
            "Edited output:  tensor([[[-23.8122, -24.2187, -27.4893,  ..., -30.6782, -29.9407, -24.9513],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6782, -29.9407, -24.9513],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6782, -29.9407, -24.9513],\n",
            "         ...,\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6781, -29.9407, -24.9512],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6781, -29.9407, -24.9512],\n",
            "         [-23.8122, -24.2187, -27.4893,  ..., -30.6781, -29.9407, -24.9512]]],\n",
            "       grad_fn=<UnsafeViewBackward0>)\n"
          ]
        }
      ],
      "source": [
        "from nnsight.util import WrapperModule\n",
        "\n",
        "class ComplexModule(torch.nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.one = WrapperModule()\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.one(x)\n",
        "\n",
        "l0 = lm.transformer.h[0]\n",
        "l0.attachment = ComplexModule()\n",
        "\n",
        "with lm.edit() as gpt2_edited:\n",
        "    acts = l0.output[0]\n",
        "    l0.output[0][:] = l0.attachment(acts, hook=True)\n",
        "\n",
        "# Get values pre editing\n",
        "with lm.trace(\"Madison Square Garden is located in the city of\"):\n",
        "    original = l0.output[0].clone().save()\n",
        "    l0.output[0][:] *= 0.\n",
        "    original_output = lm.output.logits.save()\n",
        "\n",
        "with gpt2_edited.trace(\"Madison Square Garden is located in the city of\"):\n",
        "    one = l0.attachment.one.output.clone().save()\n",
        "    l0.attachment.output *= 0.\n",
        "    edited_output = lm.output.logits.save()\n",
        "\n",
        "print(\"Original output: \", original_output)\n",
        "print(\"Edited output: \", edited_output)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Your edit call can be customized by choosing to perform edits in-place on the model andgetting access to the editor context (`nnsight.context.Tracer`).\n",
        "\n",
        "You can also choose to remove edits perfomerd on a model at a later stage."
      ],
      "metadata": {
        "id": "erNKIithZwX2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with tiny_model.edit(inplace=True, return_context=True) as editor:\n",
        "  tiny_model.layer1.output[0][:] = 0\n",
        "\n",
        "with tiny_model.trace(input):\n",
        "  l1_out = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"L1 - Out: \", l1_out)\n",
        "\n",
        "tiny_model.clear_edits()\n",
        "\n",
        "with tiny_model.trace(input):\n",
        "  l1_out = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"L1 - Out [unedited]: \", l1_out)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Iqymcqm3anWn",
        "outputId": "d1782c65-a2e1-4783-b781-ab9920410ac2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "L1 - Out:  tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])\n",
            "L1 - Out [unedited]:  tensor([[ 7.2426e-01,  2.4406e-01, -5.3356e-01,  2.4396e-01, -4.4996e-04,\n",
            "         -1.3551e-01,  1.8728e-01,  7.5127e-01, -3.3102e-01, -8.0205e-01]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VXOje4UWjxyw"
      },
      "source": [
        "Note that setting new modules with remote execution is currently not supported!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lefW2GneGGRg"
      },
      "source": [
        "## Session Context"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "51qInRPnHLnQ"
      },
      "source": [
        "`nnsight 0.3` focuses on enhancing the capabilities of our remote execution API, powered by the [NDIF](https://ndif.us) backend.\n",
        "\n",
        "To achieve this, we introduce the **Session** context: an overarching structure for efficiently handling multi-tracing experiments. This means that, multiple `Tracer` contexts can be packaged together as part of one single request to the server.\n",
        "\n",
        "The `Session` context can also be used entirely for local usage, as it enables useful functionalities and optimizes experiments."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "81FkT_fdhYqM",
        "outputId": "68cfabcb-b54b-4cfd-bc0c-21afb198ccd2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-08-30 19:58:12,517 MainProcess nnsight_remote INFO     bf66301e-f693-4146-a721-b4d93c6d3318 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 19:58:12,544 MainProcess nnsight_remote INFO     bf66301e-f693-4146-a721-b4d93c6d3318 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 19:58:12,569 MainProcess nnsight_remote INFO     bf66301e-f693-4146-a721-b4d93c6d3318 - RUNNING: Your job has started running.\n",
            "2024-08-30 19:58:12,684 MainProcess nnsight_remote INFO     bf66301e-f693-4146-a721-b4d93c6d3318 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 812kB/s]\n"
          ]
        }
      ],
      "source": [
        "with llm.session(remote=True):\n",
        "  with llm.trace(\"_\") as t1:\n",
        "    # define interventions here\n",
        "    pass\n",
        "\n",
        "  with llm.trace(\"_\") as t2:\n",
        "    # define interventions here\n",
        "    pass\n",
        "\n",
        "  with llm.trace(\"_\") as t3:\n",
        "    # define interventions here\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FYprIurafeOJ"
      },
      "source": [
        "All operations defined within a `Session` context are executed at the very end (upon exiting the overarching context) and it is conducted sequentially, strictly following the order of definition (`t2` being executed after `t1` and `t3` after `t2`)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gc_GUYeIj8sm"
      },
      "source": [
        "In a `Session`, interventions defined at any early stage can be seamlessly referenced."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DIOgEKPHcL3S",
        "outputId": "30e8ff7d-d3a5-4acf-a673-fc011dcd8087"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-08-30 19:56:44,378 MainProcess nnsight_remote INFO     e353b8a7-14a1-4ebe-a477-3ec6ae719726 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 19:56:44,413 MainProcess nnsight_remote INFO     e353b8a7-14a1-4ebe-a477-3ec6ae719726 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 19:56:44,496 MainProcess nnsight_remote INFO     e353b8a7-14a1-4ebe-a477-3ec6ae719726 - RUNNING: Your job has started running.\n",
            "2024-08-30 19:56:45,043 MainProcess nnsight_remote INFO     e353b8a7-14a1-4ebe-a477-3ec6ae719726 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 13.6M/13.6M [00:02<00:00, 6.66MB/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "T1 - Prediction:   Paris\n",
            "T2 - Prediction:   \n"
          ]
        }
      ],
      "source": [
        "with llm.session(remote=True) as session:\n",
        "  with llm.trace(\"The Eiffel Tower is in the city of\") as t1:\n",
        "    hs_11 = llm.model.layers[-1].output[0][:, -1, :] # no .save()\n",
        "    t1_tokens_out = llm.output.save()\n",
        "\n",
        "  with llm.trace(\"Buckingham Palace is in the city of\") as t2:\n",
        "    llm.model.layers[-2].output[0][:, -1, :] = hs_11[:]\n",
        "    t2_tokens_out = llm.output.save()\n",
        "\n",
        "print(\"\\nT1 - Prediction: \", llm.tokenizer.decode(t1_tokens_out[\"logits\"].argmax(dim=-1)[0][-1]))\n",
        "print(\"T2 - Prediction: \", llm.tokenizer.decode(t2_tokens_out[\"logits\"].argmax(dim=-1)[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Coq-bmrjRtx"
      },
      "source": [
        "In the example above, we are interested in patching the hidden state of a later layer into an earlier one. This experiment can only be conducted with two `Tracer` contexts; since we are using a `Session`, it is not required to save the hidden state from Tracer 1 to reference it in Tracer 2."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gIEaHC1HwJ71"
      },
      "source": [
        "The `Session` context can also be terminated early."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HhXN05Y0vSa2",
        "outputId": "59790acd-fb35-43ac-d9fd-bd02ebc446a4"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-08-30 19:58:17,651 MainProcess nnsight_remote INFO     546f8b1b-0f80-47dd-9160-2f6c3d23f154 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 19:58:17,679 MainProcess nnsight_remote INFO     546f8b1b-0f80-47dd-9160-2f6c3d23f154 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 19:58:17,694 MainProcess nnsight_remote INFO     546f8b1b-0f80-47dd-9160-2f6c3d23f154 - RUNNING: Your job has started running.\n",
            "2024-08-30 19:58:17,713 MainProcess nnsight_remote INFO     546f8b1b-0f80-47dd-9160-2f6c3d23f154 - LOG: -- Early Stop --\n",
            "2024-08-30 19:58:17,750 MainProcess nnsight_remote INFO     546f8b1b-0f80-47dd-9160-2f6c3d23f154 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 2.47MB/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "List:  [0, 1]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "\n",
        "with llm.session(remote=True) as session:\n",
        "  l = nnsight.list().save()\n",
        "\n",
        "  l.append(0)\n",
        "  l.append(1)\n",
        "  nnsight.log(\"-- Early Stop --\")\n",
        "  session.exit()\n",
        "  l.append(2)\n",
        "\n",
        "print(\"List: \", l)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ehauEV6jG9l1"
      },
      "source": [
        "## Iterator Context"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w0Bxt8S_u7ZN"
      },
      "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?\n",
        "\n",
        "If you create a `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 by introducing the **Iterator** context: an intervention loop that iteratively executes a single intervention graph with an updated parameter.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0g5rkm7TsT2w",
        "outputId": "5b6eff08-82c2-43a4-e8e6-ba86507e1bcc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-08-30 20:17:59,080 MainProcess nnsight_remote INFO     c086f993-3e4a-4d50-8375-319a1912ad4f - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 20:17:59,109 MainProcess nnsight_remote INFO     c086f993-3e4a-4d50-8375-319a1912ad4f - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 20:17:59,135 MainProcess nnsight_remote INFO     c086f993-3e4a-4d50-8375-319a1912ad4f - RUNNING: Your job has started running.\n",
            "2024-08-30 20:17:59,512 MainProcess nnsight_remote INFO     c086f993-3e4a-4d50-8375-319a1912ad4f - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 5.65M/5.65M [00:01<00:00, 3.42MB/s]\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "\n",
        "with llm.session(remote=True) as session:\n",
        "\n",
        "  prompts = nnsight.list([\"This is nnsight 0.3\",\n",
        "                          \"It works with NDIF\",\n",
        "                          \"pip install it now!\"])\n",
        "  results = nnsight.list().save()\n",
        "  with session.iter(prompts) as prompt:\n",
        "\n",
        "    with llm.trace(prompt):\n",
        "      results.append(llm.lm_head.output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i8BKgcwasXxr"
      },
      "source": [
        "Use a `Session` to define the `Iterator` context and pass in a sequence of items that you want to loop over at each executed iteration.\n",
        "\n",
        "The sequence must be iterable or be a Proxy with an iterable value.\n",
        "\n",
        "The iterable's item can be referenced in the inner intervention body of the `Iterator`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ghUE_Cr63ZBy"
      },
      "source": [
        "### loop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sZsc3tjAKL4_"
      },
      "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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oz1cSM5nnAO0",
        "outputId": "79dc6289-22ae-4197-bcc8-7bfd9c3ab9bd"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-08-30 20:49:46,956 MainProcess nnsight_remote INFO     5c4f92fb-8961-4c41-becf-7d56c1bb58ba - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 20:49:46,993 MainProcess nnsight_remote INFO     5c4f92fb-8961-4c41-becf-7d56c1bb58ba - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 20:49:47,021 MainProcess nnsight_remote INFO     5c4f92fb-8961-4c41-becf-7d56c1bb58ba - RUNNING: Your job has started running.\n",
            "2024-08-30 20:49:47,042 MainProcess nnsight_remote INFO     5c4f92fb-8961-4c41-becf-7d56c1bb58ba - LOG: 0\n",
            "2024-08-30 20:49:47,060 MainProcess nnsight_remote INFO     5c4f92fb-8961-4c41-becf-7d56c1bb58ba - LOG: 1\n",
            "2024-08-30 20:49:47,081 MainProcess nnsight_remote INFO     5c4f92fb-8961-4c41-becf-7d56c1bb58ba - LOG: 2\n",
            "2024-08-30 20:49:47,126 MainProcess nnsight_remote INFO     5c4f92fb-8961-4c41-becf-7d56c1bb58ba - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 992/992 [00:00<00:00, 6.86kB/s]\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "\n",
        "with llm.session(remote=True) as session:\n",
        "  l = nnsight.list()\n",
        "  [l.append(num) for num in range(0, 3)] # adding 0, 1, 2 to l\n",
        "  with session.iter(l) as item: # with session.iter([0, 1, 2]) also works!\n",
        "    nnsight.log(item)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GcJsufQaosfU"
      },
      "source": [
        "You can create nested `Iterator` contexts:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TIS5ckryu_yM",
        "outputId": "48e21c35-dd88-4ce6-84a4-a33e9139a2bf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "List:  [10, 10, 10, 10, 10]\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "\n",
        "with llm.session() as session:\n",
        "  l = nnsight.list([[10]] * 5)\n",
        "\n",
        "  l2 = nnsight.list().save()\n",
        "  with session.iter(l) as item:\n",
        "    with session.iter(item) as item_2:\n",
        "      l2.append(item_2)\n",
        "\n",
        "print(\"List: \", l2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m-6BvvmSsDjS"
      },
      "source": [
        "You can skip some iterations:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "120FlvGesNeu",
        "outputId": "2e8ba741-bb62-4028-c4f5-92de45b54e29"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-08-30 21:03:19,288 MainProcess nnsight_remote INFO     9587a57a-6375-405a-9439-fda3e86b5f4f - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 21:03:19,315 MainProcess nnsight_remote INFO     9587a57a-6375-405a-9439-fda3e86b5f4f - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 21:03:19,341 MainProcess nnsight_remote INFO     9587a57a-6375-405a-9439-fda3e86b5f4f - RUNNING: Your job has started running.\n",
            "2024-08-30 21:03:19,360 MainProcess nnsight_remote INFO     9587a57a-6375-405a-9439-fda3e86b5f4f - LOG: 0\n",
            "2024-08-30 21:03:19,379 MainProcess nnsight_remote INFO     9587a57a-6375-405a-9439-fda3e86b5f4f - LOG: 2\n",
            "2024-08-30 21:03:19,425 MainProcess nnsight_remote INFO     9587a57a-6375-405a-9439-fda3e86b5f4f - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 992/992 [00:00<00:00, 543kB/s]\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "\n",
        "with llm.session(remote=True) as session:\n",
        "\n",
        "  with session.iter([0, 1, 2, 3], return_context=True) as (item, iterator):\n",
        "    with iterator.cond(item % 2 == 0):\n",
        "      nnsight.log(item)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U3xOSDRWpG6Q"
      },
      "source": [
        "Or, you can choose to `break` out of the Iteration loop early:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FXnNW3d5pPRW",
        "outputId": "3379cf40-be81-4811-b624-909efd889464"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "2024-08-30 21:03:37,096 MainProcess nnsight_remote INFO     4269ef5e-0e34-4070-8f9d-cec8460f5069 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2024-08-30 21:03:37,126 MainProcess nnsight_remote INFO     4269ef5e-0e34-4070-8f9d-cec8460f5069 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2024-08-30 21:03:37,152 MainProcess nnsight_remote INFO     4269ef5e-0e34-4070-8f9d-cec8460f5069 - RUNNING: Your job has started running.\n",
            "2024-08-30 21:03:37,175 MainProcess nnsight_remote INFO     4269ef5e-0e34-4070-8f9d-cec8460f5069 - LOG: 0\n",
            "2024-08-30 21:03:37,194 MainProcess nnsight_remote INFO     4269ef5e-0e34-4070-8f9d-cec8460f5069 - LOG: 1\n",
            "2024-08-30 21:03:37,234 MainProcess nnsight_remote INFO     4269ef5e-0e34-4070-8f9d-cec8460f5069 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 992/992 [00:00<00:00, 2.23MB/s]\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "\n",
        "with llm.session(remote=True) as session:\n",
        "\n",
        "  with session.iter([0, 1, 2, 3], return_context=True) as (item, iterator):\n",
        "      with iterator.cond(item == 2):\n",
        "        iterator.exit()\n",
        "\n",
        "      nnsight.log(item)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qHGQrOrP4gt3"
      },
      "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!"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true,
      "collapsed_sections": [
        "aqLF_dP_Fb_B",
        "p14OBYvg7fbK",
        "VV5ROMYuGTCF",
        "BmtjvU4vhzMK",
        "hVScJAjpSbZd",
        "UuKa3RZ1RLKW"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}