update jax version and re-enable some tests
# What does this PR do?

Update the `JAX` version in `setup.py` to not install `jax==0.3.2`. And also re-enable some skipped tests.
_The documentation is not available anymore as the PR was closed or merged._