

TITLE
[Flax] [WIP] allow loading head model with base model weights

COMMENTS
1

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
# 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")
```