Visualization of attention weights in FlatFormer

[bunnyveil]

I would like to express my gratitude for your work on FlatFormer. I am currently exploring the model and would like to visualize the attention weights generated during inference. Could you please guide me on how to implement the visualization of attention weights within FlatFormer? Specifically, I am looking for existing code snippets that could help me generate visualizations similar to what might be seen in transformer models.
Any assistance or direction on this matter would be greatly appreciated. Thank you in advance for your time and support.