From f0bc8b814f697fe220f871640ef57d6390f7c347 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Mon, 21 Oct 2024 18:37:45 -0400
Subject: [PATCH] add dependencies

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 pyproject.toml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 4a99cfaf8b..26f8e5ae6e 100644
--- 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]