From c1d92590b8591ed7ea60f249648ac94ffeaaa14d Mon Sep 17 00:00:00 2001
From: ben0i0d <ben0i0d@foxmail.com>
Date: Thu, 6 Feb 2025 13:17:33 +0800
Subject: [PATCH] add jax to pyai

---
 ScienceCompute/pyai/cpu/Dockerfile | 1 +
 ScienceCompute/pyai/gpu/Dockerfile | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/ScienceCompute/pyai/cpu/Dockerfile b/ScienceCompute/pyai/cpu/Dockerfile
index 154d454..fbbd4ee 100644
--- a/ScienceCompute/pyai/cpu/Dockerfile
+++ b/ScienceCompute/pyai/cpu/Dockerfile
@@ -2,5 +2,6 @@ FROM eoelab.org:1027/eoeair/jupyter:py-c
 
 # install torch
 RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu && \
+    pip install jax[cpu] && \
     pip cache purge
 
diff --git a/ScienceCompute/pyai/gpu/Dockerfile b/ScienceCompute/pyai/gpu/Dockerfile
index fb423dc..4192f2a 100644
--- a/ScienceCompute/pyai/gpu/Dockerfile
+++ b/ScienceCompute/pyai/gpu/Dockerfile
@@ -1,6 +1,6 @@
 FROM eoelab.org:1027/eoeair/jupyter:scipy-g
 
 # install torch transformers
-RUN pip install torch torchvision torchaudio && \
+RUN pip install torch torchvision torchaudio "jax[cuda12]" && \
     pip cache purge