build(jax): split pinned cpu and gpu groups

[njzjz-bot]

Problem

• CPU and CUDA CI need different pinned JAX dependency groups.
• The public jax extra should stay unchanged.

Change

• Keep the existing jax optional dependency as plain jax.
• Split the pinned groups into pin_jax_cpu and pin_jax_gpu.
• Use pin_jax_cpu in CPU jobs and pin_jax_gpu in the CUDA job.

Validation

• git diff --check
• uv pip compile pyproject.toml --group pin_jax_cpu --python-version 3.10
• uv pip compile pyproject.toml --group pin_jax_gpu --python-version 3.10

Authored by OpenClaw (model: gpt-5.5)

Summary by CodeRabbit

• Chores
  • Split JAX dependency pins into separate CPU and GPU variants.
  • Updated CI workflows to install the appropriate JAX package group for CPU or GPU runs.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR splits the unified pin_jax dependency into pin_jax_cpu (jax==0.5.0) and pin_jax_gpu (jax[cuda12]==0.5.0) in pyproject.toml. CI workflows for C++ and Python tests now use pin_jax_cpu; the CUDA workflow uses pin_jax_gpu.

Changes

JAX Dependency CPU/GPU Split

Layer / File(s)	Summary
JAX Dependency Groups pyproject.toml	Removed unified pin_jax; added pin_jax_cpu (jax==0.5.0) and pin_jax_gpu (jax[cuda12]==0.5.0).
CPU Workflow Updates .github/workflows/test_cc.yml, .github/workflows/test_python.yml	Changed pip install group from pin_jax to pin_jax_cpu in CPU-oriented CI jobs.
GPU Workflow Updates .github/workflows/test_cuda.yml	Changed pip install group from pin_jax/explicit spec to pin_jax_gpu in CUDA CI job.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested labels

build

Suggested reviewers

• wanghan-iapcm

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: splitting JAX dependency groups into separate CPU and GPU variants across workflows and pyproject.toml.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.47%. Comparing base (5bd0889) to head (2d024de).
⚠️ Report is 9 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #5436   +/-   ##
=======================================
  Coverage   82.47%   82.47%           
=======================================
  Files         825      825           
  Lines       87721    87721           
  Branches     4206     4206           
=======================================
+ Hits        72344    72345    +1     
+ Misses      14094    14093    -1     
  Partials     1283     1283           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@pyproject.toml`:
- Around line 178-183: Add explicit user-facing extras "jax-cpu" and "jax-gpu"
under [tool.deepmd_build_backend.optional-dependencies], mapping "jax-cpu" to
the pin_jax_cpu group and "jax-gpu" to the pin_jax_gpu group (so CI pinning and
user extras stay aligned); remove the duplicated entries from the existing "jax"
extra and make "jax" a single, deduplicated reference (or a thin alias) to the
CPU variant as appropriate to avoid the uv#8601 workaround duplication. Ensure
the unique identifiers pin_jax_cpu, pin_jax_gpu and the existing jax extra are
updated so that the new jax-cpu and jax-gpu extras appear in package metadata.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 358a2778-1d17-4e4e-ae3e-2d49ac2b28fa

📥 Commits

Reviewing files that changed from the base of the PR and between 6bbe50b and 2d024de.

📒 Files selected for processing (4)

• .github/workflows/test_cc.yml
• .github/workflows/test_cuda.yml
• .github/workflows/test_python.yml
• pyproject.toml

🚧 Files skipped from review as they are similar to previous changes (3)

• .github/workflows/test_python.yml
• .github/workflows/test_cuda.yml
• .github/workflows/test_cc.yml

[Copilot]

Pull request overview

Splits the CI JAX dependency pins into separate CPU and CUDA groups so CPU and CUDA workflows can each install the appropriate pinned JAX variant while keeping the public jax extra unchanged.

Changes:

• Replaced the single pin_jax dependency group with pin_jax_cpu and pin_jax_gpu in pyproject.toml.
• Updated CPU workflows to install --group pin_jax_cpu.
• Updated the CUDA workflow to install --group pin_jax_gpu (pinning jax[cuda12]).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
pyproject.toml	Splits the pinned JAX dependency group into CPU/GPU variants.
.github/workflows/test_python.yml	Uses pin_jax_cpu for CPU Python test jobs.
.github/workflows/test_cuda.yml	Uses pin_jax_gpu for CUDA jobs (pins jax[cuda12]).
.github/workflows/test_cc.yml	Uses pin_jax_cpu for CPU C++/Python dependency install step.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.