From c87df0a74c3874062d4bd97906027f5818083264 Mon Sep 17 00:00:00 2001
From: ojh31 <oskar.hollinsworth@gmail.com>
Date: Sun, 17 Sep 2023 11:45:47 +0100
Subject: [PATCH 1/9] Added default dimension args to text_neuron_activations

---
 python/circuitsvis/activations.py             |  4 ++++
 .../src/activations/TextNeuronActivations.tsx | 19 ++++++++++++++++---
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/python/circuitsvis/activations.py b/python/circuitsvis/activations.py
index 1f634c21..98a5f904 100644
--- a/python/circuitsvis/activations.py
+++ b/python/circuitsvis/activations.py
@@ -13,6 +13,8 @@ def text_neuron_activations(
     second_dimension_name: Optional[str] = "Neuron",
     first_dimension_labels: Optional[List[str]] = None,
     second_dimension_labels: Optional[List[str]] = None,
+    first_dimension_default: Optional[int] = 0,
+    second_dimension_default: Optional[int] = 0,
 ) -> RenderedHTML:
     """Show activations (colored by intensity) for each token in a text or set
     of texts.
@@ -54,4 +56,6 @@ def text_neuron_activations(
         secondDimensionName=second_dimension_name,
         firstDimensionLabels=first_dimension_labels,
         secondDimensionLabels=second_dimension_labels,
+        firstDimensionDefault=first_dimension_default,
+        secondDimensionDefault=second_dimension_default,
     )
diff --git a/react/src/activations/TextNeuronActivations.tsx b/react/src/activations/TextNeuronActivations.tsx
index 7540b0f0..7d8562eb 100644
--- a/react/src/activations/TextNeuronActivations.tsx
+++ b/react/src/activations/TextNeuronActivations.tsx
@@ -36,7 +36,9 @@ export function TextNeuronActivations({
   firstDimensionName = "Layer",
   secondDimensionName = "Neuron",
   firstDimensionLabels,
-  secondDimensionLabels
+  secondDimensionLabels,
+  firstDimensionDefault = 0,
+  secondDimensionDefault = 0
 }: TextNeuronActivationsProps) {
   // If there is only one sample (i.e. if tokens is an array of strings), cast tokens and activations to an array with
   // a single element
@@ -68,8 +70,8 @@ export function TextNeuronActivations({
   const [sampleNumbers, setSampleNumbers] = useState<number[]>([
     ...Array(samplesPerPage).keys()
   ]);
-  const [layerNumber, setLayerNumber] = useState<number>(0);
-  const [neuronNumber, setNeuronNumber] = useState<number>(0);
+  const [layerNumber, setLayerNumber] = useState<number>(firstDimensionDefault);
+  const [neuronNumber, setNeuronNumber] = useState<number>(secondDimensionDefault);
 
   useEffect(() => {
     // When the user changes the samplesPerPage, update the sampleNumbers
@@ -218,4 +220,15 @@ export interface TextNeuronActivationsProps {
    * Labels for the second dimension
    */
   secondDimensionLabels?: string[];
+
+  /**
+   * Default index for the first dimension
+   */
+  firstDimensionDefault?: number;
+
+  /**
+   * Default index for the second dimension
+   */
+  secondDimensionDefault?: number;
+
 }

From a78056c106b046321021edfad8c765f59ad64e48 Mon Sep 17 00:00:00 2001
From: ojh31 <oskar.hollinsworth@gmail.com>
Date: Tue, 19 Sep 2023 17:13:53 +0100
Subject: [PATCH 2/9] Attempted to add show_selectors arg

---
 python/circuitsvis/activations.py             |   2 +
 .../src/activations/TextNeuronActivations.tsx | 109 ++++++++----------
 2 files changed, 52 insertions(+), 59 deletions(-)

diff --git a/python/circuitsvis/activations.py b/python/circuitsvis/activations.py
index 98a5f904..e6547c02 100644
--- a/python/circuitsvis/activations.py
+++ b/python/circuitsvis/activations.py
@@ -15,6 +15,7 @@ def text_neuron_activations(
     second_dimension_labels: Optional[List[str]] = None,
     first_dimension_default: Optional[int] = 0,
     second_dimension_default: Optional[int] = 0,
+    show_selectors: bool = True,
 ) -> RenderedHTML:
     """Show activations (colored by intensity) for each token in a text or set
     of texts.
