I don;t understand why the tests fail randomly agter I add some tests...

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Nov 8, 2024 · befc0c7 · befc0c7
1 parent 19f798c
commit befc0c7
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/source/tests/jax/jax2tf/__init__.py b/source/tests/jax/jax2tf/__init__.py
@@ -1,6 +1,11 @@
 # 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)
+from ...common import (
+    DP_TEST_TF2_ONLY,
+)
+
+if DP_TEST_TF2_ONLY:
+    # limit the number of threads
+    tf.config.threading.set_inter_op_parallelism_threads(1)
+    tf.config.threading.set_intra_op_parallelism_threads(1)
diff --git a/source/tests/jax/jax2tf/test_nlist.py b/source/tests/jax/jax2tf/test_nlist.py
@@ -2,13 +2,14 @@
 
 import pytest
 import tensorflow as tf
-import tensorflow.experimental.numpy as tnp
 
 from ...utils import (
     DP_TEST_TF2_ONLY,
 )
 
 if DP_TEST_TF2_ONLY:
+    import tensorflow.experimental.numpy as tnp
+
     from deepmd.jax.jax2tf.nlist import (
         build_neighbor_list,
         extend_coord_with_ghosts,

diff --git a/source/tests/jax/jax2tf/test_region.py b/source/tests/jax/jax2tf/test_region.py
@@ -3,7 +3,6 @@
 
 import pytest
 import tensorflow as tf
-import tensorflow.experimental.numpy as tnp
 
 from ...seed import (
     GLOBAL_SEED,
@@ -13,6 +12,8 @@
 )
 
 if DP_TEST_TF2_ONLY:
+    import tensorflow.experimental.numpy as tnp
+
     from deepmd.jax.jax2tf.region import (
         inter2phys,
         to_face_distance,