From caa402d0b9d82e6674ed48e3c99bc1ff374dfdca Mon Sep 17 00:00:00 2001
From: ebortz <ebortz@bu.edu>
Date: Thu, 21 Nov 2024 10:27:31 -0500
Subject: [PATCH 1/8] Fixed typos in walkthrough

---
 .../tutorials/activation_patching.ipynb       | 184 +++++++++---------
 .../tutorials/attribution_patching.ipynb      |  32 +--
 .../notebooks/tutorials/logit_lens.ipynb      |  16 +-
 .../notebooks/tutorials/walkthrough.ipynb     |  95 +++++----
 docs/source/start.rst                         |   9 +-
 5 files changed, 175 insertions(+), 161 deletions(-)

diff --git a/docs/source/notebooks/tutorials/activation_patching.ipynb b/docs/source/notebooks/tutorials/activation_patching.ipynb
index 796aab28..7b12ad0c 100644
--- a/docs/source/notebooks/tutorials/activation_patching.ipynb
+++ b/docs/source/notebooks/tutorials/activation_patching.ipynb
@@ -109,7 +109,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": 2,
       "metadata": {},
       "outputs": [
         {
@@ -137,7 +137,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": 3,
       "metadata": {
         "id": "l8WfCzQhwFM0"
       },
@@ -152,7 +152,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 4,
       "metadata": {
         "id": "qLmX2tdZmgiz"
       },
@@ -165,7 +165,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 5,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -226,7 +226,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 6,
       "metadata": {
         "id": "7FrduLBCmhOp"
       },
@@ -249,7 +249,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 7,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -314,7 +314,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 8,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -381,7 +381,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 10,
+      "execution_count": 9,
       "metadata": {
         "id": "JodVtd5VAuo2"
       },
@@ -418,7 +418,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 11,
+      "execution_count": 10,
       "metadata": {
         "id": "1vuK3rEMAug_"
       },
@@ -477,7 +477,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 13,
       "metadata": {
         "id": "G60gi4ZUmxlG"
       },
@@ -566,7 +566,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 13,
+      "execution_count": 11,
       "metadata": {
         "id": "EaHSh6ctnJFl"
       },
@@ -601,7 +601,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 14,
+      "execution_count": 12,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -715,13 +715,13 @@
                       0,
                       0,
                       0,
-                      0.997469961643219,
-                      0.0017654040129855275,
-                      -0.0001825046056183055,
-                      0.000893437594641,
-                      0.00002027829032158479,
-                      -0.0010914491722360253,
-                      -0.00047117203939706087
+                      0.997462809085846,
+                      0.0017713640118017793,
+                      -0.00017296147416345775,
+                      0.0008970139897428453,
+                      0.000027435267838882282,
+                      -0.001084289513528347,
+                      -0.00046759238466620445
                     ],
                     [
                       0,
@@ -733,13 +733,13 @@
                       0,
                       0,
                       0,
-                      0.9971312284469604,
-                      0.0026075495406985283,
-                      -0.0004866789386142045,
-                      0.0006250484730117023,
-                      -0.00006799191032769158,
-                      -0.00043896533316001296,
-                      -0.0009244514512829483
+                      0.9971288442611694,
+                      0.002613507444038987,
+                      -0.0004866778035648167,
+                      0.0006274326588027179,
+                      -0.00006202756048878655,
+                      -0.0004341929452493787,
+                      -0.0009220635402016342
                     ],
                     [
                       0,
@@ -751,13 +751,13 @@
                       0,
                       0,
                       0,
-                      0.9950997829437256,
-                      0.00456261495128274,
-                      -0.0002385681145824492,
-                      0.0005880704266019166,
-                      0.0000596420286456123,
-                      -0.0003256454656366259,
-                      -0.0022007909137755632
+                      0.9950998425483704,
+                      0.0045661828480660915,
+                      -0.00023498902737628669,
+                      0.0005928403697907925,
+                      0.00006202756048878655,
+                      -0.0003220661892555654,
+                      -0.002197207184508443
                     ],
                     [
                       0,
@@ -769,13 +769,13 @@
                       0,
                       0,
                       0,
-                      0.98686683177948,
-                      0.020100556313991547,
-                      0.0006405554013326764,
-                      0.0011200773296877742,
-                      0.00014791222929488868,
-                      0.0006286269635893404,
-                      -0.0019180875970050693
+                      0.9868704080581665,
+                      0.02010408788919449,
+                      0.0006465180194936693,
+                      0.0011272316332906485,
+                      0.00015387606981676072,
+                      0.0006274326588027179,
+                      -0.0019156973576173186
                     ],
                     [
                       0,
@@ -787,13 +787,13 @@
                       0,
                       0,
                       0,
-                      0.9520764350891113,
-                      0.07812032103538513,
-                      0.002577728359028697,
-                      0.0021280276123434305,
-                      0.00045924363075755537,
-                      0.00031610275618731976,
-                      -0.0014457227662205696
+                      0.9520777463912964,
+                      0.07812251895666122,
+                      0.0025824937038123608,
+                      0.00213160109706223,
+                      0.000459242524811998,
+                      0.0003208733396604657,
+                      -0.0014481049729511142
                     ],
                     [
                       0,
@@ -805,13 +805,13 @@
                       0,
                       0,
                       0,
-                      0.9528243541717529,
-                      0.07888732105493546,
-                      0.001983693800866604,
-                      0.00022902539058122784,
-                      0.00021232562721706927,
-                      0.00024333946930710226,
-                      -0.0049741449765861034
+                      0.952825665473938,
+                      0.07888712733983994,
+                      0.0019884605426341295,
+                      0.0002314105222467333,
+                      0.00021351795294322073,
+                      0.00024333890178240836,
+                      -0.004971747752279043
                     ],
                     [
                       0,
@@ -823,13 +823,13 @@
                       0,
                       0,
                       0,
-                      0.9087942242622375,
-                      0.10825863480567932,
-                      0.0039960159920156,
-                      -0.0003805161395575851,
-                      0.00021828983153682202,
-                      0.0005618279101327062,
-                      0.009811113588511944
+                      0.9087968468666077,
+                      0.10826075822114944,
+                      0.003999584820121527,
+                      -0.0003757438971661031,
+                      0.0002147107879864052,
+                      0.0005642122705467045,
+                      0.009820632636547089
                     ],
                     [
                       0,
@@ -841,13 +841,13 @@
                       0,
                       0,
                       0,
-                      0.685833215713501,
-                      0.03583889454603195,
-                      0.0005129214259795845,
-                      -0.0008170958026312292,
-                      0.000056063505326164886,
-                      0.0013133174506947398,
-                      0.4153614044189453
+                      0.6858399510383606,
+                      0.0358411967754364,
+                      0.0005200772429816425,
+                      -0.0008099368424154818,
+                      0.000058449048083275557,
+                      0.0013109286082908511,
+                      0.41536638140678406
                     ],
                     [
                       0,
@@ -859,13 +859,13 @@
                       0,
                       0,
                       0,
-                      0.10977353900671005,
-                      0.024679871276021004,
-                      0.00030059582786634564,
-                      -0.0003542736521922052,
-                      0.0002743533405009657,
-                      0.0013980091316625476,
-                      0.8544579148292542
+                      0.10977685451507568,
+                      0.024682197719812393,
+                      0.0003005951002705842,
+                      -0.0003554656286723912,
+                      0.0002815097104758024,
+                      0.001400391454808414,
+                      0.8544618487358093
                     ],
                     [
                       0,
@@ -877,13 +877,13 @@
                       0,
                       0,
                       0,
-                      0.019866758957505226,
-                      0.015398378483951092,
-                      -0.00018489029025658965,
-                      -0.00010258428665110841,
-                      0.0000238568118220428,
-                      0.0006512909312732518,
-                      0.9079329967498779
+                      0.019870290532708168,
+                      0.01540430635213852,
+                      -0.0001848898536991328,
+                      -0.00009423417941434309,
+                      0.000027435267838882282,
+                      0.0006560607580468059,
+                      0.907935619354248
                     ],
                     [
                       0,
@@ -895,13 +895,13 @@
                       0,
                       0,
                       0,
-                      0.019976500421762466,
-                      0.006294619757682085,
-                      0.00020874709298368543,
-                      -0.000741946860216558,
-                      -0.00004532794264378026,
-                      0.0005689849494956434,
-                      0.9069035649299622
+                      0.01998003199696541,
+                      0.006296990439295769,
+                      0.00020993943326175213,
+                      -0.0007359808660112321,
+                      -0.00004174932109890506,
+                      0.0005713692517019808,
+                      0.9069038033485413
                     ],
                     [
                       0,
@@ -1825,9 +1825,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"9d8ac3c7-fd17-40f7-8071-75c5f59936dc\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"9d8ac3c7-fd17-40f7-8071-75c5f59936dc\")) {                    Plotly.newPlot(                        \"9d8ac3c7-fd17-40f7-8071-75c5f59936dc\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"After_0\",\" John_1\",\" and_2\",\" Mary_3\",\" went_4\",\" to_5\",\" the_6\",\" store_7\",\",_8\",\" Mary_9\",\" gave_10\",\" a_11\",\" bottle_12\",\" of_13\",\" milk_14\",\" to_15\"],\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.997469961643219,0.0017654040129855275,-0.0001825046056183055,0.000893437594641,2.027829032158479e-05,-0.0010914491722360253,-0.00047117203939706087],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9971312284469604,0.0026075495406985283,-0.0004866789386142045,0.0006250484730117023,-6.799191032769158e-05,-0.00043896533316001296,-0.0009244514512829483],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9950997829437256,0.00456261495128274,-0.0002385681145824492,0.0005880704266019166,5.96420286456123e-05,-0.0003256454656366259,-0.0022007909137755632],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.98686683177948,0.020100556313991547,0.0006405554013326764,0.0011200773296877742,0.00014791222929488868,0.0006286269635893404,-0.0019180875970050693],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9520764350891113,0.07812032103538513,0.002577728359028697,0.0021280276123434305,0.00045924363075755537,0.00031610275618731976,-0.0014457227662205696],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9528243541717529,0.07888732105493546,0.001983693800866604,0.00022902539058122784,0.00021232562721706927,0.00024333946930710226,-0.0049741449765861034],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9087942242622375,0.10825863480567932,0.0039960159920156,-0.0003805161395575851,0.00021828983153682202,0.0005618279101327062,0.009811113588511944],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.685833215713501,0.03583889454603195,0.0005129214259795845,-0.0008170958026312292,5.6063505326164886e-05,0.0013133174506947398,0.4153614044189453],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.10977353900671005,0.024679871276021004,0.00030059582786634564,-0.0003542736521922052,0.0002743533405009657,0.0013980091316625476,0.8544579148292542],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.019866758957505226,0.015398378483951092,-0.00018489029025658965,-0.00010258428665110841,2.38568118220428e-05,0.0006512909312732518,0.9079329967498779],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.019976500421762466,0.006294619757682085,0.00020874709298368543,-0.000741946860216558,-4.532794264378026e-05,0.0005689849494956434,0.9069035649299622],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Patching GPT-2-small Residual Stream on IOI task\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"e11ca4d0-9f73-4090-9919-2553b6cf1514\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"e11ca4d0-9f73-4090-9919-2553b6cf1514\")) {                    Plotly.newPlot(                        \"e11ca4d0-9f73-4090-9919-2553b6cf1514\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"After_0\",\" John_1\",\" and_2\",\" Mary_3\",\" went_4\",\" to_5\",\" the_6\",\" store_7\",\",_8\",\" Mary_9\",\" gave_10\",\" a_11\",\" bottle_12\",\" of_13\",\" milk_14\",\" to_15\"],\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.997462809085846,0.0017713640118017793,-0.00017296147416345775,0.0008970139897428453,2.7435267838882282e-05,-0.001084289513528347,-0.00046759238466620445],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9971288442611694,0.002613507444038987,-0.0004866778035648167,0.0006274326588027179,-6.202756048878655e-05,-0.0004341929452493787,-0.0009220635402016342],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9950998425483704,0.0045661828480660915,-0.00023498902737628669,0.0005928403697907925,6.202756048878655e-05,-0.0003220661892555654,-0.002197207184508443],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9868704080581665,0.02010408788919449,0.0006465180194936693,0.0011272316332906485,0.00015387606981676072,0.0006274326588027179,-0.0019156973576173186],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9520777463912964,0.07812251895666122,0.0025824937038123608,0.00213160109706223,0.000459242524811998,0.0003208733396604657,-0.0014481049729511142],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.952825665473938,0.07888712733983994,0.0019884605426341295,0.0002314105222467333,0.00021351795294322073,0.00024333890178240836,-0.004971747752279043],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9087968468666077,0.10826075822114944,0.003999584820121527,-0.0003757438971661031,0.0002147107879864052,0.0005642122705467045,0.009820632636547089],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.6858399510383606,0.0358411967754364,0.0005200772429816425,-0.0008099368424154818,5.8449048083275557e-05,0.0013109286082908511,0.41536638140678406],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.10977685451507568,0.024682197719812393,0.0003005951002705842,-0.0003554656286723912,0.0002815097104758024,0.001400391454808414,0.8544618487358093],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.019870290532708168,0.01540430635213852,-0.0001848898536991328,-9.423417941434309e-05,2.7435267838882282e-05,0.0006560607580468059,0.907935619354248],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.01998003199696541,0.006296990439295769,0.00020993943326175213,-0.0007359808660112321,-4.174932109890506e-05,0.0005713692517019808,0.9069038033485413],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Patching GPT-2-small Residual Stream on IOI task\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('9d8ac3c7-fd17-40f7-8071-75c5f59936dc');\n",
+              "var gd = document.getElementById('e11ca4d0-9f73-4090-9919-2553b6cf1514');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
@@ -4374,7 +4374,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.10.15"
+      "version": "3.10.13"
     }
   },
   "nbformat": 4,
diff --git a/docs/source/notebooks/tutorials/attribution_patching.ipynb b/docs/source/notebooks/tutorials/attribution_patching.ipynb
index 66fc9bdf..c1f48859 100644
--- a/docs/source/notebooks/tutorials/attribution_patching.ipynb
+++ b/docs/source/notebooks/tutorials/attribution_patching.ipynb
@@ -58,7 +58,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 23,
+      "execution_count": 12,
       "metadata": {
         "id": "qM0t9saED43y"
       },
@@ -98,7 +98,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 24,
+      "execution_count": 13,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -238,7 +238,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 25,
+      "execution_count": 14,
       "metadata": {
         "id": "bcyK0AoeObY_"
       },
@@ -295,7 +295,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 26,
+      "execution_count": 15,
       "metadata": {
         "id": "05T3hf3LOpFG"
       },
@@ -320,7 +320,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 27,
+      "execution_count": 16,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -367,7 +367,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 28,
+      "execution_count": 17,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -433,7 +433,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 29,
+      "execution_count": 18,
       "metadata": {
         "id": "QdSrOsXvsjjb"
       },
@@ -490,7 +490,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 30,
+      "execution_count": 19,
       "metadata": {
         "id": "cxC8Xs8HsqWm"
       },
@@ -517,7 +517,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 31,
+      "execution_count": 20,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -1667,9 +1667,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"f2bf5d96-2759-4fa9-ac62-e9371c361705\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"f2bf5d96-2759-4fa9-ac62-e9371c361705\")) {                    Plotly.newPlot(                        \"f2bf5d96-2759-4fa9-ac62-e9371c361705\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-6.957763253012672e-05,-6.909317562531214e-09,-4.827359225600958e-05,-1.4337681022880133e-05,6.04442448093323e-06,-5.384786163631361e-06,-0.0002756676112767309,-2.1065989130875096e-05,-1.3769823453912977e-05,8.502480341121554e-06,-8.981670544017106e-05,5.582411904470064e-05],[8.960292325355113e-05,-0.00040236394852399826,-0.00018017760885413736,0.00014379735512193292,-6.819632835686207e-05,-6.25044631306082e-05,-4.3686544813681394e-05,-0.000808828801382333,0.00023929437156766653,0.00015937883290462196,-0.0001314240216743201,9.500579471932724e-05],[1.4101449778536335e-05,-0.00030104274628683925,-0.00037168117705732584,-0.00034691643668338656,-1.0458330507390201e-05,-0.00015961707686074078,-0.000304880493786186,9.84568105195649e-05,0.00026875207549892366,0.0001521407684776932,0.0005039765383116901,-0.0002034870849456638],[-1.1517602615640499e-05,-0.00033706481917761266,-0.00024615979054942727,6.271059100981802e-05,1.354238520434592e-05,0.00034766620956361294,-0.0005866399733349681,-0.0006978895980864763,0.00012641398643609136,-8.654002158436924e-05,-6.751238834112883e-05,0.0002554001403041184],[6.932910764589906e-05,0.00015388225438073277,3.707451105583459e-05,-0.00017302615742664784,9.010283974930644e-05,0.00022782250016462058,-1.3750723155681044e-05,-0.0014550655614584684,-9.826988389249891e-05,0.0005861636018380523,5.1824528782162815e-05,0.0005829725996591151],[1.2706756024272181e-05,-6.604211648664204e-06,0.0010883781360462308,0.0004620166728273034,-0.001068884739652276,-8.976992830866948e-05,4.2892796045634896e-05,-0.00045684745418839157,0.00039122329326346517,9.987624071072787e-05,-0.0018775209318846464,0.0008228498045355082],[0.002574358833953738,0.01968141831457615,-0.0006720172241330147,0.0014327410608530045,-0.0003379757108632475,-0.00020742780179716647,0.0014370887074619532,-0.0007567756110802293,0.0012054562103003263,-0.00020089393365196884,8.37489205878228e-05,4.493692540563643e-05],[-0.00022981743677519262,-0.001819567522034049,-9.668200073065236e-05,0.1371454894542694,-8.344589150510728e-05,-0.009004496037960052,0.0023542074486613274,-0.0004521115333773196,-0.00800877995789051,0.2027597427368164,-0.00011211885430384427,-0.00033840813557617366],[-0.0020026592537760735,-3.5826669773086905e-06,-0.0030021658167243004,-0.0004367050714790821,0.0009537620353512466,-0.0027750192675739527,0.29613298177719116,0.002325658919289708,0.0028100069612264633,0.0005649963277392089,0.27499300241470337,0.0029178475961089134],[0.005365221295505762,-0.0014320836635306478,0.005969846621155739,0.020304212346673012,-0.028698338195681572,0.015889175236225128,-0.05884050577878952,0.057027921080589294,-0.0011342937359586358,0.28792738914489746,0.0005685207433998585,-0.00238024047575891],[0.08650688081979752,0.022697942331433296,0.0005626287311315536,0.007188204675912857,-0.0018677962943911552,0.0002627074427437037,0.04765014350414276,-0.3637135326862335,-0.00011220996384508908,0.00923435389995575,0.053251881152391434,9.949051309376955e-05],[0.007055167108774185,0.0017161732539534569,-0.09743084013462067,0.006113084033131599,-3.512430703267455e-05,-0.0008988302433863282,-0.011853177100419998,0.004282553680241108,0.006519943010061979,-0.007987556979060173,-0.19988438487052917,0.001217787154018879]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Attention Heads\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"adb1d92b-75d9-40e5-9a23-e6141c8f1a77\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"adb1d92b-75d9-40e5-9a23-e6141c8f1a77\")) {                    Plotly.newPlot(                        \"adb1d92b-75d9-40e5-9a23-e6141c8f1a77\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-6.957763253012672e-05,-6.909317562531214e-09,-4.827359225600958e-05,-1.4337681022880133e-05,6.04442448093323e-06,-5.384786163631361e-06,-0.0002756676112767309,-2.1065989130875096e-05,-1.3769823453912977e-05,8.502480341121554e-06,-8.981670544017106e-05,5.582411904470064e-05],[8.960292325355113e-05,-0.00040236394852399826,-0.00018017760885413736,0.00014379735512193292,-6.819632835686207e-05,-6.25044631306082e-05,-4.3686544813681394e-05,-0.000808828801382333,0.00023929437156766653,0.00015937883290462196,-0.0001314240216743201,9.500579471932724e-05],[1.4101449778536335e-05,-0.00030104274628683925,-0.00037168117705732584,-0.00034691643668338656,-1.0458330507390201e-05,-0.00015961707686074078,-0.000304880493786186,9.84568105195649e-05,0.00026875207549892366,0.0001521407684776932,0.0005039765383116901,-0.0002034870849456638],[-1.1517602615640499e-05,-0.00033706481917761266,-0.00024615979054942727,6.271059100981802e-05,1.354238520434592e-05,0.00034766620956361294,-0.0005866399733349681,-0.0006978895980864763,0.00012641398643609136,-8.654002158436924e-05,-6.751238834112883e-05,0.0002554001403041184],[6.932910764589906e-05,0.00015388225438073277,3.707451105583459e-05,-0.00017302615742664784,9.010283974930644e-05,0.00022782250016462058,-1.3750723155681044e-05,-0.0014550655614584684,-9.826988389249891e-05,0.0005861636018380523,5.1824528782162815e-05,0.0005829725996591151],[1.2706756024272181e-05,-6.604211648664204e-06,0.0010883781360462308,0.0004620166728273034,-0.001068884739652276,-8.976992830866948e-05,4.2892796045634896e-05,-0.00045684745418839157,0.00039122329326346517,9.987624071072787e-05,-0.0018775209318846464,0.0008228498045355082],[0.002574358833953738,0.01968141831457615,-0.0006720172241330147,0.0014327410608530045,-0.0003379757108632475,-0.00020742780179716647,0.0014370887074619532,-0.0007567756110802293,0.0012054562103003263,-0.00020089393365196884,8.37489205878228e-05,4.493692540563643e-05],[-0.00022981743677519262,-0.001819567522034049,-9.668200073065236e-05,0.1371454894542694,-8.344589150510728e-05,-0.009004496037960052,0.0023542074486613274,-0.0004521115333773196,-0.00800877995789051,0.2027597427368164,-0.00011211885430384427,-0.00033840813557617366],[-0.0020026592537760735,-3.5826669773086905e-06,-0.0030021658167243004,-0.0004367050714790821,0.0009537620353512466,-0.0027750192675739527,0.29613298177719116,0.002325658919289708,0.0028100069612264633,0.0005649963277392089,0.27499300241470337,0.0029178475961089134],[0.005365221295505762,-0.0014320836635306478,0.005969846621155739,0.020304212346673012,-0.028698338195681572,0.015889175236225128,-0.05884050577878952,0.057027921080589294,-0.0011342937359586358,0.28792738914489746,0.0005685207433998585,-0.00238024047575891],[0.08650688081979752,0.022697942331433296,0.0005626287311315536,0.007188204675912857,-0.0018677962943911552,0.0002627074427437037,0.04765014350414276,-0.3637135326862335,-0.00011220996384508908,0.00923435389995575,0.053251881152391434,9.949051309376955e-05],[0.007055167108774185,0.0017161732539534569,-0.09743084013462067,0.006113084033131599,-3.512430703267455e-05,-0.0008988302433863282,-0.011853177100419998,0.004282553680241108,0.006519943010061979,-0.007987556979060173,-0.19988438487052917,0.001217787154018879]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Attention Heads\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('f2bf5d96-2759-4fa9-ac62-e9371c361705');\n",
+              "var gd = document.getElementById('adb1d92b-75d9-40e5-9a23-e6141c8f1a77');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
@@ -1746,7 +1746,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 32,
+      "execution_count": 21,
       "metadata": {
         "id": "NswiDYEisyw-"
       },
@@ -1771,7 +1771,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 33,
+      "execution_count": 22,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -2892,9 +2892,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"351ace5b-0df3-4d2a-ae14-888e41a01fc6\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"351ace5b-0df3-4d2a-ae14-888e41a01fc6\")) {                    Plotly.newPlot(                        \"351ace5b-0df3-4d2a-ae14-888e41a01fc6\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0006664255633950233,-0.0011972333304584026,-0.0003045523480977863,-0.000452825624961406,-0.00046752975322306156],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00022265268489718437,0.0002414086484350264,0.000127053092001006,0.000524361792486161,-0.0009701023809611797],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,7.685436867177486e-05,0.0002426218707114458,0.0001615851215319708,-0.0004987375577911735,-0.000660655670799315],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1091008186340332,0.0028298422694206238,0.0007662934949621558,-0.0003112038830295205,-0.0012275907211005688],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.01121988520026207,0.020876925438642502,0.0006870393408462405,-0.000619351863861084,5.905964644625783e-05],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.29051679372787476,0.0017006462439894676,0.00023508492449764162,0.0001731544325593859,-0.000579683284740895],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0843031108379364,0.02572883665561676,0.003018079325556755,0.0015348526649177074,0.024284658953547478],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.008083650842308998,0.01003817655146122,0.0003211219154763967,-8.188765787053853e-05,0.3221139907836914],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012348066084086895,0.013155452907085419,0.0004423986538313329,0.0011563682928681374,0.5724781155586243],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0011972422944381833,0.004177889320999384,-1.7678325093584135e-05,0.0003485495108179748,0.3005668520927429],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012172549031674862,0.0010988371213898063,0.00026228593196719885,-9.570646216161549e-05,-0.1382392942905426],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.291185200214386]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Token Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Token Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Token Position\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"3a0bc859-eccc-46fb-ab21-231253cfa4af\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"3a0bc859-eccc-46fb-ab21-231253cfa4af\")) {                    Plotly.newPlot(                        \"3a0bc859-eccc-46fb-ab21-231253cfa4af\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0006664255633950233,-0.0011972333304584026,-0.0003045523480977863,-0.000452825624961406,-0.00046752975322306156],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00022265268489718437,0.0002414086484350264,0.000127053092001006,0.000524361792486161,-0.0009701023809611797],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,7.685436867177486e-05,0.0002426218707114458,0.0001615851215319708,-0.0004987375577911735,-0.000660655670799315],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1091008186340332,0.0028298422694206238,0.0007662934949621558,-0.0003112038830295205,-0.0012275907211005688],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.01121988520026207,0.020876925438642502,0.0006870393408462405,-0.000619351863861084,5.905964644625783e-05],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.29051679372787476,0.0017006462439894676,0.00023508492449764162,0.0001731544325593859,-0.000579683284740895],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0843031108379364,0.02572883665561676,0.003018079325556755,0.0015348526649177074,0.024284658953547478],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.008083650842308998,0.01003817655146122,0.0003211219154763967,-8.188765787053853e-05,0.3221139907836914],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012348066084086895,0.013155452907085419,0.0004423986538313329,0.0011563682928681374,0.5724781155586243],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0011972422944381833,0.004177889320999384,-1.7678325093584135e-05,0.0003485495108179748,0.3005668520927429],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012172549031674862,0.0010988371213898063,0.00026228593196719885,-9.570646216161549e-05,-0.1382392942905426],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.291185200214386]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Token Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Token Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Token Position\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('351ace5b-0df3-4d2a-ae14-888e41a01fc6');\n",
+              "var gd = document.getElementById('3a0bc859-eccc-46fb-ab21-231253cfa4af');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
@@ -2966,7 +2966,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.10.15"
+      "version": "3.10.13"
     },
     "widgets": {
       "application/vnd.jupyter.widget-state+json": {
diff --git a/docs/source/notebooks/tutorials/logit_lens.ipynb b/docs/source/notebooks/tutorials/logit_lens.ipynb
index ef57a3e4..90c8b89e 100644
--- a/docs/source/notebooks/tutorials/logit_lens.ipynb
+++ b/docs/source/notebooks/tutorials/logit_lens.ipynb
@@ -63,7 +63,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": 6,
       "metadata": {
         "id": "_VsZHFh7Fi4X"
       },
@@ -82,7 +82,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 2,
+      "execution_count": 7,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -174,7 +174,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": 8,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -245,7 +245,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": 9,
       "metadata": {
         "id": "a2OkUOXQjmHL"
       },
@@ -307,7 +307,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 10,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -1610,9 +1610,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"1ee45ddf-5c1f-427c-8a04-32189d693375\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"1ee45ddf-5c1f-427c-8a04-32189d693375\")) {                    Plotly.newPlot(                        \"1ee45ddf-5c1f-427c-8a04-32189d693375\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"texttemplate\":\"%{text}\",\"x\":[\"The\",\" E\",\"iff\",\"el\",\" Tower\",\" is\",\" in\",\" the\",\" city\",\" of\"],\"y\":[0,1,2,3,4,5,6,7,8,9,10,11],\"z\":[[0.024983787909150124,0.12145508080720901,0.7902563214302063,0.27752962708473206,0.9977331161499023,0.32375672459602356,0.415214478969574,0.7520487904548645,0.54457688331604,0.6899943947792053],[0.04746942222118378,0.08490755409002304,0.34766101837158203,0.5692889094352722,0.983742356300354,0.2637466788291931,0.6872087717056274,0.8453984260559082,0.6062489151954651,0.8129417896270752],[0.026062989607453346,0.3097169101238251,0.7040867209434509,0.2206706404685974,0.9181663990020752,0.3488497734069824,0.6373876929283142,0.9504272937774658,0.5847728848457336,0.7201948165893555],[0.02635170891880989,0.38433972001075745,0.5039842128753662,0.2591570019721985,0.6246108412742615,0.38496726751327515,0.27967578172683716,0.9723502993583679,0.2764837443828583,0.6877986788749695],[0.02613661251962185,0.30091190338134766,0.21217022836208344,0.18413318693637848,0.7355331778526306,0.48657187819480896,0.3781237006187439,0.9603129029273987,0.41407859325408936,0.4144521951675415],[0.025683913379907608,0.13498641550540924,0.26062414050102234,0.27143916487693787,0.7304731011390686,0.3886910378932953,0.39428195357322693,0.6723183989524841,0.2978633642196655,0.057537395507097244],[0.025425400584936142,0.2042204886674881,0.5427748560905457,0.7381370663642883,0.3724276125431061,0.43281611800193787,0.39778199791908264,0.6245157718658447,0.18699254095554352,0.05666721239686012],[0.02533440850675106,0.21125738322734833,0.7032177448272705,0.9907069206237793,0.23639222979545593,0.4001739025115967,0.46343299746513367,0.7643018960952759,0.24115046858787537,0.13559488952159882],[0.0253128819167614,0.11660420149564743,0.44579413533210754,0.9970935583114624,0.14586897194385529,0.6360733509063721,0.1932840794324875,0.8479137420654297,0.3974367380142212,0.1916261911392212],[0.025318123400211334,0.06472006440162659,0.8146643042564392,0.9981274008750916,0.253418505191803,0.6489886045455933,0.13464125990867615,0.6901037693023682,0.6442651748657227,0.28538209199905396],[0.02550206519663334,0.044399384409189224,0.9033223390579224,0.9942384958267212,0.16611218452453613,0.19855716824531555,0.19015662372112274,0.2531339228153229,0.35554105043411255,0.18341107666492462],[0.013599246740341187,0.04670924320816994,0.9958765506744385,0.9699326753616333,0.09401834011077881,0.1313934326171875,0.24015067517757416,0.11476960778236389,0.18497605621814728,0.0699693113565445]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Input Tokens: %{x}\\u003cbr\\u003eLayers: %{y}\\u003cbr\\u003eProbability: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"text\":[[\"\\\\n\",\" E\",\"iff\",\"el\",\" Tower\",\" not\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"Z\",\"ness\",\"ter\",\" Tower\",\" not\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"ter\",\" Tower\",\" now\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"stra\",\" Tower\",\" also\",\" the\",\" same\",\" hall\",\" the\"],[\" the\",\"leg\",\"ness\",\"horn\",\" Tower\",\" also\",\" front\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Towers\",\" also\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" also\",\" front\",\" midst\",\"scape\",\" East\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Ing\"],[\" the\",\"TS\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Rome\"],[\" the\",\"FP\",\"el\",\" Tower\",\" tower\",\" located\",\" London\",\" midst\",\" centre\",\" London\"],[\" the\",\"nerg\",\"el\",\" Tower\",\" is\",\" located\",\" the\",\" midst\",\"'s\",\" Paris\"],[\"\\\\n\",\"-\",\"el\",\" Tower\",\" is\",\" a\",\" the\",\" middle\",\" centre\",\" Paris\"]]}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Input Tokens\"},\"tickangle\":0},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layers\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Probability\"}},\"colorscale\":[[0.0,\"rgb(49,54,149)\"],[0.1,\"rgb(69,117,180)\"],[0.2,\"rgb(116,173,209)\"],[0.3,\"rgb(171,217,233)\"],[0.4,\"rgb(224,243,248)\"],[0.5,\"rgb(255,255,191)\"],[0.6,\"rgb(254,224,144)\"],[0.7,\"rgb(253,174,97)\"],[0.8,\"rgb(244,109,67)\"],[0.9,\"rgb(215,48,39)\"],[1.0,\"rgb(165,0,38)\"]],\"cmid\":0.5},\"margin\":{\"t\":60},\"title\":{\"text\":\"Logit Lens Visualization\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"b3911d4b-d13c-4c2c-be99-472396f13a45\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"b3911d4b-d13c-4c2c-be99-472396f13a45\")) {                    Plotly.newPlot(                        \"b3911d4b-d13c-4c2c-be99-472396f13a45\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"texttemplate\":\"%{text}\",\"x\":[\"The\",\" E\",\"iff\",\"el\",\" Tower\",\" is\",\" in\",\" the\",\" city\",\" of\"],\"y\":[0,1,2,3,4,5,6,7,8,9,10,11],\"z\":[[0.024983787909150124,0.12145508080720901,0.7902563214302063,0.27752962708473206,0.9977331161499023,0.32375672459602356,0.415214478969574,0.7520487904548645,0.54457688331604,0.6899943947792053],[0.04746942222118378,0.08490755409002304,0.34766101837158203,0.5692889094352722,0.983742356300354,0.2637466788291931,0.6872087717056274,0.8453984260559082,0.6062489151954651,0.8129417896270752],[0.026062989607453346,0.3097169101238251,0.7040867209434509,0.2206706404685974,0.9181663990020752,0.3488497734069824,0.6373876929283142,0.9504272937774658,0.5847728848457336,0.7201948165893555],[0.02635170891880989,0.38433972001075745,0.5039842128753662,0.2591570019721985,0.6246108412742615,0.38496726751327515,0.27967578172683716,0.9723502993583679,0.2764837443828583,0.6877986788749695],[0.02613661251962185,0.30091190338134766,0.21217022836208344,0.18413318693637848,0.7355331778526306,0.48657187819480896,0.3781237006187439,0.9603129029273987,0.41407859325408936,0.4144521951675415],[0.025683913379907608,0.13498641550540924,0.26062414050102234,0.27143916487693787,0.7304731011390686,0.3886910378932953,0.39428195357322693,0.6723183989524841,0.2978633642196655,0.057537395507097244],[0.025425400584936142,0.2042204886674881,0.5427748560905457,0.7381370663642883,0.3724276125431061,0.43281611800193787,0.39778199791908264,0.6245157718658447,0.18699254095554352,0.05666721239686012],[0.02533440850675106,0.21125738322734833,0.7032177448272705,0.9907069206237793,0.23639222979545593,0.4001739025115967,0.46343299746513367,0.7643018960952759,0.24115046858787537,0.13559488952159882],[0.0253128819167614,0.11660420149564743,0.44579413533210754,0.9970935583114624,0.14586897194385529,0.6360733509063721,0.1932840794324875,0.8479137420654297,0.3974367380142212,0.1916261911392212],[0.025318123400211334,0.06472006440162659,0.8146643042564392,0.9981274008750916,0.253418505191803,0.6489886045455933,0.13464125990867615,0.6901037693023682,0.6442651748657227,0.28538209199905396],[0.02550206519663334,0.044399384409189224,0.9033223390579224,0.9942384958267212,0.16611218452453613,0.19855716824531555,0.19015662372112274,0.2531339228153229,0.35554105043411255,0.18341107666492462],[0.013599246740341187,0.04670924320816994,0.9958765506744385,0.9699326753616333,0.09401834011077881,0.1313934326171875,0.24015067517757416,0.11476960778236389,0.18497605621814728,0.0699693113565445]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Input Tokens: %{x}\\u003cbr\\u003eLayers: %{y}\\u003cbr\\u003eProbability: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"text\":[[\"\\\\n\",\" E\",\"iff\",\"el\",\" Tower\",\" not\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"Z\",\"ness\",\"ter\",\" Tower\",\" not\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"ter\",\" Tower\",\" now\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"stra\",\" Tower\",\" also\",\" the\",\" same\",\" hall\",\" the\"],[\" the\",\"leg\",\"ness\",\"horn\",\" Tower\",\" also\",\" front\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Towers\",\" also\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" also\",\" front\",\" midst\",\"scape\",\" East\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Ing\"],[\" the\",\"TS\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Rome\"],[\" the\",\"FP\",\"el\",\" Tower\",\" tower\",\" located\",\" London\",\" midst\",\" centre\",\" London\"],[\" the\",\"nerg\",\"el\",\" Tower\",\" is\",\" located\",\" the\",\" midst\",\"'s\",\" Paris\"],[\"\\\\n\",\"-\",\"el\",\" Tower\",\" is\",\" a\",\" the\",\" middle\",\" centre\",\" Paris\"]]}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Input Tokens\"},\"tickangle\":0},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layers\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Probability\"}},\"colorscale\":[[0.0,\"rgb(49,54,149)\"],[0.1,\"rgb(69,117,180)\"],[0.2,\"rgb(116,173,209)\"],[0.3,\"rgb(171,217,233)\"],[0.4,\"rgb(224,243,248)\"],[0.5,\"rgb(255,255,191)\"],[0.6,\"rgb(254,224,144)\"],[0.7,\"rgb(253,174,97)\"],[0.8,\"rgb(244,109,67)\"],[0.9,\"rgb(215,48,39)\"],[1.0,\"rgb(165,0,38)\"]],\"cmid\":0.5},\"margin\":{\"t\":60},\"title\":{\"text\":\"Logit Lens Visualization\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('1ee45ddf-5c1f-427c-8a04-32189d693375');\n",
+              "var gd = document.getElementById('b3911d4b-d13c-4c2c-be99-472396f13a45');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
@@ -1695,7 +1695,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.10.15"
+      "version": "3.10.13"
     },
     "widgets": {
       "application/vnd.jupyter.widget-state+json": {
diff --git a/docs/source/notebooks/tutorials/walkthrough.ipynb b/docs/source/notebooks/tutorials/walkthrough.ipynb
index 978e4c13..ba6297ee 100644
--- a/docs/source/notebooks/tutorials/walkthrough.ipynb
+++ b/docs/source/notebooks/tutorials/walkthrough.ipynb
@@ -14,22 +14,21 @@
       },
       "source": [
         "An interactive version of this walkthrough can be found\n",
-        "[here](https://colab.research.google.com/github/ndif-team/nnsight/blob/main/NNsight_Walkthrough.ipynb)\n",
+        "[here](https://colab.research.google.com/github/ndif-team/nnsight/blob/main/NNsight_Walkthrough.ipynb).\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",
+        "In this era of large-scale deep learning, often the most interesting AI models are\n",
+        "massive black boxes that are hard to run. Ordinary commercial inference APIs \n",
+        "let us interact with huge models, but they do not let us access model internals.\n",
         "\n",
-        "The nnsight library is different: it provides full access to all the neural\n",
+        "The `nnsight` library is different: it provides full access to all 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",
+        "[National Deep Inference Fabric](https://ndif.us/)\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"
+        "big insights on this profoundly important field.\n"
       ]
     },
     {
@@ -311,7 +310,7 @@
         "id": "sBz5qfwuR-6F"
       },
       "source": [
-        "Oh no an error! \"Accessing Proxy value before it's been set.\"\n",
+        "Oh no, an error! `ValueError: Accessing value before it's been set.`\n",
         "\n",
         "Why doesn't our `output` have a `value`?\n",
         "\n",
@@ -377,8 +376,8 @@
         "\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",
+        "Just like we saved the model's final output, we can save the output of\n",
+        "any of its submodules. Using normal Python attribute syntax, we can discover\n",
         "how to access them by name by printing out the model:\n"
       ]
     },
@@ -450,7 +449,7 @@
       },
       "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",
+        "the `as tracer`, as we won't be needing the `tracer` object for a few\n",
         "sections:\n"
       ]
     },
@@ -597,7 +596,7 @@
         "id": "BGPUJvWq_bOp"
       },
       "source": [
-        "Nice! That worked seamlessly, but hold on, how come we didn't need to call\n",
+        "Nice! That worked seamlessly, but why didn't we 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",
@@ -724,9 +723,7 @@
       },
       "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",
+        "really insightful, and a number of ML techniques do exactly that. However, we often not only want to view the computation of a model, but also to influence it.\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",
@@ -850,9 +847,9 @@
       "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",
+        "How can we find what went wrong? 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",
+        "Enter \"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:"
@@ -1181,7 +1178,7 @@
       "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. "
+        "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 example where this could be particularly useful is working with SAEs - we could train an SAE on one layer and then stop model execution. "
       ]
     },
     {
@@ -1266,7 +1263,7 @@
       "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."
+        "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."
       ]
     },
     {
@@ -1785,7 +1782,7 @@
         "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",
+        "In a case like this, the underlying model is called more than once, so 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",
@@ -2036,7 +2033,7 @@
         "id": "mvXLOh65MXmF"
       },
       "source": [
-        "This only needs to be run once as it will save this api key as the default in a\n",
+        "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",
@@ -2210,7 +2207,7 @@
       "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.)."
+        "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.)."
       ]
     },
     {
@@ -2249,7 +2246,7 @@
       "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! "
+        "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 could swap activations between vanilla and instruct versions of the Llama model and compare their outputs. The `session` context can also be used to run experiments entirely locally! "
       ]
     },
     {
@@ -2263,7 +2260,7 @@
       "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",
+        "We mention earlier that the `session` context enables multi-tracing execution. But how could we optimize a process that requires running an intervention graph in a loop? If we create a simple `for` loop with a **Tracer context** inside, this would create 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",
@@ -2306,7 +2303,7 @@
       "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:"
+        "The `Iterator` context extends all the `nnsight` graph-based functionalities but also closely mimics the conventional `for` loop statement in Python, allowing it to support all kind of iterative operations with a use of `as item` syntax:"
       ]
     },
     {
@@ -2434,11 +2431,11 @@
       "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",
+        "The **Iterator** context is a nice 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",
+        "But what kinds of experimental scenarios would make this useful?\n",
         "\n",
-        "In the next section, we delve into a powerful use case of the `Iterator` context and see how it enables it!"
+        "In the next section, we see how the `Iterator` context enables a powerful parameter-efficient fine-tuning tool, Low-Rank Adaptation (LoRA)!"
       ]
     },
     {
@@ -2452,7 +2449,14 @@
       "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."
+        "For more background information about LoRA, check out [this paper](https://arxiv.org/abs/2106.09685)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "Training a LoRA will integrate 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."
       ]
     },
     {
@@ -2510,7 +2514,7 @@
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "Let's define all the variables to use in LORA training. "
+        "Let's define all the variables to use in LoRA training. "
       ]
     },
     {
@@ -2559,7 +2563,7 @@
       "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."
+        "It's time to run the LoRA training loop! We using the **Session** and the **Iterator** contexts to achieve this."
       ]
     },
     {
@@ -2785,6 +2789,16 @@
         "print(llama.tokenizer.batch_decode(out.value))\n"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# Next Steps\n",
+        "\n",
+        "Check out [nnsight.net/tutorials](https://nnsight.net/tutorials) for more walkthroughs\n",
+        "implementating classic interpretability techniques using `nnsight`. "
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -2800,21 +2814,16 @@
         "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",
+        "Note that both `nnsight` and `NDIF` are in active development, \n",
+        "so changes may be made and errors may arise during use.\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",
+        "[NDIF discord](https://discord.gg/6uFJmCSwW7). We'd love to hear about your work using `nnsight`! \n",
         "\n",
-        "You can also follow us on X/Twitter [@ndif_team](https://x.com/ndif_team).\n",
+        "You can also follow us on [LinkedIn](https://www.linkedin.com/company/national-deep-inference-fabric) and X/Twitter: [@ndif_team](https://x.com/ndif_team).\n",
         "\n",
-        "💟\n"
+        "💟"
       ]
     }
   ],
@@ -2837,7 +2846,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.12.5"
+      "version": "3.10.13"
     },
     "widgets": {
       "application/vnd.jupyter.widget-state+json": {
diff --git a/docs/source/start.rst b/docs/source/start.rst
index 28b12e04..0b0b6071 100644
--- a/docs/source/start.rst
+++ b/docs/source/start.rst
@@ -19,7 +19,7 @@ Please give the project a :ndif:`star on Github` to support the project. NNsight
 First Steps
 -----------
 
-.. grid:: 2 2 3 3
+.. grid:: 2 2 2 2
    :gutter: 2
 
    .. grid-item-card:: Walkthrough
@@ -27,6 +27,11 @@ First Steps
 
       Walk through the basic functionality of the package.
 
+   .. grid-item-card:: Remote Access
+      :link: notebooks/features/remote_execution.ipynb
+
+      Configure API access for remote model execution.
+
    .. grid-item-card:: Features
       :link: features
       :link-type: doc
@@ -37,7 +42,7 @@ First Steps
       :link: tutorials
       :link-type: doc
 
-      See implementations of :bdg-primary:`nnsight` through common interpretability techniques.
+      See :bdg-primary:`nnsight` implementations of common interpretability techniques.
 
 
 

From 2e83cbfd4cacd816183bd6255c0eadac71f640b1 Mon Sep 17 00:00:00 2001
From: ebortz <ebortz@bu.edu>
Date: Thu, 21 Nov 2024 10:36:53 -0500
Subject: [PATCH 2/8] another typo fix

---
 docs/source/notebooks/tutorials/walkthrough.ipynb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/notebooks/tutorials/walkthrough.ipynb b/docs/source/notebooks/tutorials/walkthrough.ipynb
index ba6297ee..598f988c 100644
--- a/docs/source/notebooks/tutorials/walkthrough.ipynb
+++ b/docs/source/notebooks/tutorials/walkthrough.ipynb
@@ -2563,7 +2563,7 @@
       "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."
+        "It's time to run the LoRA training loop! We will be using the **Session** and the **Iterator** contexts to achieve this."
       ]
     },
     {

From f0c3021a15b4bbfa566cf6628e7477ab7ae5dd14 Mon Sep 17 00:00:00 2001
From: ebortz <ebortz@bu.edu>
Date: Thu, 21 Nov 2024 10:51:50 -0500
Subject: [PATCH 3/8] adjusted LoRA icon

---
 docs/source/tutorials.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst
index 19258fd7..a092dd19 100644
--- a/docs/source/tutorials.rst
+++ b/docs/source/tutorials.rst
@@ -139,7 +139,7 @@ Tutorials
 
          <div class="d-flex align-items-center">
             <div class="d-flex justify-content-center" style="min-width: 50px; margin-right: 15px; height: 100%;">
-               <i class="fa-solid fa-arrow-down-a-z fa-2x"></i>
+               <i class="fa-solid fa-sliders fa-2x"></i>
             </div> 
             <div>
                <h5 class="card-title">LoRA</h5>

From 13721d24fd3491f1fa82cd2ee2086a2d6a98d078 Mon Sep 17 00:00:00 2001
From: Adam Belfki <belfki.adam@gmail.com>
Date: Mon, 2 Dec 2024 11:48:40 -0500
Subject: [PATCH 4/8] bug: fixing visualization + reducing memory usage

---
 .../tutorials/activation_patching.ipynb       | 247 ++++++++++--------
 1 file changed, 132 insertions(+), 115 deletions(-)

diff --git a/docs/source/notebooks/tutorials/activation_patching.ipynb b/docs/source/notebooks/tutorials/activation_patching.ipynb
index 7b12ad0c..cb22820f 100644
--- a/docs/source/notebooks/tutorials/activation_patching.ipynb
+++ b/docs/source/notebooks/tutorials/activation_patching.ipynb
@@ -104,19 +104,35 @@
       "metadata": {},
       "outputs": [],
       "source": [
-        "from IPython.display import clear_output"
+        "try:\n",
+        "    import google.colab\n",
+        "    is_colab = True\n",
+        "except ImportError:\n",
+        "    is_colab = False\n",
+        "\n",
+        "if is_colab:\n",
+        "    !pip install -U nnsight"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": 2,
       "metadata": {},
+      "outputs": [],
+      "source": [
+        "from IPython.display import clear_output"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {},
       "outputs": [
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+            "/Users/adam/NDIF/mib-env/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
             "  from .autonotebook import tqdm as notebook_tqdm\n"
           ]
         }
@@ -137,7 +153,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": 4,
       "metadata": {
         "id": "l8WfCzQhwFM0"
       },
@@ -145,14 +161,14 @@
       "source": [
         "import plotly.express as px\n",
         "import plotly.io as pio\n",
-        "pio.renderers.default = \"plotly_mimetype+notebook_connected+colab+notebook\"\n",
+        "pio.renderers.default = \"colab\" if is_colab else \"plotly_mimetype+notebook_connected+colab+notebook\"\n",
         "from nnsight import LanguageModel, util\n",
         "from nnsight.tracing.Proxy import Proxy"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": 5,
       "metadata": {
         "id": "qLmX2tdZmgiz"
       },
@@ -165,7 +181,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 6,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -226,7 +242,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 7,
       "metadata": {
         "id": "7FrduLBCmhOp"
       },
@@ -249,7 +265,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 8,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -314,7 +330,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 9,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -381,7 +397,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 10,
       "metadata": {
         "id": "JodVtd5VAuo2"
       },
@@ -418,7 +434,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 10,
+      "execution_count": 11,
       "metadata": {
         "id": "1vuK3rEMAug_"
       },
@@ -477,7 +493,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 13,
+      "execution_count": 12,
       "metadata": {
         "id": "G60gi4ZUmxlG"
       },
@@ -541,7 +557,7 @@
         "                    clean_logit_diff - corrupted_logit_diff\n",
         "                )\n",
         "\n",
-        "                _ioi_patching_results.append(patched_result.save())\n",
+        "                _ioi_patching_results.append(patched_result.item().save())\n",
         "\n",
         "        ioi_patching_results.append(_ioi_patching_results)"
       ]
@@ -566,7 +582,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 11,
+      "execution_count": 13,
       "metadata": {
         "id": "EaHSh6ctnJFl"
       },
@@ -577,7 +593,7 @@
         "                              x_labels,\n",
         "                              plot_title=\"Normalized Logit Difference After Patching Residual Stream on the IOI Task\"):\n",
         "\n",
-        "    ioi_patching_results = util.apply(ioi_patching_results, lambda x: x.value.item(), Proxy)\n",
+        "    ioi_patching_results = util.apply(ioi_patching_results, lambda x: x.value, Proxy)\n",
         "    fig = px.imshow(\n",
         "        ioi_patching_results,\n",
         "        color_continuous_midpoint=0.0,\n",
@@ -601,7 +617,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 14,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -715,13 +731,13 @@
                       0,
                       0,
                       0,
-                      0.997462809085846,
-                      0.0017713640118017793,
-                      -0.00017296147416345775,
-                      0.0008970139897428453,
-                      0.000027435267838882282,
-                      -0.001084289513528347,
-                      -0.00046759238466620445
+                      0.9974616169929504,
+                      0.0017701691249385476,
+                      -0.0001777326106093824,
+                      0.0008922415436245501,
+                      0.000023856726329540834,
+                      -0.0010866738157346845,
+                      -0.0004663989820983261
                     ],
                     [
                       0,
@@ -733,13 +749,13 @@
                       0,
                       0,
                       0,
-                      0.9971288442611694,
-                      0.002613507444038987,
-                      -0.0004866778035648167,
-                      0.0006274326588027179,
-                      -0.00006202756048878655,
-                      -0.0004341929452493787,
-                      -0.0009220635402016342
+                      0.9971300363540649,
+                      0.002608733018860221,
+                      -0.00048309870180673897,
+                      0.0006298175430856645,
+                      -0.00006441315781557932,
+                      -0.00043419242138043046,
+                      -0.0009184839436784387
                     ],
                     [
                       0,
@@ -752,12 +768,12 @@
                       0,
                       0,
                       0.9950998425483704,
-                      0.0045661828480660915,
-                      -0.00023498902737628669,
-                      0.0005928403697907925,
-                      0.00006202756048878655,
-                      -0.0003220661892555654,
-                      -0.002197207184508443
+                      0.004563791677355766,
+                      -0.0002361815859330818,
+                      0.0005916468217037618,
+                      0.00006441315781557932,
+                      -0.0003244514809921384,
+                      -0.0021983974147588015
                     ],
                     [
                       0,
@@ -769,13 +785,13 @@
                       0,
                       0,
                       0,
-                      0.9868704080581665,
-                      0.02010408788919449,
-                      0.0006465180194936693,
-                      0.0011272316332906485,
-                      0.00015387606981676072,
-                      0.0006274326588027179,
-                      -0.0019156973576173186
+                      0.9868668913841248,
+                      0.020101677626371384,
+                      0.0006417459226213396,
+                      0.0011212661629542708,
+                      0.0001526830455986783,
+                      0.000627431902103126,
+                      -0.001915695145726204
                     ],
                     [
                       0,
@@ -787,13 +803,13 @@
                       0,
                       0,
                       0,
-                      0.9520777463912964,
-                      0.07812251895666122,
-                      0.0025824937038123608,
-                      0.00213160109706223,
-                      0.000459242524811998,
-                      0.0003208733396604657,
-                      -0.0014481049729511142
+                      0.9520730376243591,
+                      0.07812362164258957,
+                      0.0025824906770139933,
+                      0.0021339841187000275,
+                      0.00046401331201195717,
+                      0.00031968014081940055,
+                      -0.0014445247361436486
                     ],
                     [
                       0,
@@ -805,13 +821,13 @@
                       0,
                       0,
                       0,
-                      0.952825665473938,
-                      0.07888712733983994,
-                      0.0019884605426341295,
-                      0.0002314105222467333,
-                      0.00021351795294322073,
-                      0.00024333890178240836,
-                      -0.004971747752279043
+                      0.9528281092643738,
+                      0.07888465374708176,
+                      0.001984879607334733,
+                      0.0002278317406307906,
+                      0.0002147105406038463,
+                      0.0002433386107441038,
+                      -0.004974127281457186
                     ],
                     [
                       0,
@@ -823,13 +839,13 @@
                       0,
                       0,
                       0,
-                      0.9087968468666077,
-                      0.10826075822114944,
-                      0.003999584820121527,
-                      -0.0003757438971661031,
-                      0.0002147107879864052,
-                      0.0005642122705467045,
-                      0.009820632636547089
+                      0.9087921380996704,
+                      0.10825704783201218,
+                      0.003999580163508654,
+                      -0.0003733577614184469,
+                      0.0002206747158197686,
+                      0.0005677900626324117,
+                      0.00981227122247219
                     ],
                     [
                       0,
@@ -841,13 +857,13 @@
                       0,
                       0,
                       0,
-                      0.6858399510383606,
-                      0.0358411967754364,
-                      0.0005200772429816425,
-                      -0.0008099368424154818,
-                      0.000058449048083275557,
-                      0.0013109286082908511,
-                      0.41536638140678406
+                      0.6858367323875427,
+                      0.03584115207195282,
+                      0.0005200766026973724,
+                      -0.0008170928922481835,
+                      0.00005487047019414604,
+                      0.0013109270948916674,
+                      0.41536349058151245
                     ],
                     [
                       0,
@@ -859,13 +875,13 @@
                       0,
                       0,
                       0,
-                      0.10977685451507568,
-                      0.024682197719812393,
-                      0.0003005951002705842,
-                      -0.0003554656286723912,
-                      0.0002815097104758024,
-                      0.001400391454808414,
-                      0.8544618487358093
+                      0.10977672785520554,
+                      0.02468455396592617,
+                      0.00030059475102461874,
+                      -0.000356658041710034,
+                      0.000279123691143468,
+                      0.0014027755241841078,
+                      0.8544548749923706
                     ],
                     [
                       0,
@@ -877,13 +893,13 @@
                       0,
                       0,
                       0,
-                      0.019870290532708168,
-                      0.01540430635213852,
-                      -0.0001848898536991328,
-                      -0.00009423417941434309,
-                      0.000027435267838882282,
-                      0.0006560607580468059,
-                      0.907935619354248
+                      0.019870266318321228,
+                      0.01540309563279152,
+                      -0.00018250395078212023,
+                      -0.00009781257540453225,
+                      0.00002982090700243134,
+                      0.0006572527927346528,
+                      0.9079344868659973
                     ],
                     [
                       0,
@@ -895,13 +911,13 @@
                       0,
                       0,
                       0,
-                      0.01998003199696541,
-                      0.006296990439295769,
-                      0.00020993943326175213,
-                      -0.0007359808660112321,
-                      -0.00004174932109890506,
-                      0.0005713692517019808,
-                      0.9069038033485413
+                      0.01998358592391014,
+                      0.006295789964497089,
+                      0.0002087463508360088,
+                      -0.0007407513330690563,
+                      -0.00004294210521038622,
+                      0.0005737542524002492,
+                      0.9069039225578308
                     ],
                     [
                       0,
@@ -1825,9 +1841,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"e11ca4d0-9f73-4090-9919-2553b6cf1514\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"e11ca4d0-9f73-4090-9919-2553b6cf1514\")) {                    Plotly.newPlot(                        \"e11ca4d0-9f73-4090-9919-2553b6cf1514\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"After_0\",\" John_1\",\" and_2\",\" Mary_3\",\" went_4\",\" to_5\",\" the_6\",\" store_7\",\",_8\",\" Mary_9\",\" gave_10\",\" a_11\",\" bottle_12\",\" of_13\",\" milk_14\",\" to_15\"],\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.997462809085846,0.0017713640118017793,-0.00017296147416345775,0.0008970139897428453,2.7435267838882282e-05,-0.001084289513528347,-0.00046759238466620445],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9971288442611694,0.002613507444038987,-0.0004866778035648167,0.0006274326588027179,-6.202756048878655e-05,-0.0004341929452493787,-0.0009220635402016342],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9950998425483704,0.0045661828480660915,-0.00023498902737628669,0.0005928403697907925,6.202756048878655e-05,-0.0003220661892555654,-0.002197207184508443],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9868704080581665,0.02010408788919449,0.0006465180194936693,0.0011272316332906485,0.00015387606981676072,0.0006274326588027179,-0.0019156973576173186],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9520777463912964,0.07812251895666122,0.0025824937038123608,0.00213160109706223,0.000459242524811998,0.0003208733396604657,-0.0014481049729511142],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.952825665473938,0.07888712733983994,0.0019884605426341295,0.0002314105222467333,0.00021351795294322073,0.00024333890178240836,-0.004971747752279043],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9087968468666077,0.10826075822114944,0.003999584820121527,-0.0003757438971661031,0.0002147107879864052,0.0005642122705467045,0.009820632636547089],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.6858399510383606,0.0358411967754364,0.0005200772429816425,-0.0008099368424154818,5.8449048083275557e-05,0.0013109286082908511,0.41536638140678406],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.10977685451507568,0.024682197719812393,0.0003005951002705842,-0.0003554656286723912,0.0002815097104758024,0.001400391454808414,0.8544618487358093],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.019870290532708168,0.01540430635213852,-0.0001848898536991328,-9.423417941434309e-05,2.7435267838882282e-05,0.0006560607580468059,0.907935619354248],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.01998003199696541,0.006296990439295769,0.00020993943326175213,-0.0007359808660112321,-4.174932109890506e-05,0.0005713692517019808,0.9069038033485413],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Patching GPT-2-small Residual Stream on IOI task\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"8451ae81-59bf-4461-a1e8-0f5b0b76cc45\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"8451ae81-59bf-4461-a1e8-0f5b0b76cc45\")) {                    Plotly.newPlot(                        \"8451ae81-59bf-4461-a1e8-0f5b0b76cc45\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"After_0\",\" John_1\",\" and_2\",\" Mary_3\",\" went_4\",\" to_5\",\" the_6\",\" store_7\",\",_8\",\" Mary_9\",\" gave_10\",\" a_11\",\" bottle_12\",\" of_13\",\" milk_14\",\" to_15\"],\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9974616169929504,0.0017701691249385476,-0.0001777326106093824,0.0008922415436245501,2.3856726329540834e-05,-0.0010866738157346845,-0.0004663989820983261],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9971300363540649,0.002608733018860221,-0.00048309870180673897,0.0006298175430856645,-6.441315781557932e-05,-0.00043419242138043046,-0.0009184839436784387],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9950998425483704,0.004563791677355766,-0.0002361815859330818,0.0005916468217037618,6.441315781557932e-05,-0.0003244514809921384,-0.0021983974147588015],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9868668913841248,0.020101677626371384,0.0006417459226213396,0.0011212661629542708,0.0001526830455986783,0.000627431902103126,-0.001915695145726204],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9520730376243591,0.07812362164258957,0.0025824906770139933,0.0021339841187000275,0.00046401331201195717,0.00031968014081940055,-0.0014445247361436486],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9528281092643738,0.07888465374708176,0.001984879607334733,0.0002278317406307906,0.0002147105406038463,0.0002433386107441038,-0.004974127281457186],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.9087921380996704,0.10825704783201218,0.003999580163508654,-0.0003733577614184469,0.0002206747158197686,0.0005677900626324117,0.00981227122247219],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.6858367323875427,0.03584115207195282,0.0005200766026973724,-0.0008170928922481835,5.487047019414604e-05,0.0013109270948916674,0.41536349058151245],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.10977672785520554,0.02468455396592617,0.00030059475102461874,-0.000356658041710034,0.000279123691143468,0.0014027755241841078,0.8544548749923706],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.019870266318321228,0.01540309563279152,-0.00018250395078212023,-9.781257540453225e-05,2.982090700243134e-05,0.0006572527927346528,0.9079344868659973],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.01998358592391014,0.006295789964497089,0.0002087463508360088,-0.0007407513330690563,-4.294210521038622e-05,0.0005737542524002492,0.9069039225578308],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Patching GPT-2-small Residual Stream on IOI task\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('e11ca4d0-9f73-4090-9919-2553b6cf1514');\n",
+              "var gd = document.getElementById('8451ae81-59bf-4461-a1e8-0f5b0b76cc45');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
@@ -1933,7 +1949,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 15,
+      "execution_count": 17,
       "metadata": {
         "id": "2AHdnHc0l8Cj"
       },
@@ -1942,15 +1958,16 @@
         "# from google.colab import userdata\n",
         "from nnsight import CONFIG\n",
         "\n",
-        "CONFIG.set_default_api_key('YOUR_API_KEY')\n",
+        "if is_colab:\n",
+        "    CONFIG.set_default_api_key('YOUR_API_KEY')\n",
+        "    %huggingface-cli login -token 'YOUR_HF_TOKEN'\n",
         "\n",
-        "%huggingface-cli login -token 'YOUR_HF_TOKEN'\n",
         "clear_output()"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 16,
+      "execution_count": 18,
       "metadata": {
         "id": "fppg2PwAOZOh"
       },
@@ -1963,7 +1980,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 17,
+      "execution_count": 19,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -2024,7 +2041,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 18,
+      "execution_count": 20,
       "metadata": {
         "id": "q9NSOwiL1k2V"
       },
@@ -2047,7 +2064,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 19,
+      "execution_count": 21,
       "metadata": {
         "id": "MMmQOD1P1fRP"
       },
@@ -2061,7 +2078,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 20,
+      "execution_count": 22,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -2116,7 +2133,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 21,
+      "execution_count": 23,
       "metadata": {
         "id": "Eava539mV7zr"
       },
@@ -2136,7 +2153,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 22,
+      "execution_count": 24,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -2149,15 +2166,15 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "2024-11-12 11:26:27,428 6ca48f09-4bb3-4c74-8ff1-71d063d270a1 - RECEIVED: Your job has been received and is waiting approval.\n",
-            "2024-11-12 11:26:27,429 6ca48f09-4bb3-4c74-8ff1-71d063d270a1 - APPROVED: Your job was approved and is waiting to be run.\n",
-            "2024-11-12 11:26:27,467 6ca48f09-4bb3-4c74-8ff1-71d063d270a1 - RUNNING: Your job has started running.\n",
-            "2024-11-12 11:26:42,620 6ca48f09-4bb3-4c74-8ff1-71d063d270a1 - COMPLETED: Your job has been completed.\n",
-            "Downloading result:   0%|          | 0.00/276k [00:00<?, ?B/s]huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
+            "2024-12-02 11:45:39,137 6053779a-6455-4c56-ab27-ba744fefda09 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2024-12-02 11:45:39,140 6053779a-6455-4c56-ab27-ba744fefda09 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2024-12-02 11:45:39,141 6053779a-6455-4c56-ab27-ba744fefda09 - RUNNING: Your job has started running.\n",
+            "2024-12-02 11:45:54,353 6053779a-6455-4c56-ab27-ba744fefda09 - COMPLETED: Your job has been completed.\n",
+            "Downloading result:   0%|          | 0.00/34.5k [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%|██████████| 276k/276k [00:00<00:00, 3.45MB/s]\n"
+            "Downloading result: 100%|██████████| 34.5k/34.5k [00:00<00:00, 5.99MB/s]\n"
           ]
         }
       ],
@@ -2232,7 +2249,7 @@
         "                    clean_logit_diff - corrupted_logit_diff\n",
         "                )\n",
         "\n",
-        "                _ioi_patching_results.append(patched_result.save())\n",
+        "                _ioi_patching_results.append(patched_result.item().save())\n",
         "\n",
         "        ioi_patching_results.append(_ioi_patching_results)"
       ]
@@ -2257,7 +2274,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 23,
+      "execution_count": 25,
       "metadata": {
         "id": "RH2ms_7dUoFs"
       },
@@ -4311,9 +4328,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"df7fda37-94c8-45a5-a4c1-5d1cce43115b\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"df7fda37-94c8-45a5-a4c1-5d1cce43115b\")) {                    Plotly.newPlot(                        \"df7fda37-94c8-45a5-a4c1-5d1cce43115b\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"Head 0\",\"Head 1\",\"Head 2\",\"Head 3\",\"Head 4\",\"Head 5\",\"Head 6\",\"Head 7\",\"Head 8\",\"Head 9\",\"Head 10\",\"Head 11\",\"Head 12\",\"Head 13\",\"Head 14\",\"Head 15\",\"Head 16\",\"Head 17\",\"Head 18\",\"Head 19\",\"Head 20\",\"Head 21\",\"Head 22\",\"Head 23\",\"Head 24\",\"Head 25\",\"Head 26\",\"Head 27\",\"Head 28\",\"Head 29\",\"Head 30\",\"Head 31\"],\"z\":[[0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0130615234375,-0.00653076171875,-0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.0130615234375,0.00653076171875,0.00653076171875],[0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.00653076171875,-0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.0],[0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,-0.00653076171875,0.00653076171875],[0.00653076171875,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.0130615234375],[0.0,0.0130615234375,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.11767578125,0.0,0.0,0.0,0.0130615234375,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0130615234375],[0.0,0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.0,0.0130615234375,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,-0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875],[0.00653076171875,0.0,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.0,-0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0],[0.00653076171875,0.00653076171875,0.0,0.0,-0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0],[0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,-0.0196533203125,0.0,-0.00653076171875,0.0,0.0,0.0,0.0,0.0130615234375,0.0,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,-0.00653076171875,0.0,0.0,0.00653076171875],[0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.0130615234375,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875],[0.026123046875,-0.026123046875,-0.0196533203125,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,-0.00653076171875,-0.00653076171875,-0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,-0.0130615234375,-0.00653076171875,0.0,0.00653076171875,0.00653076171875],[0.0,0.0,-0.00653076171875,0.0,0.0,0.00653076171875,-0.0130615234375,0.0,-0.00653076171875,0.00653076171875,-0.0130615234375,-0.00653076171875,0.0,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875],[-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0130615234375,0.00653076171875,0.0,0.00653076171875,0.0,0.0196533203125,-0.0130615234375,-0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0130615234375,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875],[-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.0,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0130615234375,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.0,0.0],[0.0,0.0,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.045654296875,0.0,0.0,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.058837890625,0.0,0.00653076171875,0.07177734375,0.0,0.0196533203125,-0.00653076171875,0.0,0.0,0.0,-0.00653076171875,0.0,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875],[0.0,0.0,0.0,-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.1240234375,0.0,-0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,-0.00653076171875],[-0.0130615234375,0.058837890625,0.03271484375,-0.0196533203125,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,-0.058837890625,0.03271484375,0.0,0.0,0.0,0.0,0.0849609375,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.039306640625,0.00653076171875,-0.03271484375,0.0],[0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.039306640625,-0.0196533203125,0.0,0.169921875,0.0,0.0,0.00653076171875,0.00653076171875],[0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0,-0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0130615234375,0.0,0.00653076171875,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,-0.0130615234375,0.09814453125,0.00653076171875,0.0],[0.0,0.0,0.00653076171875,-0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0,0.0,0.0,-0.045654296875,0.0,0.0,0.0849609375,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0],[0.00653076171875,0.0,0.00653076171875,0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,-0.111328125,0.0849609375,-0.248046875,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875],[0.0,0.0,0.0196533203125,-0.03271484375,0.0,0.0,0.00653076171875,0.0,-0.00653076171875,0.0,0.03271484375,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875],[0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.03271484375,0.0,-0.0196533203125,0.00653076171875,0.00653076171875,0.0,-0.03271484375,0.169921875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0130615234375,0.00653076171875],[0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0130615234375,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.0],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.039306640625,0.00653076171875,0.00653076171875,0.00653076171875,0.0],[0.0,0.0,0.00653076171875,0.0,0.0,0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1044921875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0],[0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1376953125,0.0,0.0,-0.0654296875,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0],[0.00653076171875,0.0,0.0,0.0,0.0,0.0,-0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1630859375,0.0130615234375,0.0196533203125,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0],[0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0130615234375,-0.05224609375,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0196533203125,0.0,-0.1044921875,0.039306640625,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0130615234375,0.09130859375,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,-0.07861328125,0.0130615234375,0.0196533203125,0.00653076171875,0.0,0.00653076171875,0.00653076171875],[0.0,0.00653076171875,-0.00653076171875,0.0130615234375,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,-0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Patching Llama Attention Heads on IOI task\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"c9d299b2-9469-4a3a-9581-3fbed40c4ba3\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"c9d299b2-9469-4a3a-9581-3fbed40c4ba3\")) {                    Plotly.newPlot(                        \"c9d299b2-9469-4a3a-9581-3fbed40c4ba3\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"x\":[\"Head 0\",\"Head 1\",\"Head 2\",\"Head 3\",\"Head 4\",\"Head 5\",\"Head 6\",\"Head 7\",\"Head 8\",\"Head 9\",\"Head 10\",\"Head 11\",\"Head 12\",\"Head 13\",\"Head 14\",\"Head 15\",\"Head 16\",\"Head 17\",\"Head 18\",\"Head 19\",\"Head 20\",\"Head 21\",\"Head 22\",\"Head 23\",\"Head 24\",\"Head 25\",\"Head 26\",\"Head 27\",\"Head 28\",\"Head 29\",\"Head 30\",\"Head 31\"],\"z\":[[0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0130615234375,-0.00653076171875,-0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.0130615234375,0.00653076171875,0.00653076171875],[0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.00653076171875,-0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.0],[0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,-0.00653076171875,0.00653076171875],[0.00653076171875,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.0130615234375],[0.0,0.0130615234375,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.11767578125,0.0,0.0,0.0,0.0130615234375,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0130615234375],[0.0,0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.0,0.0130615234375,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,-0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875],[0.00653076171875,0.0,0.0130615234375,0.0,0.00653076171875,0.00653076171875,0.0,-0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0],[0.00653076171875,0.00653076171875,0.0,0.0,-0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0],[0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,-0.0196533203125,0.0,-0.00653076171875,0.0,0.0,0.0,0.0,0.0130615234375,0.0,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,-0.00653076171875,0.0,0.0,0.00653076171875],[0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,0.0,-0.0130615234375,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875],[0.026123046875,-0.026123046875,-0.0196533203125,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,-0.00653076171875,-0.00653076171875,-0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0,0.00653076171875,-0.0130615234375,-0.00653076171875,0.0,0.00653076171875,0.00653076171875],[0.0,0.0,-0.00653076171875,0.0,0.0,0.00653076171875,-0.0130615234375,0.0,-0.00653076171875,0.00653076171875,-0.0130615234375,-0.00653076171875,0.0,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.00653076171875,0.00653076171875,-0.00653076171875,0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875],[-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0130615234375,0.00653076171875,0.0,0.00653076171875,0.0,0.0196533203125,-0.0130615234375,-0.00653076171875,0.00653076171875,0.0,-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0130615234375,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875],[-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.0,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0130615234375,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.0,0.0],[0.0,0.0,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.045654296875,0.0,0.0,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.058837890625,0.0,0.00653076171875,0.07177734375,0.0,0.0196533203125,-0.00653076171875,0.0,0.0,0.0,-0.00653076171875,0.0,-0.00653076171875,0.0,0.00653076171875,-0.00653076171875],[0.0,0.0,0.0,-0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0130615234375,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.1240234375,0.0,-0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,-0.00653076171875],[-0.0130615234375,0.058837890625,0.03271484375,-0.0196533203125,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,-0.058837890625,0.03271484375,0.0,0.0,0.0,0.0,0.0849609375,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.039306640625,0.00653076171875,-0.03271484375,0.0],[0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.039306640625,-0.0196533203125,0.0,0.169921875,0.0,0.0,0.00653076171875,0.00653076171875],[0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,-0.00653076171875,0.0,-0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0130615234375,0.0,0.00653076171875,-0.00653076171875,-0.00653076171875,0.0,0.00653076171875,0.00653076171875,-0.0130615234375,0.09814453125,0.00653076171875,0.0],[0.0,0.0,0.00653076171875,-0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,-0.00653076171875,0.0,0.0,0.0,0.0,-0.045654296875,0.0,0.0,0.0849609375,0.00653076171875,0.0,0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0],[0.00653076171875,0.0,0.00653076171875,0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,-0.111328125,0.0849609375,-0.248046875,0.0,0.00653076171875,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875],[0.0,0.0,0.0196533203125,-0.03271484375,0.0,0.0,0.00653076171875,0.0,-0.00653076171875,0.0,0.03271484375,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875],[0.0,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.03271484375,0.0,-0.0196533203125,0.00653076171875,0.00653076171875,0.0,-0.03271484375,0.169921875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0130615234375,0.00653076171875],[0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0130615234375,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0130615234375,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.0,0.0],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.00653076171875,0.00653076171875,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.00653076171875,0.039306640625,0.00653076171875,0.00653076171875,0.00653076171875,0.0],[0.0,0.0,0.00653076171875,0.0,0.0,0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1044921875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0],[0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1376953125,0.0,0.0,-0.0654296875,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0],[0.00653076171875,0.0,0.0,0.0,0.0,0.0,-0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1630859375,0.0130615234375,0.0196533203125,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0],[0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0130615234375,-0.05224609375,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0196533203125,0.0,-0.1044921875,0.039306640625,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0130615234375,0.09130859375,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00653076171875,-0.07861328125,0.0130615234375,0.0196533203125,0.00653076171875,0.0,0.00653076171875,0.00653076171875],[0.0,0.00653076171875,-0.00653076171875,0.0130615234375,0.0,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,0.00653076171875,0.0,0.00653076171875,0.0,0.0,0.0,0.0,0.0,-0.0130615234375,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Patching Llama Attention Heads on IOI task\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('df7fda37-94c8-45a5-a4c1-5d1cce43115b');\n",
+              "var gd = document.getElementById('c9d299b2-9469-4a3a-9581-3fbed40c4ba3');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
@@ -4374,7 +4391,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.10.13"
+      "version": "3.10.14"
     }
   },
   "nbformat": 4,

From df545a72b68b5a4edf7d187769c8159ea66a639b Mon Sep 17 00:00:00 2001
From: Adam Belfki <belfki.adam@gmail.com>
Date: Mon, 2 Dec 2024 12:01:38 -0500
Subject: [PATCH 5/8] remove cell

---
 .../notebooks/tutorials/activation_patching.ipynb      | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/docs/source/notebooks/tutorials/activation_patching.ipynb b/docs/source/notebooks/tutorials/activation_patching.ipynb
index cb22820f..e4e2401d 100644
--- a/docs/source/notebooks/tutorials/activation_patching.ipynb
+++ b/docs/source/notebooks/tutorials/activation_patching.ipynb
@@ -88,16 +88,6 @@
         "# Set Up"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "If using Colab, install NNsight:\n",
-        "```\n",
-        "!pip install -U nnsight\n",
-        "```"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": 1,

From 7e29967fe20de16a090630b95181115b9790af17 Mon Sep 17 00:00:00 2001
From: ebortz <ebortz@bu.edu>
Date: Fri, 31 Jan 2025 16:59:26 -0500
Subject: [PATCH 6/8] nnsight 0.4 walkthroughs

---
 NNsight_Walkthrough.ipynb                     | 4913 +++++++++++-----
 NNsight_v0_4_guide.ipynb                      | 4223 ++++++++++++++
 .../tutorials/activation_patching.ipynb       |    2 +-
 .../tutorials/attribution_patching.ipynb      |   30 +-
 .../notebooks/tutorials/logit_lens.ipynb      |   14 +-
 .../notebooks/tutorials/walkthrough.ipynb     | 4956 ++++++++++++-----
 6 files changed, 11411 insertions(+), 2727 deletions(-)
 create mode 100644 NNsight_v0_4_guide.ipynb

diff --git a/NNsight_Walkthrough.ipynb b/NNsight_Walkthrough.ipynb
index 3b3fc64a..60db8609 100644
--- a/NNsight_Walkthrough.ipynb
+++ b/NNsight_Walkthrough.ipynb
@@ -2,9 +2,11 @@
   "cells": [
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "i4WnS0_mkkww"
+      },
       "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"
+        "<a href=\"https://colab.research.google.com/github/ndif-team/nnsight/blob/main/NNsight_Walkthrough.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
     {
@@ -24,7 +26,7 @@
         "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",
+        "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",
@@ -32,7 +34,7 @@
         "\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"
+        "big insights on a profoundly important field."
       ]
     },
     {
@@ -46,22 +48,15 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {
-        "id": "WyLMmhrAKTNM"
-      },
+      "metadata": {},
       "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"
+        "## Setup"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 1,
-      "metadata": {
-        "id": "R0RJijD0eXwf"
-      },
+      "execution_count": null,
+      "metadata": {},
       "outputs": [],
       "source": [
         "# Install nnsight\n",
@@ -73,19 +68,31 @@
         "clear_output()"
       ]
     },
+    {
+      "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": "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",
+        "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,
+      "execution_count": 1,
       "metadata": {
         "id": "2pX2Wg8Ceo6N"
       },
@@ -115,16 +122,25 @@
       },
       "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"
+        "PyTorch model to enable investigation of its internal parameters.\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 2,
       "metadata": {
         "id": "8H9R_ynTJI5y"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+            "  from .autonotebook import tqdm as notebook_tqdm\n"
+          ]
+        }
+      ],
       "source": [
         "import nnsight\n",
         "from nnsight import NNsight\n",
@@ -138,19 +154,19 @@
         "id": "2NfISlQ_Ilvp"
       },
       "source": [
-        "Printing a Pytorch model shows a named hierarchy of modules which is very useful\n",
+        "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,
+      "execution_count": 3,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "bYtnbJHvlGZV",
-        "outputId": "5eb7c572-451c-43c6-956f-af224ae1867c"
+        "outputId": "9bec5a3e-647c-4c49-b4ec-e1420ea0c964"
       },
       "outputs": [
         {
@@ -181,7 +197,7 @@
         "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",
+        "The most common application is opening files as in the following example:\n",
         "\n",
         "```python\n",
         "with open('myfile.txt', 'r') as file:\n",
@@ -192,7 +208,9 @@
         "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",
+        "read from the file.\n",
+        "\n",
+        "Simple enough! Now we can discuss how `nnsight` uses\n",
         "contexts to enable intuitive access into the internals of a neural network.\n"
       ]
     },
@@ -212,7 +230,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 4,
       "metadata": {
         "id": "qEXQ4auPSL-m"
       },
@@ -270,20 +288,19 @@
         "`.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"
+        "grab it for me and put the value into its corresponding Proxy object. Let's try it:\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 5,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 303
+          "height": 329
         },
         "id": "cYe8-r9ptGaG",
-        "outputId": "adf25efe-d2cf-4f65-9eaf-15b155762f4b"
+        "outputId": "5a1ebb19-be61-4eea-9ec1-141a62ebc313"
       },
       "outputs": [
         {
@@ -293,9 +310,10 @@
           "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",
+            "Cell \u001b[0;32mIn[5], 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;28;43mprint\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43moutput\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/proxy.py:70\u001b[0m, in \u001b[0;36mProxy.__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     66\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     68\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---> 70\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     72\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[39mtarget\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;124m\"\u001b[39m\n",
+            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/proxy.py:64\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     57\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     58\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     59\u001b[0m \n\u001b[1;32m     60\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     61\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     62\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 64\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/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/node.py:143\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    133\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    134\u001b[0m \n\u001b[1;32m    135\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or was destroyed).\u001b[39;00m\n\u001b[1;32m    140\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    142\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--> 143\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    145\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."
           ]
         }
@@ -305,40 +323,40 @@
         "\n",
         "    output = tiny_model.output\n",
         "\n",
-        "print(output.value)"
+        "print(output)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "sBz5qfwuR-6F"
+        "id": "4CGa5zSI6kkI"
       },
       "source": [
-        "Oh no an error! \"Accessing Proxy value before it's been set.\"\n",
+        "Oh no an error! \"Accessing 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"
+        "error:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 67,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "y_bXRd5dvsBu",
-        "outputId": "a8e8985c-2661-46a8-975b-d43f67cfc162"
+        "outputId": "feded172-02f5-429b-f9a9-766eed3a787d"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "tensor([[-0.2687, -0.3314]])\n"
+            "tensor([[-0.1301, -0.4906]])\n"
           ]
         }
       ],
@@ -347,7 +365,7 @@
         "\n",
         "    output = tiny_model.output.save()\n",
         "\n",
-        "print(output.value)"
+        "print(output)"
       ]
     },
     {
@@ -387,13 +405,13 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 68,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "0WcVUSP-0CJi",
-        "outputId": "7fd9d6ac-152d-4431-9a4d-0b60321638eb"
+        "outputId": "f40600e1-719d-46c2-f420-e5a9826dd315"
       },
       "outputs": [
         {
@@ -413,28 +431,30 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "MXn9XRQqkkw1"
+      },
       "source": [
-        "Let's access the output of the first layer (non-coincidentally named 'layer1'):"
+        "Let's access the output of the first layer (which we've named 'layer1'):"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 69,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "akWr-cNqy-9O",
-        "outputId": "68b9e071-c609-4b70-906d-2f2cd93b4495"
+        "outputId": "13faddb0-fead-4fd0-a500-2bff00c095b3"
       },
       "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"
+            "tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -443,7 +463,7 @@
         "\n",
         "    l1_output = tiny_model.layer1.output.save()\n",
         "\n",
-        "print(l1_output.value)"
+        "print(l1_output)"
       ]
     },
     {
@@ -459,21 +479,21 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 10,
+      "execution_count": 70,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "3EHEN38N3nXR",
-        "outputId": "5afd17a7-0bcc-4d1a-c131-7038bde9455d"
+        "outputId": "74f25b48-89fd-4b15-9dab-cead54f2386b"
       },
       "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"
+            "tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -482,7 +502,7 @@
         "\n",
         "    l2_input = tiny_model.layer2.input.save()\n",
         "\n",
-        "print(l2_input.value)"
+        "print(l2_input)"
       ]
     },
     {
@@ -512,24 +532,32 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "d5DE1ZJAkkw1"
+      },
       "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",
+        "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": {},
+      "execution_count": 71,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "oZzS86tzkkw1",
+        "outputId": "1c1f3a94-eb4d-48c7-c69f-9241f2fa08f5"
+      },
       "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"
+            "Layer 1 - out:  tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -563,13 +591,13 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 72,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "c5XCiSZn2p3k",
-        "outputId": "97d34dee-7bb5-4735-8c53-6aac6377af9e"
+        "outputId": "c6c5060b-4cd9-4154-8bc6-adf4f5aec5cc"
       },
       "outputs": [
         {
@@ -587,7 +615,7 @@
         "    # 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",
+        "    # 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",
@@ -613,20 +641,20 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 13,
+      "execution_count": 73,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "FIcOYeEjFuln",
-        "outputId": "d836abd3-b713-44c5-aedc-95bf85abf11a"
+        "outputId": "074fd78f-bc18-48b8-ccef-7db9df756ea6"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "tensor(-1.3797)\n"
+            "tensor(0.9377)\n"
           ]
         }
       ],
@@ -645,39 +673,49 @@
       },
       "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",
+        "the given `input`. When the output of `tiny_model.layer1` is computed, take its sum. Then do\n",
+        "the same for `tiny_model.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": {},
+      "metadata": {
+        "id": "kLmzsHaMkkw2"
+      },
       "source": [
         "## Custom Functions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "GJFCuZHtkkw2"
+      },
       "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",
+        "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",
+        "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": {},
+      "execution_count": 74,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "DF7MPuhckkw2",
+        "outputId": "c263468b-affe-4f69-a045-d674e07333d3"
+      },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "tensor(-0.7796) tensor(-0.7796)\n"
+            "tensor(1.5584) tensor(1.5584)\n"
           ]
         }
       ],
@@ -688,27 +726,29 @@
         "    total = 0\n",
         "    for element in flat:\n",
         "        total += element.item()\n",
-        "    \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",
+        "    # Specify the function name and its arguments (in a comma-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",
         "\n",
         "print(custom_sum, sum)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "02UKwk_Lkkw2"
+      },
       "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."
+        "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."
       ]
     },
     {
@@ -727,34 +767,32 @@
       },
       "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",
+        "really insightful, and a number of ML techniques do exactly that. However, often we not only want to view the computation of a model, but also to influence it.\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:"
+        "this really easy using the '=' operator:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 15,
+      "execution_count": 75,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "-6y2wzJqOz3a",
-        "outputId": "6274e5c5-c5fe-4e07-87c4-a119cfd09742"
+        "outputId": "13de53a9-c396-4167-e2f6-950957233c8e"
       },
       "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"
+            "Before: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n",
+            "After: tensor([[ 0.0000,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -781,52 +819,38 @@
         "id": "oZz9SMs3Y_iS"
       },
       "source": [
-        "Seems our change was reflected. Now the same for the last dimension:\n"
+        "Seems our change was reflected. Now let's do the same for the last dimension:\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 16,
+      "execution_count": 76,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 363
+          "height": 305
         },
         "id": "qwlqvHFcxld2",
-        "outputId": "28ae7175-5af0-4bd0-f9f0-21680a4fd748"
+        "outputId": "b1ac58fb-dbbd-4375-fcd7-3362762f00fe"
       },
       "outputs": [
         {
-          "ename": "IndexError",
-          "evalue": "Above exception when execution Node: 'setitem_0' in Graph: '6063279136'",
+          "ename": "NNsightError",
+          "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)",
-            "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'"
+            "Traceback (most recent call last):",
+            "  File \"/Users/emmabortz/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py\", line 297, in execute",
+            "    output = self.target(*args, **kwargs)",
+            "IndexError: index 10 is out of bounds for dimension 1 with size 10",
+            "",
+            "During handling of the above exception, another exception occurred:",
+            "",
+            "Traceback (most recent call last):",
+            "  File \"/var/folders/rx/0nl_h2cd54q90chs9hf0n5qr0000gn/T/ipykernel_80599/3404137504.py\", line 8, in <module>",
+            "    tiny_model.layer1.output[:, hidden_dims] = 0",
+            "",
+            "NNsightError: index 10 is out of bounds for dimension 1 with size 10"
           ]
         }
       ],
@@ -849,121 +873,260 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "DKaN1oi76djp"
+      },
       "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",
+        "Oh no, we are getting an error! Ah of course, we needed to index at `hidden_dims - 1` not `hidden_dims`.\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:"
+        "If you've been using `nnsight`, you are probably familiar with error messages that can be quite difficult to troubleshoot. In `nnsight 0.4` we've now improved error messaging to be descriptive and line-specific, as you should see in the above example!"
       ]
     },
     {
-      "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"
-          ]
-        }
-      ],
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Yz-XpUpZ36yd"
+      },
       "source": [
-        "# turn on scan and validate\n",
-        "with tiny_model.trace(input, scan=True, validate=True):\n",
+        "<details>\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",
+        "<summary>\n",
+        "Old NNsight error messaging\n",
+        "</summary>\n",
         "\n",
-        "print(\"Before:\", l1_output_before)\n",
-        "print(\"After:\", l1_output_after)"
+        "If you've been using NNsight prior to the NNsight 0.4 release, you will be familiar with the following non-descriptive error messaging. If you choose to turn off NNsight 0.4's new error messaging feature, this is how errors within the tracing context will appear.\n",
+        "\n",
+        "\n",
+        "```\n",
+        "---------------------------------------------------------------------------\n",
+        "IndexError                                Traceback (most recent call last)\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in execute(self)\n",
+        "    379                 # Call the target to get value.\n",
+        "--> 380                 output = self.target(*args, **kwargs)\n",
+        "    381\n",
+        "\n",
+        "IndexError: index 10 is out of bounds for dimension 1 with size 10\n",
+        "\n",
+        "The above exception was the direct cause of the following exception:\n",
+        "\n",
+        "IndexError                                Traceback (most recent call last)\n",
+        "20 frames\n",
+        "<ipython-input-16-5c81de91fb1f> in <cell line: 0>()\n",
+        "----> 1 with tiny_model.trace(input):\n",
+        "      2\n",
+        "      3     # Save the output before the edit to compare.\n",
+        "      4     # Notice we apply .clone() before saving as the setting operation is in-place.\n",
+        "      5     l1_output_before = tiny_model.layer1.output.clone().save()\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/Tracer.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
+        "    100\n",
+        "    101\n",
+        "--> 102         super().__exit__(exc_type, exc_val, exc_tb)\n",
+        "    103\n",
+        "    104     def invoke(self, *inputs: Any, **kwargs) -> Invoker:\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/GraphBasedContext.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
+        "    215             raise exc_val\n",
+        "    216\n",
+        "--> 217         self.backend(self)\n",
+        "    218\n",
+        "    219     ### BACKENDS ########\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/backends/LocalBackend.py in __call__(self, obj)\n",
+        "     25     def __call__(self, obj: LocalMixin):\n",
+        "     26\n",
+        "---> 27         obj.local_backend_execute()\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/Tracer.py in local_backend_execute(self)\n",
+        "    144         self.graph.execute()\n",
+        "    145\n",
+        "--> 146         self.model.interleave(\n",
+        "    147             self.model._execute,\n",
+        "    148             self.graph,\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in interleave(self, fn, intervention_graph, *inputs, **kwargs)\n",
+        "    467         module_paths = InterventionProtocol.get_interventions(intervention_graph).keys()\n",
+        "    468\n",
+        "--> 469         with HookHandler(\n",
+        "    470             self._model,\n",
+        "    471             list(module_paths),\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
+        "    579\n",
+        "    580         if isinstance(exc_val, Exception):\n",
+        "--> 581             raise exc_val\n",
+        "    582\n",
+        "    583\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in interleave(self, fn, intervention_graph, *inputs, **kwargs)\n",
+        "    478         ):\n",
+        "    479             try:\n",
+        "--> 480                 fn(*inputs, **kwargs)\n",
+        "    481             except protocols.EarlyStopProtocol.EarlyStopException:\n",
+        "    482                 # TODO: Log.\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in _execute(self, *prepared_inputs, **kwargs)\n",
+        "    585             pass\n",
+        "    586\n",
+        "--> 587         return self._model(\n",
+        "    588             *prepared_inputs,\n",
+        "    589             **kwargs,\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)\n",
+        "   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]\n",
+        "   1735         else:\n",
+        "-> 1736             return self._call_impl(*args, **kwargs)\n",
+        "   1737\n",
+        "   1738     # torchrec tests the code consistency with the following code\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)\n",
+        "   1842\n",
+        "   1843         try:\n",
+        "-> 1844             return inner()\n",
+        "   1845         except Exception:\n",
+        "   1846             # run always called hooks if they have not already been run\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in inner()\n",
+        "   1788                 args = bw_hook.setup_input_hook(args)\n",
+        "   1789\n",
+        "-> 1790             result = forward_call(*args, **kwargs)\n",
+        "   1791             if _global_forward_hooks or self._forward_hooks:\n",
+        "   1792                 for hook_id, hook in (\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/container.py in forward(self, input)\n",
+        "    248     def forward(self, input):\n",
+        "    249         for module in self:\n",
+        "--> 250             input = module(input)\n",
+        "    251         return input\n",
+        "    252\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)\n",
+        "   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]\n",
+        "   1735         else:\n",
+        "-> 1736             return self._call_impl(*args, **kwargs)\n",
+        "   1737\n",
+        "   1738     # torchrec tests the code consistency with the following code\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)\n",
+        "   1842\n",
+        "   1843         try:\n",
+        "-> 1844             return inner()\n",
+        "   1845         except Exception:\n",
+        "   1846             # run always called hooks if they have not already been run\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in inner()\n",
+        "   1801                         hook_result = hook(self, args, kwargs, result)\n",
+        "   1802                     else:\n",
+        "-> 1803                         hook_result = hook(self, args, result)\n",
+        "   1804\n",
+        "   1805                     if hook_result is not None:\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in output_hook(module, input, output, module_path)\n",
+        "    564\n",
+        "    565                 def output_hook(module, input, output, module_path=module_path):\n",
+        "--> 566                     return self.output_hook(output, module_path)\n",
+        "    567\n",
+        "    568                 self.handles.append(\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in <lambda>(activations, module_path)\n",
+        "    473                 activations, module_path, \"input\", intervention_handler\n",
+        "    474             ),\n",
+        "--> 475             output_hook=lambda activations, module_path: InterventionProtocol.intervene(\n",
+        "    476                 activations, module_path, \"output\", intervention_handler\n",
+        "    477             ),\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in intervene(cls, activations, module_path, key, intervention_handler)\n",
+        "    454\n",
+        "    455                 # Value injection.\n",
+        "--> 456                 node.set_value(value)\n",
+        "    457\n",
+        "    458                 # Check if through the previous value injection, there was a 'swap' intervention.\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in set_value(self, value)\n",
+        "    408\n",
+        "    409             if listener.fulfilled() and not self.graph.sequential:\n",
+        "--> 410                 listener.execute()\n",
+        "    411\n",
+        "    412         for dependency in self.arg_dependencies:\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in execute(self)\n",
+        "    385         except Exception as e:\n",
+        "    386\n",
+        "--> 387             raise type(e)(\n",
+        "    388                 f\"Above exception when execution Node: '{self.name}' in Graph: '{self.graph.id}'\"\n",
+        "    389             ) from e\n",
+        "\n",
+        "IndexError: Above exception when execution Node: 'setitem_0' in Graph: '132147685816016'\n",
+        "\n",
+        "\n",
+        "```\n",
+        "\n",
+        "\n",
+        "</details>\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "JCNR_Jkpxr8l"
+        "id": "Lo7oHC0yzI_p"
       },
       "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",
+        "The error messaging feature can be toggled using `nnsight.CONFIG.APP.DEBUG` which defaults to true.\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",
+        "<summary>\n",
+        "Toggle Error Messaging \n",
+        "</summary>\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",
+        "Turn off debugging:\n",
+        "```\n",
+        "import nnsight\n",
         "\n",
-        "---\n",
+        "nnsight.CONFIG.APP.DEBUG = False\n",
+        "nnsight.CONFIG.save()\n",
+        "```\n",
         "\n",
-        "</details>\n",
+        "Turn on debugging:\n",
+        "```\n",
+        "import nnsight\n",
         "\n",
-        "Let's try again with the correct indexing, and view the shape of the output\n",
+        "nnsight.CONFIG.APP.DEBUG = True\n",
+        "nnsight.CONFIG.save()\n",
+        "```\n",
+        "</details>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JCNR_Jkpxr8l"
+      },
+      "source": [
+        "Now that we know more about NNsight's error messaging, let's try our setting operation again with the correct indexing and view the shape of the output\n",
         "before leaving the tracing context:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 18,
+      "execution_count": 77,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "Wf8ugJKB2mru",
-        "outputId": "4eed6c19-cb8c-4deb-e0e1-25aa80959c64"
+        "outputId": "f076fa36-0271-4185-8583-9e7190e1e8e9"
       },
       "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"
+            "Layer 1 output shape: InterventionProxy (fetch_attr)\n",
+            "Before: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n",
+            "After: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0000]])\n"
           ]
         }
       ],
@@ -989,50 +1152,64 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "9DH7PeL13-WU"
+        "id": "6R3wj1BKkkw5"
       },
       "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"
+        "## Scan and Validate\n",
+        "Error codes are helpful, but sometimes you may want to quickly troubleshoot your code without actually running it. \n",
+        "\n",
+        "Enter \"Scanning\" and \"Validating\"! We can enable this features by setting the `scan=True` and `validate=True` flag in the `trace` method."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\"Scanning\" runs \"fake\" inputs throught the model to collect information like shapes and types (i.e., scanning will populate all called `.inputs` and `.outputs`). \n",
+        "\n",
+        "\"Validating\" attempts to execute the intervention proxies with \"fake\" inputs to check if they work (i.e., executes all interventions in your code with fake tensors).\n",
+        "\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. Let's try it out on our example above."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 19,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "AP5wRmz_4YE7",
-        "outputId": "4c30dc1d-232d-4900-bfdd-0c93a9b23ad7"
-      },
+      "execution_count": 78,
+      "metadata": {},
       "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"
+          "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[78], 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;01mwith\u001b[39;00m tiny_model\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;28minput\u001b[39m, scan\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, validate\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\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",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/contexts/interleaving.py:96\u001b[0m, in \u001b[0;36mInterleavingTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     92\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     94\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m---> 96\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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/tracer.py:25\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mglobals\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m GlobalTracingContext\n\u001b[1;32m     23\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mtry_deregister(\u001b[38;5;28mself\u001b[39m)\n\u001b[0;32m---> 25\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/base.py:72\u001b[0m, in \u001b[0;36mContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     69\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\u001b[38;5;241m.\u001b[39mstack\u001b[38;5;241m.\u001b[39mpop()\n\u001b[1;32m     71\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---> 72\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n\u001b[1;32m     74\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madd(graph\u001b[38;5;241m.\u001b[39mstack[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m], graph, \u001b[38;5;241m*\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[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkwargs)\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbackend \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",
+            "Cell \u001b[0;32mIn[78], 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~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:126\u001b[0m, in \u001b[0;36mProxy.__setitem__\u001b[0;34m(self, key, value)\u001b[0m\n\u001b[1;32m    125\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[Self, 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--> 126\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    127\u001b[0m \u001b[43m        \u001b[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    128\u001b[0m \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\n\u001b[1;32m    129\u001b[0m \u001b[43m        \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    130\u001b[0m \u001b[43m        \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    131\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py:250\u001b[0m, in \u001b[0;36mNode.create\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    247\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m value\n\u001b[1;32m    249\u001b[0m \u001b[38;5;66;03m# Otherwise just create the Node on the Graph like normal.\u001b[39;00m\n\u001b[0;32m--> 250\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    251\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    252\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    253\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/graph.py:131\u001b[0m, in \u001b[0;36mGraph.create\u001b[0;34m(self, target, redirect, *args, **kwargs)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;66;03m# Redirection.\u001b[39;00m\n\u001b[1;32m    129\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstack[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m] \u001b[38;5;28;01mif\u001b[39;00m redirect \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstack \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[0;32m--> 131\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mproxy_class(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode_class\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\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[43mgraph\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~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:118\u001b[0m, in \u001b[0;36mValidatingInterventionNode.__init__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m    114\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattached\n\u001b[1;32m    115\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfake_value \u001b[38;5;129;01mis\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39m_empty\n\u001b[1;32m    116\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m Protocol\u001b[38;5;241m.\u001b[39mis_protocol(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtarget)\n\u001b[1;32m    117\u001b[0m ):\n\u001b[0;32m--> 118\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfake_value \u001b[38;5;241m=\u001b[39m \u001b[43mvalidate\u001b[49m\u001b[43m(\u001b[49m\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[43m \u001b[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[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[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:147\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m FakeTensorMode(\n\u001b[1;32m    142\u001b[0m     allow_non_fake_inputs\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    143\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    144\u001b[0m ) \u001b[38;5;28;01mas\u001b[39;00m fake_mode:\n\u001b[1;32m    145\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m FakeCopyMode(fake_mode):\n\u001b[0;32m--> 147\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m GlobalTracingContext\u001b[38;5;241m.\u001b[39mexit_global_tracing_context():\n\u001b[1;32m    149\u001b[0m             \u001b[38;5;28;01mif\u001b[39;00m backwards_check(target, \u001b[38;5;241m*\u001b[39margs):\n\u001b[1;32m    150\u001b[0m                 \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/globals.py:100\u001b[0m, in \u001b[0;36mGlobalTracingContext.GlobalTracingExit.__exit__\u001b[0;34m(self, exc_type, exc_val, traceback)\u001b[0m\n\u001b[1;32m     96\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     98\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--> 100\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:156\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m    150\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    152\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m InterventionNode\u001b[38;5;241m.\u001b[39mprepare_inputs(\n\u001b[1;32m    153\u001b[0m     (args, kwargs), fake\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    154\u001b[0m )\n\u001b[0;32m--> 156\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtarget\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    157\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    158\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    159\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:2350\u001b[0m, in \u001b[0;36mFakeCopyMode.__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m   2348\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   2349\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39mDisableTorchFunctionSubclass():\n\u001b[0;32m-> 2350\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": [
-        "with tiny_model.trace(input):\n",
+        "# turn on scan and validate\n",
+        "with tiny_model.trace(input, scan=True, validate=True):\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",
+        "    # the error is happening here\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",
@@ -1043,13 +1220,42 @@
       "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."
+        "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. [Here](https://nnsight.net/notebooks/features/scan_validate/) is a more detailed example of scan and validate in action!\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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "45E98sIAVc5A"
+      },
+      "source": [
+        "We can also use the `.scan()` method to get the shape of a module without having to fully run the model. If scan  is enabled, our input is run though the model 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!"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 20,
-      "metadata": {},
+      "execution_count": 79,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 382
+        },
+        "id": "5J79HytJkkw5",
+        "outputId": "3b11405b-7bd7-45b6-dc16-11c8456b277b"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1067,6 +1273,17 @@
         "print(dim)"
       ]
     },
+    {
+      "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": "markdown",
       "metadata": {
@@ -1082,21 +1299,21 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 21,
+      "execution_count": 80,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "L-cYUNNyRDn1",
-        "outputId": "d03dafa1-ef1e-40f1-de28-4dc2f77fe4db"
+        "outputId": "c1521deb-34d7-4ea5-c80f-0c5c873252d9"
       },
       "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 1 output gradient: tensor([[-0.1296,  0.4562, -0.1182, -0.3536,  0.0703,  0.1411, -0.3201, -0.4890,\n",
+            "          0.0196, -0.0452]])\n",
             "Layer 2 output gradient: tensor([[1., 1.]])\n"
           ]
         }
@@ -1134,13 +1351,13 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 22,
+      "execution_count": 81,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "9bFoaJpOWlRb",
-        "outputId": "1c026534-9ffe-40e7-d016-e4565ddde8ac"
+        "outputId": "d3c78c79-4b0f-496f-f722-573f6db30e3a"
       },
       "outputs": [
         {
@@ -1175,29 +1392,39 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "SJ6dAgzQkkw6"
+      },
       "source": [
         "## Early Stopping"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "FZLrnIxAkkw6"
+      },
       "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. "
+        "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": {},
+      "execution_count": 82,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "MwAeiZhGkkw6",
+        "outputId": "85e6c10d-50dd-4517-b34e-2a9757e29c80"
+      },
       "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"
+            "L1 - Output:  tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -1212,7 +1439,9 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "eO_8oZPXkkw6"
+      },
       "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",
@@ -1223,8 +1452,15 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 24,
-      "metadata": {},
+      "execution_count": 83,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 347
+        },
+        "id": "uPJ6j2Zwkkw6",
+        "outputId": "aeff829d-e973-489e-c24e-1d5cdc80c7d0"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1240,10 +1476,10 @@
           "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",
+            "Cell \u001b[0;32mIn[83], 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~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:70\u001b[0m, in \u001b[0;36mProxy.__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     66\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     68\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---> 70\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     72\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[39mtarget\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;124m\"\u001b[39m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:64\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     57\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     58\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     59\u001b[0m \n\u001b[1;32m     60\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     61\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     62\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 64\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~/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py:143\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    133\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    134\u001b[0m \n\u001b[1;32m    135\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or was destroyed).\u001b[39;00m\n\u001b[1;32m    140\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    142\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--> 143\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    145\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."
           ]
         }
@@ -1258,24 +1494,34 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
-      "source": [
+      "metadata": {
+        "id": "i1iJOI29kkw6"
+      },
+      "source": [
         "## Conditional Interventions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "NvC2f8qqkkw6"
+      },
       "source": [
-        "Interventions can also be made conditional. \n",
+        "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": {},
+      "execution_count": 84,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "F_SIV2CDkkw6",
+        "outputId": "2ba0cfb8-dfd8-4403-e2e9-db3de2045ec7"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1299,22 +1545,23 @@
     },
     {
       "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": {},
+      "metadata": {
+        "id": "CZLQecmrkkw6"
+      },
       "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": {},
+      "execution_count": 21,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "qZsFtC5jkkw6",
+        "outputId": "4bbad29b-2d9d-4816-ed73-859a829e5ba5"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1328,12 +1575,200 @@
         "with tiny_model.trace(input) as tracer:\n",
         "\n",
         "  non_rand_int = 8\n",
-        "  \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": "hl99yP3kkkw6"
+      },
+      "source": [
+        "With `nnsight 0.4` we can now also use Python `if` statements within the tracing context to create a conditional context!\n",
+        "\n",
+        "*Note: Colab behaves a little strangely with this feature the first time you run it - expect some lagging and warnings*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 85,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "52a1vwtCfdIS",
+        "outputId": "51680d97-0f1f-4ead-ab97-26ec81587fff"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Random Integer  tensor([2])  is Even\n"
+          ]
+        }
+      ],
+      "source": [
+        "with tiny_model.trace(input) as tracer:\n",
+        "\n",
+        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
+        "\n",
+        "  # Since this if statement is inside the tracing context the if will\n",
+        "  # create a conditional context and will only execute the intervention\n",
+        "  # if this condition is met\n",
+        "  if rand_int % 2 == 0:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
+        "\n",
+        "  if rand_int % 2 == 1:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tEML88wTf9Lv"
+      },
+      "source": [
+        "`elif` statements should also work as `if` statements within the tracing context:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 86,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "3njmT1szf-Lq",
+        "outputId": "b1088864-8e0e-4526-da98-b6275adbabdf"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Random Integer  tensor([-3])  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",
+        "  # Since this if statement is inside the tracing context the if will\n",
+        "  # create a conditional context and will only execute the intervention\n",
+        "  # if this condition is met\n",
+        "  if rand_int % 2 == 0:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
+        "  elif rand_int % 2 == 1:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ONzsQcHtxKkf"
+      },
+      "source": [
+        "## Iterative Interventions"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NrwpkfYXxRBO"
+      },
+      "source": [
+        "With the iterator context, you can now run an intervention loop at scale. It iteratively executes and updates a single intervention graph. Use a `.session()` to define the Iterator context and pass in a sequence of items that you want to loop over at each iteration\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 87,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "AYWJVln7xP85",
+        "outputId": "3b7a1271-b199-4078-c70e-0fefc4cd5d7b"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        }
+      ],
+      "source": [
+        "with tiny_model.session() as session:\n",
+        "\n",
+        "  li = nnsight.list() # an NNsight built-in list object\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": {
+        "id": "5Dbnsw6fx8xL"
+      },
+      "source": [
+        "With `nnsight 0.4` we can now also use Python `for` loops within a tracer context at scale.\n",
+        "\n",
+        "*NOTE: inline for loops (i.e., `[x for x in <Proxy object>`]) are not currently supported.*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 88,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "14rgY7cZyNUr",
+        "outputId": "a028118d-cf04-45d3-ea38-71dd04405a9b"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        }
+      ],
+      "source": [
+        "# New: Using Python for loops for iterative interventions\n",
+        "with tiny_model.session() as session:\n",
+        "\n",
+        "    li = nnsight.list()\n",
+        "    [li.append([num]) for num in range(0, 3)]\n",
+        "    li2 = nnsight.list().save()\n",
+        "\n",
+        "    # Using regular for loops\n",
+        "    for item in li:\n",
+        "        for item_2 in item: # for loops can be nested!\n",
+        "            li2.append(item_2)\n",
+        "\n",
+        "print(\"\\nList: \", li2)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1354,8 +1789,10 @@
         "\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"
+        "complex and powerful types of models, and we're excited to see what can be done\n",
+        "to leverage its features!\n",
+        "\n",
+        "However, if you'd like to load a Language Model from HuggingFace with its tokenizer, the`LanguageModel` subclass greatly simplifies this process."
       ]
     },
     {
@@ -1383,60 +1820,71 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 27,
+      "execution_count": 6,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 769,
+          "height": 713,
           "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"
+            "26dcd7b75feb41a6b86d4260ecd64d06",
+            "39c9f1effe664e899125748e97fad198",
+            "383b893224684aee916e455c0d3d8c06",
+            "af24a12bd05b44148b5d6f91ac2445a0",
+            "d00750501f2548d0b1c69f4753aec877",
+            "98d28dddf75a45a2b1e05131b0b865a1",
+            "6e81d6f31f2446c999df5511e8afd8d0",
+            "68b205f60f9d4599a7d1f78583d0b93c",
+            "69e8c149f676457d95029f21e9a2956a",
+            "ba90058201ee48dcae9411ded457e85f",
+            "346ecc4734014cb4931bea53d5a32bf1",
+            "79a221c9228d459aa17402c9fef1ad1d",
+            "9c13a3d768eb47c8b26a7c6bfa8ec957",
+            "e97c052113f54dddbb619ed228f6b4ef",
+            "41fc21ee95c14d29943a56f2b63f4989",
+            "5250cf68037a43e389493733c5b9d4ac",
+            "e24033c5ba7b4e4fa4f4375d85b03bad",
+            "4b798a784ae748b1a0aef025ead976fc",
+            "19c0e121e52040b18393aee887730644",
+            "2702f30b5d3b4b0b8184fec58863750b",
+            "42fd0ea3be3c4b3f98ca7818efee3669",
+            "184b5c75820b4c42ae861669ac1a47e5",
+            "d8feaceb6d854cb2b2ff19a0309d9a7b",
+            "fd84ffef58604ce3a16ef8f9868072f3",
+            "64b713188f054535a3063c079bceead0",
+            "9cea52b7835b40628fc0a65199c5e82f",
+            "87a5c035bc6345aea6441b3db4b94321",
+            "7f4f63d94f8942fcb3e4b286d5848f0c",
+            "a310703500c5481aa37b698c659e1ca6",
+            "7afff81060f0481aa95a56daa570e1b8",
+            "a038f9037ed540f2ba43f3f85920b17e",
+            "ec4f48d5a9ef4afa8fa44078dfb236c6",
+            "405308ec76074c7ea319e1c1358baf39",
+            "7479281e090d4eeba4c21a9c15b02165",
+            "f4b25b3aac414499b653404ea3c066da",
+            "0383a2798a0c4fcfbd1f068817770fd6",
+            "385291e96814453b8f805cbe8a5e594d",
+            "f7a41e267f79401681bc4693c52e55fa",
+            "caff23247fb9432db8d5860a8132e1b3",
+            "f47609606aa14f6fa2407aa26edae8e3",
+            "cef2fe4eb51146afabe820140a3b0df3",
+            "52281c8594ad465397e6079385c558a1",
+            "e5302eb0481c4e529fc68353eed4b79f",
+            "0c46bce8cc3e4393898b025ba6b4e543",
+            "14008810744549208f147a132b419d20",
+            "cb4be83440e44e1c9e664ac1874b830d",
+            "f54c77317ca04861a958c8f982b25653",
+            "db24a855fb5846e28c98dcfe9afff259",
+            "647c861fcaa44148a90da061b1d5c99c",
+            "e4e0b1de76ff40a88b71e0b5a9f7bc39",
+            "504038c4efd3439daa958333f108b5fb",
+            "ded8c188917b424dbf264ce494fe0532",
+            "76fe13ab697c46cc9fcbe46112fb45b1",
+            "a1c25910a31a4b5a9c48f1a7fd8047c5",
+            "5d877979f2db4dbfa7581d2fe4381cf8"
           ]
         },
         "id": "1OD2z7d3HQJU",
-        "outputId": "7a53864c-6965-4c25-9668-4b0d56db9f6b"
+        "outputId": "3a8f135a-ac39-47fe-da78-6680d42572ed"
       },
       "outputs": [
         {
@@ -1469,17 +1917,11 @@
             "    (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",
+            "  (generator): Generator(\n",
+            "    (streamer): Streamer()\n",
+            "  )\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": [
@@ -1522,10 +1964,9 @@
         "\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",
+        "Let's now apply some of the features that we used on the small model to `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",
+        "simpler (i.e., you can send prompts to the model 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",
@@ -1534,22 +1975,40 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 28,
+      "execution_count": 90,
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/"
+          "base_uri": "https://localhost:8080/",
+          "height": 152,
+          "referenced_widgets": [
+            "2fedb8b529c34049932f07b577b515d4",
+            "bd1a11654c2d47a09610e4cc23b8bc32",
+            "26acae3d321743d2ad8efbf8727dbd05",
+            "f42e44122a2d496cb7b2ea85c2858a71",
+            "c1586aeef6814447ad46c1f8ddaae625",
+            "5e091b15cc9748469676b165533f0c44",
+            "f143977aeaf24798b7744681069762d4",
+            "a4b5c225f5ac4719ade8e6aafcff78d2",
+            "ee27d18a555c46e78c92b282bfa6dc9d",
+            "a3154c60744d4a1a87b1c86b89a142b9",
+            "de023d4526524f48a34883a287aa8a4f",
+            "a15895ad1cb541ca89fa9def63779104",
+            "ee60494714c2485f802298dce918edfe",
+            "2a0f2484053f41298aceacea5c1cc195",
+            "57aebe107749473390af2532806920ce",
+            "ff44e896c55048c5a7d92a56aa9eedfc",
+            "7fdb623e5b2d4eb1b6cee27da2314062",
+            "95832c5076ab4c80810d3c443e061156",
+            "08fb0620499c44eb9d160a0f21263451",
+            "287a8fe3ddf54d36b7603171ec7a2688",
+            "55c0bb35417e454d842e12bd2dfe62df",
+            "d5821d69e1fa4673bb39ab70a7f60466"
+          ]
         },
         "id": "mLSapaMCgLNU",
-        "outputId": "192b8998-ecca-4f25-ecf2-5196c0d2d672"
+        "outputId": "107769db-2d2e-4b52-ed59-31555ebcbfcb"
       },
       "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",
@@ -1583,13 +2042,12 @@
         "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",
+        "We just ran a little intervention on a much more complex model with many more\n",
+        "parameters! However, we're missing an important piece of information: what the\n",
+        "prediction would have looked 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"
+        "We could just run two tracing contexts and compare the outputs. However, this would require two forward passes through the model. `NNsight` can do\n",
+        "better than that with batching.\n"
       ]
     },
     {
@@ -1609,44 +2067,35 @@
         "id": "_Gc-Zr6NmEOn"
       },
       "source": [
-        "It's time to bring back the `Tracer` object we dropped before. \n",
+        "Batching is a way to process multiple inputs in one forward pass. To better understand how batching works, we're going to bring back the `Tracer` object that 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",
+        "When we call `.trace(...)`, it's actually creating two different contexts behind the scenes. The first one is the tracing context that we've discussed previously, and the second one is the invoker context. The invoker context defines the values of the `.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",
+        "If we call `.trace(...)` with some input, the input is passed on to the invoker. As 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",
+        "If we call `.trace()` without an input, then we can call `tracer.invoke(input1)` to manually create the invoker context with an input, `input1`. We can also repeatedly call `tracer.invoke(...)` to create the invoker context for additional inputs. Every subsequent time we call\n",
+        "`.invoke(...)`, interventions within its context will only refer to the input in that particular invoke statement.\n",
+        "\n",
+        "When exiting the tracing context, the inputs from all of the invokers will be batched together, and they will be executed in one forward pass! To test this out, let's do the same ablation experiment, but also add a 'control' output for comparison:\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",
+        "Note that when injecting data to only the relevant invoker interventions, `nnsight` tries, but can't guarantee, to narrow the data into the right\n",
+        "batch indices. Thus, 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",
+        "Just like `.trace(...)` created a `Tracer` object, `.invoke(...)` creates an `Invoker` object. For `LanguageModel` 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",
+        "In a case where we pass 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",
+        "Keyword arguments given to `.invoke(..)` make their way to the input pre-processing.  \n",
+        "`LanguageModel` has keyword arguments `max_length` and `truncation` used for tokenization which can be\n",
+        "passed to the invoker. If we want to pass keyword arguments to the invoker for a single-input `.trace(...)`, we can pass `invoker_args` as a dictionary of invoker keyword arguments.\n",
         "\n",
         "Here is an example to demonstrate everything we've described:\n",
         "\n",
@@ -1673,15 +2122,22 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 29,
+      "execution_count": 91,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "Kdcq4oCNmEua",
-        "outputId": "e6d33c07-d22a-41af-b84a-91c9681c7524"
+        "outputId": "33702cb9-46df-4c60-f108-f7ad3c1cd446"
       },
       "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",
@@ -1711,6 +2167,7 @@
         "        # Get the output for only the original batch.\n",
         "        token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
         "\n",
+        "\n",
         "print(\"Original token IDs:\", token_ids_original)\n",
         "print(\"Modified token IDs:\", token_ids_intervention)\n",
         "\n",
@@ -1724,30 +2181,32 @@
         "id": "8BsybgsK2Bbr"
       },
       "source": [
-        "So it did end up affecting what the model predicted. That's pretty neat!\n",
+        "Based on our control results, our ablation did end up affecting what the model predicted. That's pretty neat!\n",
+        "\n",
+        "Another cool thing with multiple invokes is that Proxies can interact between them.\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"
+        "Here, we transfer the token embeddings from a real prompt into another placeholder prompt. Therefore the latter prompt produces the output of the former prompt:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 30,
+      "execution_count": 92,
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/"
+          "base_uri": "https://localhost:8080/",
+          "height": 519
         },
         "id": "syKTI_KhpvCY",
-        "outputId": "b2238c5c-d2d8-47bf-b73a-730f1e03cf99"
+        "outputId": "35afac9a-a34a-4e60-edc8-4c1b404e4389"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
+            "original prediction shape torch.Size([])\n",
             "Original prediction:  _\n",
+            "modified prediction shape torch.Size([])\n",
             "Modified prediction:  Paris\n"
           ]
         }
@@ -1756,30 +2215,38 @@
         "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",
+        "      token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
         "\n",
-        "        token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
-        "\n",
+        "print(\"original prediction shape\", token_ids_original[0][-1].shape)\n",
         "print(\"Original prediction:\", llm.tokenizer.decode(token_ids_original[0][-1]))\n",
+        "\n",
+        "print(\"modified prediction shape\", token_ids_intervention[0][-1].shape)\n",
         "print(\"Modified prediction:\", llm.tokenizer.decode(token_ids_intervention[0][-1]))"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_bXvUGGqbwv1"
+      },
+      "source": [
+        "## Multiple Token Generation"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "vvWA-CWqQtah"
       },
       "source": [
-        "## .next()\n",
+        "### .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",
@@ -1790,114 +2257,296 @@
         "\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",
+        "`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",
+        "Each module has a call index 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"
+        "graph only at the iteration that matches the call index.\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 31,
+      "execution_count": 93,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "Yy5Z9NE1GkaN",
-        "outputId": "0b4c1bd3-a6f6-4ba1-9cb8-6019c222311f"
+        "outputId": "3827231c-86b1-482e-bd2f-ff17191692e5"
       },
       "outputs": [
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
+            "Setting `pad_token_id` to `eos_token_id`:None 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"
+            "torch.Size([1, 10, 768])\n",
+            "torch.Size([1, 1, 768])\n",
+            "tensor([[ 464,  412,  733,  417, 8765,  318,  287,  262, 1748,  286, 6342,   11,\n",
+            "          290]], device='mps:0')\n"
           ]
         }
       ],
       "source": [
-        "with llm.generate(\"The Eiffel Tower is in the city of\", max_new_tokens=3):\n",
+        "with llm.generate('The Eiffel Tower is in the city of', max_new_tokens=3) as tracer:\n",
         "\n",
-        "    token_ids_1 = llm.lm_head.output.argmax(dim=-1).save()\n",
+        "    hidden_states1 = llm.transformer.h[-1].output[0].save()\n",
         "\n",
-        "    token_ids_2 = llm.lm_head.next().output.argmax(dim=-1).save()\n",
+        "    # use module.next() to access the next intervention\n",
+        "    hidden_states2 = llm.transformer.h[-1].next().output[0].save()\n",
         "\n",
-        "    token_ids_3 = llm.lm_head.next().output.argmax(dim=-1).save()\n",
+        "    # saving the output allows you to save the hidden state across the initial prompt\n",
+        "    out = llm.generator.output.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"
+        "print(hidden_states1.shape)\n",
+        "print(hidden_states2.shape)\n",
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "oHy6FqyNRBzm"
+      },
       "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",
+        "### New! using .all()\n",
         "\n",
-        "Let's explore how to use the **Editor** context to make a simple persistent change to a model:"
+        "With `nnsight 0.4` you can now use `.all()` to recursively apply interventions to a model. Calling `.all()` on a module within a model will recursively apply its `.input` and `.output` across all iterations. Previously, we'd need to loop across each new generated token, saving the intervention for every generated token and calling `.next()` to move forward."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 32,
-      "metadata": {},
+      "execution_count": 94,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "C1C1KmclTJDH",
+        "outputId": "97f4b5ba-3b94-4541-df98-4025b273ba97"
+      },
       "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
+          ]
+        },
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "\n",
-            "Original Prediction:   Rome\n",
-            "Modified Prediction:   Paris\n"
+            "Hidden state length:  3\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",
+        "# Old approach:\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = llm.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "hidden_states = []\n",
+        "with llm.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    for i in range(n_new_tokens):\n",
+        "        # Apply intervention - set first layer output to zero\n",
+        "        layers[0].output[0][:] = 0\n",
+        "\n",
+        "        # Append desired hidden state post-intervention\n",
+        "        hidden_states.append(layers[-1].output.save())\n",
+        "\n",
+        "        # Move to next generated token\n",
+        "        layers[0].next()\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tuBLYSx_XKXK"
+      },
+      "source": [
+        "We can use also `.all()` to streamline the multiple token generation process. We simply call `.all` on the module where we are applying the intervention (in this case GPT-2's layers), apply our intervention, and append our hidden states (stored in an `nnsight.list()` object).\n",
+        "<br> <br>\n",
+        "\n",
+        "Let's test this out for the multiple token generation case:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 95,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "IgxXlvxDX1zF",
+        "outputId": "9588942b-19b8-4d8b-90a3-9b8ef463b53f"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Hidden state length:  3\n"
+          ]
+        }
+      ],
+      "source": [
+        "# using .all():\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = llm.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with llm.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() to apply intervention to each new token\n",
+        "    layers.all()\n",
+        "\n",
+        "    # Apply intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(layers[-1].output) # no need to call .save\n",
+        "    # Don't need to loop or call .next()!\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Z3dMEb6HX-Wp"
+      },
+      "source": [
+        "Easy! Note that because `.all()` is recursive, it will only work to append outputs called on children of the module that `.all()` was called on. See example below for more information. TL;DR: apply `.all()` on the highest-level accessed module if interventions and outputs have different hierarchies within model structure."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3J_IeO96jVkZ"
+      },
+      "source": [
+        "<details>\n",
+        "<summary>Recursive properties of .all()</summary>\n",
+        "\n",
+        "`.all()` recursively acts on model components. In the below code example, only the first token generation is saved, because `.all()` applied to `layers`, while the saved variable `hidden_states` is produced from `model.lm_head`, which is not a child of `layers`.\n",
+        "\n",
+        "```\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() on layers\n",
+        "    layers.all()\n",
+        "\n",
+        "    # Apply same intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(model.lm_head.output) # no need to call .save, it's already initialized\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states)) # length is 1, meaning it only saved the first token generation\n",
+        "```\n",
+        "\n",
+        "If you want to apply an intervention during multiple token generation while saving the state of a model component that isn't a child of that module, you can instead apply `.all()` to the full model:\n",
+        "\n",
+        "```\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() on model\n",
+        "    model.all()\n",
+        "\n",
+        "    # Apply same intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(model.lm_head.output) # no need to call .save\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states)) # length is 3, as expected!\n",
+        "```\n",
+        "\n",
+        "</details>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ty-luCwCkkw7"
+      },
+      "source": [
+        "## Model Editing"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hgrZVQLEkkw8"
+      },
+      "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": 96,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "-_oAprpwkkw8",
+        "outputId": "388b31e3-9585-4898-98be-7a1a043fa49c"
+      },
+      "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",
+        "\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",
@@ -1909,17 +2558,25 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "S8J6cVw0kkw8"
+      },
       "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",
+        "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": {},
+      "execution_count": 97,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "b1atxBp8kkw8",
+        "outputId": "d176ce44-a861-4407-fbf3-5c4411bd4343"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1943,15 +2600,23 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "nVeZO96jkkw8"
+      },
       "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. "
+        "If you've made in-place edits to your model and need to revert these changes, you can apply `.clear_edits()`. This method removes all edits applied to the model, effectively restoring it to its original state."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 34,
-      "metadata": {},
+      "execution_count": 98,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "_SHK4JzPkkw8",
+        "outputId": "c6da4d26-e5b8-415e-aa01-071fbdbd4802"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1985,10 +2650,9 @@
         "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",
+        "`NNsight` is only one part of our project to democratize access to AI internals. The other half is the National Deep Inference Fabric, or `NDIF`. `NDIF` hosts large models for shared access using `NNsight`, so you don't have to worry about any of the headaches of hosting large models yourself!\n",
         "\n",
-        "The interaction between the two is fairly straightforward. The\n",
+        "The interaction between `NDIF` and `NNsight` 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",
@@ -2014,15 +2678,15 @@
       },
       "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",
+        "run the rest of this walkthrough, you need one of your own. To get one, simply\n",
+        "register at https://login.ndif.us.\n",
         "\n",
-        "With a valid API key, you then can configure `nnsight` by doing the following:\n"
+        "With a valid API key, you then can configure `nnsight` as follows:\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 35,
+      "execution_count": 100,
       "metadata": {
         "id": "ax1NWoS9MJeZ"
       },
@@ -2030,7 +2694,7 @@
       "source": [
         "from nnsight import CONFIG\n",
         "\n",
-        "CONFIG.set_default_api_key(\"YOUR_API_KEY\") \n"
+        "CONFIG.set_default_api_key(\"YOUR_API_KEY\")"
       ]
     },
     {
@@ -2039,16 +2703,16 @@
         "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",
+        "If you're running in a local IDE, this only needs to be run once as it will save the API key as the default in a\n",
+        ".config file along with your `nnsight` installation. You can also add your API key to Google Colab secrets.\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"
+        "context with `Llama-3.1-8b`!\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 36,
+      "execution_count": 101,
       "metadata": {
         "id": "Q1N04sJPZnJt"
       },
@@ -2056,51 +2720,56 @@
       "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"
+        "# Llama 3.1 70b is a gated model, so you need to apply for access on HuggingFace and include your token.\n",
+        "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\""
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 37,
+      "execution_count": 8,
       "metadata": {
-        "id": "oqVzjNoyNGc7"
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "oqVzjNoyNGc7",
+        "outputId": "abf3f296-c41b-475a-946b-6e0b046333e8"
       },
       "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",
+            "2025-01-31 15:47:57,111 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:47:57,317 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:47:57,519 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:47:58,468 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - COMPLETED: Your job has been completed.\n",
+            "Downloading result:   0%|          | 0.00/4.37M [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]"
+            "Downloading result: 100%|██████████| 4.37M/4.37M [00:00<00:00, 16.5MB/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",
+            "(tensor([[[ 1.7734,  2.6875,  0.8047,  ..., -1.8594,  2.2344,  3.2656],\n",
+            "         [ 0.0312, -0.0352, -2.8750,  ..., -0.8906, -0.0547,  1.6172],\n",
+            "         [ 1.3594, -2.0156,  1.7031,  ..., -1.7031, -0.7422,  1.4375],\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",
+            "         [ 1.0000,  0.3203, -0.2656,  ..., -0.0723, -0.2559,  0.2090],\n",
+            "         [ 0.4707, -0.3496,  0.2422,  ...,  0.7344, -0.0078,  0.1133],\n",
+            "         [-0.0566, -0.3496,  0.4746,  ...,  0.9844,  0.6797, -0.8750]]],\n",
+            "       dtype=torch.bfloat16),)\n",
+            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
+            "         [-2.4375, -1.7266, -2.0156,  ..., -9.1250, -9.1250, -9.1250],\n",
+            "         [ 9.6875,  4.5625,  5.8750,  ..., -3.3906, -3.3906, -3.3906],\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"
+            "         [ 2.3281,  1.0703, -0.3203,  ..., -7.1562, -7.1562, -7.1562],\n",
+            "         [11.1875,  6.0312,  4.9062,  ..., -3.5156, -3.5156, -3.5156],\n",
+            "         [ 8.0000,  5.2500,  4.3750,  ..., -3.9844, -3.9844, -3.9844]]],\n",
+            "       dtype=torch.bfloat16)\n"
           ]
         },
         {
@@ -2112,9 +2781,11 @@
         }
       ],
       "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",
+        "from nnsight import LanguageModel\n",
         "\n",
+        "# We'll never actually load the parameters locally, so no need to specify a device_map.\n",
+        "\n",
+        "llama = LanguageModel(\"meta-llama/Meta-Llama-3.1-8B\")\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",
@@ -2139,39 +2810,48 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "lfRDZox5kkw8"
+      },
       "source": [
         "## Sessions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "33XLOvtukkw8"
+      },
       "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",
+        "This offers the following benefits:\n",
+        "1.   All interventions within a session will be executed one after another without additional wait in the NDIF queue\n",
+        "2.   All intermediate outputs for 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": {},
+      "execution_count": 9,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "KA6GROiTR3gf",
+        "outputId": "dfefe423-9caf-48ce-cfe5-83c2f808a6b2"
+      },
       "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]"
+            "2025-01-31 15:48:17,378 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:17,558 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:17,755 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:18,205 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 1.62k/1.62k [00:00<00:00, 6.57MB/s]"
           ]
         },
         {
@@ -2180,7 +2860,7 @@
           "text": [
             "\n",
             "T1 - Original Prediction:   Paris\n",
-            "T2 - Modified Prediction:   Paris\n"
+            "T2 - Modified Prediction:  ://\n"
           ]
         },
         {
@@ -2195,12 +2875,12 @@
         "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",
+        "    # capture the hidden state from layer 32 at the last token\n",
+        "    hs_31 = llama.model.layers[31].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",
+        "    llama.model.layers[1].output[0][:, -1, :] = hs_31[:]\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",
@@ -2209,7 +2889,9 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "End2DR7qkkw8"
+      },
       "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",
@@ -2218,122 +2900,209 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "M-Nq8q0Qkkw8"
+      },
       "source": [
         "The `session` context object has its own methods to log values and be terminated early."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 39,
-      "metadata": {},
+      "execution_count": 104,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 311
+        },
+        "id": "Q3c6M5Jwkkw8",
+        "outputId": "0d937244-d022-4745-f3ea-f10b305bc3e3"
+      },
       "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"
+            "2025-01-31 15:17:51,636 3d5144b4-e0b6-4106-912f-5abec8e78e7f - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:17:51,804 3d5144b4-e0b6-4106-912f-5abec8e78e7f - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:17:52,023 3d5144b4-e0b6-4106-912f-5abec8e78e7f - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:17:52,028 3d5144b4-e0b6-4106-912f-5abec8e78e7f - LOG: -- Early Stop --\n",
+            "2025-01-31 15:17:52,364 3d5144b4-e0b6-4106-912f-5abec8e78e7f - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 7.25MB/s]\n"
           ]
         }
       ],
       "source": [
         "with llama.session(remote=True) as session:\n",
-        "\n",
         "  session.log(\"-- Early Stop --\")\n",
-        "  session.exit()"
+        "  nnsight.stop"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "5Z-tTpQekkw8"
+      },
       "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! "
+        "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 could swap activations between base and instruct versions of the Llama model and compare their outputs. And `session` can also be used to run similar experiments entirely locally!"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "H6hxSZ8zTFvT"
+      },
       "source": [
-        "## Looping"
+        "## Streaming"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "4bDUYaluTRCM"
+      },
       "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",
+        "Streaming enables users apply functions and datasets locally during remote model execution. This allows users to stream results for immediate consumption (i.e., seeing tokens as they are generated) or applying non-whitelisted functions such as model tokenizers, large local datasets, and more!\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:"
+        "*   `nnsight.local()` context sends values immediately to user's local machine from server\n",
+        "*   Intervention graph is executed locally on downstream nodes\n",
+        "*   Exiting local context uploads data back to server\n",
+        "*   `@nnsight.trace` function decorator enables custom functions to be added to intervention graph when using `nnsight.local()`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fBDttzi7TT26"
+      },
+      "source": [
+        "### `nnsight.local()`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NRsfBs5MTW6f"
+      },
+      "source": [
+        "You may sometimes want to locally access and manipulate values during remote execution. Using `.local()` on a proxy, you can send remote content to your local machine and apply local functions. The intervention graph is then executed locally on downstream nodes (until you send execution back to the remote server by exiting the `.local()` context).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gN8yWbc4TXZc"
+      },
+      "source": [
+        "There are a few use cases for streaming with `.local()`, including live chat generation and applying large datasets or non-whitelisted local functions to the intervention graph."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fFxwkDC6TZ8C"
+      },
+      "source": [
+        "Now let's explore how streaming works. We'll start by grabbing some hidden states of the model and printing their value using `tracer.log()`. Without calling `nnsight.local()`, these operations will all occur remotely."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 40,
-      "metadata": {},
+      "execution_count": 120,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "QKMwwmoUTcCI",
+        "outputId": "f1400168-1f44-43c4-91c1-14deaeedb4d0"
+      },
       "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"
+            "2025-01-31 15:21:29,379 94c37060-d3f6-4f21-b55f-0e72f03d08da - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:29,592 94c37060-d3f6-4f21-b55f-0e72f03d08da - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:29,786 94c37060-d3f6-4f21-b55f-0e72f03d08da - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:21:29,822 94c37060-d3f6-4f21-b55f-0e72f03d08da - LOG: tensor(1.7656, device='cuda:0')\n",
+            "2025-01-31 15:21:30,161 94c37060-d3f6-4f21-b55f-0e72f03d08da - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 514k/514k [00:00<00:00, 4.80MB/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
+            "         [10.2500,  2.1094,  2.8281,  ..., -8.2500, -8.2500, -8.2500]]],\n",
+            "       dtype=torch.bfloat16)\n"
           ]
         }
       ],
       "source": [
-        "with llama.session(remote=True) as session:\n",
+        "# This will give you a remote LOG response because it's coming from the remote server\n",
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
         "\n",
-        "  with session.iter([0, 1, 2]) as item:\n",
-        "    # define intervention body here ...\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
         "\n",
-        "    with llama.trace(\"_\"):\n",
-        "      # define interventions here ...\n",
-        "      pass\n",
+        "    tracer.log(hs[0,0,0])\n",
         "\n",
-        "    with llama.trace(\"_\"):\n",
-        "      # define interventions here ...\n",
-        "      pass"
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "1IjPUn8oTouV"
+      },
       "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:"
+        "Now, let's try the same operation using the `nnsight.local()` context. This will send the operations to get and print the hidden states to your local machine, changing how the logging message is formatted (local formatting instead of remote)."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 41,
-      "metadata": {},
+      "execution_count": 121,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "x9byLVl_Tqy-",
+        "outputId": "e1514fdd-61b5-4cbe-c577-1bdab9605adc"
+      },
       "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]"
+            "2025-01-31 15:21:32,674 0414b4df-a160-485b-89dd-4e70a05a43e2 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:32,830 0414b4df-a160-485b-89dd-4e70a05a43e2 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:33,092 0414b4df-a160-485b-89dd-4e70a05a43e2 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:21:33,352 0414b4df-a160-485b-89dd-4e70a05a43e2 - COMPLETED: Your job has been completed.\n"
           ]
         },
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "\n",
-            "List:  [0, 1, 2]\n"
+            "tensor(1.7656, dtype=torch.bfloat16)\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Downloading result: 100%|██████████| 514k/514k [00:00<00:00, 4.28MB/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
+            "         [10.2500,  2.1094,  2.8281,  ..., -8.2500, -8.2500, -8.2500]]],\n",
+            "       dtype=torch.bfloat16)\n"
           ]
         },
         {
@@ -2345,506 +3114,2213 @@
         }
       ],
       "source": [
+        "# This will print locally because it's already local\n",
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
         "\n",
-        "with llama.session(remote=True) as session:\n",
+        "    with nnsight.local():\n",
+        "        hs = llama.model.layers[-1].output[0]\n",
+        "        tracer.log(hs[0,0,0])\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",
+        "    out =  llama.lm_head.output.save()\n",
         "\n",
-        "print(\"\\nList: \", li2)"
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "tNbmIZSCTwo2"
+      },
       "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>"
+        "### `@nnsight.trace` function decorator"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "xiQOxe9_TzqI"
+      },
       "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:"
+        "We can also use function decorators to create custom functions to be used during `.local` calls. This is a handy way to enable live streaming of a chat or to train probing classifiers on model hidden states."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "X7A3IjhLT2TE"
+      },
+      "source": [
+        "Let's try out `@nnsight.trace` and `nnsight.local()` to access a custom function during remote execution."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 42,
-      "metadata": {},
+      "execution_count": 122,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "P8ghjV-7TzPj",
+        "outputId": "6e5a2efb-8edd-40ff-a4de-84c535bbbd3d"
+      },
       "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"
+            "2025-01-31 15:21:35,853 3eeee112-3074-4b3e-adea-001a85ce554d - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:36,035 3eeee112-3074-4b3e-adea-001a85ce554d - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:36,258 3eeee112-3074-4b3e-adea-001a85ce554d - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:21:36,653 3eeee112-3074-4b3e-adea-001a85ce554d - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 258k/258k [00:00<00:00, 2.91MB/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()"
+        "# first, let's define our function\n",
+        "@nnsight.trace # decorator that enables this function to be added to the intervention graph\n",
+        "def my_local_fn(value):\n",
+        "    return value * 0\n",
+        "\n",
+        "# We use a local function to ablate some hidden states\n",
+        "# This downloads the data for the .local context, and then uploads it back to set the value.\n",
+        "with llama.generate(\"hello\", remote=True) as tracer:\n",
+        "\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
+        "\n",
+        "    with nnsight.local():\n",
+        "\n",
+        "        hs = my_local_fn(hs)\n",
+        "\n",
+        "    llama.model.layers[-1].output[0][:] = hs\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "_7d5AdU6T72j"
+      },
       "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",
+        "Note that without calling `.local`, the remote API does not know about `my_local_fn` and will throw a whitelist error. A whitelist error occurs because you are being allowed access to the function."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 123,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 472
+        },
+        "id": "n4OIIqQhT98b",
+        "outputId": "54e6ad3d-a882-4892-aac0-7bdba52cac40"
+      },
+      "outputs": [
+        {
+          "ename": "FunctionWhitelistError",
+          "evalue": "Function with name `__main__.my_local_fn` not in function whitelist.",
+          "output_type": "error",
+          "traceback": [
+            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+            "\u001b[0;31mFunctionWhitelistError\u001b[0m                    Traceback (most recent call last)",
+            "Cell \u001b[0;32mIn[123], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m llama\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhello\u001b[39m\u001b[38;5;124m\"\u001b[39m, remote\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m tracer:\n\u001b[1;32m      3\u001b[0m     hs \u001b[38;5;241m=\u001b[39m llama\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39mlayers[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m.\u001b[39moutput[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m      5\u001b[0m     hs \u001b[38;5;241m=\u001b[39m my_local_fn(hs) \u001b[38;5;66;03m# no .local - will cause an error\u001b[39;00m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/contexts/interleaving.py:96\u001b[0m, in \u001b[0;36mInterleavingTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     92\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     94\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m---> 96\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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/tracer.py:25\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mglobals\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m GlobalTracingContext\n\u001b[1;32m     23\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mtry_deregister(\u001b[38;5;28mself\u001b[39m)\n\u001b[0;32m---> 25\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/base.py:82\u001b[0m, in \u001b[0;36mContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     78\u001b[0m graph \u001b[38;5;241m=\u001b[39m graph\u001b[38;5;241m.\u001b[39mstack\u001b[38;5;241m.\u001b[39mpop()\n\u001b[1;32m     80\u001b[0m graph\u001b[38;5;241m.\u001b[39malive \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m---> 82\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[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:77\u001b[0m, in \u001b[0;36mRemoteBackend.__call__\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m     72\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, graph: Graph):\n\u001b[1;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mblocking:\n\u001b[1;32m     75\u001b[0m \n\u001b[1;32m     76\u001b[0m         \u001b[38;5;66;03m# Do blocking request.\u001b[39;00m\n\u001b[0;32m---> 77\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mblocking_request\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     80\u001b[0m \n\u001b[1;32m     81\u001b[0m         \u001b[38;5;66;03m# Otherwise we are getting the status / result of the existing job.\u001b[39;00m\n\u001b[1;32m     82\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnon_blocking_request(graph)\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:289\u001b[0m, in \u001b[0;36mRemoteBackend.blocking_request\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m    280\u001b[0m sio\u001b[38;5;241m.\u001b[39mconnect(\n\u001b[1;32m    281\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mws_address,\n\u001b[1;32m    282\u001b[0m     socketio_path\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/ws/socket.io\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    283\u001b[0m     transports\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwebsocket\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m    284\u001b[0m     wait_timeout\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m,\n\u001b[1;32m    285\u001b[0m )\n\u001b[1;32m    287\u001b[0m remote_graph \u001b[38;5;241m=\u001b[39m preprocess(graph)\n\u001b[0;32m--> 289\u001b[0m data, headers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrequest\u001b[49m\u001b[43m(\u001b[49m\u001b[43mremote_graph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    291\u001b[0m headers[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msession_id\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m sio\u001b[38;5;241m.\u001b[39msid\n\u001b[1;32m    293\u001b[0m \u001b[38;5;66;03m# Submit request via\u001b[39;00m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:60\u001b[0m, in \u001b[0;36mRemoteBackend.request\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m     58\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrequest\u001b[39m(\u001b[38;5;28mself\u001b[39m, graph: Graph) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[\u001b[38;5;28mbytes\u001b[39m, Dict[\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mstr\u001b[39m]]:\n\u001b[0;32m---> 60\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[43mRequestModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mserialize\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mformat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mzlib\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     62\u001b[0m     headers \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m     63\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_key\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel_key,\n\u001b[1;32m     64\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mformat\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mformat,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     67\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msent-timestamp\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mstr\u001b[39m(time\u001b[38;5;241m.\u001b[39mtime()),\n\u001b[1;32m     68\u001b[0m     }\n\u001b[1;32m     70\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m data, headers\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/request.py:43\u001b[0m, in \u001b[0;36mRequestModel.serialize\u001b[0;34m(graph, format, _zlib)\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mserialize\u001b[39m(graph: Graph, \u001b[38;5;28mformat\u001b[39m:\u001b[38;5;28mstr\u001b[39m, _zlib:\u001b[38;5;28mbool\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mbytes\u001b[39m:\n\u001b[1;32m     41\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mformat\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mjson\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 43\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[43mRequestModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     45\u001b[0m         json \u001b[38;5;241m=\u001b[39m data\u001b[38;5;241m.\u001b[39mmodel_dump(mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mjson\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     47\u001b[0m         data \u001b[38;5;241m=\u001b[39m msgspec\u001b[38;5;241m.\u001b[39mjson\u001b[38;5;241m.\u001b[39mencode(json)\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/request.py:30\u001b[0m, in \u001b[0;36mRequestModel.__init__\u001b[0;34m(self, memo, *args, **kwargs)\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, memo: Dict \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m---> 30\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__init__\u001b[39;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[43mmemo\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmemo\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\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     32\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m memo \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     34\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmemo \u001b[38;5;241m=\u001b[39m {\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mMEMO}\n",
+            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:276\u001b[0m, in \u001b[0;36mGraphModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    273\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m    274\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value: Graph) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Self:\n\u001b[0;32m--> 276\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mGraphModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnodes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnodes\u001b[49m\u001b[43m)\u001b[49m\n",
+            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:77\u001b[0m, in \u001b[0;36mmemoized.<locals>.inner\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minner\u001b[39m(value):\n\u001b[0;32m---> 77\u001b[0m     model \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m     _id \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mid\u001b[39m(value)\n\u001b[1;32m     81\u001b[0m     MEMO[_id] \u001b[38;5;241m=\u001b[39m model\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:101\u001b[0m, in \u001b[0;36mNodeModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m     98\u001b[0m \u001b[38;5;129m@memoized\u001b[39m\n\u001b[1;32m     99\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value: Node) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Self:\n\u001b[0;32m--> 101\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mNodeModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtarget\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:244\u001b[0m, in \u001b[0;36mFunctionModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    239\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m    240\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value:FUNCTION):\n\u001b[1;32m    242\u001b[0m     model \u001b[38;5;241m=\u001b[39m FunctionModel(function_name\u001b[38;5;241m=\u001b[39mget_function_name(value))\n\u001b[0;32m--> 244\u001b[0m     \u001b[43mFunctionModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheck_function_whitelist\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunction_name\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    246\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m model\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:251\u001b[0m, in \u001b[0;36mFunctionModel.check_function_whitelist\u001b[0;34m(cls, qualname)\u001b[0m\n\u001b[1;32m    248\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m    249\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcheck_function_whitelist\u001b[39m(\u001b[38;5;28mcls\u001b[39m, qualname: \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mstr\u001b[39m:\n\u001b[1;32m    250\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m qualname \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m FUNCTIONS_WHITELIST:\n\u001b[0;32m--> 251\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m FunctionWhitelistError(\n\u001b[1;32m    252\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFunction with name `\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mqualname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m` not in function whitelist.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    253\u001b[0m         )\n\u001b[1;32m    255\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m qualname\n",
+            "\u001b[0;31mFunctionWhitelistError\u001b[0m: Function with name `__main__.my_local_fn` not in function whitelist."
+          ]
+        }
+      ],
+      "source": [
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
         "\n",
-        "But in what kind of experimental scenario would someone even need to use it?\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
         "\n",
-        "In the next section, we delve into a powerful use case of the `Iterator` context and see how it enables it!"
+        "    hs = my_local_fn(hs) # no .local - will cause an error\n",
+        "\n",
+        "    llama.model.layers[-1].output[0][:] = hs * 2\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "v_bdhze-UBec"
+      },
       "source": [
-        "## Training a LoRA"
+        "### Example: Live-streaming remote chat\n",
+        "\n"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "hguvYgvVUE4b"
+      },
       "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."
+        "Now that we can access data within the tracing context on our local computer, we can apply non-whitelisted functions, such as the model's tokenizer, within our tracing context.\n",
+        "\n",
+        "Let's build a decoding function that will decode tokens into words and print the result."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 43,
-      "metadata": {},
+      "execution_count": 124,
+      "metadata": {
+        "id": "8OAXweC5UHBJ"
+      },
       "outputs": [],
+      "source": [
+        "@nnsight.trace\n",
+        "def my_decoding_function(tokens, model, max_length=80, state=None):\n",
+        "    # Initialize state if not provided\n",
+        "    if state is None:\n",
+        "        state = {'current_line': '', 'current_line_length': 0}\n",
+        "\n",
+        "    token = tokens[-1] # only use last token\n",
+        "\n",
+        "    # Decode the token\n",
+        "    decoded_token = llama.tokenizer.decode(token).encode(\"unicode_escape\").decode()\n",
+        "\n",
+        "    if decoded_token == '\\\\n':  # Handle explicit newline tokens\n",
+        "        # Print the current line and reset state\n",
+        "        print('',flush=True)\n",
+        "        state['current_line'] = ''\n",
+        "        state['current_line_length'] = 0\n",
+        "    else:\n",
+        "        # Check if adding the token would exceed the max length\n",
+        "        if state['current_line_length'] + len(decoded_token) > max_length:\n",
+        "            print('',flush=True)\n",
+        "            state['current_line'] = decoded_token  # Start a new line with the current token\n",
+        "            state['current_line_length'] = len(decoded_token)\n",
+        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
+        "        else:\n",
+        "            # Add a space if the line isn't empty and append the token\n",
+        "            if state['current_line']:\n",
+        "                state['current_line'] += decoded_token\n",
+        "            else:\n",
+        "                state['current_line'] = decoded_token\n",
+        "            state['current_line_length'] += len(decoded_token)\n",
+        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
+        "\n",
+        "    return state"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Uf-Ir1YgUKJt"
+      },
+      "source": [
+        "Now we can decode and print our model outputs throughout token generation by accessing our decoding function through `nnsight.local()`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 125,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "aT56JuDzULyJ",
+        "outputId": "bfa7593d-01c2-46e0-8c67-d8dbc1c7aa93"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Prompt:  A press release is an official statement delivered to members of the news media for the purpose of \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:21:48,944 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:49,121 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:49,363 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - RUNNING: Your job has started running.\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            " providing information, an official statement, or making an announcement.A press release is also written or recorded communication directed at members of the news media for the purpose of announcing something ostensibly newsworthy. Typically, they are mailed, faxed, or e"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:21:52,277 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 258k/258k [00:00<00:00, 3.03MB/s]\n"
+          ]
+        }
+      ],
       "source": [
         "import torch\n",
         "\n",
-        "from nnsight.envoy import Envoy # \n",
+        "nnsight.CONFIG.APP.REMOTE_LOGGING = False\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",
+        "prompt = \"A press release is an official statement delivered to members of the news media for the purpose of\"\n",
+        "# prompt = \"Your favorite board game is\"\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",
+        "print(\"Prompt: \",prompt,'\\n', end =\"\")\n",
         "\n",
-        "    # The Call method defines how to actually apply the LORA.\n",
-        "    def __call__(self, alpha: float = 1.0):\n",
-        "        \"\"\"Call.\n",
+        "# Initialize the state for decoding\n",
+        "state = {'current_line': '', 'current_line_length': 0}\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",
+        "with llama.generate(prompt, remote=True, max_new_tokens = 50) as generator:\n",
+        "    # Call .all() to apply to each new token\n",
+        "    llama.all()\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",
+        "    all_tokens = nnsight.list().save()\n",
         "\n",
-        "        # LORA is additive\n",
-        "        h = BA_x + self.module.output\n",
+        "    # Access model output\n",
+        "    out = llama.lm_head.output.save()\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",
+        "    # Apply softmax to obtain probabilities and save the result\n",
+        "    probs = torch.nn.functional.softmax(out, dim=-1)\n",
+        "    max_probs = torch.max(probs, dim=-1)\n",
+        "    tokens = max_probs.indices.cpu().tolist()\n",
+        "    all_tokens.append(tokens[0]).save()\n",
         "\n",
-        "    def parameters(self):\n",
-        "        # Some way to get all the parameters.\n",
-        "        return [self.WA, self.WB]"
+        "    with nnsight.local():\n",
+        "        state = my_decoding_function(tokens[0], llama, max_length=20, state=state)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "Let's define all the variables to use in LORA training. "
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 44,
-      "metadata": {},
-      "outputs": [],
+      "metadata": {
+        "id": "DUCWX1BDkkw9"
+      },
       "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"
+        "## Looping across sessions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "mWifpUjfkkw9"
+      },
       "source": [
-        "We can use the `.scan()` method to get the shape of the module without having to fully run the model."
+        "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, 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": 46,
-      "metadata": {},
+      "execution_count": 126,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "to-hAyuKkkw9",
+        "outputId": "ab21cb42-adba-4630-de5e-639cb6af3a35"
+      },
       "outputs": [
         {
-          "name": "stdout",
+          "name": "stderr",
           "output_type": "stream",
           "text": [
-            "8192\n"
+            "2025-01-31 15:22:03,714 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:22:03,869 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:22:04,120 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:22:05,086 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 1.82MB/s]\n"
           ]
         }
       ],
       "source": [
-        "with llama.scan(\" \"):\n",
-        "    dim = module.output.shape[-1]\n",
+        "with llama.session(remote=True) as session:\n",
         "\n",
-        "print(dim)"
+        "  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": {},
+      "metadata": {
+        "id": "L14XEU8jkkw9"
+      },
       "source": [
-        "It's time to run the LORA training loop! We using the **Session** and the **Iterator** contexts to achieve this."
+        "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": 47,
-      "metadata": {},
+      "execution_count": 127,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Rd-PcFJakkw9",
+        "outputId": "a5b20d0c-5f0d-4065-a874-68bfe9ec696e"
+      },
       "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"
+            "2025-01-31 15:22:07,235 570729dd-1d6d-4cd6-9876-9a096378119c - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:22:07,440 570729dd-1d6d-4cd6-9876-9a096378119c - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:22:07,674 570729dd-1d6d-4cd6-9876-9a096378119c - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:22:10,270 570729dd-1d6d-4cd6-9876-9a096378119c - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 1.47MB/s]"
           ]
-        }
-      ],
-      "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",
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "with llama.session(remote=True) as session:\n",
         "\n",
-        "    # Create our LORA on the last mlp\n",
-        "    lora = LORA(module, dim, lora_dim)\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",
-        "    # Create an optimizer. Use the parameters from LORA\n",
-        "    optimizer = torch.optim.AdamW(lora.parameters(), lr=3)\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",
-        "    # Iterate over dataloader using .iter. \n",
-        "    with session.iter(dataloader, return_context=True) as (batch, iterator):\n",
+        "print(\"\\nList: \", li2)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CZTVwF7Fkkw9"
+      },
+      "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",
-        "        prompt = batch[0]\n",
-        "        correct_token = batch[1]\n",
+        "<details>\n",
+        "<summary>A full list of NNsight Built-ins</summary>\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",
+        "`nnsight.bool()` creates a traceable Boolean\n",
         "\n",
-        "            # Get logits\n",
-        "            logits = llama.lm_head.output\n",
+        "`nnsight.bytes()` creates a traceable Bytes\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",
+        "`nnsight.int()` creates a traceable Integer\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"
+        "`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": {},
+      "metadata": {
+        "id": "pDF_eijqkkw9"
+      },
       "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. "
+        "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": 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"
-          ]
+      "execution_count": 128,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 365
         },
+        "id": "IhnG03yckkw9",
+        "outputId": "fe9ef3ca-e1b3-474d-b15f-f08d5dcd7e90"
+      },
+      "outputs": [
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "\n"
+            "2025-01-31 15:22:12,730 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:22:12,900 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:22:13,237 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:22:13,240 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 0\n",
+            "2025-01-31 15:22:13,242 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 2\n",
+            "2025-01-31 15:22:13,242 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 1\n",
+            "2025-01-31 15:22:13,759 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 4.74MB/s]\n"
           ]
         }
       ],
       "source": [
-        "# With lora. Should produce \"Hello Paris\"\n",
-        "with llama.generate(\"Hello\", remote=True) as generator:\n",
-        "\n",
-        "    lora()\n",
+        "with llama.session(remote=True) as session:\n",
         "\n",
-        "    out = llama.generator.output.save()\n",
+        "  # with session.iter([0, 1, 2, 3], return_context=True) as (item, iterator):\n",
+        "  with session.iter([0, 1, 2, 3]) as item:\n",
         "\n",
-        "print(llama.tokenizer.batch_decode(out.value))\n",
+        "      nnsight.log(item)\n",
         "\n",
-        "# Then without. Should produce \"Hello,\"\n",
-        "with llama.generate(\"Hello\", remote=True) as generator:\n",
+        "      with nnsight.cond(item == 2):\n",
+        "        nnsight.stop()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VvoZz5OGkkw9"
+      },
+      "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",
-        "    out = llama.generator.output.save()\n",
+        "But in what kind of experimental scenario would someone need iterators?\n",
         "\n",
-        "print(llama.tokenizer.batch_decode(out.value))\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": {
-        "id": "P3zRm-7VRRov"
+        "id": "1OcolbxAkkw9"
       },
       "source": [
-        "# Getting Involved!\n"
+        "## Training a LoRA"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "BnCc9xgjvxEP"
+        "id": "jzyt5d6Dkkw9"
+      },
+      "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. Using session and iterator contexts, we're going apply a very simple fine-tuning approach called low-rank adaptation (LoRA).\n",
+        "\n",
+        "Let's try training a LoRA that, when applied, makes our model always predict \"Paris\" no matter what."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "id": "e9YlPQ-lkkw9"
       },
+      "outputs": [],
       "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",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import nnsight\n",
+        "# from nnsight.envoy import Envoy # this moved in 0.4\n",
+        "from nnsight 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",
-        "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",
+        "        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",
-        "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",
+        "        # LORA is additive\n",
+        "        h = BA_x + self.module.output\n",
         "\n",
-        "You can also follow us on X/Twitter [@ndif_team](https://x.com/ndif_team).\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",
-        "💟\n"
+        "    def parameters(self):\n",
+        "        # Some way to get all the parameters.\n",
+        "        return [self.WA, self.WB]"
       ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "provenance": [],
-      "toc_visible": true
     },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_Ub6ApqAkkw9"
+      },
+      "source": [
+        "Let's define all the variables to use in LoRA training."
+      ]
     },
-    "language_info": {
-      "codemirror_mode": {
-        "name": "ipython",
-        "version": 3
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "id": "yToy3qp1kkw9"
       },
-      "file_extension": ".py",
-      "mimetype": "text/x-python",
-      "name": "python",
-      "nbconvert_exporter": "python",
-      "pygments_lexer": "ipython3",
-      "version": "3.12.5"
+      "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"
+      ]
     },
-    "widgets": {
-      "application/vnd.jupyter.widget-state+json": {
-        "0a676b8f2dc34fed9731f19a03e4bfa1": {
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NNriyaBikkw9"
+      },
+      "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": 12,
+      "metadata": {
+        "id": "emv134jkkkw9"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "4096\n"
+          ]
+        }
+      ],
+      "source": [
+        "with llama.scan(\" \"):\n",
+        "    dim = module.output.shape[-1]\n",
+        "\n",
+        "print(dim)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "63aWV_ZTkkw9"
+      },
+      "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": 13,
+      "metadata": {
+        "id": "Q2SYDaawkkw9"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:48:38,456 810d88f8-f1fc-43fd-b16d-77e11841eecd - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:38,634 810d88f8-f1fc-43fd-b16d-77e11841eecd - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:38,859 810d88f8-f1fc-43fd-b16d-77e11841eecd - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:39,402 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.7539,  0.0732, -0.4023, -0.0615],\n",
+            "        [-0.5312,  0.5391,  0.5195, -1.1328],\n",
+            "        [-1.5781, -0.2480,  0.6953,  0.3535],\n",
+            "        ...,\n",
+            "        [ 0.1670,  0.5391, -0.5703, -0.6289],\n",
+            "        [-0.0762, -1.3438,  0.8320,  1.2656],\n",
+            "        [ 0.2500,  1.2188,  0.2891, -1.2578]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,410 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.7305,  0.0713, -0.3906, -0.0598],\n",
+            "        [-0.5156,  0.5234,  0.5039, -1.1016],\n",
+            "        [-1.5312, -0.2402,  0.6758,  0.3438],\n",
+            "        ...,\n",
+            "        [ 0.1621,  0.5234, -0.5547, -0.6094],\n",
+            "        [-0.0737, -1.3047,  0.8086,  1.2266],\n",
+            "        [ 0.2422,  1.1797,  0.2812, -1.2188]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,448 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.7070,  0.0688, -0.3789, -0.0586],\n",
+            "        [-0.5000,  0.5078,  0.4883, -1.0703],\n",
+            "        [-1.4844, -0.2334,  0.6562,  0.3340],\n",
+            "        ...,\n",
+            "        [ 0.1572,  0.5078, -0.5391, -0.5898],\n",
+            "        [-0.0718, -1.2656,  0.7852,  1.1875],\n",
+            "        [ 0.2344,  1.1406,  0.2734, -1.1797]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,485 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.6836,  0.0645, -0.3633, -0.0598],\n",
+            "        [-0.4883,  0.4902,  0.4746, -1.0391],\n",
+            "        [-1.4375, -0.2285,  0.6406,  0.3203],\n",
+            "        ...,\n",
+            "        [ 0.1543,  0.4941, -0.5234, -0.5703],\n",
+            "        [-0.0718, -1.2266,  0.7617,  1.1484],\n",
+            "        [ 0.2227,  1.1094,  0.2695, -1.1484]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,521 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-0.9648, -1.5703,  0.6094, -1.6875],\n",
+            "        [-2.0938, -1.1562,  1.3906, -2.6406],\n",
+            "        [-3.0156, -1.8516,  1.5938, -1.3047],\n",
+            "        ...,\n",
+            "        [ 1.7578,  2.0781, -1.1406,  1.0625],\n",
+            "        [-1.6875, -2.7969,  1.4688, -0.5117],\n",
+            "        [-1.4219, -0.5508,  1.3516, -2.7344]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,558 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-2.7031e+00, -3.6719e-01, -9.6484e-01, -1.0938e+00],\n",
+            "        [-3.7969e+00,  2.1606e-02, -1.9238e-01, -2.0156e+00],\n",
+            "        [-4.6875e+00, -6.4453e-01,  3.3875e-03, -7.1484e-01],\n",
+            "        ...,\n",
+            "        [ 3.4688e+00,  8.6719e-01,  4.1602e-01,  4.8047e-01],\n",
+            "        [-3.4062e+00, -1.5703e+00, -1.0645e-01,  5.7129e-02],\n",
+            "        [-3.1406e+00,  6.1719e-01, -2.3633e-01, -2.0938e+00]],\n",
+            "       requires_grad=True)\n",
+            "2025-01-31 15:48:39,630 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-4.7812,  1.3125, -2.5625,  0.5117],\n",
+            "        [-5.8125,  1.6953, -1.8203, -0.3789],\n",
+            "        [-6.6875,  1.0469, -1.6250,  0.8789],\n",
+            "        ...,\n",
+            "        [ 5.5000, -0.8320,  2.0312, -1.1016],\n",
+            "        [-5.4375,  0.1494, -1.7266,  1.6406],\n",
+            "        [-5.1875,  2.2656, -1.8594, -0.4629]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,697 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-3.2812,  2.9062, -4.0938,  2.1094],\n",
+            "        [-4.3125,  3.2812, -3.3750,  1.2422],\n",
+            "        [-5.1562,  2.6406, -3.1719,  2.4531],\n",
+            "        ...,\n",
+            "        [ 4.0000, -2.4531,  3.5625, -2.6875],\n",
+            "        [-3.9375,  1.7812, -3.2656,  3.2031],\n",
+            "        [-3.6875,  3.8438, -3.4062,  1.1719]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,764 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-2.0000,  4.2500, -5.3750,  3.4688],\n",
+            "        [-3.0156,  4.6250, -4.6875,  2.6406],\n",
+            "        [-3.8125,  4.0000, -4.5000,  3.7969],\n",
+            "        ...,\n",
+            "        [ 2.6875, -3.8438,  4.8750, -4.0312],\n",
+            "        [-2.6250,  3.1719, -4.5938,  4.5312],\n",
+            "        [-2.3906,  5.1875, -4.7188,  2.5625]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,830 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-0.8867,  5.4062, -6.4688,  4.6250],\n",
+            "        [-1.8750,  5.7812, -5.8125,  3.8438],\n",
+            "        [-2.6406,  5.1562, -5.6562,  4.9375],\n",
+            "        ...,\n",
+            "        [ 1.5469, -5.0312,  6.0000, -5.1875],\n",
+            "        [-1.4922,  4.3750, -5.7188,  5.6875],\n",
+            "        [-1.2500,  6.3438, -5.8125,  3.7656]], requires_grad=True)\n",
+            "2025-01-31 15:48:40,224 810d88f8-f1fc-43fd-b16d-77e11841eecd - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 67.1k/67.1k [00:00<00:00, 1.65MB/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) as batch:\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",
+        "        nnsight.log(lora.WA)\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2FtFv2Nykkw9"
+      },
+      "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": 14,
+      "metadata": {
+        "id": "K5_F4ktlkkw9"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:48:45,111 c9fcc827-5c78-4796-87cf-0f246c75a359 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:45,281 c9fcc827-5c78-4796-87cf-0f246c75a359 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:45,535 c9fcc827-5c78-4796-87cf-0f246c75a359 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:46,013 c9fcc827-5c78-4796-87cf-0f246c75a359 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 1.24k/1.24k [00:00<00:00, 4.30MB/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "['<|begin_of_text|>Hello Paris']\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:48:46,446 9f088820-f00d-492d-9861-c0f0a7c4d784 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:46,650 9f088820-f00d-492d-9861-c0f0a7c4d784 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:46,862 9f088820-f00d-492d-9861-c0f0a7c4d784 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:47,340 9f088820-f00d-492d-9861-c0f0a7c4d784 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 1.24k/1.24k [00:00<00:00, 4.61MB/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": "IZfiTNGgu-HI"
+      },
+      "source": [
+        "# Next Steps\n",
+        "Check out [nnsight.net/tutorials](https://nnsight.net/tutorials) for more walkthroughs implementating classic interpretability techniques using `nnsight`."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "P3zRm-7VRRov"
+      },
+      "source": [
+        "# Getting Involved!\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BnCc9xgjvxEP"
+      },
+      "source": [
+        "Note that both `nnsight` and `NDIF` are in active development, so changes may be made and errors may arise during use. If you’re interested in following updates to `nnsight`, contributing, giving feedback, or finding collaborators, please join the [NDIF discord](https://discord.gg/6uFJmCSwW7). We’d love to hear about your work using nnsight!\n",
+        "\n",
+        "You can also follow us on [LinkedIn](https://www.linkedin.com/company/national-deep-inference-fabric/), Bluesky: [@ndif-team.bsky.social](https://bsky.app/profile/ndif-team.bsky.social), and X: [@ndif_team](https://x.com/ndif_team).\n",
+        "\n",
+        "💟\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "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.10.13"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "0383a2798a0c4fcfbd1f068817770fd6": {
+          "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_cef2fe4eb51146afabe820140a3b0df3",
+            "max": 456318,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_52281c8594ad465397e6079385c558a1",
+            "value": 456318
+          }
+        },
+        "08fb0620499c44eb9d160a0f21263451": {
+          "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
+          }
+        },
+        "0c46bce8cc3e4393898b025ba6b4e543": {
+          "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": ""
+          }
+        },
+        "14008810744549208f147a132b419d20": {
+          "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_cb4be83440e44e1c9e664ac1874b830d",
+              "IPY_MODEL_f54c77317ca04861a958c8f982b25653",
+              "IPY_MODEL_db24a855fb5846e28c98dcfe9afff259"
+            ],
+            "layout": "IPY_MODEL_647c861fcaa44148a90da061b1d5c99c"
+          }
+        },
+        "184b5c75820b4c42ae861669ac1a47e5": {
+          "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": ""
+          }
+        },
+        "19c0e121e52040b18393aee887730644": {
+          "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
+          }
+        },
+        "26acae3d321743d2ad8efbf8727dbd05": {
+          "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_a4b5c225f5ac4719ade8e6aafcff78d2",
+            "max": 548105171,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_ee27d18a555c46e78c92b282bfa6dc9d",
+            "value": 548105171
+          }
+        },
+        "26dcd7b75feb41a6b86d4260ecd64d06": {
+          "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_39c9f1effe664e899125748e97fad198",
+              "IPY_MODEL_383b893224684aee916e455c0d3d8c06",
+              "IPY_MODEL_af24a12bd05b44148b5d6f91ac2445a0"
+            ],
+            "layout": "IPY_MODEL_d00750501f2548d0b1c69f4753aec877"
+          }
+        },
+        "2702f30b5d3b4b0b8184fec58863750b": {
+          "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": ""
+          }
+        },
+        "287a8fe3ddf54d36b7603171ec7a2688": {
+          "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": ""
+          }
+        },
+        "2a0f2484053f41298aceacea5c1cc195": {
+          "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_08fb0620499c44eb9d160a0f21263451",
+            "max": 124,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_287a8fe3ddf54d36b7603171ec7a2688",
+            "value": 124
+          }
+        },
+        "2fedb8b529c34049932f07b577b515d4": {
+          "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_bd1a11654c2d47a09610e4cc23b8bc32",
+              "IPY_MODEL_26acae3d321743d2ad8efbf8727dbd05",
+              "IPY_MODEL_f42e44122a2d496cb7b2ea85c2858a71"
+            ],
+            "layout": "IPY_MODEL_c1586aeef6814447ad46c1f8ddaae625"
+          }
+        },
+        "346ecc4734014cb4931bea53d5a32bf1": {
+          "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": ""
+          }
+        },
+        "383b893224684aee916e455c0d3d8c06": {
+          "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_68b205f60f9d4599a7d1f78583d0b93c",
+            "max": 665,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_69e8c149f676457d95029f21e9a2956a",
+            "value": 665
+          }
+        },
+        "385291e96814453b8f805cbe8a5e594d": {
+          "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_e5302eb0481c4e529fc68353eed4b79f",
+            "placeholder": "​",
+            "style": "IPY_MODEL_0c46bce8cc3e4393898b025ba6b4e543",
+            "value": " 456k/456k [00:00&lt;00:00, 3.39MB/s]"
+          }
+        },
+        "39c9f1effe664e899125748e97fad198": {
+          "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_98d28dddf75a45a2b1e05131b0b865a1",
+            "placeholder": "​",
+            "style": "IPY_MODEL_6e81d6f31f2446c999df5511e8afd8d0",
+            "value": "config.json: 100%"
+          }
+        },
+        "405308ec76074c7ea319e1c1358baf39": {
+          "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": ""
+          }
+        },
+        "41fc21ee95c14d29943a56f2b63f4989": {
+          "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_42fd0ea3be3c4b3f98ca7818efee3669",
+            "placeholder": "​",
+            "style": "IPY_MODEL_184b5c75820b4c42ae861669ac1a47e5",
+            "value": " 26.0/26.0 [00:00&lt;00:00, 2.05kB/s]"
+          }
+        },
+        "42fd0ea3be3c4b3f98ca7818efee3669": {
+          "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
+          }
+        },
+        "4b798a784ae748b1a0aef025ead976fc": {
+          "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": ""
+          }
+        },
+        "504038c4efd3439daa958333f108b5fb": {
+          "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": ""
+          }
+        },
+        "52281c8594ad465397e6079385c558a1": {
+          "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": ""
+          }
+        },
+        "5250cf68037a43e389493733c5b9d4ac": {
+          "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
+          }
+        },
+        "55c0bb35417e454d842e12bd2dfe62df": {
+          "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
+          }
+        },
+        "57aebe107749473390af2532806920ce": {
+          "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_55c0bb35417e454d842e12bd2dfe62df",
+            "placeholder": "​",
+            "style": "IPY_MODEL_d5821d69e1fa4673bb39ab70a7f60466",
+            "value": " 124/124 [00:00&lt;00:00, 8.42kB/s]"
+          }
+        },
+        "5d877979f2db4dbfa7581d2fe4381cf8": {
+          "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": ""
+          }
+        },
+        "5e091b15cc9748469676b165533f0c44": {
+          "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
+          }
+        },
+        "647c861fcaa44148a90da061b1d5c99c": {
+          "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
+          }
+        },
+        "64b713188f054535a3063c079bceead0": {
+          "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_7afff81060f0481aa95a56daa570e1b8",
+            "max": 1042301,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_a038f9037ed540f2ba43f3f85920b17e",
+            "value": 1042301
+          }
+        },
+        "68b205f60f9d4599a7d1f78583d0b93c": {
+          "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
+          }
+        },
+        "69e8c149f676457d95029f21e9a2956a": {
+          "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": ""
+          }
+        },
+        "6e81d6f31f2446c999df5511e8afd8d0": {
+          "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": ""
+          }
+        },
+        "7479281e090d4eeba4c21a9c15b02165": {
+          "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_f4b25b3aac414499b653404ea3c066da",
+              "IPY_MODEL_0383a2798a0c4fcfbd1f068817770fd6",
+              "IPY_MODEL_385291e96814453b8f805cbe8a5e594d"
+            ],
+            "layout": "IPY_MODEL_f7a41e267f79401681bc4693c52e55fa"
+          }
+        },
+        "76fe13ab697c46cc9fcbe46112fb45b1": {
+          "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": ""
+          }
+        },
+        "79a221c9228d459aa17402c9fef1ad1d": {
+          "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_9c13a3d768eb47c8b26a7c6bfa8ec957",
+              "IPY_MODEL_e97c052113f54dddbb619ed228f6b4ef",
+              "IPY_MODEL_41fc21ee95c14d29943a56f2b63f4989"
+            ],
+            "layout": "IPY_MODEL_5250cf68037a43e389493733c5b9d4ac"
+          }
+        },
+        "7afff81060f0481aa95a56daa570e1b8": {
+          "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
+          }
+        },
+        "7f4f63d94f8942fcb3e4b286d5848f0c": {
+          "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
+          }
+        },
+        "7fdb623e5b2d4eb1b6cee27da2314062": {
+          "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
+          }
+        },
+        "87a5c035bc6345aea6441b3db4b94321": {
+          "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
+          }
+        },
+        "95832c5076ab4c80810d3c443e061156": {
+          "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": ""
+          }
+        },
+        "98d28dddf75a45a2b1e05131b0b865a1": {
+          "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
+          }
+        },
+        "9c13a3d768eb47c8b26a7c6bfa8ec957": {
+          "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_e24033c5ba7b4e4fa4f4375d85b03bad",
+            "placeholder": "​",
+            "style": "IPY_MODEL_4b798a784ae748b1a0aef025ead976fc",
+            "value": "tokenizer_config.json: 100%"
+          }
+        },
+        "9cea52b7835b40628fc0a65199c5e82f": {
+          "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_ec4f48d5a9ef4afa8fa44078dfb236c6",
+            "placeholder": "​",
+            "style": "IPY_MODEL_405308ec76074c7ea319e1c1358baf39",
+            "value": " 1.04M/1.04M [00:00&lt;00:00, 7.50MB/s]"
+          }
+        },
+        "a038f9037ed540f2ba43f3f85920b17e": {
+          "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": ""
+          }
+        },
+        "a15895ad1cb541ca89fa9def63779104": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HBoxModel",
@@ -2859,14 +5335,14 @@
             "_view_name": "HBoxView",
             "box_style": "",
             "children": [
-              "IPY_MODEL_74358c341daa44faa41a92b0eb073615",
-              "IPY_MODEL_c40e17a994ed44d79ef54421f4772858",
-              "IPY_MODEL_431ffcb0dbb1419dbbb237a29842b905"
+              "IPY_MODEL_ee60494714c2485f802298dce918edfe",
+              "IPY_MODEL_2a0f2484053f41298aceacea5c1cc195",
+              "IPY_MODEL_57aebe107749473390af2532806920ce"
             ],
-            "layout": "IPY_MODEL_c28c185a508e459b95edb29afedd9c79"
+            "layout": "IPY_MODEL_ff44e896c55048c5a7d92a56aa9eedfc"
           }
         },
-        "0fddebfb42e847e08890d09f30161984": {
+        "a1c25910a31a4b5a9c48f1a7fd8047c5": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -2918,7 +5394,7 @@
             "width": null
           }
         },
-        "11e99924bdf84739b3cfc7d676b797e5": {
+        "a310703500c5481aa37b698c659e1ca6": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "DescriptionStyleModel",
@@ -2933,7 +5409,7 @@
             "description_width": ""
           }
         },
-        "18d4c2302f424bc18f9215f2d5b83393": {
+        "a3154c60744d4a1a87b1c86b89a142b9": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -2985,43 +5461,7 @@
             "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": {
+        "a4b5c225f5ac4719ade8e6aafcff78d2": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3073,7 +5513,7 @@
             "width": null
           }
         },
-        "431ffcb0dbb1419dbbb237a29842b905": {
+        "af24a12bd05b44148b5d6f91ac2445a0": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3088,13 +5528,13 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_562a287082eb4882af864062ae13c91b",
+            "layout": "IPY_MODEL_ba90058201ee48dcae9411ded457e85f",
             "placeholder": "​",
-            "style": "IPY_MODEL_11e99924bdf84739b3cfc7d676b797e5",
-            "value": " 665/665 [00:00&lt;00:00, 12.8kB/s]"
+            "style": "IPY_MODEL_346ecc4734014cb4931bea53d5a32bf1",
+            "value": " 665/665 [00:00&lt;00:00, 36.5kB/s]"
           }
         },
-        "562a287082eb4882af864062ae13c91b": {
+        "ba90058201ee48dcae9411ded457e85f": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3146,76 +5586,7 @@
             "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": {
+        "bd1a11654c2d47a09610e4cc23b8bc32": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3230,13 +5601,65 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_7a0f3d2deeb14bbba7f0df208133a7e2",
+            "layout": "IPY_MODEL_5e091b15cc9748469676b165533f0c44",
             "placeholder": "​",
-            "style": "IPY_MODEL_94cdcf1fb0914dfe99c3c615e68b7ec8",
-            "value": "config.json: 100%"
+            "style": "IPY_MODEL_f143977aeaf24798b7744681069762d4",
+            "value": "model.safetensors: 100%"
+          }
+        },
+        "c1586aeef6814447ad46c1f8ddaae625": {
+          "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
           }
         },
-        "7a0f3d2deeb14bbba7f0df208133a7e2": {
+        "caff23247fb9432db8d5860a8132e1b3": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3288,7 +5711,7 @@
             "width": null
           }
         },
-        "7e6b41ec535648b190bfc655d9e430dd": {
+        "cb4be83440e44e1c9e664ac1874b830d": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3303,13 +5726,13 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_beb4231b49024d8a9bb7ead19647d99d",
+            "layout": "IPY_MODEL_e4e0b1de76ff40a88b71e0b5a9f7bc39",
             "placeholder": "​",
-            "style": "IPY_MODEL_26d09aa072d64af0bf4df9a89309093a",
-            "value": " 1.36M/1.36M [00:00&lt;00:00, 4.16MB/s]"
+            "style": "IPY_MODEL_504038c4efd3439daa958333f108b5fb",
+            "value": "tokenizer.json: 100%"
           }
         },
-        "8776d04bfd3748b69e10708466e9c132": {
+        "cef2fe4eb51146afabe820140a3b0df3": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3361,7 +5784,7 @@
             "width": null
           }
         },
-        "8de63414803f41c29e75242756cfd67f": {
+        "d00750501f2548d0b1c69f4753aec877": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3413,31 +5836,7 @@
             "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": {
+        "d5821d69e1fa4673bb39ab70a7f60466": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "DescriptionStyleModel",
@@ -3452,28 +5851,29 @@
             "description_width": ""
           }
         },
-        "95d303d83d3348a0adfdf6b5cb94c41c": {
+        "d8feaceb6d854cb2b2ff19a0309d9a7b": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
+          "model_name": "HBoxModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
+            "_model_name": "HBoxModel",
             "_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%"
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_fd84ffef58604ce3a16ef8f9868072f3",
+              "IPY_MODEL_64b713188f054535a3063c079bceead0",
+              "IPY_MODEL_9cea52b7835b40628fc0a65199c5e82f"
+            ],
+            "layout": "IPY_MODEL_87a5c035bc6345aea6441b3db4b94321"
           }
         },
-        "96a3b8c765374a5e98eaf713dd5e85d8": {
+        "db24a855fb5846e28c98dcfe9afff259": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3488,65 +5888,28 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_8776d04bfd3748b69e10708466e9c132",
+            "layout": "IPY_MODEL_a1c25910a31a4b5a9c48f1a7fd8047c5",
             "placeholder": "​",
-            "style": "IPY_MODEL_56490c20aef04fe4a63ba7c990e71e77",
-            "value": " 1.04M/1.04M [00:00&lt;00:00, 3.16MB/s]"
+            "style": "IPY_MODEL_5d877979f2db4dbfa7581d2fe4381cf8",
+            "value": " 1.36M/1.36M [00:00&lt;00:00, 6.83MB/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
+        "de023d4526524f48a34883a287aa8a4f": {
+          "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": ""
           }
         },
-        "b49ce0d00e5b4df097b3a9fc4c0bd8e7": {
+        "ded8c188917b424dbf264ce494fe0532": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3598,60 +5961,7 @@
             "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": {
+        "e24033c5ba7b4e4fa4f4375d85b03bad": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3703,7 +6013,7 @@
             "width": null
           }
         },
-        "c28c185a508e459b95edb29afedd9c79": {
+        "e4e0b1de76ff40a88b71e0b5a9f7bc39": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3755,46 +6065,7 @@
             "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": {
+        "e5302eb0481c4e529fc68353eed4b79f": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3846,29 +6117,31 @@
             "width": null
           }
         },
-        "cef2516aabc24f6d852079984afa8090": {
+        "e97c052113f54dddbb619ed228f6b4ef": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "HBoxModel",
+          "model_name": "FloatProgressModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "HBoxModel",
+            "_model_name": "FloatProgressModel",
             "_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"
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_19c0e121e52040b18393aee887730644",
+            "max": 26,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_2702f30b5d3b4b0b8184fec58863750b",
+            "value": 26
           }
         },
-        "d0198fa98e5f412380f2ba3e9448aef3": {
+        "ec4f48d5a9ef4afa8fa44078dfb236c6": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3920,7 +6193,7 @@
             "width": null
           }
         },
-        "d07a0a49a2924092affc0e07bf0b5ef0": {
+        "ee27d18a555c46e78c92b282bfa6dc9d": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "ProgressStyleModel",
@@ -3936,55 +6209,79 @@
             "description_width": ""
           }
         },
-        "d839da1c31a34659988e3420ba520725": {
+        "ee60494714c2485f802298dce918edfe": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "FloatProgressModel",
+          "model_name": "HTMLModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "FloatProgressModel",
+            "_model_name": "HTMLModel",
             "_view_count": null,
             "_view_module": "@jupyter-widgets/controls",
             "_view_module_version": "1.5.0",
-            "_view_name": "ProgressView",
-            "bar_style": "success",
+            "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_3731a5747724403e86ef1abaf4fb389b",
-            "max": 1042301,
-            "min": 0,
-            "orientation": "horizontal",
-            "style": "IPY_MODEL_688bf676958a4050b73cc9597f5110de",
-            "value": 1042301
+            "layout": "IPY_MODEL_7fdb623e5b2d4eb1b6cee27da2314062",
+            "placeholder": "​",
+            "style": "IPY_MODEL_95832c5076ab4c80810d3c443e061156",
+            "value": "generation_config.json: 100%"
           }
         },
-        "e4568f72971749ac92cc5577cb9c4f9c": {
+        "f143977aeaf24798b7744681069762d4": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "FloatProgressModel",
+          "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": ""
+          }
+        },
+        "f42e44122a2d496cb7b2ea85c2858a71": {
+          "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": "FloatProgressModel",
+            "_model_name": "HTMLModel",
             "_view_count": null,
             "_view_module": "@jupyter-widgets/controls",
             "_view_module_version": "1.5.0",
-            "_view_name": "ProgressView",
-            "bar_style": "success",
+            "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_ceebda1d07b545f5b605551c08e73be8",
-            "max": 456318,
-            "min": 0,
-            "orientation": "horizontal",
-            "style": "IPY_MODEL_6aa96b80e3b746b48f5c13bcafe1134e",
-            "value": 456318
+            "layout": "IPY_MODEL_a3154c60744d4a1a87b1c86b89a142b9",
+            "placeholder": "​",
+            "style": "IPY_MODEL_de023d4526524f48a34883a287aa8a4f",
+            "value": " 548M/548M [00:02&lt;00:00, 253MB/s]"
+          }
+        },
+        "f47609606aa14f6fa2407aa26edae8e3": {
+          "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": ""
           }
         },
-        "e51620404d244a4e806eeaf3b96a474f": {
+        "f4b25b3aac414499b653404ea3c066da": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3999,34 +6296,37 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_f28f557210b7402b8b2d0fd83f038dc1",
+            "layout": "IPY_MODEL_caff23247fb9432db8d5860a8132e1b3",
             "placeholder": "​",
-            "style": "IPY_MODEL_bc8190c700e949ebb6f874c8f8055161",
-            "value": "merges.txt: 100%"
+            "style": "IPY_MODEL_f47609606aa14f6fa2407aa26edae8e3",
+            "value": "merges.txt: 100%"
           }
         },
-        "e7a73cfc3fa04425afc08368532ef0be": {
+        "f54c77317ca04861a958c8f982b25653": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
+          "model_name": "FloatProgressModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
+            "_model_name": "FloatProgressModel",
             "_view_count": null,
             "_view_module": "@jupyter-widgets/controls",
             "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_b2748df9aff540df954f002310606a56",
-            "placeholder": "​",
-            "style": "IPY_MODEL_f4caf2bc58704464b6ac7bf6da31b71d",
-            "value": "vocab.json: 100%"
+            "layout": "IPY_MODEL_ded8c188917b424dbf264ce494fe0532",
+            "max": 1355256,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_76fe13ab697c46cc9fcbe46112fb45b1",
+            "value": 1355256
           }
         },
-        "eb0c489b913744a28304e5334bb66532": {
+        "f7a41e267f79401681bc4693c52e55fa": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -4078,59 +6378,28 @@
             "width": null
           }
         },
-        "f28f557210b7402b8b2d0fd83f038dc1": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
+        "fd84ffef58604ce3a16ef8f9868072f3": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
           "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
             "_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
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_7f4f63d94f8942fcb3e4b286d5848f0c",
+            "placeholder": "​",
+            "style": "IPY_MODEL_a310703500c5481aa37b698c659e1ca6",
+            "value": "vocab.json: 100%"
           }
         },
-        "f3e24aa6619e4185a3e2f53f2e473e45": {
+        "ff44e896c55048c5a7d92a56aa9eedfc": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -4181,36 +6450,6 @@
             "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": ""
-          }
         }
       }
     }
diff --git a/NNsight_v0_4_guide.ipynb b/NNsight_v0_4_guide.ipynb
new file mode 100644
index 00000000..4d6ffb33
--- /dev/null
+++ b/NNsight_v0_4_guide.ipynb
@@ -0,0 +1,4223 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JYWncbOUZzlm"
+      },
+      "source": [
+        "# `nnsight 0.4`: walkthrough\n",
+        "**We have many exciting new features in this update, including:**\n",
+        "\n",
+        "*   Descriptive error messages\n",
+        "*   .all() for multiple token generation\n",
+        "*   vLLM Integration\n",
+        "*   Streaming remote execution to local machine\n",
+        "*   Support for traditional Python syntax for `if` statements and `for` loops on proxies within tracing contexts\n",
+        "*   Ability to rename model modules\n",
+        "\n",
+        "...and more!\n",
+        "\n",
+        "The following walkthrough guides you through how to access `nnsight 0.4` and use all of its individual features. \n",
+        "\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FNqrMwiUmqOB"
+      },
+      "source": [
+        "**Breaking Changes**\n",
+        "* The InterventionGraph now follows a <u>sequential execution order</u>. Module envoys are expected to be referenced following the model’s architecture hierarchy. This means that out-of-order in-place operations will not take effect.\n",
+        "\n",
+        "* Saved node values are automatically injected into their proxy reference in the Python frames post graph execution. If you are calling `.value` in your code after tracing, this could lead to the wrong behavior."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6a2QUOA7URPF"
+      },
+      "source": [
+        "# Access `nnsight` 0.4"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "kq2qmH42sFX9"
+      },
+      "outputs": [],
+      "source": [
+        "from IPython.display import clear_output\n",
+        "!pip install nnsight\n",
+        "clear_output()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6SeRg2jrkpt0"
+      },
+      "source": [
+        "Import `nnsight` and load the GPT-2 model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "qakQzKv8kRQh"
+      },
+      "outputs": [],
+      "source": [
+        "# import packages\n",
+        "import nnsight\n",
+        "from nnsight import LanguageModel"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 713,
+          "referenced_widgets": [
+            "0f3eaadfc2c34384afd2979417d94812",
+            "a2eafc8f9262412793f35b372eb300eb",
+            "cb3a1a3578554cac9cec37d9b0390c1a",
+            "17f30f75790e4273b3e475f433c4760d",
+            "1046b0563cb64ce68fc7b1b2f5d3d65d",
+            "886f12cb69c2435d978eb2d33534763a",
+            "cb0ab1cce5364b44ba85e65ccf44d0c9",
+            "0650393001b34123a0e18663e6aa0fcb",
+            "cdfa2b53fdd14d24b9ccc0920f34bd69",
+            "2e14c72eaf6644bd9d656513c302c278",
+            "30af0b2c974d418cb40b30403a2a90cf",
+            "c3b4c47b73d544eaba67d5828f1f742d",
+            "c924f41345d34f72a00ed3b21f84de82",
+            "0e4b52697a87475882e40b3ca4a2e77d",
+            "c4c1c2ccc4014520bb265d21550fb133",
+            "3160c45575ad4426b2abcfaccd27d41c",
+            "4eb7af991ebe4bdd855d3cdd50184495",
+            "c96d870993144ed09e54de5dc00fad6d",
+            "60627b23fe934d01b4a4af1f592fbd55",
+            "a806062fe502459a93ee212e57af2e66",
+            "9b161849be8a4464a486e6270f3c0f3b",
+            "f2ca144c6f8d4537ac8e2e29c8109454",
+            "33c1fe5504d1441e805a02fb4693acd7",
+            "8bf1249ac7aa46b094bc899d432153fa",
+            "c60da3ab273d4439a0baec6ae7c216e4",
+            "9ef4dc73855c4fa69656f75110ac390d",
+            "4b1d8724a6b247e0a3dafe2f6b9de178",
+            "6909cf23963d4347820b3717b62b1cfe",
+            "f6a764157eab4ee7b24ce3dc3e317222",
+            "59f8564845764c9e9891274216be83ff",
+            "3adee7a17b2a4f50b9ff36fe10a3b701",
+            "e60f4ecadf75463ca5fc5f934954e144",
+            "d5f42dda1612435f8eb053fe29f711c0",
+            "a1595afd22fa4d999a651e1ff63f4279",
+            "5c5d0fc4f29141fc95008070c0b077d3",
+            "c410466db4ac4d878c3aac7fe4a0df5b",
+            "1ac5508e75154bd18275d6f6775242b7",
+            "6fc349dcab8d43309faaa08775f445ff",
+            "4d7f03cfa9944554b093b3e5c691db92",
+            "732fe1ca67b74026805de4688d81e6e3",
+            "3c1c03b756c1401da717b3b62842a1f9",
+            "416faee108c44dce9e1b3527e82acb9b",
+            "a22aff0a7b33470b970a32269a9a0f08",
+            "be8e4e0c8e0c482897194f08c77f82a0",
+            "d86efefd42444836b476a5caecccdb75",
+            "2f851281cd2448f8a36c9b538931e0c9",
+            "9a2747e522c94c00af3f9d0bbd54cc1d",
+            "176749c3ecd24b829a78f892dee8af91",
+            "33cc1f2b3b884a65a4a942cde1662b8f",
+            "637e15c89e6847098b27f7a4f3482a63",
+            "e039919f976f4d5fbebb5117ec53d57e",
+            "7c3d6693887b4c29b3d63afd37e4292e",
+            "6386746d591e4e4cb83a18dcc4481678",
+            "0325fe95b3e3432497fc68317b2126d4",
+            "fabbb16b5e7942f493b7dc92279715fd"
+          ]
+        },
+        "id": "0kXdv3B51oEt",
+        "outputId": "e76e7f65-56c0-4244-8e57-98a2475088a4"
+      },
+      "outputs": [
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "0f3eaadfc2c34384afd2979417d94812",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "c3b4c47b73d544eaba67d5828f1f742d",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "33c1fe5504d1441e805a02fb4693acd7",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "a1595afd22fa4d999a651e1ff63f4279",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "d86efefd42444836b476a5caecccdb75",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "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): Generator(\n",
+            "    (streamer): Streamer()\n",
+            "  )\n",
+            ")\n"
+          ]
+        }
+      ],
+      "source": [
+        "model = LanguageModel('openai-community/gpt2', device_map='auto')\n",
+        "print(model)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "yVuOTllrUjgv"
+      },
+      "source": [
+        "# Improved Error Messaging"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "45AT75G4UmRk"
+      },
+      "source": [
+        "If you've been using `nnsight`, you are probably familiar with the following type of error message:\n",
+        "\n",
+        "```\n",
+        "IndexError: Above exception when execution Node: 'setitem_0' in Graph: '6063279136'\n",
+        "```\n",
+        "It can be quite difficult to troubleshoot with these errors, so in `nnsight 0.4` we've now improved error messaging to be descriptive and line-specific! Let's check it out:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 365
+        },
+        "id": "kUtl8NUqWSbJ",
+        "outputId": "0fc43241-7f46-4ab8-831a-089121a68e42"
+      },
+      "outputs": [
+        {
+          "ename": "IndexError",
+          "evalue": "list index out of range",
+          "output_type": "error",
+          "traceback": [
+            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
+            "\u001b[0;32m<ipython-input-4-c0aea694bac2>\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0mprompt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m'The Eiffel Tower is in the city of'\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0;32mwith\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprompt\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      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;31m# try to access a layer of model that doesn't exist\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+            "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/nnsight/intervention/contexts/interleaving.py\u001b[0m in \u001b[0;36m__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     94\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_envoy\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reset\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     95\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 96\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     97\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     98\u001b[0m     \u001b[0;34m@\u001b[0m\u001b[0mclassmethod\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+            "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/nnsight/tracing/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     23\u001b[0m         \u001b[0mGlobalTracingContext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtry_deregister\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[1;32m     24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 25\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \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     26\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0miter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcollection\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.11/dist-packages/nnsight/tracing/contexts/base.py\u001b[0m in \u001b[0;36m__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     70\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     71\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[0mBaseException\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---> 72\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     73\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     74\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgraph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgraph\u001b[0m\u001b[0;34m,\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[0;34m**\u001b[0m\u001b[0mself\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<ipython-input-4-c0aea694bac2>\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0;31m# try to access a layer of model that doesn't exist\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m     \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransformer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mh\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m12\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput\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[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      7\u001b[0m     \u001b[0moutput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlm_head\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\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      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+            "\u001b[0;32m/usr/local/lib/python3.11/dist-packages/nnsight/intervention/envoy.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m    621\u001b[0m         \"\"\"\n\u001b[1;32m    622\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 623\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_children\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\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    624\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    625\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__len__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+            "\u001b[0;31mIndexError\u001b[0m: list index out of range"
+          ]
+        }
+      ],
+      "source": [
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "\n",
+        "with model.trace(prompt) as tracer:\n",
+        "\n",
+        "    # try to access a layer of model that doesn't exist\n",
+        "    model.transformer.h[12].output[0][:] = 0\n",
+        "    output = model.lm_head.output.save()\n",
+        "\n",
+        "print(\"lm_head output = \",output)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "usZHg7l_YUYL"
+      },
+      "source": [
+        "Great! Now we know that our list index was out of range within the tracing context, and if we expand to see the full message, we can tell that it's happening in line 7.\n",
+        "\n",
+        "Let's try again, now using the correct index for the final layer:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 224,
+          "referenced_widgets": [
+            "63dc1c0ce98d4539984bfc93114aa6c6",
+            "1c6092db09c04e568b23ca71eb3f5f13",
+            "54fc39e52ff94065a20594cc95a35942",
+            "581c1d23c4424ce6b50aec079df134c4",
+            "c7c13deb642f4af1bedd88c2afac85fd",
+            "dea619854cbb42a2a700f7feef77d442",
+            "c66f56814daa4745a0f14c75f85d0294",
+            "5e25783cecd346f48959de00bc2cccec",
+            "c700fe2789cb4f51a43ae1e29932d3fc",
+            "e122506e7c7c454c952b89f306360da6",
+            "53f9e552ea174246a10de4afc84272ac",
+            "4fc680f767b04af881b52fd94ea07bba",
+            "45f3295ab6e64cbda41f26c2d1581b08",
+            "78018bc8cd974345b090ad38214e75c6",
+            "56806395ddfe45f2b0245741a41e5bf5",
+            "87e140a2c262488da5b393f37b51592c",
+            "3b2050ce41594c69bf6efbacc76a574a",
+            "15332cc2b6e84cfe937eece4ec884a72",
+            "259617ed81dc4dd4b40f657ad63bec47",
+            "16f914ebda5b44d28cd6aecc5ae742b7",
+            "2130a5e5fd6d44ce8484c51d591e2d6e",
+            "25af4170e72745a095f37ef4a481fab6"
+          ]
+        },
+        "id": "Xu-WWnZsYpbw",
+        "outputId": "7f33037a-066e-49c0-fef5-4fbb0267ac02"
+      },
+      "outputs": [
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "63dc1c0ce98d4539984bfc93114aa6c6",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "model.safetensors:   0%|          | 0.00/548M [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "application/vnd.jupyter.widget-view+json": {
+              "model_id": "4fc680f767b04af881b52fd94ea07bba",
+              "version_major": 2,
+              "version_minor": 0
+            },
+            "text/plain": [
+              "generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "lm_head output =  tensor([[[ -6.3267,  -6.1134,  -8.2121,  ..., -11.1459, -10.8880,  -6.1064],\n",
+            "         [ -6.3267,  -6.1134,  -8.2121,  ..., -11.1459, -10.8880,  -6.1064],\n",
+            "         [ -6.3267,  -6.1134,  -8.2121,  ..., -11.1459, -10.8880,  -6.1064],\n",
+            "         ...,\n",
+            "         [ -6.3267,  -6.1134,  -8.2121,  ..., -11.1459, -10.8880,  -6.1064],\n",
+            "         [ -6.3267,  -6.1134,  -8.2121,  ..., -11.1459, -10.8880,  -6.1064],\n",
+            "         [ -6.3267,  -6.1134,  -8.2121,  ..., -11.1459, -10.8880,  -6.1064]]],\n",
+            "       device='cuda:0', grad_fn=<UnsafeViewBackward0>)\n"
+          ]
+        }
+      ],
+      "source": [
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "\n",
+        "with model.trace(prompt) as tracer:\n",
+        "\n",
+        "    # ablate last layer output\n",
+        "    model.transformer.h[11].output[0][:] = 0\n",
+        "    output = model.lm_head.output.save()\n",
+        "\n",
+        "print(\"lm_head output = \",output)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "yyzXpfTCn6Hu"
+      },
+      "source": [
+        "The error messaging feature can be toggled using `nnsight.CONFIG.APP.DEBUG` which defaults to true."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JTLOV5fKoIMS"
+      },
+      "outputs": [],
+      "source": [
+        "# Turn off debugging:\n",
+        "import nnsight\n",
+        "\n",
+        "nnsight.CONFIG.APP.DEBUG = False\n",
+        "nnsight.CONFIG.save()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Gl7KqYlUUKut"
+      },
+      "source": [
+        "# .all()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0Zd4oyLAykM_"
+      },
+      "source": [
+        "Sometimes you may want to recursively apply interventions to a model (e.g., when generating many tokens or for models like RNNs, where modules are called multiple times).\n",
+        "\n",
+        "\n",
+        "\n",
+        "*   Calling `.all()` on a model or its submodules will recursively apply its `.input` and `.output` across all iterations.\n",
+        "*   When generating multiple tokens with `.generate` (see: [Multiple Token Generation](https://nnsight.net/notebooks/features/multiple_token/)), using `.all()` before applying an intervention will ensure that the model undergoes the intervention for *all* new tokens generated, not just the first.\n",
+        "\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HTvh-jYt0ylW"
+      },
+      "source": [
+        "## About"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0J6oxQfG0oE4"
+      },
+      "source": [
+        "## .all() now streamlines multiple token generation"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "o8OenySKkk6m"
+      },
+      "source": [
+        "With .all, applying interventions during multiple token generation becomes much easier. Let's test this out!"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WWpVpxNk1DC5"
+      },
+      "source": [
+        "We can use `.all()` to streamline the multiple token generation process. We simply call `.all` on the module where we are applying the intervention (in this case GPT-2's layers), apply our intervention, and append our hidden states (stored in an `nnsight.list()` object)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "oj9iFEWao1Od",
+        "outputId": "ad07c368-ff8e-4b9b-b6df-aa47bc33e168"
+      },
+      "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": [
+            "Hidden state length:  3\n"
+          ]
+        }
+      ],
+      "source": [
+        "# New: using .all():\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() to apply intervention to each new token\n",
+        "    layers.all()\n",
+        "\n",
+        "    # Apply intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(layers[-1].output) # no need to call .save\n",
+        "    # Don't need to loop or call .next()!\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "pkT3YEyr4q5b"
+      },
+      "source": [
+        "Easy! Note that because `.all()` is recursive, it will only work to append outputs called on children of the module that `.all()` was called on. See example below for more information. TL;DR: apply `.all()` on the highest-level accessed module if interventions and outputs have different hierarchies within model structure."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DbaaPivB4Xj_"
+      },
+      "source": [
+        "### Note: (Old method) Applying interventions during multiple token generation without .all()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "U7LNx9aL1O5f"
+      },
+      "source": [
+        "Without `.all()`, we would need to loop across each new generated token, saving the intervention for every generated token and calling `.next()` to move forward."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Ti0nvpuI5KfS",
+        "outputId": "bdc00ff0-8980-41ad-ba2c-1bba317ae963"
+      },
+      "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": [
+            "Hidden state length:  3\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Old approach:\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "hidden_states = []\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    for i in range(n_new_tokens):\n",
+        "        # Apply intervention - set first layer output to zero\n",
+        "        layers[0].output[0][:] = 0\n",
+        "\n",
+        "        # Append desired hidden state post-intervention\n",
+        "        hidden_states.append(layers[-1].output.save())\n",
+        "\n",
+        "        # Move to next generated token\n",
+        "        layers[0].next()\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "U5dmk3dj68dV"
+      },
+      "source": [
+        "### Note: .all() recursive properties"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wGxmulNr0_WE"
+      },
+      "source": [
+        "As mentioned, `.all()` is recursive and will work to append outputs called on children of the module that `.all` was called on. In this example, calling `.all()` on the model's layer modules will not recursively affect `model.lm_head.output` as it is not a child of layers."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 8,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "pLwXPnW-5unE",
+        "outputId": "8af09e06-c563-4021-e47b-2e98f6c4ff98"
+      },
+      "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": [
+            "Hidden state length:  1\n"
+          ]
+        }
+      ],
+      "source": [
+        "# A note on .all() recursive properties:\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() on layers\n",
+        "    layers.all()\n",
+        "\n",
+        "    # Apply same intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(model.lm_head.output) # no need to call .save, it's already initialized\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states)) # length is 1, meaning it only saved the first token generation"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "IS2hNjEb6gd_"
+      },
+      "source": [
+        "So, if you want to apply an intervention during multiple token generation while saving the state of a model component that isn't a child of that module, you can apply .`all()` to the full model."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 9,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "fNWHFuCg6lpC",
+        "outputId": "ab8cb8b1-c056-4d3f-9fa7-4a37e613e868"
+      },
+      "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": [
+            "Hidden state length:  3\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Applying .all() to model fixes issue\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() on model\n",
+        "    model.all()\n",
+        "\n",
+        "    # Apply same intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(model.lm_head.output) # no need to call .save\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states)) # length is 3!"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TrAvx6K0DscN"
+      },
+      "source": [
+        "## Known Issues: `.all()`\n",
+        "\n",
+        "* IteratorEnvoy contexts can produce undesired behavior for subsequent operations defined <u>below</u> it that are not dependent on InterventionProxys.\n",
+        "\n",
+        "Example:\n",
+        "```\n",
+        "with lm.generate(\"Hello World!\", max_new_tokens=10):\n",
+        "    hs_4 = nnsight.list().save()\n",
+        "\n",
+        "    with lm.transformer.h[4].all():\n",
+        "        hs_4.append(lm.transformer.h[4].output)\n",
+        "\n",
+        "    hs_4.append(433)\n",
+        "\n",
+        "print(len(hs_4))\n",
+        "```\n",
+        "`>>> 20 # expected: 11`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "nuNMJsOdzrpF"
+      },
+      "source": [
+        "# Syntax updates"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "mJAp8_Pmj3kP"
+      },
+      "source": [
+        "With `nnsight 0.4`, we now support `if` statements and `for` loops applied to proxies with traditional Python syntax! We also remove the need to call `.value` on a proxy output."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "o3JWZkgL-7Uj"
+      },
+      "source": [
+        "## If Statements\n",
+        "\n",
+        "Previously, we would need to use `.cond()` to create a conditional context that would only execute upon meeting the logical conditions inside the `.cond()`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "cBPhM0vUj3JI",
+        "outputId": "e1c6f677-5251-4b9c-809a-c47add938906"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Random Integer  tensor([5])  is Odd\n"
+          ]
+        }
+      ],
+      "source": [
+        "import torch\n",
+        "# Old method\n",
+        "# model = LanguageModel('openai-community/gpt2', device_map='auto')\n",
+        "\n",
+        "with model.trace(\"The Eiffel Tower is in the city of\") as tracer:\n",
+        "\n",
+        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
+        "\n",
+        "  # To create the conditional context you need to put the\n",
+        "  # condition within tracer.cond()\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": {
+        "id": "B0rRtAYR_AIn"
+      },
+      "source": [
+        "Now, we can use Python `if` statements within the tracing context to create a conditional context!\n",
+        "\n",
+        "*Note: Colab may be a little strangely with this feature the first time you run it - expect some lagging and warnings.*\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 12,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "UHBoLGZt--jC",
+        "outputId": "5e4c7121-4db5-49d9-fd6a-541a27c5318c"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Random Integer  tensor([3])  is Odd\n"
+          ]
+        }
+      ],
+      "source": [
+        "with model.trace(\"The Eiffel Tower is in the city of\") as tracer:\n",
+        "\n",
+        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
+        "\n",
+        "  # Since this if statement is inside the tracing context the if will\n",
+        "  # create a conditional context and will only execute the intervention\n",
+        "  # if this condition is met\n",
+        "  if rand_int % 2 == 0:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
+        "\n",
+        "  if rand_int % 2 == 1:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "reSSkXLe71og"
+      },
+      "source": [
+        "Note: If the conditional statements are outside the tracing context, `if` operates as in base Python.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kupHJGa9v_o5"
+      },
+      "source": [
+        "`elif` statements should also work as `if` statements within the tracing context:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "4k5ubt7Fv-wY",
+        "outputId": "074349f1-6201-49cf-c73c-96c9cb751d21"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Random Integer  tensor([-9])  is Odd\n"
+          ]
+        }
+      ],
+      "source": [
+        "with model.trace(\"The Eiffel Tower is in the city of\") as tracer:\n",
+        "\n",
+        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
+        "\n",
+        "  # Since this if statement is inside the tracing context the if will\n",
+        "  # create a conditional context and will only execute the intervention\n",
+        "  # if this condition is met\n",
+        "  if rand_int % 2 == 0:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
+        "  elif rand_int % 2 == 1:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "A8dxj5cYRHjh"
+      },
+      "source": [
+        "## For Loops"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xC5xw_rb3YXJ"
+      },
+      "source": [
+        "With `nnsight 0.4`, you can now use `for` loops within a tracer context at scale. Previously, a `for` loop within a tracer context inside it resulted in creating intervention graphs over and over for each iteration - this is not scalable!"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9gui9S8J-HZK"
+      },
+      "source": [
+        "The `session.iter` context allows for scalable looping within sessions, but doesn't utilize traditional Python syntax:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 13,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "VHglWqsJSRzv",
+        "outputId": "f5adc897-b2c9-4fcc-ec48-042976a124f1"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Old Method\n",
+        "with model.session() as session:\n",
+        "\n",
+        "  li = nnsight.list() # an NNsight built-in list object\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": {
+        "id": "SOTVCf6lBS9R"
+      },
+      "source": [
+        "Now, you can use simple `for` loops within a tracer context to run an intervention loop at scale.\n",
+        "\n",
+        "*NOTE: inline for loops (i.e., `[x for x in <Proxy object>]`) are not currently supported.*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 14,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Jnpp5VhdTO--",
+        "outputId": "b7806e3c-0065-4307-a22e-4d93b9003738"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        }
+      ],
+      "source": [
+        "# New: Using Python for loops for iterative interventions\n",
+        "with model.session() as session:\n",
+        "\n",
+        "    li = nnsight.list()\n",
+        "    [li.append([num]) for num in range(0, 3)]\n",
+        "    li2 = nnsight.list().save()\n",
+        "\n",
+        "    # Using regular for loops\n",
+        "    for item in li:\n",
+        "        for item_2 in item: # for loops can be nested!\n",
+        "            li2.append(item_2)\n",
+        "\n",
+        "print(\"\\nList: \", li2)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EpGXEcHM0Bs5"
+      },
+      "source": [
+        "## `.value` injected into saved results\n",
+        "\n",
+        "Previously, directly using non-traceable functions (i.e., tokenizers) on a proxy returned from a tracing context required calling `.value` to access the proxy's numerical value. Calling traceable functions (like `print()` or `.argmax()`) on such proxies automatically returned the `.value`, making it optional to call `.value` in certain cases.\n",
+        "\n",
+        "```\n",
+        "input = \"The Eiffel Tower is in the city of\"\n",
+        "with model.trace(input):\n",
+        "\n",
+        "    l2_input = model.transformer.h[2].input.save()\n",
+        "\n",
+        "print(l2_input.value) # could optionally call .value\n",
+        "print(l2_input) # but not required for traceable functions\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6Sn1H93e-8Pd"
+      },
+      "source": [
+        "Now with `nnsight 0.4`, the proxy's value is automatically injected into the variable name, negating any needs to call `.value` on proxies. Proxy variables will automatically be populated with their value upon exiting the tracing context. This is a breaking change, and calling `.value` on a proxy will now throw an error."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 15,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 323
+        },
+        "id": "_bYVfDL80Fn5",
+        "outputId": "2236bd80-2076-4377-cc7d-63aa4b7a88fb"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "tensor([[[ 0.0386, -1.1676,  1.1246,  ..., -1.4047, -0.5742, -0.0668],\n",
+            "         [ 0.1477, -0.4208,  1.3827,  ..., -1.6436, -0.1738, -1.1567],\n",
+            "         [-0.1181, -0.5914, -0.9923,  ..., -0.8742, -0.1361,  1.0608],\n",
+            "         ...,\n",
+            "         [-1.3966,  0.8859,  0.1767,  ...,  0.0661,  0.6106, -0.6092],\n",
+            "         [-2.9199,  0.3945, -3.4569,  ...,  1.2592,  0.2188, -0.5957],\n",
+            "         [ 1.8395,  0.9940,  0.6617,  ..., -0.2802,  0.4978,  0.2308]]],\n",
+            "       device='cuda:0', grad_fn=<AddBackward0>)\n"
+          ]
+        },
+        {
+          "ename": "AttributeError",
+          "evalue": "'Tensor' object has no attribute 'value'",
+          "output_type": "error",
+          "traceback": [
+            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+            "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+            "\u001b[0;32m<ipython-input-15-02b94105ae06>\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ml2_input\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# no need to call .value\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ml2_input\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# will throw an error\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+            "\u001b[0;31mAttributeError\u001b[0m: 'Tensor' object has no attribute 'value'"
+          ]
+        }
+      ],
+      "source": [
+        "input = \"The Eiffel Tower is in the city of\"\n",
+        "with model.trace(input):\n",
+        "\n",
+        "    l2_input = model.transformer.h[2].input.save()\n",
+        "\n",
+        "print(l2_input) # no need to call .value\n",
+        "print(l2_input.value) # will throw an error"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "D9vGO0r968EC"
+      },
+      "source": [
+        "## Turning off syntactic changes"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KBzJBKN06qB7"
+      },
+      "source": [
+        "If you would like to turn off either the `if`/`for` functionality or the `.value` syntactic changes, you can apply the following changes to `nnsight.CONFIG`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM0ZFaRa65KX"
+      },
+      "outputs": [],
+      "source": [
+        "# Turn off if/for statements within tracing context:\n",
+        "import nnsight\n",
+        "\n",
+        "nnsight.CONFIG.APP.CONTROL_FLOW_HANDLING = False\n",
+        "nnsight.CONFIG.save()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "karGT1qv6qYr"
+      },
+      "outputs": [],
+      "source": [
+        "# Turn off .value injection:\n",
+        "import nnsight\n",
+        "\n",
+        "nnsight.CONFIG.APP.FRAME_INJECTION = False\n",
+        "nnsight.CONFIG.save()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "oyuEIDVJDDOi"
+      },
+      "source": [
+        "## Known Issues: Syntax Update\n",
+        "* Colab behaves a little strangely with these features the first time you run it - expect some lagging and warnings.\n",
+        "\n",
+        "* Inline Control Flow (for loops) are not supported.\n",
+        "\n",
+        "Example:\n",
+        "```\n",
+        "with lm.trace(\"Hello World!\"):\n",
+        "    foo = nnsight.list([0, 1, 2, 3]).save()\n",
+        "    [nnsight.log(item) for item in foo]\n",
+        "\n",
+        "```\n",
+        "`>>> Error`\n",
+        "\n",
+        "* Value Injection is not supported for proxies referenced within objects.\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RNliY7nF1uPT"
+      },
+      "source": [
+        "# vLLM Integration"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5tBDA0US7HAg"
+      },
+      "source": [
+        "Our new update includes support for vLLM models using `nnsight`. [vLLM](https://github.com/vllm-project/vllm) is a popular library used for fast inference. By leveraging PagedAttention, dynamic batching, and Hugging Face model integration, vLLM makes inference more efficient and scalable for real-world applications."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xFGg4be_eJVx"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OVEQeIdhxOkG"
+      },
+      "source": [
+        "You will need to install `nnsight 0.4`, `vllm`, and `triton 3.1.0` to use vLLM with NNsight."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 16,
+      "metadata": {
+        "collapsed": true,
+        "id": "3MG0_tXp_JI5"
+      },
+      "outputs": [],
+      "source": [
+        "from IPython.display import clear_output\n",
+        "# install vllm\n",
+        "!pip install vllm==0.6.4.post1\n",
+        "\n",
+        "# install triton 3.1.0\n",
+        "!pip install triton==3.1.0\n",
+        "\n",
+        "clear_output()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "lBFb58M0tZkN"
+      },
+      "source": [
+        "**NOTE: you may need to restart your Colab session before the following step to properly load the `VLLM` model wrapper.**"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cO0g7DRutzUp"
+      },
+      "source": [
+        " Next, let's load in our NNsight-supported vLLM model. You can find vLLM-supported models [here](https://docs.vllm.ai/en/latest/models/supported_models.html). For this exercise, we will use GPT-2."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9FHZysgYeTIU"
+      },
+      "outputs": [],
+      "source": [
+        "from nnsight.modeling.vllm import VLLM\n",
+        "\n",
+        "# NNsight's VLLM wrapper currently supports \"device = cuda\" and device = \"auto\"\n",
+        "vllm = VLLM(\"gpt2\", device = \"auto\", dispatch = True) # See supported models: https://docs.vllm.ai/en/latest/models/supported_models.html\n",
+        "print(vllm)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hrYn5-d8uJZb"
+      },
+      "source": [
+        "## Interventions on vLLM models\n",
+        "We now have a vLLM model that runs with `nnsight`. Let's try applying some interventions on it.\n",
+        "\n",
+        "Note that vLLM takes in sampling parameters including `temperature` and `top_p`. These parameters can be included in the `.trace()` or `.invoke()` contexts. For default model behavior, set `temperature = 0` and `top_p = 1`. For more information about parameters, reference the [vLLM documentation](https://docs.vllm.ai/en/latest/dev/sampling_params.html)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "RjSVuKBVue-m",
+        "outputId": "b96e5d67-9fb8-4330-9f3b-d85f7a23ab45"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Processed prompts: 100%|██████████| 2/2 [00:00<00:00, 38.35it/s, est. speed input: 423.14 toks/s, output: 38.46 toks/s]\n"
+          ]
+        }
+      ],
+      "source": [
+        "with vllm.trace(temperature=0.0, top_p=1.0, max_tokens=1) as tracer:\n",
+        "  with tracer.invoke(\"The Eiffel Tower is located in the city of\"):\n",
+        "    clean_logits = vllm.logits.output.save()\n",
+        "\n",
+        "  with tracer.invoke(\"The Eiffel Tower is located in the city of\"):\n",
+        "    vllm.transformer.h[-2].mlp.output[:] = 0\n",
+        "    corrupted_logits = vllm.logits.output.save()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "xMORJaB2uj7p",
+        "outputId": "744f0c11-7529-4969-9cf1-b5c55da25dd1"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "CLEAN - The Eiffel Tower is located in the city of  Paris\n",
+            "CORRUPTED - The Eiffel Tower is located in the city of  London\n"
+          ]
+        }
+      ],
+      "source": [
+        "print(\"CLEAN - The Eiffel Tower is located in the city of\", vllm.tokenizer.decode(clean_logits.argmax(dim=-1)))\n",
+        "print(\"CORRUPTED - The Eiffel Tower is located in the city of\", vllm.tokenizer.decode(corrupted_logits.argmax(dim=-1)))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9rLvuOlAv6zc"
+      },
+      "source": [
+        "We've successfully performed an intervention on our vLLM model!"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XpHPuAKTFRb3"
+      },
+      "source": [
+        "## Sampled Token Traceability\n",
+        "vLLM provides functionality to configure how each sequence samples its next token. Here's an example of how you can trace token sampling operations with the nnsight VLLM wrapper."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "xnSkk-GeFuNf",
+        "outputId": "0b38ca8f-a978-4c3e-beea-f7ec4033b0fe"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Processed prompts: 100%|██████████| 1/1 [00:00<00:00,  5.84it/s, est. speed input: 53.99 toks/s, output: 17.99 toks/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Samples:  [tensor([10346]), tensor([13]), tensor([6363])]\n",
+            "Logits:  [tensor([[-109.0625, -107.9375, -111.6875,  ..., -115.3750, -116.5625,\n",
+            "         -108.8750]], device='cuda:0', dtype=torch.float16), tensor([[-80.3125, -82.2500, -85.3750,  ..., -93.3750, -90.7500, -83.4375]],\n",
+            "       device='cuda:0', dtype=torch.float16), tensor([[-110.0000, -109.3125, -110.4375,  ..., -120.9375, -119.3750,\n",
+            "         -101.3750]], device='cuda:0', dtype=torch.float16)]\n"
+          ]
+        }
+      ],
+      "source": [
+        "import nnsight\n",
+        "with vllm.trace(\"Madison Square Garden is located in the city of\", temperature=0.8, top_p=0.95, max_tokens=3) as tracer:\n",
+        "    samples = nnsight.list().save()\n",
+        "    logits = nnsight.list().save()\n",
+        "\n",
+        "    for ii in range(3):\n",
+        "        samples.append(vllm.samples.output)\n",
+        "        vllm.samples.next()\n",
+        "        logits.append(vllm.logits.output)\n",
+        "        vllm.logits.next()\n",
+        "\n",
+        "print(\"Samples: \", samples)\n",
+        "print(\"Logits: \", logits) # different than samples with current sampling parameters"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aL5y8M16wHsL"
+      },
+      "source": [
+        "## Note: gradients are not supported with vLLM\n",
+        "vLLM speeds up inference through its paged attention mechanism. This means that accessing gradients and backward passes are not supported for vLLM models. As such, calling gradient operations when using `nnsight` vLLM wrappers will throw an error."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aiNgpkdGDdBO"
+      },
+      "source": [
+        "## Known Issues: vLLM Integration\n",
+        "* The vllm.LLM engine performs max_tokens + 1 forward passes which can lead to undesired behavior if you are running interventions on all iterations of multi-token generation.\n",
+        "\n",
+        "Example:\n",
+        "```\n",
+        "with vllm_gpt2(\"Hello World!\", max_tokens=10):\n",
+        "    logits = nnsight.list().save()\n",
+        "    with vllm_gpt2.logits.all():\n",
+        "        logits.append(vllm_gpt2.logits.output)\n",
+        "\n",
+        "print(len(logits))\n",
+        "\n",
+        "```\n",
+        "`>>> 11 # expected: 10`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xtSCr95WUN_s"
+      },
+      "source": [
+        "# Streaming"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eV1H_LzRU0eo"
+      },
+      "source": [
+        "Streaming enables users apply functions and datasets locally during remote model execution. This allows users to stream results for immediate consumption (i.e., seeing tokens as they are generated) or applying non-whitelisted functions such as model tokenizers, large local datasets, and more!\n",
+        "\n",
+        "*   `nnsight.local()` context sends values immediately to user's local machine from server\n",
+        "*   Intervention graph is executed locally on downstream nodes\n",
+        "*   Exiting local context uploads data back to server\n",
+        "*   `@nnsight.trace` function decorator enables custom functions to be added to intervention graph when using `nnsight.local()`\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "SLMF9D6KXVin"
+      },
+      "source": [
+        "## `nnsight.local()`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "LhCNjtGl-oec"
+      },
+      "source": [
+        "You may sometimes want to locally access and manipulate values during remote execution. Using `.local()` on a proxy, you can send remote content to your local machine and apply local functions. The intervention graph is then executed locally on downstream nodes until you exit the local context.\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e3ZRmJi__uXE"
+      },
+      "source": [
+        "There are a few use cases for streaming with `.local()`, including live chat generation and applying large datasets or non-whitelisted local functions to the intervention graph.\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "C7z4qaneMH7L"
+      },
+      "source": [
+        "Now let's explore how streaming works. We'll start by grabbing some hidden states of the model and printing their value using `tracer.log()`. Without calling `nnsight.local()`, these operations will all occur remotely."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PRSITdCKRBlq"
+      },
+      "outputs": [],
+      "source": [
+        "from nnsight import LanguageModel\n",
+        "\n",
+        "llama = LanguageModel(\"meta-llama/Meta-Llama-3.1-8B\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "hRA8MkG5MH7L"
+      },
+      "outputs": [],
+      "source": [
+        "# This will give you a remote LOG response because it's coming from the remote server\n",
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
+        "\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
+        "\n",
+        "    tracer.log(hs[0,0,0])\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "06nlVWNmMH7M"
+      },
+      "source": [
+        "Now, let's try the same operation using the `nnsight.local()` context. This will send the operations to get and print the hidden states to your local machine, changing how the logging message is formatted (local formatting instead of remote)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "H958j5eEMH7M"
+      },
+      "outputs": [],
+      "source": [
+        "# This will print locally because it's already local\n",
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
+        "\n",
+        "    with nnsight.local():\n",
+        "        hs = llama.model.layers[-1].output[0]\n",
+        "        tracer.log(hs[0,0,0])\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "qA347WLnyexi"
+      },
+      "source": [
+        "## `@nnsight.trace` function decorator"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bTWtoGZzykkV"
+      },
+      "source": [
+        "We can also use function decorators to create custom functions to be used during `.local` calls. This is a handy way to enable live streaming of a chat or to train probing classifiers on model hidden states."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2PhO5owrym1J"
+      },
+      "source": [
+        "Let's try out `@nnsight.trace` and `nnsight.local()` to access a custom function during remote execution."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "DsIm3De-MH7M"
+      },
+      "outputs": [],
+      "source": [
+        "# first, let's define our function\n",
+        "@nnsight.trace # decorator that enables this function to be added to the intervention graph\n",
+        "def my_local_fn(value):\n",
+        "    return value * 0\n",
+        "\n",
+        "# We use a local function to ablate some hidden states\n",
+        "# This downloads the data for the .local context, and then uploads it back to set the value.\n",
+        "with llama.generate(\"hello\", remote=True) as tracer:\n",
+        "\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
+        "\n",
+        "    with nnsight.local():\n",
+        "\n",
+        "        hs = my_local_fn(hs)\n",
+        "\n",
+        "    llama.model.layers[-1].output[0][:] = hs\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vYC1HZxOMH7M"
+      },
+      "source": [
+        "Note that without calling `.local`, the remote API does not know about `my_local_fn` and will throw a whitelist error. A whitelist error occurs because you are being allowed access to the function."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "W7vNbMLrMH7M"
+      },
+      "outputs": [],
+      "source": [
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
+        "\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
+        "\n",
+        "    hs = my_local_fn(hs) # no .local - will cause an error\n",
+        "\n",
+        "    llama.model.layers[-1].output[0][:] = hs * 2\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "M-aZ3YzOAiNr"
+      },
+      "source": [
+        "## Example: Live-streaming remote chat\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "I90m5IuuMH7M"
+      },
+      "source": [
+        "Now that we can access data within the tracing context on our local computer, we can apply non-whitelisted functions, such as the model's tokenizer, within our tracing context.\n",
+        "\n",
+        "Let's build a decoding function that will decode tokens into words and print the result."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0oub4XlnMH7M"
+      },
+      "outputs": [],
+      "source": [
+        "@nnsight.trace\n",
+        "def my_decoding_function(tokens, model, max_length=80, state=None):\n",
+        "    # Initialize state if not provided\n",
+        "    if state is None:\n",
+        "        state = {'current_line': '', 'current_line_length': 0}\n",
+        "\n",
+        "    token = tokens[-1] # only use last token\n",
+        "\n",
+        "    # Decode the token\n",
+        "    decoded_token = llama.tokenizer.decode(token).encode(\"unicode_escape\").decode()\n",
+        "\n",
+        "    if decoded_token == '\\\\n':  # Handle explicit newline tokens\n",
+        "        # Print the current line and reset state\n",
+        "        print('',flush=True)\n",
+        "        state['current_line'] = ''\n",
+        "        state['current_line_length'] = 0\n",
+        "    else:\n",
+        "        # Check if adding the token would exceed the max length\n",
+        "        if state['current_line_length'] + len(decoded_token) > max_length:\n",
+        "            print('',flush=True)\n",
+        "            state['current_line'] = decoded_token  # Start a new line with the current token\n",
+        "            state['current_line_length'] = len(decoded_token)\n",
+        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
+        "        else:\n",
+        "            # Add a space if the line isn't empty and append the token\n",
+        "            if state['current_line']:\n",
+        "                state['current_line'] += decoded_token\n",
+        "            else:\n",
+        "                state['current_line'] = decoded_token\n",
+        "            state['current_line_length'] += len(decoded_token)\n",
+        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
+        "\n",
+        "    return state"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2qbXSwMGOhd-"
+      },
+      "source": [
+        "Now we can decode and print our model outputs throughout token generation by accessing our decoding function through `nnsight.local()`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "jzjWIFbgMH7M"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "\n",
+        "nnsight.CONFIG.APP.REMOTE_LOGGING = False\n",
+        "\n",
+        "prompt = \"A press release is an official statement delivered to members of the news media for the purpose of\"\n",
+        "# prompt = \"Your favorite board game is\"\n",
+        "\n",
+        "print(\"Prompt: \",prompt,'\\n', end =\"\")\n",
+        "\n",
+        "# Initialize the state for decoding\n",
+        "state = {'current_line': '', 'current_line_length': 0}\n",
+        "\n",
+        "with llama.generate(prompt, remote=True, max_new_tokens = 50) as generator:\n",
+        "    # Call .all() to apply to each new token\n",
+        "    llama.all()\n",
+        "\n",
+        "    all_tokens = nnsight.list().save()\n",
+        "\n",
+        "    # Access model output\n",
+        "    out = llama.lm_head.output.save()\n",
+        "\n",
+        "    # Apply softmax to obtain probabilities and save the result\n",
+        "    probs = torch.nn.functional.softmax(out, dim=-1)\n",
+        "    max_probs = torch.max(probs, dim=-1)\n",
+        "    tokens = max_probs.indices.cpu().tolist()\n",
+        "    all_tokens.append(tokens[0]).save()\n",
+        "\n",
+        "    with nnsight.local():\n",
+        "        state = my_decoding_function(tokens[0], llama, max_length=20, state=state)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JY8p_7Zj81B4"
+      },
+      "source": [
+        "# General Considerations\n",
+        "* `Tracer.cond(…)` and `Tracer.iter(…)` are still supported.\n",
+        "\n",
+        "* vLLM <U>does not</u> come as a pre-installed dependency of `nnsight`.\n",
+        "\n",
+        "* `nnsight` supports `vllm==0.6.4.post1`\n",
+        "\n",
+        "* vLLM support only includes `cuda` and `auto` devices at the moment.\n",
+        "\n",
+        "* vLLM models <u>do not</u> support gradients.\n",
+        "\n",
+        "* The `@nnsight.trace` decorator does not enable user-defined operations to be executed remotely. Something coming soon for that..."
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "DbaaPivB4Xj_"
+      ],
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "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.10.13"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "0325fe95b3e3432497fc68317b2126d4": {
+          "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
+          }
+        },
+        "0650393001b34123a0e18663e6aa0fcb": {
+          "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
+          }
+        },
+        "0e4b52697a87475882e40b3ca4a2e77d": {
+          "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_60627b23fe934d01b4a4af1f592fbd55",
+            "max": 26,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_a806062fe502459a93ee212e57af2e66",
+            "value": 26
+          }
+        },
+        "0f3eaadfc2c34384afd2979417d94812": {
+          "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_a2eafc8f9262412793f35b372eb300eb",
+              "IPY_MODEL_cb3a1a3578554cac9cec37d9b0390c1a",
+              "IPY_MODEL_17f30f75790e4273b3e475f433c4760d"
+            ],
+            "layout": "IPY_MODEL_1046b0563cb64ce68fc7b1b2f5d3d65d"
+          }
+        },
+        "1046b0563cb64ce68fc7b1b2f5d3d65d": {
+          "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
+          }
+        },
+        "15332cc2b6e84cfe937eece4ec884a72": {
+          "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": ""
+          }
+        },
+        "16f914ebda5b44d28cd6aecc5ae742b7": {
+          "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": ""
+          }
+        },
+        "176749c3ecd24b829a78f892dee8af91": {
+          "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_0325fe95b3e3432497fc68317b2126d4",
+            "placeholder": "​",
+            "style": "IPY_MODEL_fabbb16b5e7942f493b7dc92279715fd",
+            "value": " 1.36M/1.36M [00:00&lt;00:00, 15.2MB/s]"
+          }
+        },
+        "17f30f75790e4273b3e475f433c4760d": {
+          "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_2e14c72eaf6644bd9d656513c302c278",
+            "placeholder": "​",
+            "style": "IPY_MODEL_30af0b2c974d418cb40b30403a2a90cf",
+            "value": " 665/665 [00:00&lt;00:00, 39.8kB/s]"
+          }
+        },
+        "1ac5508e75154bd18275d6f6775242b7": {
+          "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_a22aff0a7b33470b970a32269a9a0f08",
+            "placeholder": "​",
+            "style": "IPY_MODEL_be8e4e0c8e0c482897194f08c77f82a0",
+            "value": " 456k/456k [00:00&lt;00:00, 2.80MB/s]"
+          }
+        },
+        "1c6092db09c04e568b23ca71eb3f5f13": {
+          "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_dea619854cbb42a2a700f7feef77d442",
+            "placeholder": "​",
+            "style": "IPY_MODEL_c66f56814daa4745a0f14c75f85d0294",
+            "value": "model.safetensors: 100%"
+          }
+        },
+        "2130a5e5fd6d44ce8484c51d591e2d6e": {
+          "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
+          }
+        },
+        "259617ed81dc4dd4b40f657ad63bec47": {
+          "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
+          }
+        },
+        "25af4170e72745a095f37ef4a481fab6": {
+          "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": ""
+          }
+        },
+        "2e14c72eaf6644bd9d656513c302c278": {
+          "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
+          }
+        },
+        "2f851281cd2448f8a36c9b538931e0c9": {
+          "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_637e15c89e6847098b27f7a4f3482a63",
+            "placeholder": "​",
+            "style": "IPY_MODEL_e039919f976f4d5fbebb5117ec53d57e",
+            "value": "tokenizer.json: 100%"
+          }
+        },
+        "30af0b2c974d418cb40b30403a2a90cf": {
+          "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": ""
+          }
+        },
+        "3160c45575ad4426b2abcfaccd27d41c": {
+          "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
+          }
+        },
+        "33c1fe5504d1441e805a02fb4693acd7": {
+          "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_8bf1249ac7aa46b094bc899d432153fa",
+              "IPY_MODEL_c60da3ab273d4439a0baec6ae7c216e4",
+              "IPY_MODEL_9ef4dc73855c4fa69656f75110ac390d"
+            ],
+            "layout": "IPY_MODEL_4b1d8724a6b247e0a3dafe2f6b9de178"
+          }
+        },
+        "33cc1f2b3b884a65a4a942cde1662b8f": {
+          "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
+          }
+        },
+        "3adee7a17b2a4f50b9ff36fe10a3b701": {
+          "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": ""
+          }
+        },
+        "3b2050ce41594c69bf6efbacc76a574a": {
+          "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
+          }
+        },
+        "3c1c03b756c1401da717b3b62842a1f9": {
+          "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
+          }
+        },
+        "416faee108c44dce9e1b3527e82acb9b": {
+          "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": ""
+          }
+        },
+        "45f3295ab6e64cbda41f26c2d1581b08": {
+          "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_3b2050ce41594c69bf6efbacc76a574a",
+            "placeholder": "​",
+            "style": "IPY_MODEL_15332cc2b6e84cfe937eece4ec884a72",
+            "value": "generation_config.json: 100%"
+          }
+        },
+        "4b1d8724a6b247e0a3dafe2f6b9de178": {
+          "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
+          }
+        },
+        "4d7f03cfa9944554b093b3e5c691db92": {
+          "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
+          }
+        },
+        "4eb7af991ebe4bdd855d3cdd50184495": {
+          "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
+          }
+        },
+        "4fc680f767b04af881b52fd94ea07bba": {
+          "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_45f3295ab6e64cbda41f26c2d1581b08",
+              "IPY_MODEL_78018bc8cd974345b090ad38214e75c6",
+              "IPY_MODEL_56806395ddfe45f2b0245741a41e5bf5"
+            ],
+            "layout": "IPY_MODEL_87e140a2c262488da5b393f37b51592c"
+          }
+        },
+        "53f9e552ea174246a10de4afc84272ac": {
+          "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": ""
+          }
+        },
+        "54fc39e52ff94065a20594cc95a35942": {
+          "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_5e25783cecd346f48959de00bc2cccec",
+            "max": 548105171,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_c700fe2789cb4f51a43ae1e29932d3fc",
+            "value": 548105171
+          }
+        },
+        "56806395ddfe45f2b0245741a41e5bf5": {
+          "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_2130a5e5fd6d44ce8484c51d591e2d6e",
+            "placeholder": "​",
+            "style": "IPY_MODEL_25af4170e72745a095f37ef4a481fab6",
+            "value": " 124/124 [00:00&lt;00:00, 7.15kB/s]"
+          }
+        },
+        "581c1d23c4424ce6b50aec079df134c4": {
+          "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_e122506e7c7c454c952b89f306360da6",
+            "placeholder": "​",
+            "style": "IPY_MODEL_53f9e552ea174246a10de4afc84272ac",
+            "value": " 548M/548M [00:02&lt;00:00, 220MB/s]"
+          }
+        },
+        "59f8564845764c9e9891274216be83ff": {
+          "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
+          }
+        },
+        "5c5d0fc4f29141fc95008070c0b077d3": {
+          "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_4d7f03cfa9944554b093b3e5c691db92",
+            "placeholder": "​",
+            "style": "IPY_MODEL_732fe1ca67b74026805de4688d81e6e3",
+            "value": "merges.txt: 100%"
+          }
+        },
+        "5e25783cecd346f48959de00bc2cccec": {
+          "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
+          }
+        },
+        "60627b23fe934d01b4a4af1f592fbd55": {
+          "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
+          }
+        },
+        "637e15c89e6847098b27f7a4f3482a63": {
+          "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
+          }
+        },
+        "6386746d591e4e4cb83a18dcc4481678": {
+          "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": ""
+          }
+        },
+        "63dc1c0ce98d4539984bfc93114aa6c6": {
+          "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_1c6092db09c04e568b23ca71eb3f5f13",
+              "IPY_MODEL_54fc39e52ff94065a20594cc95a35942",
+              "IPY_MODEL_581c1d23c4424ce6b50aec079df134c4"
+            ],
+            "layout": "IPY_MODEL_c7c13deb642f4af1bedd88c2afac85fd"
+          }
+        },
+        "6909cf23963d4347820b3717b62b1cfe": {
+          "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
+          }
+        },
+        "6fc349dcab8d43309faaa08775f445ff": {
+          "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
+          }
+        },
+        "732fe1ca67b74026805de4688d81e6e3": {
+          "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": ""
+          }
+        },
+        "78018bc8cd974345b090ad38214e75c6": {
+          "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_259617ed81dc4dd4b40f657ad63bec47",
+            "max": 124,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_16f914ebda5b44d28cd6aecc5ae742b7",
+            "value": 124
+          }
+        },
+        "7c3d6693887b4c29b3d63afd37e4292e": {
+          "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
+          }
+        },
+        "87e140a2c262488da5b393f37b51592c": {
+          "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
+          }
+        },
+        "886f12cb69c2435d978eb2d33534763a": {
+          "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
+          }
+        },
+        "8bf1249ac7aa46b094bc899d432153fa": {
+          "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_6909cf23963d4347820b3717b62b1cfe",
+            "placeholder": "​",
+            "style": "IPY_MODEL_f6a764157eab4ee7b24ce3dc3e317222",
+            "value": "vocab.json: 100%"
+          }
+        },
+        "9a2747e522c94c00af3f9d0bbd54cc1d": {
+          "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_7c3d6693887b4c29b3d63afd37e4292e",
+            "max": 1355256,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_6386746d591e4e4cb83a18dcc4481678",
+            "value": 1355256
+          }
+        },
+        "9b161849be8a4464a486e6270f3c0f3b": {
+          "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
+          }
+        },
+        "9ef4dc73855c4fa69656f75110ac390d": {
+          "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_e60f4ecadf75463ca5fc5f934954e144",
+            "placeholder": "​",
+            "style": "IPY_MODEL_d5f42dda1612435f8eb053fe29f711c0",
+            "value": " 1.04M/1.04M [00:00&lt;00:00, 3.50MB/s]"
+          }
+        },
+        "a1595afd22fa4d999a651e1ff63f4279": {
+          "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_5c5d0fc4f29141fc95008070c0b077d3",
+              "IPY_MODEL_c410466db4ac4d878c3aac7fe4a0df5b",
+              "IPY_MODEL_1ac5508e75154bd18275d6f6775242b7"
+            ],
+            "layout": "IPY_MODEL_6fc349dcab8d43309faaa08775f445ff"
+          }
+        },
+        "a22aff0a7b33470b970a32269a9a0f08": {
+          "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
+          }
+        },
+        "a2eafc8f9262412793f35b372eb300eb": {
+          "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_886f12cb69c2435d978eb2d33534763a",
+            "placeholder": "​",
+            "style": "IPY_MODEL_cb0ab1cce5364b44ba85e65ccf44d0c9",
+            "value": "config.json: 100%"
+          }
+        },
+        "a806062fe502459a93ee212e57af2e66": {
+          "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": ""
+          }
+        },
+        "be8e4e0c8e0c482897194f08c77f82a0": {
+          "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": ""
+          }
+        },
+        "c3b4c47b73d544eaba67d5828f1f742d": {
+          "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_c924f41345d34f72a00ed3b21f84de82",
+              "IPY_MODEL_0e4b52697a87475882e40b3ca4a2e77d",
+              "IPY_MODEL_c4c1c2ccc4014520bb265d21550fb133"
+            ],
+            "layout": "IPY_MODEL_3160c45575ad4426b2abcfaccd27d41c"
+          }
+        },
+        "c410466db4ac4d878c3aac7fe4a0df5b": {
+          "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_3c1c03b756c1401da717b3b62842a1f9",
+            "max": 456318,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_416faee108c44dce9e1b3527e82acb9b",
+            "value": 456318
+          }
+        },
+        "c4c1c2ccc4014520bb265d21550fb133": {
+          "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_9b161849be8a4464a486e6270f3c0f3b",
+            "placeholder": "​",
+            "style": "IPY_MODEL_f2ca144c6f8d4537ac8e2e29c8109454",
+            "value": " 26.0/26.0 [00:00&lt;00:00, 1.49kB/s]"
+          }
+        },
+        "c60da3ab273d4439a0baec6ae7c216e4": {
+          "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_59f8564845764c9e9891274216be83ff",
+            "max": 1042301,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_3adee7a17b2a4f50b9ff36fe10a3b701",
+            "value": 1042301
+          }
+        },
+        "c66f56814daa4745a0f14c75f85d0294": {
+          "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": ""
+          }
+        },
+        "c700fe2789cb4f51a43ae1e29932d3fc": {
+          "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": ""
+          }
+        },
+        "c7c13deb642f4af1bedd88c2afac85fd": {
+          "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
+          }
+        },
+        "c924f41345d34f72a00ed3b21f84de82": {
+          "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_4eb7af991ebe4bdd855d3cdd50184495",
+            "placeholder": "​",
+            "style": "IPY_MODEL_c96d870993144ed09e54de5dc00fad6d",
+            "value": "tokenizer_config.json: 100%"
+          }
+        },
+        "c96d870993144ed09e54de5dc00fad6d": {
+          "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": ""
+          }
+        },
+        "cb0ab1cce5364b44ba85e65ccf44d0c9": {
+          "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": ""
+          }
+        },
+        "cb3a1a3578554cac9cec37d9b0390c1a": {
+          "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_0650393001b34123a0e18663e6aa0fcb",
+            "max": 665,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_cdfa2b53fdd14d24b9ccc0920f34bd69",
+            "value": 665
+          }
+        },
+        "cdfa2b53fdd14d24b9ccc0920f34bd69": {
+          "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": ""
+          }
+        },
+        "d5f42dda1612435f8eb053fe29f711c0": {
+          "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": ""
+          }
+        },
+        "d86efefd42444836b476a5caecccdb75": {
+          "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_2f851281cd2448f8a36c9b538931e0c9",
+              "IPY_MODEL_9a2747e522c94c00af3f9d0bbd54cc1d",
+              "IPY_MODEL_176749c3ecd24b829a78f892dee8af91"
+            ],
+            "layout": "IPY_MODEL_33cc1f2b3b884a65a4a942cde1662b8f"
+          }
+        },
+        "dea619854cbb42a2a700f7feef77d442": {
+          "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
+          }
+        },
+        "e039919f976f4d5fbebb5117ec53d57e": {
+          "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": ""
+          }
+        },
+        "e122506e7c7c454c952b89f306360da6": {
+          "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
+          }
+        },
+        "e60f4ecadf75463ca5fc5f934954e144": {
+          "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
+          }
+        },
+        "f2ca144c6f8d4537ac8e2e29c8109454": {
+          "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": ""
+          }
+        },
+        "f6a764157eab4ee7b24ce3dc3e317222": {
+          "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": ""
+          }
+        },
+        "fabbb16b5e7942f493b7dc92279715fd": {
+          "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
+}
diff --git a/docs/source/notebooks/tutorials/activation_patching.ipynb b/docs/source/notebooks/tutorials/activation_patching.ipynb
index e4e2401d..1402bb77 100644
--- a/docs/source/notebooks/tutorials/activation_patching.ipynb
+++ b/docs/source/notebooks/tutorials/activation_patching.ipynb
@@ -4381,7 +4381,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.10.14"
+      "version": "3.10.13"
     }
   },
   "nbformat": 4,
diff --git a/docs/source/notebooks/tutorials/attribution_patching.ipynb b/docs/source/notebooks/tutorials/attribution_patching.ipynb
index c1f48859..4e911a94 100644
--- a/docs/source/notebooks/tutorials/attribution_patching.ipynb
+++ b/docs/source/notebooks/tutorials/attribution_patching.ipynb
@@ -58,7 +58,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 23,
       "metadata": {
         "id": "qM0t9saED43y"
       },
@@ -98,7 +98,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 13,
+      "execution_count": 24,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -238,7 +238,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 14,
+      "execution_count": 25,
       "metadata": {
         "id": "bcyK0AoeObY_"
       },
@@ -295,7 +295,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 15,
+      "execution_count": 26,
       "metadata": {
         "id": "05T3hf3LOpFG"
       },
@@ -320,7 +320,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 16,
+      "execution_count": 27,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -367,7 +367,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 17,
+      "execution_count": 28,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -433,7 +433,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 18,
+      "execution_count": 29,
       "metadata": {
         "id": "QdSrOsXvsjjb"
       },
@@ -490,7 +490,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 19,
+      "execution_count": 30,
       "metadata": {
         "id": "cxC8Xs8HsqWm"
       },
@@ -517,7 +517,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 20,
+      "execution_count": 31,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -1667,9 +1667,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"adb1d92b-75d9-40e5-9a23-e6141c8f1a77\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"adb1d92b-75d9-40e5-9a23-e6141c8f1a77\")) {                    Plotly.newPlot(                        \"adb1d92b-75d9-40e5-9a23-e6141c8f1a77\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-6.957763253012672e-05,-6.909317562531214e-09,-4.827359225600958e-05,-1.4337681022880133e-05,6.04442448093323e-06,-5.384786163631361e-06,-0.0002756676112767309,-2.1065989130875096e-05,-1.3769823453912977e-05,8.502480341121554e-06,-8.981670544017106e-05,5.582411904470064e-05],[8.960292325355113e-05,-0.00040236394852399826,-0.00018017760885413736,0.00014379735512193292,-6.819632835686207e-05,-6.25044631306082e-05,-4.3686544813681394e-05,-0.000808828801382333,0.00023929437156766653,0.00015937883290462196,-0.0001314240216743201,9.500579471932724e-05],[1.4101449778536335e-05,-0.00030104274628683925,-0.00037168117705732584,-0.00034691643668338656,-1.0458330507390201e-05,-0.00015961707686074078,-0.000304880493786186,9.84568105195649e-05,0.00026875207549892366,0.0001521407684776932,0.0005039765383116901,-0.0002034870849456638],[-1.1517602615640499e-05,-0.00033706481917761266,-0.00024615979054942727,6.271059100981802e-05,1.354238520434592e-05,0.00034766620956361294,-0.0005866399733349681,-0.0006978895980864763,0.00012641398643609136,-8.654002158436924e-05,-6.751238834112883e-05,0.0002554001403041184],[6.932910764589906e-05,0.00015388225438073277,3.707451105583459e-05,-0.00017302615742664784,9.010283974930644e-05,0.00022782250016462058,-1.3750723155681044e-05,-0.0014550655614584684,-9.826988389249891e-05,0.0005861636018380523,5.1824528782162815e-05,0.0005829725996591151],[1.2706756024272181e-05,-6.604211648664204e-06,0.0010883781360462308,0.0004620166728273034,-0.001068884739652276,-8.976992830866948e-05,4.2892796045634896e-05,-0.00045684745418839157,0.00039122329326346517,9.987624071072787e-05,-0.0018775209318846464,0.0008228498045355082],[0.002574358833953738,0.01968141831457615,-0.0006720172241330147,0.0014327410608530045,-0.0003379757108632475,-0.00020742780179716647,0.0014370887074619532,-0.0007567756110802293,0.0012054562103003263,-0.00020089393365196884,8.37489205878228e-05,4.493692540563643e-05],[-0.00022981743677519262,-0.001819567522034049,-9.668200073065236e-05,0.1371454894542694,-8.344589150510728e-05,-0.009004496037960052,0.0023542074486613274,-0.0004521115333773196,-0.00800877995789051,0.2027597427368164,-0.00011211885430384427,-0.00033840813557617366],[-0.0020026592537760735,-3.5826669773086905e-06,-0.0030021658167243004,-0.0004367050714790821,0.0009537620353512466,-0.0027750192675739527,0.29613298177719116,0.002325658919289708,0.0028100069612264633,0.0005649963277392089,0.27499300241470337,0.0029178475961089134],[0.005365221295505762,-0.0014320836635306478,0.005969846621155739,0.020304212346673012,-0.028698338195681572,0.015889175236225128,-0.05884050577878952,0.057027921080589294,-0.0011342937359586358,0.28792738914489746,0.0005685207433998585,-0.00238024047575891],[0.08650688081979752,0.022697942331433296,0.0005626287311315536,0.007188204675912857,-0.0018677962943911552,0.0002627074427437037,0.04765014350414276,-0.3637135326862335,-0.00011220996384508908,0.00923435389995575,0.053251881152391434,9.949051309376955e-05],[0.007055167108774185,0.0017161732539534569,-0.09743084013462067,0.006113084033131599,-3.512430703267455e-05,-0.0008988302433863282,-0.011853177100419998,0.004282553680241108,0.006519943010061979,-0.007987556979060173,-0.19988438487052917,0.001217787154018879]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Attention Heads\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"f2bf5d96-2759-4fa9-ac62-e9371c361705\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"f2bf5d96-2759-4fa9-ac62-e9371c361705\")) {                    Plotly.newPlot(                        \"f2bf5d96-2759-4fa9-ac62-e9371c361705\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[-6.957763253012672e-05,-6.909317562531214e-09,-4.827359225600958e-05,-1.4337681022880133e-05,6.04442448093323e-06,-5.384786163631361e-06,-0.0002756676112767309,-2.1065989130875096e-05,-1.3769823453912977e-05,8.502480341121554e-06,-8.981670544017106e-05,5.582411904470064e-05],[8.960292325355113e-05,-0.00040236394852399826,-0.00018017760885413736,0.00014379735512193292,-6.819632835686207e-05,-6.25044631306082e-05,-4.3686544813681394e-05,-0.000808828801382333,0.00023929437156766653,0.00015937883290462196,-0.0001314240216743201,9.500579471932724e-05],[1.4101449778536335e-05,-0.00030104274628683925,-0.00037168117705732584,-0.00034691643668338656,-1.0458330507390201e-05,-0.00015961707686074078,-0.000304880493786186,9.84568105195649e-05,0.00026875207549892366,0.0001521407684776932,0.0005039765383116901,-0.0002034870849456638],[-1.1517602615640499e-05,-0.00033706481917761266,-0.00024615979054942727,6.271059100981802e-05,1.354238520434592e-05,0.00034766620956361294,-0.0005866399733349681,-0.0006978895980864763,0.00012641398643609136,-8.654002158436924e-05,-6.751238834112883e-05,0.0002554001403041184],[6.932910764589906e-05,0.00015388225438073277,3.707451105583459e-05,-0.00017302615742664784,9.010283974930644e-05,0.00022782250016462058,-1.3750723155681044e-05,-0.0014550655614584684,-9.826988389249891e-05,0.0005861636018380523,5.1824528782162815e-05,0.0005829725996591151],[1.2706756024272181e-05,-6.604211648664204e-06,0.0010883781360462308,0.0004620166728273034,-0.001068884739652276,-8.976992830866948e-05,4.2892796045634896e-05,-0.00045684745418839157,0.00039122329326346517,9.987624071072787e-05,-0.0018775209318846464,0.0008228498045355082],[0.002574358833953738,0.01968141831457615,-0.0006720172241330147,0.0014327410608530045,-0.0003379757108632475,-0.00020742780179716647,0.0014370887074619532,-0.0007567756110802293,0.0012054562103003263,-0.00020089393365196884,8.37489205878228e-05,4.493692540563643e-05],[-0.00022981743677519262,-0.001819567522034049,-9.668200073065236e-05,0.1371454894542694,-8.344589150510728e-05,-0.009004496037960052,0.0023542074486613274,-0.0004521115333773196,-0.00800877995789051,0.2027597427368164,-0.00011211885430384427,-0.00033840813557617366],[-0.0020026592537760735,-3.5826669773086905e-06,-0.0030021658167243004,-0.0004367050714790821,0.0009537620353512466,-0.0027750192675739527,0.29613298177719116,0.002325658919289708,0.0028100069612264633,0.0005649963277392089,0.27499300241470337,0.0029178475961089134],[0.005365221295505762,-0.0014320836635306478,0.005969846621155739,0.020304212346673012,-0.028698338195681572,0.015889175236225128,-0.05884050577878952,0.057027921080589294,-0.0011342937359586358,0.28792738914489746,0.0005685207433998585,-0.00238024047575891],[0.08650688081979752,0.022697942331433296,0.0005626287311315536,0.007188204675912857,-0.0018677962943911552,0.0002627074427437037,0.04765014350414276,-0.3637135326862335,-0.00011220996384508908,0.00923435389995575,0.053251881152391434,9.949051309376955e-05],[0.007055167108774185,0.0017161732539534569,-0.09743084013462067,0.006113084033131599,-3.512430703267455e-05,-0.0008988302433863282,-0.011853177100419998,0.004282553680241108,0.006519943010061979,-0.007987556979060173,-0.19988438487052917,0.001217787154018879]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Head: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Head\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Attention Heads\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('adb1d92b-75d9-40e5-9a23-e6141c8f1a77');\n",
+              "var gd = document.getElementById('f2bf5d96-2759-4fa9-ac62-e9371c361705');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
@@ -1746,7 +1746,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 21,
+      "execution_count": 32,
       "metadata": {
         "id": "NswiDYEisyw-"
       },
@@ -1771,7 +1771,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 22,
+      "execution_count": 33,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -2892,9 +2892,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"3a0bc859-eccc-46fb-ab21-231253cfa4af\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"3a0bc859-eccc-46fb-ab21-231253cfa4af\")) {                    Plotly.newPlot(                        \"3a0bc859-eccc-46fb-ab21-231253cfa4af\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0006664255633950233,-0.0011972333304584026,-0.0003045523480977863,-0.000452825624961406,-0.00046752975322306156],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00022265268489718437,0.0002414086484350264,0.000127053092001006,0.000524361792486161,-0.0009701023809611797],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,7.685436867177486e-05,0.0002426218707114458,0.0001615851215319708,-0.0004987375577911735,-0.000660655670799315],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1091008186340332,0.0028298422694206238,0.0007662934949621558,-0.0003112038830295205,-0.0012275907211005688],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.01121988520026207,0.020876925438642502,0.0006870393408462405,-0.000619351863861084,5.905964644625783e-05],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.29051679372787476,0.0017006462439894676,0.00023508492449764162,0.0001731544325593859,-0.000579683284740895],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0843031108379364,0.02572883665561676,0.003018079325556755,0.0015348526649177074,0.024284658953547478],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.008083650842308998,0.01003817655146122,0.0003211219154763967,-8.188765787053853e-05,0.3221139907836914],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012348066084086895,0.013155452907085419,0.0004423986538313329,0.0011563682928681374,0.5724781155586243],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0011972422944381833,0.004177889320999384,-1.7678325093584135e-05,0.0003485495108179748,0.3005668520927429],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012172549031674862,0.0010988371213898063,0.00026228593196719885,-9.570646216161549e-05,-0.1382392942905426],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.291185200214386]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Token Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Token Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Token Position\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"351ace5b-0df3-4d2a-ae14-888e41a01fc6\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"351ace5b-0df3-4d2a-ae14-888e41a01fc6\")) {                    Plotly.newPlot(                        \"351ace5b-0df3-4d2a-ae14-888e41a01fc6\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"z\":[[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0006664255633950233,-0.0011972333304584026,-0.0003045523480977863,-0.000452825624961406,-0.00046752975322306156],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.00022265268489718437,0.0002414086484350264,0.000127053092001006,0.000524361792486161,-0.0009701023809611797],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,7.685436867177486e-05,0.0002426218707114458,0.0001615851215319708,-0.0004987375577911735,-0.000660655670799315],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.1091008186340332,0.0028298422694206238,0.0007662934949621558,-0.0003112038830295205,-0.0012275907211005688],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.01121988520026207,0.020876925438642502,0.0006870393408462405,-0.000619351863861084,5.905964644625783e-05],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.29051679372787476,0.0017006462439894676,0.00023508492449764162,0.0001731544325593859,-0.000579683284740895],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0843031108379364,0.02572883665561676,0.003018079325556755,0.0015348526649177074,0.024284658953547478],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.008083650842308998,0.01003817655146122,0.0003211219154763967,-8.188765787053853e-05,0.3221139907836914],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012348066084086895,0.013155452907085419,0.0004423986538313329,0.0011563682928681374,0.5724781155586243],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0011972422944381833,0.004177889320999384,-1.7678325093584135e-05,0.0003485495108179748,0.3005668520927429],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.0012172549031674862,0.0010988371213898063,0.00026228593196719885,-9.570646216161549e-05,-0.1382392942905426],[0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,-0.291185200214386]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Token Position: %{x}\\u003cbr\\u003eLayer: %{y}\\u003cbr\\u003eNorm. Logit Diff: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\"}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Token Position\"}},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layer\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Norm. Logit Diff\"}},\"colorscale\":[[0.0,\"rgb(103,0,31)\"],[0.1,\"rgb(178,24,43)\"],[0.2,\"rgb(214,96,77)\"],[0.3,\"rgb(244,165,130)\"],[0.4,\"rgb(253,219,199)\"],[0.5,\"rgb(247,247,247)\"],[0.6,\"rgb(209,229,240)\"],[0.7,\"rgb(146,197,222)\"],[0.8,\"rgb(67,147,195)\"],[0.9,\"rgb(33,102,172)\"],[1.0,\"rgb(5,48,97)\"]],\"cmid\":0.0},\"title\":{\"text\":\"Attribution Patching Over Token Position\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('3a0bc859-eccc-46fb-ab21-231253cfa4af');\n",
+              "var gd = document.getElementById('351ace5b-0df3-4d2a-ae14-888e41a01fc6');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
diff --git a/docs/source/notebooks/tutorials/logit_lens.ipynb b/docs/source/notebooks/tutorials/logit_lens.ipynb
index 90c8b89e..77ad20c1 100644
--- a/docs/source/notebooks/tutorials/logit_lens.ipynb
+++ b/docs/source/notebooks/tutorials/logit_lens.ipynb
@@ -63,7 +63,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 1,
       "metadata": {
         "id": "_VsZHFh7Fi4X"
       },
@@ -82,7 +82,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 2,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -174,7 +174,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 3,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -245,7 +245,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 4,
       "metadata": {
         "id": "a2OkUOXQjmHL"
       },
@@ -307,7 +307,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 10,
+      "execution_count": 5,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -1610,9 +1610,9 @@
               }
             },
             "text/html": [
-              "<div>                            <div id=\"b3911d4b-d13c-4c2c-be99-472396f13a45\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"b3911d4b-d13c-4c2c-be99-472396f13a45\")) {                    Plotly.newPlot(                        \"b3911d4b-d13c-4c2c-be99-472396f13a45\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"texttemplate\":\"%{text}\",\"x\":[\"The\",\" E\",\"iff\",\"el\",\" Tower\",\" is\",\" in\",\" the\",\" city\",\" of\"],\"y\":[0,1,2,3,4,5,6,7,8,9,10,11],\"z\":[[0.024983787909150124,0.12145508080720901,0.7902563214302063,0.27752962708473206,0.9977331161499023,0.32375672459602356,0.415214478969574,0.7520487904548645,0.54457688331604,0.6899943947792053],[0.04746942222118378,0.08490755409002304,0.34766101837158203,0.5692889094352722,0.983742356300354,0.2637466788291931,0.6872087717056274,0.8453984260559082,0.6062489151954651,0.8129417896270752],[0.026062989607453346,0.3097169101238251,0.7040867209434509,0.2206706404685974,0.9181663990020752,0.3488497734069824,0.6373876929283142,0.9504272937774658,0.5847728848457336,0.7201948165893555],[0.02635170891880989,0.38433972001075745,0.5039842128753662,0.2591570019721985,0.6246108412742615,0.38496726751327515,0.27967578172683716,0.9723502993583679,0.2764837443828583,0.6877986788749695],[0.02613661251962185,0.30091190338134766,0.21217022836208344,0.18413318693637848,0.7355331778526306,0.48657187819480896,0.3781237006187439,0.9603129029273987,0.41407859325408936,0.4144521951675415],[0.025683913379907608,0.13498641550540924,0.26062414050102234,0.27143916487693787,0.7304731011390686,0.3886910378932953,0.39428195357322693,0.6723183989524841,0.2978633642196655,0.057537395507097244],[0.025425400584936142,0.2042204886674881,0.5427748560905457,0.7381370663642883,0.3724276125431061,0.43281611800193787,0.39778199791908264,0.6245157718658447,0.18699254095554352,0.05666721239686012],[0.02533440850675106,0.21125738322734833,0.7032177448272705,0.9907069206237793,0.23639222979545593,0.4001739025115967,0.46343299746513367,0.7643018960952759,0.24115046858787537,0.13559488952159882],[0.0253128819167614,0.11660420149564743,0.44579413533210754,0.9970935583114624,0.14586897194385529,0.6360733509063721,0.1932840794324875,0.8479137420654297,0.3974367380142212,0.1916261911392212],[0.025318123400211334,0.06472006440162659,0.8146643042564392,0.9981274008750916,0.253418505191803,0.6489886045455933,0.13464125990867615,0.6901037693023682,0.6442651748657227,0.28538209199905396],[0.02550206519663334,0.044399384409189224,0.9033223390579224,0.9942384958267212,0.16611218452453613,0.19855716824531555,0.19015662372112274,0.2531339228153229,0.35554105043411255,0.18341107666492462],[0.013599246740341187,0.04670924320816994,0.9958765506744385,0.9699326753616333,0.09401834011077881,0.1313934326171875,0.24015067517757416,0.11476960778236389,0.18497605621814728,0.0699693113565445]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Input Tokens: %{x}\\u003cbr\\u003eLayers: %{y}\\u003cbr\\u003eProbability: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"text\":[[\"\\\\n\",\" E\",\"iff\",\"el\",\" Tower\",\" not\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"Z\",\"ness\",\"ter\",\" Tower\",\" not\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"ter\",\" Tower\",\" now\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"stra\",\" Tower\",\" also\",\" the\",\" same\",\" hall\",\" the\"],[\" the\",\"leg\",\"ness\",\"horn\",\" Tower\",\" also\",\" front\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Towers\",\" also\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" also\",\" front\",\" midst\",\"scape\",\" East\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Ing\"],[\" the\",\"TS\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Rome\"],[\" the\",\"FP\",\"el\",\" Tower\",\" tower\",\" located\",\" London\",\" midst\",\" centre\",\" London\"],[\" the\",\"nerg\",\"el\",\" Tower\",\" is\",\" located\",\" the\",\" midst\",\"'s\",\" Paris\"],[\"\\\\n\",\"-\",\"el\",\" Tower\",\" is\",\" a\",\" the\",\" middle\",\" centre\",\" Paris\"]]}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Input Tokens\"},\"tickangle\":0},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layers\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Probability\"}},\"colorscale\":[[0.0,\"rgb(49,54,149)\"],[0.1,\"rgb(69,117,180)\"],[0.2,\"rgb(116,173,209)\"],[0.3,\"rgb(171,217,233)\"],[0.4,\"rgb(224,243,248)\"],[0.5,\"rgb(255,255,191)\"],[0.6,\"rgb(254,224,144)\"],[0.7,\"rgb(253,174,97)\"],[0.8,\"rgb(244,109,67)\"],[0.9,\"rgb(215,48,39)\"],[1.0,\"rgb(165,0,38)\"]],\"cmid\":0.5},\"margin\":{\"t\":60},\"title\":{\"text\":\"Logit Lens Visualization\"}},                        {\"responsive\": true}                    ).then(function(){\n",
+              "<div>                            <div id=\"1ee45ddf-5c1f-427c-8a04-32189d693375\" class=\"plotly-graph-div\" style=\"height:525px; width:100%;\"></div>            <script type=\"text/javascript\">                require([\"plotly\"], function(Plotly) {                    window.PLOTLYENV=window.PLOTLYENV || {};                                    if (document.getElementById(\"1ee45ddf-5c1f-427c-8a04-32189d693375\")) {                    Plotly.newPlot(                        \"1ee45ddf-5c1f-427c-8a04-32189d693375\",                        [{\"coloraxis\":\"coloraxis\",\"name\":\"0\",\"texttemplate\":\"%{text}\",\"x\":[\"The\",\" E\",\"iff\",\"el\",\" Tower\",\" is\",\" in\",\" the\",\" city\",\" of\"],\"y\":[0,1,2,3,4,5,6,7,8,9,10,11],\"z\":[[0.024983787909150124,0.12145508080720901,0.7902563214302063,0.27752962708473206,0.9977331161499023,0.32375672459602356,0.415214478969574,0.7520487904548645,0.54457688331604,0.6899943947792053],[0.04746942222118378,0.08490755409002304,0.34766101837158203,0.5692889094352722,0.983742356300354,0.2637466788291931,0.6872087717056274,0.8453984260559082,0.6062489151954651,0.8129417896270752],[0.026062989607453346,0.3097169101238251,0.7040867209434509,0.2206706404685974,0.9181663990020752,0.3488497734069824,0.6373876929283142,0.9504272937774658,0.5847728848457336,0.7201948165893555],[0.02635170891880989,0.38433972001075745,0.5039842128753662,0.2591570019721985,0.6246108412742615,0.38496726751327515,0.27967578172683716,0.9723502993583679,0.2764837443828583,0.6877986788749695],[0.02613661251962185,0.30091190338134766,0.21217022836208344,0.18413318693637848,0.7355331778526306,0.48657187819480896,0.3781237006187439,0.9603129029273987,0.41407859325408936,0.4144521951675415],[0.025683913379907608,0.13498641550540924,0.26062414050102234,0.27143916487693787,0.7304731011390686,0.3886910378932953,0.39428195357322693,0.6723183989524841,0.2978633642196655,0.057537395507097244],[0.025425400584936142,0.2042204886674881,0.5427748560905457,0.7381370663642883,0.3724276125431061,0.43281611800193787,0.39778199791908264,0.6245157718658447,0.18699254095554352,0.05666721239686012],[0.02533440850675106,0.21125738322734833,0.7032177448272705,0.9907069206237793,0.23639222979545593,0.4001739025115967,0.46343299746513367,0.7643018960952759,0.24115046858787537,0.13559488952159882],[0.0253128819167614,0.11660420149564743,0.44579413533210754,0.9970935583114624,0.14586897194385529,0.6360733509063721,0.1932840794324875,0.8479137420654297,0.3974367380142212,0.1916261911392212],[0.025318123400211334,0.06472006440162659,0.8146643042564392,0.9981274008750916,0.253418505191803,0.6489886045455933,0.13464125990867615,0.6901037693023682,0.6442651748657227,0.28538209199905396],[0.02550206519663334,0.044399384409189224,0.9033223390579224,0.9942384958267212,0.16611218452453613,0.19855716824531555,0.19015662372112274,0.2531339228153229,0.35554105043411255,0.18341107666492462],[0.013599246740341187,0.04670924320816994,0.9958765506744385,0.9699326753616333,0.09401834011077881,0.1313934326171875,0.24015067517757416,0.11476960778236389,0.18497605621814728,0.0699693113565445]],\"type\":\"heatmap\",\"xaxis\":\"x\",\"yaxis\":\"y\",\"hovertemplate\":\"Input Tokens: %{x}\\u003cbr\\u003eLayers: %{y}\\u003cbr\\u003eProbability: %{z}\\u003cextra\\u003e\\u003c\\u002fextra\\u003e\",\"text\":[[\"\\\\n\",\" E\",\"iff\",\"el\",\" Tower\",\" not\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"Z\",\"ness\",\"ter\",\" Tower\",\" not\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"ter\",\" Tower\",\" now\",\" order\",\" same\",\"scape\",\" the\"],[\" the\",\"leg\",\"ness\",\"stra\",\" Tower\",\" also\",\" the\",\" same\",\" hall\",\" the\"],[\" the\",\"leg\",\"ness\",\"horn\",\" Tower\",\" also\",\" front\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Towers\",\" also\",\" the\",\" same\",\"scape\",\" the\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" also\",\" front\",\" midst\",\"scape\",\" East\"],[\" the\",\"FP\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Ing\"],[\" the\",\"TS\",\"ler\",\" Tower\",\" Tower\",\" located\",\" front\",\" midst\",\" centre\",\" Rome\"],[\" the\",\"FP\",\"el\",\" Tower\",\" tower\",\" located\",\" London\",\" midst\",\" centre\",\" London\"],[\" the\",\"nerg\",\"el\",\" Tower\",\" is\",\" located\",\" the\",\" midst\",\"'s\",\" Paris\"],[\"\\\\n\",\"-\",\"el\",\" Tower\",\" is\",\" a\",\" the\",\" middle\",\" centre\",\" Paris\"]]}],                        {\"template\":{\"data\":{\"histogram2dcontour\":[{\"type\":\"histogram2dcontour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"choropleth\":[{\"type\":\"choropleth\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"histogram2d\":[{\"type\":\"histogram2d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmap\":[{\"type\":\"heatmap\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"heatmapgl\":[{\"type\":\"heatmapgl\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"contourcarpet\":[{\"type\":\"contourcarpet\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"contour\":[{\"type\":\"contour\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"surface\":[{\"type\":\"surface\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"},\"colorscale\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]]}],\"mesh3d\":[{\"type\":\"mesh3d\",\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}],\"scatter\":[{\"fillpattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2},\"type\":\"scatter\"}],\"parcoords\":[{\"type\":\"parcoords\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolargl\":[{\"type\":\"scatterpolargl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"bar\":[{\"error_x\":{\"color\":\"#2a3f5f\"},\"error_y\":{\"color\":\"#2a3f5f\"},\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"bar\"}],\"scattergeo\":[{\"type\":\"scattergeo\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterpolar\":[{\"type\":\"scatterpolar\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"histogram\":[{\"marker\":{\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"histogram\"}],\"scattergl\":[{\"type\":\"scattergl\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatter3d\":[{\"type\":\"scatter3d\",\"line\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattermapbox\":[{\"type\":\"scattermapbox\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scatterternary\":[{\"type\":\"scatterternary\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"scattercarpet\":[{\"type\":\"scattercarpet\",\"marker\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}}}],\"carpet\":[{\"aaxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"baxis\":{\"endlinecolor\":\"#2a3f5f\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"minorgridcolor\":\"white\",\"startlinecolor\":\"#2a3f5f\"},\"type\":\"carpet\"}],\"table\":[{\"cells\":{\"fill\":{\"color\":\"#EBF0F8\"},\"line\":{\"color\":\"white\"}},\"header\":{\"fill\":{\"color\":\"#C8D4E3\"},\"line\":{\"color\":\"white\"}},\"type\":\"table\"}],\"barpolar\":[{\"marker\":{\"line\":{\"color\":\"#E5ECF6\",\"width\":0.5},\"pattern\":{\"fillmode\":\"overlay\",\"size\":10,\"solidity\":0.2}},\"type\":\"barpolar\"}],\"pie\":[{\"automargin\":true,\"type\":\"pie\"}]},\"layout\":{\"autotypenumbers\":\"strict\",\"colorway\":[\"#636efa\",\"#EF553B\",\"#00cc96\",\"#ab63fa\",\"#FFA15A\",\"#19d3f3\",\"#FF6692\",\"#B6E880\",\"#FF97FF\",\"#FECB52\"],\"font\":{\"color\":\"#2a3f5f\"},\"hovermode\":\"closest\",\"hoverlabel\":{\"align\":\"left\"},\"paper_bgcolor\":\"white\",\"plot_bgcolor\":\"#E5ECF6\",\"polar\":{\"bgcolor\":\"#E5ECF6\",\"angularaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"radialaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"ternary\":{\"bgcolor\":\"#E5ECF6\",\"aaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"baxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"},\"caxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\"}},\"coloraxis\":{\"colorbar\":{\"outlinewidth\":0,\"ticks\":\"\"}},\"colorscale\":{\"sequential\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"sequentialminus\":[[0.0,\"#0d0887\"],[0.1111111111111111,\"#46039f\"],[0.2222222222222222,\"#7201a8\"],[0.3333333333333333,\"#9c179e\"],[0.4444444444444444,\"#bd3786\"],[0.5555555555555556,\"#d8576b\"],[0.6666666666666666,\"#ed7953\"],[0.7777777777777778,\"#fb9f3a\"],[0.8888888888888888,\"#fdca26\"],[1.0,\"#f0f921\"]],\"diverging\":[[0,\"#8e0152\"],[0.1,\"#c51b7d\"],[0.2,\"#de77ae\"],[0.3,\"#f1b6da\"],[0.4,\"#fde0ef\"],[0.5,\"#f7f7f7\"],[0.6,\"#e6f5d0\"],[0.7,\"#b8e186\"],[0.8,\"#7fbc41\"],[0.9,\"#4d9221\"],[1,\"#276419\"]]},\"xaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"yaxis\":{\"gridcolor\":\"white\",\"linecolor\":\"white\",\"ticks\":\"\",\"title\":{\"standoff\":15},\"zerolinecolor\":\"white\",\"automargin\":true,\"zerolinewidth\":2},\"scene\":{\"xaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"yaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2},\"zaxis\":{\"backgroundcolor\":\"#E5ECF6\",\"gridcolor\":\"white\",\"linecolor\":\"white\",\"showbackground\":true,\"ticks\":\"\",\"zerolinecolor\":\"white\",\"gridwidth\":2}},\"shapedefaults\":{\"line\":{\"color\":\"#2a3f5f\"}},\"annotationdefaults\":{\"arrowcolor\":\"#2a3f5f\",\"arrowhead\":0,\"arrowwidth\":1},\"geo\":{\"bgcolor\":\"white\",\"landcolor\":\"#E5ECF6\",\"subunitcolor\":\"white\",\"showland\":true,\"showlakes\":true,\"lakecolor\":\"white\"},\"title\":{\"x\":0.05},\"mapbox\":{\"style\":\"light\"}}},\"xaxis\":{\"anchor\":\"y\",\"domain\":[0.0,1.0],\"scaleanchor\":\"y\",\"constrain\":\"domain\",\"title\":{\"text\":\"Input Tokens\"},\"tickangle\":0},\"yaxis\":{\"anchor\":\"x\",\"domain\":[0.0,1.0],\"autorange\":\"reversed\",\"constrain\":\"domain\",\"title\":{\"text\":\"Layers\"}},\"coloraxis\":{\"colorbar\":{\"title\":{\"text\":\"Probability\"}},\"colorscale\":[[0.0,\"rgb(49,54,149)\"],[0.1,\"rgb(69,117,180)\"],[0.2,\"rgb(116,173,209)\"],[0.3,\"rgb(171,217,233)\"],[0.4,\"rgb(224,243,248)\"],[0.5,\"rgb(255,255,191)\"],[0.6,\"rgb(254,224,144)\"],[0.7,\"rgb(253,174,97)\"],[0.8,\"rgb(244,109,67)\"],[0.9,\"rgb(215,48,39)\"],[1.0,\"rgb(165,0,38)\"]],\"cmid\":0.5},\"margin\":{\"t\":60},\"title\":{\"text\":\"Logit Lens Visualization\"}},                        {\"responsive\": true}                    ).then(function(){\n",
               "                            \n",
-              "var gd = document.getElementById('b3911d4b-d13c-4c2c-be99-472396f13a45');\n",
+              "var gd = document.getElementById('1ee45ddf-5c1f-427c-8a04-32189d693375');\n",
               "var x = new MutationObserver(function (mutations, observer) {{\n",
               "        var display = window.getComputedStyle(gd).display;\n",
               "        if (!display || display === 'none') {{\n",
diff --git a/docs/source/notebooks/tutorials/walkthrough.ipynb b/docs/source/notebooks/tutorials/walkthrough.ipynb
index 598f988c..e343940f 100644
--- a/docs/source/notebooks/tutorials/walkthrough.ipynb
+++ b/docs/source/notebooks/tutorials/walkthrough.ipynb
@@ -2,9 +2,11 @@
   "cells": [
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "i4WnS0_mkkww"
+      },
       "source": [
-        "# Walkthrough\n"
+        "<a href=\"https://colab.research.google.com/github/ndif-team/nnsight/blob/main/NNsight_Walkthrough.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
       ]
     },
     {
@@ -13,22 +15,26 @@
         "id": "-Scfr942GwO4"
       },
       "source": [
-        "An interactive version of this walkthrough can be found\n",
-        "[here](https://colab.research.google.com/github/ndif-team/nnsight/blob/main/NNsight_Walkthrough.ipynb).\n",
+        "<img src=\"https://nnsight.net/_static/images/nnsight_logo.svg\" alt=\"drawing\" width=\"200\"/>\n",
+        "\n",
+        "# **NNsight**\n",
         "\n",
-        "In this era of large-scale deep learning, often the most interesting AI models are\n",
-        "massive black boxes that are hard to run. Ordinary commercial inference APIs \n",
-        "let us interact with huge models, but they do not let us access model internals.\n",
+        "## The API for a transparent science on black-box AI\n",
         "\n",
-        "The `nnsight` library is different: it provides full access to all neural\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://ndif.us/)\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 this profoundly important field.\n"
+        "big insights on a profoundly important field."
       ]
     },
     {
@@ -37,36 +43,31 @@
         "id": "N1OemD2VGyZx"
       },
       "source": [
-        "# 1️⃣ First, let's start small"
+        "# 1 First, let's start small"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {
-        "id": "WyLMmhrAKTNM"
-      },
+      "metadata": {},
       "source": [
-        "## Tracing Context\n",
+        "## Setup\n",
         "\n",
-        "To demonstrate the core functionality and syntax of nnsight, we'll define and\n",
-        "use a tiny two layer neural network.\n"
+        "Install nnsight:\n",
+        "```\n",
+        "pip install nnsight\n",
+        "```"
       ]
     },
     {
-      "cell_type": "code",
-      "execution_count": 1,
+      "cell_type": "markdown",
       "metadata": {
-        "id": "R0RJijD0eXwf"
+        "id": "WyLMmhrAKTNM"
       },
-      "outputs": [],
       "source": [
-        "# Install nnsight\n",
-        "!pip install nnsight\n",
-        "!pip install --upgrade transformers torch\n",
-        "\n",
-        "from IPython.display import clear_output\n",
+        "## Tracing Context\n",
         "\n",
-        "clear_output()"
+        "To demonstrate the core functionality and syntax of nnsight, we'll define and\n",
+        "use a tiny two layer neural network.\n"
       ]
     },
     {
@@ -75,13 +76,13 @@
         "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",
+        "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,
+      "execution_count": 1,
       "metadata": {
         "id": "2pX2Wg8Ceo6N"
       },
@@ -111,16 +112,25 @@
       },
       "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"
+        "PyTorch model to enable investigation of its internal parameters.\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 2,
       "metadata": {
         "id": "8H9R_ynTJI5y"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+            "  from .autonotebook import tqdm as notebook_tqdm\n"
+          ]
+        }
+      ],
       "source": [
         "import nnsight\n",
         "from nnsight import NNsight\n",
@@ -134,19 +144,19 @@
         "id": "2NfISlQ_Ilvp"
       },
       "source": [
-        "Printing a Pytorch model shows a named hierarchy of modules which is very useful\n",
+        "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,
+      "execution_count": 3,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "bYtnbJHvlGZV",
-        "outputId": "5eb7c572-451c-43c6-956f-af224ae1867c"
+        "outputId": "9bec5a3e-647c-4c49-b4ec-e1420ea0c964"
       },
       "outputs": [
         {
@@ -177,7 +187,7 @@
         "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",
+        "The most common application is opening files as in the following example:\n",
         "\n",
         "```python\n",
         "with open('myfile.txt', 'r') as file:\n",
@@ -188,7 +198,9 @@
         "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",
+        "read from the file.\n",
+        "\n",
+        "Simple enough! Now we can discuss how `nnsight` uses\n",
         "contexts to enable intuitive access into the internals of a neural network.\n"
       ]
     },
@@ -208,7 +220,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 4,
       "metadata": {
         "id": "qEXQ4auPSL-m"
       },
@@ -266,20 +278,19 @@
         "`.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"
+        "grab it for me and put the value into its corresponding Proxy object. Let's try it:\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 6,
+      "execution_count": 5,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 303
+          "height": 329
         },
         "id": "cYe8-r9ptGaG",
-        "outputId": "adf25efe-d2cf-4f65-9eaf-15b155762f4b"
+        "outputId": "5a1ebb19-be61-4eea-9ec1-141a62ebc313"
       },
       "outputs": [
         {
@@ -289,9 +300,10 @@
           "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",
+            "Cell \u001b[0;32mIn[5], 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;28;43mprint\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43moutput\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/proxy.py:70\u001b[0m, in \u001b[0;36mProxy.__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     66\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     68\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---> 70\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     72\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[39mtarget\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;124m\"\u001b[39m\n",
+            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/proxy.py:64\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     57\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     58\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     59\u001b[0m \n\u001b[1;32m     60\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     61\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     62\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 64\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/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/node.py:143\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    133\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    134\u001b[0m \n\u001b[1;32m    135\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or was destroyed).\u001b[39;00m\n\u001b[1;32m    140\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    142\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--> 143\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    145\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."
           ]
         }
@@ -301,40 +313,40 @@
         "\n",
         "    output = tiny_model.output\n",
         "\n",
-        "print(output.value)"
+        "print(output)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "sBz5qfwuR-6F"
+        "id": "4CGa5zSI6kkI"
       },
       "source": [
-        "Oh no, an error! `ValueError: Accessing value before it's been set.`\n",
+        "Oh no an error! \"Accessing 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"
+        "error:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 67,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "y_bXRd5dvsBu",
-        "outputId": "a8e8985c-2661-46a8-975b-d43f67cfc162"
+        "outputId": "feded172-02f5-429b-f9a9-766eed3a787d"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "tensor([[-0.2687, -0.3314]])\n"
+            "tensor([[-0.1301, -0.4906]])\n"
           ]
         }
       ],
@@ -343,7 +355,7 @@
         "\n",
         "    output = tiny_model.output.save()\n",
         "\n",
-        "print(output.value)"
+        "print(output)"
       ]
     },
     {
@@ -376,20 +388,20 @@
         "\n",
         "</details>\n",
         "\n",
-        "Just like we saved the model's final output, we can save the output of\n",
-        "any of its submodules. Using normal Python attribute syntax, we can discover\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,
+      "execution_count": 68,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "0WcVUSP-0CJi",
-        "outputId": "7fd9d6ac-152d-4431-9a4d-0b60321638eb"
+        "outputId": "f40600e1-719d-46c2-f420-e5a9826dd315"
       },
       "outputs": [
         {
@@ -409,28 +421,30 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "MXn9XRQqkkw1"
+      },
       "source": [
-        "Let's access the output of the first layer (non-coincidentally named 'layer1'):"
+        "Let's access the output of the first layer (which we've named 'layer1'):"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 69,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "akWr-cNqy-9O",
-        "outputId": "68b9e071-c609-4b70-906d-2f2cd93b4495"
+        "outputId": "13faddb0-fead-4fd0-a500-2bff00c095b3"
       },
       "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"
+            "tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -439,7 +453,7 @@
         "\n",
         "    l1_output = tiny_model.layer1.output.save()\n",
         "\n",
-        "print(l1_output.value)"
+        "print(l1_output)"
       ]
     },
     {
@@ -449,27 +463,27 @@
       },
       "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 for a few\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,
+      "execution_count": 70,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "3EHEN38N3nXR",
-        "outputId": "5afd17a7-0bcc-4d1a-c131-7038bde9455d"
+        "outputId": "74f25b48-89fd-4b15-9dab-cead54f2386b"
       },
       "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"
+            "tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -478,7 +492,7 @@
         "\n",
         "    l2_input = tiny_model.layer2.input.save()\n",
         "\n",
-        "print(l2_input.value)"
+        "print(l2_input)"
       ]
     },
     {
@@ -508,24 +522,32 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "d5DE1ZJAkkw1"
+      },
       "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",
+        "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": {},
+      "execution_count": 71,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "oZzS86tzkkw1",
+        "outputId": "1c1f3a94-eb4d-48c7-c69f-9241f2fa08f5"
+      },
       "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"
+            "Layer 1 - out:  tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -559,13 +581,13 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 72,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "c5XCiSZn2p3k",
-        "outputId": "97d34dee-7bb5-4735-8c53-6aac6377af9e"
+        "outputId": "c6c5060b-4cd9-4154-8bc6-adf4f5aec5cc"
       },
       "outputs": [
         {
@@ -583,7 +605,7 @@
         "    # 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",
+        "    # 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",
@@ -596,7 +618,7 @@
         "id": "BGPUJvWq_bOp"
       },
       "source": [
-        "Nice! That worked seamlessly, but why didn't we need to call\n",
+        "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",
@@ -609,20 +631,20 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 13,
+      "execution_count": 73,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "FIcOYeEjFuln",
-        "outputId": "d836abd3-b713-44c5-aedc-95bf85abf11a"
+        "outputId": "074fd78f-bc18-48b8-ccef-7db9df756ea6"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "tensor(-1.3797)\n"
+            "tensor(0.9377)\n"
           ]
         }
       ],
@@ -641,39 +663,49 @@
       },
       "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",
+        "the given `input`. When the output of `tiny_model.layer1` is computed, take its sum. Then do\n",
+        "the same for `tiny_model.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": {},
+      "metadata": {
+        "id": "kLmzsHaMkkw2"
+      },
       "source": [
         "## Custom Functions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "GJFCuZHtkkw2"
+      },
       "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",
+        "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",
+        "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": {},
+      "execution_count": 74,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "DF7MPuhckkw2",
+        "outputId": "c263468b-affe-4f69-a045-d674e07333d3"
+      },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "tensor(-0.7796) tensor(-0.7796)\n"
+            "tensor(1.5584) tensor(1.5584)\n"
           ]
         }
       ],
@@ -684,27 +716,29 @@
         "    total = 0\n",
         "    for element in flat:\n",
         "        total += element.item()\n",
-        "    \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",
+        "    # Specify the function name and its arguments (in a comma-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",
         "\n",
         "print(custom_sum, sum)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "02UKwk_Lkkw2"
+      },
       "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."
+        "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."
       ]
     },
     {
@@ -723,32 +757,32 @@
       },
       "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, we often not only want to view the computation of a model, but also to influence it.\n",
+        "really insightful, and a number of ML techniques do exactly that. However, often we not only want to view the computation of a model, but also to influence it.\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:"
+        "this really easy using the '=' operator:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 15,
+      "execution_count": 75,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "-6y2wzJqOz3a",
-        "outputId": "6274e5c5-c5fe-4e07-87c4-a119cfd09742"
+        "outputId": "13de53a9-c396-4167-e2f6-950957233c8e"
       },
       "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"
+            "Before: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n",
+            "After: tensor([[ 0.0000,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -775,52 +809,38 @@
         "id": "oZz9SMs3Y_iS"
       },
       "source": [
-        "Seems our change was reflected. Now the same for the last dimension:\n"
+        "Seems our change was reflected. Now let's do the same for the last dimension:\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 16,
+      "execution_count": 76,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 363
+          "height": 305
         },
         "id": "qwlqvHFcxld2",
-        "outputId": "28ae7175-5af0-4bd0-f9f0-21680a4fd748"
+        "outputId": "b1ac58fb-dbbd-4375-fcd7-3362762f00fe"
       },
       "outputs": [
         {
-          "ename": "IndexError",
-          "evalue": "Above exception when execution Node: 'setitem_0' in Graph: '6063279136'",
+          "ename": "NNsightError",
+          "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)",
-            "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'"
+            "Traceback (most recent call last):",
+            "  File \"/Users/emmabortz/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py\", line 297, in execute",
+            "    output = self.target(*args, **kwargs)",
+            "IndexError: index 10 is out of bounds for dimension 1 with size 10",
+            "",
+            "During handling of the above exception, another exception occurred:",
+            "",
+            "Traceback (most recent call last):",
+            "  File \"/var/folders/rx/0nl_h2cd54q90chs9hf0n5qr0000gn/T/ipykernel_80599/3404137504.py\", line 8, in <module>",
+            "    tiny_model.layer1.output[:, hidden_dims] = 0",
+            "",
+            "NNsightError: index 10 is out of bounds for dimension 1 with size 10"
           ]
         }
       ],
@@ -843,121 +863,260 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "DKaN1oi76djp"
+      },
       "source": [
-        "Oh no, we are getting an error. Looks like it's happening when we are setting the output. \n",
+        "Oh no, we are getting an error! Ah of course, we needed to index at `hidden_dims - 1` not `hidden_dims`.\n",
         "\n",
-        "How can we find what went wrong? Is there an easy way to debug this?\n",
-        "\n",
-        "Enter \"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:"
+        "If you've been using `nnsight`, you are probably familiar with error messages that can be quite difficult to troubleshoot. In `nnsight 0.4` we've now improved error messaging to be descriptive and line-specific, as you should see in the above example!"
       ]
     },
     {
-      "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"
-          ]
-        }
-      ],
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Yz-XpUpZ36yd"
+      },
       "source": [
-        "# turn on scan and validate\n",
-        "with tiny_model.trace(input, scan=True, validate=True):\n",
+        "<details>\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",
+        "<summary>\n",
+        "Old NNsight error messaging\n",
+        "</summary>\n",
         "\n",
-        "print(\"Before:\", l1_output_before)\n",
-        "print(\"After:\", l1_output_after)"
+        "If you've been using NNsight prior to the NNsight 0.4 release, you will be familiar with the following non-descriptive error messaging. If you choose to turn off NNsight 0.4's new error messaging feature, this is how errors within the tracing context will appear.\n",
+        "\n",
+        "\n",
+        "```\n",
+        "---------------------------------------------------------------------------\n",
+        "IndexError                                Traceback (most recent call last)\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in execute(self)\n",
+        "    379                 # Call the target to get value.\n",
+        "--> 380                 output = self.target(*args, **kwargs)\n",
+        "    381\n",
+        "\n",
+        "IndexError: index 10 is out of bounds for dimension 1 with size 10\n",
+        "\n",
+        "The above exception was the direct cause of the following exception:\n",
+        "\n",
+        "IndexError                                Traceback (most recent call last)\n",
+        "20 frames\n",
+        "<ipython-input-16-5c81de91fb1f> in <cell line: 0>()\n",
+        "----> 1 with tiny_model.trace(input):\n",
+        "      2\n",
+        "      3     # Save the output before the edit to compare.\n",
+        "      4     # Notice we apply .clone() before saving as the setting operation is in-place.\n",
+        "      5     l1_output_before = tiny_model.layer1.output.clone().save()\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/Tracer.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
+        "    100\n",
+        "    101\n",
+        "--> 102         super().__exit__(exc_type, exc_val, exc_tb)\n",
+        "    103\n",
+        "    104     def invoke(self, *inputs: Any, **kwargs) -> Invoker:\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/GraphBasedContext.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
+        "    215             raise exc_val\n",
+        "    216\n",
+        "--> 217         self.backend(self)\n",
+        "    218\n",
+        "    219     ### BACKENDS ########\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/backends/LocalBackend.py in __call__(self, obj)\n",
+        "     25     def __call__(self, obj: LocalMixin):\n",
+        "     26\n",
+        "---> 27         obj.local_backend_execute()\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/Tracer.py in local_backend_execute(self)\n",
+        "    144         self.graph.execute()\n",
+        "    145\n",
+        "--> 146         self.model.interleave(\n",
+        "    147             self.model._execute,\n",
+        "    148             self.graph,\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in interleave(self, fn, intervention_graph, *inputs, **kwargs)\n",
+        "    467         module_paths = InterventionProtocol.get_interventions(intervention_graph).keys()\n",
+        "    468\n",
+        "--> 469         with HookHandler(\n",
+        "    470             self._model,\n",
+        "    471             list(module_paths),\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
+        "    579\n",
+        "    580         if isinstance(exc_val, Exception):\n",
+        "--> 581             raise exc_val\n",
+        "    582\n",
+        "    583\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in interleave(self, fn, intervention_graph, *inputs, **kwargs)\n",
+        "    478         ):\n",
+        "    479             try:\n",
+        "--> 480                 fn(*inputs, **kwargs)\n",
+        "    481             except protocols.EarlyStopProtocol.EarlyStopException:\n",
+        "    482                 # TODO: Log.\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in _execute(self, *prepared_inputs, **kwargs)\n",
+        "    585             pass\n",
+        "    586\n",
+        "--> 587         return self._model(\n",
+        "    588             *prepared_inputs,\n",
+        "    589             **kwargs,\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)\n",
+        "   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]\n",
+        "   1735         else:\n",
+        "-> 1736             return self._call_impl(*args, **kwargs)\n",
+        "   1737\n",
+        "   1738     # torchrec tests the code consistency with the following code\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)\n",
+        "   1842\n",
+        "   1843         try:\n",
+        "-> 1844             return inner()\n",
+        "   1845         except Exception:\n",
+        "   1846             # run always called hooks if they have not already been run\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in inner()\n",
+        "   1788                 args = bw_hook.setup_input_hook(args)\n",
+        "   1789\n",
+        "-> 1790             result = forward_call(*args, **kwargs)\n",
+        "   1791             if _global_forward_hooks or self._forward_hooks:\n",
+        "   1792                 for hook_id, hook in (\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/container.py in forward(self, input)\n",
+        "    248     def forward(self, input):\n",
+        "    249         for module in self:\n",
+        "--> 250             input = module(input)\n",
+        "    251         return input\n",
+        "    252\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)\n",
+        "   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]\n",
+        "   1735         else:\n",
+        "-> 1736             return self._call_impl(*args, **kwargs)\n",
+        "   1737\n",
+        "   1738     # torchrec tests the code consistency with the following code\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)\n",
+        "   1842\n",
+        "   1843         try:\n",
+        "-> 1844             return inner()\n",
+        "   1845         except Exception:\n",
+        "   1846             # run always called hooks if they have not already been run\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in inner()\n",
+        "   1801                         hook_result = hook(self, args, kwargs, result)\n",
+        "   1802                     else:\n",
+        "-> 1803                         hook_result = hook(self, args, result)\n",
+        "   1804\n",
+        "   1805                     if hook_result is not None:\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in output_hook(module, input, output, module_path)\n",
+        "    564\n",
+        "    565                 def output_hook(module, input, output, module_path=module_path):\n",
+        "--> 566                     return self.output_hook(output, module_path)\n",
+        "    567\n",
+        "    568                 self.handles.append(\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in <lambda>(activations, module_path)\n",
+        "    473                 activations, module_path, \"input\", intervention_handler\n",
+        "    474             ),\n",
+        "--> 475             output_hook=lambda activations, module_path: InterventionProtocol.intervene(\n",
+        "    476                 activations, module_path, \"output\", intervention_handler\n",
+        "    477             ),\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in intervene(cls, activations, module_path, key, intervention_handler)\n",
+        "    454\n",
+        "    455                 # Value injection.\n",
+        "--> 456                 node.set_value(value)\n",
+        "    457\n",
+        "    458                 # Check if through the previous value injection, there was a 'swap' intervention.\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in set_value(self, value)\n",
+        "    408\n",
+        "    409             if listener.fulfilled() and not self.graph.sequential:\n",
+        "--> 410                 listener.execute()\n",
+        "    411\n",
+        "    412         for dependency in self.arg_dependencies:\n",
+        "\n",
+        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in execute(self)\n",
+        "    385         except Exception as e:\n",
+        "    386\n",
+        "--> 387             raise type(e)(\n",
+        "    388                 f\"Above exception when execution Node: '{self.name}' in Graph: '{self.graph.id}'\"\n",
+        "    389             ) from e\n",
+        "\n",
+        "IndexError: Above exception when execution Node: 'setitem_0' in Graph: '132147685816016'\n",
+        "\n",
+        "\n",
+        "```\n",
+        "\n",
+        "\n",
+        "</details>\n"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "JCNR_Jkpxr8l"
+        "id": "Lo7oHC0yzI_p"
       },
       "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",
+        "The error messaging feature can be toggled using `nnsight.CONFIG.APP.DEBUG` which defaults to true.\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",
+        "<summary>\n",
+        "Toggle Error Messaging \n",
+        "</summary>\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",
+        "Turn off debugging:\n",
+        "```\n",
+        "import nnsight\n",
         "\n",
-        "---\n",
+        "nnsight.CONFIG.APP.DEBUG = False\n",
+        "nnsight.CONFIG.save()\n",
+        "```\n",
         "\n",
-        "</details>\n",
+        "Turn on debugging:\n",
+        "```\n",
+        "import nnsight\n",
         "\n",
-        "Let's try again with the correct indexing, and view the shape of the output\n",
+        "nnsight.CONFIG.APP.DEBUG = True\n",
+        "nnsight.CONFIG.save()\n",
+        "```\n",
+        "</details>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JCNR_Jkpxr8l"
+      },
+      "source": [
+        "Now that we know more about NNsight's error messaging, let's try our setting operation again with the correct indexing and view the shape of the output\n",
         "before leaving the tracing context:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 18,
+      "execution_count": 77,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "Wf8ugJKB2mru",
-        "outputId": "4eed6c19-cb8c-4deb-e0e1-25aa80959c64"
+        "outputId": "f076fa36-0271-4185-8583-9e7190e1e8e9"
       },
       "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"
+            "Layer 1 output shape: InterventionProxy (fetch_attr)\n",
+            "Before: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n",
+            "After: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0000]])\n"
           ]
         }
       ],
@@ -983,50 +1142,64 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "9DH7PeL13-WU"
+        "id": "6R3wj1BKkkw5"
       },
       "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"
+        "## Scan and Validate\n",
+        "Error codes are helpful, but sometimes you may want to quickly troubleshoot your code without actually running it. \n",
+        "\n",
+        "Enter \"Scanning\" and \"Validating\"! We can enable this features by setting the `scan=True` and `validate=True` flag in the `trace` method."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "\"Scanning\" runs \"fake\" inputs throught the model to collect information like shapes and types (i.e., scanning will populate all called `.inputs` and `.outputs`). \n",
+        "\n",
+        "\"Validating\" attempts to execute the intervention proxies with \"fake\" inputs to check if they work (i.e., executes all interventions in your code with fake tensors).\n",
+        "\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. Let's try it out on our example above."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 19,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "AP5wRmz_4YE7",
-        "outputId": "4c30dc1d-232d-4900-bfdd-0c93a9b23ad7"
-      },
+      "execution_count": 78,
+      "metadata": {},
       "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"
+          "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[78], 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;01mwith\u001b[39;00m tiny_model\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;28minput\u001b[39m, scan\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, validate\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\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",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/contexts/interleaving.py:96\u001b[0m, in \u001b[0;36mInterleavingTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     92\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     94\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m---> 96\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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/tracer.py:25\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mglobals\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m GlobalTracingContext\n\u001b[1;32m     23\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mtry_deregister(\u001b[38;5;28mself\u001b[39m)\n\u001b[0;32m---> 25\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/base.py:72\u001b[0m, in \u001b[0;36mContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     69\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\u001b[38;5;241m.\u001b[39mstack\u001b[38;5;241m.\u001b[39mpop()\n\u001b[1;32m     71\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---> 72\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n\u001b[1;32m     74\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madd(graph\u001b[38;5;241m.\u001b[39mstack[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m], graph, \u001b[38;5;241m*\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[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkwargs)\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbackend \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",
+            "Cell \u001b[0;32mIn[78], 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~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:126\u001b[0m, in \u001b[0;36mProxy.__setitem__\u001b[0;34m(self, key, value)\u001b[0m\n\u001b[1;32m    125\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[Self, 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--> 126\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    127\u001b[0m \u001b[43m        \u001b[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    128\u001b[0m \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\n\u001b[1;32m    129\u001b[0m \u001b[43m        \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    130\u001b[0m \u001b[43m        \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    131\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py:250\u001b[0m, in \u001b[0;36mNode.create\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    247\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m value\n\u001b[1;32m    249\u001b[0m \u001b[38;5;66;03m# Otherwise just create the Node on the Graph like normal.\u001b[39;00m\n\u001b[0;32m--> 250\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    251\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    252\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    253\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/graph.py:131\u001b[0m, in \u001b[0;36mGraph.create\u001b[0;34m(self, target, redirect, *args, **kwargs)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;66;03m# Redirection.\u001b[39;00m\n\u001b[1;32m    129\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstack[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m] \u001b[38;5;28;01mif\u001b[39;00m redirect \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstack \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[0;32m--> 131\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mproxy_class(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode_class\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\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[43mgraph\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~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:118\u001b[0m, in \u001b[0;36mValidatingInterventionNode.__init__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m    114\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattached\n\u001b[1;32m    115\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfake_value \u001b[38;5;129;01mis\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39m_empty\n\u001b[1;32m    116\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m Protocol\u001b[38;5;241m.\u001b[39mis_protocol(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtarget)\n\u001b[1;32m    117\u001b[0m ):\n\u001b[0;32m--> 118\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfake_value \u001b[38;5;241m=\u001b[39m \u001b[43mvalidate\u001b[49m\u001b[43m(\u001b[49m\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[43m \u001b[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[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[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:147\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m FakeTensorMode(\n\u001b[1;32m    142\u001b[0m     allow_non_fake_inputs\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    143\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    144\u001b[0m ) \u001b[38;5;28;01mas\u001b[39;00m fake_mode:\n\u001b[1;32m    145\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m FakeCopyMode(fake_mode):\n\u001b[0;32m--> 147\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m GlobalTracingContext\u001b[38;5;241m.\u001b[39mexit_global_tracing_context():\n\u001b[1;32m    149\u001b[0m             \u001b[38;5;28;01mif\u001b[39;00m backwards_check(target, \u001b[38;5;241m*\u001b[39margs):\n\u001b[1;32m    150\u001b[0m                 \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/globals.py:100\u001b[0m, in \u001b[0;36mGlobalTracingContext.GlobalTracingExit.__exit__\u001b[0;34m(self, exc_type, exc_val, traceback)\u001b[0m\n\u001b[1;32m     96\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     98\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--> 100\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:156\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m    150\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    152\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m InterventionNode\u001b[38;5;241m.\u001b[39mprepare_inputs(\n\u001b[1;32m    153\u001b[0m     (args, kwargs), fake\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    154\u001b[0m )\n\u001b[0;32m--> 156\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtarget\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    157\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    158\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    159\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:2350\u001b[0m, in \u001b[0;36mFakeCopyMode.__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m   2348\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   2349\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39mDisableTorchFunctionSubclass():\n\u001b[0;32m-> 2350\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": [
-        "with tiny_model.trace(input):\n",
+        "# turn on scan and validate\n",
+        "with tiny_model.trace(input, scan=True, validate=True):\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",
+        "    # the error is happening here\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",
@@ -1037,13 +1210,42 @@
       "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."
+        "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. [Here](https://nnsight.net/notebooks/features/scan_validate/) is a more detailed example of scan and validate in action!\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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "45E98sIAVc5A"
+      },
+      "source": [
+        "We can also use the `.scan()` method to get the shape of a module without having to fully run the model. If scan  is enabled, our input is run though the model 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!"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 20,
-      "metadata": {},
+      "execution_count": 79,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 382
+        },
+        "id": "5J79HytJkkw5",
+        "outputId": "3b11405b-7bd7-45b6-dc16-11c8456b277b"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1061,6 +1263,17 @@
         "print(dim)"
       ]
     },
+    {
+      "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": "markdown",
       "metadata": {
@@ -1076,21 +1289,21 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 21,
+      "execution_count": 80,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "L-cYUNNyRDn1",
-        "outputId": "d03dafa1-ef1e-40f1-de28-4dc2f77fe4db"
+        "outputId": "c1521deb-34d7-4ea5-c80f-0c5c873252d9"
       },
       "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 1 output gradient: tensor([[-0.1296,  0.4562, -0.1182, -0.3536,  0.0703,  0.1411, -0.3201, -0.4890,\n",
+            "          0.0196, -0.0452]])\n",
             "Layer 2 output gradient: tensor([[1., 1.]])\n"
           ]
         }
@@ -1128,13 +1341,13 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 22,
+      "execution_count": 81,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "9bFoaJpOWlRb",
-        "outputId": "1c026534-9ffe-40e7-d016-e4565ddde8ac"
+        "outputId": "d3c78c79-4b0f-496f-f722-573f6db30e3a"
       },
       "outputs": [
         {
@@ -1169,29 +1382,39 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "SJ6dAgzQkkw6"
+      },
       "source": [
         "## Early Stopping"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "FZLrnIxAkkw6"
+      },
       "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 example where this could be particularly useful is working with SAEs - we could train an SAE on one layer and then stop model execution. "
+        "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": {},
+      "execution_count": 82,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "MwAeiZhGkkw6",
+        "outputId": "85e6c10d-50dd-4517-b34e-2a9757e29c80"
+      },
       "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"
+            "L1 - Output:  tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
+            "         -0.0748,  0.0088]])\n"
           ]
         }
       ],
@@ -1206,7 +1429,9 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "eO_8oZPXkkw6"
+      },
       "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",
@@ -1217,8 +1442,15 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 24,
-      "metadata": {},
+      "execution_count": 83,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 347
+        },
+        "id": "uPJ6j2Zwkkw6",
+        "outputId": "aeff829d-e973-489e-c24e-1d5cdc80c7d0"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1234,10 +1466,10 @@
           "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",
+            "Cell \u001b[0;32mIn[83], 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~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:70\u001b[0m, in \u001b[0;36mProxy.__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     66\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     68\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---> 70\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     72\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[39mtarget\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;124m\"\u001b[39m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:64\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     57\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     58\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     59\u001b[0m \n\u001b[1;32m     60\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     61\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     62\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 64\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~/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py:143\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    133\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    134\u001b[0m \n\u001b[1;32m    135\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or was destroyed).\u001b[39;00m\n\u001b[1;32m    140\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    142\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--> 143\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    145\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."
           ]
         }
@@ -1252,26 +1484,36 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "i1iJOI29kkw6"
+      },
       "source": [
         "## Conditional Interventions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "NvC2f8qqkkw6"
+      },
       "source": [
-        "Interventions can also be made conditional. \n",
+        "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."
+        "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": [
-        {
+      "execution_count": 84,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "F_SIV2CDkkw6",
+        "outputId": "2ba0cfb8-dfd8-4403-e2e9-db3de2045ec7"
+      },
+      "outputs": [
+        {
           "name": "stdout",
           "output_type": "stream",
           "text": [
@@ -1293,22 +1535,23 @@
     },
     {
       "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": {},
+      "metadata": {
+        "id": "CZLQecmrkkw6"
+      },
       "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": {},
+      "execution_count": 21,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "qZsFtC5jkkw6",
+        "outputId": "4bbad29b-2d9d-4816-ed73-859a829e5ba5"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1322,19 +1565,207 @@
         "with tiny_model.trace(input) as tracer:\n",
         "\n",
         "  non_rand_int = 8\n",
-        "  \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": "hl99yP3kkkw6"
+      },
+      "source": [
+        "With `nnsight 0.4` we can now also use Python `if` statements within the tracing context to create a conditional context!\n",
+        "\n",
+        "*Note: Colab behaves a little strangely with this feature the first time you run it - expect some lagging and warnings*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 85,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "52a1vwtCfdIS",
+        "outputId": "51680d97-0f1f-4ead-ab97-26ec81587fff"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Random Integer  tensor([2])  is Even\n"
+          ]
+        }
+      ],
+      "source": [
+        "with tiny_model.trace(input) as tracer:\n",
+        "\n",
+        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
+        "\n",
+        "  # Since this if statement is inside the tracing context the if will\n",
+        "  # create a conditional context and will only execute the intervention\n",
+        "  # if this condition is met\n",
+        "  if rand_int % 2 == 0:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
+        "\n",
+        "  if rand_int % 2 == 1:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tEML88wTf9Lv"
+      },
+      "source": [
+        "`elif` statements should also work as `if` statements within the tracing context:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 86,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "3njmT1szf-Lq",
+        "outputId": "b1088864-8e0e-4526-da98-b6275adbabdf"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Random Integer  tensor([-3])  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",
+        "  # Since this if statement is inside the tracing context the if will\n",
+        "  # create a conditional context and will only execute the intervention\n",
+        "  # if this condition is met\n",
+        "  if rand_int % 2 == 0:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
+        "  elif rand_int % 2 == 1:\n",
+        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ONzsQcHtxKkf"
+      },
+      "source": [
+        "## Iterative Interventions"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NrwpkfYXxRBO"
+      },
+      "source": [
+        "With the iterator context, you can now run an intervention loop at scale. It iteratively executes and updates a single intervention graph. Use a `.session()` to define the Iterator context and pass in a sequence of items that you want to loop over at each iteration\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 87,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "AYWJVln7xP85",
+        "outputId": "3b7a1271-b199-4078-c70e-0fefc4cd5d7b"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        }
+      ],
+      "source": [
+        "with tiny_model.session() as session:\n",
+        "\n",
+        "  li = nnsight.list() # an NNsight built-in list object\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": {
+        "id": "5Dbnsw6fx8xL"
+      },
+      "source": [
+        "With `nnsight 0.4` we can now also use Python `for` loops within a tracer context at scale.\n",
+        "\n",
+        "*NOTE: inline for loops (i.e., `[x for x in <Proxy object>`]) are not currently supported.*"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 88,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "14rgY7cZyNUr",
+        "outputId": "a028118d-cf04-45d3-ea38-71dd04405a9b"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        }
+      ],
+      "source": [
+        "# New: Using Python for loops for iterative interventions\n",
+        "with tiny_model.session() as session:\n",
+        "\n",
+        "    li = nnsight.list()\n",
+        "    [li.append([num]) for num in range(0, 3)]\n",
+        "    li2 = nnsight.list().save()\n",
+        "\n",
+        "    # Using regular for loops\n",
+        "    for item in li:\n",
+        "        for item_2 in item: # for loops can be nested!\n",
+        "            li2.append(item_2)\n",
+        "\n",
+        "print(\"\\nList: \", li2)"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "gvQ1nZgYQDG3"
       },
       "source": [
-        "# 2️⃣ Bigger"
+        "# 2️ Bigger"
       ]
     },
     {
@@ -1348,8 +1779,10 @@
         "\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"
+        "complex and powerful types of models, and we're excited to see what can be done\n",
+        "to leverage its features!\n",
+        "\n",
+        "However, if you'd like to load a Language Model from HuggingFace with its tokenizer, the`LanguageModel` subclass greatly simplifies this process."
       ]
     },
     {
@@ -1377,60 +1810,71 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 27,
+      "execution_count": 6,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 769,
+          "height": 713,
           "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"
+            "26dcd7b75feb41a6b86d4260ecd64d06",
+            "39c9f1effe664e899125748e97fad198",
+            "383b893224684aee916e455c0d3d8c06",
+            "af24a12bd05b44148b5d6f91ac2445a0",
+            "d00750501f2548d0b1c69f4753aec877",
+            "98d28dddf75a45a2b1e05131b0b865a1",
+            "6e81d6f31f2446c999df5511e8afd8d0",
+            "68b205f60f9d4599a7d1f78583d0b93c",
+            "69e8c149f676457d95029f21e9a2956a",
+            "ba90058201ee48dcae9411ded457e85f",
+            "346ecc4734014cb4931bea53d5a32bf1",
+            "79a221c9228d459aa17402c9fef1ad1d",
+            "9c13a3d768eb47c8b26a7c6bfa8ec957",
+            "e97c052113f54dddbb619ed228f6b4ef",
+            "41fc21ee95c14d29943a56f2b63f4989",
+            "5250cf68037a43e389493733c5b9d4ac",
+            "e24033c5ba7b4e4fa4f4375d85b03bad",
+            "4b798a784ae748b1a0aef025ead976fc",
+            "19c0e121e52040b18393aee887730644",
+            "2702f30b5d3b4b0b8184fec58863750b",
+            "42fd0ea3be3c4b3f98ca7818efee3669",
+            "184b5c75820b4c42ae861669ac1a47e5",
+            "d8feaceb6d854cb2b2ff19a0309d9a7b",
+            "fd84ffef58604ce3a16ef8f9868072f3",
+            "64b713188f054535a3063c079bceead0",
+            "9cea52b7835b40628fc0a65199c5e82f",
+            "87a5c035bc6345aea6441b3db4b94321",
+            "7f4f63d94f8942fcb3e4b286d5848f0c",
+            "a310703500c5481aa37b698c659e1ca6",
+            "7afff81060f0481aa95a56daa570e1b8",
+            "a038f9037ed540f2ba43f3f85920b17e",
+            "ec4f48d5a9ef4afa8fa44078dfb236c6",
+            "405308ec76074c7ea319e1c1358baf39",
+            "7479281e090d4eeba4c21a9c15b02165",
+            "f4b25b3aac414499b653404ea3c066da",
+            "0383a2798a0c4fcfbd1f068817770fd6",
+            "385291e96814453b8f805cbe8a5e594d",
+            "f7a41e267f79401681bc4693c52e55fa",
+            "caff23247fb9432db8d5860a8132e1b3",
+            "f47609606aa14f6fa2407aa26edae8e3",
+            "cef2fe4eb51146afabe820140a3b0df3",
+            "52281c8594ad465397e6079385c558a1",
+            "e5302eb0481c4e529fc68353eed4b79f",
+            "0c46bce8cc3e4393898b025ba6b4e543",
+            "14008810744549208f147a132b419d20",
+            "cb4be83440e44e1c9e664ac1874b830d",
+            "f54c77317ca04861a958c8f982b25653",
+            "db24a855fb5846e28c98dcfe9afff259",
+            "647c861fcaa44148a90da061b1d5c99c",
+            "e4e0b1de76ff40a88b71e0b5a9f7bc39",
+            "504038c4efd3439daa958333f108b5fb",
+            "ded8c188917b424dbf264ce494fe0532",
+            "76fe13ab697c46cc9fcbe46112fb45b1",
+            "a1c25910a31a4b5a9c48f1a7fd8047c5",
+            "5d877979f2db4dbfa7581d2fe4381cf8"
           ]
         },
         "id": "1OD2z7d3HQJU",
-        "outputId": "7a53864c-6965-4c25-9668-4b0d56db9f6b"
+        "outputId": "3a8f135a-ac39-47fe-da78-6680d42572ed"
       },
       "outputs": [
         {
@@ -1463,17 +1907,11 @@
             "    (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",
+            "  (generator): Generator(\n",
+            "    (streamer): Streamer()\n",
+            "  )\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": [
@@ -1516,10 +1954,9 @@
         "\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",
+        "Let's now apply some of the features that we used on the small model to `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",
+        "simpler (i.e., you can send prompts to the model 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",
@@ -1528,22 +1965,40 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 28,
+      "execution_count": 90,
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/"
+          "base_uri": "https://localhost:8080/",
+          "height": 152,
+          "referenced_widgets": [
+            "2fedb8b529c34049932f07b577b515d4",
+            "bd1a11654c2d47a09610e4cc23b8bc32",
+            "26acae3d321743d2ad8efbf8727dbd05",
+            "f42e44122a2d496cb7b2ea85c2858a71",
+            "c1586aeef6814447ad46c1f8ddaae625",
+            "5e091b15cc9748469676b165533f0c44",
+            "f143977aeaf24798b7744681069762d4",
+            "a4b5c225f5ac4719ade8e6aafcff78d2",
+            "ee27d18a555c46e78c92b282bfa6dc9d",
+            "a3154c60744d4a1a87b1c86b89a142b9",
+            "de023d4526524f48a34883a287aa8a4f",
+            "a15895ad1cb541ca89fa9def63779104",
+            "ee60494714c2485f802298dce918edfe",
+            "2a0f2484053f41298aceacea5c1cc195",
+            "57aebe107749473390af2532806920ce",
+            "ff44e896c55048c5a7d92a56aa9eedfc",
+            "7fdb623e5b2d4eb1b6cee27da2314062",
+            "95832c5076ab4c80810d3c443e061156",
+            "08fb0620499c44eb9d160a0f21263451",
+            "287a8fe3ddf54d36b7603171ec7a2688",
+            "55c0bb35417e454d842e12bd2dfe62df",
+            "d5821d69e1fa4673bb39ab70a7f60466"
+          ]
         },
         "id": "mLSapaMCgLNU",
-        "outputId": "192b8998-ecca-4f25-ecf2-5196c0d2d672"
+        "outputId": "107769db-2d2e-4b52-ed59-31555ebcbfcb"
       },
       "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",
@@ -1577,13 +2032,12 @@
         "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",
+        "We just ran a little intervention on a much more complex model with many more\n",
+        "parameters! However, we're missing an important piece of information: what the\n",
+        "prediction would have looked 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"
+        "We could just run two tracing contexts and compare the outputs. However, this would require two forward passes through the model. `NNsight` can do\n",
+        "better than that with batching.\n"
       ]
     },
     {
@@ -1603,44 +2057,35 @@
         "id": "_Gc-Zr6NmEOn"
       },
       "source": [
-        "It's time to bring back the `Tracer` object we dropped before. \n",
+        "Batching is a way to process multiple inputs in one forward pass. To better understand how batching works, we're going to bring back the `Tracer` object that we dropped before.\n",
+        "\n",
+        "When we call `.trace(...)`, it's actually creating two different contexts behind the scenes. The first one is the tracing context that we've discussed previously, and the second one is the invoker context. The invoker context defines the values of the `.input` and `.output` Proxies.\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",
+        "If we call `.trace(...)` with some input, the input is passed on to the invoker. As there is only one input, only one invoker context is created.\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",
+        "If we call `.trace()` without an input, then we can call `tracer.invoke(input1)` to manually create the invoker context with an input, `input1`. We can also repeatedly call `tracer.invoke(...)` to create the invoker context for additional inputs. Every subsequent time we call\n",
+        "`.invoke(...)`, interventions within its context will only refer to the input in that particular invoke statement.\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",
+        "When exiting the tracing context, the inputs from all of the invokers will be batched together, and they will be executed in one forward pass! To test this out, let's do the same ablation experiment, but also add a 'control' output for comparison:\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",
+        "Note that when injecting data to only the relevant invoker interventions, `nnsight` tries, but can't guarantee, to narrow the data into the right\n",
+        "batch indices. Thus, 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",
+        "Just like `.trace(...)` created a `Tracer` object, `.invoke(...)` creates an `Invoker` object. For `LanguageModel` 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",
+        "In a case where we pass 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",
+        "Keyword arguments given to `.invoke(..)` make their way to the input pre-processing.  \n",
+        "`LanguageModel` has keyword arguments `max_length` and `truncation` used for tokenization which can be\n",
+        "passed to the invoker. If we want to pass keyword arguments to the invoker for a single-input `.trace(...)`, we can pass `invoker_args` as a dictionary of invoker keyword arguments.\n",
         "\n",
         "Here is an example to demonstrate everything we've described:\n",
         "\n",
@@ -1667,15 +2112,22 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 29,
+      "execution_count": 91,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "Kdcq4oCNmEua",
-        "outputId": "e6d33c07-d22a-41af-b84a-91c9681c7524"
+        "outputId": "33702cb9-46df-4c60-f108-f7ad3c1cd446"
       },
       "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",
@@ -1705,6 +2157,7 @@
         "        # Get the output for only the original batch.\n",
         "        token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
         "\n",
+        "\n",
         "print(\"Original token IDs:\", token_ids_original)\n",
         "print(\"Modified token IDs:\", token_ids_intervention)\n",
         "\n",
@@ -1718,30 +2171,32 @@
         "id": "8BsybgsK2Bbr"
       },
       "source": [
-        "So it did end up affecting what the model predicted. That's pretty neat!\n",
+        "Based on our control results, our ablation 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"
+        "Another cool thing with multiple invokes is that Proxies can interact between them.\n",
+        "\n",
+        "Here, we transfer the token embeddings from a real prompt into another placeholder prompt. Therefore the latter prompt produces the output of the former prompt:"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 30,
+      "execution_count": 92,
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/"
+          "base_uri": "https://localhost:8080/",
+          "height": 519
         },
         "id": "syKTI_KhpvCY",
-        "outputId": "b2238c5c-d2d8-47bf-b73a-730f1e03cf99"
+        "outputId": "35afac9a-a34a-4e60-edc8-4c1b404e4389"
       },
       "outputs": [
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
+            "original prediction shape torch.Size([])\n",
             "Original prediction:  _\n",
+            "modified prediction shape torch.Size([])\n",
             "Modified prediction:  Paris\n"
           ]
         }
@@ -1750,30 +2205,38 @@
         "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",
+        "      token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
         "\n",
-        "        token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
-        "\n",
+        "print(\"original prediction shape\", token_ids_original[0][-1].shape)\n",
         "print(\"Original prediction:\", llm.tokenizer.decode(token_ids_original[0][-1]))\n",
+        "\n",
+        "print(\"modified prediction shape\", token_ids_intervention[0][-1].shape)\n",
         "print(\"Modified prediction:\", llm.tokenizer.decode(token_ids_intervention[0][-1]))"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_bXvUGGqbwv1"
+      },
+      "source": [
+        "## Multiple Token Generation"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "vvWA-CWqQtah"
       },
       "source": [
-        "## .next()\n",
+        "### .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",
@@ -1782,117 +2245,299 @@
         "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, so the modules\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",
+        "`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",
+        "Each module has a call index 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"
+        "graph only at the iteration that matches the call index.\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 31,
+      "execution_count": 93,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "Yy5Z9NE1GkaN",
-        "outputId": "0b4c1bd3-a6f6-4ba1-9cb8-6019c222311f"
+        "outputId": "3827231c-86b1-482e-bd2f-ff17191692e5"
       },
       "outputs": [
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
+            "Setting `pad_token_id` to `eos_token_id`:None 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"
+            "torch.Size([1, 10, 768])\n",
+            "torch.Size([1, 1, 768])\n",
+            "tensor([[ 464,  412,  733,  417, 8765,  318,  287,  262, 1748,  286, 6342,   11,\n",
+            "          290]], device='mps:0')\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",
+        "with llm.generate('The Eiffel Tower is in the city of', max_new_tokens=3) as tracer:\n",
         "\n",
-        "    token_ids_2 = llm.lm_head.next().output.argmax(dim=-1).save()\n",
+        "    hidden_states1 = llm.transformer.h[-1].output[0].save()\n",
         "\n",
-        "    token_ids_3 = llm.lm_head.next().output.argmax(dim=-1).save()\n",
+        "    # use module.next() to access the next intervention\n",
+        "    hidden_states2 = llm.transformer.h[-1].next().output[0].save()\n",
         "\n",
-        "    output = llm.generator.output.save()\n",
+        "    # saving the output allows you to save the hidden state across the initial prompt\n",
+        "    out = 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"
+        "print(hidden_states1.shape)\n",
+        "print(hidden_states2.shape)\n",
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "oHy6FqyNRBzm"
+      },
       "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",
+        "### New! using .all()\n",
         "\n",
-        "Let's explore how to use the **Editor** context to make a simple persistent change to a model:"
+        "With `nnsight 0.4` you can now use `.all()` to recursively apply interventions to a model. Calling `.all()` on a module within a model will recursively apply its `.input` and `.output` across all iterations. Previously, we'd need to loop across each new generated token, saving the intervention for every generated token and calling `.next()` to move forward."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 32,
-      "metadata": {},
+      "execution_count": 94,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "C1C1KmclTJDH",
+        "outputId": "97f4b5ba-3b94-4541-df98-4025b273ba97"
+      },
       "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
+          ]
+        },
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "\n",
-            "Original Prediction:   Rome\n",
-            "Modified Prediction:   Paris\n"
+            "Hidden state length:  3\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",
+        "# Old approach:\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = llm.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "hidden_states = []\n",
+        "with llm.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    for i in range(n_new_tokens):\n",
+        "        # Apply intervention - set first layer output to zero\n",
+        "        layers[0].output[0][:] = 0\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",
+        "        # Append desired hidden state post-intervention\n",
+        "        hidden_states.append(layers[-1].output.save())\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",
+        "        # Move to next generated token\n",
+        "        layers[0].next()\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tuBLYSx_XKXK"
+      },
+      "source": [
+        "We can use also `.all()` to streamline the multiple token generation process. We simply call `.all` on the module where we are applying the intervention (in this case GPT-2's layers), apply our intervention, and append our hidden states (stored in an `nnsight.list()` object).\n",
+        "<br> <br>\n",
+        "\n",
+        "Let's test this out for the multiple token generation case:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 95,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "IgxXlvxDX1zF",
+        "outputId": "9588942b-19b8-4d8b-90a3-9b8ef463b53f"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Hidden state length:  3\n"
+          ]
+        }
+      ],
+      "source": [
+        "# using .all():\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = llm.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with llm.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() to apply intervention to each new token\n",
+        "    layers.all()\n",
+        "\n",
+        "    # Apply intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(layers[-1].output) # no need to call .save\n",
+        "    # Don't need to loop or call .next()!\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Z3dMEb6HX-Wp"
+      },
+      "source": [
+        "Easy! Note that because `.all()` is recursive, it will only work to append outputs called on children of the module that `.all()` was called on. See example below for more information. TL;DR: apply `.all()` on the highest-level accessed module if interventions and outputs have different hierarchies within model structure."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3J_IeO96jVkZ"
+      },
+      "source": [
+        "<details>\n",
+        "<summary>Recursive properties of .all()</summary>\n",
+        "\n",
+        "`.all()` recursively acts on model components. In the below code example, only the first token generation is saved, because `.all()` applied to `layers`, while the saved variable `hidden_states` is produced from `model.lm_head`, which is not a child of `layers`.\n",
+        "\n",
+        "```\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() on layers\n",
+        "    layers.all()\n",
+        "\n",
+        "    # Apply same intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(model.lm_head.output) # no need to call .save, it's already initialized\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states)) # length is 1, meaning it only saved the first token generation\n",
+        "```\n",
+        "\n",
+        "If you want to apply an intervention during multiple token generation while saving the state of a model component that isn't a child of that module, you can instead apply `.all()` to the full model:\n",
+        "\n",
+        "```\n",
+        "prompt = 'The Eiffel Tower is in the city of'\n",
+        "layers = model.transformer.h\n",
+        "n_new_tokens = 3\n",
+        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
+        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
+        "\n",
+        "    # Call .all() on model\n",
+        "    model.all()\n",
+        "\n",
+        "    # Apply same intervention - set first layer output to zero\n",
+        "    layers[0].output[0][:] = 0\n",
+        "\n",
+        "    # Append desired hidden state post-intervention\n",
+        "    hidden_states.append(model.lm_head.output) # no need to call .save\n",
+        "\n",
+        "print(\"Hidden state length: \",len(hidden_states)) # length is 3, as expected!\n",
+        "```\n",
+        "\n",
+        "</details>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ty-luCwCkkw7"
+      },
+      "source": [
+        "## Model Editing"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hgrZVQLEkkw8"
+      },
+      "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": 96,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "-_oAprpwkkw8",
+        "outputId": "388b31e3-9585-4898-98be-7a1a043fa49c"
+      },
+      "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",
@@ -1903,17 +2548,25 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "S8J6cVw0kkw8"
+      },
       "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",
+        "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": {},
+      "execution_count": 97,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "b1atxBp8kkw8",
+        "outputId": "d176ce44-a861-4407-fbf3-5c4411bd4343"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1937,15 +2590,23 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "nVeZO96jkkw8"
+      },
       "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. "
+        "If you've made in-place edits to your model and need to revert these changes, you can apply `.clear_edits()`. This method removes all edits applied to the model, effectively restoring it to its original state."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 34,
-      "metadata": {},
+      "execution_count": 98,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "_SHK4JzPkkw8",
+        "outputId": "c6da4d26-e5b8-415e-aa01-071fbdbd4802"
+      },
       "outputs": [
         {
           "name": "stdout",
@@ -1970,7 +2631,7 @@
         "id": "rpkX-LwBQZHo"
       },
       "source": [
-        "# 3️⃣ I thought you said huge models?"
+        "# 3 I thought you said huge models?"
       ]
     },
     {
@@ -1979,10 +2640,9 @@
         "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",
+        "`NNsight` is only one part of our project to democratize access to AI internals. The other half is the National Deep Inference Fabric, or `NDIF`. `NDIF` hosts large models for shared access using `NNsight`, so you don't have to worry about any of the headaches of hosting large models yourself!\n",
         "\n",
-        "The interaction between the two is fairly straightforward. The\n",
+        "The interaction between `NDIF` and `NNsight` 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",
@@ -2008,15 +2668,15 @@
       },
       "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",
+        "run the rest of this walkthrough, you need one of your own. To get one, simply\n",
+        "register at https://login.ndif.us.\n",
         "\n",
-        "With a valid API key, you then can configure `nnsight` by doing the following:\n"
+        "With a valid API key, you then can configure `nnsight` as follows:\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 35,
+      "execution_count": 100,
       "metadata": {
         "id": "ax1NWoS9MJeZ"
       },
@@ -2024,7 +2684,7 @@
       "source": [
         "from nnsight import CONFIG\n",
         "\n",
-        "CONFIG.set_default_api_key(\"YOUR_API_KEY\") \n"
+        "CONFIG.set_default_api_key(\"YOUR_API_KEY\")"
       ]
     },
     {
@@ -2033,16 +2693,16 @@
         "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",
+        "If you're running in a local IDE, this only needs to be run once as it will save the API key as the default in a\n",
+        ".config file along with your `nnsight` installation. You can also add your API key to Google Colab secrets.\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"
+        "context with `Llama-3.1-8b`!\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 36,
+      "execution_count": 101,
       "metadata": {
         "id": "Q1N04sJPZnJt"
       },
@@ -2050,51 +2710,56 @@
       "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"
+        "# Llama 3.1 70b is a gated model, so you need to apply for access on HuggingFace and include your token.\n",
+        "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\""
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 37,
+      "execution_count": 8,
       "metadata": {
-        "id": "oqVzjNoyNGc7"
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "oqVzjNoyNGc7",
+        "outputId": "abf3f296-c41b-475a-946b-6e0b046333e8"
       },
       "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",
+            "2025-01-31 15:47:57,111 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:47:57,317 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:47:57,519 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:47:58,468 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - COMPLETED: Your job has been completed.\n",
+            "Downloading result:   0%|          | 0.00/4.37M [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]"
+            "Downloading result: 100%|██████████| 4.37M/4.37M [00:00<00:00, 16.5MB/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",
+            "(tensor([[[ 1.7734,  2.6875,  0.8047,  ..., -1.8594,  2.2344,  3.2656],\n",
+            "         [ 0.0312, -0.0352, -2.8750,  ..., -0.8906, -0.0547,  1.6172],\n",
+            "         [ 1.3594, -2.0156,  1.7031,  ..., -1.7031, -0.7422,  1.4375],\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",
+            "         [ 1.0000,  0.3203, -0.2656,  ..., -0.0723, -0.2559,  0.2090],\n",
+            "         [ 0.4707, -0.3496,  0.2422,  ...,  0.7344, -0.0078,  0.1133],\n",
+            "         [-0.0566, -0.3496,  0.4746,  ...,  0.9844,  0.6797, -0.8750]]],\n",
+            "       dtype=torch.bfloat16),)\n",
+            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
+            "         [-2.4375, -1.7266, -2.0156,  ..., -9.1250, -9.1250, -9.1250],\n",
+            "         [ 9.6875,  4.5625,  5.8750,  ..., -3.3906, -3.3906, -3.3906],\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"
+            "         [ 2.3281,  1.0703, -0.3203,  ..., -7.1562, -7.1562, -7.1562],\n",
+            "         [11.1875,  6.0312,  4.9062,  ..., -3.5156, -3.5156, -3.5156],\n",
+            "         [ 8.0000,  5.2500,  4.3750,  ..., -3.9844, -3.9844, -3.9844]]],\n",
+            "       dtype=torch.bfloat16)\n"
           ]
         },
         {
@@ -2106,9 +2771,10 @@
         }
       ],
       "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",
+        "from nnsight import LanguageModel\n",
         "\n",
+        "# We'll never actually load the parameters locally, so no need to specify a device_map.\n",
+        "llama = LanguageModel(\"meta-llama/Meta-Llama-3.1-8B\")\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",
@@ -2133,39 +2799,48 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "lfRDZox5kkw8"
+      },
       "source": [
         "## Sessions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "33XLOvtukkw8"
+      },
       "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",
+        "This offers the following benefits:\n",
+        "1.   All interventions within a session will be executed one after another without additional wait in the NDIF queue\n",
+        "2.   All intermediate outputs for 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": {},
+      "execution_count": 9,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "KA6GROiTR3gf",
+        "outputId": "dfefe423-9caf-48ce-cfe5-83c2f808a6b2"
+      },
       "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]"
+            "2025-01-31 15:48:17,378 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:17,558 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:17,755 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:18,205 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 1.62k/1.62k [00:00<00:00, 6.57MB/s]"
           ]
         },
         {
@@ -2174,7 +2849,7 @@
           "text": [
             "\n",
             "T1 - Original Prediction:   Paris\n",
-            "T2 - Modified Prediction:   Paris\n"
+            "T2 - Modified Prediction:  ://\n"
           ]
         },
         {
@@ -2189,12 +2864,12 @@
         "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",
+        "    # capture the hidden state from layer 32 at the last token\n",
+        "    hs_31 = llama.model.layers[31].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",
+        "    llama.model.layers[1].output[0][:, -1, :] = hs_31[:]\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",
@@ -2203,654 +2878,2438 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "End2DR7qkkw8"
+      },
       "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.)."
+        "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": {},
+      "metadata": {
+        "id": "M-Nq8q0Qkkw8"
+      },
       "source": [
         "The `session` context object has its own methods to log values and be terminated early."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 39,
-      "metadata": {},
+      "execution_count": 104,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 311
+        },
+        "id": "Q3c6M5Jwkkw8",
+        "outputId": "0d937244-d022-4745-f3ea-f10b305bc3e3"
+      },
       "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"
+            "2025-01-31 15:17:51,636 3d5144b4-e0b6-4106-912f-5abec8e78e7f - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:17:51,804 3d5144b4-e0b6-4106-912f-5abec8e78e7f - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:17:52,023 3d5144b4-e0b6-4106-912f-5abec8e78e7f - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:17:52,028 3d5144b4-e0b6-4106-912f-5abec8e78e7f - LOG: -- Early Stop --\n",
+            "2025-01-31 15:17:52,364 3d5144b4-e0b6-4106-912f-5abec8e78e7f - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 7.25MB/s]\n"
           ]
         }
       ],
       "source": [
         "with llama.session(remote=True) as session:\n",
-        "\n",
         "  session.log(\"-- Early Stop --\")\n",
-        "  session.exit()"
+        "  nnsight.stop"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "5Z-tTpQekkw8"
+      },
       "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 could swap activations between vanilla and instruct versions of the Llama model and compare their outputs. The `session` context can also be used to run experiments entirely locally! "
+        "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 could swap activations between base and instruct versions of the Llama model and compare their outputs. And `session` can also be used to run similar experiments entirely locally!"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "H6hxSZ8zTFvT"
+      },
       "source": [
-        "## Looping"
+        "## Streaming"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "4bDUYaluTRCM"
+      },
       "source": [
-        "We mention earlier that the `session` context enables multi-tracing execution. But how could we optimize a process that requires running an intervention graph in a loop? If we create a simple `for` loop with a **Tracer context** inside, this would create 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",
+        "Streaming enables users apply functions and datasets locally during remote model execution. This allows users to stream results for immediate consumption (i.e., seeing tokens as they are generated) or applying non-whitelisted functions such as model tokenizers, large local datasets, and more!\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:"
+        "*   `nnsight.local()` context sends values immediately to user's local machine from server\n",
+        "*   Intervention graph is executed locally on downstream nodes\n",
+        "*   Exiting local context uploads data back to server\n",
+        "*   `@nnsight.trace` function decorator enables custom functions to be added to intervention graph when using `nnsight.local()`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fBDttzi7TT26"
+      },
+      "source": [
+        "### `nnsight.local()`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NRsfBs5MTW6f"
+      },
+      "source": [
+        "You may sometimes want to locally access and manipulate values during remote execution. Using `.local()` on a proxy, you can send remote content to your local machine and apply local functions. The intervention graph is then executed locally on downstream nodes (until you send execution back to the remote server by exiting the `.local()` context).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gN8yWbc4TXZc"
+      },
+      "source": [
+        "There are a few use cases for streaming with `.local()`, including live chat generation and applying large datasets or non-whitelisted local functions to the intervention graph."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fFxwkDC6TZ8C"
+      },
+      "source": [
+        "Now let's explore how streaming works. We'll start by grabbing some hidden states of the model and printing their value using `tracer.log()`. Without calling `nnsight.local()`, these operations will all occur remotely."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 40,
-      "metadata": {},
+      "execution_count": 120,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "QKMwwmoUTcCI",
+        "outputId": "f1400168-1f44-43c4-91c1-14deaeedb4d0"
+      },
       "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"
+            "2025-01-31 15:21:29,379 94c37060-d3f6-4f21-b55f-0e72f03d08da - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:29,592 94c37060-d3f6-4f21-b55f-0e72f03d08da - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:29,786 94c37060-d3f6-4f21-b55f-0e72f03d08da - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:21:29,822 94c37060-d3f6-4f21-b55f-0e72f03d08da - LOG: tensor(1.7656, device='cuda:0')\n",
+            "2025-01-31 15:21:30,161 94c37060-d3f6-4f21-b55f-0e72f03d08da - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 514k/514k [00:00<00:00, 4.80MB/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
+            "         [10.2500,  2.1094,  2.8281,  ..., -8.2500, -8.2500, -8.2500]]],\n",
+            "       dtype=torch.bfloat16)\n"
           ]
         }
       ],
       "source": [
-        "with llama.session(remote=True) as session:\n",
+        "# This will give you a remote LOG response because it's coming from the remote server\n",
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
         "\n",
-        "  with session.iter([0, 1, 2]) as item:\n",
-        "    # define intervention body here ...\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
         "\n",
-        "    with llama.trace(\"_\"):\n",
-        "      # define interventions here ...\n",
-        "      pass\n",
+        "    tracer.log(hs[0,0,0])\n",
         "\n",
-        "    with llama.trace(\"_\"):\n",
-        "      # define interventions here ...\n",
-        "      pass"
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "1IjPUn8oTouV"
+      },
       "source": [
-        "The `Iterator` context extends all the `nnsight` graph-based functionalities but also closely mimics the conventional `for` loop statement in Python, allowing it to support all kind of iterative operations with a use of `as item` syntax:"
+        "Now, let's try the same operation using the `nnsight.local()` context. This will send the operations to get and print the hidden states to your local machine, changing how the logging message is formatted (local formatting instead of remote)."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 41,
-      "metadata": {},
+      "execution_count": 121,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "x9byLVl_Tqy-",
+        "outputId": "e1514fdd-61b5-4cbe-c577-1bdab9605adc"
+      },
       "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]"
+            "2025-01-31 15:21:32,674 0414b4df-a160-485b-89dd-4e70a05a43e2 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:32,830 0414b4df-a160-485b-89dd-4e70a05a43e2 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:33,092 0414b4df-a160-485b-89dd-4e70a05a43e2 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:21:33,352 0414b4df-a160-485b-89dd-4e70a05a43e2 - COMPLETED: Your job has been completed.\n"
           ]
         },
         {
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "\n",
-            "List:  [0, 1, 2]\n"
+            "tensor(1.7656, dtype=torch.bfloat16)\n"
           ]
         },
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "\n"
+            "Downloading result: 100%|██████████| 514k/514k [00:00<00:00, 4.28MB/s]"
           ]
-        }
-      ],
-      "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",
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
+            "         [10.2500,  2.1094,  2.8281,  ..., -8.2500, -8.2500, -8.2500]]],\n",
+            "       dtype=torch.bfloat16)\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        }
+      ],
+      "source": [
+        "# This will print locally because it's already local\n",
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
         "\n",
-        "print(\"\\nList: \", li2)"
+        "    with nnsight.local():\n",
+        "        hs = llama.model.layers[-1].output[0]\n",
+        "        tracer.log(hs[0,0,0])\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "tNbmIZSCTwo2"
+      },
       "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>"
+        "### `@nnsight.trace` function decorator"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "xiQOxe9_TzqI"
+      },
       "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:"
+        "We can also use function decorators to create custom functions to be used during `.local` calls. This is a handy way to enable live streaming of a chat or to train probing classifiers on model hidden states."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "X7A3IjhLT2TE"
+      },
+      "source": [
+        "Let's try out `@nnsight.trace` and `nnsight.local()` to access a custom function during remote execution."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 42,
-      "metadata": {},
+      "execution_count": 122,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "P8ghjV-7TzPj",
+        "outputId": "6e5a2efb-8edd-40ff-a4de-84c535bbbd3d"
+      },
       "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"
+            "2025-01-31 15:21:35,853 3eeee112-3074-4b3e-adea-001a85ce554d - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:36,035 3eeee112-3074-4b3e-adea-001a85ce554d - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:36,258 3eeee112-3074-4b3e-adea-001a85ce554d - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:21:36,653 3eeee112-3074-4b3e-adea-001a85ce554d - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 258k/258k [00:00<00:00, 2.91MB/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()"
+        "# first, let's define our function\n",
+        "@nnsight.trace # decorator that enables this function to be added to the intervention graph\n",
+        "def my_local_fn(value):\n",
+        "    return value * 0\n",
+        "\n",
+        "# We use a local function to ablate some hidden states\n",
+        "# This downloads the data for the .local context, and then uploads it back to set the value.\n",
+        "with llama.generate(\"hello\", remote=True) as tracer:\n",
+        "\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
+        "\n",
+        "    with nnsight.local():\n",
+        "\n",
+        "        hs = my_local_fn(hs)\n",
+        "\n",
+        "    llama.model.layers[-1].output[0][:] = hs\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "_7d5AdU6T72j"
+      },
+      "source": [
+        "Note that without calling `.local`, the remote API does not know about `my_local_fn` and will throw a whitelist error. A whitelist error occurs because you are being allowed access to the function."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 123,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 472
+        },
+        "id": "n4OIIqQhT98b",
+        "outputId": "54e6ad3d-a882-4892-aac0-7bdba52cac40"
+      },
+      "outputs": [
+        {
+          "ename": "FunctionWhitelistError",
+          "evalue": "Function with name `__main__.my_local_fn` not in function whitelist.",
+          "output_type": "error",
+          "traceback": [
+            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+            "\u001b[0;31mFunctionWhitelistError\u001b[0m                    Traceback (most recent call last)",
+            "Cell \u001b[0;32mIn[123], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m llama\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhello\u001b[39m\u001b[38;5;124m\"\u001b[39m, remote\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m tracer:\n\u001b[1;32m      3\u001b[0m     hs \u001b[38;5;241m=\u001b[39m llama\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39mlayers[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m.\u001b[39moutput[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m      5\u001b[0m     hs \u001b[38;5;241m=\u001b[39m my_local_fn(hs) \u001b[38;5;66;03m# no .local - will cause an error\u001b[39;00m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/contexts/interleaving.py:96\u001b[0m, in \u001b[0;36mInterleavingTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     92\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     94\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m---> 96\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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/tracer.py:25\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mglobals\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m GlobalTracingContext\n\u001b[1;32m     23\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mtry_deregister(\u001b[38;5;28mself\u001b[39m)\n\u001b[0;32m---> 25\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \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~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/base.py:82\u001b[0m, in \u001b[0;36mContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     78\u001b[0m graph \u001b[38;5;241m=\u001b[39m graph\u001b[38;5;241m.\u001b[39mstack\u001b[38;5;241m.\u001b[39mpop()\n\u001b[1;32m     80\u001b[0m graph\u001b[38;5;241m.\u001b[39malive \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m---> 82\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[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:77\u001b[0m, in \u001b[0;36mRemoteBackend.__call__\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m     72\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, graph: Graph):\n\u001b[1;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mblocking:\n\u001b[1;32m     75\u001b[0m \n\u001b[1;32m     76\u001b[0m         \u001b[38;5;66;03m# Do blocking request.\u001b[39;00m\n\u001b[0;32m---> 77\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mblocking_request\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     80\u001b[0m \n\u001b[1;32m     81\u001b[0m         \u001b[38;5;66;03m# Otherwise we are getting the status / result of the existing job.\u001b[39;00m\n\u001b[1;32m     82\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnon_blocking_request(graph)\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:289\u001b[0m, in \u001b[0;36mRemoteBackend.blocking_request\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m    280\u001b[0m sio\u001b[38;5;241m.\u001b[39mconnect(\n\u001b[1;32m    281\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mws_address,\n\u001b[1;32m    282\u001b[0m     socketio_path\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/ws/socket.io\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    283\u001b[0m     transports\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwebsocket\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m    284\u001b[0m     wait_timeout\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m,\n\u001b[1;32m    285\u001b[0m )\n\u001b[1;32m    287\u001b[0m remote_graph \u001b[38;5;241m=\u001b[39m preprocess(graph)\n\u001b[0;32m--> 289\u001b[0m data, headers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrequest\u001b[49m\u001b[43m(\u001b[49m\u001b[43mremote_graph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    291\u001b[0m headers[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msession_id\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m sio\u001b[38;5;241m.\u001b[39msid\n\u001b[1;32m    293\u001b[0m \u001b[38;5;66;03m# Submit request via\u001b[39;00m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:60\u001b[0m, in \u001b[0;36mRemoteBackend.request\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m     58\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrequest\u001b[39m(\u001b[38;5;28mself\u001b[39m, graph: Graph) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[\u001b[38;5;28mbytes\u001b[39m, Dict[\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mstr\u001b[39m]]:\n\u001b[0;32m---> 60\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[43mRequestModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mserialize\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mformat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mzlib\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     62\u001b[0m     headers \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m     63\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_key\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel_key,\n\u001b[1;32m     64\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mformat\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mformat,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     67\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msent-timestamp\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mstr\u001b[39m(time\u001b[38;5;241m.\u001b[39mtime()),\n\u001b[1;32m     68\u001b[0m     }\n\u001b[1;32m     70\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m data, headers\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/request.py:43\u001b[0m, in \u001b[0;36mRequestModel.serialize\u001b[0;34m(graph, format, _zlib)\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mserialize\u001b[39m(graph: Graph, \u001b[38;5;28mformat\u001b[39m:\u001b[38;5;28mstr\u001b[39m, _zlib:\u001b[38;5;28mbool\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mbytes\u001b[39m:\n\u001b[1;32m     41\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mformat\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mjson\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 43\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[43mRequestModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     45\u001b[0m         json \u001b[38;5;241m=\u001b[39m data\u001b[38;5;241m.\u001b[39mmodel_dump(mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mjson\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     47\u001b[0m         data \u001b[38;5;241m=\u001b[39m msgspec\u001b[38;5;241m.\u001b[39mjson\u001b[38;5;241m.\u001b[39mencode(json)\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/request.py:30\u001b[0m, in \u001b[0;36mRequestModel.__init__\u001b[0;34m(self, memo, *args, **kwargs)\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, memo: Dict \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m---> 30\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__init__\u001b[39;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[43mmemo\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmemo\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\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     32\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m memo \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     34\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmemo \u001b[38;5;241m=\u001b[39m {\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mMEMO}\n",
+            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:276\u001b[0m, in \u001b[0;36mGraphModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    273\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m    274\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value: Graph) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Self:\n\u001b[0;32m--> 276\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mGraphModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnodes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnodes\u001b[49m\u001b[43m)\u001b[49m\n",
+            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:77\u001b[0m, in \u001b[0;36mmemoized.<locals>.inner\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minner\u001b[39m(value):\n\u001b[0;32m---> 77\u001b[0m     model \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m     _id \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mid\u001b[39m(value)\n\u001b[1;32m     81\u001b[0m     MEMO[_id] \u001b[38;5;241m=\u001b[39m model\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:101\u001b[0m, in \u001b[0;36mNodeModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m     98\u001b[0m \u001b[38;5;129m@memoized\u001b[39m\n\u001b[1;32m     99\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value: Node) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Self:\n\u001b[0;32m--> 101\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mNodeModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtarget\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
+            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:244\u001b[0m, in \u001b[0;36mFunctionModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    239\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m    240\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value:FUNCTION):\n\u001b[1;32m    242\u001b[0m     model \u001b[38;5;241m=\u001b[39m FunctionModel(function_name\u001b[38;5;241m=\u001b[39mget_function_name(value))\n\u001b[0;32m--> 244\u001b[0m     \u001b[43mFunctionModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheck_function_whitelist\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunction_name\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    246\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m model\n",
+            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:251\u001b[0m, in \u001b[0;36mFunctionModel.check_function_whitelist\u001b[0;34m(cls, qualname)\u001b[0m\n\u001b[1;32m    248\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m    249\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcheck_function_whitelist\u001b[39m(\u001b[38;5;28mcls\u001b[39m, qualname: \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mstr\u001b[39m:\n\u001b[1;32m    250\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m qualname \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m FUNCTIONS_WHITELIST:\n\u001b[0;32m--> 251\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m FunctionWhitelistError(\n\u001b[1;32m    252\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFunction with name `\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mqualname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m` not in function whitelist.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    253\u001b[0m         )\n\u001b[1;32m    255\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m qualname\n",
+            "\u001b[0;31mFunctionWhitelistError\u001b[0m: Function with name `__main__.my_local_fn` not in function whitelist."
+          ]
+        }
+      ],
       "source": [
-        "The **Iterator** context is a nice piece of functionality that allows you to define a bunch of basic code operations that can now be \"traceable\" by `nnsight`.\n",
+        "with llama.trace(\"hello\", remote=True) as tracer:\n",
         "\n",
-        "But what kinds of experimental scenarios would make this useful?\n",
+        "    hs = llama.model.layers[-1].output[0]\n",
         "\n",
-        "In the next section, we see how the `Iterator` context enables a powerful parameter-efficient fine-tuning tool, Low-Rank Adaptation (LoRA)!"
+        "    hs = my_local_fn(hs) # no .local - will cause an error\n",
+        "\n",
+        "    llama.model.layers[-1].output[0][:] = hs * 2\n",
+        "\n",
+        "    out =  llama.lm_head.output.save()\n",
+        "\n",
+        "print(out)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "v_bdhze-UBec"
+      },
       "source": [
-        "## Training a LoRA"
+        "### Example: Live-streaming remote chat\n",
+        "\n"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "hguvYgvVUE4b"
+      },
+      "source": [
+        "Now that we can access data within the tracing context on our local computer, we can apply non-whitelisted functions, such as the model's tokenizer, within our tracing context.\n",
+        "\n",
+        "Let's build a decoding function that will decode tokens into words and print the result."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 124,
+      "metadata": {
+        "id": "8OAXweC5UHBJ"
+      },
+      "outputs": [],
       "source": [
-        "For more background information about LoRA, check out [this paper](https://arxiv.org/abs/2106.09685)."
+        "@nnsight.trace\n",
+        "def my_decoding_function(tokens, model, max_length=80, state=None):\n",
+        "    # Initialize state if not provided\n",
+        "    if state is None:\n",
+        "        state = {'current_line': '', 'current_line_length': 0}\n",
+        "\n",
+        "    token = tokens[-1] # only use last token\n",
+        "\n",
+        "    # Decode the token\n",
+        "    decoded_token = llama.tokenizer.decode(token).encode(\"unicode_escape\").decode()\n",
+        "\n",
+        "    if decoded_token == '\\\\n':  # Handle explicit newline tokens\n",
+        "        # Print the current line and reset state\n",
+        "        print('',flush=True)\n",
+        "        state['current_line'] = ''\n",
+        "        state['current_line_length'] = 0\n",
+        "    else:\n",
+        "        # Check if adding the token would exceed the max length\n",
+        "        if state['current_line_length'] + len(decoded_token) > max_length:\n",
+        "            print('',flush=True)\n",
+        "            state['current_line'] = decoded_token  # Start a new line with the current token\n",
+        "            state['current_line_length'] = len(decoded_token)\n",
+        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
+        "        else:\n",
+        "            # Add a space if the line isn't empty and append the token\n",
+        "            if state['current_line']:\n",
+        "                state['current_line'] += decoded_token\n",
+        "            else:\n",
+        "                state['current_line'] = decoded_token\n",
+        "            state['current_line_length'] += len(decoded_token)\n",
+        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
+        "\n",
+        "    return state"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "Uf-Ir1YgUKJt"
+      },
       "source": [
-        "Training a LoRA will integrate 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."
+        "Now we can decode and print our model outputs throughout token generation by accessing our decoding function through `nnsight.local()`."
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 43,
-      "metadata": {},
-      "outputs": [],
+      "execution_count": 125,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "aT56JuDzULyJ",
+        "outputId": "bfa7593d-01c2-46e0-8c67-d8dbc1c7aa93"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Prompt:  A press release is an official statement delivered to members of the news media for the purpose of \n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:21:48,944 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:21:49,121 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:21:49,363 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - RUNNING: Your job has started running.\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            " providing information, an official statement, or making an announcement.A press release is also written or recorded communication directed at members of the news media for the purpose of announcing something ostensibly newsworthy. Typically, they are mailed, faxed, or e"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:21:52,277 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 258k/258k [00:00<00:00, 3.03MB/s]\n"
+          ]
+        }
+      ],
       "source": [
         "import torch\n",
         "\n",
-        "from nnsight.envoy import Envoy # \n",
+        "nnsight.CONFIG.APP.REMOTE_LOGGING = False\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",
+        "prompt = \"A press release is an official statement delivered to members of the news media for the purpose of\"\n",
+        "# prompt = \"Your favorite board game is\"\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",
+        "print(\"Prompt: \",prompt,'\\n', end =\"\")\n",
         "\n",
-        "    # The Call method defines how to actually apply the LORA.\n",
-        "    def __call__(self, alpha: float = 1.0):\n",
-        "        \"\"\"Call.\n",
+        "# Initialize the state for decoding\n",
+        "state = {'current_line': '', 'current_line_length': 0}\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",
+        "with llama.generate(prompt, remote=True, max_new_tokens = 50) as generator:\n",
+        "    # Call .all() to apply to each new token\n",
+        "    llama.all()\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",
+        "    all_tokens = nnsight.list().save()\n",
         "\n",
-        "        # LORA is additive\n",
-        "        h = BA_x + self.module.output\n",
+        "    # Access model output\n",
+        "    out = llama.lm_head.output.save()\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",
+        "    # Apply softmax to obtain probabilities and save the result\n",
+        "    probs = torch.nn.functional.softmax(out, dim=-1)\n",
+        "    max_probs = torch.max(probs, dim=-1)\n",
+        "    tokens = max_probs.indices.cpu().tolist()\n",
+        "    all_tokens.append(tokens[0]).save()\n",
         "\n",
-        "    def parameters(self):\n",
-        "        # Some way to get all the parameters.\n",
-        "        return [self.WA, self.WB]"
+        "    with nnsight.local():\n",
+        "        state = my_decoding_function(tokens[0], llama, max_length=20, state=state)"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "Let's define all the variables to use in LoRA training. "
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 44,
-      "metadata": {},
-      "outputs": [],
+      "metadata": {
+        "id": "DUCWX1BDkkw9"
+      },
       "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"
+        "## Looping across sessions"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "mWifpUjfkkw9"
+      },
       "source": [
-        "We can use the `.scan()` method to get the shape of the module without having to fully run the model."
+        "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, 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": 46,
-      "metadata": {},
+      "execution_count": 126,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "to-hAyuKkkw9",
+        "outputId": "ab21cb42-adba-4630-de5e-639cb6af3a35"
+      },
       "outputs": [
         {
-          "name": "stdout",
+          "name": "stderr",
           "output_type": "stream",
           "text": [
-            "8192\n"
+            "2025-01-31 15:22:03,714 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:22:03,869 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:22:04,120 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:22:05,086 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 1.82MB/s]\n"
           ]
         }
       ],
       "source": [
-        "with llama.scan(\" \"):\n",
-        "    dim = module.output.shape[-1]\n",
+        "with llama.session(remote=True) as session:\n",
         "\n",
-        "print(dim)"
+        "  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": {},
+      "metadata": {
+        "id": "L14XEU8jkkw9"
+      },
       "source": [
-        "It's time to run the LoRA training loop! We will be using the **Session** and the **Iterator** contexts to achieve this."
+        "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": 47,
-      "metadata": {},
+      "execution_count": 127,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Rd-PcFJakkw9",
+        "outputId": "a5b20d0c-5f0d-4065-a874-68bfe9ec696e"
+      },
       "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"
+            "2025-01-31 15:22:07,235 570729dd-1d6d-4cd6-9876-9a096378119c - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:22:07,440 570729dd-1d6d-4cd6-9876-9a096378119c - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:22:07,674 570729dd-1d6d-4cd6-9876-9a096378119c - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:22:10,270 570729dd-1d6d-4cd6-9876-9a096378119c - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 1.47MB/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "List:  [0, 1, 2]\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\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",
+        "  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",
-        "    # Create a dataloader from it.\n",
-        "    dataloader = DataLoader(dataset, batch_size=10)\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",
-        "    # Create our LORA on the last mlp\n",
-        "    lora = LORA(module, dim, lora_dim)\n",
+        "print(\"\\nList: \", li2)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "CZTVwF7Fkkw9"
+      },
+      "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",
-        "    # Create an optimizer. Use the parameters from LORA\n",
-        "    optimizer = torch.optim.AdamW(lora.parameters(), lr=3)\n",
+        "<details>\n",
+        "<summary>A full list of NNsight Built-ins</summary>\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",
+        "`nnsight.bool()` creates a traceable Boolean\n",
         "\n",
-        "        # Run .trace with prompt\n",
-        "        with llama.trace(prompt) as tracer:\n",
+        "`nnsight.bytes()` creates a traceable Bytes\n",
         "\n",
-        "            # Apply LORA to intervention graph just by calling it with .trace\n",
-        "            lora()\n",
+        "`nnsight.int()` creates a traceable Integer\n",
         "\n",
-        "            # Get logits\n",
-        "            logits = llama.lm_head.output\n",
+        "`nnsight.float()` creates a traceable Float\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",
+        "`nnsight.str()` creates a traceable String\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"
+        "`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": {},
+      "metadata": {
+        "id": "pDF_eijqkkw9"
+      },
       "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. "
+        "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": 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"
-          ]
+      "execution_count": 128,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 365
         },
+        "id": "IhnG03yckkw9",
+        "outputId": "fe9ef3ca-e1b3-474d-b15f-f08d5dcd7e90"
+      },
+      "outputs": [
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "\n"
+            "2025-01-31 15:22:12,730 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:22:12,900 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:22:13,237 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:22:13,240 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 0\n",
+            "2025-01-31 15:22:13,242 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 2\n",
+            "2025-01-31 15:22:13,242 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 1\n",
+            "2025-01-31 15:22:13,759 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 4.74MB/s]\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",
+        "with llama.session(remote=True) as session:\n",
         "\n",
-        "# Then without. Should produce \"Hello,\"\n",
-        "with llama.generate(\"Hello\", remote=True) as generator:\n",
+        "  # with session.iter([0, 1, 2, 3], return_context=True) as (item, iterator):\n",
+        "  with session.iter([0, 1, 2, 3]) as item:\n",
         "\n",
-        "    out = llama.generator.output.save()\n",
+        "      nnsight.log(item)\n",
         "\n",
-        "print(llama.tokenizer.batch_decode(out.value))\n"
+        "      with nnsight.cond(item == 2):\n",
+        "        nnsight.stop()"
       ]
     },
     {
       "cell_type": "markdown",
-      "metadata": {},
+      "metadata": {
+        "id": "VvoZz5OGkkw9"
+      },
       "source": [
-        "# Next Steps\n",
+        "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",
-        "Check out [nnsight.net/tutorials](https://nnsight.net/tutorials) for more walkthroughs\n",
-        "implementating classic interpretability techniques using `nnsight`. "
+        "But in what kind of experimental scenario would someone need iterators?\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": {
-        "id": "P3zRm-7VRRov"
+        "id": "1OcolbxAkkw9"
       },
       "source": [
-        "# Getting Involved!\n"
+        "## Training a LoRA"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "BnCc9xgjvxEP"
+        "id": "jzyt5d6Dkkw9"
+      },
+      "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. Using session and iterator contexts, we're going apply a very simple fine-tuning approach called low-rank adaptation (LoRA).\n",
+        "\n",
+        "Let's try training a LoRA that, when applied, makes our model always predict \"Paris\" no matter what."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 10,
+      "metadata": {
+        "id": "e9YlPQ-lkkw9"
       },
+      "outputs": [],
       "source": [
-        "Note that both `nnsight` and `NDIF` are in active development, \n",
-        "so changes may be made and errors may arise during use.\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import nnsight\n",
+        "# from nnsight.envoy import Envoy # this moved in 0.4\n",
+        "from nnsight 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",
-        "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). We'd love to hear about your work using `nnsight`! \n",
+        "        # LORA is additive\n",
+        "        h = BA_x + self.module.output\n",
         "\n",
-        "You can also follow us on [LinkedIn](https://www.linkedin.com/company/national-deep-inference-fabric) and X/Twitter: [@ndif_team](https://x.com/ndif_team).\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]"
       ]
-    }
-  ],
-  "metadata": {
-    "colab": {
-      "provenance": [],
-      "toc_visible": true
     },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_Ub6ApqAkkw9"
+      },
+      "source": [
+        "Let's define all the variables to use in LoRA training."
+      ]
     },
-    "language_info": {
-      "codemirror_mode": {
-        "name": "ipython",
-        "version": 3
+    {
+      "cell_type": "code",
+      "execution_count": 11,
+      "metadata": {
+        "id": "yToy3qp1kkw9"
       },
-      "file_extension": ".py",
-      "mimetype": "text/x-python",
-      "name": "python",
-      "nbconvert_exporter": "python",
-      "pygments_lexer": "ipython3",
-      "version": "3.10.13"
+      "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"
+      ]
     },
-    "widgets": {
-      "application/vnd.jupyter.widget-state+json": {
-        "0a676b8f2dc34fed9731f19a03e4bfa1": {
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NNriyaBikkw9"
+      },
+      "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": 12,
+      "metadata": {
+        "id": "emv134jkkkw9"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "4096\n"
+          ]
+        }
+      ],
+      "source": [
+        "with llama.scan(\" \"):\n",
+        "    dim = module.output.shape[-1]\n",
+        "\n",
+        "print(dim)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "63aWV_ZTkkw9"
+      },
+      "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": 13,
+      "metadata": {
+        "id": "Q2SYDaawkkw9"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:48:38,456 810d88f8-f1fc-43fd-b16d-77e11841eecd - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:38,634 810d88f8-f1fc-43fd-b16d-77e11841eecd - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:38,859 810d88f8-f1fc-43fd-b16d-77e11841eecd - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:39,402 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.7539,  0.0732, -0.4023, -0.0615],\n",
+            "        [-0.5312,  0.5391,  0.5195, -1.1328],\n",
+            "        [-1.5781, -0.2480,  0.6953,  0.3535],\n",
+            "        ...,\n",
+            "        [ 0.1670,  0.5391, -0.5703, -0.6289],\n",
+            "        [-0.0762, -1.3438,  0.8320,  1.2656],\n",
+            "        [ 0.2500,  1.2188,  0.2891, -1.2578]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,410 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.7305,  0.0713, -0.3906, -0.0598],\n",
+            "        [-0.5156,  0.5234,  0.5039, -1.1016],\n",
+            "        [-1.5312, -0.2402,  0.6758,  0.3438],\n",
+            "        ...,\n",
+            "        [ 0.1621,  0.5234, -0.5547, -0.6094],\n",
+            "        [-0.0737, -1.3047,  0.8086,  1.2266],\n",
+            "        [ 0.2422,  1.1797,  0.2812, -1.2188]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,448 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.7070,  0.0688, -0.3789, -0.0586],\n",
+            "        [-0.5000,  0.5078,  0.4883, -1.0703],\n",
+            "        [-1.4844, -0.2334,  0.6562,  0.3340],\n",
+            "        ...,\n",
+            "        [ 0.1572,  0.5078, -0.5391, -0.5898],\n",
+            "        [-0.0718, -1.2656,  0.7852,  1.1875],\n",
+            "        [ 0.2344,  1.1406,  0.2734, -1.1797]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,485 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[ 0.6836,  0.0645, -0.3633, -0.0598],\n",
+            "        [-0.4883,  0.4902,  0.4746, -1.0391],\n",
+            "        [-1.4375, -0.2285,  0.6406,  0.3203],\n",
+            "        ...,\n",
+            "        [ 0.1543,  0.4941, -0.5234, -0.5703],\n",
+            "        [-0.0718, -1.2266,  0.7617,  1.1484],\n",
+            "        [ 0.2227,  1.1094,  0.2695, -1.1484]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,521 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-0.9648, -1.5703,  0.6094, -1.6875],\n",
+            "        [-2.0938, -1.1562,  1.3906, -2.6406],\n",
+            "        [-3.0156, -1.8516,  1.5938, -1.3047],\n",
+            "        ...,\n",
+            "        [ 1.7578,  2.0781, -1.1406,  1.0625],\n",
+            "        [-1.6875, -2.7969,  1.4688, -0.5117],\n",
+            "        [-1.4219, -0.5508,  1.3516, -2.7344]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,558 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-2.7031e+00, -3.6719e-01, -9.6484e-01, -1.0938e+00],\n",
+            "        [-3.7969e+00,  2.1606e-02, -1.9238e-01, -2.0156e+00],\n",
+            "        [-4.6875e+00, -6.4453e-01,  3.3875e-03, -7.1484e-01],\n",
+            "        ...,\n",
+            "        [ 3.4688e+00,  8.6719e-01,  4.1602e-01,  4.8047e-01],\n",
+            "        [-3.4062e+00, -1.5703e+00, -1.0645e-01,  5.7129e-02],\n",
+            "        [-3.1406e+00,  6.1719e-01, -2.3633e-01, -2.0938e+00]],\n",
+            "       requires_grad=True)\n",
+            "2025-01-31 15:48:39,630 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-4.7812,  1.3125, -2.5625,  0.5117],\n",
+            "        [-5.8125,  1.6953, -1.8203, -0.3789],\n",
+            "        [-6.6875,  1.0469, -1.6250,  0.8789],\n",
+            "        ...,\n",
+            "        [ 5.5000, -0.8320,  2.0312, -1.1016],\n",
+            "        [-5.4375,  0.1494, -1.7266,  1.6406],\n",
+            "        [-5.1875,  2.2656, -1.8594, -0.4629]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,697 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-3.2812,  2.9062, -4.0938,  2.1094],\n",
+            "        [-4.3125,  3.2812, -3.3750,  1.2422],\n",
+            "        [-5.1562,  2.6406, -3.1719,  2.4531],\n",
+            "        ...,\n",
+            "        [ 4.0000, -2.4531,  3.5625, -2.6875],\n",
+            "        [-3.9375,  1.7812, -3.2656,  3.2031],\n",
+            "        [-3.6875,  3.8438, -3.4062,  1.1719]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,764 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-2.0000,  4.2500, -5.3750,  3.4688],\n",
+            "        [-3.0156,  4.6250, -4.6875,  2.6406],\n",
+            "        [-3.8125,  4.0000, -4.5000,  3.7969],\n",
+            "        ...,\n",
+            "        [ 2.6875, -3.8438,  4.8750, -4.0312],\n",
+            "        [-2.6250,  3.1719, -4.5938,  4.5312],\n",
+            "        [-2.3906,  5.1875, -4.7188,  2.5625]], requires_grad=True)\n",
+            "2025-01-31 15:48:39,830 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
+            "tensor([[-0.8867,  5.4062, -6.4688,  4.6250],\n",
+            "        [-1.8750,  5.7812, -5.8125,  3.8438],\n",
+            "        [-2.6406,  5.1562, -5.6562,  4.9375],\n",
+            "        ...,\n",
+            "        [ 1.5469, -5.0312,  6.0000, -5.1875],\n",
+            "        [-1.4922,  4.3750, -5.7188,  5.6875],\n",
+            "        [-1.2500,  6.3438, -5.8125,  3.7656]], requires_grad=True)\n",
+            "2025-01-31 15:48:40,224 810d88f8-f1fc-43fd-b16d-77e11841eecd - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 67.1k/67.1k [00:00<00:00, 1.65MB/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) as batch:\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",
+        "        nnsight.log(lora.WA)\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2FtFv2Nykkw9"
+      },
+      "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": 14,
+      "metadata": {
+        "id": "K5_F4ktlkkw9"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:48:45,111 c9fcc827-5c78-4796-87cf-0f246c75a359 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:45,281 c9fcc827-5c78-4796-87cf-0f246c75a359 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:45,535 c9fcc827-5c78-4796-87cf-0f246c75a359 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:46,013 c9fcc827-5c78-4796-87cf-0f246c75a359 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 1.24k/1.24k [00:00<00:00, 4.30MB/s]\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "['<|begin_of_text|>Hello Paris']\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "2025-01-31 15:48:46,446 9f088820-f00d-492d-9861-c0f0a7c4d784 - RECEIVED: Your job has been received and is waiting approval.\n",
+            "2025-01-31 15:48:46,650 9f088820-f00d-492d-9861-c0f0a7c4d784 - APPROVED: Your job was approved and is waiting to be run.\n",
+            "2025-01-31 15:48:46,862 9f088820-f00d-492d-9861-c0f0a7c4d784 - RUNNING: Your job has started running.\n",
+            "2025-01-31 15:48:47,340 9f088820-f00d-492d-9861-c0f0a7c4d784 - COMPLETED: Your job has been completed.\n",
+            "Downloading result: 100%|██████████| 1.24k/1.24k [00:00<00:00, 4.61MB/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": "IZfiTNGgu-HI"
+      },
+      "source": [
+        "# Next Steps\n",
+        "Check out [nnsight.net/tutorials](https://nnsight.net/tutorials) for more walkthroughs implementating classic interpretability techniques using `nnsight`."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "P3zRm-7VRRov"
+      },
+      "source": [
+        "# Getting Involved!\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BnCc9xgjvxEP"
+      },
+      "source": [
+        "Note that both `nnsight` and `NDIF` are in active development, so changes may be made and errors may arise during use. If you’re interested in following updates to `nnsight`, contributing, giving feedback, or finding collaborators, please join the [NDIF discord](https://discord.gg/6uFJmCSwW7). We’d love to hear about your work using nnsight!\n",
+        "\n",
+        "You can also follow us on [LinkedIn](https://www.linkedin.com/company/national-deep-inference-fabric/), Bluesky: [@ndif-team.bsky.social](https://bsky.app/profile/ndif-team.bsky.social), and X: [@ndif_team](https://x.com/ndif_team).\n",
+        "\n",
+        "💟\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "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.10.13"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "0383a2798a0c4fcfbd1f068817770fd6": {
+          "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_cef2fe4eb51146afabe820140a3b0df3",
+            "max": 456318,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_52281c8594ad465397e6079385c558a1",
+            "value": 456318
+          }
+        },
+        "08fb0620499c44eb9d160a0f21263451": {
+          "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
+          }
+        },
+        "0c46bce8cc3e4393898b025ba6b4e543": {
+          "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": ""
+          }
+        },
+        "14008810744549208f147a132b419d20": {
+          "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_cb4be83440e44e1c9e664ac1874b830d",
+              "IPY_MODEL_f54c77317ca04861a958c8f982b25653",
+              "IPY_MODEL_db24a855fb5846e28c98dcfe9afff259"
+            ],
+            "layout": "IPY_MODEL_647c861fcaa44148a90da061b1d5c99c"
+          }
+        },
+        "184b5c75820b4c42ae861669ac1a47e5": {
+          "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": ""
+          }
+        },
+        "19c0e121e52040b18393aee887730644": {
+          "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
+          }
+        },
+        "26acae3d321743d2ad8efbf8727dbd05": {
+          "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_a4b5c225f5ac4719ade8e6aafcff78d2",
+            "max": 548105171,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_ee27d18a555c46e78c92b282bfa6dc9d",
+            "value": 548105171
+          }
+        },
+        "26dcd7b75feb41a6b86d4260ecd64d06": {
+          "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_39c9f1effe664e899125748e97fad198",
+              "IPY_MODEL_383b893224684aee916e455c0d3d8c06",
+              "IPY_MODEL_af24a12bd05b44148b5d6f91ac2445a0"
+            ],
+            "layout": "IPY_MODEL_d00750501f2548d0b1c69f4753aec877"
+          }
+        },
+        "2702f30b5d3b4b0b8184fec58863750b": {
+          "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": ""
+          }
+        },
+        "287a8fe3ddf54d36b7603171ec7a2688": {
+          "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": ""
+          }
+        },
+        "2a0f2484053f41298aceacea5c1cc195": {
+          "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_08fb0620499c44eb9d160a0f21263451",
+            "max": 124,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_287a8fe3ddf54d36b7603171ec7a2688",
+            "value": 124
+          }
+        },
+        "2fedb8b529c34049932f07b577b515d4": {
+          "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_bd1a11654c2d47a09610e4cc23b8bc32",
+              "IPY_MODEL_26acae3d321743d2ad8efbf8727dbd05",
+              "IPY_MODEL_f42e44122a2d496cb7b2ea85c2858a71"
+            ],
+            "layout": "IPY_MODEL_c1586aeef6814447ad46c1f8ddaae625"
+          }
+        },
+        "346ecc4734014cb4931bea53d5a32bf1": {
+          "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": ""
+          }
+        },
+        "383b893224684aee916e455c0d3d8c06": {
+          "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_68b205f60f9d4599a7d1f78583d0b93c",
+            "max": 665,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_69e8c149f676457d95029f21e9a2956a",
+            "value": 665
+          }
+        },
+        "385291e96814453b8f805cbe8a5e594d": {
+          "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_e5302eb0481c4e529fc68353eed4b79f",
+            "placeholder": "​",
+            "style": "IPY_MODEL_0c46bce8cc3e4393898b025ba6b4e543",
+            "value": " 456k/456k [00:00&lt;00:00, 3.39MB/s]"
+          }
+        },
+        "39c9f1effe664e899125748e97fad198": {
+          "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_98d28dddf75a45a2b1e05131b0b865a1",
+            "placeholder": "​",
+            "style": "IPY_MODEL_6e81d6f31f2446c999df5511e8afd8d0",
+            "value": "config.json: 100%"
+          }
+        },
+        "405308ec76074c7ea319e1c1358baf39": {
+          "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": ""
+          }
+        },
+        "41fc21ee95c14d29943a56f2b63f4989": {
+          "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_42fd0ea3be3c4b3f98ca7818efee3669",
+            "placeholder": "​",
+            "style": "IPY_MODEL_184b5c75820b4c42ae861669ac1a47e5",
+            "value": " 26.0/26.0 [00:00&lt;00:00, 2.05kB/s]"
+          }
+        },
+        "42fd0ea3be3c4b3f98ca7818efee3669": {
+          "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
+          }
+        },
+        "4b798a784ae748b1a0aef025ead976fc": {
+          "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": ""
+          }
+        },
+        "504038c4efd3439daa958333f108b5fb": {
+          "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": ""
+          }
+        },
+        "52281c8594ad465397e6079385c558a1": {
+          "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": ""
+          }
+        },
+        "5250cf68037a43e389493733c5b9d4ac": {
+          "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
+          }
+        },
+        "55c0bb35417e454d842e12bd2dfe62df": {
+          "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
+          }
+        },
+        "57aebe107749473390af2532806920ce": {
+          "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_55c0bb35417e454d842e12bd2dfe62df",
+            "placeholder": "​",
+            "style": "IPY_MODEL_d5821d69e1fa4673bb39ab70a7f60466",
+            "value": " 124/124 [00:00&lt;00:00, 8.42kB/s]"
+          }
+        },
+        "5d877979f2db4dbfa7581d2fe4381cf8": {
+          "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": ""
+          }
+        },
+        "5e091b15cc9748469676b165533f0c44": {
+          "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
+          }
+        },
+        "647c861fcaa44148a90da061b1d5c99c": {
+          "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
+          }
+        },
+        "64b713188f054535a3063c079bceead0": {
+          "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_7afff81060f0481aa95a56daa570e1b8",
+            "max": 1042301,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_a038f9037ed540f2ba43f3f85920b17e",
+            "value": 1042301
+          }
+        },
+        "68b205f60f9d4599a7d1f78583d0b93c": {
+          "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
+          }
+        },
+        "69e8c149f676457d95029f21e9a2956a": {
+          "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": ""
+          }
+        },
+        "6e81d6f31f2446c999df5511e8afd8d0": {
+          "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": ""
+          }
+        },
+        "7479281e090d4eeba4c21a9c15b02165": {
+          "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_f4b25b3aac414499b653404ea3c066da",
+              "IPY_MODEL_0383a2798a0c4fcfbd1f068817770fd6",
+              "IPY_MODEL_385291e96814453b8f805cbe8a5e594d"
+            ],
+            "layout": "IPY_MODEL_f7a41e267f79401681bc4693c52e55fa"
+          }
+        },
+        "76fe13ab697c46cc9fcbe46112fb45b1": {
+          "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": ""
+          }
+        },
+        "79a221c9228d459aa17402c9fef1ad1d": {
+          "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_9c13a3d768eb47c8b26a7c6bfa8ec957",
+              "IPY_MODEL_e97c052113f54dddbb619ed228f6b4ef",
+              "IPY_MODEL_41fc21ee95c14d29943a56f2b63f4989"
+            ],
+            "layout": "IPY_MODEL_5250cf68037a43e389493733c5b9d4ac"
+          }
+        },
+        "7afff81060f0481aa95a56daa570e1b8": {
+          "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
+          }
+        },
+        "7f4f63d94f8942fcb3e4b286d5848f0c": {
+          "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
+          }
+        },
+        "7fdb623e5b2d4eb1b6cee27da2314062": {
+          "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
+          }
+        },
+        "87a5c035bc6345aea6441b3db4b94321": {
+          "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
+          }
+        },
+        "95832c5076ab4c80810d3c443e061156": {
+          "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": ""
+          }
+        },
+        "98d28dddf75a45a2b1e05131b0b865a1": {
+          "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
+          }
+        },
+        "9c13a3d768eb47c8b26a7c6bfa8ec957": {
+          "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_e24033c5ba7b4e4fa4f4375d85b03bad",
+            "placeholder": "​",
+            "style": "IPY_MODEL_4b798a784ae748b1a0aef025ead976fc",
+            "value": "tokenizer_config.json: 100%"
+          }
+        },
+        "9cea52b7835b40628fc0a65199c5e82f": {
+          "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_ec4f48d5a9ef4afa8fa44078dfb236c6",
+            "placeholder": "​",
+            "style": "IPY_MODEL_405308ec76074c7ea319e1c1358baf39",
+            "value": " 1.04M/1.04M [00:00&lt;00:00, 7.50MB/s]"
+          }
+        },
+        "a038f9037ed540f2ba43f3f85920b17e": {
+          "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": ""
+          }
+        },
+        "a15895ad1cb541ca89fa9def63779104": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HBoxModel",
@@ -2865,14 +5324,14 @@
             "_view_name": "HBoxView",
             "box_style": "",
             "children": [
-              "IPY_MODEL_74358c341daa44faa41a92b0eb073615",
-              "IPY_MODEL_c40e17a994ed44d79ef54421f4772858",
-              "IPY_MODEL_431ffcb0dbb1419dbbb237a29842b905"
+              "IPY_MODEL_ee60494714c2485f802298dce918edfe",
+              "IPY_MODEL_2a0f2484053f41298aceacea5c1cc195",
+              "IPY_MODEL_57aebe107749473390af2532806920ce"
             ],
-            "layout": "IPY_MODEL_c28c185a508e459b95edb29afedd9c79"
+            "layout": "IPY_MODEL_ff44e896c55048c5a7d92a56aa9eedfc"
           }
         },
-        "0fddebfb42e847e08890d09f30161984": {
+        "a1c25910a31a4b5a9c48f1a7fd8047c5": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -2924,7 +5383,7 @@
             "width": null
           }
         },
-        "11e99924bdf84739b3cfc7d676b797e5": {
+        "a310703500c5481aa37b698c659e1ca6": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "DescriptionStyleModel",
@@ -2939,7 +5398,7 @@
             "description_width": ""
           }
         },
-        "18d4c2302f424bc18f9215f2d5b83393": {
+        "a3154c60744d4a1a87b1c86b89a142b9": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -2991,43 +5450,7 @@
             "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": {
+        "a4b5c225f5ac4719ade8e6aafcff78d2": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3079,7 +5502,7 @@
             "width": null
           }
         },
-        "431ffcb0dbb1419dbbb237a29842b905": {
+        "af24a12bd05b44148b5d6f91ac2445a0": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3094,13 +5517,13 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_562a287082eb4882af864062ae13c91b",
+            "layout": "IPY_MODEL_ba90058201ee48dcae9411ded457e85f",
             "placeholder": "​",
-            "style": "IPY_MODEL_11e99924bdf84739b3cfc7d676b797e5",
-            "value": " 665/665 [00:00&lt;00:00, 12.8kB/s]"
+            "style": "IPY_MODEL_346ecc4734014cb4931bea53d5a32bf1",
+            "value": " 665/665 [00:00&lt;00:00, 36.5kB/s]"
           }
         },
-        "562a287082eb4882af864062ae13c91b": {
+        "ba90058201ee48dcae9411ded457e85f": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3152,76 +5575,7 @@
             "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": {
+        "bd1a11654c2d47a09610e4cc23b8bc32": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3236,13 +5590,65 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_7a0f3d2deeb14bbba7f0df208133a7e2",
+            "layout": "IPY_MODEL_5e091b15cc9748469676b165533f0c44",
             "placeholder": "​",
-            "style": "IPY_MODEL_94cdcf1fb0914dfe99c3c615e68b7ec8",
-            "value": "config.json: 100%"
+            "style": "IPY_MODEL_f143977aeaf24798b7744681069762d4",
+            "value": "model.safetensors: 100%"
+          }
+        },
+        "c1586aeef6814447ad46c1f8ddaae625": {
+          "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
           }
         },
-        "7a0f3d2deeb14bbba7f0df208133a7e2": {
+        "caff23247fb9432db8d5860a8132e1b3": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3294,7 +5700,7 @@
             "width": null
           }
         },
-        "7e6b41ec535648b190bfc655d9e430dd": {
+        "cb4be83440e44e1c9e664ac1874b830d": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3309,13 +5715,13 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_beb4231b49024d8a9bb7ead19647d99d",
+            "layout": "IPY_MODEL_e4e0b1de76ff40a88b71e0b5a9f7bc39",
             "placeholder": "​",
-            "style": "IPY_MODEL_26d09aa072d64af0bf4df9a89309093a",
-            "value": " 1.36M/1.36M [00:00&lt;00:00, 4.16MB/s]"
+            "style": "IPY_MODEL_504038c4efd3439daa958333f108b5fb",
+            "value": "tokenizer.json: 100%"
           }
         },
-        "8776d04bfd3748b69e10708466e9c132": {
+        "cef2fe4eb51146afabe820140a3b0df3": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3367,7 +5773,7 @@
             "width": null
           }
         },
-        "8de63414803f41c29e75242756cfd67f": {
+        "d00750501f2548d0b1c69f4753aec877": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3419,31 +5825,7 @@
             "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": {
+        "d5821d69e1fa4673bb39ab70a7f60466": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "DescriptionStyleModel",
@@ -3458,28 +5840,29 @@
             "description_width": ""
           }
         },
-        "95d303d83d3348a0adfdf6b5cb94c41c": {
+        "d8feaceb6d854cb2b2ff19a0309d9a7b": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
+          "model_name": "HBoxModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
+            "_model_name": "HBoxModel",
             "_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%"
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_fd84ffef58604ce3a16ef8f9868072f3",
+              "IPY_MODEL_64b713188f054535a3063c079bceead0",
+              "IPY_MODEL_9cea52b7835b40628fc0a65199c5e82f"
+            ],
+            "layout": "IPY_MODEL_87a5c035bc6345aea6441b3db4b94321"
           }
         },
-        "96a3b8c765374a5e98eaf713dd5e85d8": {
+        "db24a855fb5846e28c98dcfe9afff259": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -3494,65 +5877,28 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_8776d04bfd3748b69e10708466e9c132",
+            "layout": "IPY_MODEL_a1c25910a31a4b5a9c48f1a7fd8047c5",
             "placeholder": "​",
-            "style": "IPY_MODEL_56490c20aef04fe4a63ba7c990e71e77",
-            "value": " 1.04M/1.04M [00:00&lt;00:00, 3.16MB/s]"
+            "style": "IPY_MODEL_5d877979f2db4dbfa7581d2fe4381cf8",
+            "value": " 1.36M/1.36M [00:00&lt;00:00, 6.83MB/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
+        "de023d4526524f48a34883a287aa8a4f": {
+          "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": ""
           }
         },
-        "b49ce0d00e5b4df097b3a9fc4c0bd8e7": {
+        "ded8c188917b424dbf264ce494fe0532": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3604,60 +5950,7 @@
             "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": {
+        "e24033c5ba7b4e4fa4f4375d85b03bad": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3709,7 +6002,7 @@
             "width": null
           }
         },
-        "c28c185a508e459b95edb29afedd9c79": {
+        "e4e0b1de76ff40a88b71e0b5a9f7bc39": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3761,46 +6054,7 @@
             "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": {
+        "e5302eb0481c4e529fc68353eed4b79f": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3852,29 +6106,31 @@
             "width": null
           }
         },
-        "cef2516aabc24f6d852079984afa8090": {
+        "e97c052113f54dddbb619ed228f6b4ef": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "HBoxModel",
+          "model_name": "FloatProgressModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "HBoxModel",
+            "_model_name": "FloatProgressModel",
             "_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"
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_19c0e121e52040b18393aee887730644",
+            "max": 26,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_2702f30b5d3b4b0b8184fec58863750b",
+            "value": 26
           }
         },
-        "d0198fa98e5f412380f2ba3e9448aef3": {
+        "ec4f48d5a9ef4afa8fa44078dfb236c6": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -3926,7 +6182,7 @@
             "width": null
           }
         },
-        "d07a0a49a2924092affc0e07bf0b5ef0": {
+        "ee27d18a555c46e78c92b282bfa6dc9d": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "ProgressStyleModel",
@@ -3942,55 +6198,79 @@
             "description_width": ""
           }
         },
-        "d839da1c31a34659988e3420ba520725": {
+        "ee60494714c2485f802298dce918edfe": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "FloatProgressModel",
+          "model_name": "HTMLModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "FloatProgressModel",
+            "_model_name": "HTMLModel",
             "_view_count": null,
             "_view_module": "@jupyter-widgets/controls",
             "_view_module_version": "1.5.0",
-            "_view_name": "ProgressView",
-            "bar_style": "success",
+            "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_3731a5747724403e86ef1abaf4fb389b",
-            "max": 1042301,
-            "min": 0,
-            "orientation": "horizontal",
-            "style": "IPY_MODEL_688bf676958a4050b73cc9597f5110de",
-            "value": 1042301
+            "layout": "IPY_MODEL_7fdb623e5b2d4eb1b6cee27da2314062",
+            "placeholder": "​",
+            "style": "IPY_MODEL_95832c5076ab4c80810d3c443e061156",
+            "value": "generation_config.json: 100%"
           }
         },
-        "e4568f72971749ac92cc5577cb9c4f9c": {
+        "f143977aeaf24798b7744681069762d4": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "FloatProgressModel",
+          "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": ""
+          }
+        },
+        "f42e44122a2d496cb7b2ea85c2858a71": {
+          "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": "FloatProgressModel",
+            "_model_name": "HTMLModel",
             "_view_count": null,
             "_view_module": "@jupyter-widgets/controls",
             "_view_module_version": "1.5.0",
-            "_view_name": "ProgressView",
-            "bar_style": "success",
+            "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_ceebda1d07b545f5b605551c08e73be8",
-            "max": 456318,
-            "min": 0,
-            "orientation": "horizontal",
-            "style": "IPY_MODEL_6aa96b80e3b746b48f5c13bcafe1134e",
-            "value": 456318
+            "layout": "IPY_MODEL_a3154c60744d4a1a87b1c86b89a142b9",
+            "placeholder": "​",
+            "style": "IPY_MODEL_de023d4526524f48a34883a287aa8a4f",
+            "value": " 548M/548M [00:02&lt;00:00, 253MB/s]"
+          }
+        },
+        "f47609606aa14f6fa2407aa26edae8e3": {
+          "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": ""
           }
         },
-        "e51620404d244a4e806eeaf3b96a474f": {
+        "f4b25b3aac414499b653404ea3c066da": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
           "model_name": "HTMLModel",
@@ -4005,34 +6285,37 @@
             "_view_name": "HTMLView",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_f28f557210b7402b8b2d0fd83f038dc1",
+            "layout": "IPY_MODEL_caff23247fb9432db8d5860a8132e1b3",
             "placeholder": "​",
-            "style": "IPY_MODEL_bc8190c700e949ebb6f874c8f8055161",
-            "value": "merges.txt: 100%"
+            "style": "IPY_MODEL_f47609606aa14f6fa2407aa26edae8e3",
+            "value": "merges.txt: 100%"
           }
         },
-        "e7a73cfc3fa04425afc08368532ef0be": {
+        "f54c77317ca04861a958c8f982b25653": {
           "model_module": "@jupyter-widgets/controls",
           "model_module_version": "1.5.0",
-          "model_name": "HTMLModel",
+          "model_name": "FloatProgressModel",
           "state": {
             "_dom_classes": [],
             "_model_module": "@jupyter-widgets/controls",
             "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
+            "_model_name": "FloatProgressModel",
             "_view_count": null,
             "_view_module": "@jupyter-widgets/controls",
             "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
             "description": "",
             "description_tooltip": null,
-            "layout": "IPY_MODEL_b2748df9aff540df954f002310606a56",
-            "placeholder": "​",
-            "style": "IPY_MODEL_f4caf2bc58704464b6ac7bf6da31b71d",
-            "value": "vocab.json: 100%"
+            "layout": "IPY_MODEL_ded8c188917b424dbf264ce494fe0532",
+            "max": 1355256,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_76fe13ab697c46cc9fcbe46112fb45b1",
+            "value": 1355256
           }
         },
-        "eb0c489b913744a28304e5334bb66532": {
+        "f7a41e267f79401681bc4693c52e55fa": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -4084,59 +6367,28 @@
             "width": null
           }
         },
-        "f28f557210b7402b8b2d0fd83f038dc1": {
-          "model_module": "@jupyter-widgets/base",
-          "model_module_version": "1.2.0",
-          "model_name": "LayoutModel",
+        "fd84ffef58604ce3a16ef8f9868072f3": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_module_version": "1.5.0",
+          "model_name": "HTMLModel",
           "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
             "_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
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_7f4f63d94f8942fcb3e4b286d5848f0c",
+            "placeholder": "​",
+            "style": "IPY_MODEL_a310703500c5481aa37b698c659e1ca6",
+            "value": "vocab.json: 100%"
           }
         },
-        "f3e24aa6619e4185a3e2f53f2e473e45": {
+        "ff44e896c55048c5a7d92a56aa9eedfc": {
           "model_module": "@jupyter-widgets/base",
           "model_module_version": "1.2.0",
           "model_name": "LayoutModel",
@@ -4187,36 +6439,6 @@
             "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": ""
-          }
         }
       }
     }

From e3da68680c68beb557345466bfe0e92bc424053b Mon Sep 17 00:00:00 2001
From: ebortz <ebortz@bu.edu>
Date: Fri, 31 Jan 2025 17:01:19 -0500
Subject: [PATCH 7/8] Llama 8b remote

---
 NNsight_Walkthrough.ipynb                         | 2 +-
 docs/source/notebooks/tutorials/walkthrough.ipynb | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/NNsight_Walkthrough.ipynb b/NNsight_Walkthrough.ipynb
index 60db8609..b3803a52 100644
--- a/NNsight_Walkthrough.ipynb
+++ b/NNsight_Walkthrough.ipynb
@@ -2720,7 +2720,7 @@
       "source": [
         "import os\n",
         "\n",
-        "# Llama 3.1 70b is a gated model, so you need to apply for access on HuggingFace and include your token.\n",
+        "# Llama 3.1 8b is a gated model, so you need to apply for access on HuggingFace and include your token.\n",
         "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\""
       ]
     },
diff --git a/docs/source/notebooks/tutorials/walkthrough.ipynb b/docs/source/notebooks/tutorials/walkthrough.ipynb
index e343940f..b596fddb 100644
--- a/docs/source/notebooks/tutorials/walkthrough.ipynb
+++ b/docs/source/notebooks/tutorials/walkthrough.ipynb
@@ -2710,7 +2710,7 @@
       "source": [
         "import os\n",
         "\n",
-        "# Llama 3.1 70b is a gated model, so you need to apply for access on HuggingFace and include your token.\n",
+        "# Llama 3.1 8b is a gated model, so you need to apply for access on HuggingFace and include your token.\n",
         "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\""
       ]
     },

From 60af87009641d84c3d1a2eb26d75e1318bb4759d Mon Sep 17 00:00:00 2001
From: ebortz <ebortz@bu.edu>
Date: Fri, 31 Jan 2025 17:17:08 -0500
Subject: [PATCH 8/8] walkthrough title edits

---
 .../notebooks/tutorials/walkthrough.ipynb     | 60 +++++++++----------
 1 file changed, 28 insertions(+), 32 deletions(-)

diff --git a/docs/source/notebooks/tutorials/walkthrough.ipynb b/docs/source/notebooks/tutorials/walkthrough.ipynb
index b596fddb..51c7237a 100644
--- a/docs/source/notebooks/tutorials/walkthrough.ipynb
+++ b/docs/source/notebooks/tutorials/walkthrough.ipynb
@@ -15,9 +15,7 @@
         "id": "-Scfr942GwO4"
       },
       "source": [
-        "<img src=\"https://nnsight.net/_static/images/nnsight_logo.svg\" alt=\"drawing\" width=\"200\"/>\n",
-        "\n",
-        "# **NNsight**\n",
+        "# Walkthrough\n",
         "\n",
         "## The API for a transparent science on black-box AI\n",
         "\n",
@@ -43,14 +41,14 @@
         "id": "N1OemD2VGyZx"
       },
       "source": [
-        "# 1 First, let's start small"
+        "## 1 First, let's start small"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {},
       "source": [
-        "## Setup\n",
+        "### Setup\n",
         "\n",
         "Install nnsight:\n",
         "```\n",
@@ -64,7 +62,7 @@
         "id": "WyLMmhrAKTNM"
       },
       "source": [
-        "## Tracing Context\n",
+        "### 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"
@@ -249,7 +247,7 @@
         "id": "xMfBpYzDPMoB"
       },
       "source": [
-        "## Getting\n"
+        "### Getting\n"
       ]
     },
     {
@@ -562,7 +560,7 @@
         "id": "p7Xo_PHyPr4p"
       },
       "source": [
-        "## Functions, Methods, and Operations\n"
+        "### Functions, Methods, and Operations\n"
       ]
     },
     {
@@ -674,7 +672,7 @@
         "id": "kLmzsHaMkkw2"
       },
       "source": [
-        "## Custom Functions"
+        "### Custom Functions"
       ]
     },
     {
@@ -747,7 +745,7 @@
         "id": "m_5qH5gHPOT_"
       },
       "source": [
-        "## Setting\n"
+        "### Setting\n"
       ]
     },
     {
@@ -1145,7 +1143,7 @@
         "id": "6R3wj1BKkkw5"
       },
       "source": [
-        "## Scan and Validate\n",
+        "### Scan and Validate\n",
         "Error codes are helpful, but sometimes you may want to quickly troubleshoot your code without actually running it. \n",
         "\n",
         "Enter \"Scanning\" and \"Validating\"! We can enable this features by setting the `scan=True` and `validate=True` flag in the `trace` method."
@@ -1280,7 +1278,7 @@
         "id": "s016CelFP8sx"
       },
       "source": [
-        "## Gradients\n",
+        "### 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",
@@ -1386,7 +1384,7 @@
         "id": "SJ6dAgzQkkw6"
       },
       "source": [
-        "## Early Stopping"
+        "### Early Stopping"
       ]
     },
     {
@@ -1488,7 +1486,7 @@
         "id": "i1iJOI29kkw6"
       },
       "source": [
-        "## Conditional Interventions"
+        "### Conditional Interventions"
       ]
     },
     {
@@ -1664,7 +1662,7 @@
         "id": "ONzsQcHtxKkf"
       },
       "source": [
-        "## Iterative Interventions"
+        "### Iterative Interventions"
       ]
     },
     {
@@ -1765,7 +1763,7 @@
         "id": "gvQ1nZgYQDG3"
       },
       "source": [
-        "# 2️ Bigger"
+        "## 2️ Bigger"
       ]
     },
     {
@@ -1791,7 +1789,7 @@
         "id": "5TJDblHiQpp1"
       },
       "source": [
-        "## LanguageModel\n"
+        "### LanguageModel\n"
       ]
     },
     {
@@ -2046,9 +2044,7 @@
         "id": "vLjHnRyRPXjp"
       },
       "source": [
-        "<a name=\"batching-id\"></a>\n",
-        "\n",
-        "## Batching\n"
+        "### Batching \n"
       ]
     },
     {
@@ -2227,7 +2223,7 @@
         "id": "_bXvUGGqbwv1"
       },
       "source": [
-        "## Multiple Token Generation"
+        "### Multiple Token Generation"
       ]
     },
     {
@@ -2236,7 +2232,7 @@
         "id": "vvWA-CWqQtah"
       },
       "source": [
-        "### .next()\n",
+        "#### .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",
@@ -2305,7 +2301,7 @@
         "id": "oHy6FqyNRBzm"
       },
       "source": [
-        "### New! using .all()\n",
+        "#### using .all()\n",
         "\n",
         "With `nnsight 0.4` you can now use `.all()` to recursively apply interventions to a model. Calling `.all()` on a module within a model will recursively apply its `.input` and `.output` across all iterations. Previously, we'd need to loop across each new generated token, saving the intervention for every generated token and calling `.next()` to move forward."
       ]
@@ -2484,7 +2480,7 @@
         "id": "ty-luCwCkkw7"
       },
       "source": [
-        "## Model Editing"
+        "### Model Editing"
       ]
     },
     {
@@ -2631,7 +2627,7 @@
         "id": "rpkX-LwBQZHo"
       },
       "source": [
-        "# 3 I thought you said huge models?"
+        "## 3 I thought you said huge models?"
       ]
     },
     {
@@ -2658,7 +2654,7 @@
         "id": "65Ks1LUvQaER"
       },
       "source": [
-        "## Remote execution\n"
+        "### Remote execution\n"
       ]
     },
     {
@@ -2803,7 +2799,7 @@
         "id": "lfRDZox5kkw8"
       },
       "source": [
-        "## Sessions"
+        "### Sessions"
       ]
     },
     {
@@ -2942,7 +2938,7 @@
         "id": "H6hxSZ8zTFvT"
       },
       "source": [
-        "## Streaming"
+        "### Streaming"
       ]
     },
     {
@@ -3403,7 +3399,7 @@
         "id": "DUCWX1BDkkw9"
       },
       "source": [
-        "## Looping across sessions"
+        "### Looping across sessions"
       ]
     },
     {
@@ -3623,7 +3619,7 @@
         "id": "1OcolbxAkkw9"
       },
       "source": [
-        "## Training a LoRA"
+        "### Training a LoRA"
       ]
     },
     {
@@ -3994,7 +3990,7 @@
         "id": "IZfiTNGgu-HI"
       },
       "source": [
-        "# Next Steps\n",
+        "## Next Steps\n",
         "Check out [nnsight.net/tutorials](https://nnsight.net/tutorials) for more walkthroughs implementating classic interpretability techniques using `nnsight`."
       ]
     },
@@ -4004,7 +4000,7 @@
         "id": "P3zRm-7VRRov"
       },
       "source": [
-        "# Getting Involved!\n"
+        "## Getting Involved!\n"
       ]
     },
     {