Remove boiler plate code
Author: patrickvonplaten
Pull request: yes
Labels: name: Flax

# What does this PR do?

In Flax, the logic of every model has to be built purely on a tree of `flax.linen.Module` classes so that the whole model, *e.g.* `FlaxBertForMaskedLMModule`, can automatically be recast as an explicit function. This allows for easy jitting, data parallelism, and model parallelism. 
However this also means that no weight parameters can be stored in `FlaxBertForMaskedLMModule`, which is a bit problematic since we need to store the weights from `load_from_pretrained`. This forces us to have two Flax classes: 
- `FlaxBertForMaskedLMModule`, 
- `FlaxBertForMaskedLM`,

whereas `FlaxBertForMaskedLM` takes care of loading/saving the weights and `FlaxBertForMaskedLMModule` defines the logic (function) of the model. This has led to a lot of boilerplate code since `FlaxBertModel`, `FlaxBertForMaskedLM`, `FlaxBertForPretraining`, ... are essentially all the same:

- they take the same input and pass it to the corresponding `flax.linen.Module` forward function
- they initialize a `FlaxPretrainedModel` with the correct module to inherit the loading/saving functionality
- they make use of the same `init_weights` function. 

For BERT, the `__call__` functions take identical inputs across different classes. This assumes that this holds true for more or less all BERT-like models in Flax. However, there are some exceptions:

- A `BertForCausalLM` (needed when we add `FlaxEncoderDecoderModel`) would have to overwrite the `__call__` method as it also takes `encoder_hidden_states`, `encoder_attention_mask` as an input.
- This design makes less sense for *e.g.* T5 since `T5Encoder` takes very much different inputs as `T5ForConditionalGeneration`. Here the `__call__` and `init_weights` methods would then not be placed in `T5PreTrainedModel`.

Overall, I think however that removing this much boilerplate is cleaner and will allow us to implement the important models faster. What do you think @sgugger @LysandreJik @avital ?


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->


Comments: