From f143330fc0bff8c00e8d466aa908015c44253731 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 2 Jul 2026 02:35:09 +0000 Subject: [PATCH 1/2] Mask non-deterministic chart data in visual (Argos) tests Visual regression snapshots run against real NDIF, whose float outputs drift run-to-run and cause pixel noise in heatmap cell colors and plotted line curves. Add a NEXT_PUBLIC_VISUAL_TEST flag that, when set, neutralizes just the noisy values before they reach the nnsightful widgets: - Logit-lens heatmap: blank every cell except the final prediction (last token of the last layer) and flatten tracked/entropy so the remaining cell color and layer skyline are stable. - Activation-patching line plot: flatten the plotted series so the drawn lines are constant; axes, mode bar and token selector still render. The real NDIF request/response path is unchanged; only rendering is masked. The flag is enabled in the e2e workflow before the build. Co-authored-by: Jonathan Bell --- .github/workflows/e2e.yml | 1 + .../components/ActivationPatchingDisplay.tsx | 13 ++- .../[chartId]/components/Lens2Display.tsx | 13 ++- workbench/_web/src/lib/visualTest.ts | 82 +++++++++++++++++++ 4 files changed, 105 insertions(+), 4 deletions(-) create mode 100644 workbench/_web/src/lib/visualTest.ts diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index dd2ad0d5..e3fe8409 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -91,6 +91,7 @@ jobs: "NEXT_PUBLIC_BASE_URL=http://localhost:3000" \ "NEXT_PUBLIC_BACKEND_URL=http://localhost:8000" \ "HF_TOKEN=${HF_TOKEN}" \ + "NEXT_PUBLIC_VISUAL_TEST=true" \ > .env cp .env workbench/_web/.env # Backend reads its own .env; mirror the secrets so CONFIG=e2e diff --git a/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx b/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx index fc200a10..b89e70c5 100644 --- a/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx +++ b/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx @@ -1,6 +1,6 @@ "use client"; -import { useState, useCallback, useRef, useEffect } from "react"; +import { useState, useCallback, useRef, useEffect, useMemo } from "react"; import { useParams } from "next/navigation"; import { useQuery, useIsMutating } from "@tanstack/react-query"; import { getChartById, getConfigForChart } from "@/lib/queries/chartQueries"; @@ -16,6 +16,7 @@ import { useUpdateChartConfig } from "@/lib/api/configApi"; import { NotebookExporter } from "@/components/NotebookExporter"; import { ChartModelPill } from "@/components/charts/ChartModelPill"; import { chartModelFromConfig, isChartStale } from "@/lib/configModelDiff"; +import { isVisualTestMode, maskActivationPatchingDataForVisualTest } from "@/lib/visualTest"; import { useModelsQuery } from "@/lib/api/modelsApi"; import { useWorkspace } from "@/stores/useWorkspace"; @@ -79,6 +80,14 @@ export function ActivationPatchingDisplay() { const hasData = patchingChart?.data && "lines" in patchingChart.data && patchingChart.data.lines.length > 0; + // In visual-test mode, flatten the noisy plotted series so Argos snapshots + // are deterministic against real NDIF (chart chrome stays, lines ignored). + const widgetData = useMemo(() => { + const raw = patchingChart?.data; + if (!raw || !hasData) return raw; + return isVisualTestMode() ? maskActivationPatchingDataForVisualTest(raw) : raw; + }, [patchingChart?.data, hasData]); + // Get the chart's saved name (treat "Untitled Chart" default as empty) const rawChartName = patchingChart?.name || ""; const chartName = rawChartName === "Untitled Chart" ? "" : rawChartName; @@ -280,7 +289,7 @@ export function ActivationPatchingDisplay() { {/* Chart area */}
{ + const raw = lens2Chart?.data as LogitLensData | undefined; + if (!raw || !hasData) return raw; + return isVisualTestMode() ? maskLogitLensDataForVisualTest(raw) : raw; + }, [lens2Chart?.data, hasData]); + // ── Persist heatmap UI state (pins, selection, layer window, appearance) // into the chart config, mirroring ActivationPatchingDisplay. Debounced // because the widget emits on every interaction; restored on mount via @@ -218,7 +227,7 @@ export function Lens2Display() { {stale && chartModel && }
+ layerRow.map((cell, pos) => + layerIdx === lastLayer && pos === lastPos ? cell : [""], + ), + ); + + // Flatten every tracked trajectory to a constant. The final cell reads its + // probability from here, so this pins its color; blanked cells reference + // tokens absent from `tracked` and render at probability 0. + const tracked = data.tracked.map((posMap) => { + const flat: Record = {}; + for (const [token, trajectory] of Object.entries(posMap)) { + flat[token] = trajectory.map(() => 1); + } + return flat; + }); + + const entropy = data.entropy?.map((layerRow) => layerRow.map(() => 0)); + + return { ...data, topk, tracked, ...(entropy ? { entropy } : {}) }; +} + +/** + * Flatten the activation-patching series so the plotted lines are constant and + * therefore deterministic. Layer count, token labels and chart chrome are left + * intact — only the curve values are ignored. + */ +export function maskActivationPatchingDataForVisualTest( + data: ActivationPatchingData, +): ActivationPatchingData { + const flatten = (grid: number[][]) => grid.map((row) => row.map(() => 0)); + return { + ...data, + lines: flatten(data.lines), + ranks: flatten(data.ranks), + prob_diffs: flatten(data.prob_diffs), + }; +} From 620574732e5de6c18039b76e92fc73e0b0ff3db6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 3 Jul 2026 00:14:34 +0000 Subject: [PATCH 2/2] Apply prettier formatting to visualTest.ts Co-authored-by: Jonathan Bell --- workbench/_web/src/lib/visualTest.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/workbench/_web/src/lib/visualTest.ts b/workbench/_web/src/lib/visualTest.ts index 29d971c7..6619a410 100644 --- a/workbench/_web/src/lib/visualTest.ts +++ b/workbench/_web/src/lib/visualTest.ts @@ -43,9 +43,7 @@ export function maskLogitLensDataForVisualTest(data: LogitLensData): LogitLensDa // topk is indexed [layer][position]; keep only the final-prediction cell. const topk = data.topk.map((layerRow, layerIdx) => - layerRow.map((cell, pos) => - layerIdx === lastLayer && pos === lastPos ? cell : [""], - ), + layerRow.map((cell, pos) => (layerIdx === lastLayer && pos === lastPos ? cell : [""])), ); // Flatten every tracked trajectory to a constant. The final cell reads its