[JAX] Add an MoE Block (Layer) that compound router, permutation, groupedGEMM and communication#2912
Draft
tdophung wants to merge 2 commits intoNVIDIA:mainfrom
Draft
[JAX] Add an MoE Block (Layer) that compound router, permutation, groupedGEMM and communication#2912tdophung wants to merge 2 commits intoNVIDIA:mainfrom
tdophung wants to merge 2 commits intoNVIDIA:mainfrom
Commits
Commits on Apr 14, 2026
Commits on Apr 21, 2026
- committed