Skip to content

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

Open

Description

@njzjz-bot

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions