[Flax] Serialization, Design changes
This PR is a proposition for some changes of the flax design. @mfuntowicz - I'm trying to get a better understanding of how to best use the flax library with Transformers' philosophy. Would be super happy about some comments from your side :-) 

1) I think Flax's `from_pretrained()` should default to the flax serialization and not the PyTorch one. Flax's serialization didn't work previously and the model was loaded from PyTorch by default. This PR changes this to Flax and makes `from_pretrained()` and `save_pretrained()` work. I uploaded Bert's and Roberta's flax model weights to the model hub (I noticed that I accidently overwrote an existing Flax `bert-base-cased` - hope that was fine @mfuntowicz - it's doesn't break anything on master since PT was loaded by default)

2) Not sure why we have the `model_class` class attribute in Flax - I don't think we need it no? @mfuntowicz - It would be nice to avoid it for simplicity IMO.

3) I added a `FlaxBertPretrainedModel` class, just as it's done for PyTorch. IMO, ideally we should stay as close as possible to the design of PyTorch. Not sure at all if something like this could work:

```python 
class FlaxBertForMaskedLM(FlaxBertPretrainedModel):

         def __init__(self, config, state, seed, **kwargs):
               self.bert = FlaxBertModel(config, state[self.base_model_prefix], seed, **kwargs)  # pass bert relevant 
        
        @nn.compact
        def __call__(....):
               last_hidden_states = self.bert(hidden_states)[0]
               logits = FlaxBertLMPredictionHead(vocab_size=self.vocab_size, name="mlm", dtype=self.dtype)(last_hidden_states)     
```
=> What do you think @mfuntowicz ?

Would be awesome if we could do some flax library design discussions here @mfuntowicz @LysandreJik @sgugger @thomwolf 
I think I'll take a Flax 101 class before reviewing this PR.<|||||>Requires more  discussion