Port Jax weights to PyTorch
Author: jxiw
Pull request: no

I have a question when convert JAX weights to torch weights.

I follow this to convert the weights. 
https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flax_pytorch_utils.py#L331-L441

However, there are two parameters that i cannot find in JAX weights [‘cls.predictions.decoder.bias’, ‘cls.predictions.decoder.weight’] in my model msgpack

And I found BERT also has this problem.

```
bert=FlaxBertForMaskedLM.from_pretrained("bert-large-uncased")
bert.params['cls']['predictions'].keys()
dict_keys(['bias', 'transform'])
```

So there is no decoder variable here. But i can find this in the class 

https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_flax_bert.py#L707

Do you know why it is? And do i need to also copy the weights of this decoder into pytorch model convert JAX MaskedLM to Pytorch MaskedLM?

@sanchit-gandhi helped me to fix this! I just post it here! 

Comments:
Hey! This is most probably because the weights for this are shared with the embedding layer. Would recommend you to dive in the code and check for `tied_weigth_keys` or if the  config has `tie_word_embedding` attribute 😉 
Some [info on the forum](https://discuss.huggingface.co/t/what-is-the-tie-word-embeddings-option-exactly-doing/8483) should also be easy to access! 

See here how we initialise a linear LM head under the attribute `self.decoder`: https://github.com/huggingface/transformers/blob/ab0ddc99e853c974949d823dbfaa732202696f3e/src/transformers/models/bert/modeling_flax_bert.py#L707
However if we pass `shared_embedding`, then we use the kernel/bias weights from this shared embedding that is passed as an argument: https://github.com/huggingface/transformers/blob/ab0ddc99e853c974949d823dbfaa732202696f3e/src/transformers/models/bert/modeling_flax_bert.py#L713-L714
Now if you look at `FlaxBertForMaskedLMModule`, you see that we pass the input word embeddings as our `shared_embedding`: https://github.com/huggingface/transformers/blob/ab0ddc99e853c974949d823dbfaa732202696f3e/src/transformers/models/bert/modeling_flax_bert.py#L1176-L1177
=> this means that if we tie the word embeddings (`config.tie_word_embeddings=True)`, then we never use the `self.decoder` kernel/bias weights in the LM head, but rather the kernel/bias weights from our input word embeddings!

So in short, if you have the word embeddings tied (`config.tie_word_embeddings=True`), then the `self.decoder` weights can be ignored, since we use the input word embedding ones here.

Closing since this was discussed offline with @jxiw and the issue is now resolved 🤗

you said it is solved . so could you tell me how to convert a .msgpack (flax/jax format) model weights to .bin (huggingface format) or .pth (pytorch format)? thanks a lot !!!!!!!

You just need to instantiate a model like `FlaxBertModel` using `from_pretrained(...., from_pt = True)` or `BertModel.from_pretrained(....)` which will automatically convert the `msg` checkpoint and then save the model. 

Load a PyTorch model (`.bin`) into Flax:

```python
from transformers import FlaxBertModel

model = FlaxBertModel("hf-internal-testing/tiny-random-bert", from_pt=True)

# save the flax model
model.save_pretrained("./tiny-random-bert-flax")
```

Load the Flax model (`.msgpack`) into PyTorch:
```python
from transformers import BertModel

model = BertModel("./tiny-random-bert-flax", from_flax=True)
```