[Flax] [WIP] allow loading head model with base model weights
# What does this PR do?
Allows loading flax head model with base model weights.

Right now it's not possible to load a flax head model using weights from the base model as weights `dict` of the base model does not contain the `base_model_prefix` key. To reproduce

```python
from transformers import BertConfig, FlaxBertModel, FlaxBertForSequenceClassification
config = BertConfig(hidden_size=64, intermediate_size=128, max_position_embeddings=128, num_attention_heads=8, num_hidden_layers=8)
base_model = FlaxBertModel(config)
base_model.save_pretrained("base")
head_model = FlaxBertForSequenceClassification.from_pretrained("base")
```
Thanks for fixing this! Could we maybe also add some tests that make sure that loading/saving works correctly for all models?