xla

Add jax_device context manager to control the device target

#9382

Merged