Add unified trainer (PC + backprop + autoregressive)

[aprotsenko24]

What

One self-contained, algorithm-parameterized trainer that consolidates the three legacy trainers (train.py / train_backprop.py / train_autoregressive.py) behind a single API:

train(params, structure, loader, optimizer, config, key,
      algorithm="pc"|"backprop", autoregressive=False, ...)
evaluate(..., algorithm=..., autoregressive=False, ...)

Design

• The two learning rules differ in exactly three funnels — _forward / compute_gradients / _eval_step. Everything else (batch conversion, clamping, optimizer update, the epoch/batch driver, sharding, eval reduction) is shared.
• Autoregressive is an orthogonal boolean flag, not a third algorithm: its only footprint is a causal-mask clamp + loss/perplexity reporting; it works with either algorithm.
• The gradient math is called verbatim from fabricpc.core (compute_local_weight_gradients, run_inference, initialize_graph_state) — not reimplemented.
• The three legacy trainer files are left untouched as the parity baseline (this is a new, standalone file).
• algorithm / autoregressive / use_causal_mask / metric_name are explicit signature parameters (not config keys); config holds only num_epochs.

Numerical parity (single device)

• scripts/_parity_check_unified.py — PC & backprop, train + eval vs the legacy trainers, exact to 1e-6.
• scripts/_parity_check_ar.py — AR-PC / AR-backprop vs legacy: eager gradients bit-identical (atol=0); jitted params ≤2e-3 (a float artifact — the legacy AR step computes an unused cross-entropy inside its jit, which the unified version omits); includes a ragged-batch eval assertion.
• tests/test_unified_trainer.py — PC & backprop train + eval parity, 1e-12 (pytest gate).

All green on CPU (JAX_PLATFORMS=cpu).

Known AR caveat (do not "fix")

AR eager gradients are bit-identical to the legacy formula; only the jitted params drift ≤2e-3, purely because the legacy AR step computes a dead cross-entropy inside the jit (changes XLA float-scheduling; Adam amplifies it). The unified step is cleaner and run-to-run deterministic.

Deferred (separate atomic step, after sign-off)

Wiring into training/__init__.py and converting the three legacy files to deprecation shims that delegate to this module.