NNX train#3442

Draft

charlesli640 wants to merge 7 commits intoAI-Hypercomputer:mainfrom

CIeNET-International:charlesli/nnx_train

Commits on Mar 23, 2026

Add nnx_train
Charles Li
committed
Combine loss_fn for both train and evel
Charles Li
committed
Support gradient_accumulation and align to latest train.py
Charles Li
committed
Support muon
Charles Li
committed
NNX train_compile
Charles Li
committed
Fix unit test failures
Charles Li
committed
Chagne the config base.yml to set nnx flags to True
Charles Li
committed