Address comments

Author: tengyifei

test/test_assume_pure_spmd.py

@@ -12,10 +12,12 @@
 
 class AssumePureSpmdTest(unittest.TestCase):
 
-  def setUp(self):
+  @classmethod
+  def setUpClass(cls):
     # Activate SPMD
     xr.use_spmd()
 
+  def setUp(self):
     # Set up a simple SPMD mesh for these tests.
     self.spmd_mesh = get_1d_mesh(axis_name="model")
     set_global_mesh(self.spmd_mesh)

test/test_jax_interop_spmd.py

@@ -11,11 +11,14 @@
 
 class TestJaxInteropSpmd(unittest.TestCase):
 
-  def setUp(self):
-    xb._JAX_TO_XLA_COMPUTATION_CACHE.clear()
+  @classmethod
+  def setUpClass(cls):
     # Activate SPMD
     xr.use_spmd()
 
+  def setUp(self):
+    # Clear cached HLO between test cases.
+    xb._JAX_TO_XLA_COMPUTATION_CACHE.clear()
     # Set up a simple SPMD mesh for these tests.
     self.spmd_mesh = get_1d_mesh(axis_name="model")
     set_global_mesh(self.spmd_mesh)