feat(jax): train non-energy task models with matching losses

[njzjz-bot]

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

[njzjz-bot]

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