Add configuration options jax_cuda_visible_devices and jax_rocm_visible_devices.
Add configuration options jax_cuda_visible_devices and jax_rocm_visible_devices.

Example usage:
jax.config.update("jax_cuda_visible_devices", "1,3")

google/jax,google/jax,google/jax,google/jax