Skip failing tests on TPU v6+

PiperOrigin-RevId: 741515935

Author: ayaka14732

tests/lax_numpy_reducers_test.py

@@ -905,8 +905,8 @@ def testCumulativeSumBool(self):
   @jtu.ignore_warning(category=NumpyComplexWarning)
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises mixed type promotion
   def testCumulativeProd(self, shape, axis, dtype, out_dtype, include_initial):
-    if jtu.is_device_tpu(6):
-      raise unittest.SkipTest("TODO(b/364258243): Test fails on TPU v6e")
+    if jtu.is_device_tpu_at_least(6):
+      raise unittest.SkipTest("TODO(b/364258243): Test fails on TPU v6+")
     rng = jtu.rand_some_zero(self.rng())
 
     # We currently "cheat" to ensure we have JAX arrays, not NumPy arrays as

tests/lax_scipy_test.py

@@ -339,8 +339,8 @@ def scipy_fun(z):
   )
   @jtu.ignore_warning(category=DeprecationWarning, message=".*scipy.special.lpmn.*")
   def testLpmn(self, l_max, shape, dtype):
-    if jtu.is_device_tpu(6, "e"):
-      self.skipTest("TODO(b/364258243): fails on TPU v6e")
+    if jtu.is_device_tpu_at_least(6):
+      self.skipTest("TODO(b/364258243): fails on TPU v6+")
     rng = jtu.rand_uniform(self.rng(), low=-0.2, high=0.9)
     args_maker = lambda: [rng(shape, dtype)]
 
@@ -461,8 +461,8 @@ def testSphHarmOrderOneDegreeOne(self):
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmForJitAndAgainstNumpy(self, l_max, num_z, dtype):
     """Tests against JIT compatibility and Numpy."""
-    if jtu.is_device_tpu(6, "e"):
-      self.skipTest("TODO(b/364258243): fails on TPU v6e")
+    if jtu.is_device_tpu_at_least(6):
+      self.skipTest("TODO(b/364258243): fails on TPU v6+")
     n_max = l_max
     shape = (num_z,)
     rng = jtu.rand_int(self.rng(), -l_max, l_max + 1)
@@ -508,8 +508,8 @@ def testSphHarmCornerCaseWithWrongNmax(self):
   )
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmY(self, l_max, num_z, dtype):
-    if jtu.is_device_tpu(6, "e"):
-      self.skipTest("TODO(b/364258243): fails on TPU v6e")
+    if jtu.is_device_tpu_at_least(6):
+      self.skipTest("TODO(b/364258243): fails on TPU v6+")
     n_max = l_max
     shape = (num_z,)
     rng = jtu.rand_int(self.rng(), -l_max, l_max + 1)