set xla flags before any imports

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Nov 8, 2024 · 050c20c · 050c20c
1 parent 31f9e87
commit 050c20c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/source/tests/__init__.py b/source/tests/__init__.py
@@ -1 +1,11 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import os
+
+# set XLA FLAGS before any jax import
+os.environ["XLA_FLAGS"] = " ".join(
+    (
+        "--xla_cpu_multi_thread_eigen=false",
+        "intra_op_parallelism_threads=1",
+        "inter_op_parallelism_threads=1",
+    )
+)
diff --git a/source/tests/jax/__init__.py b/source/tests/jax/__init__.py
@@ -1,10 +1 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import os
-
-os.environ["XLA_FLAGS"] = " ".join(
-    (
-        "--xla_cpu_multi_thread_eigen=false",
-        "intra_op_parallelism_threads=1",
-        "inter_op_parallelism_threads=1",
-    )
-)