deepmodeling / deepmodeling/deepmd-kit

feat(jax): train non-energy task models with matching losses

Open
#5,865 1 comment 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.