can't load flax weights in PyTorch if flax model is saved with dtype `bfloat16`
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version:
- Platform: linux
- Python version: 3,8
- PyTorch version (GPU?):
- Tensorflow version (GPU?):
- Using GPU in script?:
- Using distributed or parallel set-up in script?:

### Who can help

- research_projects/r./run_clm_flax.py : @patrickvonplaten @patil-suraj


 -->

## Information

Model I am using (gpt2):

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

Script:

1. run conversion script
2. receive error

`path = "./romanian-gpt2_80000/ckpt-80000"config = AutoConfig.from_pretrained(path)
model = AutoModelForCausalLM.from_config(config)
load_flax_checkpoint_in_pytorch_model(model, path+"/flax_model.msgpack")
model.save_pretrained("./romanian-gpt2-large_80000"). `

Converting FLAX model to Pytorch gives the following error:
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.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
Convert flax to pytorch

<!-- A clear and concise description of what you would expect to happen. -->

Adding argument --dtype="bfloat16" to run_clm_flax.py   converts some of parameters to fp16, but it gives error when trying to convert the saved flax model to pytorch. A workaround is to convert all parameters of flax model to fp32 and convert then convert the flax model to pytorch.

`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)`
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.