[Flax] Add attention weights outputs to all models
# 🚀 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`.

Also, adding this feature will require us to bump up the Flax dependency to `>=0.3.4` for `flax` in https://github.com/huggingface/transformers/blob/master/setup.py<|||||>I am starting with Bert.<|||||>Hi @patrickvonplaten @patil-suraj https://github.com/huggingface/transformers/pull/11918