Update app.py

app.py

@@ -168,7 +168,7 @@ def __(ChessBoard, dropdown_fen, dropdown_moves):
 
 @app.cell
 def __(torch):
-    def rollout(x, skip_connection=True, parse="min"):
+    def rollout(x, skip_last_layers=0, skip_connection=True, parse="min"):
         attns = []
         for k, v in x.items():
             v = v[0, :, ::-1, :]
@@ -179,10 +179,9 @@ def __(torch):
                 item = torch.max(v, dim=0).values
             elif parse == "mean":
                 item = torch.mean(v, dim=0)
-        attns.append(item + int(skip_connection) * torch.eye(item.shape[0]))
-        roll = torch.prod(torch.stack(attns), dim=0)
+            attns.append(item + int(skip_connection) * torch.eye(item.shape[0]))
+        roll = torch.prod(torch.stack(attns)[:skip_last_layers], dim=0)
         return roll
-
     return (rollout,)
 
 
@@ -202,30 +201,68 @@ def __(mo):
 
 
 @app.cell
-def __(switch):
-    switch.value
-    return
+def __(mo):
+    max_value_switch = mo.ui.switch(value=False, label="use rollout layers with max value")
+    max_value_switch
+    return (max_value_switch,)
+
+
+@app.cell
+def __(mo):
+    highlight_squares_switch = mo.ui.switch(value=False, label="Use rollout to highlight squares")
+    highlight_squares_switch
+    return (highlight_squares_switch,)
 
 
 @app.cell
-def __(dropdown_method, focus_square, rollout, switch, torch):
+def __(
+    dropdown_layer,
+    dropdown_method,
+    focus_square,
+    highlight_squares_switch,
+    max_value_switch,
+    rollout,
+    switch,
+    torch,
+):
     import chess 
     from global_data import global_data
 
+    def find_max(a):
+        ar = a.reshape(a.shape[0], -1)
+        i = torch.max(ar, dim=1).values
+        im = torch.argmax(i[1:])
+        return a[im + 1]
+
     def parse_activations(act, layer_number=None):
         if dropdown_method.value == "Attention visualization":
             layer_key = [k for k in act.keys() if "0" in k][0].replace("0", f"{layer_number}")
             a = act[layer_key][0, :, ::-1 , :]
         elif dropdown_method.value == "Attention rollout (MIN)":
-            a = rollout(act, skip_connection=switch.value, parse="min")
-            a = torch.stack([a for _ in range(32)]).numpy()
+            if not max_value_switch.value:
+                a = rollout(act, skip_last_layers=int(dropdown_layer.value), skip_connection=switch.value, parse="min")
+            else:
+                a = torch.stack([rollout(act, skip_last_layers=_i, skip_connection=switch.value, parse="min") for _i in range(0, 15)], dim=0)
+                a = find_max(a)
         elif dropdown_method.value == "Attention rollout (MAX)":
-            a = rollout(act, skip_connection=switch.value, parse="max")
-            a = torch.stack([a for _ in range(32)]).numpy()
+            if not max_value_switch.value:
+                a = rollout(act, skip_last_layers=int(dropdown_layer.value), skip_connection=switch.value, parse="max")
+            else:
+                a = torch.stack([rollout(act, skip_last_layers=_i, skip_connection=switch.value, parse="max") for _i in range(0, 15)], dim=0)
+                a = find_max(a)
         elif dropdown_method.value == "Attention rollout (MEAN)":
-            a = rollout(act, skip_connection=switch.value, parse="mean")
-            a = torch.stack([a for _ in range(32)]).numpy()
+            if not max_value_switch.value:
+                a = rollout(act, skip_last_layers=int(dropdown_layer.value), skip_connection=switch.value, parse="mean")
+            else:
+                a = torch.stack([rollout(act, skip_last_layers=_i, skip_connection=switch.value, parse="mean") for _i in range(0, 15)], dim=0)
+                a = find_max(a)
+        if dropdown_method.value != "Attention visualization":
+            if highlight_squares_switch.value:
+                a = a.max(dim=0).values
+                a = torch.stack([a for _ in range(64)], dim=0)
+            a = torch.stack([a for _ in range(32)]).numpy() 
         return a
+
     focus_square_ind = 8 * (int(focus_square.value[1]) - 1) + ord(focus_square.value[0]) - ord("a")
 
     def set_plotting_parameters(act, fen):
@@ -246,6 +283,7 @@ def __(dropdown_method, focus_square, rollout, switch, torch):
         global_data.show_colorscale = False
     return (
         chess,
+        find_max,
         focus_square_ind,
         global_data,
         parse_activations,