feat(jax): apply training.mixed_precision to the JAX train step

## Summary

The JAX trainer reads `training.mixed_precision`, but the configuration is not applied to model execution, loss computation, optimizer updates, or checkpoint behavior.

## Scope

- Define compute, parameter, output, and reduction dtypes from the normalized mixed-precision configuration.
- Apply the policy consistently to the JAX model and train step.
- Add loss scaling or non-finite update handling when required by the selected low-precision mode.
- Preserve numerically sensitive reductions in an appropriate dtype.
- Serialize enough policy and optimizer state for deterministic restart.
- Reject unsupported dtype/device combinations before training.

## Acceptance criteria

- Supported mixed-precision configurations change the actual computation dtype.
- Energy, force, loss, and gradient results remain within documented tolerances of full precision.
- Non-finite gradients do not silently corrupt checkpoints.
- Checkpoint/resume preserves the policy and reproduces the next update.
- Tests cover at least one supported low-precision mode, unsupported configurations, and full-precision regression behavior.

Refs #5756.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh