[JAX] Remove shard_map from MoEBlock to support quant before FSDP AG using Grouped quant+GEMM custom partitioning rules#3131
Draft
jberchtold-nvidia wants to merge 38 commits into
Commits
Commits on May 26, 2026
Commits on Jun 1, 2026
- committed
Commits on Jun 2, 2026
- committed
- committed
- committed
- committed
- committed
- committed
- committed
![pre-commit-ci[bot]](https://avatars.githubusercontent.com/in/68672?v=4&size=32){kind=avatar}
Commits on Jun 3, 2026
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- authored andcommitted
- committed
- committed
- committed
- committed
- committed
- committed
