I think what's happening with ValueError: torch_xla device ID [1 2 3] not found in available JAX devices is:

torch_xla.devices() has 4 devices but jax.devices() only have 1. One possibility is that jax[cuda] was not installed so jax.devices() returnts one device and that is the CPU device.

Last time I tried to add the install and hit a different error. I am OK with disabling the test for CUDA until later.

I think what's happening with ValueError: torch_xla device ID [1 2 3] not found in available JAX devices is [...]

That is a great point!

This made me realize that if we land this PR as-is, then not only will call_jax not work on GPUs, it also won't work under multi-slice TPUs. That in turn means shard_as won't work, and scan won't work either, preventing training e.g. Llama 3.1 405B.

I think I'll split out the call_jax part of this PR into a separate one, and unfortunately that one can't be landed unless we fix the PJRT client sharing between PyTorch/XLA and JAX.

Draft for registering Jax contextual mesh in call_jax: #9043 (informational only)