No way to unlock focus on `cv.attention.attention_patterns`?

[JasonGross]

I'm looking at https://arena-ch1-transformers.streamlit.app/[1.2]_Intro_to_Mech_Interp / https://colab.research.google.com/drive/1w9zCWpE7xd1sDuMT_rsjARfFozeWiKF4, in particular the "Visualising Attention Heads" section, with the code

print(type(gpt2_cache))
attention_pattern = gpt2_cache["pattern", 0, "attn"]
print(attention_pattern.shape)
gpt2_str_tokens = gpt2_small.to_str_tokens(gpt2_text)

print("Layer 0 Head Attention Patterns:")
display(cv.attention.attention_patterns(
    tokens=gpt2_str_tokens,
    attention=attention_pattern,
    attention_head_names=[f"L0H{i}" for i in range(12)],
))

And it seems like once I click on a head and/or token to lock the focus, there's no way to unlock the focus and get back the averaged value. There should be a way to do this, and the visualization should signpost this.

(Also, "Tokens (click to focus)" should probably be "Tokens (hover to focus, click to lock)" much like "Head selector (hover to focus, click to lock)")

[callummcdougall]

+1, this would be quite useful if not too inconvenient to create. (Although as one solution, you can just try and hover over heads rather than actually clicking on them - but that can be difficult if you're showing a lot of heads at once.)