[JAX] Bump jax lib (#12053)

Commit

4 years ago

[JAX] Bump jax lib (#12053) * fix_torch_device_generate_test * remove @ * bump up jax lib