Get core dump after import FlaxRobertaModel
## Environment info
I tried to run "transformers-cli env" but I get also this error message:
2021-07-03 05:37:33.144628: F external/org_tensorflow/tensorflow/core/tpu/tpu_executor_init_fns.inc:110] TpuTransferManager_ReadDynamicShapes not available in this library.
Aborted (core dumped)

- `transformers` version: 4.9.0.dev0
- Platform: TPU vm from the Jax/Flax event
- Python version: 3.8
- PyTorch version (GPU?): not installed or pytorch 1.8.1, the same issue 
- Tensorflow version (GPU?): not installed or tensorflow 2.4.1, the same issue 
- Using GPU in script?:
- Using distributed or parallel set-up in script?: no

### Who can help

@patrickvonplaten 

## Information

I followed the instruction https://github.com/huggingface/transformers/tree/master/examples/research_projects/jax-projects#how-to-install-relevant-libraries in a TPU vm for the JAX/Flax event. But after the installation of transformers and the datasets, I get following error message when I import FlaxRobertaModel:
```
$ python
Python 3.8.10 (default, Jun  4 2021, 15:09:15) 
[GCC 7.5.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from transformers import FlaxRobertaModel
2021-07-03 05:55:17.729321: F external/org_tensorflow/tensorflow/core/tpu/tpu_executor_init_fns.inc:110] TpuTransferManager_ReadDynamicShapes not available in this library.
Aborted (core dumped)

```
or run the transformers-cli:
```
$ transformers-cli env
WARNING:tensorflow:From /home/cahya/transformers/src/transformers/commands/env.py:50: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.list_physical_devices('GPU')` instead.
2021-07-03 05:56:23.495538: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE4.1 SSE4.2 AVX AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-07-03 05:56:23.555453: F external/org_tensorflow/tensorflow/core/tpu/tpu_executor_init_fns.inc:110] TpuTransferManager_ReadDynamicShapes not available in this library.
Aborted (core dumped)

```
The problem arises when using:
* [ x] the official example scripts: (see above)

The tasks I am working on is:
* [ ] JAX/Flax project
* [ ] 

## To reproduce
see above


## Expected behavior
There should be no core dump
It seems that the jax version 0.2.16 installed automatically during transformers installation make this issue. I can fix it after reinstalling the jax using following command:
`pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html`

