[Flax] Error converting model to PyTorch from Flax
Author: w11wo
Pull request: no

Hi, I followed the causal language modeling in Flax tutorial notebook provided [here](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/causal_language_modeling_flax.ipynb) in Colab. And at the end of the training, I'd like to get a working PyTorch model from the JAX/Flax weights, hence I did this:

```python
from transformers import GPT2LMHeadModel

mdl_path = "w11wo/sundanese-gpt2-base"

pt_model = GPT2LMHeadModel.from_pretrained(mdl_path, from_flax=True)
pt_model.save_pretrained(mdl_path)
```

But during the conversion, it raised this error

```python
/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
  pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-44-08a8ff6e575c> in <module>()
      7 tokenizer.save_pretrained(mdl_path)
      8 
----> 9 pt_model = GPT2LMHeadModel.from_pretrained(mdl_path, from_flax=True)
     10 pt_model.save_pretrained(mdl_path)
     11 

2 frames
/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py in load_flax_weights_in_pytorch_model(pt_model, flax_state)
    199                 # add weight to pytorch dict
    200                 flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor
--> 201                 pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
    202                 # remove from missing keys
    203                 missing_keys.remove(flax_key)

TypeError: can't convert np.ndarray of type bfloat16. The only supported types are: float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool.
```

I think this issue occurred because the model I instantiated used `bfloat16` -- just as the tutorial showed. Specifically this block

```python
from transformers import FlaxAutoModelForCausalLM

model = FlaxAutoModelForCausalLM.from_config(config, seed=training_seed, dtype=jnp.dtype("bfloat16"))
```

I'd like to know if there's a workaround to this problem. Thanks!

