Repository navigation
feat(jax): train non-energy task models with matching losses #5865
Description
Summary
JAX registers DOS, property, dipole, and polar fitting/model classes, but the training loss builder only accepts loss.type: ener. As a result, these model types can be constructed but cannot be trained through the JAX backend.
Scope
- Add JAX loss dispatch for DOS, property, dipole, and polar/tensor tasks.
- Reuse backend-neutral loss definitions and output conventions where possible.
- Wire the required labels, masks, atom/frame reductions, and model outputs into the JAX train step.
- Preserve single-task and multi-task behavior.
- Reject incompatible model/loss combinations during configuration or trainer construction with actionable errors.
Acceptance criteria
- Representative DOS, property, dipole, and polar configurations reach a JAX training step.
- Loss values and gradients agree with the shared loss definitions on deterministic fixtures.
- Missing labels and incompatible model/loss pairs fail before the main training loop.
- Checkpoint save/resume works for each newly supported task.
- Regression tests preserve existing energy-training behavior.
Refs #5756.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Activity
njzjz-bot commentedon Jul 29, 2026
Tracking note for the documentation work:\n\n- #5915 adds backend-support labels and support matrices to the argument documentation.\n- #5929 documents JAX non-energy fitting/model availability separately from the current energy-only training-loss support.\n- Both PRs are currently open and document the limitation; they do not implement the non-energy losses tracked here.\n\nWhen #5915 and/or #5929 are merged, please re-check this issue against the latest master and update its description and references to reflect the merged documentation. The implementation scope should remain open until the non-energy loss dispatch and the acceptance criteria above are completed.\n\nCoding agent: Codex\nCodex version: codex-cli 0.144.6\nModel: gpt-5.6-sol\nReasoning effort: xhigh