deepmodeling / deepmodeling/deepmd-kit

feat(jax): support or reject common composite model variants

Open
#5,868 0 comments 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

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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.