jax is only available in Python 3.10

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

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -133,7 +133,7 @@ cu12 = [
     "nvidia-cuda-nvcc-cu12",
 ]
 jax = [
-    "jax>=0.4.33",
+    'jax>=0.4.33;python_version>="3.10"',
 ]
 
 [tool.deepmd_build_backend.scripts]

diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py
@@ -37,8 +37,8 @@
 INSTALLED_PT = Backend.get_backend("pytorch")().is_available()
 INSTALLED_JAX = Backend.get_backend("jax")().is_available()
 
-if os.environ.get("CI") and not (INSTALLED_TF and INSTALLED_PT and INSTALLED_JAX):
-    raise ImportError("TensorFlow or PyTorch or JAX should be tested in the CI")
+if os.environ.get("CI") and not (INSTALLED_TF and INSTALLED_PT):
+    raise ImportError("TensorFlow or PyTorch should be tested in the CI")
 
 
 if INSTALLED_TF: