Move the sharding axes from dimensions that need replication to batch dimensions, such that we replace an `all-gather` with an `all-to-all`.#21735

Merged

copybara-service[bot] merged 1 commit intomainfrom test_718546009

Jan 23, 2025

+40-30

Commits on Jan 23, 2025

Move the sharding axes from dimensions that need replication to batch dimensions, such that we replace an all-gather with an all-to-all.

ZixuanJiang
authored and
Google-ML-Automation
committed