jax

268c6842 - [JAX] In McJAX, support JIT compilation/execution on only the devices attached to a single host, with other hosts not participating.