feat: framework swap to JAX/Equinox/Optax (v2 stack) - #115

Open

thomas-schweich wants to merge 134 commits into

mainfrom

Commits on May 23, 2026

docs(jax-migration): add from-scratch parity-bar migration plan
thomas-schweich
committed
6a3fb21
docs(jax-migration): add Invocation block for /review-driven-development --resume
thomas-schweich
committed
d78c652
feat(jax-migration): [jax-migration S1] foundation — JAX deps + orientation map
thomas-schweich
committed
d7d6d54
feat(jax-migration): [jax-migration S2] JAX core — config + supernet + model + checkpoint
thomas-schweich
committed
6f6089c
feat(jax-migration): [jax-migration S3] pydantic run_config — keystone
thomas-schweich
committed
14703bc
feat(jax-migration): [jax-migration S4] MetricsLogger — torch-free JSONL writer
thomas-schweich
committed
2090ebf
feat(jax-migration): [jax-migration S5] Corpus + Lichess data
thomas-schweich
committed
c77ed1e
feat(jax-migration): [jax-migration S6] pretraining trainer — lax.scan + supernet joint loss
thomas-schweich
committed
d7cb0aa
feat(jax-migration): [jax-migration S7] adapter trainer + 8 strategies
thomas-schweich
committed
aafb65f
feat(jax-migration): [jax-migration S8] eval surface — accuracy + probes + 5 diagnostics + Elo strat + edge cases
thomas-schweich
committed
0a77ee8
feat(jax-migration): [jax-migration S11] PAD-init engine token buffers
thomas-schweich
committed
3e9702e
feat(jax-migration): [jax-migration S12] lifecycle — SIGTERM + HF push + --resume
thomas-schweich
committed
e37ac58
feat(jax-migration): [jax-migration S10] compatibility loader — v1 PyTorch → v2 JAX
thomas-schweich
committed
f1aa7da
feat(jax-migration): [jax-migration S9] sweep + lab + dashboard + wandb
thomas-schweich
committed
dc3016b
feat(jax-migration): [jax-migration S13] scripts + data tools
thomas-schweich
committed
816ba86
feat(jax-migration): [jax-migration S14] docs + deploy v2 examples
thomas-schweich
committed
12d1213
feat(jax-migration): [jax-migration S15] comprehensive tests — full suite green
thomas-schweich
committed
e20f120
feat(jax-migration): [jax-migration S16-prep] port benchmark.py + DEFERRALS + final_smoke
thomas-schweich
committed
b89d4fa
feat(jax-migration): [jax-migration S16-fix] adapter val_loss + nested v1 config + sweep CLI
thomas-schweich
committed
0840b14
chore(jax-migration): [jax-migration S16-cleanup] delete v1 GONE-BY-DESIGN modules + pyright green
thomas-schweich
committed
014b436
fix(jax-migration): [jax-migration S16-review] address multi-lane review findings
thomas-schweich
committed
6bf5dc7
fix(jax-migration): [jax-migration S16-pr-review] address PR #115 code-review comments
thomas-schweich
committed
92d618b
fix(jax-migration): [jax-migration S16-loop-r1] address loop round-1 review findings
thomas-schweich
committed
e0462dd
fix(jax-migration): [jax-migration S16-loop-r2] daemon executor + timeout/error split + dtype defense
thomas-schweich
committed
0c6fc72
fix(jax-migration): [jax-migration S16-loop-r3] pop workers from _threads_queues + explicit TimeoutError + python-int dtype norm
thomas-schweich
committed
187eb99
fix(jax-migration): [jax-migration S16-loop-r4] normalise numpy-scalar template dtypes
thomas-schweich
committed
eee69b7

Commits on May 24, 2026

Commits on May 25, 2026

Commits on May 27, 2026

Commits on May 28, 2026

Commits on May 29, 2026

Commits on May 30, 2026

Commits on May 31, 2026

Commits on Jun 1, 2026

Commits on Jun 5, 2026

fix(v2): untie output head by default — fixes single-step output collapse
thomas-schweich
committed
18a0047

Commits on Jun 6, 2026

Commits on Jun 8, 2026

Commits on Jun 9, 2026

Commits on Jun 10, 2026