feat(jax): support epoch-based and multi-task training horizons #5869

Summary

The common training schema accepts training.numb_epoch and its aliases, plus multi-task training.num_epoch_dict, while the JAX path still relies on a resolved numb_steps value. Epoch-only inputs therefore lack a complete JAX execution contract.

Scope

• Resolve single-task numb_epoch after dataset and batch-size configuration determine steps per epoch.
• Resolve multi-task num_epoch_dict into per-task sampling/horizon semantics and a deterministic total step count.
• Define epoch transitions, reshuffling, and restart state.
• Ensure all distributed ranks use the same optimizer-step count.
• Preserve explicit numb_steps behavior.

A shared backend-neutral horizon resolver should be used where feasible so the common schema has consistent semantics across backends.

Acceptance criteria

• A normalized JAX input containing only an epoch count trains without requiring numb_steps.
• Multi-task epoch dictionaries produce documented probabilities and total steps.
• Dataset filtering and batch-size resolution occur before epoch length is computed.
• Restart resumes at the same epoch and intra-epoch position.
• Single-rank and distributed tests cover deterministic reshuffling and uneven dataset tails.
• Existing fixed-step training is unchanged.

Refs #5756.

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