deepmodeling / deepmodeling/deepmd-kit

feat(jax): support spin models or add backend-aware validation

Open
#5,866 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 schema accepts the top-level `model.spin` block, while JAX model construction currently raises `NotImplementedError` when spin is present. Users therefore pass common-schema validation and fail later in a backend-specific path.

## Scope

Choose and implement one explicit JAX contract:

1. support spin-aware model construction, data requirements, losses, training, serialization, and inference; or
2. mark spin as unsupported during backend-aware validation before model construction.

If implementation is selected, reuse shared spin semantics and fixtures so JAX does not introduce a backend-specific interpretation of virtual/real atom types, force outputs, or magnetic labels.

## Acceptance criteria

- JAX configurations with `model.spin` either train and infer correctly or fail during backend-aware validation with a clear message.
- Unsupported configurations never reach the current late `NotImplementedError`.
- Serialization and configuration round-trips preserve the selected behavior.
- Tests cover a minimal spin configuration and a non-spin regression case.
- Documentation distinguishes common-schema availability from JAX support.

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.