[Flax] Add FlaxBart models
Author: stancld
Pull request: yes

# What does this PR do?

This PR adds Flax implementation of BART and classes for various downstream tasks. Fixes #11478.

Most of the code is inspired by the Flax implementation of BERT and PyTorch implementation of BART.

From Suraj:
A couple of important points to note:

- The seq2seq API is slightly different from the PyTorch BART model in that the `__call__` method of `FlaxBART` does not accept `encoder_outputs`. In PT if `encoder_outputs` is passed then the encoder is skipped and only the decoder is called. This is not supported in `FlaxBART` to prevent any unintended issues during JIT compilation as skipping module or passing different inputs to function causes re-compilation. Also, the idiomatic way of accessing intermediate modules in Flax models is to expose explicit methods.  So the API is as follows 
   - the `__call__` method expects both the encoder and decoder inputs and does forward pass through both the modules
   - Every model has a `encode` and `decode` method which should be called if one wants to run just the encoder or decoder. The `decode` method only returns the decoder outputs and for `*ForConditionalGeneration` modules it also returns the `logits`
   
```python
# runs encoder and decoder
model(input_ids, decoder_input_ids)

# just run the encoder
encoder_outputs = model.encode(input_ids)

# run the decoder
decoder_outputs = model.decode(decoder_input_ids, encoder_outputs)
```

- For now `past_key_values` caching is only implemented in the self-attention layer in the decoder i.e it's not implemented for the cross-attention layer. 

<hr>

**Reviewers:** @patrickvonplaten @sgugger @patil-suraj (and whoever else in the community)

Comments:
Hey @stancld,

One important thing we should implement before merging is the caching mechanismh similar to how it's done in GPT2: https://github.com/huggingface/transformers/blob/0b0a598452b02278075a75f84b5ca7bb457224ad/src/transformers/models/gpt2/modeling_flax_gpt2.py#L139. @patil-suraj, could you help here maybe? :-) 

Sure @patrickvonplaten !

@stancld  let me know if you need help here or I could take this if you are okay with it :)

@patrickvonplaten  a couple of last questions about the API

- Right now the `FlaxBartPretrainedModel.__call__` method also accepts the `encoder_outputs`, but think it would be cleaner to not do that as we already have `encode/decode` methods and let the user use `encode` or `decode` if they want to run just one part of the model. So we could make `decode` method available for every model (right now it's only available for the `ForConditionalGeneration` model) and it'll return the decoder outputs and for `*ForConditionalGeneration` models, it'll also return the `logits`.

- The `decode` method returns `FlaxSeq2SeqLM` output, which includes both the encoder and decoder outputs, but when calling  `decode` the user already has `encoder_outputs`, so maybe we should just return decoder outputs since the `decode` method only runs the decoder.

What do you think? 


IMO:

1) Agree, happy to remove `encoder_outptus` as an input argument from `call` & make `decode` available for all models
2) Yes, it doesn't make too much sense to include all the encoder relevant output here!