add dependencies

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Oct 21, 2024 · f0bc8b8 · f0bc8b8
1 parent 5024f70
commit f0bc8b8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -138,6 +138,8 @@ cu12 = [
 jax = [
     'jax>=0.4.33;python_version>="3.10"',
     'flax>=0.8.0;python_version>="3.10"',
+    'orbax-checkpoint;python_version>="3.10"',
+    'jax-ai-stack;python_version>="3.10"',
 ]
 
 [tool.deepmd_build_backend.scripts]