text: 
  
  TITLE
  [Flax] Add attention weights outputs to all models
  
  COMMENTS
  3
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  # 🚀 Feature request
  
  At the moment we cannot return a list of attention weight outputs in Flax as we can do in PyTorch. 
  
  In PyTorch, there is a `output_attentions` boolean in the forward call of every function, see [here](https://github.com/huggingface/transformers/blob/42fe0dc23e4a7495ebd08185f5850315a1a12dc0/src/transformers/models/bert/modeling_bert.py#L528) which when set to True collects all attention weights and returns them as a tuple. 
  In PyTorch, the attention weights are returned (if `output_attentions=True`) from the self-attention layer *e.g.* here: https://github.com/huggingface/transformers/blob/42fe0dc23e4a7495ebd08185f5850315a1a12dc0/src/transformers/models/bert/modeling_bert.py#L331 and then passed with the outputs. 
  
  Currently, this is not implemented in Flax and needs to be done. At the moment the function [`dot_product_attention`](https://github.com/google/flax/blob/6fb839c640de80f887580a533b222c6dddf04c0d/flax/linen/attention.py#L109) is used in every Flax model which makes it impossible to retrieve the attention weights, see [here](https://github.com/huggingface/transformers/blob/42fe0dc23e4a7495ebd08185f5850315a1a12dc0/src/transformers/models/bert/modeling_flax_bert.py#L244). However recently the Flax authors refactored this function into a smaller one called [`dot_product_attention_weights`](https://github.com/google/flax/blob/6fb839c640de80f887580a533b222c6dddf04c0d/flax/linen/attention.py#L37) which would allow us to correctly retrieve the attention weights if needed. To do so all `dot_product_attention` functions should be replaced by `dot_product_attention_weights`, followed by a `jnp.einsum`, see [here](https://github.com/google/flax/blob/6fb839c640de80f887580a533b222c6dddf04c0d/flax/linen/attention.py#L162) so that we can retrieve the attention weights. 
  Next, the whole `output_attentions` logic should be implemented for all Flax models analog to `output_hidden_states`.
  