By default, to("jax") should go to TPU (#9468)

Author: zzzwen

torchax/test/test_interop.py

@@ -132,6 +132,16 @@ def forward(self, x):
   def test_to_jax_device(self):
     a = torch.ones(3, 3)
 
+    if is_tpu_available():
+      # by default if tpu is available, to jax will be to tpu
+      e = a.to("jax")
+      self.assertEqual(e.jax_device.platform, "tpu")
+      self.assertEqual(e.device.type, "jax")
+    else:
+      e = a.to("jax")
+      self.assertEqual(e.jax_device.platform, "cpu")
+      self.assertEqual(e.device.type, "jax")
+
     with jax_device("cpu"):
       # move torch.tensor to torchax.tensor CPU
       b = a.to("jax")

torchax/torchax/tensor.py

@@ -330,7 +330,7 @@ def __init__(self, configuration=None):
     self._prng_key = mutable_array(
         jax.random.key(torch.initial_seed() % (1 << 63)))
     self.autocast_dtype = None
-    self._target_device = "cpu"
+    self._target_device = jax.local_devices()[0].platform
 
   @property
   def target_device(self):