fix tests on py38

Signed-off-by: Jinzhe Zeng <[email protected]>
deepmodeling · Sep 23, 2024 · a89f1eb · a89f1eb
1 parent b8a33b1
commit a89f1eb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/source/tests/consistent/test_activation.py b/source/tests/consistent/test_activation.py
@@ -2,7 +2,6 @@
 import sys
 import unittest
 
-import array_api_strict as xp
 import numpy as np
 
 from deepmd.common import (
@@ -69,6 +68,8 @@ def test_pt_consistent_with_ref(self):
         sys.version_info >= (3, 9), "array_api_strict doesn't support Python<=3.8"
     )
     def test_arary_api_strict(self):
+        import array_api_strict as xp
+
         xp.set_array_api_strict_flags(
             api_version=get_activation_fn_dp.array_api_version
         )

diff --git a/source/tests/consistent/test_type_embedding.py b/source/tests/consistent/test_type_embedding.py
@@ -36,6 +36,8 @@
         jnp,
     )
     from deepmd.jax.utils.type_embed import TypeEmbedNet as TypeEmbedNetJAX
+else:
+    TypeEmbedNetJAX = object
 
 
 @parameterized(
-Original file line number
+Diff line change
@@ Expand Up / @@ -36,6 +36,8 @@ @@
             jnp,
         )
         from deepmd.jax.utils.type_embed import TypeEmbedNet as TypeEmbedNetJAX
+    else:
+        TypeEmbedNetJAX = object
     @parameterized(
@@ Expand Down @@