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

## 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