Merge pull request #21108 from justinjfu/skip_pallas_test_64

Skip float64 test_nextafter on TPU.
google · May 9, 2024 · eb0b1b0 · eb0b1b0
2 parents 671fb12 + 080b3f1
commit eb0b1b0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/pallas/pallas_test.py b/tests/pallas/pallas_test.py
@@ -1512,6 +1512,8 @@ def kernel(x_ref, y_ref, o_ref):
 
   @parameterized.parameters("float32", "float64")
   def test_nextafter(self, dtype):
+    if jtu.test_device_matches(["tpu"]) and dtype == "float64":
+        self.skipTest("float64 disabled on TPU.")
     @functools.partial(
         self.pallas_call, out_shape=jax.ShapeDtypeStruct((4,), dtype), grid=1
     )