[torchax][distributed] Use local device for default jax device (#8667)

Co-authored-by: Siyuan Liu <[email protected]>

Author: lsy323

torchax/torchax/tensor.py

@@ -304,7 +304,7 @@ def get_as_jax_device(self, device: Any):
 
       if device == 'cpu':
         return jax.devices('cpu')[0]
-      return jax.devices()[0]
+      return jax.local_devices()[0]
 
 
     def load_ops(self):