Disable a pallas export compatibility test that fails on TPU v6e.

PiperOrigin-RevId: 673295487
jax-ml · Sep 11, 2024 · 49dd6ed · 49dd6ed
1 parent 808003b
commit 49dd6ed
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/pallas/export_back_compat_pallas_test.py b/tests/pallas/export_back_compat_pallas_test.py
@@ -62,6 +62,8 @@ def add_one(x_ref, o_ref):
 
   @jax.default_matmul_precision("bfloat16")
   def test_mosaic_matmul(self):
+    if jtu.is_device_tpu(6, "e"):
+      self.skipTest("TODO(apaszke): Test fails on TPU v6e")
     dtype = jnp.float32
     def func():
       # Build the inputs here, to reduce the size of the golden inputs.