maxtext icon indicating copy to clipboard operation
maxtext copied to clipboard

Introduce `gather()` and `scatter_add()` reference functions for token-routing in Mixture-of-Experts models.

Open copybara-service[bot] opened this issue 3 months ago • 2 comments

Introduce gather() and scatter_add() reference functions for token-routing in Mixture-of-Experts models.

copybara-service[bot] avatar Oct 16 '25 17:10 copybara-service[bot]

@jesselu-google do you want to include some simple moe test to cover this gather/scatter_add?

suexu1025 avatar Oct 17 '25 20:10 suexu1025

unit tests are included in an internal-only test file, let me know if you think it's worth it to replicate

jesselu-google avatar Oct 20 '25 14:10 jesselu-google

This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.

github-actions[bot] avatar Dec 08 '25 16:12 github-actions[bot]

This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it.

github-actions[bot] avatar Dec 16 '25 16:12 github-actions[bot]