Gather expert/context co-sharded axes in FSDP weight prefetching (EP support) by darisoy · Pull Request #5848 · AI-Hypercomputer/maxtext

✅ All contributors are covered under a CLA with Google

See https://cla.developers.google.com/ for more info about Google's Contributor License Agreement (CLA).

ℹ️ Googlers: Go here to view more details and manage scans for this pull request.

@darisoy

The following contributors were found for this pull request:

✅ 71fc655 Author: @darisoy <34922984+darisoy@users.noreply.github.com>

(Only the first commit for a unique contributor is listed.)