Set env var for torax sim no compile test in test.

PiperOrigin-RevId: 686089551

torax/tests/sim_no_compile.py

@@ -14,6 +14,8 @@
 
 """Tests that TORAX can be run with compilation disabled."""
 
+import copy
+import os
 from typing import Optional, Sequence
 
 from absl.testing import absltest
@@ -28,6 +30,17 @@
 class SimTest(sim_test_case.SimTestCase):
   """No-compilation integration tests for torax.sim."""
 
+  def setUp(self):
+    self.original_flag = copy.deepcopy(os.environ)
+    compilation_flag = os.environ.get('TORAX_COMPILATION_ENABLED', None)
+    if compilation_flag is None or compilation_flag in ['1', 'True', 'true']:
+      os.environ['TORAX_COMPILATION_ENABLED'] = 'False'
+    super().setUp()
+
+  def tearDown(self):
+    os.environ = self.original_flag
+    super().tearDown()
+
   @parameterized.named_parameters(
       (
           'test_implicit_optimizer_no_compile',