Remove masks and update gmm function in moe.py#4082

NuojCheng wants to merge 4 commits into

mainAI-Hypercomputer/maxtext:mainfrom

chengnuojin-remove-maskAI-Hypercomputer/maxtext:chengnuojin-remove-maskCopy head branch name to clipboard