Repository navigation
Restore Flax compatibility and perform real local CLI training - #2
Merged
jadenfix merged 2 commits intoOct 13, 2026
Merged
Conversation
jadenfix
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
NanoLM/Projection used nnx.List unavailable at the declared Flax minimum, and the advertised CLI only performed a forward pass while ignoring trainer.max_steps. Preserve compatible graph containers, wire SFT/DPO through real Optax updates, bound short-run warmup, stamp actual optimizer steps into fairness fingerprints, and emit training/decode receipts. Unsupported RL-rollout and scale CLI paths now fail explicitly instead of succeeding as dispatch stubs.
Validation on CPU JAX0.4.30/Flax0.10.0: full38test suite passed, including joint CTS backbone/projection/critic updates; latest targeted3CLI tests passed, checking real parameter changes, decode shape, receipt/log output, and rejecting the rollout stub. Full Ruff check/format passed. Changed-file Pyright:0errors; full-repository Pyright still has3pre-existing errors and6warnings. Fixtures and tiny random model are computational validation, not paper-quality results. Teacher APIs, task-specific on-policy rollout/reward builders, Tunix/MaxText/cloud execution remain unverified.