Re-enable TPU tests now that TPU thread stack sizes have been increased.

Reverts 4d9ff5b

PiperOrigin-RevId: 840845684

Author: danielsuo

tests/multiprocess/pjit_test.py

@@ -527,9 +527,6 @@ def f(x):
       self.assertEqual(output(), "")
 
   def test_print_in_multihost_shard_map(self):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     devices = jax.devices()
     mesh = jax.sharding.Mesh(devices, ("i",))
     num_devices = jax.local_device_count()

tests/pallas/ops_test.py

@@ -2756,9 +2756,6 @@ class OpsInterpretTest(OpsTest):
   INTERPRET = True
 
   def test_debug_print(self):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     @functools.partial(
         self.pallas_call,
         out_shape=jax.ShapeDtypeStruct((2,), jnp.float32),

tests/pjit_test.py

@@ -4247,9 +4247,6 @@ def test_in_out_shardings_unconstrained_error(self):
               in_shardings=NamedSharding(mesh, P(P.UNCONSTRAINED, 'x')))
 
   def test_empty_io_callback_under_shard_map(self):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     mesh = jtu.create_mesh((4,), 'i')
 
     def empty_callback(x):

tests/python_callback_test.py

@@ -53,7 +53,6 @@
 )
 
 
-@unittest.skipIf(jtu.is_cloud_tpu(), "TODO: b/465504705")
 class PythonCallbackTest(jtu.JaxTestCase):
 
   def setUp(self):
@@ -670,7 +669,6 @@ def f(x):
     np.testing.assert_array_equal(x, result)
 
 
-@unittest.skipIf(jtu.is_cloud_tpu(), "TODO: b/465504705")
 class PureCallbackTest(jtu.JaxTestCase):
 
   def setUp(self):
@@ -1152,9 +1150,6 @@ def tearDown(self):
     dispatch.runtime_tokens.clear()
 
   def test_io_callback_can_mutate_state(self):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     x = 0
     def cb():
       nonlocal x
@@ -1171,9 +1166,6 @@ def f():
     self.assertEqual(x, 2)
 
   def test_io_callback_can_be_batched_if_unordered(self):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     _mut = 0
     def cb(x):
       nonlocal _mut
@@ -1282,9 +1274,6 @@ def f(x, y):
   def test_can_use_io_callback_in_pjit(
       self, *, ordered: bool, with_sharding: bool
   ):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     devices = jax.devices()
     mesh = jax.sharding.Mesh(np.array(devices), ['dev'])
 
@@ -1345,9 +1334,6 @@ def f(x):
   @jtu.ignore_warning(message='.*Please use `jax.jit` instead.*',
                       category=DeprecationWarning)
   def test_sequence_pjit_io_callback_ordered(self):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     if jtu.is_device_tpu(7, 'x'):
       self.skipTest('TODO(b/453664256): Failing on TPU 7x.')
 
@@ -1409,8 +1395,6 @@ def f_base(i, x):
          single_device=True)
   )
   def test_can_shard_io_callback_manually(self, single_device: bool):
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
 
     devices = jax.devices()
     if single_device:
@@ -1445,9 +1429,6 @@ def f(shard_ids, x):
 
   def test_batching_with_side_effects(self):
     # https://github.com/jax-ml/jax/issues/20628#issuecomment-2050800195
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     x_lst = []
     def append_x(x):
       nonlocal x_lst
@@ -1464,9 +1445,6 @@ def f(x):
 
   def test_batching_with_side_effects_while_loop(self):
     # https://github.com/jax-ml/jax/issues/20628#issuecomment-2050921219
-    if jtu.is_cloud_tpu():
-      self.skipTest("TODO: b/465504705")
-
     x_lst = []
     def append_x(x):
       nonlocal x_lst