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

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