@@ -58,4 +59,5 @@ def text_neuron_activations(
         secondDimensionLabels=second_dimension_labels,
         firstDimensionDefault=first_dimension_default,
         secondDimensionDefault=second_dimension_default,
+        showSelectors=show_selectors,
     )
diff --git a/react/src/activations/TextNeuronActivations.tsx b/react/src/activations/TextNeuronActivations.tsx
index 7d8562eb..9e4f1b7b 100644
--- a/react/src/activations/TextNeuronActivations.tsx
+++ b/react/src/activations/TextNeuronActivations.tsx
@@ -38,7 +38,8 @@ export function TextNeuronActivations({
   firstDimensionLabels,
   secondDimensionLabels,
   firstDimensionDefault = 0,
-  secondDimensionDefault = 0
+  secondDimensionDefault = 0,
+  showSelectors = true
 }: TextNeuronActivationsProps) {
   // If there is only one sample (i.e. if tokens is an array of strings), cast tokens and activations to an array with
   // a single element
@@ -98,77 +99,62 @@ export function TextNeuronActivations({
 
   return (
     <Container fluid>
-      <Row>
-        <Col>
-          <Row style={selectRowStyle}>
-            <Col>
-              <label htmlFor="layer-selector" style={{ marginRight: 15 }}>
-                {firstDimensionName}:
-              </label>
-              <NumberSelector
-                id="layer-selector"
-                largestNumber={numberOfLayers! - 1}
-                currentValue={layerNumber}
-                setCurrentValue={setLayerNumber}
-                labels={firstDimensionLabels}
-              />
-            </Col>
-          </Row>
-          <Row style={selectRowStyle}>
-            <Col>
-              <label htmlFor="neuron-selector" style={{ marginRight: 15 }}>
-                {secondDimensionName}:
-              </label>
-              <NumberSelector
-                id="neuron-selector"
-                largestNumber={numberOfNeurons! - 1}
-                currentValue={neuronNumber}
-                setCurrentValue={setNeuronNumber}
-                labels={secondDimensionLabels}
-              />
-            </Col>
-          </Row>
-          {/* Only show the sample selector if there is more than one sample */}
-          {numberOfSamples > 1 && (
+      {showSelectors && (
+        <Row>
+          <Col>
             <Row style={selectRowStyle}>
               <Col>
-                <label htmlFor="sample-selector" style={{ marginRight: 15 }}>
-                  Samples:
+                <label htmlFor="layer-selector" style={{ marginRight: 15 }}>
+                  {firstDimensionName}:
                 </label>
-                <RangeSelector
-                  id="sample-selector"
-                  largestNumber={numberOfSamples - 1}
-                  currentRangeArr={sampleNumbers}
-                  setCurrentValue={setSampleNumbers}
-                  numValsInRange={samplesPerPage}
+                <NumberSelector
+                  id="layer-selector"
+                  largestNumber={numberOfLayers! - 1}
+                  currentValue={layerNumber}
+                  setCurrentValue={setLayerNumber}
+                  labels={firstDimensionLabels}
                 />
               </Col>
             </Row>
-          )}
-        </Col>
-        <Col>
-          {/* Only show the sample per page selector if there is more than one sample */}
-          {numberOfSamples > 1 && (
             <Row style={selectRowStyle}>
               <Col>
-                <label
-                  htmlFor="samples-per-page-selector"
-                  style={{ marginRight: 15 }}
-                >
-                  Samples per page:
+                <label htmlFor="neuron-selector" style={{ marginRight: 15 }}>
+                  {secondDimensionName}:
                 </label>
                 <NumberSelector
-                  id="samples-per-page-selector"
-                  smallestNumber={1}
-                  largestNumber={numberOfSamples}
-                  currentValue={samplesPerPage}
-                  setCurrentValue={setSamplesPerPage}
+                  id="neuron-selector"
+                  largestNumber={numberOfNeurons! - 1}
+                  currentValue={neuronNumber}
+                  setCurrentValue={setNeuronNumber}
+                  labels={secondDimensionLabels}
                 />
               </Col>
             </Row>
-          )}
-        </Col>
-      </Row>
+          </Col>
+          <Col>
+            {/* Only show the sample per page selector if there is more than one sample */}
+            {numberOfSamples > 1 && (
+              <Row style={selectRowStyle}>
+                <Col>
+                  <label
+                    htmlFor="samples-per-page-selector"
+                    style={{ marginRight: 15 }}
+                  >
+                    Samples per page:
+                  </label>
+                  <NumberSelector
+                    id="samples-per-page-selector"
+                    smallestNumber={1}
+                    largestNumber={numberOfSamples}
+                    currentValue={samplesPerPage}
+                    setCurrentValue={setSamplesPerPage}
+                  />
+                </Col>
+              </Row>
+            )}
+          </Col>
+        </Row>
+      )}
       <Row>
         <Col>
           <SampleItems
@@ -231,4 +217,9 @@ export interface TextNeuronActivationsProps {
    */
   secondDimensionDefault?: number;
 
+
+    /**
+   * Whether to show the selector dropdowns
+   */
+  showSelectors?: boolean;
 }

From 576ba90855d6e1265e8a8d7f00ff73366a790f21 Mon Sep 17 00:00:00 2001
From: ojh31 <oskar.hollinsworth@gmail.com>
Date: Wed, 20 Sep 2023 14:01:12 +0100
Subject: [PATCH 3/9] Labelled show_selectors as optional

---
 python/circuitsvis/activations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/circuitsvis/activations.py b/python/circuitsvis/activations.py
index e6547c02..8abd4277 100644
--- a/python/circuitsvis/activations.py
+++ b/python/circuitsvis/activations.py
@@ -15,7 +15,7 @@ def text_neuron_activations(
     second_dimension_labels: Optional[List[str]] = None,
     first_dimension_default: Optional[int] = 0,
     second_dimension_default: Optional[int] = 0,
-    show_selectors: bool = True,
+    show_selectors: Optional[bool] = True,
 ) -> RenderedHTML:
     """Show activations (colored by intensity) for each token in a text or set
     of texts.

From 87af8ca78ca76da7c18a07f9f4aaf7b4b0aee2a2 Mon Sep 17 00:00:00 2001
From: skar0 <oskar.hollinsworth@gmail.com>
Date: Mon, 20 Nov 2023 14:55:05 +0000
Subject: [PATCH 4/9] Updated snap_test_activations for new args

---
 .../tests/snapshots/snap_test_activations.py  | 24 +++++--------------
 1 file changed, 6 insertions(+), 18 deletions(-)

diff --git a/python/circuitsvis/tests/snapshots/snap_test_activations.py b/python/circuitsvis/tests/snapshots/snap_test_activations.py
index aaaf8501..8f9fad15 100644
--- a/python/circuitsvis/tests/snapshots/snap_test_activations.py
+++ b/python/circuitsvis/tests/snapshots/snap_test_activations.py
@@ -7,22 +7,10 @@
 
 snapshots = Snapshot()
 
-snapshots['TestTextNeuronActivations.test_multi_matches_snapshot 1'] = '''<div id="circuits-vis-mock" style="margin: 15px 0;"/>
-    <script crossorigin type="module">
-    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";
-    render(
-      "circuits-vis-mock",
-      TextNeuronActivations,
-      {"tokens": [["a", "b"], ["c", "d", "e"]], "activations": [[[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]], [[0, 1, 1], [1, 1, 1]]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron"}
-    )
-    </script>'''
+snapshots[
+    "TestTextNeuronActivations.test_multi_matches_snapshot 1"
+] = """<div id="circuits-vis-mock" style="margin: 15px 0;"/>\n    <script crossorigin type="module">\n    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";\n    render(\n      "circuits-vis-mock",\n      TextNeuronActivations,\n      {"tokens": [["a", "b"], ["c", "d", "e"]], "activations": [[[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]], [[0, 1, 1], [1, 1, 1]]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron", "firstDimensionDefault": 0, "secondDimensionDefault": 0, "showSelectors": true}\n    )\n    </script>"""
 
-snapshots['TestTextNeuronActivations.test_single_matches_snapshot 1'] = '''<div id="circuits-vis-mock" style="margin: 15px 0;"/>
-    <script crossorigin type="module">
-    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";
-    render(
-      "circuits-vis-mock",
-      TextNeuronActivations,
-      {"tokens": ["a", "b"], "activations": [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron"}
-    )
-    </script>'''
+snapshots[
+    "TestTextNeuronActivations.test_single_matches_snapshot 1"
+] = """<div id="circuits-vis-mock" style="margin: 15px 0;"/>\n    <script crossorigin type="module">\n    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";\n    render(\n      "circuits-vis-mock",\n      TextNeuronActivations,\n      {"tokens": ["a", "b"], "activations": [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron", "firstDimensionDefault": 0, "secondDimensionDefault": 0, "showSelectors": true}\n    )\n    </script>"""

From 441dd90d3729c40c9da9273a6b4d08e5896d5713 Mon Sep 17 00:00:00 2001
From: skar0 <oskar.hollinsworth@gmail.com>
Date: Mon, 20 Nov 2023 15:06:54 +0000
Subject: [PATCH 5/9] Fixed eslint in textneuronactivations

---
 react/src/activations/TextNeuronActivations.tsx | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/react/src/activations/TextNeuronActivations.tsx b/react/src/activations/TextNeuronActivations.tsx
index 9e4f1b7b..5cacae94 100644
--- a/react/src/activations/TextNeuronActivations.tsx
+++ b/react/src/activations/TextNeuronActivations.tsx
@@ -72,7 +72,9 @@ export function TextNeuronActivations({
     ...Array(samplesPerPage).keys()
   ]);
   const [layerNumber, setLayerNumber] = useState<number>(firstDimensionDefault);
-  const [neuronNumber, setNeuronNumber] = useState<number>(secondDimensionDefault);
+  const [neuronNumber, setNeuronNumber] = useState<number>(
+    secondDimensionDefault
+  );
 
   useEffect(() => {
     // When the user changes the samplesPerPage, update the sampleNumbers
@@ -217,8 +219,7 @@ export interface TextNeuronActivationsProps {
    */
   secondDimensionDefault?: number;
 
-
-    /**
+  /**
    * Whether to show the selector dropdowns
    */
   showSelectors?: boolean;

From 2964a4dd0181e83599ebb123ada874f139e96798 Mon Sep 17 00:00:00 2001
From: skar0 <oskar.hollinsworth@gmail.com>
Date: Mon, 20 Nov 2023 15:15:40 +0000
Subject: [PATCH 6/9] Removed RangeSelector import

---
 react/src/activations/TextNeuronActivations.tsx | 1 -
 1 file changed, 1 deletion(-)

diff --git a/react/src/activations/TextNeuronActivations.tsx b/react/src/activations/TextNeuronActivations.tsx
index 5cacae94..127db70e 100644
--- a/react/src/activations/TextNeuronActivations.tsx
+++ b/react/src/activations/TextNeuronActivations.tsx
@@ -2,7 +2,6 @@ import { Rank, tensor, Tensor1D, Tensor3D } from "@tensorflow/tfjs";
 import React, { useState, useEffect } from "react";
 import { Container, Row, Col } from "react-grid-system";
 import { SampleItems } from "../shared/SampleItems";
-import { RangeSelector } from "../shared/RangeSelector";
 import { NumberSelector } from "../shared/NumberSelector";
 import { minMaxInNestedArray } from "../utils/arrayOps";
 

From 036d3ba152176f56e0cd380fab770ef897897c58 Mon Sep 17 00:00:00 2001
From: skar0 <oskar.hollinsworth@gmail.com>
Date: Mon, 20 Nov 2023 15:19:22 +0000
Subject: [PATCH 7/9] Added type hint for tensor input to topk

---
 python/circuitsvis/topk_samples.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/python/circuitsvis/topk_samples.py b/python/circuitsvis/topk_samples.py
index 92d772ff..1d9a6b15 100644
--- a/python/circuitsvis/topk_samples.py
+++ b/python/circuitsvis/topk_samples.py
@@ -1,12 +1,16 @@
 """Activations visualizations"""
-from typing import List, Optional
+from typing import List, Optional, Union
+from jaxtyping import Float
+from torch import Tensor
 
 from circuitsvis.utils.render import RenderedHTML, render
 
 
 def topk_samples(
     tokens: List[List[List[List[str]]]],
-    activations: List[List[List[List[float]]]],
+    activations: Union[
+        List[List[List[List[float]]]], Float[Tensor, "layer neuron sample token"]
+    ],
     zeroth_dimension_name: Optional[str] = "Layer",
     first_dimension_name: Optional[str] = "Neuron",
     zeroth_dimension_labels: Optional[List[str]] = None,

From 6b67eac973f89aefdb0e2a254b6a7ab3f50e805a Mon Sep 17 00:00:00 2001
From: skar0 <oskar.hollinsworth@gmail.com>
Date: Mon, 20 Nov 2023 15:37:28 +0000
Subject: [PATCH 8/9] Removed jaxtyping import

---
 .gitignore                         | 1 +
 python/circuitsvis/topk_samples.py | 7 ++-----
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index 6e0a7cce..56b311f0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,6 +11,7 @@ __pycache__
 venv
 *-checkpoint.ipynb
 requirements.txt
+python/circuitsvis/.vscode
 
 # Node
 node_modules/
diff --git a/python/circuitsvis/topk_samples.py b/python/circuitsvis/topk_samples.py
index 1d9a6b15..45058671 100644
--- a/python/circuitsvis/topk_samples.py
+++ b/python/circuitsvis/topk_samples.py
@@ -1,16 +1,13 @@
 """Activations visualizations"""
+import torch
 from typing import List, Optional, Union
-from jaxtyping import Float
-from torch import Tensor
 
 from circuitsvis.utils.render import RenderedHTML, render
 
 
 def topk_samples(
     tokens: List[List[List[List[str]]]],
-    activations: Union[
-        List[List[List[List[float]]]], Float[Tensor, "layer neuron sample token"]
-    ],
+    activations: Union[List[List[List[List[float]]]], torch.Tensor],
     zeroth_dimension_name: Optional[str] = "Layer",
     first_dimension_name: Optional[str] = "Neuron",
     zeroth_dimension_labels: Optional[List[str]] = None,

From 35d72d68107d0be7bc912100483e714e6ba31cb4 Mon Sep 17 00:00:00 2001
From: skar0 <oskar.hollinsworth@gmail.com>
Date: Mon, 20 Nov 2023 15:49:33 +0000
Subject: [PATCH 9/9] Added back the sample-selector

---
 .../src/activations/TextNeuronActivations.tsx  | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/react/src/activations/TextNeuronActivations.tsx b/react/src/activations/TextNeuronActivations.tsx
index 127db70e..e645462f 100644
--- a/react/src/activations/TextNeuronActivations.tsx
+++ b/react/src/activations/TextNeuronActivations.tsx
@@ -2,6 +2,7 @@ import { Rank, tensor, Tensor1D, Tensor3D } from "@tensorflow/tfjs";
 import React, { useState, useEffect } from "react";
 import { Container, Row, Col } from "react-grid-system";
 import { SampleItems } from "../shared/SampleItems";
+import { RangeSelector } from "../shared/RangeSelector";
 import { NumberSelector } from "../shared/NumberSelector";
 import { minMaxInNestedArray } from "../utils/arrayOps";
 
@@ -131,6 +132,23 @@ export function TextNeuronActivations({
                 />
               </Col>
             </Row>
+            {/* Only show the sample selector if there is more than one sample */}
+            {numberOfSamples > 1 && (
+              <Row style={selectRowStyle}>
+                <Col>
+                  <label htmlFor="sample-selector" style={{ marginRight: 15 }}>
+                    Samples:
+                  </label>
+                  <RangeSelector
+                    id="sample-selector"
+                    largestNumber={numberOfSamples - 1}
+                    currentRangeArr={sampleNumbers}
+                    setCurrentValue={setSampleNumbers}
+                    numValsInRange={samplesPerPage}
+                  />
+                </Col>
+              </Row>
+            )}
           </Col>
           <Col>
             {/* Only show the sample per page selector if there is more than one sample */}