avoid downgrading jax 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.