Transformers not working with the new Jax 0.4.14 due to API deprecation
### System Info

- `transformers` version: 4.31.0
- Platform: Linux-5.15.90.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.0+cu117 (True)
- Tensorflow version (GPU?): 2.11.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.7.0 (gpu)
- Jax version: 0.4.14
- JaxLib version: 0.4.14
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @younesbelkada @sanchit-gandhi

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

For example, loading a Bart model from Flax using msgpack_restore (modeling_flax_bart.py) raises `AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'`



### Expected behavior

Load the state dict correctly
Hi @SystemPanic 

`transformers` currently only support `jax<=0.4.13` and `jaxlib<=0.4.13`. You can see that in

https://github.com/huggingface/transformers/blob/400e76ef11d94a12c255fe1a598966e1d6021511/setup.py#L127-L128