Feature request: Visualize attention rollout from the output to the input

Hi,

I am trying to implement the attention rollout for the CrossFormer to visualize the attention from the output token to the input
space. Please take a look at the image below. Do you think this feature is something that you could consider adding?

<img width="197" alt="image" src="https://github.com/user-attachments/assets/d72a3b5c-80f9-4050-8075-341e99dfe0a0">
