deepmodeling / deepmodeling/deepmd-kit
feat(jax): support or reject common composite model variants
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
The common model schema exposes `frozen`, top-level `pairtab`, and `linear_ener`, but the JAX model registry does not currently provide matching constructors. These variants remain discoverable through the unified schema and fail only after JAX backend selection.
## Scope
For each variant, either:
- implement the JAX model/composition path with serialization and inference support; or
- add backend-aware validation that marks it unsupported before registry lookup.
Where implementation is practical, reuse dpmodel composition and serialization semantics rather than creating JAX-only formats.
## Acceptance criteria
- Every listed model variant has an explicit JAX support status.
- Supported variants construct, serialize, deserialize, and run deterministic inference.
- Unsupported variants fail during backend-aware validation with the variant name and backend in the error.
- Composite parameter/type-map compatibility is validated consistently.
- Tests cover each supported or rejected path and preserve existing JAX model registrations.
Refs #5756.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Contributor guide
Research direction
Start at the JAX model registry and the unified model schema handling for frozen, top-level pairtab, and linear_ener. Compare the existing dpmodel composition and serialization semantics, then inspect the backend-aware validation path and related tests. Done means each variant has an explicit support status, supported paths round-trip and infer deterministically, unsupported paths report the variant and backend, and compatibility tests pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend-api-design, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100