Feature request: optionally sow attention weight in dot_product_attention

[JyChang012]

It seems that there is no good way to retrieve attention weight using MultiHeadDotProductAttention. Probably we can add an optional mdl parameter to dot_product_attention and dot_product_attention_weights, and sow the weight inside. This should have no effect when intermediates is not mutable. Here is my implementation. I can create a PR if you like.

Update:
Implemented in #3529

[JyChang012]

Any ideas on how to extract attention weight other than rewriting the attention module? @cgarciae

[cgarciae]

Hey @JyChang012, sorry for the delay. I see that Tensorflow's MultiHeadAttention has a return_attention_scores flag which would solve your use case, however, Flax doesn't have an equivalent feature even though it would be easy to add.

In the mean time, we do encourage users forking our Modules for their own needs, we try to maintain them readable for this purpose.

@marcvanzee WDYT? Should we consider adding this flag?

[JyChang012]

This implementation is applied in PR #3529