Skip activation offload tests on GPU and CPU until they are supported

PiperOrigin-RevId: 604698627
jax-ml · Feb 6, 2024 · 299b983 · 299b983
1 parent 87d1670
commit 299b983
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/memories_test.py b/tests/memories_test.py
@@ -1099,6 +1099,11 @@ def f_bwd(res, tx):
 
 class ActivationOffloadingTest(jtu.JaxTestCase):
 
+  def setUp(self):
+    if not jtu.test_device_matches(["tpu"]):
+      self.skipTest("Memories do not work on CPU and GPU backends yet.")
+    super().setUp()
+
   def test_remat_jaxpr_offloadable(self):
     mesh = jtu.create_global_mesh((2,), ("x",))
     inp = jax.device_put(np.arange(16.), NamedSharding(mesh, P("x")))