[Flax] Add FlaxBlenderbot
Author: stancld
Pull request: yes

# What does this PR do?
This PR adds flax implementation of Blenderbot.


## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?

<hr>

### TODOs:

- fix PT-Flax model equivalence


## Who can review?

@patrickvonplaten  @patil-suraj

Comments:
@patrickvonplaten I would like to kindly ping for a review. :) I've been struggling to achieve the pt-flax equivalence, however, I cannot find that difference/bug in this new flax implementation.

Thanks a lot! :) 

Hey @stancld, 

Thanks a lot for the PR! The difference between PT and Flax in your PR is very close actually < 0.1 so it might also very well be that the implementation is correct! 

I'll try to take a deeper look at the end of next week. Could you try one last thing:

add print statements such as:

`print("PT", hidden_states.sum())` in PyTorch

and 

`print("FX", hidden_states.sum())` in Flax

before the word embeddings, after the word embeddings, each encoder transformer layer, before the decoder word embeddings, the decoder attention layers, ... to see when the activations start to diverge. If it happens gradually it might very well be that the model is correct and there is a difference. If it haapens all of a sudden at some point then there might be a subtle bug.

@patrickvonplaten Thank you for the tip! I'll have a look :) 

Hello @patrickvonplaten, I ran a few tests it seems and one output is below. There is some level of divergence, but not sure if it's too severe. I'm gonna check the Flax code today once again :)

```
===PyTorch===
---Encoder---
PT first hidden-states:  tensor(-1.2589)
PT encoder after self-attn:  tensor(0.5862)
PT encoder: tensor(-0.7895)
PT encoder after self-attn:  tensor(0.0465)
PT encoder last hidden states before norm:  tensor(-0.2601)
PT encoder last hidden states after norm:  tensor(0.)
---Decoder---
PT decoder after self-attn:  tensor(1.1000)
PT decoder after cross-attn:  tensor(0.1547)
PT decoder: tensor(-0.0142)
PT decoder after self-attn:  tensor(0.9638)
PT decoder after cross-attn:  tensor(1.7759)
PT decoder: tensor(2.7198)
PT decoder last hidden states before norm:  tensor(2.7198)
PT decoder last hidden states after norm:  tensor(-5.7220e-06)
PT output:  tensor(-5.7220e-06)

===Flax===
---Encoder---
FX first hidden-states:  -1.2589027
FX encoder after self-attn:  0.59013414
FX encoder: -0.7862803
FX encoder after self-attn:  0.04762589
FX encoder last hidden states before norm:  -0.25001374
FX encoder last hidden states after norm:  4.053116e-06
---Decoder---
FX decoder after self-attn:  1.1029385
FX decoder after cross-attn:  0.15325405
FX decoder: -0.013041288
FX decoder after self-attn:  0.96520036
FX decoder after cross-attn:  1.7912248
FX decoder last hidden states before norm:  2.735363
FX decoder last hidden states after norm:  -1.1697412e-06
FX output:  -1.1697412e-06
```

@patrickvonplaten Thank you very much for spotting the problem! :] 

Tests on master seem to be broken currently :-/ 

But I think the PR is good to go. @patil-suraj could you maybe take a look once you're back (and maybe rebase to master with @stancld to fix the circli ci runner) 

Awesome - I let you merge @patil-suraj once you're back :-)