limit threads during tests

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

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/source/tests/jax/__init__.py b/source/tests/jax/__init__.py
@@ -1 +1,10 @@
 # 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",
+    )
+)
diff --git a/source/tests/jax/jax2tf/__init__.py b/source/tests/jax/jax2tf/__init__.py
@@ -1 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import tensorflow as tf
+
+# limit the number of threads
+tf.config.threading.set_inter_op_parallelism_threads(1)
+tf.config.threading.set_intra_op_parallelism_threads(1)