feat(jax): consume the common optimizer configuration #5867

Summary

JAX training currently constructs a fixed Optax Adam chain and does not consume the normalized top-level optimizer block. Optimizer choices and parameters accepted by the common schema can therefore be ignored.

HybridMuon for DPA4 is tracked separately in #5771; this issue covers the general optimizer configuration and dispatch contract.

Scope

• Read and validate the normalized optimizer configuration.
• Add a JAX optimizer factory for the optimizer types that JAX can support.
• Wire optimizer-specific parameters into the selected Optax transformation.
• Fail explicitly for unsupported optimizer types instead of silently using Adam.
• Preserve optimizer state through checkpoints and restarts.
• Keep learning-rate scheduling separate from optimizer selection.

Acceptance criteria

• Selecting a supported optimizer changes the transformation actually used by training.
• Supported optimizer parameters affect deterministic one-step update tests.
• Unsupported schema variants produce an actionable backend-specific error.
• Adam's current behavior remains covered by regression tests.
• Checkpoint/resume restores all optimizer state and reproduces the next update.
• The factory provides the integration point required by Add HybridMuon optimizer support to JAX and TF2 for DPA4 training #5771 without duplicating configuration parsing.

Refs #5756.

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