avoid downgrading jax version
https://github.com/huggingface/datasets/blob/main/setup.py#L236
Currently has the highest version at 0.3.25; Not sure if there is any reason for this, other than that was the tested version?
### Feature request

Whenever I `pip install datasets[jax]` it downgrades jax to version 0.3.25. I seem to be able to install this library first then upgrade jax back to version 0.4.13.

### Motivation

It would be nice to not overwrite currently installed version of jax if possible.

### Your contribution

I would be willing to beta test. Or maybe write some code if I could get pointed in the right direction, I'm not super familiar with this codebase.