[FlaxBert] Non-broadcastable attention mask in batched forward-pass
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version:
- Platform: Linux
- Python version: 3.8
- JAX version:
  - jax==0.2.6
  - jaxlib==0.1.57+cuda110
  - flax==0.2.2
- PyTorch version (GPU?): n/a
- Tensorflow version (GPU?): n/a
- Using GPU in script?: **yes** (cuda 11.0)
- Using distributed or parallel set-up in script?: **no**

### Who can help
@mfuntowicz 
@avital 
@LysandreJik 

## Information

I ran the script from the recent Twitter [post](https://twitter.com/huggingface/status/1331255460033400834):

![](https://pbs.twimg.com/media/EnmRwGDW4AA0j26?format=jpg)

The only thing I changed was that I fed in multiple sentences:

```python
from transformers import FlaxBertModel, BertTokenizerFast, TensorType

tokenizer = BertTokenizerFast.from_pretrained('bert-base-cased')
model = FlaxBertModel.from_pretrained('bert-base-cased')
# apply_fn = jax.jit(model.model.apply)

sentences = ["this is an example sentence", "this is another", "and a third one"]
encodings = tokenizer(sentences, return_tensors=TensorType.JAX, padding=True, truncation=True)
tokens, pooled = model(**encodings)
```
> ValueError: Incompatible shapes for broadcasting: ((3, 12, 7, 7), (1, 1, 3, 7))

See full stack trace: https://pastebin.com/sPUSjGVi
