build(jax): split cpu and gpu extras

Add explicit JAX CPU and CUDA extras so CPU workflows avoid pulling GPU runtime dependencies while CUDA jobs request jax[cuda12] through the dedicated GPU extra.

Authored by OpenClaw (model: gpt-5.5)

Author: njzjz-bot

.github/workflows/test_cc.yml

@@ -44,9 +44,9 @@ jobs:
       - run: python -m pip install uv
       - name: Install Python dependencies
         run: |
-          source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --group pin_jax --torch-backend cpu
+          source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --group pin_jax_cpu --torch-backend cpu
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-          source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax] mpi4py mpich
+          source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax-cpu] mpi4py mpich
       - name: Convert models
         run: source/tests/infer/convert-models.sh
       # https://github.com/actions/runner-images/issues/9491

.github/workflows/test_cuda.yml

@@ -43,12 +43,12 @@ jobs:
           && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3
         if: false # skip as we use nvidia image
       - run: python -m pip install -U uv
-      - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax "jax[cuda12]"
+      - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax_gpu
       - run: |
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
           pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn "paddlepaddle-gpu==3.4.0.dev20260310"
-          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
+          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax-gpu] mpi4py --reinstall-package deepmd-kit
           # See https://github.com/jax-ml/jax/issues/29042
           source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
         env:

.github/workflows/test_python.yml

@@ -31,7 +31,7 @@ jobs:
           source/install/uv_with_retry.sh pip install --system openmpi --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
-          source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py --group pin_jax
+          source/install/uv_with_retry.sh pip install --system -e .[test,jax-cpu] mpi4py --group pin_jax_cpu
           source/install/uv_with_retry.sh pip install --system --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn paddlepaddle==3.4.0.dev20260310
         env:
           # Please note that uv has some issues with finding

pyproject.toml

@@ -132,10 +132,35 @@ cu12 = [
     "nvidia-cuda-nvcc-cu12",
 ]
 jax = [
+    # Backward-compatible alias for CPU JAX.
     # below is a funny workaround for
     # https://github.com/astral-sh/uv/issues/8601
-    'jax>=0.4.33;python_version>="3.10"',
-    'jax>=0.4.33;python_version>="3.10"',
+    'jax[cpu]>=0.4.33;python_version>="3.10"',
+    'jax[cpu]>=0.4.33;python_version>="3.10"',
+    'flax>=0.10.0;python_version>="3.10"',
+    'flax>=0.10.0;python_version>="3.10"',
+    'orbax-checkpoint;python_version>="3.10"',
+    'orbax-checkpoint;python_version>="3.10"',
+    # The pinning of ml_dtypes may conflict with TF
+    # 'jax-ai-stack;python_version>="3.10"',
+]
+jax-cpu = [
+    # below is a funny workaround for
+    # https://github.com/astral-sh/uv/issues/8601
+    'jax[cpu]>=0.4.33;python_version>="3.10"',
+    'jax[cpu]>=0.4.33;python_version>="3.10"',
+    'flax>=0.10.0;python_version>="3.10"',
+    'flax>=0.10.0;python_version>="3.10"',
+    'orbax-checkpoint;python_version>="3.10"',
+    'orbax-checkpoint;python_version>="3.10"',
+    # The pinning of ml_dtypes may conflict with TF
+    # 'jax-ai-stack;python_version>="3.10"',
+]
+jax-gpu = [
+    # below is a funny workaround for
+    # https://github.com/astral-sh/uv/issues/8601
+    'jax[cuda12]>=0.4.33;python_version>="3.10"',
+    'jax[cuda12]>=0.4.33;python_version>="3.10"',
     'flax>=0.10.0;python_version>="3.10"',
     'flax>=0.10.0;python_version>="3.10"',
     'orbax-checkpoint;python_version>="3.10"',
@@ -175,8 +200,11 @@ pin_pytorch_cpu = [
 pin_pytorch_gpu = [
   "torch==2.10.0",
 ]
-pin_jax = [
-  "jax==0.5.0;python_version>='3.10'",
+pin_jax_cpu = [
+  "jax[cpu]==0.5.0;python_version>='3.10'",
+]
+pin_jax_gpu = [
+  "jax[cuda12]==0.5.0;python_version>='3.10'",
 ]
 
 [tool.setuptools_scm]