Comments:
Seeing how we are able to convert `float32` weights to `bfloat16` as shown [here](https://github.com/huggingface/transformers/issues/12534#issue-937864452), I tried to the reverse, i.e. converting `bfloat16` to `float32`. I executed the following lines:

```python
from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel

pretrained = "w11wo/sundanese-gpt2-base"
tmp_path = "sundanese-gpt-base"

model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)
tokenizer = GPT2Tokenizer.from_pretrained(pretrained)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)
pt_model.save_pretrained(tmp_path)
```

And the model converted returned the following messages,

```python
All Flax model weights were used when initializing GPT2LMHeadModel.

Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.5.attn.masked_bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.0.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.8.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.9.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.10.attn.masked_bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.11.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.2.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.11.attn.masked_bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

I suspect that there are conversion errors, but am clueless as to the reason behind the error.

Further, I tried using the converted PyTorch model for text-generation via the following pipeline:

```python
from transformers import pipeline

nlp = pipeline(
    "text-generation",
    model=tmp_path,
    tokenizer=tokenizer
)

nlp("Nami abdi Budi")
>> Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
[{'generated_text': "Testipit ayeuna ayeuna409 pillowsKrKr bunderan Platform sipil poé Summary Sakitu Enak maké bunderan sipilthough papatong hubungiashmina Stock protésnyieun248duta terhadap Mm �issinginjeumdoi ' 'Dealer Studio gunarior floridabodas ' békénAndroid Holland majalah dot mangaruhanumpahinting"}]
```
Unsurprisingly, the output is very much jibberish, despite the model trained down to about 3.66 validation loss. I think the conversion is somehow incorrect, hence the incorrect porting of weights, and thus the jibberish output.

I will take a look later today!

Also cc @patil-suraj 

Thanks @patrickvonplaten! Really appreciate what you and your 🤗 team are doing! 

One more thing, I trained RoBERTa using a different script for the Flax community week; hub repo [here](https://huggingface.co/flax-community/indonesian-roberta-base). It seems that converting the model to PyTorch raises a possibly related error:

```python
from transformers import RobertaForMaskedLM
model = RobertaForMaskedLM.from_pretrained("flax-community/indonesian-roberta-base", from_flax=True)
```
```
/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
  pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
All Flax model weights were used when initializing RobertaForMaskedLM.

Some weights of RobertaForMaskedLM were not initialized from the Flax model and are newly initialized: ['lm_head.decoder.bias', 'roberta.embeddings.position_ids', 'lm_head.decoder.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

> Seeing how we are able to convert `float32` weights to `bfloat16` as shown [here](https://github.com/huggingface/transformers/issues/12534#issue-937864452), I tried to the reverse, i.e. converting `bfloat16` to `float32`. I executed the following lines:
> 
> ```python
> from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
> 
> pretrained = "w11wo/sundanese-gpt2-base"
> tmp_path = "sundanese-gpt-base"
> 
> model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)
> tokenizer = GPT2Tokenizer.from_pretrained(pretrained)
> 
> def to_f32(t):
>     return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
> 
> model.params = to_f32(model.params)
> model.save_pretrained(tmp_path)
> 
> pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)
> pt_model.save_pretrained(tmp_path)
> ```
> 
> And the model converted returned the following messages,
> 
> ```python
> All Flax model weights were used when initializing GPT2LMHeadModel.
> 
> Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.5.attn.masked_bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.0.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.8.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.9.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.10.attn.masked_bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.11.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.2.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.11.attn.masked_bias']
> You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
> ```
> 
> I suspect that there are conversion errors, but am clueless as to the reason behind the error.
> 
> Further, I tried using the converted PyTorch model for text-generation via the following pipeline:
> 
> ```python
> from transformers import pipeline
> 
> nlp = pipeline(
>     "text-generation",
>     model=tmp_path,
>     tokenizer=tokenizer
> )
> 
> nlp("Nami abdi Budi")
> >> Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
> [{'generated_text': "Testipit ayeuna ayeuna409 pillowsKrKr bunderan Platform sipil poé Summary Sakitu Enak maké bunderan sipilthough papatong hubungiashmina Stock protésnyieun248duta terhadap Mm �issinginjeumdoi ' 'Dealer Studio gunarior floridabodas ' békénAndroid Holland majalah dot mangaruhanumpahinting"}]
> ```
> 
> Unsurprisingly, the output is very much jibberish, despite the model trained down to about 3.66 validation loss. I think the conversion is somehow incorrect, hence the incorrect porting of weights, and thus the jibberish output.

I think this is actually the correct way how you should convert bfloat16 to float32. Having followed this way of conversion can you verify that the PT and Flax model give similar outputs (ideally on CPU since TPU uses approximations)

> One more thing, I trained RoBERTa using a different script for the Flax community week; hub repo [here](https://huggingface.co/flax-community/indonesian-roberta-base). It seems that converting the model to PyTorch raises a possibly related error:
> 
> ```python
> from transformers import RobertaForMaskedLM
> model = RobertaForMaskedLM.from_pretrained("flax-community/indonesian-roberta-base", from_flax=True)
> ```
> 
> ```
> /usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
>   pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
> All Flax model weights were used when initializing RobertaForMaskedLM.
> 
> Some weights of RobertaForMaskedLM were not initialized from the Flax model and are newly initialized: ['lm_head.decoder.bias', 'roberta.embeddings.position_ids', 'lm_head.decoder.weight']
> You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
> ```

This is expected can you check with https://github.com/huggingface/transformers/issues/12554

```python
from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
import torch
import numpy as np
import jax
import jax.numpy as jnp

pretrained = "w11wo/sundanese-gpt2-base"
tmp_path = "./"

model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

model_pt = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)

input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
input_ids_pt = torch.tensor(input_ids)

logits_pt = model_pt(input_ids_pt).logits
print(logits_pt)
logits_fx = model(input_ids).logits
print(logits_fx)
```

=> here you can see that you are correctly converting the weighst. The two models give the same results when doing a forward pass. This probably means that the training didn't work very well. Also not that vanilla generation often doesn't work very well. One should use `do_sample=True`

> > One more thing, I trained RoBERTa using a different script for the Flax community week; hub repo [here](https://huggingface.co/flax-community/indonesian-roberta-base). It seems that converting the model to PyTorch raises a possibly related error:
> > ```python
> > from transformers import RobertaForMaskedLM
> > model = RobertaForMaskedLM.from_pretrained("flax-community/indonesian-roberta-base", from_flax=True)
> > ```
> > 
> > 
> >     
> >       
> >     
> > 
> >       
> >     
> > 
> >     
> >   
> > ```
> > /usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
> >   pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
> > All Flax model weights were used when initializing RobertaForMaskedLM.
> > 
> > Some weights of RobertaForMaskedLM were not initialized from the Flax model and are newly initialized: ['lm_head.decoder.bias', 'roberta.embeddings.position_ids', 'lm_head.decoder.weight']
> > You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
> > ```
> 
> This is expected can you check with #12554

Checked, and it does return the same results. Thanks!

> ```python
> from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
> import torch
> import numpy as np
> import jax
> import jax.numpy as jnp
> 
> pretrained = "w11wo/sundanese-gpt2-base"
> tmp_path = "./"
> 
> model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)
> 
> def to_f32(t):
>     return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
> 
> model.params = to_f32(model.params)
> model.save_pretrained(tmp_path)
> 
> pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)
> 
> input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
> input_ids_pt = torch.tensor(input_ids)
> 
> logits_pt = model_pt(input_ids_pt).logits
> print(logits_pt)
> logits_fx = model(input_ids).logits
> print(logits_fx)
> ```

This doesn't seem to return the same results on my end.
```
All Flax model weights were used when initializing GPT2LMHeadModel.

Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.9.attn.masked_bias', 'transformer.h.8.attn.bias', 'transformer.h.11.attn.masked_bias', 'transformer.h.2.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.5.attn.masked_bias', 'transformer.h.11.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.0.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.masked_bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
tensor([[[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
         [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
         [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
         ...,
         [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
         [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
         [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]],

        [[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
         [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
         [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
         ...,
         [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
         [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
         [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]]],
       grad_fn=<AddBackward0>)
[[[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
    0.63595504]
  [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
    0.46392882]
  [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
    0.21045035]
  ...
  [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
   -0.27594602]
  [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
   -0.2215142 ]
  [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
   -0.2000624 ]]

 [[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
    0.63595504]
  [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
    0.46392882]
  [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
    0.21045035]
  ...
  [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
   -0.27594602]
  [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
   -0.2215142 ]
  [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
   -0.2000624 ]]]
```

> > ```python
> > from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
> > import torch
> > import numpy as np
> > import jax
> > import jax.numpy as jnp
> > 
> > pretrained = "w11wo/sundanese-gpt2-base"
> > tmp_path = "./"
> > 
> > model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)
> > 
> > def to_f32(t):
> >     return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
> > 
> > model.params = to_f32(model.params)
> > model.save_pretrained(tmp_path)
> > 
> > pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)
> > 
> > input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
> > input_ids_pt = torch.tensor(input_ids)
> > 
> > logits_pt = model_pt(input_ids_pt).logits
> > print(logits_pt)
> > logits_fx = model(input_ids).logits
> > print(logits_fx)
> > ```
> 
> This doesn't seem to return the same results on my end.
> 
> ```
> All Flax model weights were used when initializing GPT2LMHeadModel.
> 
> Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.9.attn.masked_bias', 'transformer.h.8.attn.bias', 'transformer.h.11.attn.masked_bias', 'transformer.h.2.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.5.attn.masked_bias', 'transformer.h.11.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.0.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.masked_bias']
> You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
> tensor([[[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
>          [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
>          [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
>          ...,
>          [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
>          [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
>          [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]],
> 
>         [[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
>          [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
>          [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
>          ...,
>          [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
>          [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
>          [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]]],
>        grad_fn=<AddBackward0>)
> [[[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
>     0.63595504]
>   [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
>     0.46392882]
>   [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
>     0.21045035]
>   ...
>   [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
>    -0.27594602]
>   [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
>    -0.2215142 ]
>   [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
>    -0.2000624 ]]
> 
>  [[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
>     0.63595504]
>   [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
>     0.46392882]
>   [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
>     0.21045035]
>   ...
>   [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
>    -0.27594602]
>   [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
>    -0.2215142 ]
>   [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
>    -0.2000624 ]]]
> ```

weird it gives the same results for me

Here's my environment @patrickvonplaten, in case it'll help

- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.13
- JaxLib version: 0.1.66
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

That's my output:

```bash
Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.9.attn.masked_bias', 'transformer.h.3.attn.bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.2.attn.bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.11.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.10.attn.masked_bias', 'lm_head.weight', 'transformer.h.11.attn.masked_bias', 'transformer.h.0.attn.bias', 'transformer.h.10.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.4.attn.bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.7.attn.bias', 'transformer.h.5.attn.masked_bias', 'transformer.h.9.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.6.attn.bias', 'transformer.h.8.attn.bias', 'transformer.h.1.attn.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
tensor([[[-0.1830, -0.3880,  0.5160,  ..., -0.0358,  0.3377,  0.6360],
         [-0.2227, -0.4381,  0.5417,  ..., -0.0942,  0.1144,  0.4639],
         [-0.2117, -0.4112,  0.5285,  ..., -0.1005, -0.0111,  0.2105],
         ...,
         [-0.1392, -0.1115,  0.2263,  ..., -0.0606,  0.6304, -0.2759],
         [-0.2090, -0.1306,  0.2280,  ..., -0.0729,  0.6694, -0.2215],
         [-0.1788, -0.1188,  0.2094,  ..., -0.0635,  0.6529, -0.2001]],

        [[-0.1830, -0.3880,  0.5160,  ..., -0.0358,  0.3377,  0.6360],
         [-0.2227, -0.4381,  0.5417,  ..., -0.0942,  0.1144,  0.4639],
         [-0.2117, -0.4112,  0.5285,  ..., -0.1005, -0.0111,  0.2105],
         ...,
         [-0.1392, -0.1115,  0.2263,  ..., -0.0606,  0.6304, -0.2759],
         [-0.2090, -0.1306,  0.2280,  ..., -0.0729,  0.6694, -0.2215],
         [-0.1788, -0.1188,  0.2094,  ..., -0.0635,  0.6529, -0.2001]]],
       grad_fn=<UnsafeViewBackward>)
2021-07-07 14:27:19.167439: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
[[[-0.18300146 -0.38216865  0.51648384 ... -0.03733653  0.33876732
    0.63263166]
  [-0.217145   -0.4353617   0.5423262  ... -0.09400246  0.11277512
    0.46030286]
  [-0.20957293 -0.4070356   0.52822715 ... -0.10254323 -0.01395297
    0.20638129]
  ...
  [-0.13773967 -0.10961649  0.22636533 ... -0.065424    0.6344244
   -0.27795315]
  [-0.20452675 -0.1309557   0.23095486 ... -0.07456343  0.66748977
   -0.21973252]
  [-0.17408411 -0.11795001  0.21041611 ... -0.06386398  0.6547011
   -0.1990372 ]]

 [[-0.18300146 -0.38216865  0.51648384 ... -0.03733653  0.33876732
    0.63263166]
  [-0.217145   -0.4353617   0.5423262  ... -0.09400246  0.11277512
    0.46030286]
  [-0.20957293 -0.4070356   0.52822715 ... -0.10254323 -0.01395297
    0.20638129]
  ...
  [-0.13773967 -0.10961649  0.22636533 ... -0.065424    0.6344244
   -0.27795315]
  [-0.20452675 -0.1309557   0.23095486 ... -0.07456343  0.66748977
   -0.21973252]
  [-0.17408411 -0.11795001  0.21041611 ... -0.06386398  0.6547011
   -0.1990372 ]]]

```

My env is:

```
- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

```

Can you update jax, flax and jaxlib and also check on TPU?

> Can you update jax, flax and jaxlib and also check on TPU?

Got it, I'll do it in a bit. Thanks for the help 👍 

> Can you update jax, flax and jaxlib and also check on TPU?

Okay, I checked several setups. I upgraded both jax and jaxlib in Google Colab's **CPU** runtime.

I ran this line to upgrade
```
!pip install -Uq jax jaxlib git+https://github.com/huggingface/transformers.git tokenizers datasets flax git+https://github.com/deepmind/optax.git
```

- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

Then this

> ```python
> from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
> import torch
> import numpy as np
> import jax
> import jax.numpy as jnp
> 
> pretrained = "w11wo/sundanese-gpt2-base"
> tmp_path = "./"
> 
> model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)
> 
> def to_f32(t):
>     return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)
> 
> model.params = to_f32(model.params)
> model.save_pretrained(tmp_path)
> 
> model_pt = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)
> 
> input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
> input_ids_pt = torch.tensor(input_ids)
> 
> logits_pt = model_pt(input_ids_pt).logits
> print(logits_pt)
> logits_fx = model(input_ids).logits
> print(logits_fx)
> ```

It worked as intended. It returned the correct tensors.

---

Finally, I checked the **TPU** setup in Colab. Ran the same upgrade command, plus this block.

```python
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
jax.device_count()
```
```
8
```

Then I ran the same test as above. ~~Unfortunately, it returns a different error message (didn't appear in **CPU** runtime):~~
EDIT: Both the PyTorch and JAX models returned the same tensors after restarting the TPU runtime after library upgrades and JAX Colab TPU setup.

Here's my final environment on Colab's TPU

- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

@patrickvonplaten I think I'll close the issue now. Upgrading both jax and jaxlib solved the problem. Many thanks!