TopkSamples; Consistent color scale for layer and neuron

TransformerLensOrg · Feb 20, 2023 · e845aa4 · e845aa4
1 parent d11b40e
commit e845aa4
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/react/src/topk/TopkSamples.tsx b/react/src/topk/TopkSamples.tsx
@@ -17,9 +17,6 @@ export function TopkSamples({
   zerothDimensionLabels,
   firstDimensionLabels
 }: TopkSamplesProps) {
-  // Obtain min and max activations for a consistent color scale across all samples
-  const [minValue, maxValue] = minMaxInNestedArray(activations);
-
   const numberOfLayers = activations.length;
   const numberOfNeurons = activations[0].length;
   const numberOfSamples = activations[0][0].length;
@@ -46,6 +43,11 @@ export function TopkSamples({
     return tokens[layerNumber][neuronNumber][sampleNumber];
   });
 
+  // For a consistent color scale across all samples in this layer and neuron
+  const [minValue, maxValue] = minMaxInNestedArray(
+    activations[layerNumber][neuronNumber]
+  );
+
   const selectRowStyle = {
     paddingTop: 5,
     paddingBottom: 5