feat(jax): support shared_dict in multi-task training

[njzjz]

Summary

• move shared_dict preprocessing and shared-link traversal into dpmodel utilities
• reuse the common shared-link traversal from pt_expt and JAX backends
• enable JAX multi-task shared_dict preprocessing and apply shared descriptor/fitting links before optimizer setup

Tests

• ruff check .
• ruff format deepmd/dpmodel/utils/multi_task.py deepmd/jax/utils/multi_task.py deepmd/jax/entrypoints/train.py deepmd/jax/train/trainer.py deepmd/pt_expt/utils/multi_task.py deepmd/pt_expt/train/wrapper.py source/tests/jax/test_training.py
• pytest source/tests/jax/test_training.py -q
• pytest source/tests/pt_expt/test_multitask.py::TestMultiTaskSeA::test_multitask_train source/tests/pt_expt/test_multitask.py::TestMultiTaskSeAShareFit::test_multitask_train -q

Note

• ruff format --check . currently reports existing unrelated formatting drift in dpa_adapt/cli.py.

Summary by CodeRabbit

• New Features
  • Multi-task models can now declare shared_dict to reuse descriptor and fitting settings across model branches.
  • Shared-parameter “links” are created during config preprocessing and then applied when building the JAX training model.
• Bug Fixes
  • Improved validation and resume-aware merging of shared descriptor and fitting statistics during training.
• Tests
  • Added JAX regression coverage for shared_dict preprocessing, link propagation, and correct shared-stat/parameter behavior.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: efc0e511-3b7b-4ba9-9888-509e3f7ac7de

📥 Commits

Reviewing files that changed from the base of the PR and between 0b0d3ba and 5b327a7.

📒 Files selected for processing (1)

• deepmd/dpmodel/utils/multi_task.py

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

• deepmd/dpmodel/utils/multi_task.py

---

📝 Walkthrough

Walkthrough

This PR adds backend-neutral multi-task shared-parameter preprocessing and application, integrates shared_dict support into JAX training, reuses the utilities in pt_expt, and adds JAX regression tests for sharing and statistic merging.

Changes

Multi-task shared parameter support

Layer / File(s)	Summary
Common multi-task shared-link utilities deepmd/dpmodel/utils/multi_task.py	Adds shared_dict preprocessing, link traversal, descriptor accessors, atomic-model lookup, link sanitization, and probability-ratio handling.
JAX entrypoint and preprocessing wiring deepmd/jax/utils/multi_task.py, deepmd/jax/entrypoints/train.py	Adds JAX class resolution and preprocessing, tracks shared_links, and passes them into DPTrainer.
JAX trainer shared-parameter application deepmd/jax/train/trainer.py	Applies descriptor and fitting links during setup, including resume-aware sharing and weighted statistic merging.
pt_expt refactor to reuse shared helpers deepmd/pt_expt/train/wrapper.py, deepmd/pt_expt/utils/multi_task.py	Delegates shared-link application and multi-task preprocessing to common utilities.
JAX multitask training tests source/tests/jax/test_training.py	Adds coverage for preprocessing, link propagation, object sharing, and weighted statistic merges.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5397: Related multi-task shared-parameter implementation using the same descriptor and fitting sharing mechanisms.
• deepmodeling/deepmd-kit#5603: Related changes to JAX entrypoint handling for multi-task shared_dict configurations.

Suggested labels: enhancement

Suggested reviewers: wanghan-iapcm, iProzd

🚥 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 clearly summarizes the main change: adding shared_dict support to multi-task training, especially for JAX.
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.

✨ 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.}

[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 `@deepmd/jax/entrypoints/train.py`:
- Around line 109-117: The shared-task preprocessing in `TrainEntrypoint` can
leave `self.shared_links` stale because `preprocess_shared_params()` runs before
the later `init_model`/`use_pretrain_script` model replacement in the same flow.
Update the `is_multi_task` handling so the final `config["model"]` is
established before calling `preprocess_shared_params()`, or explicitly clear and
recompute `self.shared_links` after any model overwrite, ensuring the value
passed to `DPTrainer` matches the actual model.

🪄 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: 6af13361-7181-40a4-bfcf-d87b6ac5eccc

📥 Commits

Reviewing files that changed from the base of the PR and between 3d73c67 and 568b0cd.

📒 Files selected for processing (7)

• deepmd/dpmodel/utils/multi_task.py
• deepmd/jax/entrypoints/train.py
• deepmd/jax/train/trainer.py
• deepmd/jax/utils/multi_task.py
• deepmd/pt_expt/train/wrapper.py
• deepmd/pt_expt/utils/multi_task.py
• source/tests/jax/test_training.py

[codecov]

Codecov Report

❌ Patch coverage is 78.18930% with 53 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.53%. Comparing base (3d73c67) to head (5b327a7).
⚠️ Report is 15 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/dpmodel/utils/multi_task.py	75.93%	32 Missing ⚠️
deepmd/jax/train/trainer.py	75.64%	19 Missing ⚠️
deepmd/jax/entrypoints/train.py	87.50%	1 Missing ⚠️
deepmd/jax/utils/multi_task.py	93.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5740      +/-   ##
==========================================
- Coverage   81.36%   79.53%   -1.83%     
==========================================
  Files         991     1017      +26     
  Lines      111198   115900    +4702     
  Branches     4235     4277      +42     
==========================================
+ Hits        90474    92183    +1709     
- Misses      19199    22173    +2974     
- Partials     1525     1544      +19     

☔ View full report in Codecov by Harness.
📢 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.