Merge pull request #28935 from jburnim:jburnim_interpret_test

Google-ML-Automation · Google-ML-Automation · commit 5dba0cfee271 · 2025-05-22T08:23:42.000-07:00
PiperOrigin-RevId: 761974061
diff --git a/tests/pallas/tpu_pallas_interpret_distributed_test.py b/tests/pallas/tpu_pallas_interpret_distributed_test.py
@@ -38,6 +38,7 @@
 P = jax.sharding.PartitionSpec
 
 
+@jtu.thread_unsafe_test_class()
 class InterpretDistributedTest(jtu.JaxTestCase):
   def setUp(self):
     super().setUp()
diff --git a/tests/pallas/tpu_pallas_interpret_test.py b/tests/pallas/tpu_pallas_interpret_test.py
@@ -82,6 +82,7 @@ def grid_points(self):
     return self._grid_points
 
 
+@jtu.thread_unsafe_test_class()
 class InterpretTest(jtu.JaxTestCase):
 
   def setUp(self):
