deepmodeling / deepmodeling/deepmd-kit
feat(jax): train non-energy task models with matching losses
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
JAX registers DOS, property, dipole, and polar fitting/model classes, but the training loss builder only accepts `loss.type: ener`. As a result, these model types can be constructed but cannot be trained through the JAX backend.
## Scope
- Add JAX loss dispatch for DOS, property, dipole, and polar/tensor tasks.
- Reuse backend-neutral loss definitions and output conventions where possible.
- Wire the required labels, masks, atom/frame reductions, and model outputs into the JAX train step.
- Preserve single-task and multi-task behavior.
- Reject incompatible model/loss combinations during configuration or trainer construction with actionable errors.
## Acceptance criteria
- Representative DOS, property, dipole, and polar configurations reach a JAX training step.
- Loss values and gradients agree with the shared loss definitions on deterministic fixtures.
- Missing labels and incompatible model/loss pairs fail before the main training loop.
- Checkpoint save/resume works for each newly supported task.
- Regression tests preserve existing energy-training behavior.
Refs #5756.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Contributor guide
Assessment
This issue has not been assessed yet.