You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
feat(jax): support spin models or add backend-aware validation #5866
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:
support spin-aware model construction, data requirements, losses, training, serialization, and inference; or
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.
Activity