From d68aa73fcfaf931a2c52dd892b21ffbc503d4b68 Mon Sep 17 00:00:00 2001
From: jlarson4 <jonahalarson@comcast.net>
Date: Wed, 21 Oct 2026 10:10:03 -0500
Subject: [PATCH] Fixing the flaky backwards lens test

---
 demos/Backward_Lens_Demo.ipynb | 33 +++++++++++++++++----------------
 1 file changed, 17 insertions(+), 16 deletions(-)

diff --git a/demos/Backward_Lens_Demo.ipynb b/demos/Backward_Lens_Demo.ipynb
index 070e7f984..94942c9a6 100644
--- a/demos/Backward_Lens_Demo.ipynb
+++ b/demos/Backward_Lens_Demo.ipynb
@@ -843,49 +843,49 @@
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>0</td>\n",
-       "      <td>6.557e-01</td>\n",
+       "      <td>6.56e-01</td>\n",
        "      <td>adr, adian, umn, itia, igmatic</td>\n",
        "      <td>adr, adian, umn, itia, igmatic</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>2</td>\n",
-       "      <td>1.442e+00</td>\n",
+       "      <td>1.44e+00</td>\n",
        "      <td>EStreamFrame, ORPG,  Dise,  misconception, senal</td>\n",
        "      <td>EStreamFrame, ORPG,  Dise,  misconception, senal</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>4</td>\n",
-       "      <td>1.329e+00</td>\n",
+       "      <td>1.33e+00</td>\n",
        "      <td>Diseases,  Firstly,  Feet, heses,  Differences</td>\n",
        "      <td>Diseases,  Firstly,  Feet, heses,  Differences</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
        "      <td>6</td>\n",
-       "      <td>1.042e+00</td>\n",
+       "      <td>1.04e+00</td>\n",
        "      <td>Tycoon, 76561,  Osiris,  Modes,  Memories</td>\n",
        "      <td>Tycoon, 76561,  Osiris,  Modes,  Memories</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>4</th>\n",
        "      <td>8</td>\n",
-       "      <td>4.872e-01</td>\n",
+       "      <td>4.87e-01</td>\n",
        "      <td>Paris, Paris,  Cologne,  Copenhagen,  Rome</td>\n",
        "      <td>Paris, Paris,  Cologne,  Copenhagen,  Rome</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>5</th>\n",
        "      <td>10</td>\n",
-       "      <td>2.749e-01</td>\n",
+       "      <td>2.75e-01</td>\n",
        "      <td>Paris, Paris,  Copenhagen,  Cologne,  Kyoto</td>\n",
        "      <td>Paris, Paris,  Copenhagen,  Cologne,  Kyoto</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>6</th>\n",
        "      <td>11</td>\n",
-       "      <td>2.609e-01</td>\n",
+       "      <td>2.61e-01</td>\n",
        "      <td>Paris, Paris,  Copenhagen,  Marse,  France</td>\n",
        "      <td>Paris, Paris,  Copenhagen,  Marse,  France</td>\n",
        "    </tr>\n",
@@ -894,14 +894,14 @@
        "</div>"
       ],
       "text/plain": [
-       "   layer   VJP norm                             raw FF2 bottom tokens  \\\n",
-       "0      0  6.557e-01                    adr, adian, umn, itia, igmatic   \n",
-       "1      2  1.442e+00  EStreamFrame, ORPG,  Dise,  misconception, senal   \n",
-       "2      4  1.329e+00    Diseases,  Firstly,  Feet, heses,  Differences   \n",
-       "3      6  1.042e+00         Tycoon, 76561,  Osiris,  Modes,  Memories   \n",
-       "4      8  4.872e-01        Paris, Paris,  Cologne,  Copenhagen,  Rome   \n",
-       "5     10  2.749e-01       Paris, Paris,  Copenhagen,  Cologne,  Kyoto   \n",
-       "6     11  2.609e-01        Paris, Paris,  Copenhagen,  Marse,  France   \n",
+       "   layer  VJP norm                             raw FF2 bottom tokens  \\\n",
+       "0      0  6.56e-01                    adr, adian, umn, itia, igmatic   \n",
+       "1      2  1.44e+00  EStreamFrame, ORPG,  Dise,  misconception, senal   \n",
+       "2      4  1.33e+00    Diseases,  Firstly,  Feet, heses,  Differences   \n",
+       "3      6  1.04e+00         Tycoon, 76561,  Osiris,  Modes,  Memories   \n",
+       "4      8  4.87e-01        Paris, Paris,  Cologne,  Copenhagen,  Rome   \n",
+       "5     10  2.75e-01       Paris, Paris,  Copenhagen,  Cologne,  Kyoto   \n",
+       "6     11  2.61e-01        Paris, Paris,  Copenhagen,  Marse,  France   \n",
        "\n",
        "                       normalized FF2 bottom tokens  \n",
        "0                    adr, adian, umn, itia, igmatic  \n",
@@ -932,10 +932,11 @@
     "        row = (layer, position, token, joined(ff1_top[position]), joined(ff2_bottom[position]))\n",
     "        layer_position_rows.append(row)\n",
     "    normalized = vocab_logits(model, delta[-1], normalized=True)\n",
+    "    # fp32 backends disagree in the fifth significant digit, so the norm shows three.\n",
     "    norm_comparison_rows.append(\n",
     "        (\n",
     "            layer,\n",
-    "            f\"{float(delta[-1].norm()):.3e}\",\n",
+    "            f\"{float(delta[-1].norm()):.2e}\",\n",
     "            joined(ff2[-1].topk(5, largest=False).indices),\n",
     "            joined(normalized.topk(5, largest=False).indices),\n",
     "        )\n",