From 8605b162ce19b24efda3389492ac4f2e25b33fb1 Mon Sep 17 00:00:00 2001
From: UFO-101 <47218308+UFO-101@users.noreply.github.com>
Date: Tue, 1 Aug 2023 18:03:49 +0100
Subject: [PATCH] Fix AttentionHead resizing bug (#63)

---
 react/src/attention/AttentionHeads.tsx   | 13 ++-----
 react/src/attention/AttentionPattern.tsx | 46 +++++++++++++++++++-----
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/react/src/attention/AttentionHeads.tsx b/react/src/attention/AttentionHeads.tsx
index f9c816c..6c601d0 100644
--- a/react/src/attention/AttentionHeads.tsx
+++ b/react/src/attention/AttentionHeads.tsx
@@ -143,17 +143,7 @@ export function AttentionHeads({
       <Row>
         <Col xs={12}>
           <h3 style={{ marginBottom: 10 }}>{headNames[focused]} Zoomed</h3>
-
-          <div
-            style={{
-              position: "relative",
-              // Set the maximum width such that a head with just a few tokens
-              // doesn't have crazy large boxes per token. Note this is the
-              // width of the full chart (including axis labels) so it also
-              // needs a sensible lowest maximum.
-              maxWidth: `${Math.max(Math.round(tokens.length * 2.4), 20)}em`
-            }}
-          >
+          <div>
             <h2
               style={{
                 position: "absolute",
@@ -174,6 +164,7 @@ export function AttentionHeads({
               minValue={minValue}
               negativeColor={negativeColor}
               positiveColor={positiveColor}
+              zoomed={true}
               tokens={tokens}
             />
           </div>
diff --git a/react/src/attention/AttentionPattern.tsx b/react/src/attention/AttentionPattern.tsx
index 823428d..0bdcfd7 100644
--- a/react/src/attention/AttentionPattern.tsx
+++ b/react/src/attention/AttentionPattern.tsx
@@ -61,6 +61,7 @@ export function AttentionPattern({
   positiveColor,
   upperTriColor = DefaultUpperTriColor,
   showAxisLabels = true,
+  zoomed = false,
   tokens
 }: AttentionPatternProps) {
   // Tokens must be unique (for the categories), so we add an index prefix
@@ -164,15 +165,37 @@ export function AttentionPattern({
 
   return (
     <Col>
-      <Row style={{ aspectRatio: showAxisLabels ? undefined : "1/1" }}>
-        <Chart
-          type="matrix"
-          options={options}
-          data={data}
-          width={1000}
-          height={1000}
-          updateMode="none"
-        />
+      <Row>
+        <div
+          style={{
+            // Chart.js charts resizing is weird.
+            // Responsive chart elements (which all are by default) require the
+            // parent element to have position: 'relative' and no sibling elements.
+            // There were previously issues that only occured at particular display
+            // sizes and zoom levels. See:
+            // https://github.com/alan-cooney/CircuitsVis/pull/63
+            // https://www.chartjs.org/docs/latest/configuration/responsive.html#important-note
+            // https://stackoverflow.com/a/48770978/7086623
+            position: "relative",
+            // Set the maximum width of zoomed heads such that a head with just a
+            // few tokens doesn't have crazy large boxes per token and the chart
+            // doesn't overflow the screen. Other heads fill their width.
+            maxWidth: zoomed
+              ? `min(100%, ${Math.round(tokens.length * 8)}em)`
+              : "initial",
+            width: zoomed ? "initial" : "100%",
+            aspectRatio: "1/1"
+          }}
+        >
+          <Chart
+            type="matrix"
+            options={options}
+            data={data}
+            width={1000}
+            height={1000}
+            updateMode="none"
+          />
+        </div>
       </Row>
     </Col>
   );
@@ -255,6 +278,11 @@ export interface AttentionPatternProps {
    */
   showAxisLabels?: boolean;
 
+  /**
+   * Is this a zoomed in view?
+   */
+  zoomed?: boolean;
+
   /**
    * List of tokens
    *