Fix Pathwaysutils tests to run with the JAX cpu platform.

PiperOrigin-RevId: 726536737

Author: lukebaumann

pathwaysutils/test/persistence_test.py

@@ -14,7 +14,11 @@ class PersistenceTest(absltest.TestCase):
   name = "name"
   dtype = np.dtype(np.int32)
   shape = [8, 4]
-  timeout = datetime.timedelta(seconds=30)
+  timeout = datetime.timedelta(seconds=3)
+
+  def setUp(self):
+    jax.config.update("jax_platforms", "cpu")
+    super().setUp()
 
   def test_get_read_request(self):
     devices = jax.devices()

pathwaysutils/test/plugin_executable_test.py

@@ -13,6 +13,10 @@
 
 class PluginExecutableTest(absltest.TestCase):
 
+  def setUp(self):
+    jax.config.update("jax_platforms", "cpu")
+    super().setUp()
+
   def test_bad_json_program(self):
     with self.assertRaisesRegex(XlaRuntimeError, "INVALID_ARGUMENT"):
       PluginExecutable('{"printTextRequest":{"badParamName":"foo"}}')